scripts/evaluate_dropbear_deterministic.py

#!/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__")