From bbd2c3609165dc968e7332fcd656545e1b25872b Mon Sep 17 00:00:00 2001
From: meelstorm <meelstorm@gmail.com>
Date: Fri, 17 Jul 2026 17:17:29 +0000
Subject: [PATCH] feat(training): export SB3 PPO policy to ONNX with parity
 test

---
 Training/export_onnx.py            | 67 +++++++++++++++++++++++++++++
 Training/tests/test_export_onnx.py | 68 ++++++++++++++++++++++++++++++
 2 files changed, 135 insertions(+)
 create mode 100644 Training/export_onnx.py
 create mode 100644 Training/tests/test_export_onnx.py

diff --git a/Training/export_onnx.py b/Training/export_onnx.py
new file mode 100644
index 0000000..330bf0a
--- /dev/null
+++ b/Training/export_onnx.py
@@ -0,0 +1,67 @@
+"""Export a trained SB3 PPO policy to ONNX for the .NET backend.
+
+The exported model has a single float input of shape [batch, 7] (the observation
+vector) and a single float output of shape [batch, 4] (action logits). Argmax
+over the last axis gives the action.
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+from stable_baselines3 import PPO
+
+
+class OnnxablePolicy(torch.nn.Module):
+    """Wraps the SB3 policy's actor path for ONNX export.
+
+    Standard SB3 MlpPolicy for a discrete action space:
+        features = policy.extract_features(obs)
+        latent_pi = policy.mlp_extractor.forward_actor(features)
+        logits = policy.action_net(latent_pi)
+    """
+
+    def __init__(self, policy: torch.nn.Module) -> None:
+        super().__init__()
+        self.policy = policy
+
+    def forward(self, obs: torch.Tensor) -> torch.Tensor:
+        features = self.policy.extract_features(obs)
+        latent_pi = self.policy.mlp_extractor.forward_actor(features)
+        return self.policy.action_net(latent_pi)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", type=Path, required=True,
+                        help="path to SB3 .zip checkpoint")
+    parser.add_argument("--out", type=Path, required=True,
+                        help="destination .onnx path")
+    parser.add_argument("--opset", type=int, default=17)
+    args = parser.parse_args()
+
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+
+    model = PPO.load(str(args.checkpoint), device="cpu")
+    model.policy.eval()
+
+    onnxable = OnnxablePolicy(model.policy)
+    onnxable.eval()
+
+    dummy = torch.randn(1, 7, dtype=torch.float32)
+    torch.onnx.export(
+        onnxable,
+        dummy,
+        str(args.out),
+        input_names=["obs"],
+        output_names=["logits"],
+        dynamic_axes={"obs": {0: "batch"}, "logits": {0: "batch"}},
+        opset_version=args.opset,
+        dynamo=False,
+    )
+    print(f"exported ONNX policy to {args.out}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/Training/tests/test_export_onnx.py b/Training/tests/test_export_onnx.py
new file mode 100644
index 0000000..1dd2146
--- /dev/null
+++ b/Training/tests/test_export_onnx.py
@@ -0,0 +1,68 @@
+"""Verify export_onnx.py produces an ONNX file that matches the SB3 policy."""
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+CHECKPOINT = REPO_ROOT / "checkpoints" / "ppo_lander_final.zip"
+ONNX_OUT = REPO_ROOT / "models" / "ppo_lander.onnx"
+
+
+@pytest.fixture(scope="module")
+def exported_onnx():
+    if not CHECKPOINT.exists():
+        pytest.skip(f"no checkpoint at {CHECKPOINT} — run train.py first")
+    # export
+    import subprocess
+    import sys
+    training_dir = REPO_ROOT / "Training"
+    result = subprocess.run(
+        [sys.executable, str(training_dir / "export_onnx.py"),
+         "--checkpoint", str(CHECKPOINT), "--out", str(ONNX_OUT)],
+        capture_output=True, text=True, cwd=str(REPO_ROOT),
+    )
+    assert result.returncode == 0, f"export failed: {result.stderr}"
+    assert ONNX_OUT.exists(), f"ONNX file not produced at {ONNX_OUT}"
+    return ONNX_OUT
+
+
+def test_onnx_file_exists(exported_onnx):
+    assert exported_onnx.stat().st_size > 0
+
+
+def test_onnx_input_output_shapes(exported_onnx):
+    import onnx
+    model = onnx.load(str(exported_onnx))
+    assert len(model.graph.input) == 1
+    in_shape = [d.dim_value for d in model.graph.input[0].type.tensor_type.shape.dim]
+    # (batch, 7)
+    assert in_shape[-1] == 7
+    assert len(model.graph.output) >= 1
+    out_shape = [d.dim_value for d in model.graph.output[0].type.tensor_type.shape.dim]
+    assert out_shape[-1] == 4  # 4 discrete actions
+
+
+def test_onnx_output_matches_sb3(exported_onnx):
+    """Sanity: ONNX inference for 100 random obs must argmax to the same action as SB3."""
+    import onnxruntime
+    from stable_baselines3 import PPO
+
+    model = PPO.load(str(CHECKPOINT))
+    session = onnxruntime.InferenceSession(str(exported_onnx))
+    input_name = session.get_inputs()[0].name
+
+    rng = np.random.default_rng(1234)
+    obs_batch = rng.normal(size=(100, 7)).astype(np.float32)
+
+    # ONNX: (100, 4) logits
+    onnx_logits = session.run(None, {input_name: obs_batch})[0]
+    onnx_actions = onnx_logits.argmax(axis=-1)
+
+    # SB3 deterministic prediction
+    sb3_actions, _ = model.predict(obs_batch, deterministic=True)
+
+    mismatches = int((onnx_actions != sb3_actions).sum())
+    # Allow up to 2 mismatches out of 100 for float rounding on marginal states.
+    assert mismatches <= 2, f"{mismatches}/100 argmax mismatches between ONNX and SB3"