| """Closed-loop MuJoCo evaluation for the trained phase-imitation policy.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| from pathlib import Path |
| import sys |
|
|
| import mujoco |
| import numpy as np |
| import onnxruntime as ort |
|
|
|
|
| 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_POLICY, |
| PHYSICS_HZ, |
| ContactDiagnostics, |
| MultiDuckPolicy, |
| ) |
|
|
|
|
| def phase_command(time_seconds: float, duration: float) -> np.ndarray: |
| phase = (time_seconds % duration) / duration |
| result = np.zeros(13, dtype=np.float32) |
| result[0] = math.cos(2.0 * math.pi * phase) |
| result[1] = math.sin(2.0 * math.pi * phase) |
| return result |
|
|
|
|
| def run(args: argparse.Namespace) -> dict: |
| output = args.output.resolve() |
| output.parent.mkdir(parents=True, exist_ok=True) |
| scene = output.parent / "phase_policy_eval_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) |
| indexing = MultiDuckPolicy(model, data, args.teacher_policy.resolve(), grid=1) |
| indexing.reset() |
|
|
| session = ort.InferenceSession(str(args.policy.resolve())) |
| input_name = session.get_inputs()[0].name |
| output_name = session.get_outputs()[0].name |
| metadata = session.get_modelmeta().custom_metadata_map |
| action_scale = float(metadata["action_scale"]) |
| default_pose = np.array( |
| [float(value) for value in metadata["default_joint_pos"].split(",")], |
| dtype=np.float32, |
| ) |
| if session.get_inputs()[0].shape[-1] != 61 or session.get_outputs()[0].shape[-1] != 14: |
| raise ValueError("Phase policy must use the 61D -> 14D deployment contract") |
| indexing.default_pose = default_pose |
| indexing.data.ctrl[:] = default_pose |
| indexing.last_actions.fill(0.0) |
| mujoco.mj_forward(model, data) |
|
|
| for _ in range(int(round(args.warmup * CONTROL_HZ))): |
| command = phase_command(0.0, args.duration) |
| observation = indexing.observation(0, command) |
| action = session.run([output_name], {input_name: observation[None, :]})[0][0] |
| indexing.last_actions[0] = action |
| data.ctrl[indexing.actuator_ids[0]] = default_pose + action * action_scale |
| for _ in range(DECIMATION): |
| mujoco.mj_step(model, data) |
|
|
| initial_xy = data.qpos[indexing.free_qpos[0]:indexing.free_qpos[0] + 2].copy() |
| diagnostics = ContactDiagnostics(indexing, initial_xy[None, :]) |
| reference = np.load(args.reference.resolve(), allow_pickle=True) |
| steps = int(round(args.duration * CONTROL_HZ)) |
| rollout_joint = np.empty((steps, 14), dtype=np.float32) |
| rollout_action = np.empty((steps, 14), dtype=np.float32) |
| observations = np.empty((steps, 61), dtype=np.float32) |
|
|
| for index in range(steps): |
| command = phase_command(index / CONTROL_HZ, args.duration) |
| observation = indexing.observation(0, command) |
| action = session.run([output_name], {input_name: observation[None, :]})[0][0] |
| indexing.last_actions[0] = action |
| data.ctrl[indexing.actuator_ids[0]] = default_pose + action * action_scale |
| for _ in range(DECIMATION): |
| mujoco.mj_step(model, data) |
| diagnostics.update() |
| observations[index] = observation |
| rollout_action[index] = action |
| rollout_joint[index] = data.qpos[indexing.joint_qpos[0]] |
|
|
| reference_joint = reference["joint_pos"][:steps] |
| joint_rmse = np.sqrt(np.mean((rollout_joint - reference_joint) ** 2, axis=0)) |
| report = diagnostics.report(args.duration) |
| report["closed_loop_joint_rmse_rad"] = float(np.sqrt(np.mean((rollout_joint - reference_joint) ** 2))) |
| report["closed_loop_joint_rmse_per_joint_rad"] = joint_rmse.round(7).tolist() |
| report["policy_path"] = str(args.policy) |
| report["reference_path"] = str(args.reference) |
| report["passed"] = ( |
| not report["fallen_ducks"] |
| and report["minimum_foot_liftoffs"] >= 2 |
| and report["maximum_airborne_fraction"] <= 0.20 |
| and np.isfinite(rollout_joint).all() |
| ) |
| np.savez_compressed( |
| output, |
| time_seconds=np.arange(steps, dtype=np.float32) / CONTROL_HZ, |
| joint_pos=rollout_joint, |
| action=rollout_action, |
| observation_61d=observations, |
| reference_joint_pos=reference_joint, |
| ) |
| output.with_suffix(".json").write_text(json.dumps(report, indent=2), encoding="utf-8") |
| print(json.dumps(report, indent=2)) |
| if not report["passed"]: |
| raise SystemExit("Closed-loop phase-policy validation failed") |
| return report |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--policy", type=Path, required=True) |
| parser.add_argument("--reference", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--teacher-policy", type=Path, default=DEFAULT_POLICY) |
| parser.add_argument("--duration", type=float, default=SHOW_SECONDS) |
| parser.add_argument("--warmup", type=float, default=1.5) |
| run(parser.parse_args()) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
