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