Files
Experiment_ReinforcementLea…/docs/superpowers/plans/2026-07-17-02-training-pipeline.md
2026-07-17 18:23:16 +00:00

24 KiB

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
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
[pytest]
testpaths = tests
python_files = test_*.py
addopts = -v
  • Step 6: Write Training/tests/__init__.py (empty file)
touch Training/tests/__init__.py
  • Step 7: Verify pytest can discover an empty test suite

From Training/ directory:

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
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:

"""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
cd Training && pytest && cd ..

Expected: import error — lander_cli_env module does not exist.

  • Step 3: Write Training/lander_cli_env.py
"""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:

dotnet publish GameCli/GameCli.csproj -c Release -o publish/GameCli
  • Step 5: Run tests to verify they pass
cd Training && pytest tests/test_lander_cli_env.py && cd ..

Expected: 8 passed, 0 failed.

  • Step 6: Commit
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

"""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)
dotnet publish GameCli/GameCli.csproj -c Release -o publish/GameCli
  • Step 3: Smoke-run training for a small step count
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
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:

"""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
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
"""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
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
ls -la models/ppo_lander.onnx

Expected: file exists, non-zero size.

  • Step 6: Commit
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.

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

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:

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.