826 lines
34 KiB
Python
826 lines
34 KiB
Python
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 |