307 lines
11 KiB
Python
307 lines
11 KiB
Python
import random
|
|
import numpy
|
|
import numpy as np
|
|
import torch
|
|
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
|
|
import isaaclab.envs.mdp as mdp
|
|
|
|
|
|
# --- 1. 自定义 MDP 逻辑函数 ---
|
|
|
|
def standing_with_feet_reward(
|
|
env: ManagerBasedRLEnv,
|
|
min_head_height: float,
|
|
min_pelvis_height: float,
|
|
sensor_cfg: SceneEntityCfg,
|
|
force_threshold: float = 20.0,
|
|
max_v_z: float = 0.5
|
|
) -> torch.Tensor:
|
|
|
|
head_idx, _ = env.scene["robot"].find_bodies("H2")
|
|
pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
|
|
|
|
curr_head_h = torch.clamp(env.scene["robot"].data.body_state_w[:, head_idx[0], 2], 0.0, 2.0)
|
|
curr_pelvis_h = torch.clamp(env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2], 0.0, 2.0)
|
|
|
|
head_score = torch.tanh(curr_head_h / (min_head_height + 1e-6) * 2.0)
|
|
pelvis_score = torch.tanh(curr_pelvis_h / (min_pelvis_height + 1e-6) * 2.0)
|
|
height_reward = (head_score + pelvis_score) / 2.0
|
|
|
|
contact_sensor = env.scene.sensors.get(sensor_cfg.name)
|
|
foot_forces_z = torch.sum(contact_sensor.data.net_forces_w[:, :, 2], dim=-1)
|
|
force_weight = torch.sigmoid((foot_forces_z - force_threshold) / 5.0)
|
|
root_vel_z = env.scene["robot"].data.root_lin_vel_w[:, 2]
|
|
vel_penalty = torch.exp(-2.0 * torch.clamp(torch.abs(root_vel_z) - max_v_z, min=0.0))
|
|
|
|
influence_weight = torch.clamp((curr_pelvis_h - 0.2) / 0.4, min=0.0, max=1.0)
|
|
combined_reward = height_reward * ((1.0 - influence_weight) + influence_weight * force_weight * vel_penalty)
|
|
|
|
return combined_reward
|
|
|
|
|
|
def universal_arm_support_reward(
|
|
env: ManagerBasedRLEnv,
|
|
sensor_cfg: SceneEntityCfg,
|
|
height_threshold: float = 0.60,
|
|
min_force: float = 2.0
|
|
) -> torch.Tensor:
|
|
"""
|
|
通用手臂支撑奖励：同时支持仰卧起坐支撑和俯卧撑起。
|
|
逻辑：只要手臂有向上的推力，且身体正在向上移动，就给奖。
|
|
"""
|
|
# 1. 获取传感器数据
|
|
contact_sensor = env.scene.sensors.get(sensor_cfg.name)
|
|
if contact_sensor is None:
|
|
return torch.zeros(env.num_envs, device=env.device)
|
|
|
|
# 获取所有定义的手臂/手部 link 的垂直总受力 (World Z)
|
|
# net_forces_w 形状: (num_envs, num_bodies, 3)
|
|
arm_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
|
|
# 取所有受力点的最大值或平均值，代表支撑强度
|
|
max_arm_force = torch.max(arm_forces_z, dim=-1)[0]
|
|
|
|
# 2. 获取状态数据
|
|
pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
|
|
pelvis_pos_z = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
|
|
root_vel_z = env.scene["robot"].data.root_lin_vel_w[:, 2]
|
|
|
|
# 3. 计算奖励项
|
|
# A. 受力奖励：鼓励手部与地面产生大于 min_force 的推力
|
|
# 使用 tanh 归一化，防止力矩过大导致奖励爆炸 (NaN 风险)
|
|
force_reward = torch.tanh(torch.clamp(max_arm_force - min_force, min=0.0) / 50.0)
|
|
|
|
# B. 速度引导：只有当机器人正在“向上起”时，支撑奖励才翻倍
|
|
# 这样可以防止它趴在地上乱按手骗分
|
|
velocity_factor = torch.clamp(root_vel_z, min=0.0, max=2.0)
|
|
|
|
# C. 姿态惩罚回避：
|
|
# 不再检查手是否在盆骨下方，而是检查手是否“在干活”
|
|
# 只要受力足够大，就认为是在支撑
|
|
is_supporting = (max_arm_force > min_force).float()
|
|
|
|
# 4. 阶段性退出机制 (Curriculum)
|
|
# 当盆骨高度超过 height_threshold (0.6m) 时，奖励线性消失
|
|
# 强迫机器人最终依靠腿部力量平衡，而不是一直扶着地
|
|
height_fade = torch.clamp((height_threshold - pelvis_pos_z) / 0.15, min=0.0, max=1.0)
|
|
|
|
# 最终组合
|
|
# 逻辑：受力 * (1 + 垂直速度) * 高度衰减
|
|
total_reward = force_reward * (1.0 + 2.0 * velocity_factor) * is_supporting * height_fade
|
|
|
|
return total_reward
|
|
|
|
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
|
|
|
|
def joint_deviation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
|
|
"""计算关节相对于默认姿态(default_joint_pos)的偏差平方和"""
|
|
# 获取当前关节位置相对于默认位置的差值
|
|
# mdp.joint_pos_rel 返回的是 (current_pos - default_pos)
|
|
diff = mdp.joint_pos_rel(env, asset_cfg)
|
|
return torch.sum(torch.square(diff), dim=-1)
|
|
|
|
# --- 2. 配置类 ---
|
|
|
|
T1_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',
|
|
# 手臂
|
|
'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'
|
|
]
|
|
|
|
|
|
@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)
|
|
root_pos = ObsTerm(func=mdp.root_pos_w)
|
|
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": (-1.57, 1.57), # 左右侧卧
|
|
"pitch": tuple(numpy.array([1.4, 1.6], dtype=np.float32) * random.choice([-1 , 1])), # 仰卧/俯卧
|
|
"yaw": (-3.14, 3.14), # 全向旋转
|
|
"x": (0.0, 0.0),
|
|
"y": (0.0, 0.0),
|
|
"z": (0.1, 0.2),
|
|
},
|
|
"velocity_range": {},
|
|
},
|
|
mode="reset",
|
|
)
|
|
|
|
|
|
@configclass
|
|
class T1ActionCfg:
|
|
"""关键修改：降低 scale 让动作变丝滑，增大阻尼效果"""
|
|
joint_pos = JointPositionActionCfg(
|
|
asset_name="robot",
|
|
joint_names=T1_JOINT_NAMES,
|
|
scale=0.5,
|
|
use_default_offset=True
|
|
)
|
|
|
|
|
|
@configclass
|
|
class T1GetUpRewardCfg:
|
|
# 1. 姿态基础奖 (引导身体变正)
|
|
upright = RewTerm(func=mdp.flat_orientation_l2, weight=2.0)
|
|
|
|
# 2. 【条件高度奖】：双高度判定（头+盆骨），且必须脚踩地
|
|
height_with_feet = RewTerm(
|
|
func=standing_with_feet_reward,
|
|
weight=20.0, # 作为核心引导，增加权重
|
|
params={
|
|
"min_head_height": 1.10,
|
|
"min_pelvis_height": 0.7,
|
|
"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]),
|
|
"force_threshold": 20.0,
|
|
"max_v_z": 0.3
|
|
}
|
|
)
|
|
|
|
# 3. 手臂撑地奖：辅助脱离地面阶段
|
|
arm_push_support = RewTerm(
|
|
func=universal_arm_support_reward,
|
|
weight=20.0, # 显著增加权重（从 3.0 提到 15.0），让它成为起步的关键
|
|
params={
|
|
"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_hand_link", "AL3", "AR3"]),
|
|
"height_threshold": 0.65, # 躯干升到 0.6m 前都鼓励手臂用力
|
|
"min_force": 3.0 # 只要有 15N 的力就触发
|
|
}
|
|
)
|
|
|
|
# 4. 关节限位惩罚 (新增：防止关节撞死导致数值问题)
|
|
joint_limits = RewTerm(
|
|
func=mdp.joint_pos_limits,
|
|
weight=-1.0,
|
|
params={"asset_cfg": SceneEntityCfg("robot")}
|
|
)
|
|
|
|
# 5. 时间惩罚 (强制效率)
|
|
time_penalty = RewTerm(
|
|
func=mdp.is_alive,
|
|
weight=-1.2
|
|
)
|
|
|
|
# 6. 成功终极大奖
|
|
is_success = RewTerm(
|
|
func=lambda env, keys: env.termination_manager.get_term(keys).float(),
|
|
weight=500.0,
|
|
params={"keys": "standing_success"}
|
|
)
|
|
|
|
# 7. 手臂关节活跃度奖 (诱导摆动)
|
|
arm_movement_exploration = RewTerm(
|
|
func=mdp.joint_vel_l2,
|
|
weight=2, # 权重不要太高，防止变成“风扇”
|
|
params={"asset_cfg": SceneEntityCfg("robot",joint_names=["Left_Shoulder.*", "Left_Elbow.*", "Right_Shoulder.*", "Right_Elbow.*"])}
|
|
)
|
|
|
|
# 8. 手臂位置多样性奖 (离开默认折叠姿态)
|
|
arm_deviation_bonus = RewTerm(
|
|
func=joint_deviation_l2,
|
|
weight=1,
|
|
params={"asset_cfg": SceneEntityCfg("robot",joint_names=["Left_Shoulder.*", "Left_Elbow.*", "Right_Shoulder.*", "Right_Elbow.*"])}
|
|
)
|
|
|
|
@configclass
|
|
class T1GetUpTerminationsCfg:
|
|
time_out = DoneTerm(func=mdp.time_out)
|
|
|
|
# 失败判定：躯干倾斜超过 45 度重置
|
|
#base_crash = DoneTerm(func=mdp.bad_orientation, params={"limit_angle": 0.785})
|
|
|
|
# 成功判定：双高度 + 稳定
|
|
standing_success = DoneTerm(
|
|
func=is_standing_still,
|
|
params={
|
|
"min_head_height": 1.05,
|
|
"min_pelvis_height": 0.75,
|
|
"max_angle_error": 0.3,
|
|
"standing_time": 0.2,
|
|
"velocity_threshold": 0.5
|
|
}
|
|
)
|
|
|
|
|
|
@configclass
|
|
class T1EnvCfg(ManagerBasedRLEnvCfg):
|
|
scene = T1SceneCfg(num_envs=8192, env_spacing=2.5) # 5090 性能全开
|
|
|
|
def __post_init__(self):
|
|
super().__post_init__()
|
|
self.scene.robot.init_state.pos = (0.0, 0.0, 0.2)
|
|
|
|
observations = T1ObservationCfg()
|
|
rewards = T1GetUpRewardCfg()
|
|
terminations = T1GetUpTerminationsCfg()
|
|
events = T1EventCfg()
|
|
actions = T1ActionCfg()
|
|
|
|
episode_length_s = 6.0
|
|
decimation = 4 |