Focus project on JAX PPO

This commit is contained in:
2026-07-14 20:09:03 +09:00
parent 79273f7eb3
commit ef4b9d82b0
44 changed files with 302 additions and 503 deletions
+25
View File
@@ -0,0 +1,25 @@
# Legacy research stacks
The supported project path is the JAX rules engine and JAX PPO tooling exposed
by `lost-cities-jax-ppo`. The historical Deep CFR and ISMCTS implementations
are preserved for research reproduction, but they are not current training
recipes and no longer have public console-script entry points.
## What is preserved
- Deep CFR sources: `src/coolrl_lost_cities/games/classic/deep_cfr/`
- ISMCTS sources: `src/coolrl_lost_cities/games/classic/ismcts/`
- Their retired configurations: `legacy/deep-cfr/configs/` and
`legacy/ismcts/configs/`
- Earlier plans and dated findings: `docs/plans/archive/`, `docs/archive/`,
and `docs/research/`
These files remain in the repository so a historical experiment can be read or
reproduced against its original commit. They should not be chosen for new
models, benchmarks, or product work.
## Current workflow
Start from `configs/jax_ppo/`, run `lost-cities-jax-ppo`, and use the browser
client in `web/`. The root [README](../README.md) contains the short training,
evaluation, and static-web recipes.
+8 -5
View File
@@ -1,7 +1,10 @@
# Deep CFR Performance Notes
# Legacy Deep CFR Performance Notes
This document tracks current runtime bottlenecks for the active Deep CFR
training path. The numbers below are observational, not a benchmark contract.
> Historical record only. The supported training stack is JAX PPO; see the
> root [README](../README.md) and [legacy notes](legacy.md).
This document records the Deep CFR runtime bottlenecks observed before the JAX
PPO transition. The numbers below are historical, not a benchmark contract.
## Current Default Runtime
@@ -11,7 +14,7 @@ Source run:
runs/tmp/2026-05-07_171535_deep-cfr-default/metrics.jsonl
```
The run used `configs/deep_cfr/default.yaml` with CUDA enabled. At the time of
The run used `legacy/deep-cfr/configs/default.yaml` with CUDA enabled. At the time of
inspection, completed metrics covered iterations 70 through 95. The training
process was still running, so later rows may differ.
@@ -88,7 +91,7 @@ This gives 560 traversals per iteration, split into 70 worker batches.
## Device Use
The trainer constructs the advantage and strategy networks on `run.device`.
`configs/deep_cfr/default.yaml` sets:
`legacy/deep-cfr/configs/default.yaml` sets:
```yaml
run:
@@ -117,7 +117,7 @@ operator applies the patch.
- ✅ AGENTS.md "Docs & Experiment Workflow" section landed
(commit `09d5815`, 2026-05-07).
- ✅ Plan drafted at `docs/plans/librarian.md` (this file).
- ✅ Plan drafted at `docs/plans/archive/librarian.md` (this file).
- ✅ Prompt moved: `.claude/agents/librarian.md`
`scripts/librarian-prompt.md`. Claude-specific subagent registration
removed.
@@ -1,11 +1,11 @@
# Plan: Model-Size Experiment (Keystone for Model-Scale Optimizations)
> **⚠️ 스택 주의 — 이 문서는 Deep CFR / PyTorch 스택 전용이다.**
> `input_dim=365`, `configs/deep_cfr/`, `DeepCFRMLP` 기준으로 쓰였다.
> `input_dim=365`, `legacy/deep-cfr/configs/`, `DeepCFRMLP` 기준으로 쓰였다.
> 현행 JAX PPO 스택(`OBS_DIM=454`, `configs/jax_ppo/`, `ActorCritic`)에는
> **적용되지 않는다.** PPO 쪽 모델 크기 문제는
> [jax-ppo-model-size-ab.md](jax-ppo-model-size-ab.md)와
> [../reports/jax-ppo-model-capacity-2026-07-12.md](../reports/jax-ppo-model-capacity-2026-07-12.md)를 볼 것.
> [jax-ppo-model-size-ab.md](../jax-ppo-model-size-ab.md)와
> [../../reports/jax-ppo-model-capacity-2026-07-12.md](../../reports/jax-ppo-model-capacity-2026-07-12.md)를 볼 것.
**Status:** Ready to execute
**Owner:** Operator (runs grid on `home`); Codex (adds configs and runner script)
+2 -2
View File
@@ -7,8 +7,8 @@
**Background:** [docs/reports/jax-ppo-model-capacity-2026-07-12.md](../reports/jax-ppo-model-capacity-2026-07-12.md)
> 이 계획은 은퇴한 Deep CFR/PyTorch 스택용
> [model_size_experiment.md](model_size_experiment.md)를 **대체하지 않는다** —
> 그쪽은 다른 스택(`input_dim=365`, `configs/deep_cfr/`) 이야기다. 서로 무관하다.
> [model_size_experiment.md](archive/model_size_experiment.md)를 **대체하지 않는다** —
> 그쪽은 다른 스택(`input_dim=365`, `legacy/deep-cfr/configs/`) 이야기다. 서로 무관하다.
## 가설
@@ -80,4 +80,4 @@ Validation performed before final report:
- `uv run pytest -q` passed: 257 passed, 1 skipped.
- `uv run lost-cities-jax-ppo play --help` passed.
- `uv run lost-cities-jax-ppo human-play summarize --log-dir /tmp/nonexistent-human-play-log` passed.
- `scripts/librarian.sh` found no link or code-citation errors; it still exits non-zero on the known pre-existing `docs/plans/deep-cfr-selectivity.md` 500-line soft cap.
- `scripts/librarian.sh` found no link or code-citation errors; it still exits non-zero on the known pre-existing `docs/plans/archive/deep-cfr-selectivity.md` 500-line soft cap.
@@ -23,8 +23,8 @@ Metrics: `.../final-cycles/2026-07-05/league/2026-07-05_191933_jax-ppo-final-cyc
(`main_ppo_config.json``network: {hidden_size: 512, num_layers: 3}`).
`hidden_size`를 다룬 기존 문서는 전부 은퇴한 Deep CFR/PyTorch 스택 것이다
(`input_dim=365`, `configs/deep_cfr/` 기준). 특히
[docs/plans/model_size_experiment.md](../plans/model_size_experiment.md)는
(`input_dim=365`, `legacy/deep-cfr/configs/` 기준). 특히
[docs/plans/model_size_experiment.md](../plans/archive/model_size_experiment.md)는
**현 JAX PPO 스택과 무관하다.**
→ 512×3은 실험으로 고른 값이 아니라 구 스택에서 복사돼 온 값이다.
@@ -54,5 +54,5 @@ The 8-worker interleaved path is more effective than a single-process CUDA path
- `src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py`: Implementation of AMP and training loops.
- `src/coolrl_lost_cities/games/classic/deep_cfr/networks.py`: `DeepCFRMLP` architecture.
- `configs/deep_cfr/default.yaml`: Configuration for interleaved scheduler.
- `scripts/profile_gpu_forward.py`: GPU forward pass micro-benchmarks.
- `legacy/deep-cfr/configs/default.yaml`: Configuration for interleaved scheduler.
- `scripts/profile_gpu_forward.py`: GPU forward pass micro-benchmarks.
+2 -2
View File
@@ -6,10 +6,10 @@
| Lever | 어디서 nail되는지 | 현재 상태 |
| --- | --- | --- |
| **Model size growth** (hidden ≥ 1024 / layers ≥ 6) | `docs/plans/model_size_experiment.md` | 인프라 미설치 |
| **Model size growth** (hidden ≥ 1024 / layers ≥ 6) | `docs/plans/archive/model_size_experiment.md` | 인프라 미설치 |
| **Option B** (per-worker interleaved traversal) | plan 미작성 | 미시작 |
| **AMP** trainer | `docs/plans/archive/amp_trainer.md` (구현 됨, default off) | 모델 키운 후 재측정 |
| **torch.compile** trainer | `docs/plans/torch_compile.md` | 모델 키운 후 재측정 |
| **torch.compile** trainer | `docs/plans/archive/torch_compile.md` | 모델 키운 후 재측정 |
| **TensorRT** inference | plan 미작성 | 모델 키운 + eval dense 시점 |
| **Option A re-enable** | 코드 있음 (default off) | 모델 키운 후 또는 Option B 후 |
| **Julia port** | `docs/research/julia_port_evaluation.md` | Torch.jl 결과 대기 중. Flux FAIL. |
+2 -2
View File
@@ -19,7 +19,7 @@ with torch.inference_mode():
advantages = networks[player](x).squeeze(0).detach().cpu().numpy().astype(np.float32)
```
When `traversal.inference_backend` is set to `server` in `configs/deep_cfr/default.yaml`, the `networks[player]` call is intercepted by a `NetworkProxy` (instantiated in `workers.py`, around line 91). This proxy posts a request to the `InferenceServer` and blocks until a response is received via a per-slot event.
When `traversal.inference_backend` is set to `server` in `legacy/deep-cfr/configs/default.yaml`, the `networks[player]` call is intercepted by a `NetworkProxy` (instantiated in `workers.py`, around line 91). This proxy posts a request to the `InferenceServer` and blocks until a response is received via a per-slot event.
The server's batching logic in `src/coolrl_lost_cities/games/classic/deep_cfr/inference_server.py` (around line 221) reports the realized batch size:
@@ -46,7 +46,7 @@ The "structural ceiling" is that `batch_window_us` and `max_batch` tuning cannot
## Practical implication
Option A is deferred for the current small MLP models (512x3). The `local` backend remains the default in `configs/deep_cfr/default.yaml`.
Option A is deferred for the current small MLP models (512x3). The `local` backend remains the default in `legacy/deep-cfr/configs/default.yaml`.
To unlock the projected GPU gains, the traversal must be restructured to drive batch sizes up. This leads to two primary paths:
1. **Option B (Interleaved Traversal):** Refactor the Cython traversal into a state machine that can advance multiple traversals concurrently per worker. Each worker would suspend at a policy call, batch its own requests, and resume continuations once the results return.
@@ -9,7 +9,7 @@ Why do architectural optimizations like `torch.compile` and TensorRT integration
## Code reference
The current baseline configuration is defined in `configs/deep_cfr/default.yaml`:
The historical baseline configuration is defined in `legacy/deep-cfr/configs/default.yaml`:
```yaml
network:
@@ -51,4 +51,4 @@ Avoid premature optimization with `torch.compile` or TensorRT on the current sma
- `docs/archive/deep-cfr-performance-experiments-2026-05-07.md` (Small-model regression data)
- `docs/archive/option-a-bench-result-2026-05-07.md` (Batched traversal benchmarks)
- `docs/performance.md` (Current runtime bottleneck profile)
- `docs/performance.md` (Current runtime bottleneck profile)
+1 -1
View File
@@ -46,7 +46,7 @@ sampling-mode branch. The `info_state` is computed by `_policy(state,
player, ...)` for the *current acting player*, which is the right thing in
both conventions.
`configs/deep_cfr/default.yaml` sets:
`legacy/deep-cfr/configs/default.yaml` sets:
```yaml
store_strategy_on_traverser_nodes: true