from __future__ import annotations import torch from typing import TYPE_CHECKING from isaaclab.envs import mdp from isaaclab.managers import SceneEntityCfg from isaaclab.sensors import ContactSensor from isaaclab.utils.math import ( euler_xyz_from_quat, quat_apply_inverse, quat_from_euler_xyz, quat_mul, quat_rotate_inverse, wrap_to_pi, yaw_quat, ) if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv def _to_env_device(env: ManagerBasedRLEnv, tensor: torch.Tensor) -> torch.Tensor: """Move sensor tensors back to the RL environment device when Isaac uses another CUDA device.""" return tensor.to(env.device) def _single_body_id(asset_cfg: SceneEntityCfg, term_name: str) -> int: body_ids = asset_cfg.body_ids if body_ids is None: raise ValueError(f"{term_name} requires asset_cfg with exactly one body name.") if isinstance(body_ids, int): return body_ids if isinstance(body_ids, slice) or len(body_ids) != 1: raise ValueError(f"{term_name} requires exactly one body id, got {body_ids}.") return body_ids[0] def _heading_quat_with_offset(body_quat_w: torch.Tensor, heading_yaw_offset: float = 0.0) -> torch.Tensor: roll, pitch, yaw = euler_xyz_from_quat(body_quat_w) zeros = torch.zeros_like(yaw) return quat_from_euler_xyz(zeros, zeros, yaw - heading_yaw_offset) def _heading_yaw_with_offset(body_quat_w: torch.Tensor, heading_yaw_offset: float = 0.0) -> torch.Tensor: roll, pitch, yaw = euler_xyz_from_quat(body_quat_w) return wrap_to_pi(yaw - heading_yaw_offset) def _command_is_moving( env: ManagerBasedRLEnv, command_name: str, linear_threshold: float = 0.1, angular_threshold: float = 0.1, ) -> torch.Tensor: commands = env.command_manager.get_command(command_name) return (torch.norm(commands[:, :2], dim=1) > linear_threshold) | ( torch.abs(commands[:, 2]) > angular_threshold ) def action_smoothness(env: ManagerBasedRLEnv) -> torch.Tensor: """Penalize action second-order differences to encourage smooth control.""" action_manager = env.action_manager prev_prev_action = getattr(action_manager, "_prev_prev_action", None) if prev_prev_action is None: prev_prev_action = torch.zeros_like(action_manager.action) action_manager._prev_prev_action = prev_prev_action second_diff = action_manager.action + prev_prev_action - 2.0 * action_manager.prev_action reward = torch.sum(torch.square(second_diff), dim=1) # Update action history for the next step and clear reset environments. prev_prev_action.copy_(action_manager.prev_action) reset_buf = getattr(env, "reset_buf", None) if reset_buf is not None: reset_env_ids = reset_buf.nonzero(as_tuple=False).squeeze(-1) if reset_env_ids.numel() > 0: prev_prev_action[reset_env_ids] = 0.0 return reward def _epoch_curriculum_scale(env: ManagerBasedRLEnv, start_scale: float, power: float, interval_epochs: int) -> float: """Compute epoch-based scale by exponentiating by power every interval_epochs.""" num_step = env.common_step_counter interval = max(int(interval_epochs), 1) updates = num_step // interval return float(start_scale) ** (float(power) ** updates) def action_smoothness_with_curriculum( env: ManagerBasedRLEnv, start_scale: float, power: float, interval_epochs: int ) -> torch.Tensor: """Action smoothness penalty with epoch-based curriculum scaling.""" reward = action_smoothness(env) return reward * _epoch_curriculum_scale(env, start_scale, power, interval_epochs) def action_rate_with_curriculum(env: ManagerBasedRLEnv, start_scale: float, power: float, interval_epochs: int) -> torch.Tensor: """Action rate penalty with epoch-based curriculum scaling.""" reward = mdp.action_rate_l2(env) return reward * _epoch_curriculum_scale(env, start_scale, power, interval_epochs) def energy_cost_with_curriculum( env: ManagerBasedRLEnv, start_scale: float, power: float, interval_epochs: int, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), ) -> torch.Tensor: """Energy cost penalty with epoch-based curriculum scaling.""" reward = energy_cost(env, asset_cfg=asset_cfg) return reward * _epoch_curriculum_scale(env, start_scale, power, interval_epochs) def feet_air_time_similarity( env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, scale: float = 4.0, min_air_time: float = 0.0, ) -> torch.Tensor: """Reward similar air time between two feet. The reward is computed when either foot makes a new contact. It compares the last completed air times for the two feet and rewards small differences. """ contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] body_ids = sensor_cfg.body_ids if body_ids is None or len(body_ids) != 2: raise ValueError("feet_air_time_similarity expects exactly two foot body ids in sensor_cfg.body_ids.") first_contact = _to_env_device(env, contact_sensor.compute_first_contact(env.step_dt)[:, body_ids]) last_air_time = _to_env_device(env, contact_sensor.data.last_air_time[:, body_ids]) recent_contact = torch.any(first_contact > 0.0, dim=1) valid = torch.all(last_air_time > min_air_time, dim=1) diff = torch.abs(last_air_time[:, 0] - last_air_time[:, 1]) reward = torch.exp(-diff * scale) return reward * (recent_contact & valid) def feet_air_time( env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float ) -> torch.Tensor: """Reward long steps while keeping contact-sensor tensors on the env device.""" contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] first_contact = _to_env_device(env, contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]) last_air_time = _to_env_device(env, contact_sensor.data.last_air_time[:, sensor_cfg.body_ids]) reward = torch.sum((last_air_time - threshold) * first_contact, dim=1) reward *= torch.norm(env.command_manager.get_command(command_name)[:, :2], dim=1) > 0.1 return reward def feet_air_time_positive_on_contact( env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg, min_air_time: float = 0.05, max_air_time: float = 0.25, linear_threshold: float = 0.1, angular_threshold: float = 0.1, ) -> torch.Tensor: """Reward completed swing times without penalizing short exploratory steps.""" if max_air_time <= min_air_time: raise ValueError("max_air_time must be greater than min_air_time.") contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] first_contact = _to_env_device(env, contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]) last_air_time = _to_env_device(env, contact_sensor.data.last_air_time[:, sensor_cfg.body_ids]) completed_swing = torch.clamp(last_air_time - min_air_time, min=0.0, max=max_air_time - min_air_time) reward = torch.sum(completed_swing * first_contact, dim=1) moving = _command_is_moving(env, command_name, linear_threshold, angular_threshold) return reward * moving def feet_air_time_positive_biped( env: ManagerBasedRLEnv, command_name: str, threshold: float, sensor_cfg: SceneEntityCfg, linear_threshold: float = 0.1, angular_threshold: float = 0.1, ) -> torch.Tensor: """Dense biped air-time reward with contact-sensor tensors on the env device.""" contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] air_time = _to_env_device(env, contact_sensor.data.current_air_time[:, sensor_cfg.body_ids]) contact_time = _to_env_device(env, contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids]) in_contact = contact_time > 0.0 in_mode_time = torch.where(in_contact, contact_time, air_time) single_stance = torch.sum(in_contact.int(), dim=1) == 1 reward = torch.min(torch.where(single_stance.unsqueeze(-1), in_mode_time, 0.0), dim=1)[0] reward = torch.clamp(reward, max=threshold) moving = _command_is_moving(env, command_name, linear_threshold, angular_threshold) return reward * moving def track_lin_vel_xy_yaw_frame_exp( env, sigma: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), stand_threshold: float = 0.06 ) -> torch.Tensor: """Reward tracking of linear velocity commands (xy axes) in the gravity aligned robot frame using exponential kernel.""" commands = env.command_manager.get_command(command_name) stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & ( torch.abs(commands[:, 2]) < stand_threshold ) asset = env.scene[asset_cfg.name] vel_yaw = quat_apply_inverse(yaw_quat(asset.data.root_quat_w), asset.data.root_lin_vel_w[:, :3]) lin_vel_error_square = torch.sum(torch.square(commands[:, :2] - vel_yaw[:, :2]), dim=1) lin_vel_error_abs = torch.sum(torch.abs(commands[:, :2] - vel_yaw[:, :2]), dim=1) rew_square = torch.exp(-lin_vel_error_square * sigma) rew_abs = torch.exp(-lin_vel_error_abs * sigma) return torch.where(stand_command, rew_abs, rew_square) def track_ang_vel_z_world_exp( env, command_name: str, sigma: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), stand_threshold: float = 0.06 ) -> torch.Tensor: """Reward tracking of angular velocity commands (yaw) in world frame using exponential kernel.""" commands = env.command_manager.get_command(command_name) stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & ( torch.abs(commands[:, 2]) < stand_threshold ) asset = env.scene[asset_cfg.name] ang_vel_error_square = torch.square(commands[:, 2] - asset.data.root_ang_vel_w[:, 2]) ang_vel_error_abs = torch.abs(commands[:, 2] - asset.data.root_ang_vel_w[:, 2]) rew_square = torch.exp(-ang_vel_error_square * sigma) rew_abs = torch.exp(-ang_vel_error_abs * sigma) return torch.where(stand_command, rew_abs, rew_square) def track_lin_vel_xy_yaw_frame_exp_body( env, sigma: float, command_name: str, asset_cfg: SceneEntityCfg, stand_threshold: float = 0.06, heading_yaw_offset: float = 0.0, ) -> torch.Tensor: """Track planar velocity of a configured body in its yaw-aligned frame.""" body_id = _single_body_id(asset_cfg, "track_lin_vel_xy_yaw_frame_exp_body") commands = env.command_manager.get_command(command_name) stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & ( torch.abs(commands[:, 2]) < stand_threshold ) asset = env.scene[asset_cfg.name] body_quat_w = asset.data.body_quat_w[:, body_id, :] body_lin_vel_w = asset.data.body_lin_vel_w[:, body_id, :] vel_yaw = quat_apply_inverse(_heading_quat_with_offset(body_quat_w, heading_yaw_offset), body_lin_vel_w) lin_vel_error_square = torch.sum(torch.square(commands[:, :2] - vel_yaw[:, :2]), dim=1) lin_vel_error_abs = torch.sum(torch.abs(commands[:, :2] - vel_yaw[:, :2]), dim=1) rew_square = torch.exp(-lin_vel_error_square * sigma) rew_abs = torch.exp(-lin_vel_error_abs * sigma) return torch.where(stand_command, rew_abs, rew_square) def track_ang_vel_z_world_exp_body( env, command_name: str, sigma: float, asset_cfg: SceneEntityCfg, stand_threshold: float = 0.06, ) -> torch.Tensor: """Track world-frame yaw angular velocity of a configured body.""" body_id = _single_body_id(asset_cfg, "track_ang_vel_z_world_exp_body") commands = env.command_manager.get_command(command_name) stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & ( torch.abs(commands[:, 2]) < stand_threshold ) asset = env.scene[asset_cfg.name] ang_vel_error_square = torch.square(commands[:, 2] - asset.data.body_ang_vel_w[:, body_id, 2]) ang_vel_error_abs = torch.abs(commands[:, 2] - asset.data.body_ang_vel_w[:, body_id, 2]) rew_square = torch.exp(-ang_vel_error_square * sigma) rew_abs = torch.exp(-ang_vel_error_abs * sigma) return torch.where(stand_command, rew_abs, rew_square) def track_ang_vel_z_world_exp_bodies( env, command_name: str, sigma: float, asset_cfg: SceneEntityCfg, stand_threshold: float = 0.06, ) -> torch.Tensor: """Track yaw rate across all configured bodies so internal waist motion cannot satisfy the command alone.""" body_ids = asset_cfg.body_ids if body_ids is None or isinstance(body_ids, slice): raise ValueError("track_ang_vel_z_world_exp_bodies requires one or more explicitly resolved body ids.") commands = env.command_manager.get_command(command_name) stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & ( torch.abs(commands[:, 2]) < stand_threshold ) asset = env.scene[asset_cfg.name] body_yaw_rates = asset.data.body_ang_vel_w[:, body_ids, 2] command_yaw_rate = commands[:, 2].unsqueeze(1) ang_vel_error_square = torch.mean(torch.square(command_yaw_rate - body_yaw_rates), dim=1) ang_vel_error_abs = torch.mean(torch.abs(command_yaw_rate - body_yaw_rates), dim=1) rew_square = torch.exp(-ang_vel_error_square * sigma) rew_abs = torch.exp(-ang_vel_error_abs * sigma) return torch.where(stand_command, rew_abs, rew_square) def feet_stumble( env, sensor_cfg: SceneEntityCfg, tangential_threshold: float = 2.0, normal_threshold: float = 1.0 ) -> torch.Tensor: """Penalize feet hitting vertical surfaces using contact forces. Flags a stumble when tangential force exceeds ``tangential_threshold`` while the normal force stays below ``normal_threshold``. Returns the count of stumbling feet per environment. """ contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] forces = _to_env_device(env, contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :]) tangential = torch.norm(forces[..., :2], dim=-1) > tangential_threshold small_normal = torch.abs(forces[..., 2]) < normal_threshold stumble = tangential & small_normal return stumble.sum(dim=1) def feet_contact( env, sensor_cfg: SceneEntityCfg, command_name: str, stand_threshold: float = 0.06, force_threshold: float = 5.0 ) -> torch.Tensor: """Reward valid foot contacts during walking and standing. - When the command is effectively zero (stand), reward 1 only if both feet are in contact. - Otherwise, reward 1 if any recent timestep had exactly one foot in contact. Uses contact force history if available, else falls back to the latest forces. """ commands = env.command_manager.get_command(command_name) stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & ( torch.abs(commands[:, 2]) < stand_threshold ) contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] contact_history = contact_sensor.data.net_forces_w_history if contact_history is None: contact_history = contact_sensor.data.net_forces_w.unsqueeze(1) contact_history = _to_env_device(env, contact_history) contacts = contact_history[:, :, sensor_cfg.body_ids, 2] > force_threshold contact_num_buf = torch.sum(contacts, dim=-1) reward = stand_command.float() contact_mask = (~stand_command) & torch.any(contact_num_buf == 1, dim=1) reward[contact_mask] = 1.0 return reward def feet_contact_fixed( env, sensor_cfg: SceneEntityCfg, command_name: str, stand_threshold: float = 0.06, force_threshold: float = 5.0 ) -> torch.Tensor: """Reward valid foot contacts during walking and standing. - When the command is effectively zero (stand), reward 1. - Otherwise, reward 1 if any recent timestep had exactly one foot in contact. Uses contact force history if available, else falls back to the latest forces. """ commands = env.command_manager.get_command(command_name) stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & ( torch.abs(commands[:, 2]) < stand_threshold ) contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] contact_history = contact_sensor.data.net_forces_w_history if contact_history is None: contact_history = contact_sensor.data.net_forces_w.unsqueeze(1) contact_history = _to_env_device(env, contact_history) contacts = contact_history[:, :, sensor_cfg.body_ids, 2] > force_threshold contact_num_buf = torch.sum(contacts, dim=-1) # For stand, require both feet in contact at the latest timestep. stand_contact = contact_num_buf[:, -1] == 2 reward = (stand_command & stand_contact).float() contact_mask = (~stand_command) & torch.any(contact_num_buf == 1, dim=1) reward[contact_mask] = 1.0 return reward def biped_contact_mode_reward( env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, command_name: str, force_threshold: float = 5.0, linear_threshold: float = 0.1, angular_threshold: float = 0.1, ) -> torch.Tensor: """Reward double support while standing and exactly one supporting foot while moving.""" contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] forces = _to_env_device(env, contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :]) contacts = torch.norm(forces, dim=-1) > force_threshold contact_count = torch.sum(contacts.int(), dim=1) moving = _command_is_moving(env, command_name, linear_threshold, angular_threshold) return torch.where(moving, contact_count == 1, contact_count == 2).float() def speed_conditioned_biped_contact_flight( env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, command_name: str, run_transition_start: float = 1.6, run_transition_end: float = 2.4, max_flight_time: float = 0.18, min_swing_lead_time: float = 0.05, invalid_flight_penalty: float = 1.0, excess_flight_penalty: float = 2.0, double_support_penalty: float = 0.25, linear_threshold: float = 0.1, angular_threshold: float = 0.1, ) -> torch.Tensor: """Blend walking contact rules into a conservative running flight phase. A short flight phase becomes neutral at high speed, but is never rewarded every frame. This avoids suppressing a true running gait without incentivizing jumping. """ if run_transition_end <= run_transition_start: raise ValueError("run_transition_end must exceed run_transition_start.") if max_flight_time <= 0.0: raise ValueError("max_flight_time must be positive.") body_ids = sensor_cfg.body_ids if body_ids is None or isinstance(body_ids, (int, slice)) or len(body_ids) != 2: raise ValueError("speed_conditioned_biped_contact_flight requires two ordered foot body ids.") contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] if not contact_sensor.cfg.track_air_time: raise RuntimeError("ContactSensorCfg.track_air_time must be enabled for running contact rewards.") air_time = _to_env_device(env, contact_sensor.data.current_air_time[:, body_ids]) contact_time = _to_env_device(env, contact_sensor.data.current_contact_time[:, body_ids]) in_contact = contact_time > 0.0 contact_count = torch.sum(in_contact.int(), dim=1) single_support = contact_count == 1 double_support = contact_count == 2 in_flight = contact_count == 0 commands = env.command_manager.get_command(command_name) moving = (torch.linalg.norm(commands[:, :2], dim=1) > linear_threshold) | ( torch.abs(commands[:, 2]) > angular_threshold ) run_mix = torch.clamp( (commands[:, 0] - run_transition_start) / (run_transition_end - run_transition_start), min=0.0, max=1.0, ) # min() measures the time for which both feet have simultaneously been airborne. flight_time = torch.min(air_time, dim=1).values swing_lead_time = torch.abs(air_time[:, 0] - air_time[:, 1]) valid_short_flight = ( in_flight & (flight_time <= max_flight_time) & (swing_lead_time >= min_swing_lead_time) ) # Walking penalizes any flight; running makes only short alternating flight neutral. flight_invalidity = in_flight.float() * (1.0 - run_mix) flight_invalidity += (in_flight & ~valid_short_flight).float() * run_mix excess_ratio = torch.clamp( (flight_time - max_flight_time) / max_flight_time, min=0.0, max=2.0 ) moving_reward = single_support.float() moving_reward -= invalid_flight_penalty * flight_invalidity moving_reward -= excess_flight_penalty * excess_ratio * in_flight.float() moving_reward -= double_support_penalty * run_mix * double_support.float() return torch.where(moving, moving_reward, double_support.float()) def swing_foot_clearance_reward( env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, sensor_cfg: SceneEntityCfg, command_name: str, target_height: float, std: float, force_threshold: float = 5.0, linear_threshold: float = 0.1, angular_threshold: float = 0.1, ) -> torch.Tensor: """Reward swing-foot height relative to the supporting foot.""" if std <= 0.0: raise ValueError("std must be positive.") asset = env.scene[asset_cfg.name] contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] forces = _to_env_device(env, contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :]) contacts = torch.norm(forces, dim=-1) > force_threshold single_stance = torch.sum(contacts.int(), dim=1) == 1 swing_feet = ~contacts foot_height = asset.data.body_pos_w[:, asset_cfg.body_ids, 2] if foot_height.shape[1] != contacts.shape[1]: raise ValueError("asset_cfg and sensor_cfg must resolve the same number of feet.") stance_height = torch.sum(foot_height * contacts, dim=1, keepdim=True) swing_clearance = foot_height - stance_height clearance_reward = torch.exp(-torch.square(swing_clearance - target_height) / (std * std)) clearance_reward = torch.sum(clearance_reward * swing_feet, dim=1) moving = _command_is_moving(env, command_name, linear_threshold, angular_threshold) return clearance_reward * single_stance * moving def feet_slide(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Penalize foot sliding while keeping contact tensors on the env device.""" contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] contacts = _to_env_device( env, contact_sensor.data.net_forces_w_history[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] ) > 1.0 asset = env.scene[asset_cfg.name] body_vel = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2] return torch.sum(body_vel.norm(dim=-1) * contacts, dim=1) def feet_position(env, asset_cfg: SceneEntityCfg, command_name: str, stand_threshold: float = 0.06, ankle_distance: float = 0.22, base_height_target: float = 0.82, ) -> torch.Tensor: """Reward keeping feet near a desired stance when standing; otherwise return 1.""" commands = env.command_manager.get_command(command_name) stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & ( torch.abs(commands[:, 2]) < stand_threshold ) asset = env.scene[asset_cfg.name] feet_pos_w = asset.data.body_pos_w[:, asset_cfg.body_ids, :] base_pos_w = asset.data.root_pos_w base_quat_w = asset.data.root_quat_w # isolate yaw heading; zero roll/pitch r, p, y = euler_xyz_from_quat(base_quat_w) heading_quat = quat_from_euler_xyz(torch.zeros_like(r), torch.zeros_like(p), y) feet_pos_rel = feet_pos_w - base_pos_w.unsqueeze(1) # Expand heading quaternions per foot to satisfy broadcasting expected by quat_apply_inverse. num_envs, num_feet, _ = feet_pos_rel.shape heading_quat_per_foot = heading_quat.unsqueeze(1).expand(-1, num_feet, -1).reshape(-1, 4) feet_pos_rel_flat = feet_pos_rel.reshape(-1, 3) feet_pos_heading = quat_apply_inverse(heading_quat_per_foot, feet_pos_rel_flat).reshape(num_envs, num_feet, 3) desired_x = torch.zeros((num_envs, num_feet), device=feet_pos_heading.device) desired_y = torch.cat( ( (ankle_distance * 0.5) * torch.ones((num_envs, num_feet // 2), device=feet_pos_heading.device), (-ankle_distance * 0.5) * torch.ones((num_envs, num_feet - num_feet // 2), device=feet_pos_heading.device), ), dim=1, ) desired_z = -(base_height_target - 0.045) * torch.ones((num_envs, num_feet), device=feet_pos_heading.device) desired = torch.stack((desired_x, desired_y, desired_z), dim=-1) position_error = torch.sum(torch.abs(feet_pos_heading - desired), dim=(1, 2)) reward_stand = torch.exp(-position_error * 3.0) return torch.where(stand_command, reward_stand, torch.ones_like(reward_stand)) def feet_position_relative_to_body( env, asset_cfg: SceneEntityCfg, reference_body_cfg: SceneEntityCfg, command_name: str, desired_foot_positions: tuple[tuple[float, float, float], ...], stand_threshold: float = 0.06, heading_yaw_offset: float = 0.0, scale: float = 3.0, ) -> torch.Tensor: """Reward standing foot positions relative to a configured reference body.""" reference_body_id = _single_body_id(reference_body_cfg, "feet_position_relative_to_body") commands = env.command_manager.get_command(command_name) stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & ( torch.abs(commands[:, 2]) < stand_threshold ) asset = env.scene[asset_cfg.name] reference_asset = env.scene[reference_body_cfg.name] feet_pos_w = asset.data.body_pos_w[:, asset_cfg.body_ids, :] reference_pos_w = reference_asset.data.body_pos_w[:, reference_body_id, :] reference_quat_w = reference_asset.data.body_quat_w[:, reference_body_id, :] num_envs, num_feet, _ = feet_pos_w.shape feet_pos_rel = feet_pos_w - reference_pos_w.unsqueeze(1) heading_quat = _heading_quat_with_offset(reference_quat_w, heading_yaw_offset) heading_quat_per_foot = heading_quat.unsqueeze(1).expand(-1, num_feet, -1).reshape(-1, 4) feet_pos_heading = quat_apply_inverse(heading_quat_per_foot, feet_pos_rel.reshape(-1, 3)).reshape(num_envs, num_feet, 3) desired = torch.tensor(desired_foot_positions, dtype=feet_pos_heading.dtype, device=feet_pos_heading.device) if desired.shape != (num_feet, 3): raise ValueError(f"desired_foot_positions must have shape ({num_feet}, 3), got {tuple(desired.shape)}.") position_error = torch.sum(torch.abs(feet_pos_heading - desired.unsqueeze(0)), dim=(1, 2)) reward_stand = torch.exp(-position_error * scale) return torch.where(stand_command, reward_stand, torch.ones_like(reward_stand)) def feet_regulation( env, asset_cfg: SceneEntityCfg, sensor_cfg: SceneEntityCfg, ankle_height: float = 0.045, base_height_target: float = 0.82, clearance_clip: float = 1.0, scale: float = 40.0, ) -> torch.Tensor: """Penalize feet moving fast while too close to terrain.""" asset = env.scene[asset_cfg.name] raycaster = env.scene.sensors[sensor_cfg.name] # terrain height estimate from rays ray_hits = raycaster.data.ray_hits_w if ray_hits is None: terrain_height = torch.zeros(asset.data.body_pos_w.shape[0], device=asset.data.body_pos_w.device) else: # use the highest hit point in world frame as terrain height reference terrain_height = torch.max(ray_hits[..., 2], dim=-1).values foot_pos_w = asset.data.body_pos_w[:, asset_cfg.body_ids, 2] foot_vel_xy = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2] foot_clearance = torch.clamp(foot_pos_w - ankle_height - terrain_height.unsqueeze(1), min=0.0, max=clearance_clip) speed_term = torch.square(torch.norm(foot_vel_xy, dim=-1)) height_term = torch.exp(-foot_clearance / max(base_height_target, 1e-6) * scale) reward = torch.sum(height_term * speed_term, dim=1) return reward def feet_landing_velocity( env, asset_cfg: SceneEntityCfg, sensor_cfg: SceneEntityCfg, velocity_threshold: float = 0.25, power: float = 2.0, ) -> torch.Tensor: """Penalize high downward landing speed at first contact to reduce impact noise.""" contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] first_contact = _to_env_device(env, contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]) asset = env.scene[asset_cfg.name] foot_vel_z = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, 2] landing_speed = torch.clamp(-foot_vel_z - velocity_threshold, min=0.0) penalty = torch.sum(torch.pow(landing_speed, power) * first_contact, dim=1) return penalty def feet_z_velocity( env, asset_cfg: SceneEntityCfg, sensor_cfg: SceneEntityCfg, power: float = 2.0, ) -> torch.Tensor: asset = env.scene[asset_cfg.name] foot_vel_z = asset.data.body_lin_vel_w[:, sensor_cfg.body_ids, 2] z_speed = torch.clamp(-foot_vel_z, min=0.0) penalty = torch.sum(torch.pow(z_speed, power), dim=1) return penalty def foot_sound_suppression( env, asset_cfg: SceneEntityCfg, sensor_cfg: SceneEntityCfg, max_delta: float = 0.5, ) -> torch.Tensor: """Penalize frame-to-frame changes in foot vertical velocity to suppress landing sounds.""" asset = env.scene[asset_cfg.name] foot_vel_z = asset.data.body_lin_vel_w[:, sensor_cfg.body_ids, 2] try: prev_vel = env._prev_foot_z_vel except AttributeError: env._prev_foot_z_vel = foot_vel_z.clone() return foot_vel_z.new_zeros(foot_vel_z.shape[0]) delta = foot_vel_z - prev_vel delta = torch.clamp(delta, min=-max_delta, max=max_delta) penalty = torch.sum(delta * delta, dim=1) prev_vel.copy_(foot_vel_z) reset_env_ids = env.reset_buf.nonzero(as_tuple=False).squeeze(-1) if reset_env_ids.numel() > 0: prev_vel[reset_env_ids] = foot_vel_z[reset_env_ids] penalty[reset_env_ids] = 0.0 return penalty def base_height_tracking(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), target_height: float = 0.82) -> torch.Tensor: """Reward keeping the base height near a target height.""" asset = env.scene[asset_cfg.name] height_error = torch.abs(asset.data.root_pos_w[:, 2] - target_height) return torch.exp(-height_error * 30.0) def body_height_tracking( env, asset_cfg: SceneEntityCfg, target_height: float, scale: float = 30.0, ) -> torch.Tensor: """Reward keeping a configured body height near a target height.""" body_id = _single_body_id(asset_cfg, "body_height_tracking") asset = env.scene[asset_cfg.name] height_error = torch.abs(asset.data.body_pos_w[:, body_id, 2] - target_height) return torch.exp(-height_error * scale) def body_vertical_velocity_l2( env, asset_cfg: SceneEntityCfg, deadband: float = 0.0, ) -> torch.Tensor: """Penalize vertical body velocity outside a small natural-motion deadband.""" body_id = _single_body_id(asset_cfg, "body_vertical_velocity_l2") asset = env.scene[asset_cfg.name] vertical_speed = torch.abs(asset.data.body_lin_vel_w[:, body_id, 2]) return torch.square(torch.clamp(vertical_speed - deadband, min=0.0)) def body_roll_pitch_ang_vel_l2( env, asset_cfg: SceneEntityCfg, deadband: float = 0.0, ) -> torch.Tensor: """Penalize horizontal angular speed while allowing normal gait oscillation.""" body_id = _single_body_id(asset_cfg, "body_roll_pitch_ang_vel_l2") asset = env.scene[asset_cfg.name] horizontal_ang_speed = torch.linalg.norm(asset.data.body_ang_vel_w[:, body_id, :2], dim=1) return torch.square(torch.clamp(horizontal_ang_speed - deadband, min=0.0)) def body_forward_velocity_undertracking_l1( env, command_name: str, asset_cfg: SceneEntityCfg, heading_yaw_offset: float = 0.0, min_command_speed: float = 0.1, ) -> torch.Tensor: """Penalize only forward-speed shortfall to retain a dense signal at high commands.""" body_id = _single_body_id(asset_cfg, "body_forward_velocity_undertracking_l1") asset = env.scene[asset_cfg.name] commands = env.command_manager.get_command(command_name) body_quat_w = asset.data.body_quat_w[:, body_id, :] body_lin_vel_w = asset.data.body_lin_vel_w[:, body_id, :] body_lin_vel_heading = quat_apply_inverse( _heading_quat_with_offset(body_quat_w, heading_yaw_offset), body_lin_vel_w ) speed_shortfall = torch.clamp(commands[:, 0] - body_lin_vel_heading[:, 0], min=0.0) return speed_shortfall * (commands[:, 0] >= min_command_speed) def body_forward_velocity_undertracking_huber( env, command_name: str, asset_cfg: SceneEntityCfg, heading_yaw_offset: float = 0.0, min_command_speed: float = 0.1, delta: float = 0.25, ) -> torch.Tensor: """Penalize forward-speed shortfall with a normalized smooth-L1 loss.""" if delta <= 0.0: raise ValueError("delta must be positive.") body_id = _single_body_id(asset_cfg, "body_forward_velocity_undertracking_huber") asset = env.scene[asset_cfg.name] commands = env.command_manager.get_command(command_name) body_lin_vel_heading = quat_apply_inverse( _heading_quat_with_offset(asset.data.body_quat_w[:, body_id, :], heading_yaw_offset), asset.data.body_lin_vel_w[:, body_id, :], ) speed_shortfall = torch.clamp(commands[:, 0] - body_lin_vel_heading[:, 0], min=0.0) loss = torch.where( speed_shortfall <= delta, 0.5 * torch.square(speed_shortfall) / delta, speed_shortfall - 0.5 * delta, ) return loss * (commands[:, 0] >= min_command_speed) def joint_deviation_l1_with_deadband( env, asset_cfg: SceneEntityCfg, deadband: float = 0.0, ) -> torch.Tensor: """Penalize joint deviation from defaults while allowing a small neutral range.""" if deadband < 0.0: raise ValueError("deadband must be non-negative.") asset = env.scene[asset_cfg.name] joint_error = torch.abs( asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids] ) return torch.sum(torch.clamp(joint_error - deadband, min=0.0), dim=1) def _forward_swing_speed_scale( commands: torch.Tensor, min_forward_speed: float, full_swing_speed: float, min_swing_scale: float, ) -> torch.Tensor: """Scale arm swing with commanded forward speed while keeping slow walking visible.""" if full_swing_speed <= min_forward_speed: raise ValueError("full_swing_speed must be greater than min_forward_speed.") if not 0.0 <= min_swing_scale <= 1.0: raise ValueError("min_swing_scale must be in [0, 1].") forward_speed = torch.abs(commands[:, 0]) speed_progress = torch.clamp( (forward_speed - min_forward_speed) / (full_swing_speed - min_forward_speed), min=0.0, max=1.0, ) moving = forward_speed >= min_forward_speed moving_scale = min_swing_scale + (1.0 - min_swing_scale) * speed_progress return torch.where(moving, moving_scale, torch.zeros_like(moving_scale)) def cross_body_arm_swing_reward( env, arm_asset_cfg: SceneEntityCfg, feet_asset_cfg: SceneEntityCfg, reference_body_cfg: SceneEntityCfg, command_name: str, heading_yaw_offset: float = 0.0, min_forward_speed: float = 0.1, full_swing_speed: float = 0.6, min_swing_scale: float = 0.0, phase_distance: float = 0.25, shoulder_amplitude: float = 0.18, elbow_flexion: float = 0.15, shoulder_phase_signs: tuple[float, float] = (-1.0, -1.0), elbow_flexion_signs: tuple[float, float] = (1.0, -1.0), std: float = 0.2, joint_error_weights: tuple[float, float, float, float] | None = None, ) -> torch.Tensor: """Track speed-scaled shoulder and elbow targets that oppose the leg phase. Arm joints must be ordered as left/right shoulder pitch followed by left/right elbow flexion. Feet must be ordered left then right. """ if phase_distance <= 0.0 or std <= 0.0: raise ValueError("phase_distance and std must be positive.") joint_ids = arm_asset_cfg.joint_ids foot_ids = feet_asset_cfg.body_ids if joint_ids is None or isinstance(joint_ids, slice) or len(joint_ids) != 4: raise ValueError("cross_body_arm_swing_reward requires four ordered arm joint ids.") if foot_ids is None or isinstance(foot_ids, slice) or len(foot_ids) != 2: raise ValueError("cross_body_arm_swing_reward requires two ordered foot body ids.") reference_body_id = _single_body_id(reference_body_cfg, "cross_body_arm_swing_reward") asset = env.scene[arm_asset_cfg.name] feet_asset = env.scene[feet_asset_cfg.name] reference_asset = env.scene[reference_body_cfg.name] feet_pos_w = feet_asset.data.body_pos_w[:, foot_ids, :] reference_pos_w = reference_asset.data.body_pos_w[:, reference_body_id, :] reference_quat_w = reference_asset.data.body_quat_w[:, reference_body_id, :] feet_pos_rel = feet_pos_w - reference_pos_w.unsqueeze(1) num_envs = feet_pos_w.shape[0] heading_quat = _heading_quat_with_offset(reference_quat_w, heading_yaw_offset) heading_quat = heading_quat.unsqueeze(1).expand(-1, 2, -1).reshape(-1, 4) feet_pos_heading = quat_apply_inverse(heading_quat, feet_pos_rel.reshape(-1, 3)).reshape(num_envs, 2, 3) leg_phase = torch.clamp( (feet_pos_heading[:, 0, 0] - feet_pos_heading[:, 1, 0]) / phase_distance, min=-1.0, max=1.0, ) commands = env.command_manager.get_command(command_name) speed_scale = _forward_swing_speed_scale( commands, min_forward_speed=min_forward_speed, full_swing_speed=full_swing_speed, min_swing_scale=min_swing_scale, ) joint_pos = asset.data.joint_pos[:, joint_ids] target_pos = asset.data.default_joint_pos[:, joint_ids].clone() shoulder_signs = joint_pos.new_tensor(shoulder_phase_signs) elbow_signs = joint_pos.new_tensor(elbow_flexion_signs) target_pos[:, :2] += ( shoulder_amplitude * speed_scale * leg_phase ).unsqueeze(1) * shoulder_signs.unsqueeze(0) target_pos[:, 2:] += (elbow_flexion * speed_scale).unsqueeze(1) * elbow_signs.unsqueeze(0) square_error = torch.square(joint_pos - target_pos) if joint_error_weights is None: mean_square_error = torch.mean(square_error, dim=1) else: if any(weight < 0.0 for weight in joint_error_weights) or sum(joint_error_weights) <= 0.0: raise ValueError("joint_error_weights must be non-negative with a positive sum.") error_weights = joint_pos.new_tensor(joint_error_weights) mean_square_error = torch.sum(square_error * error_weights.unsqueeze(0), dim=1) / sum( joint_error_weights ) return torch.exp(-mean_square_error / (std * std)) def running_elbow_flexion_huber( env, arm_asset_cfg: SceneEntityCfg, feet_asset_cfg: SceneEntityCfg, reference_body_cfg: SceneEntityCfg, command_name: str, heading_yaw_offset: float = 0.0, run_transition_start: float = 1.6, full_run_speed: float = 3.0, min_elbow_flexion: float = 0.2, max_elbow_flexion: float = 0.75, min_phase_modulation: float = 0.05, max_phase_modulation: float = 0.12, phase_distance: float = 0.25, delta: float = 0.2, moving_threshold: float = 0.1, ) -> torch.Tensor: """Track mirrored running elbow flexion with a dense Smooth-L1 loss. The two J4 joints must be ordered left then right. Positive left-J4 and negative right-J4 angles represent the same physical elbow flexion on Gen2. A small phase modulation bends the forward-swinging arm slightly more. """ if full_run_speed <= run_transition_start: raise ValueError("full_run_speed must exceed run_transition_start.") if not 0.0 <= min_elbow_flexion <= max_elbow_flexion: raise ValueError("Expected 0 <= min_elbow_flexion <= max_elbow_flexion.") if not 0.0 <= min_phase_modulation <= max_phase_modulation: raise ValueError("Expected 0 <= min_phase_modulation <= max_phase_modulation.") if phase_distance <= 0.0 or delta <= 0.0: raise ValueError("phase_distance and delta must be positive.") joint_ids = arm_asset_cfg.joint_ids foot_ids = feet_asset_cfg.body_ids if joint_ids is None or isinstance(joint_ids, slice) or len(joint_ids) != 2: raise ValueError("running_elbow_flexion_huber requires two ordered elbow joint ids.") if foot_ids is None or isinstance(foot_ids, slice) or len(foot_ids) != 2: raise ValueError("running_elbow_flexion_huber requires two ordered foot body ids.") reference_body_id = _single_body_id(reference_body_cfg, "running_elbow_flexion_huber") asset = env.scene[arm_asset_cfg.name] feet_asset = env.scene[feet_asset_cfg.name] reference_asset = env.scene[reference_body_cfg.name] reference_pos_w = reference_asset.data.body_pos_w[:, reference_body_id, :] reference_quat_w = reference_asset.data.body_quat_w[:, reference_body_id, :] feet_pos_rel = feet_asset.data.body_pos_w[:, foot_ids, :] - reference_pos_w.unsqueeze(1) num_envs = reference_pos_w.shape[0] heading_quat = _heading_quat_with_offset(reference_quat_w, heading_yaw_offset) heading_quat_pairs = heading_quat.unsqueeze(1).expand(-1, 2, -1).reshape(-1, 4) feet_pos_heading = quat_apply_inverse(heading_quat_pairs, feet_pos_rel.reshape(-1, 3)).reshape( num_envs, 2, 3 ) leg_phase = torch.clamp( (feet_pos_heading[:, 0, 0] - feet_pos_heading[:, 1, 0]) / phase_distance, min=-1.0, max=1.0, ) commands = env.command_manager.get_command(command_name) forward_speed = torch.abs(commands[:, 0]) run_progress = torch.clamp( (forward_speed - run_transition_start) / (full_run_speed - run_transition_start), min=0.0, max=1.0, ) base_flexion = min_elbow_flexion + run_progress * ( max_elbow_flexion - min_elbow_flexion ) phase_modulation = min_phase_modulation + run_progress * ( max_phase_modulation - min_phase_modulation ) target_offset = torch.stack( ( base_flexion - phase_modulation * leg_phase, -(base_flexion + phase_modulation * leg_phase), ), dim=1, ) moving = forward_speed >= moving_threshold target_offset = torch.where(moving.unsqueeze(1), target_offset, torch.zeros_like(target_offset)) target_pos = asset.data.default_joint_pos[:, joint_ids] + target_offset abs_error = torch.abs(asset.data.joint_pos[:, joint_ids] - target_pos) huber = torch.where( abs_error <= delta, 0.5 * torch.square(abs_error) / delta, abs_error - 0.5 * delta, ) return torch.mean(huber, dim=1) def running_arm_target_curriculum( env, env_ids, warmup_steps: int, ramp_end_steps: int, initial_shoulder_amplitude: float, final_shoulder_amplitude: float, initial_upper_arm_phase_amplitude: float, final_upper_arm_phase_amplitude: float, initial_min_elbow_flexion: float, final_min_elbow_flexion: float, initial_max_elbow_flexion: float, final_max_elbow_flexion: float, initial_min_phase_modulation: float, final_min_phase_modulation: float, initial_max_phase_modulation: float, final_max_phase_modulation: float, initial_forearm_phase_amplitude: float, final_forearm_phase_amplitude: float, initial_cross_body_weight: float = 0.8, final_cross_body_weight: float = 1.0, initial_upper_arm_phase_weight: float = 0.5, final_upper_arm_phase_weight: float = 0.65, initial_elbow_weight: float = -0.25, final_elbow_weight: float = -0.5, initial_forearm_phase_weight: float = 0.0, final_forearm_phase_weight: float = 0.3, cross_body_term_name: str = "cross_body_arm_swing", upper_arm_phase_term_name: str = "contralateral_arm_phase", elbow_term_name: str = "elbow_flexion", forearm_phase_term_name: str = "contralateral_forearm_phase", ) -> dict[str, float]: """Progressively introduce natural running arm targets after checkpoint transfer. Reward-manager configs are updated only when an environment resets. With randomized episode lengths this produces a smooth population-level transition instead of an abrupt reward change on the first fine-tuning iteration. """ del env_ids if warmup_steps < 0 or ramp_end_steps <= warmup_steps: raise ValueError("Expected 0 <= warmup_steps < ramp_end_steps.") progress = min( max((env.common_step_counter - warmup_steps) / (ramp_end_steps - warmup_steps), 0.0), 1.0, ) def lerp(initial: float, final: float) -> float: return float(initial + progress * (final - initial)) cross_body_cfg = env.reward_manager.get_term_cfg(cross_body_term_name) cross_body_cfg.params["shoulder_amplitude"] = lerp( initial_shoulder_amplitude, final_shoulder_amplitude ) cross_body_cfg.weight = lerp(initial_cross_body_weight, final_cross_body_weight) env.reward_manager.set_term_cfg(cross_body_term_name, cross_body_cfg) upper_arm_phase_cfg = env.reward_manager.get_term_cfg(upper_arm_phase_term_name) upper_arm_phase_cfg.params["arm_phase_amplitude"] = lerp( initial_upper_arm_phase_amplitude, final_upper_arm_phase_amplitude ) upper_arm_phase_cfg.weight = lerp( initial_upper_arm_phase_weight, final_upper_arm_phase_weight ) env.reward_manager.set_term_cfg(upper_arm_phase_term_name, upper_arm_phase_cfg) elbow_cfg = env.reward_manager.get_term_cfg(elbow_term_name) elbow_cfg.params["min_elbow_flexion"] = lerp( initial_min_elbow_flexion, final_min_elbow_flexion ) elbow_cfg.params["max_elbow_flexion"] = lerp( initial_max_elbow_flexion, final_max_elbow_flexion ) elbow_cfg.params["min_phase_modulation"] = lerp( initial_min_phase_modulation, final_min_phase_modulation ) elbow_cfg.params["max_phase_modulation"] = lerp( initial_max_phase_modulation, final_max_phase_modulation ) elbow_cfg.weight = lerp(initial_elbow_weight, final_elbow_weight) env.reward_manager.set_term_cfg(elbow_term_name, elbow_cfg) forearm_phase_cfg = env.reward_manager.get_term_cfg(forearm_phase_term_name) forearm_phase_cfg.params["arm_phase_amplitude"] = lerp( initial_forearm_phase_amplitude, final_forearm_phase_amplitude ) forearm_phase_cfg.weight = lerp( initial_forearm_phase_weight, final_forearm_phase_weight ) env.reward_manager.set_term_cfg(forearm_phase_term_name, forearm_phase_cfg) return { "progress": float(progress), "shoulder_amplitude": cross_body_cfg.params["shoulder_amplitude"], "max_elbow_flexion": elbow_cfg.params["max_elbow_flexion"], "forearm_phase_amplitude": forearm_phase_cfg.params["arm_phase_amplitude"], } def contralateral_arm_phase_reward( env, arm_body_cfg: SceneEntityCfg, feet_asset_cfg: SceneEntityCfg, reference_body_cfg: SceneEntityCfg, command_name: str, heading_yaw_offset: float = 0.0, min_forward_speed: float = 0.1, full_swing_speed: float = 0.6, min_swing_scale: float = 0.0, phase_distance: float = 0.25, arm_phase_amplitude: float = 0.1, std: float = 0.05, ) -> torch.Tensor: """Reward physical left/right arm motion that is opposite to the same-side leg. Arm bodies and feet must both be ordered left then right. When the left foot is ahead, the right arm body is expected to be ahead of the left arm body, and vice versa. Using upper-arm/elbow bodies prevents wrist motion from satisfying the phase target. """ if phase_distance <= 0.0 or arm_phase_amplitude <= 0.0 or std <= 0.0: raise ValueError("phase_distance, arm_phase_amplitude, and std must be positive.") arm_body_ids = arm_body_cfg.body_ids foot_ids = feet_asset_cfg.body_ids if arm_body_ids is None or isinstance(arm_body_ids, slice) or len(arm_body_ids) != 2: raise ValueError("contralateral_arm_phase_reward requires two ordered arm body ids.") if foot_ids is None or isinstance(foot_ids, slice) or len(foot_ids) != 2: raise ValueError("contralateral_arm_phase_reward requires two ordered foot body ids.") reference_body_id = _single_body_id(reference_body_cfg, "contralateral_arm_phase_reward") arm_asset = env.scene[arm_body_cfg.name] feet_asset = env.scene[feet_asset_cfg.name] reference_asset = env.scene[reference_body_cfg.name] reference_pos_w = reference_asset.data.body_pos_w[:, reference_body_id, :] reference_quat_w = reference_asset.data.body_quat_w[:, reference_body_id, :] heading_quat = _heading_quat_with_offset(reference_quat_w, heading_yaw_offset) arm_pos_rel = arm_asset.data.body_pos_w[:, arm_body_ids, :] - reference_pos_w.unsqueeze(1) feet_pos_rel = feet_asset.data.body_pos_w[:, foot_ids, :] - reference_pos_w.unsqueeze(1) num_envs = reference_pos_w.shape[0] heading_quat_pairs = heading_quat.unsqueeze(1).expand(-1, 2, -1).reshape(-1, 4) arm_pos_heading = quat_apply_inverse(heading_quat_pairs, arm_pos_rel.reshape(-1, 3)).reshape(num_envs, 2, 3) feet_pos_heading = quat_apply_inverse(heading_quat_pairs, feet_pos_rel.reshape(-1, 3)).reshape(num_envs, 2, 3) leg_phase = torch.clamp( (feet_pos_heading[:, 0, 0] - feet_pos_heading[:, 1, 0]) / phase_distance, min=-1.0, max=1.0, ) commands = env.command_manager.get_command(command_name) speed_scale = _forward_swing_speed_scale( commands, min_forward_speed=min_forward_speed, full_swing_speed=full_swing_speed, min_swing_scale=min_swing_scale, ) physical_arm_phase = arm_pos_heading[:, 1, 0] - arm_pos_heading[:, 0, 0] target_arm_phase = arm_phase_amplitude * speed_scale * leg_phase phase_error = physical_arm_phase - target_arm_phase return torch.exp(-torch.square(phase_error) / (std * std)) def energy_cost(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Penalize energy consumption approximated by the sum of squared joint torques.""" asset = env.scene[asset_cfg.name] joint_torques = asset.data.applied_torque[:, :] joint_vel = asset.data.joint_vel[:, :] power = joint_torques * joint_vel energy = torch.sum(torch.abs(power), dim=1) return energy def feet_orientation(env, asset_cfg: SceneEntityCfg, command_name: str, stand_threshold: float = 0.06) -> torch.Tensor: """Reward aligning feet orientation; ignore yaw error while turning.""" commands = env.command_manager.get_command(command_name) yaw_command = torch.abs(commands[:, 2]) > stand_threshold asset = env.scene[asset_cfg.name] feet_quat = asset.data.body_quat_w[:, asset_cfg.body_ids, :] base_quat = asset.data.root_quat_w num_envs, num_feet, _ = feet_quat.shape feet_flat = feet_quat.reshape(-1, 4) roll, pitch, yaw = euler_xyz_from_quat(feet_flat) roll = roll.reshape(num_envs, num_feet) pitch = pitch.reshape(num_envs, num_feet) yaw = yaw.reshape(num_envs, num_feet) _, _, base_yaw = euler_xyz_from_quat(base_quat) feet_roll_pitch_error = torch.sum(torch.abs(torch.stack((roll, pitch), dim=-1)), dim=-1) feet_yaw_error = torch.abs(wrap_to_pi(yaw - base_yaw.unsqueeze(1))) rew = torch.sum(feet_roll_pitch_error + feet_yaw_error, dim=1) rew[yaw_command] = torch.sum(feet_roll_pitch_error[yaw_command], dim=1) return torch.exp(-rew * 2.0) def feet_orientation_relative_to_body( env, asset_cfg: SceneEntityCfg, reference_body_cfg: SceneEntityCfg, command_name: str, stand_threshold: float = 0.06, heading_yaw_offset: float = 0.0, foot_frame_offsets_rpy: tuple[tuple[float, float, float], ...] | None = None, scale: float = 2.0, ) -> torch.Tensor: """Reward physical sole orientation relative to a configured reference body's heading.""" reference_body_id = _single_body_id(reference_body_cfg, "feet_orientation_relative_to_body") commands = env.command_manager.get_command(command_name) yaw_command = torch.abs(commands[:, 2]) > stand_threshold asset = env.scene[asset_cfg.name] reference_asset = env.scene[reference_body_cfg.name] feet_quat = asset.data.body_quat_w[:, asset_cfg.body_ids, :] reference_quat = reference_asset.data.body_quat_w[:, reference_body_id, :] num_envs, num_feet, _ = feet_quat.shape feet_flat = feet_quat.reshape(-1, 4) if foot_frame_offsets_rpy is not None: offsets = torch.tensor(foot_frame_offsets_rpy, dtype=feet_quat.dtype, device=feet_quat.device) if offsets.shape != (num_feet, 3): raise ValueError(f"foot_frame_offsets_rpy must have shape ({num_feet}, 3), got {tuple(offsets.shape)}.") offset_quat = quat_from_euler_xyz(offsets[:, 0], offsets[:, 1], offsets[:, 2]) offset_quat = offset_quat.unsqueeze(0).expand(num_envs, -1, -1).reshape(-1, 4) feet_flat = quat_mul(feet_flat, offset_quat) roll, pitch, yaw = euler_xyz_from_quat(feet_flat) roll = roll.reshape(num_envs, num_feet) pitch = pitch.reshape(num_envs, num_feet) yaw = yaw.reshape(num_envs, num_feet) reference_yaw = _heading_yaw_with_offset(reference_quat, heading_yaw_offset) feet_roll_pitch_error = torch.sum(torch.abs(torch.stack((roll, pitch), dim=-1)), dim=-1) feet_yaw_error = torch.abs(wrap_to_pi(yaw - reference_yaw.unsqueeze(1))) rew = torch.sum(feet_roll_pitch_error + feet_yaw_error, dim=1) rew[yaw_command] = torch.sum(feet_roll_pitch_error[yaw_command], dim=1) return torch.exp(-rew * scale) def base_orientation(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Reward keeping the base roll/pitch near zero.""" asset = env.scene[asset_cfg.name] roll, pitch, yaw = euler_xyz_from_quat(asset.data.root_quat_w) base_euler = torch.stack((roll, pitch, yaw), dim=-1) return torch.exp(-torch.sum(torch.abs(base_euler[:, :2]), dim=-1) * 10.0) def body_orientation(env, asset_cfg: SceneEntityCfg, scale: float = 10.0) -> torch.Tensor: """Reward keeping a configured body's roll/pitch near zero.""" body_id = _single_body_id(asset_cfg, "body_orientation") asset = env.scene[asset_cfg.name] roll, pitch, yaw = euler_xyz_from_quat(asset.data.body_quat_w[:, body_id, :]) return torch.exp(-torch.sum(torch.abs(torch.stack((roll, pitch), dim=-1)), dim=-1) * scale) def body_yaw_alignment( env, asset_cfg: SceneEntityCfg, reference_body_cfg: SceneEntityCfg, heading_yaw_offset: float = 0.0, reference_heading_yaw_offset: float = 0.0, scale: float = 4.0, ) -> torch.Tensor: """Reward keeping one body's heading aligned with another body's heading.""" body_id = _single_body_id(asset_cfg, "body_yaw_alignment") reference_body_id = _single_body_id(reference_body_cfg, "body_yaw_alignment") asset = env.scene[asset_cfg.name] reference_asset = env.scene[reference_body_cfg.name] body_yaw = _heading_yaw_with_offset(asset.data.body_quat_w[:, body_id, :], heading_yaw_offset) reference_yaw = _heading_yaw_with_offset( reference_asset.data.body_quat_w[:, reference_body_id, :], reference_heading_yaw_offset ) return torch.exp(-torch.abs(wrap_to_pi(body_yaw - reference_yaw)) * scale) def body_yaw_rate_difference_l2( env, asset_cfg: SceneEntityCfg, reference_body_cfg: SceneEntityCfg, ) -> torch.Tensor: """Penalize relative world-frame yaw rate between two configured bodies.""" body_id = _single_body_id(asset_cfg, "body_yaw_rate_difference_l2") reference_body_id = _single_body_id(reference_body_cfg, "body_yaw_rate_difference_l2") asset = env.scene[asset_cfg.name] reference_asset = env.scene[reference_body_cfg.name] yaw_rate = asset.data.body_ang_vel_w[:, body_id, 2] reference_yaw_rate = reference_asset.data.body_ang_vel_w[:, reference_body_id, 2] return torch.square(yaw_rate - reference_yaw_rate) def reward_waist_pos( env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), joint_names: list[str] | str | None = None, scale: float = 5.0, ) -> torch.Tensor: """Penalize motion on specific waist joints; joint_names provided by user.""" asset = env.scene[asset_cfg.name] if joint_names is not None: joint_ids = asset.find_joints(joint_names, preserve_order=True)[0] elif asset_cfg.joint_ids is not None: joint_ids = asset_cfg.joint_ids else: raise ValueError("penalize_waist_joint_motion requires joint_names or asset_cfg.joint_ids.") pos_err = torch.abs(asset.data.joint_pos[:, joint_ids]) return torch.exp(-torch.sum(pos_err, dim=1) * scale) def penalize_foot_stumble(env, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] contacts = _to_env_device( env, contact_sensor.data.net_forces_w_history[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] ) > 1.0 asset = env.scene[asset_cfg.name] body_vel = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2] return torch.sum(body_vel.norm(dim=-1) * contacts, dim=1) def joint_deviation_exp( env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), tolerance: float = 0.1, scale: float = 3.0, max_err: float = 50.0, ) -> torch.Tensor: """Penalize joint positions deviating from defaults, in an exponential, configurable way.""" asset: Articulation = env.scene[asset_cfg.name] joint_ids = asset_cfg.joint_ids if asset_cfg.joint_ids is not None else slice(None) joint_pos = asset.data.joint_pos[:, joint_ids] default_pos = getattr(asset.data, "default_joint_pos", None) if default_pos is not None: default_pos = default_pos[:, joint_ids] else: default_pos = torch.zeros_like(joint_pos) print("Warning: joint_deviation_exp reward called but default_joint_pos not set in asset data; assuming zeros.") joint_error = torch.norm(joint_pos - default_pos, dim=1) joint_error = torch.clamp(joint_error - tolerance, min=0.0, max=max_err) return torch.exp(-joint_error * scale)