feat(training): add executable masked PPO training and evaluation

Author: nosam1998

.github/workflows/training.yml

@@ -0,0 +1,40 @@
+name: Training CI
+on:
+  pull_request:
+  push:
+    branches: [main]
+permissions:
+  contents: read
+concurrency:
+  group: training-ci-${{ github.ref }}
+  cancel-in-progress: true
+jobs:
+  masked-ppo:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+          cache-dependency-path: pyproject.toml
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2
+      - run: python -m pip install "torch>=2.3,<3" --index-url https://download.pytorch.org/whl/cpu
+      - run: python -m pip install "maturin>=1.9,<2" "pytest>=8,<9" "numpy>=1.24,<3"
+      - run: maturin build --release --locked --out dist
+      - name: Install wheel with training dependencies
+        run: |
+          python - <<'PY'
+          import glob
+          import subprocess
+          import sys
+          subprocess.check_call([sys.executable, "-m", "pip", "install", glob.glob("dist/*.whl")[0] + "[training]"])
+          PY
+      - run: python -m pytest python/tests/test_ppo_example.py python/tests/test_training_mask.py -q
+      - run: python examples/train_ppo.py --steps 4096 --envs 2 --rollout 64 --epochs 2 --eval-games 8 --output training-output
+      - uses: actions/upload-artifact@v4
+        with:
+          name: masked-ppo-smoke
+          path: training-output/

.gitignore

@@ -12,3 +12,5 @@ __pycache__/
 
 /playwright-report/
 /test-results/
+
+training-output/

crates/apps/gfa-py/src/lib.rs

@@ -111,6 +111,11 @@ impl NativeEnv {
         .map_err(error)
     }
 
+    fn action_mask<'py>(&self, py: Python<'py>, seat: u8) -> PyResult<Bound<'py, PyArray1<bool>>> {
+        let mask = py.detach(|| self.inner.action_mask(seat)).map_err(error)?;
+        Ok(PyArray1::from_vec(py, mask))
+    }
+
     fn spec_json(&self) -> PyResult<String> {
         serde_json::to_string(&self.inner.spec()).map_err(error)
     }

examples/README.md

@@ -0,0 +1,24 @@
+# Training examples
+
+Install a wheel with optional training dependencies, or run `python -m pip install '.[training]'` from the repository root. The environment itself does not require a training framework or host model weights.
+
+```sh
+python examples/train_ppo.py --game connect4 --steps 250000 --envs 8 --output training-output
+```
+
+The example uses [SB3-Contrib MaskablePPO](https://sb3-contrib.readthedocs.io/en/master/modules/ppo_mask.html) with a small CPU policy and GamesForAI's Gymnasium action masks. It trains both seats against the configured opponent. Replace the policy architecture or use the same environment in your own trainer; model training stays in your process.
+
+`--opponent minimax:3` trains against that native search baseline. `--eval-opponent minimax:3` selects the comparison baseline; evaluation also runs against random play. Fixed held-out seeds, balanced seats, and two random opening plies by default create repeatable evaluation positions. `--opening-plies 0` evaluates only standard starts. Reports include sample size, outcomes, win-rate intervals, raw mean returns, package versions, seeds, and training configuration. Small smoke-run estimates do not establish model strength.
+
+`policy.zip` is saved and loaded before evaluation. `report.json` records changed parameters, exact checkpoint verification, and a repeated evaluation check. CI runs actual training with a small budget and uploads those artifacts; the larger PRD acceptance run against level 3 is separate from this integration test. The example uses SB3's environment batching; native VectorEnv is available for custom batched self-play trainers.
+
+```python
+from sb3_contrib import MaskablePPO
+from gamesforai import make
+
+model = MaskablePPO.load("training-output/policy.zip", device="cpu")
+env = make("connect4", opponent="minimax:3")
+obs, info = env.reset(seed=123)
+action, _ = model.predict(obs, action_masks=env.action_masks(), deterministic=True)
+obs, reward, terminated, truncated, info = env.step(int(action))
+```

examples/train_ppo.py

@@ -0,0 +1,153 @@
+"""Train and evaluate a masked policy using an installed gamesforai wheel."""
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+from importlib.metadata import version
+from pathlib import Path
+
+import numpy as np
+import torch
+from sb3_contrib import MaskablePPO
+from stable_baselines3.common.vec_env import DummyVecEnv
+
+from gamesforai import NativeEnv, make
+
+
+def digest(policy):
+    result = hashlib.sha256()
+    for name, tensor in sorted(policy.state_dict().items()):
+        result.update(name.encode())
+        result.update(tensor.detach().cpu().contiguous().numpy().tobytes())
+    return result.hexdigest()
+
+
+def opening(game, seed, plies):
+    env = NativeEnv(game, seed=seed)
+    rng = np.random.default_rng(seed)
+    for _ in range(plies):
+        actors = env.current_players()
+        if not actors:
+            raise ValueError("opening unexpectedly ended the game")
+        env.step(actors[0], int(rng.choice(np.flatnonzero(env.action_mask(actors[0])))))
+    return json.dumps(json.loads(env.get_state())["state"])
+
+
+def evaluate(model, game, opponent, episodes, seed, opening_plies):
+    results = []
+    for episode in range(episodes):
+        env = make(game, seat=episode % 2, opponent=opponent)
+        obs, info = env.reset(seed=seed + episode, options={
+            "position": opening(game, seed + episode, opening_plies)
+        })
+        total = 0.0
+        for _ in range(42):
+            action, _ = model.predict(obs, action_masks=env.action_masks(), deterministic=True)
+            obs, reward, terminated, truncated, info = env.step(int(action))
+            total += reward
+            if terminated or truncated:
+                break
+        else:
+            raise RuntimeError("evaluation episode exceeded game bounds")
+        env.close()
+        results.append(total)
+    wins = sum(reward > 0 for reward in results)
+    draws = sum(reward == 0 for reward in results)
+    probability, z = wins / episodes, 1.959963984540054
+    denominator = 1 + z * z / episodes
+    center = (probability + z * z / (2 * episodes)) / denominator
+    radius = z * math.sqrt(probability * (1 - probability) / episodes + z * z / (4 * episodes**2)) / denominator
+    return {
+        "opponent": opponent, "games": episodes, "wins": wins, "draws": draws,
+        "losses": episodes - wins - draws, "win_rate": probability,
+        "win_rate_wilson95": [max(0.0, center - radius), min(1.0, center + radius)],
+        "score": (wins + 0.5 * draws) / episodes, "mean_return": float(np.mean(results)),
+        "evaluation_seed": seed, "random_opening_plies": opening_plies,
+    }
+
+
+def run(args):
+    if args.game not in ("tictactoe", "connect4"):
+        raise ValueError("example supports Tic-Tac-Toe and Connect Four")
+    if not 1 <= args.envs <= 64 or not 1 <= args.threads <= 64:
+        raise ValueError("envs and threads must be 1..64")
+    if args.steps < 1 or args.eval_games < 2 or not 0 <= args.opening_plies <= 2:
+        raise ValueError("require positive steps, >=2 evaluation games, and 0..2 opening plies")
+    if args.rollout < 2 or args.epochs < 1 or args.batch_size < 2:
+        raise ValueError("invalid PPO rollout, batch, or epoch sizes")
+    if args.batch_size > args.rollout * args.envs:
+        raise ValueError("batch size exceeds rollout size")
+    output = Path(args.output)
+    output.mkdir(parents=True, exist_ok=True)
+    torch.set_num_threads(args.threads)
+    torch.use_deterministic_algorithms(True)
+    environments = DummyVecEnv([
+        (lambda seat=i % 2: make(args.game, seat=seat, opponent=args.opponent))
+        for i in range(args.envs)
+    ])
+    environments.seed(args.seed)
+    model = MaskablePPO(
+        "MlpPolicy", environments, n_steps=args.rollout, batch_size=args.batch_size,
+        n_epochs=args.epochs, learning_rate=3e-4, gamma=0.99, ent_coef=0.01,
+        policy_kwargs={"net_arch": [64, 64]}, seed=args.seed, device="cpu", verbose=1,
+    )
+    evaluation_seed = args.seed + 1_000_000
+    opponents = list(dict.fromkeys(["random", args.eval_opponent]))
+    initial_digest = digest(model.policy)
+    before = [
+        evaluate(model, args.game, opponent, args.eval_games, evaluation_seed, args.opening_plies)
+        for opponent in opponents
+    ]
+    model.learn(total_timesteps=args.steps)
+    trained_digest = digest(model.policy)
+    if trained_digest == initial_digest:
+        raise RuntimeError("training did not change policy parameters")
+    model_path = output / "policy.zip"
+    model.save(model_path)
+    restored = MaskablePPO.load(model_path, device="cpu")
+    if digest(restored.policy) != trained_digest:
+        raise RuntimeError("saved policy parameters differ after loading")
+    after = [
+        evaluate(restored, args.game, opponent, args.eval_games, evaluation_seed, args.opening_plies)
+        for opponent in opponents
+    ]
+    repeat = evaluate(restored, args.game, opponents[-1], args.eval_games, evaluation_seed, args.opening_plies)
+    if repeat != after[-1]:
+        raise RuntimeError("fixed evaluation did not reproduce")
+    report = {
+        "game": args.game, "training_seed": args.seed, "requested_steps": args.steps,
+        "actual_steps": model.num_timesteps, "training_opponent": args.opponent,
+        "parameters_changed": True, "checkpoint_verified": True,
+        "evaluation_reproduced": True, "policy_sha256": trained_digest,
+        "initial_policy_sha256": initial_digest, "before": before, "after": after,
+        "versions": {name: version(name) for name in ("gamesforai", "sb3-contrib", "stable-baselines3", "torch", "numpy")},
+        "configuration": vars(args),
+    }
+    (output / "report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
+    environments.close()
+    print(json.dumps(report, indent=2))
+    return report
+
+
+def parser():
+    result = argparse.ArgumentParser(description=__doc__)
+    result.add_argument("--game", choices=["tictactoe", "connect4"], default="connect4")
+    result.add_argument("--steps", type=int, default=250_000)
+    result.add_argument("--envs", type=int, default=8)
+    result.add_argument("--rollout", type=int, default=128)
+    result.add_argument("--batch-size", type=int, default=64)
+    result.add_argument("--epochs", type=int, default=4)
+    result.add_argument("--threads", type=int, default=1)
+    result.add_argument("--seed", type=int, default=42)
+    result.add_argument("--opponent", default="random")
+    result.add_argument("--eval-opponent", default="minimax:3")
+    result.add_argument("--eval-games", type=int, default=100)
+    result.add_argument("--opening-plies", type=int, default=2)
+    result.add_argument("--output", default="training-output")
+    return result
+
+
+if __name__ == "__main__":
+    run(parser().parse_args())

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = ["numpy>=1.24,<3", "gymnasium>=1.2,<2", "pettingzoo>=1.25,<2"]
 
 [project.optional-dependencies]
 test = ["pytest>=8,<9"]
+training = ["sb3-contrib>=2.7,<3", "torch>=2.3,<3"]
 
 [tool.maturin]
 manifest-path = "crates/apps/gfa-py/Cargo.toml"

python/gamesforai/_native.pyi

@@ -11,6 +11,7 @@ class NativeEnv:
     def step(self, seat: int, action: int) -> tuple[list[float], bool, bool]: ...
     def step_string(self, seat: int, action: str) -> tuple[list[float], bool, bool]: ...
     def builtin_action(self, seat: int, algorithm: str, level: int, seed: int) -> int: ...
+    def action_mask(self, seat: int) -> Any: ...
     def spec_json(self) -> str: ...
     def current_players(self) -> list[int]: ...
     def flags(self) -> tuple[bool, bool]: ...

python/gamesforai/gym_env.py

@@ -146,6 +146,12 @@ def step(self, action: int) -> tuple[np.ndarray, float, bool, bool, dict[str, An
         observation, info = self._frame(self.seat)
         return observation, float(reward), terminated, truncated, info
 
+    def action_masks(self) -> np.ndarray:
+        """Boolean legal-action mask used by MaskablePPO and other trainers."""
+        if self._needs_reset:
+            raise ResetNeeded("Call reset before requesting an active policy mask")
+        return self.native.action_mask(self.seat)
+
     def render(self) -> str:
         return self.native.frame(self.seat)["text"]
 

python/tests/test_ppo_example.py

@@ -0,0 +1,22 @@
+from pathlib import Path
+import runpy
+
+import pytest
+
+pytest.importorskip("sb3_contrib")
+
+
+def test_real_ppo_updates_saves_loads_and_repeats_evaluation(tmp_path):
+    example = runpy.run_path(str(Path(__file__).parents[2] / "examples/train_ppo.py"))
+    args = example["parser"]().parse_args([
+        "--game", "tictactoe", "--steps", "512", "--envs", "2", "--rollout", "32",
+        "--batch-size", "32", "--epochs", "2", "--eval-games", "4",
+        "--eval-opponent", "minimax:1", "--output", str(tmp_path),
+    ])
+    report = example["run"](args)
+    assert report["parameters_changed"] and report["checkpoint_verified"]
+    assert report["evaluation_reproduced"]
+    assert report["actual_steps"] >= 512
+    assert (tmp_path / "policy.zip").is_file()
+    assert (tmp_path / "report.json").is_file()
+    assert all(row["games"] == row["wins"] + row["draws"] + row["losses"] for row in report["after"])

python/tests/test_training_mask.py

@@ -0,0 +1,19 @@
+import numpy as np
+import pytest
+from gymnasium.error import ResetNeeded
+
+from gamesforai import make
+
+
+def test_maskable_training_api_is_boolean_owned_and_requires_active_episode():
+    env = make("connect4")
+    with pytest.raises(ResetNeeded):
+        env.action_masks()
+    _, info = env.reset(seed=1)
+    mask = env.action_masks()
+    assert mask.dtype == np.bool_
+    np.testing.assert_array_equal(mask, info["action_mask"])
+    mask[:] = False
+    assert env.action_masks().all()
+    _, _, _, _, info = env.step(3)
+    np.testing.assert_array_equal(env.action_masks(), info["action_mask"])