# general-bots-training

Reinforcement learning training for the [generals.bot](https://github.com/strakam/generals-bots) competition. Trains a policy with PPO against the `generals` JAX environment.

## Status

**v1** — minimal working PPO loop. 4×4 grid, fog of war, composite reward shaping, configurable static opponents, and two-sided self-play. The code is structured so planned extensions such as opponent leagues, curriculum, and the competition ruleset fit without rewriting the core.

## Quick start

```bash
# install (editable, picks up src/general_bots_training)
uv sync

# train (uses GPU if available, falls back to CPU)
uv run python scripts/train.py
```

In a network-restricted sandbox, `uv run` may fail to re-resolve the build backend; use the venv directly:

```bash
.venv/bin/python scripts/train.py
```

Checkpoints are written to `ppo_model.eqx` by default. Override this with `checkpoint_path=...`.

## Configuration

[`scripts/train.py`](scripts/train.py) defines a typed OmegaConf configuration. Override defaults with `key=value` arguments; CLI values are merged over the structured defaults:

```bash
uv run python scripts/train.py num_envs=32 rollout_steps=64 lr=1e-4
uv run python scripts/train.py grid_dims=[6,6] opponent=expander
uv run python scripts/train.py opponent=hunter checkpoint_path=models/ppo-hunter.eqx
uv run python scripts/train.py resume_from=models/ppo-hunter.eqx opponent=expander
uv run python scripts/train.py opponent=self_play
```

Unknown keys and incompatible value types are rejected.

| key                           | default         | notes                                          |
| ----------------------------- | --------------- | ---------------------------------------------- |
| `grid_dims`                   | `[4, 4]`        | start small; curriculum to larger grids later  |
| `truncation`                  | `500`           | max turns before a game is scored as a draw    |
| `num_envs`                    | `256`           | parallel games; tune to GPU VRAM               |
| `rollout_steps`               | `256`           | steps per rollout before a PPO update          |
| `num_iterations`              | `500`           | PPO update count                               |
| `num_epochs`                  | `1`             | epochs over each rollout buffer                |
| `minibatch_size`              | `256`           |                                                |
| `lr`                          | `3e-4`          | Adam                                           |
| `gamma` / `lam`               | `0.99` / `0.95` | GAE                                            |
| `clip`                        | `0.2`           | PPO ratio clip                                 |
| `value_coef` / `entropy_coef` | `0.5` / `0.01`  |                                                |
| `log_every`                   | `10`            | iterations between progress logs               |
| `checkpoint_path`             | `ppo_model.eqx` | output model path                              |
| `resume_from`                 | `null`          | optional model checkpoint to continue from     |
| `opponent`                    | `random`        | `random`, `expander`, `hunter`, or `self_play` |
| `seed`                        | `0`             | JAX random seed                                |

`resume_from` restores the policy/value network weights. Existing checkpoints do not contain optimizer state, so Adam starts with fresh moments and the configured learning rate.

## Layout

```
src/general_bots_training/
  network.py   # equinox conv policy-value net + observation encoding
  mcts.py      # competition observation adapter and particle PUCT search
  ppo.py       # reusable GAE, clipped PPO loss, and optimizer helpers
  opponents.py # opponent interfaces and named strategy selection
  rollout.py   # jitted static-opponent and two-sided self-play collection
scripts/
  train.py      # executable config, training loop, logging, checkpointing
  mcts_agent.py # competition stdio inference entrypoint
agents/mcts/
  run.sh        # local matchup wrapper for ppo_model.eqx
```

## Architecture

### Network (`network.py`)

`PolicyValueNetwork` is an equinox module:

- **Backbone**: 4 conv layers (3×3, padding=1) over a 14-channel normalized observation. Armies, army-counts, and timestep are log-normalized; scalar values are broadcast to spatial planes so a plain conv stack can consume them.
- **Policy head**: 1×1 conv to 9 channels = 4 full-move directions + 4 half-move (split) directions + a spatial pass score. Move channels are flattened and the pass scores are spatially pooled into one global pass action, yielding `8*H*W+1` logits. Invalid moves are masked to −1e9 via `compute_valid_move_mask`; pass is always available.
- **Value head**: 1×1 conv → global average pool → 2-layer MLP → scalar. Global pooling makes the network grid-size-agnostic, so the same architecture extends to larger boards without reshaping linear layers.

`obs_to_tensor` encodes a `generals.Observation` into the `(14, H, W)` float32 input.

### Rollout (`rollout.py`)

`make_collect_rollout(env, num_steps, opponent)` accepts a stateless JAX-compatible agent or the self-play marker and returns a jitted function `(states, pool, network, key) -> (states, transitions, (key, last_next_obs))`. Each step:

1. Observe both players from the current state.
2. Sample p0's action from the policy network. For a static opponent, obtain p1's action from `RandomAgent`, `ExpanderAgent`, or `HunterAgent`; in `self_play`, sample p1 independently from the same current network.
3. Step the env (vmapped), which auto-resets from the pool on done.
4. Compute the shaped reward for p0 with `composite_reward_fn` from the pre-step and post-step observations. The post-step observation is taken from `timestep.last_state` (the state _before_ auto-reset) so terminal and shaping rewards are computed against the actual end-of-episode board.
5. Record `(obs, mask, action, logprob, value, reward, done, winner)`.

The bootstrap observation for the critic is threaded through the `lax.scan` carry; only the final step's is returned (as `last_next_obs`) to avoid storing T copies. Static-opponent rollouts produce `N` trajectories per step. Self-play produces `2N`, with observations, actions, shaped rewards, values, and log-probabilities from both player perspectives included in the same PPO update.

### PPO (`ppo.py`)

- `compute_gae`: GAE via reverse `lax.scan`, bootstrapping from the critic value of the post-rollout state (zeroed on done steps).
- `ppo_loss`: clipped surrogate + value loss + entropy bonus.
- `make_train_epoch`: flattens `(T, N)` → `(T*N)`, shuffles, minibatches with `eqx.filter_grad`.

### Training loop (`scripts/train.py`)

Non-mutating warmup (compile) → per iteration: collect rollout → GAE → compute returns from raw advantages → normalize policy advantages → PPO update → log (loss, reward, episodes, win/loss, SPS) → checkpoint at the end.

### Competition PUCT (`mcts.py`)

Run the local stdio bot directly through the bundled matchup driver:

```bash
PYTHONPATH=src:generals-bots .venv/bin/python generals-bots/competition/matchup.py \
  agents/mcts/run.sh \
  generals-bots/competition/agents/expander_python/run.sh \
  --mode competition
```

The bot performs deadline-bounded root PUCT using only the perspective-relative wire observation. Each simulation samples a hidden-state determinization consistent with visible ownership and global opponent totals, samples a simultaneous opponent action from the same policy, applies build-castles and deathtouch transitions, and evaluates the resulting leaf with the critic. Network move/pass logits provide priors; affordable build actions are added with exact legality and heuristic priors so existing checkpoints remain compatible.

The handshake warmup compiles all board-shape-dependent paths before the first action. On a pinned Ryzen 5800X core, a 21×21 search configured for 125 ms completed in approximately 111 ms with seven depth-2 simulations. Results depend on CPU and position complexity.

This is a conservative first particle search, not full information-set MCTS: particles are regenerated from each current observation and do not yet maintain a persistent history belief. Also, the published competition environment manifest includes JAX but not Equinox or `generals-bots`; `agents/mcts/run.sh` is therefore a local evaluation wrapper. A submitted bot must bundle those dependencies or export the network/simulator to the sandbox's available runtime.

## Key correctness choices

These differ from the experimental reference in `generals-bots/examples/_experimental/ppo/`:

- **Bootstrap GAE from the post-step critic value**, not 0. Done steps are zeroed via the done mask, so a fresh reset state's value doesn't contaminate the advantage.
- **Post-step observation from `timestep.last_state`** (pre-auto-reset) so terminal/shaping rewards are correct. The env's auto-reset overwrites the state with a fresh board; using that for reward shaping would attribute the reset board's counts to the just-finished episode.
- **Thread the pool explicitly** through `env.step` (vmapped) rather than capturing it as a constant, so it isn't baked into the JIT trace.
- **`jax.vmap(network, in_axes=(0,0,None,0))`** for batched forward — the network is the vmapped callable, so its weight leaves are batched alongside the data. This composes correctly with `eqx.filter_grad`; `eqx.filter_vmap` on a closure capturing the network does not.

## Validation

Validated end-to-end on CPU (the sandbox has no GPU):

- Compiles in ~20s, ~320 SPS on 32 envs / 200-step rollouts.
- Episodes complete, win/loss counting works, checkpoints save.
- An untrained network wins ~20–44% vs random (random also wins some by accident) — a sensible starting point.

On a 4080 / rented GPU, throughput should be substantially higher; tune `NUM_ENVS` and `ROLLOUT_STEPS` to VRAM.

## Notes

- The `generals` package is pinned via git in `[tool.uv.sources]`; the `generals-bots/` subdir is a clone for reference and is not part of the build.
- GPU isn't visible from the Zed sandbox (`cuInit` fails → CPU fallback). Run `scripts/train.py` from your local machine or a GPU host for CUDA.

## Roadmap

Planned extensions, in rough priority order:

1. **Opponent league** — extend current-policy self-play with frozen historical snapshots to reduce strategy collapse.
2. **Curriculum** — step up from 4×4 to larger grids, then to `GeneralsEnv(mode="competition")` (variable 18–21 grids, 1200-step truncation, `build_castles` + `deathtouch` modifiers).
3. **Algorithm swap** — the PPO logic is isolated in `ppo.py`; REINFORCE or another algorithm can replace it without touching the rollout or network.
4. **Evaluation harness** — match the trained policy against the bundled `ExpanderAgent` and the competition's stdio bots via `competition/matchup.py`.