249 lines
8.7 KiB
Python
249 lines
8.7 KiB
Python
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 is_standing_still(
|
||
env: ManagerBasedRLEnv,
|
||
minimum_height: float,
|
||
max_angle_error: float,
|
||
standing_time: float,
|
||
velocity_threshold: float = 0.2
|
||
) -> torch.Tensor:
|
||
"""
|
||
判定判定逻辑：高度达标 + 躯干垂直 + 几乎静止 + 维持时间。
|
||
增加了速度判定，彻底杜绝起跳瞬间触发。
|
||
"""
|
||
# 1. 获取 Body 索引
|
||
head_idx, _ = env.scene["robot"].find_bodies("H2")
|
||
|
||
# 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) &
|
||
(gravity_error < max_angle_error) &
|
||
(root_vel_norm < velocity_threshold) &
|
||
(root_ang_vel_norm < velocity_threshold * 2.0)
|
||
)
|
||
|
||
# 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"])
|
||
)
|
||
|
||
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. 配置类定义 ---
|
||
|
||
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'
|
||
]
|
||
|
||
|
||
@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)})
|
||
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": (0, 0),#(-1.57, 1.57),
|
||
"pitch": (-1.57, 1.57),#(-1.57, 1.57),
|
||
"yaw": (0, 0),#(-3.14, 3.14),
|
||
"x": (0.0, 0.0),
|
||
"y": (0.0, 0.0),
|
||
"z": (0.0, 0.0),
|
||
},
|
||
"velocity_range": {},
|
||
},
|
||
mode="reset",
|
||
)
|
||
|
||
|
||
@configclass
|
||
class T1ActionCfg:
|
||
"""关键修改：降低 scale 让动作变丝滑，增大阻尼效果"""
|
||
joint_pos = JointPositionActionCfg(
|
||
asset_name="robot",
|
||
joint_names=T1_JOINT_NAMES,
|
||
scale=0.2, # 从 0.5 降到 0.2，防止电机暴力抽搐
|
||
use_default_offset=True
|
||
)
|
||
|
||
|
||
@configclass
|
||
class T1GetUpRewardCfg:
|
||
# 1. 姿态奖 (核心引导)
|
||
upright = RewTerm(func=mdp.flat_orientation_l2, weight=5.0)
|
||
|
||
# 2. 只有脚着地时才给的高度奖（模拟逻辑）
|
||
# 修正：直接使用 root_height 配合强力的速度惩罚
|
||
height_tracking = RewTerm(
|
||
func=mdp.root_height_below_minimum,
|
||
weight=2.0,
|
||
params={
|
||
"minimum_height": 1.05,
|
||
"asset_cfg": SceneEntityCfg("robot", body_names="H2"),
|
||
}
|
||
)
|
||
|
||
# 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. 终极奖励
|
||
is_success = RewTerm(
|
||
func=get_success_reward,
|
||
weight=200.0, # 成功一次给大奖，但判定条件极严
|
||
params={"term_keys": "standing_success"}
|
||
)
|
||
|
||
@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={
|
||
"minimum_height": 1.05, # H2 高度
|
||
"max_angle_error": 0.15, # 极小角度误差（约 8.6 度）
|
||
"standing_time": 0.8, # 必须保持 0.8 秒（起跳不可能在空中停这么久）
|
||
"velocity_threshold": 0.15 # 速度必须极低
|
||
}
|
||
)
|
||
|
||
|
||
@configclass
|
||
class T1EnvCfg(ManagerBasedRLEnvCfg):
|
||
scene = T1SceneCfg(num_envs=16384, env_spacing=2.5) # 建议先用 4096 验证逻辑
|
||
|
||
def __post_init__(self):
|
||
super().__post_init__()
|
||
self.scene.robot.init_state.pos = (0.0, 0.0, 0.4) # 初始稍微抬高一点点，防止卡地
|
||
|
||
observations = T1ObservationCfg()
|
||
rewards = T1GetUpRewardCfg()
|
||
terminations = T1GetUpTerminationsCfg()
|
||
events = T1EventCfg()
|
||
actions = T1ActionCfg()
|
||
|
||
episode_length_s = 5.0 # 缩短时长，增加训练效率
|
||
decimation = 4 |