"""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()