initial commit

scripts/export_fire_ant_policy.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+"""Export the RSL-RL FireAnt actor from a training checkpoint to ONNX."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import onnx
+import onnxruntime as ort
+import torch
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_CHECKPOINT = (
+    PROJECT_ROOT.parent
+    / "rl_training"
+    / "logs"
+    / "rsl_rl"
+    / "deeprobotics_m20_support"
+    / "2026-09-12_17-52-15_stage2_multiterrain"
+    / "model_39899.pt"
+)
+DEFAULT_OUTPUT = PROJECT_ROOT / "policy" / "fire_ant_policy.onnx"
+
+
+def build_actor() -> torch.nn.Module:
+    return torch.nn.Sequential(
+        torch.nn.Linear(63, 512),
+        torch.nn.ELU(),
+        torch.nn.Linear(512, 256),
+        torch.nn.ELU(),
+        torch.nn.Linear(256, 128),
+        torch.nn.ELU(),
+        torch.nn.Linear(128, 8),
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
+    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if not args.checkpoint.is_file():
+        raise FileNotFoundError(args.checkpoint)
+
+    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+    state_dict = checkpoint["model_state_dict"]
+    actor_state = {
+        name.removeprefix("actor."): tensor
+        for name, tensor in state_dict.items()
+        if name.startswith("actor.")
+    }
+    actor = build_actor().eval()
+    actor.load_state_dict(actor_state, strict=True)
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    sample = torch.zeros(1, 63, dtype=torch.float32)
+    with torch.no_grad():
+        torch_output = actor(sample).cpu().numpy()
+    torch.onnx.export(
+        actor,
+        sample,
+        args.output,
+        input_names=["obs"],
+        output_names=["actions"],
+        dynamic_axes={"obs": {0: "batch"}, "actions": {0: "batch"}},
+        opset_version=17,
+        do_constant_folding=True,
+        dynamo=False,
+    )
+
+    model = onnx.load(args.output)
+    onnx.checker.check_model(model)
+    session = ort.InferenceSession(str(args.output), providers=["CPUExecutionProvider"])
+    ort_output = session.run(["actions"], {"obs": sample.numpy()})[0]
+    max_error = float(np.max(np.abs(torch_output - ort_output)))
+    if max_error > 1e-5:
+        raise RuntimeError(f"ONNX validation failed: max error {max_error:.3e}")
+    print(f"Exported {args.checkpoint}")
+    print(f"ONNX: {args.output} (input 63, output 8, max error {max_error:.3e})")
+
+
+if __name__ == "__main__":
+    main()

scripts/pt_to_onnx.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Compatibility entry point; use export_fire_ant_policy.py for new deployments."""
+
+from export_fire_ant_policy import main
+
+
+if __name__ == "__main__":
+    main()