| """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 |
| from generate_scene import build_scene |
| from physics_show import ( |
| 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() |
|
|
