code/physics_show.py

"""Run the 25-duck Electric Slide with MuJoCo dynamics and Pollen's walking policy."""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
import subprocess

import mujoco
import numpy as np
import onnxruntime as ort

from choreography import COUNTS_PER_WALL, SHOW_SECONDS, STAND_POSE
from generate_scene import DEFAULT_OUTPUT, build_scene
from physics_choreography import command_for_time, policy_for_time
from literal_choreography import (
    LITERAL_COUNTS,
    command_vector as literal_command_vector,
    count_phase as literal_count_phase,
    joint_bias as literal_joint_bias,
    policy_for_time as literal_policy_for_time,
)
from simple_choreography import (
    COUNTS_PER_WALL as SIMPLE_COUNTS_PER_WALL,
    SHOW_SECONDS as SIMPLE_SHOW_SECONDS,
    SIMPLE_COUNTS,
    command_vector as simple_command_vector,
    count_phase as simple_count_phase,
    joint_bias as simple_joint_bias,
    policy_for_time as simple_policy_for_time,
)


DUCK_JOINTS = 14
CONTROL_HZ = 50
PHYSICS_HZ = 200
DECIMATION = PHYSICS_HZ // CONTROL_HZ
DEFAULT_POLICY = (
    Path(__file__).resolve().parent
    / "official/app/public/policies/BEST_alpha_walking.onnx"
)
DEFAULT_KICK_LEFT_POLICY = DEFAULT_POLICY.with_name("ball_kick_left.onnx")
DEFAULT_KICK_RIGHT_POLICY = DEFAULT_POLICY.with_name("ball_kick_right.onnx")


def quat_rotate_inverse(quaternion: np.ndarray, vector: np.ndarray) -> np.ndarray:
    w = quaternion[0]
    xyz = quaternion[1:4]
    cross = np.cross(xyz, vector) * 2.0
    return vector - w * cross + np.cross(xyz, cross)


def yaw_from_matrix(matrix: np.ndarray) -> float:
    rotation = matrix.reshape(3, 3)
    return math.atan2(float(rotation[1, 0]), float(rotation[0, 0]))


def pitch_from_matrix(matrix: np.ndarray) -> float:
    rotation = matrix.reshape(3, 3)
    return math.asin(float(np.clip(-rotation[2, 0], -1.0, 1.0)))


class MultiDuckPolicy:
    """One shared ONNX network with independent state for every duck."""

    def __init__(
        self,
        model: mujoco.MjModel,
        data: mujoco.MjData,
        policy_path: Path,
        grid: int,
        fixed_command: tuple[float, float, float] | None = None,
        kick_left_path: Path = DEFAULT_KICK_LEFT_POLICY,
        kick_right_path: Path = DEFAULT_KICK_RIGHT_POLICY,
        synchronized_actions: bool = False,
        literal: bool = False,
        simple: bool = False,
    ) -> None:
        self.model = model
        self.data = data
        self.grid = grid
        self.count = grid * grid
        self.fixed_command = fixed_command
        self.synchronized_actions = synchronized_actions
        self.literal = literal
        self.simple = simple
        if literal and simple:
            raise ValueError("Literal and simple choreography modes are mutually exclusive")
        self.count_specs = SIMPLE_COUNTS if simple else LITERAL_COUNTS
        self.dance_counts_per_wall = SIMPLE_COUNTS_PER_WALL if simple else COUNTS_PER_WALL
        if model.nu != self.count * DUCK_JOINTS:
            raise ValueError(f"Expected {self.count * DUCK_JOINTS} actuators, got {model.nu}")

        options = ort.SessionOptions()
        options.intra_op_num_threads = 1
        options.inter_op_num_threads = 1
        self.sessions = {
            "walk": ort.InferenceSession(str(policy_path), sess_options=options),
            "kick_left": ort.InferenceSession(str(kick_left_path), sess_options=options),
            "kick_right": ort.InferenceSession(str(kick_right_path), sess_options=options),
        }
        self.input_name = self.sessions["walk"].get_inputs()[0].name
        self.output_name = self.sessions["walk"].get_outputs()[0].name
        for name, session in self.sessions.items():
            if session.get_inputs()[0].shape != [1, 61]:
                raise ValueError(f"Official {name} policy input must be [1, 61]")
            if session.get_outputs()[0].shape != [1, 14]:
                raise ValueError(f"Official {name} policy output must be [1, 14]")

        metadata = self.sessions["walk"].get_modelmeta().custom_metadata_map
        self.action_scale = float(metadata.get("action_scale", "1.0"))
        pose_text = metadata.get("default_joint_pos")
        if pose_text:
            self.default_pose = np.array([float(value) for value in pose_text.split(",")], dtype=np.float32)
        else:
            self.default_pose = np.array(STAND_POSE, dtype=np.float32)
        if self.default_pose.shape != (DUCK_JOINTS,):
            raise ValueError(f"Policy default pose has shape {self.default_pose.shape}")

        self.actuator_ids = np.arange(model.nu, dtype=np.int32).reshape(self.count, DUCK_JOINTS)
        self.joint_qpos = np.empty((self.count, DUCK_JOINTS), dtype=np.int32)
        self.joint_qvel = np.empty((self.count, DUCK_JOINTS), dtype=np.int32)
        self.joint_lower = np.empty((self.count, DUCK_JOINTS), dtype=np.float32)
        self.joint_upper = np.empty((self.count, DUCK_JOINTS), dtype=np.float32)
        self.body_ids = np.empty(self.count, dtype=np.int32)
        self.free_qpos = np.empty(self.count, dtype=np.int32)
        self.free_qvel = np.empty(self.count, dtype=np.int32)
        self.imu_sensor_addresses = np.empty(self.count, dtype=np.int32)
        self.foot_geom_ids = np.empty((self.count, 2), dtype=np.int32)

        for duck in range(self.count):
            prefix = f"duck_{duck:02d}_"
            expected_first = prefix + "left_hip_yaw"
            actual_first = mujoco.mj_id2name(
                model,
                mujoco.mjtObj.mjOBJ_ACTUATOR,
                int(self.actuator_ids[duck, 0]),
            )
            if actual_first != expected_first:
                raise ValueError(f"Actuator ordering mismatch: expected {expected_first}, got {actual_first}")

            for local, actuator_id in enumerate(self.actuator_ids[duck]):
                joint_id = int(model.actuator_trnid[int(actuator_id), 0])
                self.joint_qpos[duck, local] = int(model.jnt_qposadr[joint_id])
                self.joint_qvel[duck, local] = int(model.jnt_dofadr[joint_id])
                self.joint_lower[duck, local] = float(model.jnt_range[joint_id, 0])
                self.joint_upper[duck, local] = float(model.jnt_range[joint_id, 1])

            body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, prefix + "trunk_base")
            free_id = mujoco.mj_name2id(
                model,
                mujoco.mjtObj.mjOBJ_JOINT,
                prefix + "trunk_base_freejoint",
            )
            sensor_id = mujoco.mj_name2id(
                model,
                mujoco.mjtObj.mjOBJ_SENSOR,
                prefix + "imu_ang_vel",
            )
            if min(body_id, free_id, sensor_id) < 0:
                raise ValueError(f"Missing body/freejoint/IMU for {prefix}")
            self.body_ids[duck] = body_id
            self.free_qpos[duck] = int(model.jnt_qposadr[free_id])
            self.free_qvel[duck] = int(model.jnt_dofadr[free_id])
            self.imu_sensor_addresses[duck] = int(model.sensor_adr[sensor_id])
            self.foot_geom_ids[duck, 0] = mujoco.mj_name2id(
                model,
                mujoco.mjtObj.mjOBJ_GEOM,
                prefix + "left_foot_collision",
            )
            self.foot_geom_ids[duck, 1] = mujoco.mj_name2id(
                model,
                mujoco.mjtObj.mjOBJ_GEOM,
                prefix + "right_foot_collision",
            )

        self.last_actions = np.zeros((self.count, DUCK_JOINTS), dtype=np.float32)
        self.policy_steps = {name: 0 for name in self.sessions}
        self.literal_initial_yaw = np.zeros(self.count, dtype=np.float32)

    @property
    def dance_active(self) -> bool:
        return self.literal or self.simple

    def dance_count_phase(self, show_time: float) -> tuple[int, int, float]:
        return simple_count_phase(show_time) if self.simple else literal_count_phase(show_time)

    def reset(self) -> None:
        key_id = mujoco.mj_name2id(
            self.model,
            mujoco.mjtObj.mjOBJ_KEY,
            "STAND_5X5",
        )
        if key_id < 0:
            raise ValueError("Generated scene is missing STAND_5X5")
        mujoco.mj_resetDataKeyframe(self.model, self.data, key_id)
        self.data.ctrl[:] = np.tile(self.default_pose, self.count)
        self.last_actions.fill(0.0)
        mujoco.mj_forward(self.model, self.data)
        self.literal_initial_yaw = np.array(
            [yaw_from_matrix(self.data.xmat[int(body_id)]) for body_id in self.body_ids],
            dtype=np.float32,
        )

    def _literal_heading_action(
        self,
        show_time: float,
        duck: int,
        kick_action: np.ndarray,
        base_command: np.ndarray,
    ) -> np.ndarray:
        """Blend the swing primitive with a physical heading correction gait."""
        wall, count, _ = self.dance_count_phase(show_time)
        if count == self.dance_counts_per_wall - 1:
            return kick_action
        current_yaw = yaw_from_matrix(self.data.xmat[int(self.body_ids[duck])])
        target_yaw = float(self.literal_initial_yaw[duck]) + wall * (math.pi / 2.0)
        yaw_error = math.atan2(
            math.sin(target_yaw - current_yaw),
            math.cos(target_yaw - current_yaw),
        )
        correction_command = base_command.copy()
        correction_command[2] = float(np.clip(7.0 * yaw_error, -5.0, 5.0))
        walk_observation = self.observation(duck, correction_command)
        walk_action = self.sessions["walk"].run(
            [self.output_name],
            {self.input_name: walk_observation.reshape(1, -1)},
        )[0][0].astype(np.float32)
        self.policy_steps["walk"] += 1
        return kick_action * 0.72 + walk_action * 0.28

    def command_vector(
        self,
        show_time: float | None,
        duck: int,
        policy_name: str = "walk",
    ) -> np.ndarray:
        if show_time is None:
            return np.zeros(13, dtype=np.float32)
        if self.dance_active:
            command = simple_command_vector(show_time) if self.simple else literal_command_vector(show_time)
            wall, count, _ = self.dance_count_phase(show_time)
            if count != self.dance_counts_per_wall - 1:
                current_yaw = yaw_from_matrix(self.data.xmat[int(self.body_ids[duck])])
                target_yaw = float(self.literal_initial_yaw[duck]) + wall * (math.pi / 2.0)
                yaw_error = math.atan2(
                    math.sin(target_yaw - current_yaw),
                    math.cos(target_yaw - current_yaw),
                )
                gain = 4.0 if count == 0 and wall > 0 else 7.0
                limit = 3.0 if count == 0 and wall > 0 else 5.0
                command[2] = float(np.clip(gain * yaw_error, -limit, limit))
            return command
        if policy_name != "walk":
            return np.zeros(13, dtype=np.float32)
        if self.fixed_command is not None:
            return np.array((*self.fixed_command, *([0.0] * 10)), dtype=np.float32)
        command = command_for_time(show_time, duck)
        return np.array(
            (*command.twist, *command.head_pose, *command.body_pose),
            dtype=np.float32,
        )

    def observation(self, duck: int, command: np.ndarray) -> np.ndarray:
        sensor_address = int(self.imu_sensor_addresses[duck])
        angular_velocity = self.data.sensordata[
            sensor_address : sensor_address + 3
        ].astype(np.float32)
        quaternion = self.data.xquat[int(self.body_ids[duck])].astype(np.float32)
        projected_gravity = quat_rotate_inverse(
            quaternion,
            np.array([0.0, 0.0, -1.0], dtype=np.float32),
        ).astype(np.float32)
        joint_position = self.data.qpos[self.joint_qpos[duck]].astype(np.float32)
        joint_velocity = self.data.qvel[self.joint_qvel[duck]].astype(np.float32)
        observation = np.concatenate(
            (
                angular_velocity,
                projected_gravity,
                joint_position - self.default_pose,
                joint_velocity,
                self.last_actions[duck],
                command,
            )
        ).astype(np.float32)
        if observation.shape != (61,):
            raise RuntimeError(f"Observation shape for duck {duck}: {observation.shape}")
        return observation

    def infer_and_apply(self, show_time: float | None) -> None:
        policy_name = "walk"
        if show_time is not None and self.fixed_command is None:
            if self.simple:
                policy_name = simple_policy_for_time(show_time)
            elif self.literal:
                policy_name = literal_policy_for_time(show_time)
            else:
                policy_name = policy_for_time(show_time)
        session = self.sessions[policy_name]
        self.policy_steps[policy_name] += 1
        bias = np.zeros(DUCK_JOINTS, dtype=np.float32)
        if show_time is not None and self.fixed_command is None:
            if self.simple:
                bias = simple_joint_bias(show_time)
            elif self.literal:
                bias = literal_joint_bias(show_time)
        if self.synchronized_actions:
            leader = (self.count - 1) // 2
            command = self.command_vector(show_time, leader, policy_name)
            observation = self.observation(leader, command)
            action = session.run(
                [self.output_name],
                {self.input_name: observation.reshape(1, -1)},
            )[0][0].astype(np.float32)
            self.last_actions[:] = action
            target = self.default_pose + action * self.action_scale + bias
            target = np.clip(target, self.joint_lower[0], self.joint_upper[0])
            self.data.ctrl[:] = np.tile(target, self.count)
            return
        for duck in range(self.count):
            command = self.command_vector(show_time, duck, policy_name)
            observation = self.observation(duck, command)
            action = session.run(
                [self.output_name],
                {self.input_name: observation.reshape(1, -1)},
            )[0][0].astype(np.float32)
            action = np.clip(action, -1.8, 1.8)
            if self.dance_active and show_time is not None and policy_name != "walk":
                action = self._literal_heading_action(
                    show_time,
                    duck,
                    action,
                    command,
                )
            self.last_actions[duck] = action
            target = self.default_pose + action * self.action_scale + bias
            self.data.ctrl[self.actuator_ids[duck]] = np.clip(
                target,
                self.joint_lower[duck],
                self.joint_upper[duck],
            )


class ContactDiagnostics:
    def __init__(self, controller: MultiDuckPolicy, initial_xy: np.ndarray) -> None:
        self.controller = controller
        self.initial_xy = initial_xy
        self.initial_yaw = np.array(
            [yaw_from_matrix(controller.data.xmat[int(body_id)]) for body_id in controller.body_ids]
        )
        self.focus_duck = (controller.count - 1) // 2
        self.previous_yaw = self.initial_yaw.copy()
        self.cumulative_yaw = np.zeros(controller.count)
        self.floor_id = mujoco.mj_name2id(
            controller.model,
            mujoco.mjtObj.mjOBJ_GEOM,
            "dance_floor",
        )
        self.foot_lookup = {
            int(geom): (duck, side)
            for duck in range(controller.count)
            for side, geom in enumerate(controller.foot_geom_ids[duck])
        }
        self.previous = np.zeros((controller.count, 2), dtype=bool)
        self.strikes = np.zeros((controller.count, 2), dtype=np.int32)
        self.liftoffs = np.zeros((controller.count, 2), dtype=np.int32)
        self.airborne_frames = np.zeros(controller.count, dtype=np.int32)
        self.frames = 0
        self.minimum_trunk_z = np.full(controller.count, np.inf)
        self.first_fall_time_seconds: list[float | None] = [None] * controller.count
        self.maximum_excursion = np.zeros(controller.count)
        self.maximum_abs_action = 0.0
        self.action_abs_sum = 0.0
        self.action_samples = 0
        self.minimum_pairwise_separation = np.inf
        self.geom_duck_owner = np.full(controller.model.ngeom, -1, dtype=np.int32)
        for geom_id in range(controller.model.ngeom):
            name = mujoco.mj_id2name(
                controller.model,
                mujoco.mjtObj.mjOBJ_GEOM,
                geom_id,
            )
            if name and name.startswith("duck_") and len(name) >= 7:
                try:
                    self.geom_duck_owner[geom_id] = int(name[5:7])
                except ValueError:
                    pass
        self.inter_duck_contact_count = 0
        self.inter_duck_contact_frames = 0
        self.count_liftoffs = np.zeros((controller.dance_counts_per_wall, 2), dtype=np.int32)
        self.count_strikes = np.zeros((controller.dance_counts_per_wall, 2), dtype=np.int32)
        self.focus_count_liftoffs = np.zeros((controller.dance_counts_per_wall, 2), dtype=np.int32)
        self.focus_count_strikes = np.zeros((controller.dance_counts_per_wall, 2), dtype=np.int32)
        self.active_focus_liftoffs = np.zeros(2, dtype=np.int32)
        self.active_focus_strikes = np.zeros(2, dtype=np.int32)
        self.count_events: list[dict] = []
        self.focus_trace: list[dict] = []
        self.active_wall = 0
        self.active_count = 0
        self.count_start_xy = initial_xy.copy()
        self.count_start_yaw = self.initial_yaw.copy()
        focus_body = int(controller.body_ids[self.focus_duck])
        focus_pitch = pitch_from_matrix(controller.data.xmat[focus_body])
        focus_joints = controller.data.qpos[controller.joint_qpos[self.focus_duck]].copy()
        self.active_focus_pitch_min = focus_pitch
        self.active_focus_pitch_max = focus_pitch
        self.active_focus_joint_min = focus_joints.copy()
        self.active_focus_joint_max = focus_joints.copy()

    def _finish_count(self, wall: int, count: int, elapsed: float) -> None:
        current_xy = np.array(
            [
                self.controller.data.qpos[int(address) : int(address) + 2]
                for address in self.controller.free_qpos
            ]
        )
        current_yaw = np.array(
            [yaw_from_matrix(self.controller.data.xmat[int(body_id)]) for body_id in self.controller.body_ids]
        )
        world_delta = current_xy - self.count_start_xy
        cosine = np.cos(self.count_start_yaw)
        sine = np.sin(self.count_start_yaw)
        local_forward = cosine * world_delta[:, 0] + sine * world_delta[:, 1]
        local_left = -sine * world_delta[:, 0] + cosine * world_delta[:, 1]
        yaw_delta = np.arctan2(
            np.sin(current_yaw - self.count_start_yaw),
            np.cos(current_yaw - self.count_start_yaw),
        )
        focus = self.focus_duck
        self.count_events.append(
            {
                "wall": int(wall),
                "count": int(count + 1),
                "label": self.controller.count_specs[count].label,
                "elapsed_seconds": float(elapsed),
                "mean_local_displacement_m": [
                    float(local_forward.mean()),
                    float(local_left.mean()),
                ],
                "mean_yaw_change_radians": float(yaw_delta.mean()),
                "center_duck_index": int(focus),
                "center_duck_start_xy_m": self.count_start_xy[focus].astype(float).tolist(),
                "center_duck_end_xy_m": current_xy[focus].astype(float).tolist(),
                "center_duck_local_displacement_m": [
                    float(local_forward[focus]),
                    float(local_left[focus]),
                ],
                "center_duck_yaw_change_radians": float(yaw_delta[focus]),
                "center_duck_foot_liftoffs_left_right": self.active_focus_liftoffs.tolist(),
                "center_duck_foot_strikes_left_right": self.active_focus_strikes.tolist(),
                "center_duck_body_pitch_range_degrees": [
                    math.degrees(self.active_focus_pitch_min),
                    math.degrees(self.active_focus_pitch_max),
                ],
                "center_duck_joint_min_14d": self.active_focus_joint_min.astype(float).tolist(),
                "center_duck_joint_max_14d": self.active_focus_joint_max.astype(float).tolist(),
            }
        )
        self.count_start_xy = current_xy.copy()
        self.count_start_yaw = current_yaw.copy()
        focus_body = int(self.controller.body_ids[focus])
        focus_pitch = pitch_from_matrix(self.controller.data.xmat[focus_body])
        focus_joints = self.controller.data.qpos[self.controller.joint_qpos[focus]].copy()
        self.active_focus_liftoffs.fill(0)
        self.active_focus_strikes.fill(0)
        self.active_focus_pitch_min = focus_pitch
        self.active_focus_pitch_max = focus_pitch
        self.active_focus_joint_min = focus_joints.copy()
        self.active_focus_joint_max = focus_joints.copy()

    def update(self, show_time: float | None = None) -> None:
        if self.controller.dance_active and show_time is not None:
            wall, count, _ = self.controller.dance_count_phase(show_time)
            if (wall, count) != (self.active_wall, self.active_count):
                self._finish_count(self.active_wall, self.active_count, show_time)
                self.active_wall, self.active_count = wall, count
        contacts = np.zeros_like(self.previous)
        inter_duck_this_frame = False
        for index in range(self.controller.data.ncon):
            contact = self.controller.data.contact[index]
            geom1 = int(contact.geom1)
            geom2 = int(contact.geom2)
            if geom1 in self.foot_lookup and geom2 == self.floor_id:
                duck, side = self.foot_lookup[geom1]
                contacts[duck, side] = True
            elif geom2 in self.foot_lookup and geom1 == self.floor_id:
                duck, side = self.foot_lookup[geom2]
                contacts[duck, side] = True
            owner1 = int(self.geom_duck_owner[geom1])
            owner2 = int(self.geom_duck_owner[geom2])
            if owner1 >= 0 and owner2 >= 0 and owner1 != owner2:
                self.inter_duck_contact_count += 1
                inter_duck_this_frame = True
        self.inter_duck_contact_frames += int(inter_duck_this_frame)
        strikes = np.logical_and(contacts, np.logical_not(self.previous))
        liftoffs = np.logical_and(np.logical_not(contacts), self.previous)
        self.strikes += strikes
        self.liftoffs += liftoffs
        if self.controller.dance_active and show_time is not None:
            self.count_strikes[self.active_count] += strikes.sum(axis=0)
            self.count_liftoffs[self.active_count] += liftoffs.sum(axis=0)
            focus = self.focus_duck
            self.focus_count_strikes[self.active_count] += strikes[focus]
            self.focus_count_liftoffs[self.active_count] += liftoffs[focus]
            self.active_focus_strikes += strikes[focus]
            self.active_focus_liftoffs += liftoffs[focus]
        self.airborne_frames += np.logical_not(np.any(contacts, axis=1))
        self.previous = contacts
        self.frames += 1
        positions_xy = np.array(
            [
                self.controller.data.qpos[int(address) : int(address) + 2]
                for address in self.controller.free_qpos
            ]
        )
        deltas = positions_xy[:, None, :] - positions_xy[None, :, :]
        distances = np.linalg.norm(deltas, axis=2)
        distances[np.diag_indices_from(distances)] = np.inf
        self.minimum_pairwise_separation = min(
            self.minimum_pairwise_separation,
            float(distances.min()),
        )
        current_yaw = np.array(
            [
                yaw_from_matrix(self.controller.data.xmat[int(body_id)])
                for body_id in self.controller.body_ids
            ]
        )
        yaw_delta = np.arctan2(
            np.sin(current_yaw - self.previous_yaw),
            np.cos(current_yaw - self.previous_yaw),
        )
        self.cumulative_yaw += yaw_delta
        self.previous_yaw = current_yaw
        if self.controller.dance_active and show_time is not None:
            focus = self.focus_duck
            focus_body = int(self.controller.body_ids[focus])
            focus_pitch = pitch_from_matrix(self.controller.data.xmat[focus_body])
            focus_joints = self.controller.data.qpos[self.controller.joint_qpos[focus]].copy()
            self.active_focus_pitch_min = min(self.active_focus_pitch_min, focus_pitch)
            self.active_focus_pitch_max = max(self.active_focus_pitch_max, focus_pitch)
            self.active_focus_joint_min = np.minimum(self.active_focus_joint_min, focus_joints)
            self.active_focus_joint_max = np.maximum(self.active_focus_joint_max, focus_joints)
            wall, count, phase = self.controller.dance_count_phase(show_time)
            focus_address = int(self.controller.free_qpos[focus])
            self.focus_trace.append(
                {
                    "time_seconds": float(show_time),
                    "wall": int(wall + 1),
                    "count": int(count + 1),
                    "phase": float(phase),
                    "body_xy_m": self.controller.data.qpos[focus_address : focus_address + 2].astype(float).tolist(),
                    "body_z_m": float(self.controller.data.qpos[focus_address + 2]),
                    "yaw_radians": float(current_yaw[focus]),
                    "body_pitch_degrees": math.degrees(focus_pitch),
                    "foot_contacts_left_right": contacts[focus].tolist(),
                }
            )
        for duck, qpos_address in enumerate(self.controller.free_qpos):
            current_xy = self.controller.data.qpos[
                int(qpos_address) : int(qpos_address) + 2
            ]
            self.maximum_excursion[duck] = max(
                self.maximum_excursion[duck],
                float(np.linalg.norm(current_xy - self.initial_xy[duck])),
            )
            self.minimum_trunk_z[duck] = min(
                self.minimum_trunk_z[duck],
                float(self.controller.data.qpos[int(qpos_address) + 2]),
            )
            if (
                show_time is not None
                and self.first_fall_time_seconds[duck] is None
                and self.controller.data.qpos[int(qpos_address) + 2] < 0.07
            ):
                self.first_fall_time_seconds[duck] = float(show_time)
        absolute_actions = np.abs(self.controller.last_actions)
        self.maximum_abs_action = max(
            self.maximum_abs_action,
            float(absolute_actions.max()),
        )
        self.action_abs_sum += float(absolute_actions.sum())
        self.action_samples += int(absolute_actions.size)

    def report(self, duration: float) -> dict:
        if self.controller.dance_active:
            self._finish_count(self.active_wall, self.active_count, duration)
        final_xy = np.array(
            [
                self.controller.data.qpos[int(address) : int(address) + 2]
                for address in self.controller.free_qpos
            ]
        )
        displacement = np.linalg.norm(final_xy - self.initial_xy, axis=1)
        airborne = self.airborne_frames / max(1, self.frames)
        return {
            "physics_timestep_seconds": float(self.controller.model.opt.timestep),
            "control_hz": CONTROL_HZ,
            "duration_seconds": duration,
            "root_pose_writes_after_reset": 0,
            "shared_policy_actions": self.controller.synchronized_actions,
            "controller_mode": (
                "simple_16_count"
                if self.controller.simple
                else "literal_18_count"
                if self.controller.literal
                else "learned_expression"
            ),
            "minimum_trunk_z_m": self.minimum_trunk_z.tolist(),
            "fallen_ducks": np.flatnonzero(self.minimum_trunk_z < 0.07).tolist(),
            "first_fall_time_seconds": self.first_fall_time_seconds,
            "foot_strikes_per_duck": self.strikes.sum(axis=1).tolist(),
            "minimum_foot_strikes": int(self.strikes.sum(axis=1).min()),
            "foot_liftoffs_per_duck": self.liftoffs.sum(axis=1).tolist(),
            "minimum_foot_liftoffs": int(self.liftoffs.sum(axis=1).min()),
            "airborne_fraction_per_duck": airborne.tolist(),
            "maximum_airborne_fraction": float(airborne.max()),
            "net_displacement_m": displacement.tolist(),
            "mean_net_displacement_m": float(displacement.mean()),
            "yaw_change_radians": self.cumulative_yaw.tolist(),
            "mean_yaw_change_radians": float(self.cumulative_yaw.mean()),
            "yaw_change_spread_radians": float(np.ptp(self.cumulative_yaw)),
            "minimum_pairwise_separation_m": self.minimum_pairwise_separation,
            "inter_duck_contact_count": self.inter_duck_contact_count,
            "inter_duck_contact_frames": self.inter_duck_contact_frames,
            "maximum_excursion_m": self.maximum_excursion.tolist(),
            "minimum_maximum_excursion_m": float(self.maximum_excursion.min()),
            "maximum_abs_policy_action": self.maximum_abs_action,
            "mean_abs_policy_action": self.action_abs_sum / max(1, self.action_samples),
            "policy_control_steps": self.controller.policy_steps,
            "count_liftoffs_left_right": self.count_liftoffs.tolist(),
            "count_strikes_left_right": self.count_strikes.tolist(),
            "center_duck_index": int(self.focus_duck),
            "center_duck_count_liftoffs_left_right": self.focus_count_liftoffs.tolist(),
            "center_duck_count_strikes_left_right": self.focus_count_strikes.tolist(),
            "center_duck_trace_50hz": self.focus_trace,
            "count_events": self.count_events,
        }


def make_camera(distance: float, azimuth: float, elevation: float) -> mujoco.MjvCamera:
    camera = mujoco.MjvCamera()
    mujoco.mjv_defaultCamera(camera)
    camera.type = mujoco.mjtCamera.mjCAMERA_FREE
    camera.distance = distance
    camera.azimuth = azimuth
    camera.elevation = elevation
    return camera


def update_cameras(
    controller: MultiDuckPolicy,
    wide_camera: mujoco.MjvCamera,
    close_camera: mujoco.MjvCamera,
) -> None:
    positions = controller.data.xpos[controller.body_ids]
    center = positions.mean(axis=0)
    wide_camera.lookat[:] = (float(center[0]), float(center[1]), 0.23)

    focus_duck = (controller.count - 1) // 2
    focus_body = int(controller.body_ids[focus_duck])
    focus = controller.data.xpos[focus_body]
    close_camera.lookat[:] = (float(focus[0]), float(focus[1]), float(focus[2]) + 0.015)
    yaw = yaw_from_matrix(controller.data.xmat[focus_body])
    # MuJoCo's free camera sits on the ray selected by azimuth and looks toward
    # the subject. Offset by 180° so the lens is in front of the duck rather
    # than following the back of its head.
    close_camera.azimuth = math.degrees(yaw) + 180.0


def mux_audio(silent_video: Path, audio: Path, output: Path, offset: float, duration: float) -> None:
    import imageio_ffmpeg

    command = [
        imageio_ffmpeg.get_ffmpeg_exe(),
        "-y",
        "-ss",
        str(offset),
        "-i",
        str(audio),
        "-i",
        str(silent_video),
        "-t",
        str(duration),
        "-map",
        "1:v:0",
        "-map",
        "0:a:0",
        "-c:v",
        "copy",
        "-c:a",
        "aac",
        "-b:a",
        "192k",
        "-shortest",
        str(output),
    ]
    subprocess.run(command, check=True)


def run(args: argparse.Namespace) -> dict:
    scene = args.scene.resolve()
    if not scene.exists():
        build_scene(scene, args.grid, args.spacing)
    model = mujoco.MjModel.from_xml_path(str(scene))
    model.opt.timestep = 1.0 / PHYSICS_HZ
    data = mujoco.MjData(model)
    fixed_command = tuple(args.fixed_command) if args.fixed_command is not None else None
    controller = MultiDuckPolicy(
        model,
        data,
        args.policy.resolve(),
        args.grid,
        fixed_command=fixed_command,
        kick_left_path=args.kick_left_policy.resolve(),
        kick_right_path=args.kick_right_policy.resolve(),
        synchronized_actions=args.synchronized_actions,
        literal=args.literal,
        simple=args.simple,
    )
    controller.reset()

    for _ in range(int(round(args.warmup * CONTROL_HZ))):
        controller.infer_and_apply(None)
        for _ in range(DECIMATION):
            mujoco.mj_step(model, data)

    initial_xy = np.array(
        [
            data.qpos[int(address) : int(address) + 2].copy()
            for address in controller.free_qpos
        ]
    )
    diagnostics = ContactDiagnostics(controller, initial_xy)

    writer = None
    renderers: list[mujoco.Renderer] = []
    wide_distance = 2.2 + args.spacing * args.grid
    wide_camera = make_camera(wide_distance, 135.0, -33.0)
    # A genuinely close front-follow keeps neighboring ducks from crossing
    # between the lens and the magenta center subject as the formation turns.
    close_camera = make_camera(0.82, 0.0, -14.0)
    next_frame_time = 0.0
    frame_dt = 1.0 / args.fps

    output = args.render.resolve() if args.render else None
    silent_output = output
    if output and args.audio:
        silent_output = output.with_name(output.stem + ".silent.mp4")
    if silent_output:
        import imageio.v2 as imageio

        silent_output.parent.mkdir(parents=True, exist_ok=True)
        pane_width = args.width // 2 if args.view == "split" else args.width
        camera_count = 2 if args.view == "split" else 1
        renderers = [
            mujoco.Renderer(model, width=pane_width, height=args.height)
            for _ in range(camera_count)
        ]
        writer = imageio.get_writer(
            silent_output,
            fps=args.fps,
            codec="libx264",
            quality=8,
            macro_block_size=1,
        )

    control_steps = int(math.ceil(args.duration * CONTROL_HZ))
    try:
        for control_index in range(control_steps):
            show_time = control_index / CONTROL_HZ
            controller.infer_and_apply(show_time)
            for _ in range(DECIMATION):
                mujoco.mj_step(model, data)
            diagnostics.update(show_time)

            if writer is not None:
                simulated_until = (control_index + 1) / CONTROL_HZ
                while next_frame_time < simulated_until and next_frame_time < args.duration:
                    update_cameras(controller, wide_camera, close_camera)
                    if args.view == "split":
                        renderers[0].update_scene(data, camera=close_camera)
                        renderers[1].update_scene(data, camera=wide_camera)
                        frame = np.concatenate(
                            (renderers[0].render(), renderers[1].render()),
                            axis=1,
                        )
                    elif args.view == "close":
                        renderers[0].update_scene(data, camera=close_camera)
                        frame = renderers[0].render()
                    else:
                        renderers[0].update_scene(data, camera=wide_camera)
                        frame = renderers[0].render()
                    writer.append_data(frame)
                    next_frame_time += frame_dt
    finally:
        if writer is not None:
            writer.close()
        for renderer in renderers:
            renderer.close()

    report = diagnostics.report(args.duration)
    if args.diagnostics:
        args.diagnostics.resolve().write_text(
            json.dumps(report, indent=2),
            encoding="utf-8",
        )
    if output and args.audio:
        mux_audio(silent_output, args.audio.resolve(), output, args.audio_offset, args.duration)
        silent_output.unlink(missing_ok=True)
    return report


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--scene", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--policy", type=Path, default=DEFAULT_POLICY)
    parser.add_argument("--kick-left-policy", type=Path, default=DEFAULT_KICK_LEFT_POLICY)
    parser.add_argument("--kick-right-policy", type=Path, default=DEFAULT_KICK_RIGHT_POLICY)
    parser.add_argument("--grid", type=int, default=5)
    parser.add_argument("--spacing", type=float, default=1.5)
    parser.add_argument("--duration", type=float)
    parser.add_argument("--warmup", type=float, default=1.5)
    parser.add_argument("--fps", type=int, default=30)
    parser.add_argument("--width", type=int, default=1600)
    parser.add_argument("--height", type=int, default=720)
    parser.add_argument("--view", choices=("split", "wide", "close"), default="split")
    parser.add_argument("--render", type=Path)
    parser.add_argument("--audio", type=Path)
    parser.add_argument("--audio-offset", type=float, default=0.0)
    parser.add_argument("--diagnostics", type=Path)
    parser.add_argument(
        "--fixed-command",
        type=float,
        nargs=3,
        metavar=("VX", "VY", "WZ"),
        help="Diagnostic override for the walking-policy twist command.",
    )
    parser.add_argument("--strict", action="store_true")
    dance_group = parser.add_mutually_exclusive_group()
    dance_group.add_argument(
        "--literal",
        action="store_true",
        help="Use the canonical literal 18-count Electric Slide retargeting.",
    )
    dance_group.add_argument(
        "--simple",
        action="store_true",
        help="Use the morphology-specific ultra-simple 16-count Electric Slide.",
    )
    parser.add_argument(
        "--synchronized-actions",
        action="store_true",
        help="Diagnostic mode: use one open-loop action for the entire flock.",
    )
    args = parser.parse_args()

    if args.duration is None:
        args.duration = SIMPLE_SHOW_SECONDS if args.simple else SHOW_SECONDS

    for policy_path in (args.policy, args.kick_left_policy, args.kick_right_policy):
        if not policy_path.exists():
            parser.error(f"Official policy not found: {policy_path}")
    if args.render and args.diagnostics is None:
        args.diagnostics = args.render.with_suffix(".physics.json")
    report = run(args)
    print(json.dumps(report, indent=2))
    if args.render:
        print(args.render.resolve())
    if args.strict and (
        report["fallen_ducks"]
        or report["minimum_foot_liftoffs"] < 2
        or report["maximum_airborne_fraction"] > 0.20
        or report["minimum_maximum_excursion_m"] < 0.03
        or report["minimum_pairwise_separation_m"] < 0.20
        or report["inter_duck_contact_count"] > 0
    ):
        raise SystemExit("Physics validation failed")


if __name__ == "__main__":
    main()