python/openduck/sim.py

from __future__ import annotations

import io
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

import numpy as np

from .spec import (
    ACTION_DIM,
    ALL_BEHAVIORS,
    BEHAVIORS,
    CONTROL_DECIMATION,
    CONTROL_HZ,
    DEFAULT_POSE,
    RECOVERY_STABLE_S,
    STATE_DIM,
    TERMINAL_STABILIZATION_S,
    BehaviorSpec,
)


BALL_OFFSET_X = 0.09
BALL_OFFSET_ABS_Y = 0.042
BALL_RADIUS = 0.035
# The generated MJCF camera quaternion points the optical axis into the head
# shell. MuJoCo cameras look down local -Z, so use the physical convention:
# robot +X is forward, world +Z is image-up, with 15 degrees of downward tilt.
HEAD_CAMERA_QUAT = np.array(
    [0.7010574, -0.0922960, 0.0922960, -0.7010574], dtype=np.float64
)
RECOVERY_INSTRUCTION = (
    "Recover from the fall and stand upright before doing anything else."
)
BALANCE_INSTRUCTION = "Stand upright, hold balance, and remain still."
AUTONOMOUS_TASK_INSTRUCTIONS = {
    "sitstand": "Perform one sit-down and rise cycle.",
    "ground_pick": "Touch the ground once with your beak.",
    "roller_crouch": "Crouch once while gliding on the rollers.",
    "roulade": "Perform one forward roll.",
}

ControlProgram = Callable[[float], tuple[str, np.ndarray, str, str]]


@dataclass
class EpisodeResult:
    images: np.ndarray
    states: np.ndarray
    observation_control_indices: np.ndarray
    actions: np.ndarray
    instruction: str
    success: bool
    metrics: dict[str, float | bool | str]
    command: np.ndarray
    subtask: str
    observation_instructions: tuple[str, ...]
    observation_phases: tuple[str, ...]
    video_frames: np.ndarray | None = None


class SessionCache:
    def __init__(self, policies_dir: Path):
        self.policies_dir = policies_dir
        self._sessions: dict[str, object] = {}

    def get(self, filename: str):
        import onnxruntime as ort

        if filename not in self._sessions:
            path = self.policies_dir / filename
            if not path.exists():
                raise FileNotFoundError(path)
            options = ort.SessionOptions()
            options.intra_op_num_threads = 1
            options.inter_op_num_threads = 1
            session = ort.InferenceSession(
                str(path), sess_options=options, providers=["CPUExecutionProvider"]
            )
            in_shape = list(session.get_inputs()[0].shape)
            out_shape = list(session.get_outputs()[0].shape)
            if in_shape[-1] != STATE_DIM or out_shape[-1] != ACTION_DIM:
                raise ValueError(
                    f"{filename}: expected {STATE_DIM}->{ACTION_DIM}, got {in_shape}->{out_shape}"
                )
            self._sessions[filename] = session
        return self._sessions[filename]


def _quat_rotate_inverse(quat: np.ndarray, vec: np.ndarray) -> np.ndarray:
    w = quat[0]
    xyz = quat[1:4]
    t = np.cross(xyz, vec) * 2
    return vec - w * t + np.cross(xyz, t)


def _instruction(spec: BehaviorSpec, command: np.ndarray, ordinal: int, subtask: str) -> str:
    variant = ordinal % 6
    if spec.kind == "recovery":
        return (
            "Recover from the fall, stand upright, and hold balance.",
            "Get back onto both feet and stabilize in a standing pose.",
            "Stand up from the ground, then remain upright and still.",
            "First recover from the fall, then balance on both feet.",
        )[ordinal % 4]
    if spec.name == "stand":
        z, roll, pitch = command[9], command[10], command[11]
        templates = (
            "Stand still with body height offset {z:+.3f} m, roll {r:+.2f} rad and pitch {p:+.2f} rad.",
            "Hold a stable standing pose: height {z:+.3f} m, roll {r:+.2f}, pitch {p:+.2f}.",
            "Keep balance while standing at body offsets z={z:+.3f}, roll={r:+.2f}, pitch={p:+.2f}.",
            "Remain upright and stationary with z {z:+.3f} m, roll {r:+.2f}, pitch {p:+.2f}.",
            "Balance in place at the requested body pose: z={z:+.3f}, r={r:+.2f}, p={p:+.2f}.",
            "Stay standing and steady; body pose z {z:+.3f}, roll {r:+.2f}, pitch {p:+.2f}.",
        )
        return templates[variant].format(z=z, r=roll, p=pitch)
    if spec.name == "walking":
        vx, vy, wz = command[:3]
        templates = (
            "Walk with forward velocity {x:+.2f} m/s, lateral velocity {y:+.2f} m/s and yaw rate {w:+.2f} rad/s.",
            "Move at vx={x:+.2f}, vy={y:+.2f} m/s while turning at {w:+.2f} rad/s.",
            "Follow the walking command: forward {x:+.2f}, sideways {y:+.2f}, rotation {w:+.2f}.",
            "Walk steadily with velocity command ({x:+.2f}, {y:+.2f}, {w:+.2f}).",
            "Travel using vx {x:+.2f} m/s, vy {y:+.2f} m/s and wz {w:+.2f} rad/s.",
            "Execute the commanded gait: x {x:+.2f}, y {y:+.2f}, yaw {w:+.2f}.",
        )
        return templates[variant].format(x=vx, y=vy, w=wz)
    if spec.name == "sitstand":
        if subtask == "stand_up":
            return (
                "Stand up gently from sitting and finish balanced.",
                "Rise from the seated pose and stabilize upright.",
                "Move from sitting to a balanced standing pose.",
            )[ordinal % 3]
        return (
            "Sit down gently, then stand back up and balance.",
            "Lower into a seated pose, rise again, and finish upright.",
            "Perform one sit-to-stand cycle and stabilize on both feet.",
        )[ordinal % 3]
    if spec.name == "ground_pick":
        return (
            "Lower your beak to touch the ground, then stand back up.",
            "Perform one ground-touch cycle and return upright.",
            "Crouch to the floor with your head and rise again.",
        )[ordinal % 3]
    if spec.name in ("kick_left", "kick_right"):
        side = spec.name.split("_")[1]
        if ordinal % 2 == 0:
            return "Kick the visible ball."
        return (
            f"Kick the ball with the {side} foot.",
            f"Use your {side} leg to kick the ball forward.",
            f"Perform a {side}-foot ball kick.",
        )[ordinal % 3]
    if spec.name == "roller":
        vx = command[0]
        return (
            f"Skate straight with push command {vx:+.2f}.",
            f"Roll forward on the skates at command {vx:+.2f}.",
            f"Use the rollers to move straight, drive {vx:+.2f}.",
        )[ordinal % 3]
    if spec.name == "roller_crouch":
        return (
            "Crouch while gliding on the rollers, then stand again.",
            "Perform one roller crouch cycle and recover upright.",
            "Lower into a skating crouch and rise while rolling.",
        )[ordinal % 3]
    return (
        "Perform a forward roll and recover to standing.",
        "Roll forward over your head, then stand up.",
        "Do one forward roulade and finish upright.",
    )[ordinal % 3]


class MicroDuckEpisode:
    def __init__(
        self,
        spec: BehaviorSpec,
        microduck_rl: Path,
        sessions: SessionCache,
        seed: int,
        ordinal: int,
        image_size: int = 256,
        randomize: bool = True,
        start_fallen: bool = False,
        control_program: ControlProgram | None = None,
        instruction_override: str | None = None,
        ball_forward_offset: float | None = None,
    ):
        import mujoco

        self.spec = spec
        self.sessions = sessions
        self.rng = np.random.default_rng(seed)
        self.ordinal = ordinal
        self.image_size = image_size
        self.control_program = control_program
        self.ball_forward_offset = ball_forward_offset
        self.started_fallen = bool(start_fallen or spec.kind == "recovery")
        self.recovery_active = self.started_fallen
        self.recovery_completed = False
        self.recovery_enter_step = 0 if self.started_fallen else None
        self.recovery_exit_step: int | None = None
        self.recovery_upright_ticks = 0
        self.task_start_step: int | None = None if self.started_fallen else 0
        self.task_control_steps = 0
        self.fall_direction = ""
        xml = (
            microduck_rl
            / "src/mjlab_microduck/robot/microduck"
            / spec.scene_file
        )
        if not xml.exists():
            raise FileNotFoundError(xml)
        self.model = mujoco.MjModel.from_xml_path(str(xml))
        self.model.opt.timestep = 1.0 / 200.0
        self.data = mujoco.MjData(self.model)
        self.renderer = mujoco.Renderer(self.model, height=image_size, width=image_size)
        self._setup_indices()
        self.subtask = ""
        self.command = np.zeros(13, np.float32)
        self._reset(randomize)
        self._sample_command()
        self.instruction = _instruction(spec, self.command, ordinal, self.subtask)
        if instruction_override is not None:
            self.instruction = instruction_override
        self.autonomous_instruction = AUTONOMOUS_TASK_INSTRUCTIONS.get(
            spec.name, self.instruction
        )
        if self.started_fallen:
            self.force_fall()

    def _setup_indices(self) -> None:
        import mujoco

        m = self.model
        self.trunk_body = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "trunk_base")
        self.trunk_joint = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_JOINT, "trunk_base_freejoint")
        self.trunk_qpos = int(m.jnt_qposadr[self.trunk_joint])
        self.trunk_qvel = int(m.jnt_dofadr[self.trunk_joint])
        self.imu_sensor = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SENSOR, "imu_ang_vel")
        self.head_camera = mujoco.mj_name2id(
            m, mujoco.mjtObj.mjOBJ_CAMERA, "head_camera"
        )
        if self.head_camera < 0:
            raise ValueError("scene does not contain head_camera")
        m.cam_quat[self.head_camera] = HEAD_CAMERA_QUAT
        self.joint_qpos = np.asarray(
            [m.jnt_qposadr[m.actuator_trnid[i, 0]] for i in range(m.nu)], dtype=np.int64
        )
        self.joint_qvel = np.asarray(
            [m.jnt_dofadr[m.actuator_trnid[i, 0]] for i in range(m.nu)], dtype=np.int64
        )
        if m.nu != ACTION_DIM:
            raise ValueError(f"scene has {m.nu} actuators, expected {ACTION_DIM}")
        ball_joint = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_JOINT, "ball_free")
        self.ball_qpos = int(m.jnt_qposadr[ball_joint]) if ball_joint >= 0 else None
        self.ball_qvel = int(m.jnt_dofadr[ball_joint]) if ball_joint >= 0 else None

    def _reset(self, randomize: bool) -> None:
        import mujoco

        key_name = self.spec.start_keyframe
        if self.spec.name == "sitstand" and self.ordinal % 2:
            key_name = "SIT"
            self.subtask = "stand_up"
        elif self.spec.name == "sitstand":
            self.subtask = "sit_then_stand"
        key = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_KEY, key_name)
        if key >= 0:
            mujoco.mj_resetDataKeyframe(self.model, self.data, key)
        else:
            mujoco.mj_resetData(self.model, self.data)
            self.data.qpos[self.trunk_qpos : self.trunk_qpos + 7] = [
                0.0,
                0.0,
                0.1385 if "roller" in self.spec.name else 0.125,
                1.0,
                0.0,
                0.0,
                0.0,
            ]
            self.data.qpos[self.joint_qpos] = DEFAULT_POSE
            self.data.ctrl[:] = DEFAULT_POSE
        if randomize:
            self.data.qpos[self.joint_qpos] += self.rng.uniform(
                -math.radians(1.5), math.radians(1.5), ACTION_DIM
            )
            yaw = self.rng.uniform(-math.radians(8), math.radians(8))
            self.data.qpos[self.trunk_qpos + 3 : self.trunk_qpos + 7] = [
                math.cos(yaw / 2),
                0.0,
                0.0,
                math.sin(yaw / 2),
            ]
            self._randomize_model_visuals()
        self.data.ctrl[:] = DEFAULT_POSE
        self._set_wheel_friction()
        mujoco.mj_forward(self.model, self.data)
        if self.spec.kind == "kick":
            self._place_ball()

    def _randomize_model_visuals(self) -> None:
        import mujoco

        if self.model.nlight:
            self.model.light_diffuse[:] *= self.rng.uniform(0.75, 1.2)
            self.model.light_ambient[:] = np.clip(
                self.model.light_ambient + self.rng.uniform(-0.08, 0.08), 0.0, 1.0
            )
        floor = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "floor")
        if floor >= 0:
            rgb = self.rng.uniform(0.25, 0.75, 3)
            self.model.geom_rgba[floor, :3] = rgb
        if self.head_camera >= 0:
            self.model.cam_pos[self.head_camera] += self.rng.uniform(-0.003, 0.003, 3)
            self.model.cam_fovy[self.head_camera] = np.clip(
                self.model.cam_fovy[self.head_camera] + self.rng.uniform(-3.0, 3.0),
                30.0,
                90.0,
            )

    def _set_wheel_friction(self) -> None:
        import mujoco

        for j in range(self.model.njnt):
            name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_JOINT, j)
            if name and name.startswith("passive_") and "wheel" in name:
                self.model.dof_frictionloss[self.model.jnt_dofadr[j]] = 0.003

    def _sample_command(self) -> None:
        if self.spec.name == "stand":
            caps = np.array([0.35, 0.35, 0.45, 0.2], np.float32)
            self.command[3:7] = self.rng.uniform(-caps, caps)
            self.command[9] = self.rng.uniform(-0.015, 0.015)
            self.command[10:12] = self.rng.uniform(-0.15, 0.15, 2)
        elif self.spec.name == "walking":
            candidates_x = np.array([-0.25, -0.15, 0.15, 0.25])
            candidates_y = np.array([-0.15, 0.0, 0.15])
            candidates_w = np.array([-1.0, -0.5, 0.0, 0.5, 1.0])
            self.command[:3] = [
                self.rng.choice(candidates_x),
                self.rng.choice(candidates_y),
                self.rng.choice(candidates_w),
            ]
        elif self.spec.name == "sitstand":
            self.command[0] = 1.0 if self.subtask == "sit_then_stand" else 0.0
        elif self.spec.name == "roller":
            self.command[0] = self.rng.uniform(0.25, 0.55)

    def _place_ball(self) -> None:
        if self.ball_qpos is None or self.ball_qvel is None:
            raise ValueError("kick scene does not contain ball_free")
        q = self.data.qpos[self.trunk_qpos + 3 : self.trunk_qpos + 7]
        qw, qx, qy, qz = q
        yaw = math.atan2(2 * (qw * qz + qx * qy), 1 - 2 * (qy * qy + qz * qz))
        off_y = (
            BALL_OFFSET_ABS_Y
            if self.spec.name.endswith("left")
            else -BALL_OFFSET_ABS_Y
        )
        off_y += self.rng.uniform(-0.004, 0.004)
        nominal_x = (
            BALL_OFFSET_X
            if self.ball_forward_offset is None
            else self.ball_forward_offset
        )
        off_x = nominal_x + self.rng.uniform(-0.004, 0.004)
        x, y = self.data.qpos[self.trunk_qpos : self.trunk_qpos + 2]
        bx = x + math.cos(yaw) * off_x - math.sin(yaw) * off_y
        by = y + math.sin(yaw) * off_x + math.cos(yaw) * off_y
        self.data.qpos[self.ball_qpos : self.ball_qpos + 7] = [
            bx,
            by,
            BALL_RADIUS,
            1,
            0,
            0,
            0,
        ]
        self.data.qvel[self.ball_qvel : self.ball_qvel + 6] = 0

    def force_fall(self) -> None:
        """Place the robot in a physically valid fallen pose for recovery data."""
        import mujoco

        directions = (
            ("front", 1, 1.0),
            ("back", 1, -1.0),
            ("left", 0, 1.0),
            ("right", 0, -1.0),
        )
        name, axis, sign = directions[int(self.rng.integers(0, len(directions)))]
        angle = sign * (math.pi / 2 + self.rng.uniform(-0.12, 0.12))
        quat = np.zeros(4, np.float64)
        quat[0] = math.cos(angle / 2)
        quat[1 + axis] = math.sin(angle / 2)
        self.data.qpos[self.trunk_qpos + 2] = self.rng.uniform(0.048, 0.062)
        self.data.qpos[self.trunk_qpos + 3 : self.trunk_qpos + 7] = quat
        self.data.qpos[self.joint_qpos] = DEFAULT_POSE + self.rng.uniform(
            -0.08, 0.08, ACTION_DIM
        )
        self.data.qvel[:] = 0
        self.data.ctrl[:] = DEFAULT_POSE
        mujoco.mj_forward(self.model, self.data)
        self.fall_direction = name

    def _observation(self, last_action: np.ndarray, command: np.ndarray) -> np.ndarray:
        adr = self.model.sensor_adr[self.imu_sensor]
        gyro = self.data.sensordata[adr : adr + 3].astype(np.float32).copy()
        quat = self.data.xquat[self.trunk_body].astype(np.float32).copy()
        gravity = _quat_rotate_inverse(quat, np.array([0, 0, -1], np.float32))
        joint_pos = self.data.qpos[self.joint_qpos].astype(np.float32) - DEFAULT_POSE
        joint_vel = self.data.qvel[self.joint_qvel].astype(np.float32).copy()
        obs = np.concatenate([gyro, gravity, joint_pos, joint_vel, last_action, command])
        if obs.shape != (STATE_DIM,) or not np.isfinite(obs).all():
            raise ValueError(f"invalid observation: shape={obs.shape}, finite={np.isfinite(obs).all()}")
        return obs.astype(np.float32)

    def _is_upright(self, gravity_z: float | None = None) -> bool:
        if gravity_z is None:
            quat = self.data.xquat[self.trunk_body].astype(np.float32).copy()
            gravity_z = float(
                _quat_rotate_inverse(quat, np.array([0, 0, -1], np.float32))[2]
            )
        height = float(self.data.qpos[self.trunk_qpos + 2])
        return bool(gravity_z < -0.82 and height > 0.09)

    def _recovery_detection_allowed(self, task_elapsed: float) -> bool:
        if self.spec.kind not in ("kick", "roulade", "phase"):
            return True
        handoff = self.spec.handoff_after_s
        return handoff is None or task_elapsed >= handoff

    def _base_policy_filename_and_command(self, task_elapsed: float):
        if self.control_program is not None:
            return self.control_program(task_elapsed)
        filename = self.spec.policy_file
        command = self.command.copy()
        if self.spec.name == "sitstand":
            command[:] = 0
            if self.subtask == "sit_then_stand" and task_elapsed < 2.0:
                command[0] = 1.0
        if self.spec.phase_period_s is not None and (
            self.spec.handoff_after_s is None or task_elapsed < self.spec.handoff_after_s
        ):
            phase = task_elapsed / self.spec.phase_period_s
            command[:] = 0
            command[0] = math.cos(2 * math.pi * phase)
            command[1] = math.sin(2 * math.pi * phase)
        if (
            self.spec.handoff_after_s is not None
            and task_elapsed >= self.spec.handoff_after_s
        ):
            filename = self.spec.handoff_policy_file or filename
            command[:] = 0
            return filename, command, BALANCE_INSTRUCTION, "balance"
        if self.spec.kind == "recovery":
            return "alpha_stand.onnx", np.zeros_like(command), BALANCE_INSTRUCTION, "balance"
        return filename, command, self.instruction, "task"

    def _control_context(self, step: int):
        elapsed = step / CONTROL_HZ
        task_elapsed = (
            0.0
            if self.task_start_step is None
            else max(0.0, (step - self.task_start_step) / CONTROL_HZ)
        )
        quat = self.data.xquat[self.trunk_body].astype(np.float32).copy()
        gravity_z = float(
            _quat_rotate_inverse(quat, np.array([0, 0, -1], np.float32))[2]
        )
        height = float(self.data.qpos[self.trunk_qpos + 2])
        upright = self._is_upright(gravity_z)

        if self.recovery_active:
            self.recovery_upright_ticks = (
                self.recovery_upright_ticks + 1 if upright else 0
            )
            required = round(RECOVERY_STABLE_S * CONTROL_HZ)
            if self.recovery_upright_ticks >= required:
                self.recovery_active = False
                self.recovery_completed = True
                self.recovery_exit_step = step
                self.task_start_step = step
                self.recovery_upright_ticks = 0
                task_elapsed = 0.0
        elif (
            elapsed < self.spec.duration_s - TERMINAL_STABILIZATION_S
            and self._recovery_detection_allowed(task_elapsed)
            and (gravity_z > -0.58 or height < 0.073)
        ):
            self.recovery_active = True
            self.recovery_enter_step = step
            self.recovery_upright_ticks = 0

        if self.recovery_active:
            return (
                "alpha_stand.onnx",
                np.zeros_like(self.command),
                RECOVERY_INSTRUCTION,
                "recovery",
            )
        if elapsed >= self.spec.duration_s - TERMINAL_STABILIZATION_S:
            return (
                "alpha_stand.onnx",
                np.zeros_like(self.command),
                BALANCE_INSTRUCTION,
                "balance",
            )
        context = self._base_policy_filename_and_command(task_elapsed)
        if context[3] == "task":
            self.task_control_steps += 1
        return context

    def run(self, delay_ticks: int = 0) -> EpisodeResult:
        import mujoco

        total_steps = round(self.spec.duration_s * CONTROL_HZ)
        last_action = np.zeros(ACTION_DIM, np.float32)
        delay = [np.zeros(ACTION_DIM, np.float32) for _ in range(delay_ticks + 1)]
        images: list[np.ndarray] = []
        states: list[np.ndarray] = []
        obs_indices: list[int] = []
        observation_instructions: list[str] = []
        observation_phases: list[str] = []
        actions: list[np.ndarray] = []
        trunk_xyz: list[np.ndarray] = []
        gravity_z: list[float] = []
        cumulative_abs_pitch_rate = 0.0
        ball_start = (
            self.data.qpos[self.ball_qpos : self.ball_qpos + 3].copy()
            if self.ball_qpos is not None
            else None
        )
        previous_program_policy: str | None = None

        for step in range(total_steps):
            filename, command, active_instruction, active_phase = self._control_context(step)
            if (
                self.control_program is not None
                and previous_program_policy is not None
                and filename != previous_program_policy
            ):
                # The source ONNX teachers are stateless policies trained from
                # their own neutral previous-action convention. Reset only the
                # observation feature at a composed teacher handoff; physical
                # controls still pass continuously through the delay buffer.
                last_action = np.zeros(ACTION_DIM, np.float32)
            previous_program_policy = filename
            session = self.sessions.get(filename)
            obs = self._observation(last_action, command)
            inp = session.get_inputs()[0].name
            out = session.get_outputs()[0].name
            action = session.run([out], {inp: obs[None]})[0][0].astype(np.float32)
            if action.shape != (ACTION_DIM,) or not np.isfinite(action).all():
                raise ValueError(f"invalid action from {self.spec.name}: {action}")
            if step % CONTROL_DECIMATION == 0:
                self.renderer.update_scene(self.data, camera="head_camera")
                image = self.renderer.render().copy()
                if image.shape != (self.image_size, self.image_size, 3):
                    raise ValueError(f"unexpected render shape {image.shape}")
                images.append(image)
                states.append(obs)
                obs_indices.append(step)
                observation_instructions.append(active_instruction)
                observation_phases.append(active_phase)
            actions.append(action)
            delay.append(action.copy())
            applied = delay.pop(0)
            self.data.ctrl[:] = DEFAULT_POSE + applied
            for _ in range(4):
                mujoco.mj_step(self.model, self.data)
            last_action = action
            trunk_xyz.append(self.data.qpos[self.trunk_qpos : self.trunk_qpos + 3].copy())
            gravity_z.append(float(self._observation(last_action, command)[5]))
            cumulative_abs_pitch_rate += abs(float(self.data.qvel[self.trunk_qvel + 4])) / CONTROL_HZ

        return self._finish_result(
            images,
            states,
            obs_indices,
            actions,
            trunk_xyz,
            gravity_z,
            cumulative_abs_pitch_rate,
            ball_start,
            observation_instructions,
            observation_phases,
            None,
        )

    def run_chunk_policy(
        self,
        policy,
        delay_ticks: int = 0,
        record_external: bool = False,
        implicit_phase_instruction: bool = False,
    ) -> EpisodeResult:
        """Run a vision policy that replans at the 10 Hz observation rate.

        ``policy(image, state, instruction, control_index)`` returns at least
        five raw 14D joint-offset actions. The simulator executes the first
        five at 50 Hz, then captures a fresh head-camera observation.
        """
        import mujoco

        total_steps = round(self.spec.duration_s * CONTROL_HZ)
        last_action = np.zeros(ACTION_DIM, np.float32)
        delay = [np.zeros(ACTION_DIM, np.float32) for _ in range(delay_ticks + 1)]
        images: list[np.ndarray] = []
        states: list[np.ndarray] = []
        obs_indices: list[int] = []
        observation_instructions: list[str] = []
        observation_phases: list[str] = []
        video_frames: list[np.ndarray] = []
        actions: list[np.ndarray] = []
        trunk_xyz: list[np.ndarray] = []
        gravity_z: list[float] = []
        cumulative_abs_pitch_rate = 0.0
        ball_start = (
            self.data.qpos[self.ball_qpos : self.ball_qpos + 3].copy()
            if self.ball_qpos is not None
            else None
        )
        active_chunk: np.ndarray | None = None
        tracking_camera = None
        if record_external:
            tracking_camera = mujoco.MjvCamera()
            mujoco.mjv_defaultCamera(tracking_camera)
            tracking_camera.type = mujoco.mjtCamera.mjCAMERA_TRACKING
            tracking_camera.trackbodyid = self.trunk_body
            tracking_camera.distance = 0.62
            tracking_camera.azimuth = 135.0
            tracking_camera.elevation = -18.0

        for step in range(total_steps):
            _, command, active_instruction, active_phase = self._control_context(step)
            policy_instruction = (
                self.autonomous_instruction
                if implicit_phase_instruction
                else active_instruction
            )
            obs = self._observation(last_action, command)
            if step % CONTROL_DECIMATION == 0:
                self.renderer.update_scene(self.data, camera="head_camera")
                image = self.renderer.render().copy()
                active_chunk = np.asarray(
                    policy(image, obs, policy_instruction, step), dtype=np.float32
                )
                if (
                    active_chunk.ndim != 2
                    or active_chunk.shape[0] < CONTROL_DECIMATION
                    or active_chunk.shape[1] != ACTION_DIM
                    or not np.isfinite(active_chunk).all()
                ):
                    raise ValueError(f"vision policy returned invalid chunk {active_chunk.shape}")
                images.append(image)
                states.append(obs)
                obs_indices.append(step)
                observation_instructions.append(policy_instruction)
                observation_phases.append(active_phase)
                if tracking_camera is not None:
                    self.renderer.update_scene(self.data, camera=tracking_camera)
                    video_frames.append(self.renderer.render().copy())
            assert active_chunk is not None
            action = active_chunk[step % CONTROL_DECIMATION]
            actions.append(action.copy())
            delay.append(action.copy())
            applied = delay.pop(0)
            self.data.ctrl[:] = DEFAULT_POSE + applied
            for _ in range(4):
                mujoco.mj_step(self.model, self.data)
            last_action = action
            trunk_xyz.append(self.data.qpos[self.trunk_qpos : self.trunk_qpos + 3].copy())
            gravity_z.append(float(self._observation(last_action, command)[5]))
            cumulative_abs_pitch_rate += (
                abs(float(self.data.qvel[self.trunk_qvel + 4])) / CONTROL_HZ
            )

        return self._finish_result(
            images,
            states,
            obs_indices,
            actions,
            trunk_xyz,
            gravity_z,
            cumulative_abs_pitch_rate,
            ball_start,
            observation_instructions,
            observation_phases,
            np.asarray(video_frames, np.uint8) if video_frames else None,
        )

    def _finish_result(
        self,
        images,
        states,
        obs_indices,
        actions,
        trunk_xyz,
        gravity_z,
        cumulative_abs_pitch_rate,
        ball_start,
        observation_instructions,
        observation_phases,
        video_frames,
    ) -> EpisodeResult:
        self.renderer.close()
        xyz = np.asarray(trunk_xyz)
        upright = (np.asarray(gravity_z) < -0.72) & (xyz[:, 2] > 0.09)
        terminal_steps = max(1, round(TERMINAL_STABILIZATION_S * CONTROL_HZ / 2))
        terminal_upright_fraction = float(upright[-terminal_steps:].mean())
        terminal_stable = bool(terminal_upright_fraction >= 0.8 and upright[-1])
        final_upright = bool(upright[-1])
        planar_displacement = float(np.linalg.norm(xyz[-1, :2] - xyz[0, :2]))
        planar_path_length = float(
            np.linalg.norm(np.diff(xyz[:, :2], axis=0), axis=1).sum()
        )
        min_height = float(xyz[:, 2].min())
        ball_displacement = 0.0
        if ball_start is not None:
            ball_end = self.data.qpos[self.ball_qpos : self.ball_qpos + 3]
            ball_displacement = float(np.linalg.norm(ball_end[:2] - ball_start[:2]))

        if self.spec.kind == "recovery":
            success = bool(self.recovery_completed and terminal_stable)
        elif self.started_fallen:
            success = bool(
                self.recovery_completed
                and self.task_control_steps >= CONTROL_HZ // 2
                and terminal_stable
            )
        elif self.spec.name == "sitstand" and self.subtask == "sit_then_stand":
            success = bool(min_height < xyz[0, 2] - 0.008 and terminal_stable)
        elif self.spec.kind == "kick":
            success = bool(ball_displacement > 0.03 and terminal_stable)
        elif self.spec.kind == "roulade":
            success = bool(cumulative_abs_pitch_rate > 2.5 and terminal_stable)
        elif self.spec.name in ("walking", "roller"):
            success = bool(planar_displacement > 0.025 and terminal_stable)
        elif self.spec.kind == "phase":
            success = bool(min_height < xyz[0, 2] - 0.01 and terminal_stable)
        else:
            success = bool(terminal_stable)

        metrics: dict[str, float | bool | str] = {
            "success": success,
            "final_upright": final_upright,
            "planar_displacement_m": planar_displacement,
            "planar_path_length_m": planar_path_length,
            "initial_trunk_height_m": float(xyz[0, 2]),
            "minimum_trunk_height_m": min_height,
            "final_trunk_height_m": float(xyz[-1, 2]),
            "ball_displacement_m": ball_displacement,
            "integrated_abs_pitch_rate_rad": cumulative_abs_pitch_rate,
            "subtask": self.subtask,
            "terminal_upright_fraction": terminal_upright_fraction,
            "terminal_stable": terminal_stable,
            "started_fallen": self.started_fallen,
            "fall_direction": self.fall_direction,
            "recovery_completed": self.recovery_completed,
            "recovery_time_s": (
                None
                if self.recovery_exit_step is None
                else self.recovery_exit_step / CONTROL_HZ
            ),
            "task_control_steps": self.task_control_steps,
        }
        return EpisodeResult(
            images=np.asarray(images, np.uint8),
            states=np.asarray(states, np.float32),
            observation_control_indices=np.asarray(obs_indices, np.int32),
            actions=np.asarray(actions, np.float32),
            instruction=self.instruction,
            success=success,
            metrics=metrics,
            command=self.command,
            subtask=self.subtask,
            observation_instructions=tuple(observation_instructions),
            observation_phases=tuple(observation_phases),
            video_frames=video_frames,
        )


def encode_episode_metadata(metadata: dict) -> np.ndarray:
    def convert(value):
        if isinstance(value, np.generic):
            return value.item()
        if isinstance(value, np.ndarray):
            return value.tolist()
        raise TypeError(f"cannot JSON-encode {type(value).__name__}")

    return np.asarray(json.dumps(metadata, ensure_ascii=False, default=convert))


def decode_episode_metadata(value: np.ndarray) -> dict:
    return json.loads(str(value.item()))