Gym_GPU/rl_game/get_up/config/t1_env_cfg.py

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. 自定义 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 = env.scene["robot"].data.body_state_w[:, head_idx[0], 2]
    curr_pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]

    # 归一化高度评分
    head_score = torch.clamp(curr_head_h / min_head_height, 0.0, 1.2)
    pelvis_score = torch.clamp(curr_pelvis_h / min_pelvis_height, 0.0, 1.2)
    height_reward = (head_score + pelvis_score) / 2.0

    # 足部受力判定
    contact_sensor = env.scene.sensors.get(sensor_cfg.name)
    if contact_sensor is None: return torch.zeros(env.num_envs, device=env.device)

    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(-torch.abs(root_vel_z) / max_v_z)

    return height_reward * (0.5 + 0.5 * force_weight * vel_penalty)


def dynamic_getup_strategy_reward(env: ManagerBasedRLEnv) -> torch.Tensor:
    """
    全姿态对称起立策略：
    1. 核心蜷缩 (Spring Loading)：无论仰卧还是俯卧，只要高度低，就必须强制收腿。
    2. 仰卧支撑 (Back-Pushing)：在仰卧状态下，鼓励手臂向后发力并抬高盆骨。
    3. 协同爆发 (Explosive Jump)：蜷缩状态下产生的向上动量获得最高倍率奖励。
    """
    # --- 1. 获取物理状态 ---
    gravity_z = env.scene["robot"].data.projected_gravity_b[:, 2]  # 1:仰卧, -1:俯卧
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
    curr_pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
    root_vel_z = env.scene["robot"].data.root_lin_vel_w[:, 2]

    # 关节索引：11,12髋, 17,18膝 (确保与T1模型一致)
    knee_joints = [17, 18]
    hip_pitch_joints = [11, 12]
    joint_pos = env.scene["robot"].data.joint_pos

    # --- 2. 核心蜷缩评分 (Crouch Score) ---
    # 无论仰俯，蜷缩是起立的绝对前提。目标是让脚尽可能靠近质心。
    # 提高膝盖弯曲目标 (1.5 rad)，引导更深度的折叠
    knee_flex_err = torch.abs(joint_pos[:, knee_joints] - 1.5).sum(dim=-1)
    hip_flex_err = torch.abs(joint_pos[:, hip_pitch_joints] - 1.2).sum(dim=-1)
    crouch_score = torch.exp(-(knee_flex_err + hip_flex_err) * 0.6)

    # 基础蜷缩奖励 (Spring Base) - 权重加大
    crouch_trigger = torch.clamp(0.6 - curr_pelvis_h, min=0.0)
    base_crouch_reward = crouch_trigger * crouch_score * 40.0

    # --- 3. 支撑力奖励 (Support Force) ---
    push_reward = torch.zeros_like(curr_pelvis_h)
    contact_sensor = env.scene.sensors.get("contact_sensor")
    if contact_sensor is not None:
        # 监测非足部Link（手、臂）的受力
        # 无论正反，只要手能提供垂直向上的推力，就是好手
        arm_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
        push_reward = torch.tanh(torch.max(arm_forces_z, dim=-1)[0] / 30.0)

    # --- 4. 姿态特定引导 (Orientation-Neutral) ---
    is_back = torch.clamp(gravity_z, min=0.0)  # 仰卧程度
    is_belly = torch.clamp(-gravity_z, min=0.0)  # 俯卧程度

    # A. 仰卧直接起立逻辑：
    # 在仰卧时，如果能把盆骨撑起来 (curr_pelvis_h 增加)，给予重奖
    # 配合crouch_score，鼓励“收腿-撑地-挺髋”的动作链
    back_lift_reward = is_back * torch.clamp(curr_pelvis_h - 0.15, min=0.0) * crouch_score * 50.0

    # B. 俯卧/翻身辅助逻辑 (保留一定的翻身倾向，但不再是唯一路径)
    flip_reward = is_back * (1.0 - gravity_z) * 5.0  # 权重降低，仅作为备选

    # --- 5. 最终爆发项 (The Jump) ---
    # 核心公式：蜷缩程度 * 向上速度 * 支撑力感应
    # 这是一个通用的“起跳”奖励，无论正反面，只要满足“缩得紧、跳得快、手有撑”，奖励就爆炸
    explosion_reward = crouch_score * torch.clamp(root_vel_z, min=0.0) * (0.5 + 0.5 * push_reward) * 80.0

    # --- 6. 汇总 ---
    total_reward = (
            base_crouch_reward +  # 必须缩腿
            back_lift_reward +  # 仰卧挺髋
            flip_reward +  # 翻身尝试
            explosion_reward  # 终极爆发
    )

    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


# --- 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)
        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(np.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.3, 0.4),
            },
            "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=1.2,  # 给了手臂相对充裕的自由度去摸索
        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.6,
        use_default_offset=True
    )


@configclass
class T1GetUpRewardCfg:
    # 1. 核心阶段性引导 (翻身 -> 蜷缩 -> 支撑)
    dynamic_strategy = RewTerm(
        func=dynamic_getup_strategy_reward,
        weight=1.5
    )

    # 2. 站立质量奖励 (强化双脚受力)
    height_with_feet = RewTerm(
        func=standing_with_feet_reward,
        weight=40.0,  # 大权重
        params={
            "min_head_height": 1.1,
            "min_pelvis_height": 0.7,
            "sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]),
            "force_threshold": 40.0,  # 必须达到一定压力，防止脚尖点地作弊
            "max_v_z": 0.2
        }
    )

    # 3. 惩罚项：防止钻空子
    # 严厉惩罚：如果躯干(Trunk)或头(H2)直接接触地面，扣大分
    body_contact_penalty = RewTerm(
        func=mdp.contact_forces,
        weight=-20.0,
        params={
            "sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["Trunk", "H2"]),
            "threshold": 1.0
        }
    )

    # 4. 关节功耗惩罚 (防止高频抽搐)
    action_rate = RewTerm(
        func=mdp.action_rate_l2,
        weight=-0.01
    )

    # 5. 成功维持奖励
    is_success_maintain = RewTerm(
        func=is_standing_still,
        weight=1000.0,  # 巨大的成功奖励
        params={
            "min_head_height": 1.08,
            "min_pelvis_height": 0.72,
            "max_angle_error": 0.2,
            "standing_time": 0.4,  # 必须站稳 0.4s
            "velocity_threshold": 0.3
        }
    )


@configclass
class T1GetUpTerminationsCfg:
    time_out = DoneTerm(func=mdp.time_out)
    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)

    observations = T1ObservationCfg()
    rewards = T1GetUpRewardCfg()
    terminations = T1GetUpTerminationsCfg()
    events = T1EventCfg()
    actions = T1ActionCfg()

    episode_length_s = 10.0
    decimation = 4