Add arm link rewards

This commit is contained in:
2026-03-19 09:08:57 -04:00
parent 6ca671dce5
commit 5df147b0b1
2 changed files with 120 additions and 122 deletions

View File

@@ -14,88 +14,98 @@ from rl_game.get_up.env.t1_env import T1SceneCfg
import isaaclab.envs.mdp as mdp
# --- 1. 修正后的自定义 MDP 逻辑函数 ---
# --- 1. 自定义 MDP 逻辑函数 ---
def standing_with_feet_reward(
env: ManagerBasedRLEnv,
min_head_height: float,
min_pelvis_height: float,
sensor_cfg: SceneEntityCfg,
force_threshold: float = 30.0
) -> torch.Tensor:
"""
【双高度条件奖励】:只有脚踩地,且头和躯干同时达到高度,才给予高度奖励。
"""
# 1. 获取脚部触地力判定
contact_sensor = env.scene.sensors.get(sensor_cfg.name)
foot_forces_z = torch.sum(contact_sensor.data.net_forces_w[:, :, 2], dim=-1)
is_feet_on_ground = foot_forces_z > force_threshold
# 2. 获取头部和躯干索引并提取高度
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]
# 3. 计算高度达标度 (0.0 - 1.0)
head_reward = torch.clamp(current_head_h / min_head_height, max=1.2)
pelvis_reward = torch.clamp(current_pelvis_h / min_pelvis_height, max=1.0)
# 综合高度奖励(取平均值)
combined_height_reward = (head_reward + pelvis_reward) / 2.0
# 4. 逻辑门:脚不着地,奖励为 0脚着地后根据高度给分
return torch.where(is_feet_on_ground, combined_height_reward, torch.zeros_like(combined_height_reward))
def arm_push_up_reward(
env: ManagerBasedRLEnv,
sensor_cfg: SceneEntityCfg,
height_threshold: float = 0.6
) -> torch.Tensor:
"""手臂撑地奖励:辅助机器人从趴/躺状态利用手臂反作用力起身"""
contact_sensor = env.scene.sensors.get(sensor_cfg.name)
if contact_sensor is None: return torch.zeros(env.num_envs, device=env.device)
arm_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
max_arm_force = torch.max(arm_forces_z, dim=-1)[0]
# 当躯干还很低时,鼓励手撑地
pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
current_height = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
pushing_reward = torch.clamp(max_arm_force, max=200.0) / 100.0
return torch.where(current_height < height_threshold, pushing_reward, torch.zeros_like(pushing_reward))
def is_standing_still(
env: ManagerBasedRLEnv,
minimum_height: float,
min_head_height: float,
min_pelvis_height: float,
max_angle_error: float,
standing_time: float,
velocity_threshold: float = 0.2
velocity_threshold: float = 0.15
) -> torch.Tensor:
"""
判定判定逻辑:高度达标 + 躯干垂直 + 几乎静止 + 维持时间。
增加了速度判定,彻底杜绝起跳瞬间触发。
"""
# 1. 获取 Body 索引
"""判定逻辑:双高度达标 + 躯干垂直 + 全身静止"""
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]
# 2. 状态量获取
current_head_height = env.scene["robot"].data.body_state_w[:, head_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)
root_ang_vel_norm = torch.norm(env.scene["robot"].data.root_ang_vel_w, dim=-1)
# 3. 综合判定(这里不强制检查力,改用更稳健的速度限制)
# 判定条件:头够高 且 盆骨够高 且 垂直误差小 且 速度低
is_stable_now = (
(current_head_height > minimum_height) &
(current_head_h > min_head_height) &
(current_pelvis_h > min_pelvis_height) &
(gravity_error < max_angle_error) &
(root_vel_norm < velocity_threshold) &
(root_ang_vel_norm < velocity_threshold * 2.0)
(root_vel_norm < velocity_threshold)
)
# 4. 计时器
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"])
)
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
def get_success_reward(env: ManagerBasedRLEnv, term_keys: str) -> torch.Tensor:
return env.termination_manager.get_term(term_keys)
def root_vel_z_l2_local(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
# 严厉惩罚 Z 轴正向速度(向上窜)
vel_z = env.scene[asset_cfg.name].data.root_lin_vel_w[:, 2]
return torch.square(torch.clamp(vel_z, min=0.0))
def feet_airtime_penalty_local(
env: ManagerBasedRLEnv,
sensor_cfg: SceneEntityCfg,
threshold: float = 1.0
) -> torch.Tensor:
"""
自定义滞空惩罚逻辑:
如果脚部的垂直合力小于阈值,说明脚离地了。
返回一个 Tensor离地时为 1.0,着地时为 0.0。
"""
# 1. 获取传感器对象
contact_sensor = env.scene.sensors.get(sensor_cfg.name)
if contact_sensor is None:
# 如果没搜到传感器,返回全 0防止程序崩溃
return torch.zeros(env.num_envs, device=env.device)
# 2. 获取触地力 (num_envs, num_bodies_in_sensor, 3)
# 我们取所有被监测 Body (左右脚) 的 Z 轴推力
# 如果所有脚的力都小于 threshold判定为“完全腾空”
foot_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
is_in_air = torch.all(foot_forces_z < threshold, dim=-1)
return is_in_air.float()
# --- 2. 配置类定义 ---
# --- 2. 配置类 ---
T1_JOINT_NAMES = [
'Left_Hip_Pitch', 'Right_Hip_Pitch', 'Left_Hip_Roll', 'Right_Hip_Roll',
@@ -106,22 +116,11 @@ T1_JOINT_NAMES = [
@configclass
class T1ObservationCfg:
"""观察值空间配置:严格对应你的 Robot 基类数据结构"""
@configclass
class PolicyCfg(ObsGroup):
concatenate_terms = True
enable_corruption = False
# --- 状态量 (对应你的 Robot 类属性) ---
# 1. 基体线速度 (accelerometer 相关的速度项)
base_lin_vel = ObsTerm(func=mdp.base_lin_vel)
# 2. 角速度 (对应你的 gyroscope 属性: degrees/s -> IsaacLab 默认为 rad/s)
base_ang_vel = ObsTerm(func=mdp.base_ang_vel)
# 3. 重力投影 (对应 global_orientation_euler/quat 相关的姿态感知)
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)})
@@ -134,7 +133,7 @@ class T1ObservationCfg:
@configclass
class T1EventCfg:
"""重置时的随机姿态:背躺、趴着、侧躺"""
"""随机初始化:支持趴、躺、侧身"""
reset_robot_rotation = EventTerm(
func=mdp.reset_root_state_uniform,
params={
@@ -166,78 +165,74 @@ class T1ActionCfg:
@configclass
class T1GetUpRewardCfg:
# 1. 姿态奖 (核心引导)
# 1. 姿态基础奖 (引导身体变正)
upright = RewTerm(func=mdp.flat_orientation_l2, weight=5.0)
# 2. 只有脚着地时才给的高度奖(模拟逻辑)
# 修正:直接使用 root_height 配合强力的速度惩罚
height_tracking = RewTerm(
func=mdp.root_height_below_minimum,
# 2. 【条件高度奖】:双高度判定(头+盆骨),且必须脚踩地
height_with_feet = RewTerm(
func=standing_with_feet_reward,
weight=25.0, # 作为核心引导,增加权重
params={
"min_head_height": 1.10,
"min_pelvis_height": 0.65,
"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]),
"force_threshold": 30.0
}
)
# 3. 手臂撑地奖:辅助脱离地面阶段
arm_push_support = RewTerm(
func=arm_push_up_reward,
weight=2.0,
params={"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_hand_link"])}
)
# 4. 惩罚项
undesired_contacts = RewTerm(
func=mdp.undesired_contacts,
weight=-5.0,
params={
"minimum_height": 1.05,
"asset_cfg": SceneEntityCfg("robot", body_names="H2"),
"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["H2", "Trunk", ".*_Left", ".*_Right"]),
# 注意:这里我们排除了脚部,惩罚大腿/小腿/躯干着地
"threshold": 1.0
}
)
# 3. 抑制跳跃:严厉惩罚向上窜的速度
root_vel_z_penalty = RewTerm(
func=root_vel_z_l2_local,
weight=-15.0, # 增大负权重
params={"asset_cfg": SceneEntityCfg("robot")}
)
# 4. 抑制滞空 (Airtime Penalty)
# 如果脚离开地面,按时间扣分
feet_airtime = RewTerm(
func=feet_airtime_penalty_local,
weight=-15.0,
params={
"sensor_cfg": SceneEntityCfg("feet_contact_sensor"),
"threshold": 0.2, # 超过 0.2s 离地就开始扣分
}
)
# 5. 动作平滑 (非常重要:解决视频中的高频抖动)
action_rate = RewTerm(func=mdp.action_rate_l2, weight=-1.0)
joint_vel = RewTerm(func=mdp.joint_vel_l2, weight=-0.1, params={"asset_cfg": SceneEntityCfg("robot")})
# 6. 时间惩罚:逼迫它快点站稳
time_penalty = RewTerm(func=mdp.is_alive, weight=-0.5)
# 7. 终极奖励
# 5. 成功终极大奖
is_success = RewTerm(
func=get_success_reward,
weight=200.0, # 成功一次给大奖,但判定条件极严
params={"term_keys": "standing_success"}
func=lambda env, keys: env.termination_manager.get_term(keys),
weight=1000.0,
params={"keys": "standing_success"}
)
@configclass
class T1GetUpTerminationsCfg:
time_out = DoneTerm(func=mdp.time_out)
# 严格的失败判定:躯干倾斜超过 45 度重置,不让它在地上滚
# 失败判定:躯干倾斜超过 45 度重置
base_crash = DoneTerm(func=mdp.bad_orientation, params={"limit_angle": 0.785})
# 严格的成功判定
# 成功判定:双高度 + 稳
standing_success = DoneTerm(
func=is_standing_still,
params={
"minimum_height": 1.05, # H2 高度
"max_angle_error": 0.15, # 极小角度误差(约 8.6 度)
"standing_time": 0.8, # 必须保持 0.8 秒(起跳不可能在空中停这么久)
"velocity_threshold": 0.15 # 速度必须极低
"min_head_height": 1.05,
"min_pelvis_height": 0.75,
"max_angle_error": 0.15,
"standing_time": 0.8,
"velocity_threshold": 0.15
}
)
@configclass
class T1EnvCfg(ManagerBasedRLEnvCfg):
scene = T1SceneCfg(num_envs=16384, env_spacing=2.5) # 建议先用 4096 验证逻辑
scene = T1SceneCfg(num_envs=16384, env_spacing=2.5) # 5090 性能全开
def __post_init__(self):
super().__post_init__()
self.scene.robot.init_state.pos = (0.0, 0.0, 0.4) # 初始稍微抬高一点点,防止卡地
self.scene.robot.init_state.pos = (0.0, 0.0, 0.2)
observations = T1ObservationCfg()
rewards = T1GetUpRewardCfg()
@@ -245,5 +240,5 @@ class T1EnvCfg(ManagerBasedRLEnvCfg):
events = T1EventCfg()
actions = T1ActionCfg()
episode_length_s = 5.0 # 缩短时长,增加训练效率
episode_length_s = 6.0
decimation = 4