MicroGame — Train Game-Playing Agents with SFT + On-Policy PPO

A compact, self-contained framework for training agents to play simple games:

SFT (behavior cloning) teaches basic skills first, then on-policy PPO improves them with rewards — backed by a pluggable reward registry, configurable action spaces, OCR-based screen reading, and an ONNX export path.

The framework is game-agnostic. It ships with three demo games that require no license, account, or hardware — all verified end to end. LaneBattle (mini MOBA) is the flagship example and comes with a trained model and real gameplay video; FortDefense (tower defense) demonstrates richer decision-making.

Quick Start

uv sync --extra device        # setup (uses the CUDA torch wheel in wheels/ if present)

# Watch the trained model play (recorded version in media/, live version below)
uv run scripts/predict.py --game lane_battle --weights weights/lane_battle_demo.pth

# Your first full training run in ~10 minutes (MiniCatch, CPU is fine)
uv run scripts/collect_mini.py --game catch --episodes 3 --out ./data_catch
uv run scripts/train_sft.py --data ./data_catch --encoder tiny --epochs 5 --no_pretrained
uv run scripts/train_ppo.py --env catch --init_from weights/bc_policy.pth --updates 20

Demo Games

LaneBattle — mini MOBA duel (flagship)

A single-lane 1v1: first to 3 kills wins. Everything a MOBA agent needs, compressed into a 112x112 pixel, dependency-free environment:

MOBA element	Implementation
Movement	Standard 8-direction + stop (subset of the default 10x13 factored space)
Combat	Basic attack (nearest unit, 22px) / skill (40px AoE, 100-frame cooldown)
Economy	Last-hit minions +5 gold, hero kill +50, gold rendered on the HUD
KDA	Kill/death counters, 20-frame respawn, first to 3 kills wins
Sustain	Recover action +15 HP (200-frame cooldown)
Enemy bot	Chases, retreats at low HP, heals, with movement noise
Combat feedback	Hero ring color reflects current state (attacking/taking damage/dead)

Trained model + real gameplay video (included in the repo):

• Model: weights/lane_battle_demo.pth — expert-data SFT, 10/10 win rate

• Video: media/lane_battle_demo.mp4 (3 straight wins, generated by scripts/record_video.py)

• Reproduce:

  uv run scripts/predict.py --game lane_battle --weights weights/lane_battle_demo.pth --episodes 5

FortDefense — tower defense (complex decisions)

Defend your fort through 8 waves of attackers (a boss every 4th wave), with a custom 9-movement x 4-action space (36 classes). Adds decision structure beyond LaneBattle:

Source of complexity	Mechanic
Resource allocation	Kills grant gold; "buy upgrade" converts gold into permanent attack power (increasing price) — save vs. spend now
Two-sided defense	Enemies march to your fort and chip it down; fort HP is the real loss condition
Wave pacing	A wave every 90 frames; boss waves demand cooldown and HP management
Unit priorities	Melee / ranged / boss differ in threat and value; ranged units kite — turtling loses

• Reward: dedicated fort_defense function (heavy fort-damage penalty, wave bonuses, upgrade purchases booked as instant cost)
• Expert demonstration: media/fort_defense_expert.mp4 (record_video.py --expert)
• Difficulty is configurable: MiniGameConfig(game_kwargs={"waves_target": N})

MiniCatch — your first PPO

Catch falling apples with a custom 3-movement x 2-action space (6 classes). Intuitive task, clear rewards — a tiny model on CPU shows clear evaluation improvement within minutes. Use it to run the full SFT -> PPO loop the first time.

Training Pipeline

 collect_mini.py            train_sft.py            train_ppo.py            export_onnx.py
 [frames + expert/human  ->  [behavior       ->  [on-policy PPO     ->  [ONNX + numeric
  action labels]             cloning / SFT]      improvement]           verification]
                                                      |
                                              predict.py / record_video.py

1. Collect SFT data

Two modes, same output format (data/<dir>/):

# Expert collection (recommended): the built-in expert plays and every frame
# is saved with its action label
uv run scripts/collect_mini.py --game lane_battle --episodes 8 --out ./data_lane

# Human collection (optional): play it yourself with the keyboard
# (requires GUI opencv + pynput; WASD move, J attack, K skill, L recover, Esc quit)
uv run --extra device scripts/collect_mini.py --game lane_battle --human

Output: action.jsonl (frame + flat action id per frame), full-resolution frames, and labels.json (the action-space definition — training builds matching network heads automatically). Around 8 episodes works well; too little data overfits the demonstration trajectory.

Model selection must be done by closed-loop win rate, not validation accuracy: variants trained from the same data with different seeds can vary from 0/3 to 3/3 while validation accuracy stays flat.

2. SFT (behavior cloning)

uv run scripts/train_sft.py --data ./data_lane --encoder tiny --epochs 6 --no_pretrained

Learns basic skills from demonstrations (attack minions in lane, retreat at low HP) — the policy needs fundamentals before on-policy exploration is meaningful. Backbones: mobilenetv4_conv_small (default, timm pretrained), convnextv2_nano, tiny (fast iteration). AMP is on by default.

3. PPO (on-policy fine-tuning)

uv run scripts/train_ppo.py --env catch --init_from weights/bc_policy.pth --updates 20
uv run scripts/train_ppo.py --env lane_battle --reward_fn kda --updates 60
uv run scripts/train_ppo.py --env fort_defense --gamma 0.995

• The algorithm follows rsl_rl 5.0.1: true-distribution KL, value clipping, NaN sanitization, separate actor/critic. Use --schedule fixed for single-environment small buffers (adaptive LR targets large batch regimes); use --gamma 0.995 for long-horizon games.
• Per-term reward logging ([reward] kill: ...) shows whether each signal is actually being produced.

4. Export and inference

uv run scripts/export_onnx.py --weights weights/lane_battle_demo.pth
uv run scripts/predict.py --game lane_battle --weights weights/policy.onnx
uv run scripts/record_video.py --game lane_battle --weights weights/lane_battle_demo.pth \
    --out media/my_video.mp4

.pth checkpoints are validated on load (contract version + action space); .onnx runs on onnxruntime. Deterministic argmax by default, --stochastic to sample. record_video.py renders MP4 gameplay (--expert records the built-in expert instead of a model).

Rewards

Every game ships a reward function designed for it — the framework itself has no game semantics. All rewards share one signature:

def fn(ctx: RewardContext) -> tuple[float, dict]:
    """returns (total reward, per-term breakdown for logging)"""

Name	Game	Design
kda	LaneBattle	Kill +5 / death -5 / gold 0.1 per point / damage taken -0.5 / enemy HP loss +1.0 / win ±100
fort_defense	FortDefense	Fort damage -1.0 per point / kill +1, boss +5 / wave cleared +8 / upgrade purchase booked as cost / win ±100
catch	MiniCatch	Apple +1 / bomb -1.5 / optional tracking shaping for cold start

Register your own with one decorator — RewardContext carries HP, gold, KDA, wave progress, confidence and more; missing signals are skipped automatically:

@register_reward("my_mode")
def _make_my_mode(**kwargs):
    def fn(ctx):
        ...
        return reward, terms

RewardContext.gold can be fed from the screen: LaneBattle supports use_ocr_gold=True, where an OnnxOCR engine (PPOCRv5 ONNX, recognition-only path, ~16ms per step) reads the gold counter off the rendered HUD — the same pattern you would use for a real game.

Adding a New Game

Implement four hooks plus a registry entry; the collection/training/inference scripts work unchanged:

from envs.mini.base import MiniGameEnv, MiniGameConfig, register_mini_game
from contract import ActionSpace

MY_SPACE = ActionSpace(move_labels=("Left", "Right"), act_labels=("Jump", "Idle"))

@register_mini_game("my_game")
class MyGame(MiniGameEnv):
    action_space = MY_SPACE
    KEYMAP = {"a": (0, 0), "d": (1, 0), "space": (0, 1)}

    def _reset_state(self): ...
    def _step_dynamics(self, move_idx, act_idx): ...
    def _render(self) -> np.ndarray: ...          # uint8 RGB frame
    def _context(self) -> RewardContext: ...      # reward signals for this step

Requirements for any new environment: uint8 RGB frames, an enumerable action space, and a computable reward.

Tests

uv run pytest -q    # contract / models / SFT+PPO pipeline / games / OCR / ONNX