69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
"""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"
|