"""Record a physically executed single-Microduck dance reference at 50 Hz.""" from __future__ import annotations import argparse import json from pathlib import Path import sys import mujoco import numpy as np DANCE_DIR = Path(__file__).resolve().parents[1] sys.path.insert(0, str(DANCE_DIR)) from choreography import SHOW_SECONDS # noqa: E402 from generate_scene import build_scene # noqa: E402 from physics_show import ( # noqa: E402 CONTROL_HZ, DECIMATION, DEFAULT_KICK_LEFT_POLICY, DEFAULT_KICK_RIGHT_POLICY, DEFAULT_POLICY, PHYSICS_HZ, ContactDiagnostics, MultiDuckPolicy, ) def run(args: argparse.Namespace) -> dict: output = args.output.resolve() output.parent.mkdir(parents=True, exist_ok=True) scene = output.parent / "single_duck_reference_scene.xml" build_scene(scene, grid=1, spacing=1.5) model = mujoco.MjModel.from_xml_path(str(scene)) model.opt.timestep = 1.0 / PHYSICS_HZ data = mujoco.MjData(model) controller = MultiDuckPolicy( model, data, args.policy.resolve(), grid=1, kick_left_path=args.kick_left_policy.resolve(), kick_right_path=args.kick_right_policy.resolve(), ) 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 = data.qpos[controller.free_qpos[0]:controller.free_qpos[0] + 2].copy() diagnostics = ContactDiagnostics(controller, initial_xy[None, :]) steps = int(round(args.duration * CONTROL_HZ)) joint_pos = np.empty((steps, 14), dtype=np.float32) joint_vel = np.empty((steps, 14), dtype=np.float32) action = np.empty((steps, 14), dtype=np.float32) command = np.empty((steps, 13), dtype=np.float32) root_pos = np.empty((steps, 3), dtype=np.float32) root_quat = np.empty((steps, 4), dtype=np.float32) root_lin_vel = np.empty((steps, 3), dtype=np.float32) root_ang_vel = np.empty((steps, 3), dtype=np.float32) foot_pos = np.empty((steps, 2, 3), dtype=np.float32) foot_contact = np.zeros((steps, 2), dtype=bool) policy_code = np.empty(steps, dtype=np.int8) teacher_observation = np.empty((steps, 61), dtype=np.float32) phase_observation = np.empty((steps, 61), dtype=np.float32) floor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "dance_floor") foot_ids = controller.foot_geom_ids[0] foot_sites = np.array([ mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "duck_00_left_foot"), mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "duck_00_right_foot"), ]) free_qpos = int(controller.free_qpos[0]) free_qvel = int(controller.free_qvel[0]) for index in range(steps): show_time = index / CONTROL_HZ policy_name = __import__("physics_choreography").policy_for_time(show_time) teacher_command = controller.command_vector(show_time, 0, policy_name) teacher_observation[index] = controller.observation(0, teacher_command) phase = index / steps phase_command = np.zeros(13, dtype=np.float32) phase_command[0] = np.cos(2.0 * np.pi * phase) phase_command[1] = np.sin(2.0 * np.pi * phase) phase_observation[index] = controller.observation(0, phase_command) controller.infer_and_apply(show_time) for _ in range(DECIMATION): mujoco.mj_step(model, data) diagnostics.update() joint_pos[index] = data.qpos[controller.joint_qpos[0]] joint_vel[index] = data.qvel[controller.joint_qvel[0]] action[index] = controller.last_actions[0] command[index] = controller.command_vector(show_time, 0, policy_name) root_pos[index] = data.qpos[free_qpos:free_qpos + 3] root_quat[index] = data.qpos[free_qpos + 3:free_qpos + 7] root_lin_vel[index] = data.qvel[free_qvel:free_qvel + 3] root_ang_vel[index] = data.qvel[free_qvel + 3:free_qvel + 6] foot_pos[index] = data.site_xpos[foot_sites] policy_code[index] = {"walk": 0, "kick_left": 1, "kick_right": 2}[policy_name] for contact_index in range(data.ncon): contact = data.contact[contact_index] pair = {int(contact.geom1), int(contact.geom2)} for foot in range(2): if floor_id in pair and int(foot_ids[foot]) in pair: foot_contact[index, foot] = True report = diagnostics.report(args.duration) np.savez_compressed( output, time_seconds=np.arange(steps, dtype=np.float32) / CONTROL_HZ, phase=np.arange(steps, dtype=np.float32) / steps, joint_pos=joint_pos, joint_vel=joint_vel, action=action, command_13d=command, root_pos_w=root_pos, root_quat_w=root_quat, root_lin_vel_w=root_lin_vel, root_ang_vel_w=root_ang_vel, foot_pos_w=foot_pos, foot_contact=foot_contact, policy_code=policy_code, teacher_observation_61d=teacher_observation, phase_observation_61d=phase_observation, policy_names=np.array(("walk", "kick_left", "kick_right")), default_joint_pos=controller.default_pose, action_scale=np.asarray(controller.action_scale), joint_names=np.array(( "left_hip_yaw", "left_hip_roll", "left_hip_pitch", "left_knee", "left_ankle", "neck_pitch", "head_pitch", "head_yaw", "head_roll", "right_hip_yaw", "right_hip_roll", "right_hip_pitch", "right_knee", "right_ankle", )), control_hz=np.asarray(CONTROL_HZ), physics_hz=np.asarray(PHYSICS_HZ), ) with output.with_suffix(".jsonl").open("w", encoding="utf-8") as stream: for index in range(steps): stream.write(json.dumps({ "control_index": index, "time_seconds": round(index / CONTROL_HZ, 6), "phase": round(index / steps, 8), "joint_pos_rad": [round(float(value), 7) for value in joint_pos[index]], "joint_vel_rad_s": [round(float(value), 7) for value in joint_vel[index]], "root_pos_w_m": [round(float(value), 7) for value in root_pos[index]], "root_quat_wxyz": [round(float(value), 7) for value in root_quat[index]], "foot_contact": [bool(value) for value in foot_contact[index]], "policy": ("walk", "kick_left", "kick_right")[policy_code[index]], }, separators=(",", ":")) + "\n") summary = { "schema_version": "1.0", "duration_seconds": args.duration, "control_hz": CONTROL_HZ, "physics_hz": PHYSICS_HZ, "frames": steps, "root_pose_writes_after_reset": 0, "diagnostics": report, "files": {"npz": output.name, "jsonl": output.with_suffix(".jsonl").name}, } output.with_name("physics_reference_summary.json").write_text( json.dumps(summary, indent=2), encoding="utf-8", ) print(json.dumps(summary, indent=2)) if report["fallen_ducks"] or report["minimum_foot_liftoffs"] < 2: raise SystemExit("Reference physics validation failed") return summary def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--duration", type=float, default=SHOW_SECONDS) parser.add_argument("--warmup", type=float, default=1.5) 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) run(parser.parse_args()) if __name__ == "__main__": main()