Add Plan 2: Python PPO training pipeline plan
This commit is contained in:
789
docs/superpowers/plans/2026-07-17-02-training-pipeline.md
Normal file
789
docs/superpowers/plans/2026-07-17-02-training-pipeline.md
Normal file
@@ -0,0 +1,789 @@
|
|||||||
|
# Training Pipeline Implementation Plan
|
||||||
|
|
||||||
|
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||||
|
|
||||||
|
**Goal:** Build the Python side of the design: a `LanderCliEnv` Gymnasium environment that wraps the `GameCli` binary as a subprocess, a `train.py` that trains a PPO policy on N vectorized copies of it, and an `export_onnx.py` that exports the trained network as `models/ppo_lander.onnx` for the ASP.NET backend to load. Reference spec at `docs/superpowers/specs/2026-07-17-cursor-following-lander-design.md`.
|
||||||
|
|
||||||
|
**Architecture:** All Python code lives under `Training/`. `LanderCliEnv` speaks the exact protocol the `GameCli` documents on stdio, so training uses literally the same physics as inference. Training runs a `SubprocVecEnv` of N envs → PPO batches transitions → checkpoints to `checkpoints/` → export → verify → `models/ppo_lander.onnx`.
|
||||||
|
|
||||||
|
**Tech Stack:** Python 3.11 or 3.12, `gymnasium`, `stable-baselines3`, `torch`, `onnx`, `onnxruntime`, `pytest`.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Prerequisites
|
||||||
|
|
||||||
|
- **Plan 1 complete.** The `GameCli` binary must be publishable at `publish/GameCli/GameCli` (run `dotnet publish GameCli/GameCli.csproj -c Release -o publish/GameCli` if it's not there yet). All Training/ code assumes that path.
|
||||||
|
- Python 3.11 or 3.12 available as `python3`.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## File structure
|
||||||
|
|
||||||
|
```
|
||||||
|
Experiment_ReinforcementLearning/
|
||||||
|
├── Training/
|
||||||
|
│ ├── requirements.txt
|
||||||
|
│ ├── pytest.ini
|
||||||
|
│ ├── lander_cli_env.py # LanderCliEnv(gym.Env)
|
||||||
|
│ ├── train.py # PPO training with SubprocVecEnv
|
||||||
|
│ ├── export_onnx.py # SB3 -> ONNX
|
||||||
|
│ └── tests/
|
||||||
|
│ ├── __init__.py
|
||||||
|
│ ├── test_lander_cli_env.py
|
||||||
|
│ └── test_export_onnx.py
|
||||||
|
├── models/ # gitignored; export_onnx.py writes ppo_lander.onnx here
|
||||||
|
├── checkpoints/ # gitignored; SB3 checkpoints
|
||||||
|
└── tensorboard/ # gitignored; TB logs
|
||||||
|
```
|
||||||
|
|
||||||
|
`Training/lander_cli_env.py` is the single file that touches the subprocess. `train.py` and `export_onnx.py` just consume Gymnasium/SB3 APIs.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Task 1: Scaffold Python project
|
||||||
|
|
||||||
|
**Files:**
|
||||||
|
- Create: `Training/requirements.txt`
|
||||||
|
- Create: `Training/pytest.ini`
|
||||||
|
- Create: `Training/tests/__init__.py`
|
||||||
|
|
||||||
|
- [ ] **Step 1: Check Python is available**
|
||||||
|
|
||||||
|
Run: `python3 --version`
|
||||||
|
Expected: `Python 3.11.x` or `3.12.x`. If older or missing, STOP and report BLOCKED.
|
||||||
|
|
||||||
|
- [ ] **Step 2: Create the virtualenv at repo root**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python3 -m venv .venv
|
||||||
|
source .venv/bin/activate
|
||||||
|
pip install --upgrade pip
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 3: Write `Training/requirements.txt`**
|
||||||
|
|
||||||
|
```
|
||||||
|
gymnasium>=0.29,<2
|
||||||
|
stable-baselines3>=2.3,<3
|
||||||
|
torch>=2.2,<3
|
||||||
|
onnx>=1.16,<2
|
||||||
|
onnxruntime>=1.18,<2
|
||||||
|
numpy>=1.24,<3
|
||||||
|
tensorboard>=2.15
|
||||||
|
pytest>=8
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 4: Install dependencies**
|
||||||
|
|
||||||
|
Run: `pip install -r Training/requirements.txt`
|
||||||
|
Expected: successful install. This may take a few minutes (torch is large).
|
||||||
|
|
||||||
|
- [ ] **Step 5: Write `Training/pytest.ini`**
|
||||||
|
|
||||||
|
```ini
|
||||||
|
[pytest]
|
||||||
|
testpaths = tests
|
||||||
|
python_files = test_*.py
|
||||||
|
addopts = -v
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 6: Write `Training/tests/__init__.py`** (empty file)
|
||||||
|
|
||||||
|
```bash
|
||||||
|
touch Training/tests/__init__.py
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 7: Verify pytest can discover an empty test suite**
|
||||||
|
|
||||||
|
From `Training/` directory:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd Training && pytest && cd ..
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected: `no tests ran` (0 tests, exit code 0 or 5 — both mean "no tests but no failures"). Exit code 5 is fine.
|
||||||
|
|
||||||
|
- [ ] **Step 8: Commit**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
git add Training/
|
||||||
|
git -c user.email=meelstorm@gmail.com -c user.name=meelstorm commit -m "feat(training): scaffold Python training project"
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Task 2: `LanderCliEnv` — Gymnasium environment wrapping the CLI
|
||||||
|
|
||||||
|
**Files:**
|
||||||
|
- Create: `Training/lander_cli_env.py`
|
||||||
|
- Create: `Training/tests/test_lander_cli_env.py`
|
||||||
|
|
||||||
|
### Protocol quick reference (from Plan 1)
|
||||||
|
|
||||||
|
- CLI startup emits: `{"init":{"world":[1,1],"dt":0.02,"obs_dim":7,"n_actions":4}}`
|
||||||
|
- Step input: `{"action":<int>,"target":[<float>,<float>]}`
|
||||||
|
- Reset input: `{"cmd":"reset","seed":<int>}` (seed optional)
|
||||||
|
- Output line: `{"obs":[...7 floats...],"state":{...},"reward":<float>,"done":<bool>,"step":<int>}`
|
||||||
|
- CLI binary path (relative to repo root): `publish/GameCli/GameCli`
|
||||||
|
|
||||||
|
- [ ] **Step 1: Write the failing tests**
|
||||||
|
|
||||||
|
Create `Training/tests/test_lander_cli_env.py`:
|
||||||
|
|
||||||
|
```python
|
||||||
|
"""Unit tests for LanderCliEnv."""
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
from gymnasium.utils.env_checker import check_env
|
||||||
|
|
||||||
|
# Add parent dir to sys.path so we can import lander_cli_env
|
||||||
|
import sys
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||||
|
|
||||||
|
from lander_cli_env import LanderCliEnv
|
||||||
|
|
||||||
|
|
||||||
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||||
|
CLI_BINARY = REPO_ROOT / "publish" / "GameCli" / "GameCli"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="module")
|
||||||
|
def cli_binary_exists():
|
||||||
|
if not CLI_BINARY.exists():
|
||||||
|
pytest.skip(f"CLI binary not found at {CLI_BINARY} — run `dotnet publish GameCli -c Release -o publish/GameCli`")
|
||||||
|
return CLI_BINARY
|
||||||
|
|
||||||
|
|
||||||
|
def test_env_opens_and_closes(cli_binary_exists):
|
||||||
|
env = LanderCliEnv(cli_binary=str(cli_binary_exists))
|
||||||
|
env.close() # should not raise
|
||||||
|
|
||||||
|
|
||||||
|
def test_observation_space_is_seven_dim(cli_binary_exists):
|
||||||
|
env = LanderCliEnv(cli_binary=str(cli_binary_exists))
|
||||||
|
try:
|
||||||
|
assert env.observation_space.shape == (7,)
|
||||||
|
finally:
|
||||||
|
env.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_action_space_is_discrete_four(cli_binary_exists):
|
||||||
|
env = LanderCliEnv(cli_binary=str(cli_binary_exists))
|
||||||
|
try:
|
||||||
|
assert env.action_space.n == 4
|
||||||
|
finally:
|
||||||
|
env.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_reset_returns_seven_dim_observation(cli_binary_exists):
|
||||||
|
env = LanderCliEnv(cli_binary=str(cli_binary_exists))
|
||||||
|
try:
|
||||||
|
obs, info = env.reset(seed=42)
|
||||||
|
assert isinstance(obs, np.ndarray)
|
||||||
|
assert obs.shape == (7,)
|
||||||
|
assert obs.dtype == np.float32
|
||||||
|
finally:
|
||||||
|
env.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_step_advances_and_returns_five_tuple(cli_binary_exists):
|
||||||
|
env = LanderCliEnv(cli_binary=str(cli_binary_exists))
|
||||||
|
try:
|
||||||
|
env.reset(seed=0)
|
||||||
|
obs, reward, terminated, truncated, info = env.step(0) # noop
|
||||||
|
assert obs.shape == (7,)
|
||||||
|
assert isinstance(reward, float)
|
||||||
|
assert terminated is False # never terminates
|
||||||
|
assert isinstance(truncated, bool)
|
||||||
|
finally:
|
||||||
|
env.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_reset_randomizes_target_across_episodes(cli_binary_exists):
|
||||||
|
env = LanderCliEnv(cli_binary=str(cli_binary_exists))
|
||||||
|
try:
|
||||||
|
# Two resets with different seeds should generally give different targets.
|
||||||
|
env.reset(seed=1)
|
||||||
|
target1 = env._target
|
||||||
|
env.reset(seed=2)
|
||||||
|
target2 = env._target
|
||||||
|
assert target1 != target2
|
||||||
|
finally:
|
||||||
|
env.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_truncation_at_max_steps(cli_binary_exists):
|
||||||
|
env = LanderCliEnv(cli_binary=str(cli_binary_exists), max_episode_steps=5)
|
||||||
|
try:
|
||||||
|
env.reset(seed=0)
|
||||||
|
truncated = False
|
||||||
|
for _ in range(5):
|
||||||
|
_, _, _, truncated, _ = env.step(0)
|
||||||
|
assert truncated is True
|
||||||
|
finally:
|
||||||
|
env.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_check_env_passes(cli_binary_exists):
|
||||||
|
"""Gymnasium's own env sanity checker — spaces, dtypes, reset/step contract."""
|
||||||
|
env = LanderCliEnv(cli_binary=str(cli_binary_exists), max_episode_steps=50)
|
||||||
|
try:
|
||||||
|
# skip_render_check because we don't implement rendering
|
||||||
|
check_env(env.unwrapped, skip_render_check=True)
|
||||||
|
finally:
|
||||||
|
env.close()
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 2: Run tests to verify they fail**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd Training && pytest && cd ..
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected: import error — `lander_cli_env` module does not exist.
|
||||||
|
|
||||||
|
- [ ] **Step 3: Write `Training/lander_cli_env.py`**
|
||||||
|
|
||||||
|
```python
|
||||||
|
"""Gymnasium environment that wraps the GameCli C# binary as a subprocess.
|
||||||
|
|
||||||
|
Speaks the line-delimited JSON protocol defined in:
|
||||||
|
docs/superpowers/specs/2026-07-17-cursor-following-lander-design.md
|
||||||
|
|
||||||
|
Each episode picks a random target in [0,1]^2 at reset time and holds it fixed.
|
||||||
|
Never terminates — episodes end via truncation at max_episode_steps.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import subprocess
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Optional
|
||||||
|
|
||||||
|
import gymnasium as gym
|
||||||
|
import numpy as np
|
||||||
|
from gymnasium import spaces
|
||||||
|
|
||||||
|
|
||||||
|
class LanderCliEnv(gym.Env):
|
||||||
|
"""One CLI subprocess per env instance. Not thread-safe; safe under SubprocVecEnv."""
|
||||||
|
|
||||||
|
metadata = {"render_modes": []}
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
cli_binary: str,
|
||||||
|
max_episode_steps: int = 1000,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self._cli_binary = cli_binary
|
||||||
|
self._max_episode_steps = max_episode_steps
|
||||||
|
|
||||||
|
self._proc: Optional[subprocess.Popen] = None
|
||||||
|
self._target: tuple[float, float] = (0.5, 0.5)
|
||||||
|
self._step_count: int = 0
|
||||||
|
|
||||||
|
self._start_process()
|
||||||
|
self._read_init()
|
||||||
|
|
||||||
|
self.observation_space = spaces.Box(
|
||||||
|
low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32
|
||||||
|
)
|
||||||
|
self.action_space = spaces.Discrete(4)
|
||||||
|
|
||||||
|
# -------- subprocess lifecycle --------
|
||||||
|
|
||||||
|
def _start_process(self) -> None:
|
||||||
|
self._proc = subprocess.Popen(
|
||||||
|
[self._cli_binary],
|
||||||
|
stdin=subprocess.PIPE,
|
||||||
|
stdout=subprocess.PIPE,
|
||||||
|
stderr=subprocess.PIPE,
|
||||||
|
text=True,
|
||||||
|
bufsize=1, # line-buffered
|
||||||
|
)
|
||||||
|
|
||||||
|
def _read_init(self) -> None:
|
||||||
|
assert self._proc is not None and self._proc.stdout is not None
|
||||||
|
line = self._proc.stdout.readline()
|
||||||
|
if not line:
|
||||||
|
raise RuntimeError("CLI died before emitting init handshake")
|
||||||
|
msg = json.loads(line)
|
||||||
|
if "init" not in msg:
|
||||||
|
raise RuntimeError(f"expected init handshake, got: {line!r}")
|
||||||
|
|
||||||
|
def _send(self, obj: dict[str, Any]) -> None:
|
||||||
|
assert self._proc is not None and self._proc.stdin is not None
|
||||||
|
self._proc.stdin.write(json.dumps(obj) + "\n")
|
||||||
|
self._proc.stdin.flush()
|
||||||
|
|
||||||
|
def _recv(self) -> dict[str, Any]:
|
||||||
|
assert self._proc is not None and self._proc.stdout is not None
|
||||||
|
line = self._proc.stdout.readline()
|
||||||
|
if not line:
|
||||||
|
raise RuntimeError("CLI closed stdout unexpectedly")
|
||||||
|
return json.loads(line)
|
||||||
|
|
||||||
|
# -------- gym.Env API --------
|
||||||
|
|
||||||
|
def reset(
|
||||||
|
self, *, seed: Optional[int] = None, options: Optional[dict] = None
|
||||||
|
) -> tuple[np.ndarray, dict]:
|
||||||
|
super().reset(seed=seed)
|
||||||
|
# Sample a fresh target each episode; the Gymnasium-provided
|
||||||
|
# np_random is deterministic given `seed`.
|
||||||
|
tx = float(self.np_random.uniform(0.05, 0.95))
|
||||||
|
ty = float(self.np_random.uniform(0.05, 0.95))
|
||||||
|
self._target = (tx, ty)
|
||||||
|
self._step_count = 0
|
||||||
|
|
||||||
|
cmd: dict[str, Any] = {"cmd": "reset"}
|
||||||
|
if seed is not None:
|
||||||
|
cmd["seed"] = int(seed)
|
||||||
|
self._send(cmd)
|
||||||
|
msg = self._recv()
|
||||||
|
obs = np.asarray(msg["obs"], dtype=np.float32)
|
||||||
|
# Overwrite the CLI's dummy (0.5, 0.5) target displacement with the real one.
|
||||||
|
# dx = target.X - ship.X, dy = target.Y - ship.Y.
|
||||||
|
obs[0] = tx - msg["state"]["x"]
|
||||||
|
obs[1] = ty - msg["state"]["y"]
|
||||||
|
return obs, {}
|
||||||
|
|
||||||
|
def step(self, action: int) -> tuple[np.ndarray, float, bool, bool, dict]:
|
||||||
|
self._send({"action": int(action), "target": list(self._target)})
|
||||||
|
msg = self._recv()
|
||||||
|
obs = np.asarray(msg["obs"], dtype=np.float32)
|
||||||
|
reward = float(msg["reward"])
|
||||||
|
self._step_count += 1
|
||||||
|
truncated = self._step_count >= self._max_episode_steps
|
||||||
|
terminated = False # continuous hover task
|
||||||
|
return obs, reward, terminated, truncated, {}
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
if self._proc is None:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
if self._proc.stdin is not None:
|
||||||
|
self._proc.stdin.close()
|
||||||
|
except (BrokenPipeError, OSError):
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
self._proc.wait(timeout=3)
|
||||||
|
except subprocess.TimeoutExpired:
|
||||||
|
self._proc.kill()
|
||||||
|
self._proc.wait(timeout=1)
|
||||||
|
self._proc = None
|
||||||
|
|
||||||
|
def __del__(self) -> None:
|
||||||
|
try:
|
||||||
|
self.close()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 4: Publish the CLI binary (prerequisite for tests)**
|
||||||
|
|
||||||
|
If `publish/GameCli/GameCli` doesn't exist yet, run:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
dotnet publish GameCli/GameCli.csproj -c Release -o publish/GameCli
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 5: Run tests to verify they pass**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd Training && pytest tests/test_lander_cli_env.py && cd ..
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected: 8 passed, 0 failed.
|
||||||
|
|
||||||
|
- [ ] **Step 6: Commit**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
git add Training/lander_cli_env.py Training/tests/test_lander_cli_env.py
|
||||||
|
git -c user.email=meelstorm@gmail.com -c user.name=meelstorm commit -m "feat(training): implement LanderCliEnv wrapping the GameCli subprocess"
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Task 3: `train.py` — PPO training with vectorized envs
|
||||||
|
|
||||||
|
**Files:**
|
||||||
|
- Create: `Training/train.py`
|
||||||
|
|
||||||
|
- [ ] **Step 1: Write `Training/train.py`**
|
||||||
|
|
||||||
|
```python
|
||||||
|
"""Train a PPO policy on LanderCliEnv.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python train.py --steps 100000 # short smoke run
|
||||||
|
python train.py --steps 2000000 --n-envs 8 # full training
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from stable_baselines3 import PPO
|
||||||
|
from stable_baselines3.common.callbacks import CheckpointCallback
|
||||||
|
from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
|
||||||
|
|
||||||
|
# sys.path shim so this script can be run from any CWD.
|
||||||
|
import sys
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||||
|
|
||||||
|
from lander_cli_env import LanderCliEnv
|
||||||
|
|
||||||
|
|
||||||
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
CLI_BINARY = REPO_ROOT / "publish" / "GameCli" / "GameCli"
|
||||||
|
|
||||||
|
|
||||||
|
def make_env(seed: int, max_episode_steps: int):
|
||||||
|
"""Return a thunk that SubprocVecEnv can call to construct one env."""
|
||||||
|
def _fn():
|
||||||
|
env = LanderCliEnv(
|
||||||
|
cli_binary=str(CLI_BINARY),
|
||||||
|
max_episode_steps=max_episode_steps,
|
||||||
|
)
|
||||||
|
env.reset(seed=seed)
|
||||||
|
return env
|
||||||
|
return _fn
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument("--steps", type=int, default=2_000_000,
|
||||||
|
help="total env steps to train for")
|
||||||
|
parser.add_argument("--n-envs", type=int, default=8,
|
||||||
|
help="parallel envs; each spawns one GameCli subprocess")
|
||||||
|
parser.add_argument("--max-episode-steps", type=int, default=1000)
|
||||||
|
parser.add_argument("--seed", type=int, default=0)
|
||||||
|
parser.add_argument("--checkpoint-dir", type=Path,
|
||||||
|
default=REPO_ROOT / "checkpoints")
|
||||||
|
parser.add_argument("--tb-dir", type=Path,
|
||||||
|
default=REPO_ROOT / "tensorboard")
|
||||||
|
parser.add_argument("--use-dummy", action="store_true",
|
||||||
|
help="run envs in-process (slower, easier to debug)")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
if not CLI_BINARY.exists():
|
||||||
|
raise SystemExit(
|
||||||
|
f"CLI binary not found at {CLI_BINARY}. Run: "
|
||||||
|
f"dotnet publish GameCli/GameCli.csproj -c Release -o publish/GameCli"
|
||||||
|
)
|
||||||
|
|
||||||
|
args.checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
args.tb_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
env_fns = [
|
||||||
|
make_env(seed=args.seed + i, max_episode_steps=args.max_episode_steps)
|
||||||
|
for i in range(args.n_envs)
|
||||||
|
]
|
||||||
|
vec_env_cls = DummyVecEnv if args.use_dummy else SubprocVecEnv
|
||||||
|
vec_env = vec_env_cls(env_fns)
|
||||||
|
|
||||||
|
model = PPO(
|
||||||
|
"MlpPolicy",
|
||||||
|
vec_env,
|
||||||
|
verbose=1,
|
||||||
|
seed=args.seed,
|
||||||
|
tensorboard_log=str(args.tb_dir),
|
||||||
|
policy_kwargs=dict(net_arch=[64, 64]),
|
||||||
|
)
|
||||||
|
|
||||||
|
checkpoint_cb = CheckpointCallback(
|
||||||
|
save_freq=max(args.steps // 10, 1) // args.n_envs,
|
||||||
|
save_path=str(args.checkpoint_dir),
|
||||||
|
name_prefix="ppo_lander",
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
model.learn(total_timesteps=args.steps, callback=checkpoint_cb)
|
||||||
|
final_path = args.checkpoint_dir / "ppo_lander_final.zip"
|
||||||
|
model.save(str(final_path))
|
||||||
|
print(f"saved final model to {final_path}")
|
||||||
|
finally:
|
||||||
|
vec_env.close()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 2: Publish binary (idempotent)**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
dotnet publish GameCli/GameCli.csproj -c Release -o publish/GameCli
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 3: Smoke-run training for a small step count**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
source .venv/bin/activate
|
||||||
|
cd Training && python train.py --steps 5000 --n-envs 2 --use-dummy --max-episode-steps 200 && cd ..
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected:
|
||||||
|
- SB3 prints the standard rollout table (`| rollout/ | ...`).
|
||||||
|
- Final line: `saved final model to .../checkpoints/ppo_lander_final.zip`.
|
||||||
|
- Exit code 0.
|
||||||
|
- The file `checkpoints/ppo_lander_final.zip` exists.
|
||||||
|
|
||||||
|
The `--use-dummy` flag runs envs in-process which is easier for debugging. Real training uses SubprocVecEnv (leave that flag off).
|
||||||
|
|
||||||
|
- [ ] **Step 4: Commit**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
git add Training/train.py
|
||||||
|
git -c user.email=meelstorm@gmail.com -c user.name=meelstorm commit -m "feat(training): PPO training script with vectorized LanderCliEnv"
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Task 4: `export_onnx.py` — export trained policy to ONNX
|
||||||
|
|
||||||
|
**Files:**
|
||||||
|
- Create: `Training/export_onnx.py`
|
||||||
|
- Create: `Training/tests/test_export_onnx.py`
|
||||||
|
|
||||||
|
- [ ] **Step 1: Write the failing tests**
|
||||||
|
|
||||||
|
Create `Training/tests/test_export_onnx.py`:
|
||||||
|
|
||||||
|
```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"
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 2: Run tests to verify they fail**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd Training && pytest tests/test_export_onnx.py && cd ..
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected: fail because `export_onnx.py` does not exist.
|
||||||
|
|
||||||
|
- [ ] **Step 3: Write `Training/export_onnx.py`**
|
||||||
|
|
||||||
|
```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,
|
||||||
|
)
|
||||||
|
print(f"exported ONNX policy to {args.out}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 4: Run tests to verify they pass**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd Training && pytest tests/test_export_onnx.py && cd ..
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected: 3 passed. The `test_onnx_output_matches_sb3` test may occasionally show 1-2 argmax mismatches on marginal states — that's why the tolerance is `mismatches <= 2`. If it exceeds 2, something is wrong with the export.
|
||||||
|
|
||||||
|
- [ ] **Step 5: Verify the produced ONNX file**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
ls -la models/ppo_lander.onnx
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected: file exists, non-zero size.
|
||||||
|
|
||||||
|
- [ ] **Step 6: Commit**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
git add Training/export_onnx.py Training/tests/test_export_onnx.py
|
||||||
|
git -c user.email=meelstorm@gmail.com -c user.name=meelstorm commit -m "feat(training): export SB3 PPO policy to ONNX with parity test"
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Task 5: End-to-end pipeline smoke
|
||||||
|
|
||||||
|
**Files:** (no new files)
|
||||||
|
|
||||||
|
- [ ] **Step 1: Full pipeline dry-run**
|
||||||
|
|
||||||
|
Wipes state and re-runs the entire pipeline to prove it works end-to-end.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
rm -rf checkpoints/ models/
|
||||||
|
dotnet publish GameCli/GameCli.csproj -c Release -o publish/GameCli
|
||||||
|
source .venv/bin/activate
|
||||||
|
cd Training
|
||||||
|
python train.py --steps 10000 --n-envs 2 --use-dummy --max-episode-steps 200
|
||||||
|
python export_onnx.py --checkpoint ../checkpoints/ppo_lander_final.zip --out ../models/ppo_lander.onnx
|
||||||
|
cd ..
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected:
|
||||||
|
- Training completes without error.
|
||||||
|
- `checkpoints/ppo_lander_final.zip` exists.
|
||||||
|
- `models/ppo_lander.onnx` exists and is non-empty.
|
||||||
|
- Exit code 0.
|
||||||
|
|
||||||
|
- [ ] **Step 2: Full test suite**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd Training && pytest && cd ..
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected: 11 passed, 0 failed (8 env + 3 export).
|
||||||
|
|
||||||
|
- [ ] **Step 3: No commit — this task is verification only**
|
||||||
|
|
||||||
|
If everything passed, Plan 2 is done. If not, fix whatever failed and iterate.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Definition of done for Plan 2
|
||||||
|
|
||||||
|
- `cd Training && pytest` reports 11 passed, 0 failed.
|
||||||
|
- Running `python Training/train.py --steps 10000 --n-envs 2 --use-dummy` produces a checkpoint under `checkpoints/`.
|
||||||
|
- Running `python Training/export_onnx.py --checkpoint checkpoints/ppo_lander_final.zip --out models/ppo_lander.onnx` produces a non-empty ONNX file.
|
||||||
|
- The exported ONNX file's argmax action matches SB3's `predict(deterministic=True)` on ≥98% of random observations.
|
||||||
|
|
||||||
|
The next plan (Plan 3) will build the ASP.NET backend that loads this ONNX file and drives the CLI at 50 Hz per connected browser.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Notes on real training
|
||||||
|
|
||||||
|
The smoke-run steps use tiny values (10k steps, 2 envs) to prove the pipeline works quickly during implementation. To actually train a useful policy, the spec calls for ~2M steps on 8 envs:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python Training/train.py --steps 2000000 --n-envs 8
|
||||||
|
```
|
||||||
|
|
||||||
|
This takes on the order of hours on a modern CPU. It's out of scope for Plan 2 (which is about building the pipeline); do it before Plan 3 needs a real model, or use a smoke-trained model to prototype the backend and swap in the real one later.
|
||||||
Reference in New Issue
Block a user