91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
#!/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()
|