feat(training): export SB3 PPO policy to ONNX with parity test

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()