Initial commit

scripts/export_onnx.py

@@ -0,0 +1,162 @@
+"""Export an S10 SuperMimic actor checkpoint to ONNX + deployment JSON.
+
+Example:
+  python scripts/export_onnx.py \
+    --checkpoint logs/s10_terrain_v4_20k_full/model_20000.pt \
+    --output logs/s10_terrain_v4_20k_full/model_20000.onnx
+
+The ONNX file contains only the deterministic actor (RunningNorm + MLP + tanh).
+Input ``obs`` is [batch, 265] for v4/v4-upward and for the new v7 (wheel speed
+restored); only legacy v7 checkpoints use [batch, 245]. Output is [batch, 16].
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from mjlab.utils.string import resolve_expr
+
+from s10_supermimic.core.policy import load_actor
+from s10_supermimic.robot.config import (
+  LEG_JOINT_NAMES, WHEEL_JOINT_NAMES, get_s10_robot_cfg,
+)
+
+
+def task_matches_actor(task, actor_wheel_vel: bool, frame_dim: int) -> bool:
+  return task.actor.wheel_vel == actor_wheel_vel and task.actor.frame_dim == frame_dim
+
+
+def default_joint_pos(checkpoint, actor_wheel_vel: bool, frame_dim: int) -> list[float]:
+  """Use the checkpoint task's default pose when its Actor contract matches.
+
+  Legacy 49-value v7 checkpoints predate the ±5° safe stance, so they fall
+  back to the legacy v4 ±20° pose instead of silently exporting a new target.
+  """
+  robot_cfg = get_s10_robot_cfg()
+  task_name = checkpoint.get("config", {}).get("task")
+  if task_name:
+    from s10_supermimic.tasks import get_task
+    try:
+      task = get_task(task_name)
+      if task_matches_actor(task, actor_wheel_vel, frame_dim):
+        cfg = task.build_env_cfg(
+          play=True, history=int(checkpoint["history"]),
+          seed=int(checkpoint.get("config", {}).get("seed", 0)))
+        robot_cfg = cfg.scene.entities["robot"]
+    except KeyError:
+      pass
+  names = LEG_JOINT_NAMES + WHEEL_JOINT_NAMES
+  return [float(v) for v in resolve_expr(robot_cfg.init_state.joint_pos, names, 0.0)]
+
+
+def main() -> None:
+  p = argparse.ArgumentParser(description=__doc__)
+  p.add_argument("--checkpoint", type=Path, required=True)
+  p.add_argument("--output", type=Path, default=None,
+                 help="ONNX path; defaults to <checkpoint>.onnx")
+  p.add_argument("--opset", type=int, default=17)
+  p.add_argument("--device", default="cpu")
+  args = p.parse_args()
+
+  ckpt = torch.load(args.checkpoint, map_location=args.device, weights_only=True)
+  actor = load_actor(ckpt, args.device).eval()
+  history = int(ckpt["history"])
+  if ckpt.get("format") not in ("s10-flat-policy-v1", "s10-terrain-policy-v1"):
+    raise ValueError(f"Not an S10 policy checkpoint: {ckpt.get('format')}")
+  if history < 1:
+    raise ValueError("checkpoint history must be positive")
+
+  output = args.output or args.checkpoint.with_suffix(".onnx")
+  output.parent.mkdir(parents=True, exist_ok=True)
+  dummy = torch.zeros(1, actor.frame_dim * history, device=args.device)
+
+  with torch.no_grad():
+    reference = actor(dummy).cpu().numpy()
+
+  torch.onnx.export(
+    actor,
+    (dummy,),
+    str(output),
+    input_names=["obs"],
+    output_names=["actions"],
+    opset_version=args.opset,
+    dynamo=False,
+    do_constant_folding=True,
+    dynamic_axes={"obs": {0: "batch"}, "actions": {0: "batch"}},
+  )
+
+  import onnx
+  import onnxruntime as ort
+  onnx.checker.check_model(str(output))
+  session = ort.InferenceSession(str(output), providers=["CPUExecutionProvider"])
+  onnx_out = session.run(["actions"], {"obs": dummy.cpu().numpy()})[0]
+  max_error = float(np.max(np.abs(onnx_out - reference)))
+  if max_error > 1e-5:
+    raise RuntimeError(f"ONNX/PyTorch mismatch: max error {max_error}")
+
+  revision = int(ckpt.get("config", {}).get("terrain_revision", 1))
+  leg_scales = ([0.25, 0.6, 0.8] * 4) if revision >= 2 else ([0.25] * 12)
+  wheel_scale = 35.0 if revision >= 3 else 15.0
+  task_name = ckpt.get("config", {}).get("task")
+  task = None
+  if task_name:
+    from s10_supermimic.tasks import get_task
+    try:
+      task = get_task(task_name)
+    except KeyError:
+      task = None
+  if task is not None and task_matches_actor(task, actor.actor_wheel_vel, actor.frame_dim):
+    actor_schema = task.actor.name
+  else:
+    actor_schema = "s10-v4-53" if actor.actor_wheel_vel else "s10-v7-49"
+  sidecar = {
+    "format": ckpt["format"],
+    "step": ckpt.get("step"),
+    "history": history,
+    "actor_schema": actor_schema,
+    "actor_wheel_vel": actor.actor_wheel_vel,
+    "frame_dim": actor.frame_dim,
+    "observation_dim": actor.frame_dim * history,
+    "action_dim": 16,
+    "dt": ckpt.get("dt", 0.02),
+    "actor_hidden": list(ckpt.get("actor_hidden", ())),
+    "actor_params": sum(x.numel() for x in actor.parameters()),
+    "frame_terms": [
+      {"name": "base_ang_vel", "dim": 3, "scale": 1.0},
+      {"name": "projected_gravity", "dim": 3, "scale": 1.0},
+      {"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
+      {"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
+      *([{"name": "wheel_vel_rel", "dim": 4, "scale": 0.05}]
+        if actor.actor_wheel_vel else []),
+      {"name": "last_action", "dim": 16, "scale": 1.0},
+      {"name": "command", "dim": 3, "scale": 1.0},
+    ],
+    "history_layout": "term-wise blocks, oldest to newest (mjlab flattening)",
+    "env_action_order": list(ckpt["env_action_order"]),
+    "leg_scale": leg_scales,
+    "wheel_scale": wheel_scale,
+    "wheel_clip": 35.0,
+    "default_joint_pos": default_joint_pos(
+      ckpt, actor.actor_wheel_vel, actor.frame_dim),
+    "leg_joint_names": list(LEG_JOINT_NAMES),
+    "wheel_joint_names": list(WHEEL_JOINT_NAMES),
+    "onnx": {
+      "input": "obs",
+      "output": "actions",
+      "opset": args.opset,
+      "max_abs_error_vs_torch": max_error,
+    },
+  }
+  sidecar_path = output.with_suffix(".json")
+  sidecar_path.write_text(json.dumps(sidecar, indent=2) + "\n")
+  print(f"Exported ONNX: {output}")
+  print(f"Deployment metadata: {sidecar_path}")
+  print(f"ONNX/PyTorch max error: {max_error:.3e}")
+
+
+if __name__ == "__main__":
+  main()