File size: 5,347 Bytes
32502a8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """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 # noqa: E402
from generate_scene import build_scene # noqa: E402
from physics_show import ( # noqa: E402
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()
|