Add JAX PPO static-opponent trainer

This commit is contained in:
2026-07-04 22:30:16 +09:00
parent 768f74693d
commit 9e27f42f27
12 changed files with 1682 additions and 0 deletions
+41
View File
@@ -196,6 +196,47 @@ steps_per_sec=3956598.38
- Observation scalar normalization is implementation-defined as documented
above and locked by the exported `OBS_DIM`.
## JAX PPO Static-Opponent Ladder
The first training stack above `lost_cities_jax` is exposed as:
```bash
uv run lost-cities-jax-ppo --help
```
CPU smoke:
```bash
uv run lost-cities-jax-ppo rollout-smoke --config configs/jax_ppo/smoke.yaml
uv run lost-cities-jax-ppo train --config configs/jax_ppo/smoke.yaml
```
GPU training uses optional CUDA JAX, keeping CUDA wheels out of the default
project dependency set:
```bash
tmux new-session -s coolrl-jax-ppo-discard \
-c /home/coolguy/dev/coolrl-lost-cities \
"flock -n .compute.lock uv run --with 'jax[cuda12]' lost-cities-jax-ppo train \
--config configs/jax_ppo/discard-only.yaml"
```
Random-policy baseline vs `discard_only`, measured on 2026-07-04 with
8192 games x 400 plies on GPU:
```text
return_mean=-55.582763671875
game_length_mean=69.6085205078125
max_steps_rate=0.0
play_action_rate=0.28854578733444214
opened_colors_mean=4.942626953125
positive_expeditions_mean=0.41796875
```
Large PPO artifacts are written under
`/mnt/2tbhdd/coolrl-lost-cities-artifacts/jax-ppo-static-opponents/`.
Generated checkpoints and evaluation JSON are not committed to git.
## Basic Usage
```python