68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
"""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()
|