#!/usr/bin/env python3 """Evaluate a Dropbear checkpoint with deterministic mean actions. This wraps Isaac Lab's stock RSL-RL player, pins the environment to the full command range, and aggregates completed episodes before terminating. Unlike the PPO training dashboard, these results contain no exploration-action noise. """ from __future__ import annotations import json import math import os import runpy import sys import tempfile from pathlib import Path from typing import Any import gymnasium as gym import torch import warp as wp def _early_pop_float_pair_arg(name: str) -> tuple[float, float] | None: """Consume an environment-shaping pair before task registration.""" if name not in sys.argv: return None index = sys.argv.index(name) try: low = float(sys.argv[index + 1]) high = float(sys.argv[index + 2]) except (IndexError, ValueError) as exc: raise ValueError(f"{name} requires two numbers") from exc del sys.argv[index : index + 3] return low, high def _early_pop_float_arg(name: str) -> float | None: """Consume an environment-shaping scalar before task registration.""" if name not in sys.argv: return None index = sys.argv.index(name) try: value = float(sys.argv[index + 1]) except (IndexError, ValueError) as exc: raise ValueError(f"{name} requires a number") from exc del sys.argv[index : index + 2] return value ACTOR_BASE_LIN_VEL = "--actor-base-lin-vel" in sys.argv if ACTOR_BASE_LIN_VEL: sys.argv.remove("--actor-base-lin-vel") os.environ["DROPBEAR_ACTOR_BASE_LIN_VEL"] = "1" ADAPT_COM_CHECKPOINT = "--eval-adapt-com-checkpoint" in sys.argv if ADAPT_COM_CHECKPOINT: sys.argv.remove("--eval-adapt-com-checkpoint") RESET_JOINT_POSITION_RANGE = _early_pop_float_pair_arg( "--eval-reset-joint-position-range" ) RESET_JOINT_VELOCITY_RANGE = _early_pop_float_pair_arg( "--eval-reset-joint-velocity-range" ) if RESET_JOINT_POSITION_RANGE is not None: reset_min, reset_max = RESET_JOINT_POSITION_RANGE if not 0.1 <= reset_min <= reset_max <= 2.0: raise ValueError( "--eval-reset-joint-position-range must satisfy " "0.1 <= low <= high <= 2.0" ) os.environ["DROPBEAR_RESET_JOINT_POSITION_MIN"] = f"{reset_min:g}" os.environ["DROPBEAR_RESET_JOINT_POSITION_MAX"] = f"{reset_max:g}" if RESET_JOINT_VELOCITY_RANGE is not None: reset_vel_min, reset_vel_max = RESET_JOINT_VELOCITY_RANGE if not -5.0 <= reset_vel_min <= reset_vel_max <= 5.0: raise ValueError( "--eval-reset-joint-velocity-range must satisfy " "-5.0 <= low <= high <= 5.0" ) os.environ["DROPBEAR_RESET_JOINT_VELOCITY_MIN"] = f"{reset_vel_min:g}" os.environ["DROPBEAR_RESET_JOINT_VELOCITY_MAX"] = f"{reset_vel_max:g}" RESET_POLICY_JOINTS_ONLY = "--eval-reset-policy-joints-only" in sys.argv if RESET_POLICY_JOINTS_ONLY: sys.argv.remove("--eval-reset-policy-joints-only") os.environ["DROPBEAR_RESET_POLICY_JOINTS_ONLY"] = "1" POSE_SEQUENCE = "--eval-pose-sequence" in sys.argv if POSE_SEQUENCE: sys.argv.remove("--eval-pose-sequence") os.environ["DROPBEAR_POSE_SEQUENCE"] = "1" COM_CONTROL = "--eval-com-control" in sys.argv if COM_CONTROL: sys.argv.remove("--eval-com-control") os.environ["DROPBEAR_COM_CONTROL"] = "1" GAIT_MIRROR_ACTION_ADAPTER = ( "--eval-gait-mirror-action-adapter" in sys.argv ) if GAIT_MIRROR_ACTION_ADAPTER: sys.argv.remove("--eval-gait-mirror-action-adapter") os.environ["DROPBEAR_GAIT_MIRROR_ACTION_ADAPTER"] = "1" GAIT_MIRROR_TRANSFER_STRENGTH = _early_pop_float_arg( "--eval-gait-mirror-transfer-strength" ) if GAIT_MIRROR_TRANSFER_STRENGTH is not None: if not GAIT_MIRROR_ACTION_ADAPTER: raise ValueError( "--eval-gait-mirror-transfer-strength requires " "--eval-gait-mirror-action-adapter" ) if not 0.0 <= GAIT_MIRROR_TRANSFER_STRENGTH <= 1.0: raise ValueError( "--eval-gait-mirror-transfer-strength must be between 0 and 1" ) os.environ["DROPBEAR_GAIT_MIRROR_TRANSFER_STRENGTH"] = ( f"{GAIT_MIRROR_TRANSFER_STRENGTH:g}" ) POLICY_MIRROR_PROJECTION = _early_pop_float_arg( "--eval-policy-mirror-projection" ) if POLICY_MIRROR_PROJECTION is not None and not ( 0.0 <= POLICY_MIRROR_PROJECTION <= 1.0 ): raise ValueError( "--eval-policy-mirror-projection must be between 0 and 1" ) EVAL_OBSTACLE_TERRAIN = "--eval-obstacle-terrain" in sys.argv if EVAL_OBSTACLE_TERRAIN: sys.argv.remove("--eval-obstacle-terrain") os.environ["DROPBEAR_OBSTACLE_TERRAIN"] = "1" os.environ["DROPBEAR_TERRAIN_HEIGHT_SCAN"] = "1" EVAL_GAIT_PERIOD = _early_pop_float_arg("--eval-gait-period") if EVAL_GAIT_PERIOD is None: EVAL_GAIT_PERIOD = 0.60 if not 0.20 <= EVAL_GAIT_PERIOD <= 2.0: raise ValueError("--eval-gait-period must be between 0.20 and 2.0 seconds") os.environ["DROPBEAR_GAIT_PERIOD"] = f"{EVAL_GAIT_PERIOD:g}" for cli_name, env_name, minimum, maximum in ( ("--eval-com-stand-height", "DROPBEAR_COM_STAND_HEIGHT", 0.0, 3.0), ("--eval-com-height-delta", "DROPBEAR_COM_HEIGHT_DELTA", 0.0, 1.0), ( "--eval-com-height-error-scale", "DROPBEAR_COM_HEIGHT_ERROR_SCALE", 0.001, 1.0, ), ( "--eval-com-vertical-velocity-error-scale", "DROPBEAR_COM_VERTICAL_VELOCITY_ERROR_SCALE", 0.001, 5.0, ), ): value = _early_pop_float_arg(cli_name) if value is not None: if not minimum <= value <= maximum: raise ValueError( f"{cli_name} must be between {minimum:g} and {maximum:g}" ) os.environ[env_name] = f"{value:g}" POSE_RESIDUAL_SCALE = _early_pop_float_arg("--eval-pose-residual-scale") POSE_PG_PITCH_OFFSET = _early_pop_float_arg("--eval-pose-pg-pitch-offset") POSE_KNEE_OFFSET = _early_pop_float_arg("--eval-pose-knee-offset") POSE_ANKLE67_OFFSET = _early_pop_float_arg("--eval-pose-ankle67-offset") for value, env_name, cli_name in ( ( POSE_PG_PITCH_OFFSET, "DROPBEAR_POSE_PG_PITCH_OFFSET", "--eval-pose-pg-pitch-offset", ), ( POSE_KNEE_OFFSET, "DROPBEAR_POSE_KNEE_OFFSET", "--eval-pose-knee-offset", ), ( POSE_ANKLE67_OFFSET, "DROPBEAR_POSE_ANKLE67_OFFSET", "--eval-pose-ankle67-offset", ), ): if value is not None: if not -2.0 <= value <= 2.0: raise ValueError(f"{cli_name} must be between -2 and 2") os.environ[env_name] = f"{value:g}" POSE_BASELINE_DEPTH = _early_pop_float_arg("--eval-pose-baseline-depth") if POSE_BASELINE_DEPTH is None: POSE_BASELINE_DEPTH = -0.33 if not -2.0 <= POSE_BASELINE_DEPTH <= 2.0: raise ValueError("--eval-pose-baseline-depth must be between -2 and 2") os.environ["DROPBEAR_POSE_BASELINE_DEPTH"] = f"{POSE_BASELINE_DEPTH:g}" POSE_DEPTH_AMPLITUDE = _early_pop_float_arg("--eval-pose-depth-amplitude") if POSE_DEPTH_AMPLITUDE is None: POSE_DEPTH_AMPLITUDE = 1.0 if not 0.0 <= POSE_DEPTH_AMPLITUDE <= 2.0: raise ValueError("--eval-pose-depth-amplitude must be between 0 and 2") os.environ["DROPBEAR_POSE_DEPTH_AMPLITUDE"] = f"{POSE_DEPTH_AMPLITUDE:g}" POSE_STAND_HEIGHT = _early_pop_float_arg("--eval-pose-stand-height") if POSE_STAND_HEIGHT is None: POSE_STAND_HEIGHT = 1.62 POSE_CROUCH_HEIGHT_DELTA = _early_pop_float_arg( "--eval-pose-crouch-height-delta" ) if POSE_CROUCH_HEIGHT_DELTA is None: POSE_CROUCH_HEIGHT_DELTA = 0.18 POSE_ACTION_RESIDUAL = "--eval-pose-action-residual" in sys.argv if POSE_ACTION_RESIDUAL: sys.argv.remove("--eval-pose-action-residual") if not POSE_SEQUENCE: raise ValueError( "--eval-pose-action-residual requires --eval-pose-sequence" ) os.environ["DROPBEAR_POSE_ACTION_RESIDUAL"] = "1" if POSE_RESIDUAL_SCALE is not None: if not 0.0 <= POSE_RESIDUAL_SCALE <= 1.0: raise ValueError("--eval-pose-residual-scale must be between 0 and 1") os.environ["DROPBEAR_POSE_RESIDUAL_SCALE"] = f"{POSE_RESIDUAL_SCALE:g}" elif POSE_RESIDUAL_SCALE is not None: raise ValueError( "--eval-pose-residual-scale requires --eval-pose-action-residual" ) POSE_BIAS_GRID = "--eval-pose-bias-grid" in sys.argv if POSE_BIAS_GRID: sys.argv.remove("--eval-pose-bias-grid") if not POSE_SEQUENCE: raise ValueError("--eval-pose-bias-grid requires --eval-pose-sequence") import dropbear_walk # noqa: F401 from isaaclab.utils.math import quat_apply_inverse, yaw_quat from dropbear_walk.isaaclab_asset.dropbear import DISTAL_FOOT_BODIES from dropbear_walk.mdp.observations import ( com_height_reference, mass_weighted_com_state, pose_sequence_reference, ) from dropbear_walk.mdp.symmetry import ( compute_symmetric_states, mirror_dropbear_joints, ) from rsl_rl.runners import OnPolicyRunner WORKSPACE_ROOT = Path(__file__).resolve().parents[1] PLAY_SCRIPT = WORKSPACE_ROOT / "IsaacLab" / "scripts" / "reinforcement_learning" / "rsl_rl" / "play.py" ORIGINAL_GYM_MAKE = gym.make if POLICY_MIRROR_PROJECTION is not None: _original_get_inference_policy = OnPolicyRunner.get_inference_policy class _MirrorProjectedPolicy: """Blend policy output toward exact left/right equivariance.""" def __init__(self, runner, policy, strength: float): self._runner = runner self._policy = policy self._strength = strength def __call__(self, observations): original_actions = self._policy(observations) mirrored_observations, _ = compute_symmetric_states( self._runner.env, observations, None, ) batch_size = observations.batch_size[0] mirrored_actions = self._policy( mirrored_observations[batch_size:] ) projected_actions = 0.5 * ( original_actions + mirror_dropbear_joints(mirrored_actions) ) return torch.lerp( original_actions, projected_actions, self._strength, ) def reset(self, dones=None): return self._policy.reset(dones) def _get_mirror_projected_policy(self, device=None): policy = _original_get_inference_policy(self, device=device) return _MirrorProjectedPolicy( self, policy, POLICY_MIRROR_PROJECTION, ) OnPolicyRunner.get_inference_policy = _get_mirror_projected_policy if ADAPT_COM_CHECKPOINT: _original_runner_load = OnPolicyRunner.load def _load_with_com_adapter( self, path, load_cfg=None, strict=True, map_location=None, ): checkpoint = torch.load(path, weights_only=False, map_location="cpu") def _append(state, expected_width: int, model_name: str) -> None: first_layer = state["mlp.0.weight"] old_width = int(first_layer.shape[1]) if old_width == expected_width: return if old_width + 2 != expected_width: raise ValueError( f"{model_name} COM adapter expected {old_width}+2=" f"{expected_width} inputs" ) state["mlp.0.weight"] = torch.cat( ( first_layer, torch.zeros( first_layer.shape[0], 2, dtype=first_layer.dtype, ), ), dim=1, ) for key in ( "obs_normalizer._mean", "obs_normalizer._var", "obs_normalizer._std", ): tensor = state[key] fill = 0.0 if key.endswith("_mean") else 1.0 state[key] = torch.cat( ( tensor, torch.full( (tensor.shape[0], 2), fill, dtype=tensor.dtype, ), ), dim=1, ) _append( checkpoint["actor_state_dict"], int(self.alg.actor.mlp[0].in_features), "Actor", ) _append( checkpoint["critic_state_dict"], int(self.alg.critic.mlp[0].in_features), "Critic", ) temporary_path: Path | None = None try: with tempfile.NamedTemporaryFile( suffix=".pt", delete=False ) as handle: temporary_path = Path(handle.name) torch.save(checkpoint, temporary_path) return _original_runner_load( self, str(temporary_path), load_cfg={ "actor": True, "critic": True, "optimizer": False, "iteration": True, "rnd": False, }, strict=strict, map_location=map_location, ) finally: if temporary_path is not None: temporary_path.unlink(missing_ok=True) OnPolicyRunner.load = _load_with_com_adapter def _pop_int_arg(name: str, default: int) -> int: if name not in sys.argv: return default index = sys.argv.index(name) try: value = int(sys.argv[index + 1]) except (IndexError, ValueError) as exc: raise ValueError(f"{name} requires an integer") from exc del sys.argv[index : index + 2] return value def _pop_float_arg(name: str) -> float | None: if name not in sys.argv: return None index = sys.argv.index(name) try: value = float(sys.argv[index + 1]) except (IndexError, ValueError) as exc: raise ValueError(f"{name} requires a number") from exc del sys.argv[index : index + 2] return value def _pop_path_arg(name: str) -> Path | None: if name not in sys.argv: return None index = sys.argv.index(name) try: value = Path(sys.argv[index + 1]) except IndexError as exc: raise ValueError(f"{name} requires a path") from exc del sys.argv[index : index + 2] return value TARGET_EPISODES = _pop_int_arg("--eval-episodes", 1024) if TARGET_EPISODES <= 0: raise ValueError("--eval-episodes must be positive") WARMUP_EPISODES = _pop_int_arg("--eval-warmup-episodes", 1) if WARMUP_EPISODES < 0: raise ValueError("--eval-warmup-episodes cannot be negative") OUTPUT_PATH = _pop_path_arg("--eval-output") FORWARD_SPEED = _pop_float_arg("--eval-forward-speed") BACKWARD_SPEED = _pop_float_arg("--eval-backward-speed") LATERAL_SPEED = _pop_float_arg("--eval-lateral-speed") if LATERAL_SPEED is not None and not -2.0 <= LATERAL_SPEED <= 2.0: raise ValueError("--eval-lateral-speed must be between -2 and 2 m/s") if BACKWARD_SPEED is not None and not -2.0 <= BACKWARD_SPEED < 0.0: raise ValueError("--eval-backward-speed must be negative and at least -2 m/s") YAW_RATE = _pop_float_arg("--eval-yaw-rate") if YAW_RATE is not None and not -3.0 <= YAW_RATE <= 3.0: raise ValueError("--eval-yaw-rate must be between -3 and 3 rad/s") PLANAR_CARDINAL_SWEEP = "--eval-planar-cardinal-sweep" in sys.argv if PLANAR_CARDINAL_SWEEP: sys.argv.remove("--eval-planar-cardinal-sweep") if FORWARD_SPEED is None or FORWARD_SPEED <= 0.0: raise ValueError( "--eval-planar-cardinal-sweep requires a positive " "--eval-forward-speed" ) if BACKWARD_SPEED is None: raise ValueError( "--eval-planar-cardinal-sweep requires --eval-backward-speed" ) if LATERAL_SPEED is None or LATERAL_SPEED <= 0.0: raise ValueError( "--eval-planar-cardinal-sweep requires a positive " "--eval-lateral-speed" ) if TARGET_EPISODES < 4: raise ValueError( "--eval-planar-cardinal-sweep requires at least 4 episodes" ) FEET_MIN_DISTANCE = _pop_float_arg("--eval-feet-min-distance") if FEET_MIN_DISTANCE is None: FEET_MIN_DISTANCE = 0.14 if not 0.01 <= FEET_MIN_DISTANCE <= 1.0: raise ValueError("--eval-feet-min-distance must be between 0.01 and 1 m") FEET_MIN_LATERAL_SEPARATION = _pop_float_arg( "--eval-feet-min-lateral-separation" ) if FEET_MIN_LATERAL_SEPARATION is None: FEET_MIN_LATERAL_SEPARATION = 0.08 if not 0.01 <= FEET_MIN_LATERAL_SEPARATION <= 1.0: raise ValueError( "--eval-feet-min-lateral-separation must be between 0.01 and 1 m" ) PLANE_ONLY = "--eval-plane" in sys.argv if PLANE_ONLY: sys.argv.remove("--eval-plane") DISABLE_PUSHES = "--eval-disable-pushes" in sys.argv if DISABLE_PUSHES: sys.argv.remove("--eval-disable-pushes") PROJECT_RECIPROCAL_SHOULDERS = ( "--eval-project-reciprocal-shoulders" in sys.argv ) if PROJECT_RECIPROCAL_SHOULDERS: sys.argv.remove("--eval-project-reciprocal-shoulders") SHOULDER_COUNTERWEIGHT_SCALE = _pop_float_arg( "--eval-shoulder-counterweight-scale" ) if PROJECT_RECIPROCAL_SHOULDERS: if SHOULDER_COUNTERWEIGHT_SCALE is not None: raise ValueError( "--eval-project-reciprocal-shoulders conflicts with " "--eval-shoulder-counterweight-scale" ) SHOULDER_COUNTERWEIGHT_SCALE = 0.0 if ( SHOULDER_COUNTERWEIGHT_SCALE is not None and not 0.0 <= SHOULDER_COUNTERWEIGHT_SCALE <= 1.0 ): raise ValueError( "--eval-shoulder-counterweight-scale must be between 0 and 1" ) class DeterministicEvalWrapper(gym.Wrapper): """Collect deterministic episode statistics from a vector environment.""" def __init__(self, env: gym.Env): super().__init__(env) base = env.unwrapped self._base = base self._num_envs = base.num_envs self._step_dt = float(base.step_dt) self._max_episode_seconds = float(base.max_episode_length_s) action_term = base.action_manager.get_term("joint_pos") action_index = { name: index for index, name in enumerate(action_term._joint_names) } self._shoulder_action_ids = [ action_index["LH_yaw"], action_index["RH_yaw"], ] self._returns = torch.zeros(self._num_envs, device=base.device) self._tracking = torch.zeros_like(self._returns) self._gait = torch.zeros_like(self._returns) self._body_velocity_error = torch.zeros_like(self._returns) self._yaw_velocity_error = torch.zeros_like(self._returns) self._yaw_rate_error = torch.zeros_like(self._returns) self._com_velocity_error = torch.zeros_like(self._returns) self._root_velocity_yaw_xy = torch.zeros( self._num_envs, 2, device=base.device ) self._com_velocity_yaw_xy = torch.zeros_like( self._root_velocity_yaw_xy ) self._root_yaw_rate = torch.zeros_like(self._returns) self._com_vertical_velocity_error = torch.zeros_like(self._returns) self._com_height_abs_error = torch.zeros_like(self._returns) self._arm_swing_score = torch.zeros_like(self._returns) self._arm_extension_l2 = torch.zeros_like(self._returns) self._arm_reciprocal_mode_abs = torch.zeros_like(self._returns) self._arm_counterweight_mode_abs = torch.zeros_like(self._returns) self._arm_reciprocal_sin = torch.zeros_like(self._returns) self._arm_reciprocal_cos = torch.zeros_like(self._returns) self._arm_counterweight_sin = torch.zeros_like(self._returns) self._arm_counterweight_cos = torch.zeros_like(self._returns) self._foot_phase_velocity_score = torch.zeros_like(self._returns) self._foot_min_distance = torch.full_like( self._returns, float("inf"), ) self._foot_min_lateral_separation = torch.full_like( self._returns, float("inf"), ) self._foot_clearance_violation_steps = torch.zeros_like(self._returns) self._foot_lateral_order_violation_steps = torch.zeros_like( self._returns ) self._severe_foot_overlap_steps = torch.zeros_like(self._returns) self._leg_contact_steps = torch.zeros( self._num_envs, 2, device=base.device ) self._leg_touchdown_count = torch.zeros_like( self._leg_contact_steps ) self._leg_stride_interval_sum = torch.zeros_like( self._leg_contact_steps ) self._leg_stride_interval_count = torch.zeros_like( self._leg_contact_steps ) self._leg_flight_time_sum = torch.zeros_like( self._leg_contact_steps ) self._leg_flight_count = torch.zeros_like( self._leg_contact_steps ) self._last_leg_touchdown_step = torch.full( (self._num_envs, 2), -1, dtype=torch.long, device=base.device, ) self._leg_air_steps = torch.zeros( self._num_envs, 2, dtype=torch.long, device=base.device, ) self._was_leg_contact = torch.zeros( self._num_envs, 2, dtype=torch.bool, device=base.device, ) self._knee_position = torch.zeros( self._num_envs, 2, device=base.device ) self._knee_swing_position = torch.zeros_like(self._knee_position) self._knee_stance_position = torch.zeros_like(self._knee_position) self._knee_swing_steps = torch.zeros_like(self._knee_position) self._knee_stance_steps = torch.zeros_like(self._knee_position) self._knee_min = torch.full_like( self._knee_position, float("inf"), ) self._knee_max = torch.full_like( self._knee_position, float("-inf"), ) self._phase_joint_position = torch.zeros( self._num_envs, 6, device=base.device ) self._phase_joint_swing_position = torch.zeros_like( self._phase_joint_position ) self._phase_joint_stance_position = torch.zeros_like( self._phase_joint_position ) self._phase_joint_swing_steps = torch.zeros_like( self._phase_joint_position ) self._phase_joint_stance_steps = torch.zeros_like( self._phase_joint_position ) self._torso_roll = torch.zeros_like(self._returns) self._torso_roll_abs = torch.zeros_like(self._returns) self._torso_roll_squared = torch.zeros_like(self._returns) self._torso_roll_sin = torch.zeros_like(self._returns) self._torso_roll_cos = torch.zeros_like(self._returns) self._torso_roll_positive_steps = torch.zeros_like(self._returns) self._foot_lateral_min = torch.full( (self._num_envs, 2), float("inf"), device=base.device, ) self._foot_lateral_max = torch.full( (self._num_envs, 2), float("-inf"), device=base.device, ) self._foot_position_yaw_xy = torch.zeros( self._num_envs, 2, 2, device=base.device ) self._late_swing_foot_advance = torch.zeros( self._num_envs, 2, device=base.device ) self._late_swing_foot_steps = torch.zeros_like( self._late_swing_foot_advance ) self._obstacle_crossed = torch.zeros( self._num_envs, dtype=torch.bool, device=base.device ) self._obstacle_progressed_past = torch.zeros_like( self._obstacle_crossed ) self._obstacle_zone_seen = torch.zeros_like(self._obstacle_crossed) self._obstacle_zone_max_foot_height = torch.full( (self._num_envs, 2), float("-inf"), device=base.device, ) self._max_forward_progress = torch.full_like( self._returns, float("-inf"), ) self._pose_joint_mse = torch.zeros_like(self._returns) self._pose_depth_abs_error = torch.zeros_like(self._returns) self._pose_height_abs_error = torch.zeros_like(self._returns) self._stand_achieved_depth = torch.zeros_like(self._returns) self._stand_joint_delta = torch.zeros( self._num_envs, 6, device=base.device ) self._stand_root_height = torch.zeros_like(self._returns) self._stand_com_height = torch.zeros_like(self._returns) self._stand_pose_steps = torch.zeros_like(self._returns) self._crouch_achieved_depth = torch.zeros_like(self._returns) self._crouch_joint_delta = torch.zeros( self._num_envs, 6, device=base.device ) self._crouch_root_height = torch.zeros_like(self._returns) self._crouch_com_height = torch.zeros_like(self._returns) self._crouch_pose_steps = torch.zeros_like(self._returns) self._lengths = torch.zeros(self._num_envs, dtype=torch.long, device=base.device) self._completed_per_env = torch.zeros_like(self._lengths) reward_names = base.reward_manager.active_terms self._tracking_index = reward_names.index("track_lin_vel_xy") self._gait_index = reward_names.index("gait") self._episodes = 0 self._fall_episodes = 0 self._timeout_episodes = 0 self._return_sum = 0.0 self._length_sum = 0 self._tracking_sum = 0.0 self._gait_sum = 0.0 self._body_velocity_error_sum = 0.0 self._yaw_velocity_error_sum = 0.0 self._yaw_rate_error_sum = 0.0 self._com_velocity_error_sum = 0.0 self._root_velocity_yaw_xy_sum = torch.zeros(2, device=base.device) self._com_velocity_yaw_xy_sum = torch.zeros(2, device=base.device) self._root_yaw_rate_sum = 0.0 self._com_vertical_velocity_error_sum = 0.0 self._com_height_abs_error_sum = 0.0 self._arm_swing_score_sum = 0.0 self._arm_extension_l2_sum = 0.0 self._arm_reciprocal_mode_abs_sum = 0.0 self._arm_counterweight_mode_abs_sum = 0.0 self._arm_reciprocal_sin_coefficient_sum = 0.0 self._arm_reciprocal_cos_coefficient_sum = 0.0 self._arm_counterweight_sin_coefficient_sum = 0.0 self._arm_counterweight_cos_coefficient_sum = 0.0 self._foot_phase_velocity_score_sum = 0.0 self._foot_min_distance_sum = 0.0 self._foot_min_lateral_separation_sum = 0.0 self._foot_clearance_violation_fraction_sum = 0.0 self._foot_lateral_order_violation_fraction_sum = 0.0 self._severe_foot_overlap_fraction_sum = 0.0 self._leg_contact_fraction_sum = torch.zeros(2, device=base.device) self._leg_touchdown_count_sum = torch.zeros(2, device=base.device) self._leg_stride_interval_total = torch.zeros(2, device=base.device) self._leg_stride_interval_samples = torch.zeros(2, device=base.device) self._leg_flight_time_total = torch.zeros(2, device=base.device) self._leg_flight_samples = torch.zeros(2, device=base.device) self._flight_time_abs_difference_sum = 0.0 self._contact_duty_abs_difference_sum = 0.0 self._touchdown_count_abs_difference_sum = 0.0 self._knee_position_sum = torch.zeros(2, device=base.device) self._knee_swing_position_sum = torch.zeros(2, device=base.device) self._knee_stance_position_sum = torch.zeros(2, device=base.device) self._knee_range_sum = torch.zeros(2, device=base.device) self._phase_joint_position_sum = torch.zeros(6, device=base.device) self._phase_joint_swing_position_sum = torch.zeros( 6, device=base.device ) self._phase_joint_stance_position_sum = torch.zeros( 6, device=base.device ) self._torso_roll_sum = 0.0 self._torso_roll_abs_sum = 0.0 self._torso_roll_squared_sum = 0.0 self._torso_roll_sin_coefficient_sum = 0.0 self._torso_roll_cos_coefficient_sum = 0.0 self._torso_roll_positive_fraction_sum = 0.0 self._foot_lateral_excursion_sum = torch.zeros( 2, device=base.device ) self._foot_position_yaw_xy_sum = torch.zeros( 2, 2, device=base.device ) self._late_swing_foot_advance_sum = torch.zeros( 2, device=base.device ) self._obstacle_crossed_episodes = 0 self._obstacle_progressed_past_episodes = 0 self._obstacle_zone_seen_episodes = 0 self._obstacle_zone_max_foot_height_sum = torch.zeros( 2, device=base.device ) self._max_forward_progress_sum = 0.0 self._pose_joint_mse_sum = 0.0 self._pose_depth_abs_error_sum = 0.0 self._pose_height_abs_error_sum = 0.0 self._stand_achieved_depth_sum = 0.0 self._stand_joint_delta_sum = torch.zeros(6, device=base.device) self._stand_root_height_sum = 0.0 self._stand_com_height_sum = 0.0 self._stand_pose_step_count = 0.0 self._crouch_achieved_depth_sum = 0.0 self._crouch_joint_delta_sum = torch.zeros(6, device=base.device) self._crouch_root_height_sum = 0.0 self._crouch_com_height_sum = 0.0 self._crouch_pose_step_count = 0.0 self._pose_joint_ids: list[int] = [] self._pose_joint_names: list[str] = [] self._pose_offsets: torch.Tensor | None = None self._pose_height_body_names = [ "torso_RMD_X10__1_Rotor_1", "torso_RMD_X10Rotot_1", ] self._pose_bias_by_env: torch.Tensor | None = None self._pose_bias_group_ids: torch.Tensor | None = None self._pose_bias_names: list[str] = [] self._pose_bias_episode_counts: torch.Tensor | None = None self._pose_bias_fall_counts: torch.Tensor | None = None self._pose_bias_stand_height_sums: torch.Tensor | None = None self._pose_bias_stand_step_sums: torch.Tensor | None = None self._pose_bias_crouch_height_sums: torch.Tensor | None = None self._pose_bias_crouch_step_sums: torch.Tensor | None = None self._cardinal_names = ["forward", "backward", "left", "right"] self._cardinal_group_ids: torch.Tensor | None = None self._cardinal_episode_counts: torch.Tensor | None = None self._cardinal_fall_counts: torch.Tensor | None = None self._cardinal_length_sums: torch.Tensor | None = None self._cardinal_com_velocity_sums: torch.Tensor | None = None self._cardinal_com_error_sums: torch.Tensor | None = None self._cardinal_gait_sums: torch.Tensor | None = None self._cardinal_foot_min_distance_sums: torch.Tensor | None = None self._cardinal_foot_min_lateral_separation_sums: ( torch.Tensor | None ) = None self._cardinal_foot_clearance_violation_fraction_sums: ( torch.Tensor | None ) = None self._cardinal_foot_lateral_order_violation_fraction_sums: ( torch.Tensor | None ) = None self._height_debug_steps = 0 robot = base.scene["robot"] self._shoulder_joint_ids, _ = robot.find_joints( ["LH_yaw", "RH_yaw"], preserve_order=True, ) self._knee_joint_names = [ "LL_knee_actuator_joint", "RL_knee_actuator_joint", ] self._knee_joint_ids, matched_knee_names = robot.find_joints( self._knee_joint_names, preserve_order=True, ) if matched_knee_names != self._knee_joint_names: raise ValueError( "Evaluator could not resolve the ordered knees: " f"{matched_knee_names}" ) self._phase_joint_names = [ "PG_left_leg_pitch", "PG_right_leg_pitch", "LL_knee_actuator_joint", "RL_knee_actuator_joint", "LL_Revolute67", "RL_Revolute67", ] self._phase_joint_ids, matched_phase_joint_names = robot.find_joints( self._phase_joint_names, preserve_order=True, ) if matched_phase_joint_names != self._phase_joint_names: raise ValueError( "Evaluator could not resolve ordered sagittal joints: " f"{matched_phase_joint_names}" ) self._phase_joint_leg_ids = torch.tensor( [0, 1, 0, 1, 0, 1], dtype=torch.long, device=base.device, ) self._pose_height_body_ids, _ = robot.find_bodies( self._pose_height_body_names, preserve_order=True, ) self._distal_foot_body_ids, matched_foot_names = robot.find_bodies( DISTAL_FOOT_BODIES, preserve_order=True, ) if matched_foot_names != DISTAL_FOOT_BODIES: raise ValueError( "Evaluator could not resolve the ordered distal feet: " f"{matched_foot_names}" ) gait_term_cfg = base.reward_manager.get_term_cfg("gait") contact_sensor_cfg = gait_term_cfg.params["sensor_cfg"] self._contact_sensor = base.scene.sensors[contact_sensor_cfg.name] self._contact_body_ids = contact_sensor_cfg.body_ids initial_root_pos_w = robot.data.root_pos_w if isinstance(initial_root_pos_w, wp.array): initial_root_pos_w = wp.to_torch(initial_root_pos_w) print( "[EVAL HEIGHT INIT] " f"root_z={float(initial_root_pos_w[0, 2]):.4f} " f"origin_z={float(base.scene.env_origins[0, 2]):.4f}", flush=True, ) if POSE_SEQUENCE: self._pose_joint_names = [ "PG_left_leg_pitch", "PG_right_leg_pitch", "LL_knee_actuator_joint", "RL_knee_actuator_joint", "LL_Revolute67", "RL_Revolute67", ] self._pose_joint_ids, _ = robot.find_joints( self._pose_joint_names, preserve_order=True, ) self._pose_offsets = torch.tensor( [ ( POSE_PG_PITCH_OFFSET if POSE_PG_PITCH_OFFSET is not None else -0.10 ), ( POSE_PG_PITCH_OFFSET if POSE_PG_PITCH_OFFSET is not None else -0.10 ), ( POSE_KNEE_OFFSET if POSE_KNEE_OFFSET is not None else 0.25 ), ( POSE_KNEE_OFFSET if POSE_KNEE_OFFSET is not None else 0.25 ), ( POSE_ANKLE67_OFFSET if POSE_ANKLE67_OFFSET is not None else -0.15 ), ( POSE_ANKLE67_OFFSET if POSE_ANKLE67_OFFSET is not None else -0.15 ), ], device=base.device, ) if POSE_BIAS_GRID: action_term = base.action_manager.get_term("joint_pos") action_names = list(action_term._joint_names) action_index = { name: index for index, name in enumerate(action_names) } pattern_specs: list[tuple[str, dict[str, float]]] = [ ("none", {}), ( "knees_pos_015", { "LL_knee_actuator_joint": 0.15, "RL_knee_actuator_joint": 0.15, }, ), ( "knees_pos_030", { "LL_knee_actuator_joint": 0.30, "RL_knee_actuator_joint": 0.30, }, ), ( "knees_neg_015", { "LL_knee_actuator_joint": -0.15, "RL_knee_actuator_joint": -0.15, }, ), ( "pelvis_pitch_pos_015", { "PG_left_leg_pitch": 0.15, "PG_right_leg_pitch": 0.15, }, ), ( "pelvis_pitch_neg_015", { "PG_left_leg_pitch": -0.15, "PG_right_leg_pitch": -0.15, }, ), ( "ankle67_pos_015", {"LL_Revolute67": 0.15, "RL_Revolute67": 0.15}, ), ( "ankle67_neg_015", {"LL_Revolute67": -0.15, "RL_Revolute67": -0.15}, ), ( "crouch_combo", { "PG_left_leg_pitch": -0.15, "PG_right_leg_pitch": -0.15, "LL_knee_actuator_joint": 0.30, "RL_knee_actuator_joint": 0.30, "LL_Revolute67": -0.15, "RL_Revolute67": -0.15, }, ), ] missing = sorted( { joint_name for _, spec in pattern_specs for joint_name in spec if joint_name not in action_index } ) if missing: raise ValueError( "pose-bias joints are absent from the action term: " + ", ".join(missing) ) patterns = torch.zeros( len(pattern_specs), len(action_names), device=base.device, ) for pattern_id, (_, spec) in enumerate(pattern_specs): for joint_name, bias in spec.items(): patterns[pattern_id, action_index[joint_name]] = bias self._pose_bias_names = [name for name, _ in pattern_specs] self._pose_bias_group_ids = ( torch.arange(self._num_envs, device=base.device) % len(pattern_specs) ) self._pose_bias_by_env = patterns[self._pose_bias_group_ids] group_count = len(pattern_specs) self._pose_bias_episode_counts = torch.zeros( group_count, dtype=torch.long, device=base.device ) self._pose_bias_fall_counts = torch.zeros_like( self._pose_bias_episode_counts ) self._pose_bias_stand_height_sums = torch.zeros( group_count, device=base.device ) self._pose_bias_stand_step_sums = torch.zeros( group_count, device=base.device ) self._pose_bias_crouch_height_sums = torch.zeros( group_count, device=base.device ) self._pose_bias_crouch_step_sums = torch.zeros( group_count, device=base.device ) if PLANAR_CARDINAL_SWEEP: group_count = len(self._cardinal_names) self._cardinal_group_ids = ( torch.arange(self._num_envs, device=base.device) % group_count ) self._cardinal_episode_counts = torch.zeros( group_count, dtype=torch.long, device=base.device ) self._cardinal_fall_counts = torch.zeros_like( self._cardinal_episode_counts ) self._cardinal_length_sums = torch.zeros( group_count, device=base.device ) self._cardinal_com_velocity_sums = torch.zeros( group_count, 2, device=base.device ) self._cardinal_com_error_sums = torch.zeros( group_count, device=base.device ) self._cardinal_gait_sums = torch.zeros( group_count, device=base.device ) self._cardinal_foot_min_distance_sums = torch.zeros( group_count, device=base.device ) self._cardinal_foot_min_lateral_separation_sums = torch.zeros( group_count, device=base.device ) self._cardinal_foot_clearance_violation_fraction_sums = ( torch.zeros(group_count, device=base.device) ) self._cardinal_foot_lateral_order_violation_fraction_sums = ( torch.zeros(group_count, device=base.device) ) def step(self, action): if self._pose_bias_by_env is not None: depth, _ = pose_sequence_reference(self._base) action = action.clone() action += depth.unsqueeze(1) * self._pose_bias_by_env if SHOULDER_COUNTERWEIGHT_SCALE is not None: action = action.clone() shoulder_actions = action[:, self._shoulder_action_ids] reciprocal_action = torch.mean( shoulder_actions, dim=1, keepdim=True, ) action[:, self._shoulder_action_ids] = ( reciprocal_action + SHOULDER_COUNTERWEIGHT_SCALE * (shoulder_actions - reciprocal_action) ) command_term = self._base.command_manager.get_term("base_velocity") command = command_term.vel_command_b.clone() command_xy = command[:, :2] observations, rewards, terminated, truncated, extras = self.env.step(action) step_rewards = self._base.reward_manager._step_reward robot_data = self._base.scene["robot"].data root_lin_vel_b = robot_data.root_lin_vel_b root_ang_vel_b = robot_data.root_ang_vel_b root_quat_w = robot_data.root_quat_w root_lin_vel_w = robot_data.root_lin_vel_w if isinstance(root_lin_vel_b, wp.array): root_lin_vel_b = wp.to_torch(root_lin_vel_b) if isinstance(root_ang_vel_b, wp.array): root_ang_vel_b = wp.to_torch(root_ang_vel_b) if isinstance(root_quat_w, wp.array): root_quat_w = wp.to_torch(root_quat_w) if isinstance(root_lin_vel_w, wp.array): root_lin_vel_w = wp.to_torch(root_lin_vel_w) root_lin_vel_yaw = quat_apply_inverse(yaw_quat(root_quat_w), root_lin_vel_w) body_velocity_error = torch.linalg.norm(command_xy - root_lin_vel_b[:, :2], dim=-1) yaw_velocity_error = torch.linalg.norm(command_xy - root_lin_vel_yaw[:, :2], dim=-1) yaw_rate_error = torch.abs(command[:, 2] - root_ang_vel_b[:, 2]) com_pos_w, com_lin_vel_w = mass_weighted_com_state(self._base) com_lin_vel_yaw = quat_apply_inverse( yaw_quat(root_quat_w), com_lin_vel_w, ) com_velocity_error = torch.linalg.norm( command_xy - com_lin_vel_yaw[:, :2], dim=-1, ) target_com_height, target_com_vertical_velocity = com_height_reference( self._base, float(os.environ.get("DROPBEAR_COM_STAND_HEIGHT", "1.20")), float(os.environ.get("DROPBEAR_COM_HEIGHT_DELTA", "0.0")), 16.0, ) env_origins = self._base.scene.env_origins.to(com_pos_w.device) com_height = com_pos_w[:, 2] - env_origins[:, 2] com_vertical_velocity_error = torch.abs( com_lin_vel_w[:, 2] - target_com_vertical_velocity ) com_height_abs_error = torch.abs(com_height - target_com_height) def raw_reward_term(name: str) -> torch.Tensor: term_cfg = self._base.reward_manager.get_term_cfg(name) return term_cfg.func(self._base, **term_cfg.params) arm_swing_score = raw_reward_term("natural_arm_swing") arm_extension_l2 = raw_reward_term("excessive_arm_extension") foot_phase_velocity_score = raw_reward_term("foot_phase_velocity") body_pos_w = robot_data.body_pos_w if isinstance(body_pos_w, wp.array): body_pos_w = wp.to_torch(body_pos_w) foot_delta_w = ( body_pos_w[:, self._distal_foot_body_ids[0]] - body_pos_w[:, self._distal_foot_body_ids[1]] ) foot_distance = torch.linalg.vector_norm(foot_delta_w, dim=1) foot_delta_yaw = quat_apply_inverse( yaw_quat(root_quat_w), foot_delta_w, ) foot_lateral_separation = foot_delta_yaw[:, 1] foot_clearance_violation = ( foot_distance < FEET_MIN_DISTANCE ) foot_lateral_order_violation = ( foot_lateral_separation < FEET_MIN_LATERAL_SEPARATION ) severe_foot_overlap = ( (foot_distance < 0.10) | (foot_lateral_separation < 0.0) ) current_contact_time = self._contact_sensor.data.current_contact_time if isinstance(current_contact_time, wp.array): current_contact_time = wp.to_torch(current_contact_time) foot_body_contact = ( current_contact_time[:, self._contact_body_ids] > 0 ) contact_split = foot_body_contact.shape[1] // 2 if contact_split == 0: raise ValueError( "Evaluator requires ordered contact bodies for both feet" ) leg_contact = torch.stack( ( foot_body_contact[:, :contact_split].any(dim=1), foot_body_contact[:, contact_split:].any(dim=1), ), dim=1, ) valid_touchdown = ( leg_contact & ~self._was_leg_contact & (self._leg_air_steps * self._step_dt >= 0.06) ) completed_flight_time = self._leg_air_steps.float() * self._step_dt self._leg_flight_time_sum += ( completed_flight_time * valid_touchdown.float() ) self._leg_flight_count += valid_touchdown current_step = self._lengths.unsqueeze(1).expand_as( self._last_leg_touchdown_step ) prior_touchdown = self._last_leg_touchdown_step >= 0 interval_update = valid_touchdown & prior_touchdown stride_interval = ( current_step - self._last_leg_touchdown_step ).float() * self._step_dt self._leg_stride_interval_sum += ( stride_interval * interval_update.float() ) self._leg_stride_interval_count += interval_update self._last_leg_touchdown_step = torch.where( valid_touchdown, current_step, self._last_leg_touchdown_step, ) self._leg_touchdown_count += valid_touchdown self._leg_contact_steps += leg_contact self._leg_air_steps = torch.where( leg_contact, torch.zeros_like(self._leg_air_steps), self._leg_air_steps + 1, ) self._was_leg_contact = leg_contact.clone() joint_pos = robot_data.joint_pos default_joint_pos = robot_data.default_joint_pos if isinstance(joint_pos, wp.array): joint_pos = wp.to_torch(joint_pos) if isinstance(default_joint_pos, wp.array): default_joint_pos = wp.to_torch(default_joint_pos) knee_position = joint_pos[:, self._knee_joint_ids] global_gait_phase = ( ( self._base.episode_length_buf.to(knee_position.dtype) * self._step_dt ) % EVAL_GAIT_PERIOD / EVAL_GAIT_PERIOD ).unsqueeze(1) leg_phase = ( global_gait_phase + torch.tensor( [0.0, 0.5], dtype=knee_position.dtype, device=knee_position.device, ).unsqueeze(0) ) % 1.0 knee_stance_mask = leg_phase < 0.55 knee_swing_mask = ~knee_stance_mask phase_joint_position = joint_pos[:, self._phase_joint_ids] phase_joint_swing_mask = knee_swing_mask[ :, self._phase_joint_leg_ids ] phase_joint_stance_mask = ~phase_joint_swing_mask swing_progress = torch.clamp( (leg_phase - 0.55) / (1.0 - 0.55), min=0.0, max=1.0, ) late_swing_mask = knee_swing_mask & (swing_progress >= 0.80) gravity_w = torch.zeros( self._num_envs, 3, dtype=root_quat_w.dtype, device=root_quat_w.device, ) gravity_w[:, 2] = -1.0 gravity_b = quat_apply_inverse(root_quat_w, gravity_w) torso_roll = torch.atan2(-gravity_b[:, 1], -gravity_b[:, 2]) root_pos_w = robot_data.root_pos_w if isinstance(root_pos_w, wp.array): root_pos_w = wp.to_torch(root_pos_w) foot_relative_w = body_pos_w[:, self._distal_foot_body_ids, :] - ( root_pos_w.unsqueeze(1) ) foot_relative_yaw = torch.stack( [ quat_apply_inverse( yaw_quat(root_quat_w), foot_relative_w[:, leg_index, :], ) for leg_index in range(2) ], dim=1, ) foot_lateral_position = foot_relative_yaw[:, :, 1] travel_sign = torch.where( command_xy[:, 0] >= 0.0, torch.ones_like(command_xy[:, 0]), -torch.ones_like(command_xy[:, 0]), ) foot_advance = torch.stack( ( travel_sign * ( foot_relative_yaw[:, 0, 0] - foot_relative_yaw[:, 1, 0] ), travel_sign * ( foot_relative_yaw[:, 1, 0] - foot_relative_yaw[:, 0, 0] ), ), dim=1, ) forward_progress = root_pos_w[:, 0] - env_origins[:, 0] self._max_forward_progress = torch.maximum( self._max_forward_progress, forward_progress, ) if EVAL_OBSTACLE_TERRAIN: # The anchored box spans x=[0.85, 1.15] m relative to each # terrain origin. Count a controlled crossing only when the # pelvis has passed it while upright and near the commanded # planar velocity; retain raw progress as a separate diagnostic. progressed_past = forward_progress > 1.35 self._obstacle_progressed_past |= progressed_past upright = gravity_b[:, 2] < -0.85 controlled_velocity = ( torch.linalg.vector_norm( root_lin_vel_yaw[:, :2] - command_xy, dim=1, ) < 0.35 ) self._obstacle_crossed |= ( progressed_past & upright & controlled_velocity ) foot_forward_from_origin = ( body_pos_w[:, self._distal_foot_body_ids, 0] - env_origins[:, None, 0] ) foot_in_obstacle_zone = ( (foot_forward_from_origin >= 0.75) & (foot_forward_from_origin <= 1.25) ) self._obstacle_zone_seen |= torch.any( foot_in_obstacle_zone, dim=1 ) foot_height_from_origin = ( body_pos_w[:, self._distal_foot_body_ids, 2] - env_origins[:, None, 2] ) zone_height = torch.where( foot_in_obstacle_zone, foot_height_from_origin, torch.full_like(foot_height_from_origin, float("-inf")), ) self._obstacle_zone_max_foot_height = torch.maximum( self._obstacle_zone_max_foot_height, zone_height, ) shoulder_offset = ( joint_pos[:, self._shoulder_joint_ids] - default_joint_pos[:, self._shoulder_joint_ids] ) # The USD axes are mirrored (LH -Y, RH +Y). Equal signed joint # offsets are reciprocal physical motion; opposite signed offsets are # both arms moving fore/aft together as a counterweight. reciprocal_mode = 0.5 * (shoulder_offset[:, 0] + shoulder_offset[:, 1]) counterweight_mode = 0.5 * ( shoulder_offset[:, 0] - shoulder_offset[:, 1] ) arm_phase = ( 2.0 * torch.pi * self._base.episode_length_buf.to(reciprocal_mode.dtype) * self._step_dt / EVAL_GAIT_PERIOD ) phase_sin = torch.sin(arm_phase) phase_cos = torch.cos(arm_phase) self._returns += rewards self._tracking += step_rewards[:, self._tracking_index] * self._step_dt self._gait += step_rewards[:, self._gait_index] * self._step_dt self._body_velocity_error += body_velocity_error self._yaw_velocity_error += yaw_velocity_error self._yaw_rate_error += yaw_rate_error self._com_velocity_error += com_velocity_error self._root_velocity_yaw_xy += root_lin_vel_yaw[:, :2] self._com_velocity_yaw_xy += com_lin_vel_yaw[:, :2] self._root_yaw_rate += root_ang_vel_b[:, 2] self._com_vertical_velocity_error += com_vertical_velocity_error self._com_height_abs_error += com_height_abs_error self._arm_swing_score += arm_swing_score self._arm_extension_l2 += arm_extension_l2 self._arm_reciprocal_mode_abs += torch.abs(reciprocal_mode) self._arm_counterweight_mode_abs += torch.abs(counterweight_mode) self._arm_reciprocal_sin += reciprocal_mode * phase_sin self._arm_reciprocal_cos += reciprocal_mode * phase_cos self._arm_counterweight_sin += counterweight_mode * phase_sin self._arm_counterweight_cos += counterweight_mode * phase_cos self._foot_phase_velocity_score += foot_phase_velocity_score self._foot_min_distance = torch.minimum( self._foot_min_distance, foot_distance, ) self._foot_min_lateral_separation = torch.minimum( self._foot_min_lateral_separation, foot_lateral_separation, ) self._foot_clearance_violation_steps += foot_clearance_violation self._foot_lateral_order_violation_steps += ( foot_lateral_order_violation ) self._severe_foot_overlap_steps += severe_foot_overlap self._knee_position += knee_position self._knee_swing_position += ( knee_position * knee_swing_mask.float() ) self._knee_stance_position += ( knee_position * knee_stance_mask.float() ) self._knee_swing_steps += knee_swing_mask self._knee_stance_steps += knee_stance_mask self._knee_min = torch.minimum(self._knee_min, knee_position) self._knee_max = torch.maximum(self._knee_max, knee_position) self._phase_joint_position += phase_joint_position self._phase_joint_swing_position += ( phase_joint_position * phase_joint_swing_mask.float() ) self._phase_joint_stance_position += ( phase_joint_position * phase_joint_stance_mask.float() ) self._phase_joint_swing_steps += phase_joint_swing_mask self._phase_joint_stance_steps += phase_joint_stance_mask self._torso_roll += torso_roll self._torso_roll_abs += torch.abs(torso_roll) self._torso_roll_squared += torch.square(torso_roll) self._torso_roll_sin += torso_roll * phase_sin self._torso_roll_cos += torso_roll * phase_cos self._torso_roll_positive_steps += torso_roll > 0.0 self._foot_lateral_min = torch.minimum( self._foot_lateral_min, foot_lateral_position, ) self._foot_lateral_max = torch.maximum( self._foot_lateral_max, foot_lateral_position, ) self._foot_position_yaw_xy += foot_relative_yaw[:, :, :2] self._late_swing_foot_advance += ( foot_advance * late_swing_mask.float() ) self._late_swing_foot_steps += late_swing_mask if POSE_SEQUENCE and self._pose_offsets is not None: depth, _ = pose_sequence_reference(self._base) target = ( default_joint_pos[:, self._pose_joint_ids] + depth.unsqueeze(1) * self._pose_offsets.unsqueeze(0) ) joint_delta = ( joint_pos[:, self._pose_joint_ids] - default_joint_pos[:, self._pose_joint_ids] ) achieved_depth = torch.sum( joint_delta * self._pose_offsets.unsqueeze(0), dim=1, ) / torch.sum(torch.square(self._pose_offsets)) self._pose_joint_mse += torch.mean( torch.square(joint_pos[:, self._pose_joint_ids] - target), dim=1, ) target_depth = POSE_BASELINE_DEPTH + POSE_DEPTH_AMPLITUDE * depth self._pose_depth_abs_error += torch.abs( achieved_depth - target_depth ) env_origins = self._base.scene.env_origins.to(body_pos_w.device) torso_height = torch.mean( body_pos_w[:, self._pose_height_body_ids, 2], dim=1 ) - env_origins[:, 2] com_pos_w, _ = mass_weighted_com_state(self._base) com_height = com_pos_w[:, 2] - env_origins[:, 2] if self._height_debug_steps < 5: print( "[EVAL HEIGHT STEP] " f"step={self._height_debug_steps + 1} " f"origin_z={float(env_origins[0, 2]):.4f} " f"torso_height={float(torso_height[0]):.4f}", flush=True, ) self._height_debug_steps += 1 target_height = ( POSE_STAND_HEIGHT - POSE_CROUCH_HEIGHT_DELTA * depth ) self._pose_height_abs_error += torch.abs( torso_height - target_height ) stand_mask = depth < 0.05 crouch_mask = depth > 0.95 self._stand_achieved_depth += achieved_depth * stand_mask self._stand_joint_delta += joint_delta * stand_mask.unsqueeze(1) self._stand_root_height += torso_height * stand_mask self._stand_com_height += com_height * stand_mask self._stand_pose_steps += stand_mask self._crouch_achieved_depth += achieved_depth * crouch_mask self._crouch_joint_delta += joint_delta * crouch_mask.unsqueeze(1) self._crouch_root_height += torso_height * crouch_mask self._crouch_com_height += com_height * crouch_mask self._crouch_pose_steps += crouch_mask self._lengths += 1 done = terminated | truncated done_ids = done.nonzero(as_tuple=False).squeeze(-1) if done_ids.numel() > 0: self._completed_per_env[done_ids] += 1 evaluated_ids = done_ids[self._completed_per_env[done_ids] > WARMUP_EPISODES] completed = int(evaluated_ids.numel()) self._episodes += completed if completed > 0: if ( self._cardinal_group_ids is not None and self._cardinal_episode_counts is not None and self._cardinal_fall_counts is not None and self._cardinal_length_sums is not None and self._cardinal_com_velocity_sums is not None and self._cardinal_com_error_sums is not None and self._cardinal_gait_sums is not None and self._cardinal_foot_min_distance_sums is not None and self._cardinal_foot_min_lateral_separation_sums is not None and self._cardinal_foot_clearance_violation_fraction_sums is not None and self._cardinal_foot_lateral_order_violation_fraction_sums is not None ): evaluated_groups = self._cardinal_group_ids[ evaluated_ids ] for group_id in range(len(self._cardinal_names)): group_ids = evaluated_ids[ evaluated_groups == group_id ] if group_ids.numel() == 0: continue group_lengths = self._lengths[group_ids] self._cardinal_episode_counts[group_id] += ( group_ids.numel() ) self._cardinal_fall_counts[group_id] += terminated[ group_ids ].sum() self._cardinal_length_sums[group_id] += ( group_lengths.sum() ) self._cardinal_com_velocity_sums[group_id] += ( self._com_velocity_yaw_xy[group_ids] / group_lengths.unsqueeze(1) ).sum(dim=0) self._cardinal_com_error_sums[group_id] += ( self._com_velocity_error[group_ids] / group_lengths ).sum() self._cardinal_gait_sums[group_id] += self._gait[ group_ids ].sum() self._cardinal_foot_min_distance_sums[group_id] += ( self._foot_min_distance[group_ids].sum() ) self._cardinal_foot_min_lateral_separation_sums[ group_id ] += self._foot_min_lateral_separation[group_ids].sum() self._cardinal_foot_clearance_violation_fraction_sums[ group_id ] += ( self._foot_clearance_violation_steps[group_ids] / group_lengths ).sum() self._cardinal_foot_lateral_order_violation_fraction_sums[ group_id ] += ( self._foot_lateral_order_violation_steps[group_ids] / group_lengths ).sum() if ( self._pose_bias_group_ids is not None and self._pose_bias_episode_counts is not None and self._pose_bias_fall_counts is not None and self._pose_bias_stand_height_sums is not None and self._pose_bias_stand_step_sums is not None and self._pose_bias_crouch_height_sums is not None and self._pose_bias_crouch_step_sums is not None ): evaluated_groups = self._pose_bias_group_ids[evaluated_ids] for group_id in range(len(self._pose_bias_names)): group_ids = evaluated_ids[ evaluated_groups == group_id ] if group_ids.numel() == 0: continue self._pose_bias_episode_counts[group_id] += ( group_ids.numel() ) self._pose_bias_fall_counts[group_id] += terminated[ group_ids ].sum() self._pose_bias_stand_height_sums[group_id] += ( self._stand_root_height[group_ids].sum() ) self._pose_bias_stand_step_sums[group_id] += ( self._stand_pose_steps[group_ids].sum() ) self._pose_bias_crouch_height_sums[group_id] += ( self._crouch_root_height[group_ids].sum() ) self._pose_bias_crouch_step_sums[group_id] += ( self._crouch_pose_steps[group_ids].sum() ) self._fall_episodes += int(terminated[evaluated_ids].sum().item()) self._timeout_episodes += int(truncated[evaluated_ids].sum().item()) self._return_sum += float(self._returns[evaluated_ids].sum().item()) self._length_sum += int(self._lengths[evaluated_ids].sum().item()) self._tracking_sum += float(self._tracking[evaluated_ids].sum().item()) self._gait_sum += float(self._gait[evaluated_ids].sum().item()) self._body_velocity_error_sum += float( (self._body_velocity_error[evaluated_ids] / self._lengths[evaluated_ids]).sum().item() ) self._yaw_velocity_error_sum += float( (self._yaw_velocity_error[evaluated_ids] / self._lengths[evaluated_ids]).sum().item() ) self._yaw_rate_error_sum += float( ( self._yaw_rate_error[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._com_velocity_error_sum += float( ( self._com_velocity_error[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._root_velocity_yaw_xy_sum += ( self._root_velocity_yaw_xy[evaluated_ids] / self._lengths[evaluated_ids].unsqueeze(1) ).sum(dim=0) self._com_velocity_yaw_xy_sum += ( self._com_velocity_yaw_xy[evaluated_ids] / self._lengths[evaluated_ids].unsqueeze(1) ).sum(dim=0) self._root_yaw_rate_sum += float( ( self._root_yaw_rate[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._com_vertical_velocity_error_sum += float( ( self._com_vertical_velocity_error[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._com_height_abs_error_sum += float( ( self._com_height_abs_error[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._arm_swing_score_sum += float( ( self._arm_swing_score[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._arm_extension_l2_sum += float( ( self._arm_extension_l2[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) evaluated_lengths = self._lengths[evaluated_ids] self._arm_reciprocal_mode_abs_sum += float( ( self._arm_reciprocal_mode_abs[evaluated_ids] / evaluated_lengths ).sum().item() ) self._arm_counterweight_mode_abs_sum += float( ( self._arm_counterweight_mode_abs[evaluated_ids] / evaluated_lengths ).sum().item() ) # For q(t) = A sin(wt + phi), 2*mean(q*sin(wt)) is # A*cos(phi) and 2*mean(q*cos(wt)) is A*sin(phi). self._arm_reciprocal_sin_coefficient_sum += float( ( 2.0 * self._arm_reciprocal_sin[evaluated_ids] / evaluated_lengths ).sum().item() ) self._arm_reciprocal_cos_coefficient_sum += float( ( 2.0 * self._arm_reciprocal_cos[evaluated_ids] / evaluated_lengths ).sum().item() ) self._arm_counterweight_sin_coefficient_sum += float( ( 2.0 * self._arm_counterweight_sin[evaluated_ids] / evaluated_lengths ).sum().item() ) self._arm_counterweight_cos_coefficient_sum += float( ( 2.0 * self._arm_counterweight_cos[evaluated_ids] / evaluated_lengths ).sum().item() ) self._foot_phase_velocity_score_sum += float( ( self._foot_phase_velocity_score[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._foot_min_distance_sum += float( self._foot_min_distance[evaluated_ids].sum().item() ) self._foot_min_lateral_separation_sum += float( self._foot_min_lateral_separation[evaluated_ids] .sum() .item() ) self._foot_clearance_violation_fraction_sum += float( ( self._foot_clearance_violation_steps[evaluated_ids] / self._lengths[evaluated_ids] ) .sum() .item() ) self._foot_lateral_order_violation_fraction_sum += float( ( self._foot_lateral_order_violation_steps[evaluated_ids] / self._lengths[evaluated_ids] ) .sum() .item() ) self._severe_foot_overlap_fraction_sum += float( ( self._severe_foot_overlap_steps[evaluated_ids] / self._lengths[evaluated_ids] ) .sum() .item() ) contact_fraction = ( self._leg_contact_steps[evaluated_ids] / self._lengths[evaluated_ids].unsqueeze(1) ) self._leg_contact_fraction_sum += contact_fraction.sum(dim=0) self._leg_touchdown_count_sum += self._leg_touchdown_count[ evaluated_ids ].sum(dim=0) self._leg_stride_interval_total += ( self._leg_stride_interval_sum[evaluated_ids] ).sum(dim=0) self._leg_stride_interval_samples += ( self._leg_stride_interval_count[evaluated_ids] ).sum(dim=0) self._leg_flight_time_total += self._leg_flight_time_sum[ evaluated_ids ].sum(dim=0) self._leg_flight_samples += self._leg_flight_count[ evaluated_ids ].sum(dim=0) episode_flight_mean = ( self._leg_flight_time_sum[evaluated_ids] / torch.clamp( self._leg_flight_count[evaluated_ids], min=1.0, ) ) both_flights_observed = torch.all( self._leg_flight_count[evaluated_ids] > 0, dim=1, ) self._flight_time_abs_difference_sum += float( ( torch.abs( episode_flight_mean[:, 0] - episode_flight_mean[:, 1] ) * both_flights_observed.float() ).sum().item() ) self._contact_duty_abs_difference_sum += float( torch.abs( contact_fraction[:, 0] - contact_fraction[:, 1] ).sum().item() ) self._touchdown_count_abs_difference_sum += float( torch.abs( self._leg_touchdown_count[evaluated_ids, 0] - self._leg_touchdown_count[evaluated_ids, 1] ).sum().item() ) evaluated_lengths_2d = self._lengths[ evaluated_ids ].unsqueeze(1) self._knee_position_sum += ( self._knee_position[evaluated_ids] / evaluated_lengths_2d ).sum(dim=0) self._knee_swing_position_sum += ( self._knee_swing_position[evaluated_ids] / torch.clamp( self._knee_swing_steps[evaluated_ids], min=1.0, ) ).sum(dim=0) self._knee_stance_position_sum += ( self._knee_stance_position[evaluated_ids] / torch.clamp( self._knee_stance_steps[evaluated_ids], min=1.0, ) ).sum(dim=0) self._knee_range_sum += ( self._knee_max[evaluated_ids] - self._knee_min[evaluated_ids] ).sum(dim=0) self._phase_joint_position_sum += ( self._phase_joint_position[evaluated_ids] / evaluated_lengths_2d ).sum(dim=0) self._phase_joint_swing_position_sum += ( self._phase_joint_swing_position[evaluated_ids] / torch.clamp( self._phase_joint_swing_steps[evaluated_ids], min=1.0, ) ).sum(dim=0) self._phase_joint_stance_position_sum += ( self._phase_joint_stance_position[evaluated_ids] / torch.clamp( self._phase_joint_stance_steps[evaluated_ids], min=1.0, ) ).sum(dim=0) self._torso_roll_sum += float( ( self._torso_roll[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._torso_roll_abs_sum += float( ( self._torso_roll_abs[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._torso_roll_squared_sum += float( ( self._torso_roll_squared[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._torso_roll_sin_coefficient_sum += float( ( 2.0 * self._torso_roll_sin[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._torso_roll_cos_coefficient_sum += float( ( 2.0 * self._torso_roll_cos[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._torso_roll_positive_fraction_sum += float( ( self._torso_roll_positive_steps[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._foot_lateral_excursion_sum += ( self._foot_lateral_max[evaluated_ids] - self._foot_lateral_min[evaluated_ids] ).sum(dim=0) self._foot_position_yaw_xy_sum += ( self._foot_position_yaw_xy[evaluated_ids] / self._lengths[evaluated_ids, None, None] ).sum(dim=0) self._late_swing_foot_advance_sum += ( self._late_swing_foot_advance[evaluated_ids] / torch.clamp( self._late_swing_foot_steps[evaluated_ids], min=1.0, ) ).sum(dim=0) self._max_forward_progress_sum += float( self._max_forward_progress[evaluated_ids].sum().item() ) if EVAL_OBSTACLE_TERRAIN: self._obstacle_crossed_episodes += int( self._obstacle_crossed[evaluated_ids].sum().item() ) self._obstacle_progressed_past_episodes += int( self._obstacle_progressed_past[evaluated_ids] .sum() .item() ) self._obstacle_zone_seen_episodes += int( self._obstacle_zone_seen[evaluated_ids].sum().item() ) obstacle_heights = self._obstacle_zone_max_foot_height[ evaluated_ids ] self._obstacle_zone_max_foot_height_sum += torch.where( torch.isfinite(obstacle_heights), obstacle_heights, torch.zeros_like(obstacle_heights), ).sum(dim=0) self._pose_joint_mse_sum += float( (self._pose_joint_mse[evaluated_ids] / self._lengths[evaluated_ids]).sum().item() ) self._pose_depth_abs_error_sum += float( ( self._pose_depth_abs_error[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._pose_height_abs_error_sum += float( ( self._pose_height_abs_error[evaluated_ids] / self._lengths[evaluated_ids] ).sum().item() ) self._stand_achieved_depth_sum += float( self._stand_achieved_depth[evaluated_ids].sum().item() ) self._stand_joint_delta_sum += self._stand_joint_delta[ evaluated_ids ].sum(dim=0) self._stand_root_height_sum += float( self._stand_root_height[evaluated_ids].sum().item() ) self._stand_com_height_sum += float( self._stand_com_height[evaluated_ids].sum().item() ) self._stand_pose_step_count += float( self._stand_pose_steps[evaluated_ids].sum().item() ) self._crouch_achieved_depth_sum += float( self._crouch_achieved_depth[evaluated_ids].sum().item() ) self._crouch_joint_delta_sum += self._crouch_joint_delta[ evaluated_ids ].sum(dim=0) self._crouch_root_height_sum += float( self._crouch_root_height[evaluated_ids].sum().item() ) self._crouch_com_height_sum += float( self._crouch_com_height[evaluated_ids].sum().item() ) self._crouch_pose_step_count += float( self._crouch_pose_steps[evaluated_ids].sum().item() ) self._returns[done_ids] = 0.0 self._tracking[done_ids] = 0.0 self._gait[done_ids] = 0.0 self._body_velocity_error[done_ids] = 0.0 self._yaw_velocity_error[done_ids] = 0.0 self._yaw_rate_error[done_ids] = 0.0 self._com_velocity_error[done_ids] = 0.0 self._root_velocity_yaw_xy[done_ids] = 0.0 self._com_velocity_yaw_xy[done_ids] = 0.0 self._root_yaw_rate[done_ids] = 0.0 self._com_vertical_velocity_error[done_ids] = 0.0 self._com_height_abs_error[done_ids] = 0.0 self._arm_swing_score[done_ids] = 0.0 self._arm_extension_l2[done_ids] = 0.0 self._arm_reciprocal_mode_abs[done_ids] = 0.0 self._arm_counterweight_mode_abs[done_ids] = 0.0 self._arm_reciprocal_sin[done_ids] = 0.0 self._arm_reciprocal_cos[done_ids] = 0.0 self._arm_counterweight_sin[done_ids] = 0.0 self._arm_counterweight_cos[done_ids] = 0.0 self._foot_phase_velocity_score[done_ids] = 0.0 self._foot_min_distance[done_ids] = float("inf") self._foot_min_lateral_separation[done_ids] = float("inf") self._foot_clearance_violation_steps[done_ids] = 0.0 self._foot_lateral_order_violation_steps[done_ids] = 0.0 self._severe_foot_overlap_steps[done_ids] = 0.0 self._leg_contact_steps[done_ids] = 0.0 self._leg_touchdown_count[done_ids] = 0.0 self._leg_stride_interval_sum[done_ids] = 0.0 self._leg_stride_interval_count[done_ids] = 0.0 self._leg_flight_time_sum[done_ids] = 0.0 self._leg_flight_count[done_ids] = 0.0 self._last_leg_touchdown_step[done_ids] = -1 self._leg_air_steps[done_ids] = 0 self._was_leg_contact[done_ids] = False self._knee_position[done_ids] = 0.0 self._knee_swing_position[done_ids] = 0.0 self._knee_stance_position[done_ids] = 0.0 self._knee_swing_steps[done_ids] = 0.0 self._knee_stance_steps[done_ids] = 0.0 self._knee_min[done_ids] = float("inf") self._knee_max[done_ids] = float("-inf") self._phase_joint_position[done_ids] = 0.0 self._phase_joint_swing_position[done_ids] = 0.0 self._phase_joint_stance_position[done_ids] = 0.0 self._phase_joint_swing_steps[done_ids] = 0.0 self._phase_joint_stance_steps[done_ids] = 0.0 self._torso_roll[done_ids] = 0.0 self._torso_roll_abs[done_ids] = 0.0 self._torso_roll_squared[done_ids] = 0.0 self._torso_roll_sin[done_ids] = 0.0 self._torso_roll_cos[done_ids] = 0.0 self._torso_roll_positive_steps[done_ids] = 0.0 self._foot_lateral_min[done_ids] = float("inf") self._foot_lateral_max[done_ids] = float("-inf") self._foot_position_yaw_xy[done_ids] = 0.0 self._late_swing_foot_advance[done_ids] = 0.0 self._late_swing_foot_steps[done_ids] = 0.0 self._obstacle_crossed[done_ids] = False self._obstacle_progressed_past[done_ids] = False self._obstacle_zone_seen[done_ids] = False self._obstacle_zone_max_foot_height[done_ids] = float("-inf") self._max_forward_progress[done_ids] = float("-inf") self._pose_joint_mse[done_ids] = 0.0 self._pose_depth_abs_error[done_ids] = 0.0 self._pose_height_abs_error[done_ids] = 0.0 self._stand_achieved_depth[done_ids] = 0.0 self._stand_joint_delta[done_ids] = 0.0 self._stand_root_height[done_ids] = 0.0 self._stand_com_height[done_ids] = 0.0 self._stand_pose_steps[done_ids] = 0.0 self._crouch_achieved_depth[done_ids] = 0.0 self._crouch_joint_delta[done_ids] = 0.0 self._crouch_root_height[done_ids] = 0.0 self._crouch_com_height[done_ids] = 0.0 self._crouch_pose_steps[done_ids] = 0.0 self._lengths[done_ids] = 0 target_reached = self._episodes >= TARGET_EPISODES if ( POSE_BIAS_GRID and self._pose_bias_episode_counts is not None and self._pose_bias_names ): episodes_per_group = max( 1, TARGET_EPISODES // len(self._pose_bias_names) ) target_reached = bool( torch.all( self._pose_bias_episode_counts >= episodes_per_group ).item() ) if ( PLANAR_CARDINAL_SWEEP and self._cardinal_episode_counts is not None ): episodes_per_group = max( 1, TARGET_EPISODES // len(self._cardinal_names) ) target_reached = bool( torch.all( self._cardinal_episode_counts >= episodes_per_group ).item() ) if target_reached: arm_reciprocal_sin_coefficient = ( self._arm_reciprocal_sin_coefficient_sum / self._episodes ) arm_reciprocal_cos_coefficient = ( self._arm_reciprocal_cos_coefficient_sum / self._episodes ) arm_counterweight_sin_coefficient = ( self._arm_counterweight_sin_coefficient_sum / self._episodes ) arm_counterweight_cos_coefficient = ( self._arm_counterweight_cos_coefficient_sum / self._episodes ) result = { "episodes": self._episodes, "warmup_episodes_per_env": WARMUP_EPISODES, "forward_speed_command_mps": FORWARD_SPEED, "lateral_speed_command_mps": LATERAL_SPEED, "yaw_rate_command_rad_s": YAW_RATE, "backward_speed_command_mps": BACKWARD_SPEED, "planar_cardinal_sweep": ( [ { "name": name, "command_mps": command, "episodes": int( self._cardinal_episode_counts[index].item() ), "fall_rate": float( self._cardinal_fall_counts[index].item() / self._cardinal_episode_counts[index].item() ), "mean_episode_seconds": float( self._cardinal_length_sums[index].item() * self._step_dt / self._cardinal_episode_counts[index].item() ), "mean_com_velocity_yaw_frame_mps": ( self._cardinal_com_velocity_sums[index] / self._cardinal_episode_counts[index] ).tolist(), "mean_com_velocity_error_mps": float( self._cardinal_com_error_sums[index].item() / self._cardinal_episode_counts[index].item() ), "mean_gait_reward": float( self._cardinal_gait_sums[index].item() / ( self._cardinal_episode_counts[index].item() * self._max_episode_seconds ) ), "mean_episode_min_foot_distance_m": float( self._cardinal_foot_min_distance_sums[ index ].item() / self._cardinal_episode_counts[index].item() ), "mean_episode_min_foot_lateral_separation_m": float( self._cardinal_foot_min_lateral_separation_sums[ index ].item() / self._cardinal_episode_counts[index].item() ), "mean_foot_clearance_violation_fraction": float( self._cardinal_foot_clearance_violation_fraction_sums[ index ].item() / self._cardinal_episode_counts[index].item() ), "mean_foot_lateral_order_violation_fraction": float( self._cardinal_foot_lateral_order_violation_fraction_sums[ index ].item() / self._cardinal_episode_counts[index].item() ), } for index, (name, command) in enumerate( ( ("forward", [FORWARD_SPEED, 0.0]), ("backward", [BACKWARD_SPEED, 0.0]), ("left", [0.0, LATERAL_SPEED]), ("right", [0.0, -LATERAL_SPEED]), ) ) ] if PLANAR_CARDINAL_SWEEP and self._cardinal_episode_counts is not None and self._cardinal_fall_counts is not None and self._cardinal_length_sums is not None and self._cardinal_com_velocity_sums is not None and self._cardinal_com_error_sums is not None and self._cardinal_gait_sums is not None and self._cardinal_foot_min_distance_sums is not None and self._cardinal_foot_min_lateral_separation_sums is not None and self._cardinal_foot_clearance_violation_fraction_sums is not None and self._cardinal_foot_lateral_order_violation_fraction_sums is not None else None ), "foot_clearance_body_names": DISTAL_FOOT_BODIES, "foot_min_distance_threshold_m": FEET_MIN_DISTANCE, "foot_min_lateral_separation_threshold_m": ( FEET_MIN_LATERAL_SEPARATION ), "gait_period_s": EVAL_GAIT_PERIOD, "plane_only": PLANE_ONLY, "obstacle_terrain": EVAL_OBSTACLE_TERRAIN, "obstacle_crossing_rate": ( self._obstacle_crossed_episodes / self._episodes if EVAL_OBSTACLE_TERRAIN else None ), "obstacle_progressed_past_rate": ( self._obstacle_progressed_past_episodes / self._episodes if EVAL_OBSTACLE_TERRAIN else None ), "obstacle_zone_reach_rate": ( self._obstacle_zone_seen_episodes / self._episodes if EVAL_OBSTACLE_TERRAIN else None ), "mean_obstacle_zone_max_foot_height_m": ( ( self._obstacle_zone_max_foot_height_sum / self._episodes ).tolist() if EVAL_OBSTACLE_TERRAIN else None ), "mean_episode_max_forward_progress_m": ( self._max_forward_progress_sum / self._episodes ), "pushes_disabled": DISABLE_PUSHES, "push_interval_s": ( None if DISABLE_PUSHES else float( os.environ.get("DROPBEAR_PUSH_INTERVAL_S", "5.0") ) ), "push_forward_velocity_mps": ( None if DISABLE_PUSHES else float( os.environ.get( "DROPBEAR_PUSH_FORWARD_VELOCITY", "0.5", ) ) ), "push_lateral_velocity_mps": ( None if DISABLE_PUSHES else float( os.environ.get( "DROPBEAR_PUSH_LATERAL_VELOCITY", "0.5", ) ) ), "reciprocal_shoulder_action_projection": ( SHOULDER_COUNTERWEIGHT_SCALE is not None ), "shoulder_counterweight_action_scale": ( SHOULDER_COUNTERWEIGHT_SCALE ), "actor_base_lin_vel_feedback": ACTOR_BASE_LIN_VEL, "policy_mirror_projection": POLICY_MIRROR_PROJECTION, "reset_joint_position_range": RESET_JOINT_POSITION_RANGE, "reset_joint_velocity_range": RESET_JOINT_VELOCITY_RANGE, "reset_policy_joints_only": RESET_POLICY_JOINTS_ONLY, "pose_sequence": POSE_SEQUENCE, "pose_action_residual": POSE_ACTION_RESIDUAL, "pose_baseline_depth": ( POSE_BASELINE_DEPTH if POSE_SEQUENCE else None ), "pose_depth_amplitude": ( POSE_DEPTH_AMPLITUDE if POSE_SEQUENCE else None ), "pose_stand_height_m": ( POSE_STAND_HEIGHT if POSE_SEQUENCE else None ), "pose_height_body_names": ( self._pose_height_body_names if POSE_SEQUENCE else None ), "pose_crouch_height_delta_m": ( POSE_CROUCH_HEIGHT_DELTA if POSE_SEQUENCE else None ), "pose_residual_scale": ( POSE_RESIDUAL_SCALE if POSE_ACTION_RESIDUAL and POSE_RESIDUAL_SCALE is not None else (1.0 if POSE_ACTION_RESIDUAL else None) ), "pose_adapter_offsets": ( self._pose_offsets.tolist() if POSE_SEQUENCE and self._pose_offsets is not None else None ), "pose_bias_grid": ( [ { "name": name, "episodes": int( self._pose_bias_episode_counts[index].item() ), "fall_rate": ( float( self._pose_bias_fall_counts[index].item() / self._pose_bias_episode_counts[index].item() ) if self._pose_bias_episode_counts[index] > 0 else None ), "mean_torso_height_stand_m": ( float( self._pose_bias_stand_height_sums[index].item() / self._pose_bias_stand_step_sums[index].item() ) if self._pose_bias_stand_step_sums[index] > 0 else None ), "mean_torso_height_crouch_hold_m": ( float( self._pose_bias_crouch_height_sums[index].item() / self._pose_bias_crouch_step_sums[index].item() ) if self._pose_bias_crouch_step_sums[index] > 0 else None ), } for index, name in enumerate(self._pose_bias_names) ] if POSE_BIAS_GRID else None ), "mean_return": self._return_sum / self._episodes, "mean_episode_length": self._length_sum / self._episodes, "mean_episode_seconds": self._length_sum * self._step_dt / self._episodes, "mean_tracking_reward": self._tracking_sum / (self._episodes * self._max_episode_seconds), "mean_gait_reward": self._gait_sum / (self._episodes * self._max_episode_seconds), "mean_body_velocity_error_mps": self._body_velocity_error_sum / self._episodes, "mean_yaw_velocity_error_mps": self._yaw_velocity_error_sum / self._episodes, "mean_yaw_rate_error_rad_s": ( self._yaw_rate_error_sum / self._episodes ), "mean_com_velocity_error_mps": ( self._com_velocity_error_sum / self._episodes ), "mean_root_velocity_yaw_frame_mps": ( self._root_velocity_yaw_xy_sum / self._episodes ).tolist(), "mean_com_velocity_yaw_frame_mps": ( self._com_velocity_yaw_xy_sum / self._episodes ).tolist(), "mean_root_yaw_rate_rad_s": ( self._root_yaw_rate_sum / self._episodes ), "mean_com_vertical_velocity_abs_error_mps": ( self._com_vertical_velocity_error_sum / self._episodes ), "mean_com_height_abs_error_m": ( self._com_height_abs_error_sum / self._episodes ), "mean_arm_swing_score": ( self._arm_swing_score_sum / self._episodes ), "mean_arm_extension_l2": ( self._arm_extension_l2_sum / self._episodes ), "mean_arm_reciprocal_mode_abs_rad": ( self._arm_reciprocal_mode_abs_sum / self._episodes ), "mean_arm_counterweight_mode_abs_rad": ( self._arm_counterweight_mode_abs_sum / self._episodes ), "mean_arm_reciprocal_sin_coefficient_rad": ( arm_reciprocal_sin_coefficient ), "mean_arm_reciprocal_cos_coefficient_rad": ( arm_reciprocal_cos_coefficient ), "mean_arm_reciprocal_fundamental_amplitude_rad": math.hypot( arm_reciprocal_sin_coefficient, arm_reciprocal_cos_coefficient, ), "mean_arm_reciprocal_phase_offset_rad": math.atan2( arm_reciprocal_cos_coefficient, arm_reciprocal_sin_coefficient, ), "mean_arm_counterweight_sin_coefficient_rad": ( arm_counterweight_sin_coefficient ), "mean_arm_counterweight_cos_coefficient_rad": ( arm_counterweight_cos_coefficient ), "mean_arm_counterweight_fundamental_amplitude_rad": math.hypot( arm_counterweight_sin_coefficient, arm_counterweight_cos_coefficient, ), "mean_arm_counterweight_phase_offset_rad": math.atan2( arm_counterweight_cos_coefficient, arm_counterweight_sin_coefficient, ), "mean_foot_phase_velocity_score": ( self._foot_phase_velocity_score_sum / self._episodes ), "mean_episode_min_foot_distance_m": ( self._foot_min_distance_sum / self._episodes ), "mean_episode_min_foot_lateral_separation_m": ( self._foot_min_lateral_separation_sum / self._episodes ), "mean_foot_clearance_violation_fraction": ( self._foot_clearance_violation_fraction_sum / self._episodes ), "mean_foot_lateral_order_violation_fraction": ( self._foot_lateral_order_violation_fraction_sum / self._episodes ), "mean_severe_foot_overlap_fraction": ( self._severe_foot_overlap_fraction_sum / self._episodes ), "mean_leg_contact_fraction": ( self._leg_contact_fraction_sum / self._episodes ).tolist(), "mean_contact_duty_abs_difference": ( self._contact_duty_abs_difference_sum / self._episodes ), "mean_leg_touchdowns_per_episode": ( self._leg_touchdown_count_sum / self._episodes ).tolist(), "mean_touchdown_count_abs_difference": ( self._touchdown_count_abs_difference_sum / self._episodes ), "mean_leg_stride_interval_s": ( self._leg_stride_interval_total / torch.clamp( self._leg_stride_interval_samples, min=1.0, ) ).tolist(), "mean_leg_flight_time_s": ( self._leg_flight_time_total / torch.clamp(self._leg_flight_samples, min=1.0) ).tolist(), "mean_flight_time_abs_difference_s": ( self._flight_time_abs_difference_sum / self._episodes ), "knee_joint_names": self._knee_joint_names, "mean_knee_position_rad": ( self._knee_position_sum / self._episodes ).tolist(), "mean_swing_knee_position_rad": ( self._knee_swing_position_sum / self._episodes ).tolist(), "mean_stance_knee_position_rad": ( self._knee_stance_position_sum / self._episodes ).tolist(), "mean_knee_swing_stance_delta_rad": ( ( self._knee_swing_position_sum - self._knee_stance_position_sum ) / self._episodes ).tolist(), "mean_episode_knee_range_rad": ( self._knee_range_sum / self._episodes ).tolist(), "phase_joint_names": self._phase_joint_names, "mean_phase_joint_position_rad": ( self._phase_joint_position_sum / self._episodes ).tolist(), "mean_phase_joint_swing_position_rad": ( self._phase_joint_swing_position_sum / self._episodes ).tolist(), "mean_phase_joint_stance_position_rad": ( self._phase_joint_stance_position_sum / self._episodes ).tolist(), "mean_phase_joint_swing_stance_delta_rad": ( ( self._phase_joint_swing_position_sum - self._phase_joint_stance_position_sum ) / self._episodes ).tolist(), "mean_torso_roll_rad": ( self._torso_roll_sum / self._episodes ), "mean_abs_torso_roll_rad": ( self._torso_roll_abs_sum / self._episodes ), "rms_torso_roll_rad": ( self._torso_roll_squared_sum / self._episodes ) ** 0.5, "mean_episode_foot_lateral_excursion_m": ( self._foot_lateral_excursion_sum / self._episodes ).tolist(), "mean_foot_position_yaw_frame_xy_m": ( self._foot_position_yaw_xy_sum / self._episodes ).tolist(), "mean_late_swing_foot_advance_m": ( self._late_swing_foot_advance_sum / self._episodes ).tolist(), "mean_pose_joint_rmse_rad": ( (self._pose_joint_mse_sum / self._episodes) ** 0.5 if POSE_SEQUENCE else None ), "mean_torso_roll_sin_coefficient_rad": ( self._torso_roll_sin_coefficient_sum / self._episodes ), "mean_torso_roll_cos_coefficient_rad": ( self._torso_roll_cos_coefficient_sum / self._episodes ), "mean_torso_roll_fundamental_amplitude_rad": math.hypot( self._torso_roll_sin_coefficient_sum / self._episodes, self._torso_roll_cos_coefficient_sum / self._episodes, ), "mean_torso_roll_positive_fraction": ( self._torso_roll_positive_fraction_sum / self._episodes ), "mean_pose_depth_abs_error": ( self._pose_depth_abs_error_sum / self._episodes if POSE_SEQUENCE else None ), "mean_pose_height_abs_error_m": ( self._pose_height_abs_error_sum / self._episodes if POSE_SEQUENCE else None ), "mean_achieved_depth_stand": ( self._stand_achieved_depth_sum / self._stand_pose_step_count if POSE_SEQUENCE and self._stand_pose_step_count > 0 else None ), "pose_joint_names": ( self._pose_joint_names if POSE_SEQUENCE else None ), "mean_stand_joint_delta_rad": ( ( self._stand_joint_delta_sum / self._stand_pose_step_count ).tolist() if POSE_SEQUENCE and self._stand_pose_step_count > 0 else None ), "mean_achieved_depth_crouch_hold": ( self._crouch_achieved_depth_sum / self._crouch_pose_step_count if POSE_SEQUENCE and self._crouch_pose_step_count > 0 else None ), "mean_crouch_hold_joint_delta_rad": ( ( self._crouch_joint_delta_sum / self._crouch_pose_step_count ).tolist() if POSE_SEQUENCE and self._crouch_pose_step_count > 0 else None ), "mean_root_height_stand_m": ( self._stand_root_height_sum / self._stand_pose_step_count if POSE_SEQUENCE and self._stand_pose_step_count > 0 else None ), "mean_torso_height_stand_m": ( self._stand_root_height_sum / self._stand_pose_step_count if POSE_SEQUENCE and self._stand_pose_step_count > 0 else None ), "mean_com_height_stand_m": ( self._stand_com_height_sum / self._stand_pose_step_count if POSE_SEQUENCE and self._stand_pose_step_count > 0 else None ), "mean_root_height_crouch_hold_m": ( self._crouch_root_height_sum / self._crouch_pose_step_count if POSE_SEQUENCE and self._crouch_pose_step_count > 0 else None ), "mean_torso_height_crouch_hold_m": ( self._crouch_root_height_sum / self._crouch_pose_step_count if POSE_SEQUENCE and self._crouch_pose_step_count > 0 else None ), "mean_com_height_crouch_hold_m": ( self._crouch_com_height_sum / self._crouch_pose_step_count if POSE_SEQUENCE and self._crouch_pose_step_count > 0 else None ), "fall_rate": self._fall_episodes / self._episodes, "timeout_rate": self._timeout_episodes / self._episodes, } if OUTPUT_PATH is not None: OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True) OUTPUT_PATH.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") print(f"[DETERMINISTIC EVAL] {json.dumps(result, sort_keys=True)}", flush=True) raise KeyboardInterrupt return observations, rewards, terminated, truncated, extras def make_evaluation_env(env_id: str, *args: Any, **kwargs: Any) -> gym.Env: env_cfg = kwargs.get("cfg") if env_cfg is not None: env_cfg.commands.base_velocity.ranges = env_cfg.commands.base_velocity.limit_ranges if PLANAR_CARDINAL_SWEEP: ranges = env_cfg.commands.base_velocity.ranges ranges.lin_vel_x = (BACKWARD_SPEED, FORWARD_SPEED) ranges.lin_vel_y = (-LATERAL_SPEED, LATERAL_SPEED) ranges.ang_vel_z = (0.0, 0.0) env_cfg.commands.base_velocity.cardinal_commands = True env_cfg.commands.base_velocity.cardinal_by_env = True env_cfg.commands.base_velocity.rel_standing_envs = 0.0 env_cfg.commands.base_velocity.rel_heading_envs = 0.0 env_cfg.commands.base_velocity.heading_command = False elif ( FORWARD_SPEED is not None or LATERAL_SPEED is not None or YAW_RATE is not None ): ranges = env_cfg.commands.base_velocity.ranges forward_speed = FORWARD_SPEED or 0.0 lateral_speed = LATERAL_SPEED or 0.0 ranges.lin_vel_x = (forward_speed, forward_speed) ranges.lin_vel_y = (lateral_speed, lateral_speed) yaw_rate = YAW_RATE or 0.0 ranges.ang_vel_z = (yaw_rate, yaw_rate) # Exact-command trials must never be replaced by the command # term's stochastic standing/heading modes. Otherwise the stock # 2 % standing draw can silently evaluate [0, 0, 0] instead of # the requested locomotion command. env_cfg.commands.base_velocity.rel_standing_envs = 0.0 env_cfg.commands.base_velocity.rel_heading_envs = 0.0 env_cfg.commands.base_velocity.heading_command = False env_cfg.curriculum.lin_vel_cmd_levels = None env_cfg.curriculum.terrain_levels = None if EVAL_OBSTACLE_TERRAIN: if PLANE_ONLY: raise ValueError( "--eval-obstacle-terrain conflicts with --eval-plane" ) generator = env_cfg.scene.terrain.terrain_generator generator.sub_terrains["flat"].proportion = 0.0 generator.sub_terrains["forward_box"].proportion = 1.0 generator.num_rows = 2 generator.num_cols = max( 1, math.ceil(env_cfg.scene.num_envs / generator.num_rows), ) env_cfg.scene.terrain.max_init_terrain_level = 1 if DISABLE_PUSHES: env_cfg.events.push_robot = None if PLANE_ONLY: env_cfg.scene.terrain.terrain_type = "plane" env_cfg.scene.terrain.terrain_generator = None env_cfg.scene.terrain.max_init_terrain_level = 0 elif env_cfg.scene.terrain.terrain_generator is not None: env_cfg.scene.terrain.terrain_generator.curriculum = False return DeterministicEvalWrapper(ORIGINAL_GYM_MAKE(env_id, *args, **kwargs)) if not PLAY_SCRIPT.is_file(): raise FileNotFoundError(f"Isaac Lab player not found: {PLAY_SCRIPT}") gym.make = make_evaluation_env sys.path.insert(0, str(PLAY_SCRIPT.parent)) runpy.run_path(str(PLAY_SCRIPT), run_name="__main__")