320 lines
12 KiB
Python
320 lines
12 KiB
Python
import torch
|
||
import random
|
||
import numpy as np
|
||
import isaaclab.envs.mdp as mdp
|
||
from isaaclab.assets import ArticulationCfg
|
||
from isaaclab.envs import ManagerBasedRLEnvCfg, ManagerBasedRLEnv
|
||
from isaaclab.managers import ObservationGroupCfg as ObsGroup
|
||
from isaaclab.managers import ObservationTermCfg as ObsTerm
|
||
from isaaclab.managers import RewardTermCfg as RewTerm
|
||
from isaaclab.managers import TerminationTermCfg as DoneTerm
|
||
from isaaclab.managers import EventTermCfg as EventTerm
|
||
from isaaclab.envs.mdp import JointPositionActionCfg
|
||
from isaaclab.managers import SceneEntityCfg
|
||
from isaaclab.utils import configclass
|
||
from rl_game.get_up.env.t1_env import T1SceneCfg
|
||
|
||
|
||
# ==========================================
|
||
# 1. 核心逻辑区：非线性加法引导与惩罚机制
|
||
# ==========================================
|
||
|
||
def smooth_additive_getup_reward(
|
||
env: ManagerBasedRLEnv,
|
||
min_head_height: float,
|
||
min_pelvis_height: float,
|
||
foot_sensor_cfg: SceneEntityCfg,
|
||
arm_sensor_cfg: SceneEntityCfg,
|
||
) -> torch.Tensor:
|
||
"""
|
||
指数级纯加法平滑奖励 (完全修复版)：
|
||
融合了头部高度防驼背，以及脚底受力防跪地。
|
||
"""
|
||
head_idx, _ = env.scene["robot"].find_bodies("H2")
|
||
pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
|
||
foot_indices, _ = env.scene["robot"].find_bodies(".*_foot_link")
|
||
|
||
head_h = env.scene["robot"].data.body_state_w[:, head_idx[0], 2]
|
||
pelvis_pos = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], :3]
|
||
pelvis_h = pelvis_pos[:, 2]
|
||
|
||
foot_pos = env.scene["robot"].data.body_state_w[:, foot_indices, :3]
|
||
feet_center_xy = torch.mean(foot_pos[:, :, :2], dim=1)
|
||
pelvis_xy = pelvis_pos[:, :2]
|
||
|
||
projected_gravity = env.scene["robot"].data.projected_gravity_b
|
||
|
||
# --- 基础状态比率 ---
|
||
upright_ratio = torch.clamp(1.0 - torch.norm(projected_gravity[:, :2], dim=-1), min=0.0, max=1.0)
|
||
raw_height_ratio = torch.clamp(pelvis_h / min_pelvis_height, min=0.0, max=1.0)
|
||
|
||
# 🌟 修复 1：把头部高度加回来，防驼背
|
||
head_ratio = torch.clamp(head_h / min_head_height, min=0.0, max=1.0)
|
||
|
||
# 指数级高度比例
|
||
height_ratio_sq = torch.square(raw_height_ratio)
|
||
|
||
# 1. 核心主线分：加入 head_ratio，促使它把上半身彻底挺直
|
||
core_reward = (height_ratio_sq * 3.0) + (upright_ratio * 1.0) + (head_ratio * 1.0)
|
||
|
||
# --- 辅助过渡分 ---
|
||
foot_sensor = env.scene.sensors.get(foot_sensor_cfg.name)
|
||
arm_sensor = env.scene.sensors.get(arm_sensor_cfg.name)
|
||
foot_force_z = torch.clamp(torch.sum(foot_sensor.data.net_forces_w[:, :, 2], dim=-1), min=0.0)
|
||
arm_force_z = torch.clamp(torch.sum(arm_sensor.data.net_forces_w[:, :, 2], dim=-1), min=0.0)
|
||
|
||
feet_to_pelvis_dist = torch.norm(feet_center_xy - pelvis_xy, dim=-1)
|
||
tuck_legs_reward = torch.exp(-3.0 * feet_to_pelvis_dist)
|
||
|
||
arm_force_capped = torch.clamp(arm_force_z, min=0.0, max=200.0)
|
||
arm_push_reward = arm_force_capped / 200.0
|
||
|
||
# 🌟 修复 2：把脚底受力加回来，只要脚底板踩实了就给分，引导它脱离跪姿
|
||
foot_contact_reward = (foot_force_z > 10.0).float()
|
||
|
||
ground_factor = torch.clamp(1.0 - raw_height_ratio, min=0.0, max=1.0)
|
||
|
||
# 将脚部接触分加入辅助奖励中
|
||
aux_reward = ground_factor * (tuck_legs_reward * 0.5 + arm_push_reward * 0.5) + (foot_contact_reward * 0.5)
|
||
|
||
return core_reward + aux_reward
|
||
|
||
|
||
def ground_farming_timeout(
|
||
env: ManagerBasedRLEnv,
|
||
max_time: float,
|
||
height_threshold: float
|
||
) -> torch.Tensor:
|
||
pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
|
||
pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
|
||
episode_time = env.episode_length_buf * env.step_dt
|
||
return ((episode_time > max_time) & (pelvis_h < height_threshold)).bool()
|
||
|
||
|
||
def root_height_below_minimum(env: ManagerBasedRLEnv, minimum_height: float) -> torch.Tensor:
|
||
pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
|
||
pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
|
||
return (pelvis_h < minimum_height).bool()
|
||
|
||
|
||
def is_standing_still(
|
||
env: ManagerBasedRLEnv,
|
||
min_head_height: float,
|
||
min_pelvis_height: float,
|
||
max_angle_error: float,
|
||
standing_time: float,
|
||
velocity_threshold: float = 0.15
|
||
) -> torch.Tensor:
|
||
head_idx, _ = env.scene["robot"].find_bodies("H2")
|
||
pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
|
||
current_head_h = env.scene["robot"].data.body_state_w[:, head_idx[0], 2]
|
||
current_pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
|
||
gravity_error = torch.norm(env.scene["robot"].data.projected_gravity_b[:, :2], dim=-1)
|
||
root_vel_norm = torch.norm(env.scene["robot"].data.root_lin_vel_w, dim=-1)
|
||
|
||
is_stable_now = (
|
||
(current_head_h > min_head_height) &
|
||
(current_pelvis_h > min_pelvis_height) &
|
||
(gravity_error < max_angle_error) &
|
||
(root_vel_norm < velocity_threshold)
|
||
)
|
||
|
||
if "stable_timer" not in env.extras:
|
||
env.extras["stable_timer"] = torch.zeros(env.num_envs, device=env.device)
|
||
dt = env.physics_dt * env.cfg.decimation
|
||
env.extras["stable_timer"] = torch.where(is_stable_now, env.extras["stable_timer"] + dt,
|
||
torch.zeros_like(env.extras["stable_timer"]))
|
||
return (env.extras["stable_timer"] > standing_time).bool()
|
||
|
||
|
||
def anti_flying_penalty(
|
||
env: ManagerBasedRLEnv,
|
||
foot_sensor_cfg: SceneEntityCfg,
|
||
arm_sensor_cfg: SceneEntityCfg,
|
||
) -> torch.Tensor:
|
||
pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
|
||
pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
|
||
|
||
foot_sensor = env.scene.sensors.get(foot_sensor_cfg.name)
|
||
arm_sensor = env.scene.sensors.get(arm_sensor_cfg.name)
|
||
|
||
foot_force_z = torch.clamp(torch.sum(foot_sensor.data.net_forces_w[:, :, 2], dim=-1), min=0.0)
|
||
arm_force_z = torch.clamp(torch.sum(arm_sensor.data.net_forces_w[:, :, 2], dim=-1), min=0.0)
|
||
|
||
# 必须双手双脚都悬空 (< 10N)，才算是在天上飞
|
||
is_flying = ((pelvis_h > 0.4) & (foot_force_z < 10.0) & (arm_force_z < 10.0)).float()
|
||
return is_flying
|
||
|
||
|
||
def base_ang_vel_penalty(env: ManagerBasedRLEnv) -> torch.Tensor:
|
||
ang_vel = env.scene["robot"].data.root_ang_vel_w
|
||
return torch.sum(torch.square(ang_vel), dim=-1)
|
||
|
||
|
||
def anti_kneeling_penalty(env: ManagerBasedRLEnv) -> torch.Tensor:
|
||
"""
|
||
🛠️ 修复版防跪地机制：
|
||
使用正则匹配小腿 "Shank.*"。一旦骨盆抬离地面试图站立时，严格惩罚小腿碰地！
|
||
"""
|
||
shank_indices, _ = env.scene["robot"].find_bodies("Shank.*")
|
||
shank_z = env.scene["robot"].data.body_state_w[:, shank_indices, 2]
|
||
|
||
# 判断是否有任何小腿刚体高度低于 0.15 米
|
||
is_kneeling = torch.any(shank_z < 0.15, dim=-1).float()
|
||
|
||
pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
|
||
pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
|
||
|
||
return is_kneeling * (pelvis_h > 0.35).float()
|
||
|
||
|
||
# ==========================================
|
||
# 2. 配置类区
|
||
# ==========================================
|
||
|
||
T1_JOINT_NAMES = [
|
||
'AAHead_yaw', 'Head_pitch',
|
||
'Left_Shoulder_Pitch', 'Left_Shoulder_Roll', 'Left_Elbow_Pitch', 'Left_Elbow_Yaw',
|
||
'Right_Shoulder_Pitch', 'Right_Shoulder_Roll', 'Right_Elbow_Pitch', 'Right_Elbow_Yaw',
|
||
'Waist',
|
||
'Left_Hip_Pitch', 'Right_Hip_Pitch', 'Left_Hip_Roll', 'Right_Hip_Roll',
|
||
'Left_Hip_Yaw', 'Right_Hip_Yaw', 'Left_Knee_Pitch', 'Right_Knee_Pitch',
|
||
'Left_Ankle_Pitch', 'Right_Ankle_Pitch', 'Left_Ankle_Roll', 'Right_Ankle_Roll'
|
||
]
|
||
|
||
|
||
@configclass
|
||
class T1ObservationCfg:
|
||
@configclass
|
||
class PolicyCfg(ObsGroup):
|
||
concatenate_terms = True
|
||
base_lin_vel = ObsTerm(func=mdp.base_lin_vel)
|
||
base_ang_vel = ObsTerm(func=mdp.base_ang_vel)
|
||
projected_gravity = ObsTerm(func=mdp.projected_gravity)
|
||
joint_pos = ObsTerm(func=mdp.joint_pos_rel,
|
||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=T1_JOINT_NAMES)})
|
||
joint_vel = ObsTerm(func=mdp.joint_vel_rel,
|
||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=T1_JOINT_NAMES)})
|
||
actions = ObsTerm(func=mdp.last_action)
|
||
|
||
policy = PolicyCfg()
|
||
|
||
|
||
@configclass
|
||
class T1EventCfg:
|
||
reset_robot_rotation = EventTerm(
|
||
func=mdp.reset_root_state_uniform,
|
||
params={
|
||
"asset_cfg": SceneEntityCfg("robot"),
|
||
"pose_range": {
|
||
"roll": (-3.14, 3.14),
|
||
"pitch": (-3.14, 3.14),
|
||
"yaw": (-3.14, 3.14),
|
||
"x": (0.0, 0.0),
|
||
"y": (0.0, 0.0),
|
||
"z": (0.45, 0.65),
|
||
},
|
||
"velocity_range": {},
|
||
},
|
||
mode="reset",
|
||
)
|
||
|
||
|
||
@configclass
|
||
class T1ActionCfg:
|
||
arm_action = JointPositionActionCfg(
|
||
asset_name="robot",
|
||
joint_names=[
|
||
'Left_Shoulder_Pitch', 'Left_Shoulder_Roll', 'Left_Elbow_Pitch', 'Left_Elbow_Yaw',
|
||
'Right_Shoulder_Pitch', 'Right_Shoulder_Roll', 'Right_Elbow_Pitch', 'Right_Elbow_Yaw'
|
||
],
|
||
scale=0.8, use_default_offset=True
|
||
)
|
||
torso_action = JointPositionActionCfg(
|
||
asset_name="robot",
|
||
joint_names=['Waist', 'AAHead_yaw', 'Head_pitch'],
|
||
scale=0.8, use_default_offset=True
|
||
)
|
||
leg_action = JointPositionActionCfg(
|
||
asset_name="robot",
|
||
joint_names=[
|
||
'Left_Hip_Pitch', 'Right_Hip_Pitch', 'Left_Hip_Roll', 'Right_Hip_Roll',
|
||
'Left_Hip_Yaw', 'Right_Hip_Yaw', 'Left_Knee_Pitch', 'Right_Knee_Pitch',
|
||
'Left_Ankle_Pitch', 'Right_Ankle_Pitch', 'Left_Ankle_Roll', 'Right_Ankle_Roll'
|
||
],
|
||
scale=0.8, use_default_offset=True
|
||
)
|
||
|
||
|
||
@configclass
|
||
class T1GetUpRewardCfg:
|
||
smooth_getup = RewTerm(
|
||
func=smooth_additive_getup_reward, weight=15.0,
|
||
params={
|
||
"min_head_height": 1.08, "min_pelvis_height": 0.72,
|
||
"foot_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]),
|
||
"arm_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["A[LR][23]", ".*_hand_link"])
|
||
}
|
||
)
|
||
|
||
anti_fly = RewTerm(
|
||
func=anti_flying_penalty,
|
||
weight=-5.0,
|
||
params={
|
||
"foot_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]),
|
||
"arm_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["A[LR][23]", ".*_hand_link"])
|
||
}
|
||
)
|
||
|
||
# 🌟 新增：重罚小腿触地的跪姿
|
||
anti_kneeling = RewTerm(func=anti_kneeling_penalty, weight=-3.0)
|
||
|
||
base_ang_vel = RewTerm(func=base_ang_vel_penalty, weight=-0.05)
|
||
action_rate = RewTerm(func=mdp.action_rate_l2, weight=-0.05)
|
||
joint_vel = RewTerm(func=mdp.joint_vel_l2, weight=-0.01)
|
||
|
||
# 🌟 新增：大大降低成功的门槛，诱使它敢于站立
|
||
is_success_bonus = RewTerm(
|
||
func=is_standing_still,
|
||
weight=100.0,
|
||
params={
|
||
"min_head_height": 1.05,
|
||
"min_pelvis_height": 0.75,
|
||
"max_angle_error": 0.4, # 允许一定的弯腰
|
||
"standing_time": 0.2, # 仅需保持 0.1 秒即可拿走巨额赏金！
|
||
"velocity_threshold": 0.8 # 允许身体轻微晃动
|
||
}
|
||
)
|
||
|
||
|
||
@configclass
|
||
class T1GetUpTerminationsCfg:
|
||
time_out = DoneTerm(func=mdp.time_out)
|
||
anti_farming = DoneTerm(
|
||
func=ground_farming_timeout,
|
||
params={"max_time": 5.0, "height_threshold": 0.3}
|
||
)
|
||
base_height = DoneTerm(func=root_height_below_minimum, params={"minimum_height": -0.2})
|
||
illegal_contact = DoneTerm(
|
||
func=mdp.illegal_contact,
|
||
params={"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["Trunk"]), "threshold": 5000.0}
|
||
)
|
||
standing_success = DoneTerm(
|
||
func=is_standing_still,
|
||
params={
|
||
"min_head_height": 1.05, "min_pelvis_height": 0.75,
|
||
"max_angle_error": 0.2, "standing_time": 0.5, "velocity_threshold": 0.2
|
||
}
|
||
)
|
||
|
||
|
||
@configclass
|
||
class T1EnvCfg(ManagerBasedRLEnvCfg):
|
||
scene = T1SceneCfg(num_envs=8192, env_spacing=2.5)
|
||
observations = T1ObservationCfg()
|
||
rewards = T1GetUpRewardCfg()
|
||
terminations = T1GetUpTerminationsCfg()
|
||
events = T1EventCfg()
|
||
actions = T1ActionCfg()
|
||
episode_length_s = 10.0
|
||
decimation = 4 |