mg/general-bots-training
Watch
1
0
Fork
You've already forked general-bots-training
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

checkpoint

Browse source
This commit is contained in:
Moritz Gmeiner 2026-08-07 22:06:27 +02:00
commit d0deefd8cf
35 changed files with 2543 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

160
README.md Normal file
View file

@ -0,0 +1,160 @@
# 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`.