Add JAX PPO static-opponent trainer
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user