Gym_GPU/rl_game/get_up/config/t1_env_cfg.py

import torch
from pathlib import Path
import yaml
import isaaclab.envs.mdp as mdp
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.amp.amp_rewards import amp_style_prior_reward
from rl_game.get_up.env.t1_env import T1SceneCfg

_PROJECT_ROOT = Path(__file__).resolve().parents[3]
_DEFAULT_FRONT_KEYFRAME_YAML = str(_PROJECT_ROOT / "behaviors" / "custom" / "keyframe" / "get_up" / "get_up_front.yaml")
_DEFAULT_BACK_KEYFRAME_YAML = str(_PROJECT_ROOT / "behaviors" / "custom" / "keyframe" / "get_up" / "get_up_back.yaml")

def _contact_force_z(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
    """Sum positive vertical contact force on selected bodies."""
    sensor = env.scene.sensors.get(sensor_cfg.name)
    forces_z = sensor.data.net_forces_w[:, :, 2]
    body_ids = sensor_cfg.body_ids
    if body_ids is None:
        selected_z = forces_z
    else:
        selected_z = forces_z[:, body_ids]
    return torch.clamp(torch.sum(selected_z, dim=-1), min=0.0)


def _safe_tensor(x: torch.Tensor, nan: float = 0.0, pos: float = 1e3, neg: float = -1e3) -> torch.Tensor:
    """Keep reward pipeline numerically stable."""
    return torch.nan_to_num(x, nan=nan, posinf=pos, neginf=neg)


def _resolve_path(path_like: str) -> Path:
    path = Path(path_like).expanduser()
    if path.is_absolute():
        return path
    return (_PROJECT_ROOT / path).resolve()


def _interpolate_keyframes(
        keyframes: list[dict],
        sample_dt: float,
) -> tuple[torch.Tensor, list[str]]:
    """Interpolate sparse keyframes into dense [T, M] motor angle table in radians."""
    if len(keyframes) == 0:
        raise ValueError("No keyframes in motion yaml")

    motor_names: set[str] = set()
    for frame in keyframes:
        motors = frame.get("motor_positions", {})
        if isinstance(motors, dict):
            motor_names.update(motors.keys())
    names = sorted(motor_names)
    if len(names) == 0:
        raise ValueError("No motor_positions found in keyframes")

    frame_count = len(keyframes)
    times = torch.zeros(frame_count, dtype=torch.float32)
    values = torch.full((frame_count, len(names)), float("nan"), dtype=torch.float32)
    name_to_idx = {name: idx for idx, name in enumerate(names)}

    elapsed = 0.0
    for i, frame in enumerate(keyframes):
        elapsed += max(float(frame.get("delta", 0.1)), 1e-4)
        times[i] = elapsed
        motors = frame.get("motor_positions", {})
        if not isinstance(motors, dict):
            continue
        for motor_name, motor_deg in motors.items():
            if motor_name not in name_to_idx:
                continue
            values[i, name_to_idx[motor_name]] = float(motor_deg) * (torch.pi / 180.0)

    values[0] = torch.nan_to_num(values[0], nan=0.0)
    for i in range(1, frame_count):
        values[i] = torch.where(torch.isnan(values[i]), values[i - 1], values[i])

    sample_dt = max(float(sample_dt), 1e-3)
    sample_times = torch.arange(0.0, float(times[-1]) + 1e-6, sample_dt, dtype=torch.float32)
    sample_times[0] = max(sample_times[0], 1e-6)
    sample_times = torch.clamp(sample_times, min=times[0], max=times[-1])

    upper = torch.searchsorted(times, sample_times, right=True)
    upper = torch.clamp(upper, min=1, max=frame_count - 1)
    lower = upper - 1

    t0 = times.index_select(0, lower)
    t1 = times.index_select(0, upper)
    v0 = values.index_select(0, lower)
    v1 = values.index_select(0, upper)
    alpha = ((sample_times - t0) / (t1 - t0 + 1e-6)).unsqueeze(-1)
    interp = v0 + alpha * (v1 - v0)
    return interp, names


def _motion_table_from_yaml(
        yaml_path: str,
        joint_names: list[str],
        sample_dt: float,
) -> tuple[torch.Tensor, float]:
    """
    Build [T, J] target joint motion from get-up keyframes.
    Unknown joints default to 0.0 (neutral).
    """
    path = _resolve_path(yaml_path)
    if not path.is_file():
        raise FileNotFoundError(f"Motion yaml not found: {path}")
    with path.open("r", encoding="utf-8") as f:
        payload = yaml.safe_load(f) or {}
    keyframes = payload.get("keyframes", [])
    if not isinstance(keyframes, list):
        raise ValueError(f"Invalid keyframes in yaml: {path}")

    motor_table, motor_names = _interpolate_keyframes(keyframes, sample_dt=sample_dt)
    motor_to_idx = {name: idx for idx, name in enumerate(motor_names)}
    joint_table = torch.zeros((motor_table.shape[0], len(joint_names)), dtype=torch.float32)

    sign_flip_bases = {"Shoulder_Roll", "Elbow_Yaw", "Hip_Roll", "Hip_Yaw", "Ankle_Roll"}
    for j, joint_name in enumerate(joint_names):
        alias = "Head_yaw" if joint_name == "AAHead_yaw" else joint_name
        base = alias.replace("Left_", "").replace("Right_", "")
        src = None
        if alias in motor_to_idx:
            src = alias
        elif base in motor_to_idx:
            src = base
        if src is None:
            continue
        sign = -1.0 if alias.startswith("Right_") and base in sign_flip_bases else 1.0
        joint_table[:, j] = sign * motor_table[:, motor_to_idx[src]]

    duration_s = max(float(motor_table.shape[0] - 1) * sample_dt, sample_dt)
    return joint_table, duration_s


def _get_keyframe_motion_cache(
        env: ManagerBasedRLEnv,
        front_motion_path: str,
        back_motion_path: str,
        sample_dt: float,
):
    """Cache interpolated front/back get-up motion priors on env device."""
    cache_key = "getup_keyframe_motion_cache"
    sig = (str(front_motion_path), str(back_motion_path), float(sample_dt))
    cached = env.extras.get(cache_key, None)
    if isinstance(cached, dict) and cached.get("sig") == sig:
        return cached

    front_table, front_duration = _motion_table_from_yaml(front_motion_path, T1_JOINT_NAMES, sample_dt)
    back_table, back_duration = _motion_table_from_yaml(back_motion_path, T1_JOINT_NAMES, sample_dt)

    cache = {
        "sig": sig,
        "sample_dt": float(sample_dt),
        "front_motion": front_table.to(device=env.device),
        "back_motion": back_table.to(device=env.device),
        "front_duration": float(front_duration),
        "back_duration": float(back_duration),
    }
    env.extras[cache_key] = cache
    return cache


def keyframe_motion_prior_reward(
        env: ManagerBasedRLEnv,
        front_motion_path: str = _DEFAULT_FRONT_KEYFRAME_YAML,
        back_motion_path: str = _DEFAULT_BACK_KEYFRAME_YAML,
        sample_dt: float = 0.04,
        pose_sigma: float = 0.42,
        vel_sigma: float = 1.6,
        joint_subset: str = "all",
        internal_reward_scale: float = 1.0,
) -> torch.Tensor:
    """
    DeepMimic-style dense reward from keyframe get-up motions.
    - mode=1 uses front sequence
    - mode=0 uses back sequence
    """
    motion_cache = _get_keyframe_motion_cache(
        env=env,
        front_motion_path=front_motion_path,
        back_motion_path=back_motion_path,
        sample_dt=sample_dt,
    )

    # env.extras["getup_mode"]: 1=front, 0=back
    getup_mode = env.extras.get("getup_mode", None)
    if not isinstance(getup_mode, torch.Tensor) or getup_mode.shape[0] != env.num_envs:
        getup_mode = torch.zeros(env.num_envs, device=env.device, dtype=torch.long)
        env.extras["getup_mode"] = getup_mode
    getup_mode = getup_mode.to(dtype=torch.long)
    use_front = getup_mode == 1

    step_dt = env.step_dt
    phase_time = torch.clamp(env.episode_length_buf * step_dt, min=0.0)

    front_motion = motion_cache["front_motion"]
    back_motion = motion_cache["back_motion"]
    front_idx = torch.clamp((phase_time / sample_dt).to(torch.long), min=0, max=front_motion.shape[0] - 1)
    back_idx = torch.clamp((phase_time / sample_dt).to(torch.long), min=0, max=back_motion.shape[0] - 1)

    target_front = front_motion.index_select(0, front_idx)
    target_back = back_motion.index_select(0, back_idx)
    target_pos = torch.where(use_front.unsqueeze(-1), target_front, target_back)

    robot_data = env.scene["robot"].data
    current_pos = _safe_tensor(robot_data.joint_pos)
    current_vel = _safe_tensor(robot_data.joint_vel)

    if joint_subset == "legs":
        legs_idx, _ = env.scene["robot"].find_joints(".*(Hip|Knee|Ankle).*")
        if len(legs_idx) > 0:
            ids = torch.tensor(legs_idx, device=env.device, dtype=torch.long)
            current_pos = current_pos.index_select(1, ids)
            current_vel = current_vel.index_select(1, ids)
            target_pos = target_pos.index_select(1, ids)
    elif joint_subset == "core":
        core_idx, _ = env.scene["robot"].find_joints(".*(Waist|Hip|Knee|Ankle).*")
        if len(core_idx) > 0:
            ids = torch.tensor(core_idx, device=env.device, dtype=torch.long)
            current_pos = current_pos.index_select(1, ids)
            current_vel = current_vel.index_select(1, ids)
            target_pos = target_pos.index_select(1, ids)

    target_vel = torch.zeros_like(target_pos)
    pos_mse = torch.mean(torch.square(current_pos - target_pos), dim=-1)
    vel_mse = torch.mean(torch.square(current_vel - target_vel), dim=-1)

    pose_sigma = max(float(pose_sigma), 1e-3)
    vel_sigma = max(float(vel_sigma), 1e-3)
    pose_reward = torch.exp(-pos_mse / pose_sigma)
    vel_reward = torch.exp(-vel_mse / vel_sigma)
    prior_reward = 0.75 * pose_reward + 0.25 * vel_reward
    prior_reward = _safe_tensor(prior_reward, nan=0.0, pos=1.0, neg=0.0)

    log_dict = env.extras.get("log", {})
    if isinstance(log_dict, dict):
        log_dict["keyframe_prior_mean"] = torch.mean(prior_reward).detach().item()
        log_dict["keyframe_front_ratio"] = torch.mean(use_front.float()).detach().item()
        env.extras["log"] = log_dict

    return internal_reward_scale * prior_reward


def root_height_obs(env: ManagerBasedRLEnv) -> torch.Tensor:
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
    return env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2].unsqueeze(-1)


def head_height_obs(env: ManagerBasedRLEnv) -> torch.Tensor:
    head_idx, _ = env.scene["robot"].find_bodies("H2")
    return env.scene["robot"].data.body_state_w[:, head_idx[0], 2].unsqueeze(-1)


def foot_support_force_obs(env: ManagerBasedRLEnv, foot_sensor_cfg: SceneEntityCfg, norm_force: float = 120.0) -> torch.Tensor:
    foot_force_z = _contact_force_z(env, foot_sensor_cfg)
    return torch.tanh(foot_force_z / norm_force).unsqueeze(-1)


def arm_support_force_obs(env: ManagerBasedRLEnv, arm_sensor_cfg: SceneEntityCfg, norm_force: float = 120.0) -> torch.Tensor:
    arm_force_z = _contact_force_z(env, arm_sensor_cfg)
    return torch.tanh(arm_force_z / norm_force).unsqueeze(-1)


def contact_balance_obs(
        env: ManagerBasedRLEnv,
        foot_sensor_cfg: SceneEntityCfg,
        arm_sensor_cfg: SceneEntityCfg,
) -> torch.Tensor:
    foot_force_z = _contact_force_z(env, foot_sensor_cfg)
    arm_force_z = _contact_force_z(env, arm_sensor_cfg)
    total_support = foot_force_z + arm_force_z + 1e-6
    foot_support_ratio = torch.clamp(foot_force_z / total_support, min=0.0, max=1.0)
    return foot_support_ratio.unsqueeze(-1)


def reset_root_state_bimodal_lie_pose(
        env: ManagerBasedRLEnv,
        env_ids: torch.Tensor,
        asset_cfg: SceneEntityCfg,
        roll_range: tuple[float, float],
        pitch_abs_range: tuple[float, float],
        yaw_abs_range: tuple[float, float],
        x_range: tuple[float, float],
        y_range: tuple[float, float],
        z_range: tuple[float, float],
):
    """Reset with two lying modes around +pi/2 and -pi/2."""
    robot = env.scene[asset_cfg.name]
    num_resets = len(env_ids)
    default_root_state = robot.data.default_root_state[env_ids].clone()
    env_origins = env.scene.env_origins[env_ids]

    for i, bounds in enumerate([x_range, y_range, z_range]):
        v_min, v_max = bounds
        rand_vals = torch.rand(num_resets, device=env.device)
        default_root_state[:, i] = env_origins[:, i] + v_min + rand_vals * (v_max - v_min)

    euler_angles = torch.zeros((num_resets, 3), device=env.device)

    roll_min, roll_max = roll_range
    euler_angles[:, 0] = roll_min + torch.rand(num_resets, device=env.device) * (roll_max - roll_min)

    pitch_min, pitch_max = pitch_abs_range
    pitch_mag = pitch_min + torch.rand(num_resets, device=env.device) * (pitch_max - pitch_min)
    pitch_sign = torch.where(
        torch.rand(num_resets, device=env.device) > 0.5,
        torch.ones(num_resets, device=env.device),
        -torch.ones(num_resets, device=env.device),
    )
    euler_angles[:, 1] = pitch_mag * pitch_sign
    # Cache get-up mode for motion priors: pitch<0 => front, pitch>=0 => back.
    if "getup_mode" not in env.extras or not isinstance(env.extras.get("getup_mode"), torch.Tensor):
        env.extras["getup_mode"] = torch.zeros(env.num_envs, device=env.device, dtype=torch.long)
    env.extras["getup_mode"][env_ids] = (pitch_sign < 0.0).to(torch.long)

    yaw_min, yaw_max = yaw_abs_range
    yaw_mag = yaw_min + torch.rand(num_resets, device=env.device) * (yaw_max - yaw_min)
    yaw_sign = torch.where(
        torch.rand(num_resets, device=env.device) > 0.5,
        torch.ones(num_resets, device=env.device),
        -torch.ones(num_resets, device=env.device),
    )
    euler_angles[:, 2] = yaw_mag * yaw_sign

    roll, pitch, yaw = euler_angles[:, 0], euler_angles[:, 1], euler_angles[:, 2]
    cr, sr = torch.cos(roll * 0.5), torch.sin(roll * 0.5)
    cp, sp = torch.cos(pitch * 0.5), torch.sin(pitch * 0.5)
    cy, sy = torch.cos(yaw * 0.5), torch.sin(yaw * 0.5)

    qw = cr * cp * cy + sr * sp * sy
    qx = sr * cp * cy - cr * sp * sy
    qy = cr * sp * cy + sr * cp * sy
    qz = cr * cp * sy - sr * sp * cy
    default_root_state[:, 3:7] = torch.stack([qw, qx, qy, qz], dim=-1)

    robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids)
    robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)


def smooth_additive_getup_reward(
        env: ManagerBasedRLEnv,
        min_head_height: float,
        min_pelvis_height: float,
        foot_sensor_cfg: SceneEntityCfg,
        arm_sensor_cfg: SceneEntityCfg,
        upright_gain: float = 2.4,
        pelvis_progress_gain: float = 1.8,
        head_clearance_gain: float = 1.0,
        foot_support_gain: float = 1.2,
        arm_release_gain: float = 1.2,
        knee_mid_bend_gain: float = 0.8,
        knee_target: float = 1.0,
        knee_sigma: float = 0.5,
        hip_roll_penalty_gain: float = 0.5,
        hip_roll_soft_limit: float = 0.42,
        symmetry_penalty_gain: float = 0.2,
        standing_vel_penalty_gain: float = 0.35,
        standing_vel_gate_h: float = 0.65,
        stand_core_gain: float = 2.4,
        stand_upright_threshold: float = 0.82,
        stand_foot_support_threshold: float = 0.65,
        stand_arm_support_threshold: float = 0.25,
        internal_reward_scale: float = 1.0,
) -> torch.Tensor:
    # Cache expensive regex-based index lookups once per run.
    idx_cache_key = "getup_idx_cache"
    idx_cache = env.extras.get(idx_cache_key, None)
    if not isinstance(idx_cache, dict):
        idx_cache = {}
        env.extras[idx_cache_key] = idx_cache

    def _cached_joint_ids(cache_name: str, expr: str) -> torch.Tensor:
        ids = idx_cache.get(cache_name, None)
        if isinstance(ids, torch.Tensor):
            return ids
        joint_idx, _ = env.scene["robot"].find_joints(expr)
        ids = torch.tensor(joint_idx, device=env.device, dtype=torch.long) if len(joint_idx) > 0 else torch.empty(0, device=env.device, dtype=torch.long)
        idx_cache[cache_name] = ids
        return ids

    def _cached_body_id(cache_name: str, expr: str) -> int | None:
        idx = idx_cache.get(cache_name, None)
        if isinstance(idx, int):
            return idx
        body_idx, _ = env.scene["robot"].find_bodies(expr)
        idx = int(body_idx[0]) if len(body_idx) > 0 else None
        idx_cache[cache_name] = idx
        return idx

    joint_pos = _safe_tensor(env.scene["robot"].data.joint_pos)
    head_id = _cached_body_id("head_id", "H2")
    pelvis_id = _cached_body_id("pelvis_id", "Trunk")
    head_h = env.scene["robot"].data.body_state_w[:, head_id, 2] if head_id is not None else torch.zeros(env.num_envs, device=env.device)
    pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_id, 2] if pelvis_id is not None else torch.zeros(env.num_envs, device=env.device)
    head_h = _safe_tensor(head_h, nan=0.0, pos=2.0, neg=-2.0)
    pelvis_h = _safe_tensor(pelvis_h, nan=0.0, pos=2.0, neg=-2.0)
    projected_gravity = _safe_tensor(env.scene["robot"].data.projected_gravity_b, nan=0.0, pos=2.0, neg=-2.0)

    upright_ratio = torch.clamp(1.0 - torch.norm(projected_gravity[:, :2], dim=-1), min=0.0, max=1.0)
    pelvis_progress = torch.clamp((pelvis_h - 0.20) / (min_pelvis_height - 0.20 + 1e-6), min=0.0, max=1.0)
    head_clearance = torch.clamp((head_h - pelvis_h) / 0.45, min=0.0, max=1.0)

    foot_force_z = _safe_tensor(_contact_force_z(env, foot_sensor_cfg), nan=0.0, pos=1e3, neg=0.0)
    arm_force_z = _safe_tensor(_contact_force_z(env, arm_sensor_cfg), nan=0.0, pos=1e3, neg=0.0)
    total_support = _safe_tensor(foot_force_z + arm_force_z + 1e-6, nan=1.0, pos=1e4, neg=1e-6)
    foot_support_ratio = torch.clamp(foot_force_z / total_support, min=0.0, max=1.0)
    arm_support_ratio = torch.clamp(arm_force_z / total_support, min=0.0, max=1.0)
    release_phase = torch.sigmoid((pelvis_h - 0.58) * 8.0)
    arm_release_reward = (1.0 - arm_support_ratio) * release_phase

    knee_ids = _cached_joint_ids("knee_ids", ".*Knee_Pitch")
    if knee_ids.numel() > 0:
        knee_fold = torch.mean(torch.abs(joint_pos.index_select(1, knee_ids)), dim=-1)
        knee_mid_bend = torch.exp(-0.5 * torch.square((knee_fold - knee_target) / knee_sigma))
    else:
        knee_mid_bend = torch.zeros_like(pelvis_h)

    hip_roll_ids = _cached_joint_ids("hip_roll_ids", ".*Hip_Roll")
    if hip_roll_ids.numel() > 0:
        hip_roll_abs = torch.mean(torch.abs(joint_pos.index_select(1, hip_roll_ids)), dim=-1)
    else:
        hip_roll_abs = torch.zeros_like(pelvis_h)
    hip_roll_excess = torch.clamp(hip_roll_abs - hip_roll_soft_limit, min=0.0, max=0.5)

    left_leg_ids = _cached_joint_ids("left_leg_ids", "^Left_(Hip_Pitch|Hip_Roll|Hip_Yaw|Knee_Pitch|Ankle_Pitch|Ankle_Roll)$")
    right_leg_ids = _cached_joint_ids("right_leg_ids", "^Right_(Hip_Pitch|Hip_Roll|Hip_Yaw|Knee_Pitch|Ankle_Pitch|Ankle_Roll)$")
    if left_leg_ids.numel() > 0 and right_leg_ids.numel() > 0 and left_leg_ids.numel() == right_leg_ids.numel():
        left_leg = joint_pos.index_select(1, left_leg_ids)
        right_leg = joint_pos.index_select(1, right_leg_ids)
        symmetry_penalty = torch.mean(torch.abs(left_leg - right_leg), dim=-1)
    else:
        symmetry_penalty = torch.zeros_like(pelvis_h)

    root_vel = _safe_tensor(torch.norm(env.scene["robot"].data.root_lin_vel_w, dim=-1), nan=10.0, pos=10.0, neg=0.0)
    standing_gate = torch.sigmoid((pelvis_h - standing_vel_gate_h) * 10.0)

    # Core dense signal for "stand up and stand stable".
    stand_head = torch.sigmoid((head_h - min_head_height) * 10.0)
    stand_pelvis = torch.sigmoid((pelvis_h - min_pelvis_height) * 10.0)
    stand_upright = torch.sigmoid((upright_ratio - stand_upright_threshold) * 12.0)
    stand_foot = torch.sigmoid((foot_support_ratio - stand_foot_support_threshold) * 12.0)
    stand_arm_release = torch.sigmoid((stand_arm_support_threshold - arm_support_ratio) * 12.0)
    stand_still = torch.exp(-3.0 * root_vel)
    stand_core = (
        0.22 * stand_head
        + 0.22 * stand_pelvis
        + 0.18 * stand_upright
        + 0.18 * stand_foot
        + 0.10 * stand_arm_release
        + 0.10 * stand_still
    )

    total_reward = (
        upright_gain * upright_ratio
        + pelvis_progress_gain * pelvis_progress
        + head_clearance_gain * head_clearance
        + foot_support_gain * foot_support_ratio
        + arm_release_gain * arm_release_reward
        + knee_mid_bend_gain * knee_mid_bend
        - hip_roll_penalty_gain * hip_roll_excess
        - symmetry_penalty_gain * symmetry_penalty
        - standing_vel_penalty_gain * standing_gate * root_vel
        + stand_core_gain * stand_core
    )
    total_reward = _safe_tensor(total_reward, nan=0.0, pos=100.0, neg=-100.0)

    upright_mean = torch.mean(upright_ratio).detach().item()
    foot_support_ratio_mean = torch.mean(foot_support_ratio).detach().item()
    arm_support_ratio_mean = torch.mean(arm_support_ratio).detach().item()
    hip_roll_mean = torch.mean(hip_roll_abs).detach().item()
    stand_core_mean = torch.mean(stand_core).detach().item()

    log_dict = env.extras.get("log", {})
    if isinstance(log_dict, dict):
        log_dict["upright_mean"] = upright_mean
        log_dict["foot_support_ratio_mean"] = foot_support_ratio_mean
        log_dict["arm_support_ratio_mean"] = arm_support_ratio_mean
        log_dict["hip_roll_mean"] = hip_roll_mean
        log_dict["stand_core_mean"] = stand_core_mean
        env.extras["log"] = log_dict

    return internal_reward_scale * total_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 is_supported_standing(
        env: ManagerBasedRLEnv,
        foot_sensor_cfg: SceneEntityCfg,
        arm_sensor_cfg: SceneEntityCfg,
        min_head_height: float,
        min_pelvis_height: float,
        max_angle_error: float,
        velocity_threshold: float,
        min_foot_support_force: float,
        max_arm_support_force: float,
        standing_time: float,
        timer_name: str = "stable_timer",
) -> torch.Tensor:
    head_idx, _ = env.scene["robot"].find_bodies("H2")
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
    head_h = env.scene["robot"].data.body_state_w[:, head_idx[0], 2]
    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)

    foot_force_z = _contact_force_z(env, foot_sensor_cfg)
    arm_force_z = _contact_force_z(env, arm_sensor_cfg)

    is_stable_now = (
        (head_h > min_head_height)
        & (pelvis_h > min_pelvis_height)
        & (gravity_error < max_angle_error)
        & (root_vel_norm < velocity_threshold)
        & (foot_force_z > min_foot_support_force)
        & (arm_force_z < max_arm_support_force)
    )

    if timer_name not in env.extras:
        env.extras[timer_name] = torch.zeros(env.num_envs, device=env.device)

    dt = env.physics_dt * env.cfg.decimation
    env.extras[timer_name] = torch.where(
        is_stable_now, env.extras[timer_name] + dt, torch.zeros_like(env.extras[timer_name])
    )
    return (env.extras[timer_name] > standing_time).bool()


def airborne_flip_termination(
        env: ManagerBasedRLEnv,
        foot_sensor_cfg: SceneEntityCfg,
        arm_sensor_cfg: SceneEntityCfg,
        full_support_sensor_cfg: SceneEntityCfg,
        full_support_threshold: float = 12.0,
        min_pelvis_height: float = 0.34,
        contact_force_threshold: float = 6.0,
        inverted_gravity_threshold: float = 0.45,
        flip_ang_vel_threshold: float = 5.6,
        persist_time: float = 0.18,
        timer_name: str = "airborne_flip_timer",
) -> torch.Tensor:
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
    pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
    projected_gravity = env.scene["robot"].data.projected_gravity_b
    ang_vel = env.scene["robot"].data.root_ang_vel_w
    foot_force_z = _contact_force_z(env, foot_sensor_cfg)
    arm_force_z = _contact_force_z(env, arm_sensor_cfg)
    has_no_support = (foot_force_z < contact_force_threshold) & (arm_force_z < contact_force_threshold)
    full_support_force_z = _contact_force_z(env, full_support_sensor_cfg)
    is_fully_airborne = full_support_force_z < full_support_threshold
    is_airborne = (pelvis_h > min_pelvis_height) & has_no_support & is_fully_airborne
    is_inverted = projected_gravity[:, 2] > inverted_gravity_threshold
    is_fast_spin = torch.norm(ang_vel, dim=-1) > flip_ang_vel_threshold
    bad_state = is_airborne & (is_inverted | is_fast_spin)

    if timer_name not in env.extras:
        env.extras[timer_name] = torch.zeros(env.num_envs, device=env.device)
    dt = env.physics_dt * env.cfg.decimation
    env.extras[timer_name] = torch.where(
        bad_state, env.extras[timer_name] + dt, torch.zeros_like(env.extras[timer_name])
    )
    return (env.extras[timer_name] > persist_time).bool()


def nonfinite_state_termination(env: ManagerBasedRLEnv) -> torch.Tensor:
    """Terminate envs when sim state becomes NaN/Inf."""
    robot_data = env.scene["robot"].data
    finite_joint_pos = torch.isfinite(robot_data.joint_pos).all(dim=-1)
    finite_joint_vel = torch.isfinite(robot_data.joint_vel).all(dim=-1)
    finite_root_lin = torch.isfinite(robot_data.root_lin_vel_w).all(dim=-1)
    finite_root_ang = torch.isfinite(robot_data.root_ang_vel_w).all(dim=-1)
    finite_gravity = torch.isfinite(robot_data.projected_gravity_b).all(dim=-1)
    finite_root_pos = torch.isfinite(robot_data.root_pos_w).all(dim=-1)
    is_finite = finite_joint_pos & finite_joint_vel & finite_root_lin & finite_root_ang & finite_gravity & finite_root_pos
    return ~is_finite


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
        root_height = ObsTerm(func=root_height_obs)
        head_height = ObsTerm(func=head_height_obs)
        foot_support_force = ObsTerm(
            func=foot_support_force_obs,
            params={"foot_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]), "norm_force": 120.0},
        )
        arm_support_force = ObsTerm(
            func=arm_support_force_obs,
            params={"arm_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["A[LR][23]", ".*_hand_link"]), "norm_force": 120.0},
        )
        foot_support_ratio = ObsTerm(
            func=contact_balance_obs,
            params={
                "foot_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]),
                "arm_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["A[LR][23]", ".*_hand_link"]),
            },
        )
        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=reset_root_state_bimodal_lie_pose,
        params={
            "asset_cfg": SceneEntityCfg("robot"),
            "roll_range": (-0.15, 0.15),
            "pitch_abs_range": (1.40, 1.70),
            "yaw_abs_range": (0.0, 3.14),
            "x_range": (-0.04, 0.04),
            "y_range": (-0.04, 0.04),
            "z_range": (0.10, 0.18),
        },
        mode="reset",
    )


@configclass
class T1ActionCfg:
    head_action = JointPositionActionCfg(
        asset_name="robot", 
        joint_names=[
            "AAHead_yaw", "Head_pitch",
        ], 
        scale=0.5, 
        use_default_offset=True
    )
    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.82,
        use_default_offset=True,
    )
    torso_action = JointPositionActionCfg(
        asset_name="robot", 
        joint_names=[
            "Waist"
        ], 
        scale=0.58, 
        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=1.05,
        use_default_offset=True,
    )


@configclass
class T1GetUpRewardCfg:
    smooth_getup = RewTerm(
        func=smooth_additive_getup_reward,
        weight=5.0,
        params={
            "min_head_height": 1.02,
            "min_pelvis_height": 0.78,
            "foot_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]),
            "arm_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["A[LR][23]", ".*_hand_link"]),
            "upright_gain": 2.4,
            "pelvis_progress_gain": 1.8,
            "head_clearance_gain": 1.0,
            "foot_support_gain": 1.2,
            "arm_release_gain": 1.2,
            "knee_mid_bend_gain": 0.8,
            "knee_target": 1.0,
            "knee_sigma": 0.5,
            "hip_roll_penalty_gain": 0.5,
            "hip_roll_soft_limit": 0.42,
            "symmetry_penalty_gain": 0.2,
            "standing_vel_penalty_gain": 0.35,
            "standing_vel_gate_h": 0.65,
            "stand_core_gain": 2.4,
            "stand_upright_threshold": 0.82,
            "stand_foot_support_threshold": 0.65,
            "stand_arm_support_threshold": 0.25,
            "internal_reward_scale": 1.0,
        },
    )
    is_success_bonus = RewTerm(
        func=is_supported_standing,
        weight=150.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"]),
            "min_head_height": 1.05,
            "min_pelvis_height": 0.65,
            "max_angle_error": 0.18,
            "velocity_threshold": 0.10,
            "min_foot_support_force": 36.0,
            "max_arm_support_force": 14.0,
            "standing_time": 0.70,
            "timer_name": "reward_stable_timer",
        },
    )
    keyframe_motion_prior = RewTerm(
        func=keyframe_motion_prior_reward,
        weight=0.0,
        params={
            "front_motion_path": _DEFAULT_FRONT_KEYFRAME_YAML,
            "back_motion_path": _DEFAULT_BACK_KEYFRAME_YAML,
            "sample_dt": 0.04,
            "pose_sigma": 0.42,
            "vel_sigma": 1.6,
            "joint_subset": "all",
            "internal_reward_scale": 1.0,
        },
    )
    # AMP reward is disabled by default until a discriminator model path is provided.
    amp_style_prior = RewTerm(
        func=amp_style_prior_reward,
        weight=0.0,
        params={
            "amp_enabled": False,
            "amp_model_path": "",
            "amp_train_enabled": False,
            "amp_expert_features_path": "",
            "disc_hidden_dim": 256,
            "disc_hidden_layers": 2,
            "disc_lr": 3e-4,
            "disc_weight_decay": 1e-6,
            "disc_update_interval": 4,
            "disc_batch_size": 1024,
            "disc_min_expert_samples": 2048,
            "disc_history_steps": 4,
            "feature_clip": 8.0,
            "logit_scale": 1.0,
            "amp_reward_gain": 1.0,
            "internal_reward_scale": 1.0,
        },
    )


@configclass
class T1GetUpTerminationsCfg:
    time_out = DoneTerm(func=mdp.time_out)
    nonfinite_state_abort = DoneTerm(func=nonfinite_state_termination)
    anti_farming = DoneTerm(func=ground_farming_timeout, params={"max_time": 4.5, "height_threshold": 0.48})
    illegal_contact = DoneTerm(
        func=mdp.illegal_contact,
        params={"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["Trunk"]), "threshold": 200.0},
    )
    standing_success = DoneTerm(
        func=is_supported_standing,
        params={
            "foot_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]),
            "arm_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["A[LR][23]", ".*_hand_link"]),
            "min_head_height": 1.10,
            "min_pelvis_height": 0.83,
            "max_angle_error": 0.08,
            "velocity_threshold": 0.08,
            "min_foot_support_force": 36.0,
            "max_arm_support_force": 16.0,
            "standing_time": 1.4,
            "timer_name": "term_stable_timer",
        },
    )
    joint_velocity_limit = DoneTerm(
        func=mdp.joint_vel_out_of_manual_limit,
        params={"asset_cfg": SceneEntityCfg("robot"), "max_velocity": 50.0},
    )
    airborne_flip_abort = DoneTerm(
        func=airborne_flip_termination,
        params={
            "foot_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]),
            "arm_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["A[LR][23]", ".*_hand_link"]),
            "full_support_sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*"]),
            "full_support_threshold": 12.0,
            "min_pelvis_height": 0.34,
            "contact_force_threshold": 6.0,
            "inverted_gravity_threshold": 0.45,
            "flip_ang_vel_threshold": 5.6,
            "persist_time": 0.18,
            "timer_name": "airborne_flip_timer",
        },
    )


@configclass
class T1EnvCfg(ManagerBasedRLEnvCfg):
    scene = T1SceneCfg(num_envs=8192, env_spacing=5.0)
    observations = T1ObservationCfg()
    rewards = T1GetUpRewardCfg()
    terminations = T1GetUpTerminationsCfg()
    events = T1EventCfg()
    actions = T1ActionCfg()
    episode_length_s = 10.0
    decimation = 4