SLM-Lab/benchmark

1# SLM-Lab v5.3.02 3MuJoCo Playground integration. 54 GPU-accelerated environments via JAX/MJX backend.4 5**What changed:**6- **New env backend**: MuJoCo Playground (DeepMind) — 25 DM Control Suite, 19 Locomotion (Go1, Spot, H1, G1), 10 Manipulation (Panda, ALOHA, LEAP)7- **PlaygroundVecEnv**: JAX-native vectorized env wrapper with `jax.vmap` batching and Brax auto-reset. Converts JAX arrays to numpy at the API boundary for PyTorch compatibility8- **Prefix routing**: `playground/EnvName` in specs routes to PlaygroundVecEnv instead of Gymnasium9- **Optional dependency**: `uv sync --group playground` installs `mujoco-playground`, `jax`, `brax`10- **Benchmark specs**: `slm_lab/spec/benchmark/playground/` — SAC specs for all 54 envs across 3 categories11 12<!-- TODO: Add benchmark results from DM Control Suite baseline runs (task #11) -->13 14---15 16# SLM-Lab v5.2.017 18Training path performance optimization. **+15% SAC throughput on GPU**, verified with no score regression.19 20**What changed (18 files):**21- `polyak_update`: in-place `lerp_()` replaces 3-op manual arithmetic22- `SAC`: single `log_softmax→exp` replaces dual softmax+log_softmax; cached entropy between policy/alpha loss; cached `_is_per` and `_LOG2`23- `to_torch_batch`: uint8/float16 sent directly to GPU then `.float()` — avoids 4x CPU float32 intermediate (matters for Atari 84x84x4)24- `SumTree`: iterative propagation/retrieval replaces recursion; vectorized sampling25- `forward_tails`: cached output (was called twice per step)26- `VectorFullGameStatistics`: `deque(maxlen=N)` + `np.flatnonzero` replaces list+pop(0)+loop27- `pydash→builtins`: `isinstance` over `ps.is_list/is_dict`, dict comprehensions over `ps.pick/ps.omit` in hot paths28- `PPO`: `total_loss` as plain float prevents computation graph leak across epochs29- Minor: `hasattr→is not None` in conv/recurrent forward, cached `_is_dev`, `no_decay` early exit in VarScheduler30 31**Measured gains (normalized, same hardware A/B on RTX 3090):**32- SAC MuJoCo: +15-17% fps33- SAC Atari: +14% fps34- PPO: ~0% (env-bound; most optimizations target SAC's training-heavy inner loop — PPO doesn't use polyak, replay buffer, twin Q, or entropy tuning)35 36---37 38# SLM-Lab v5.1.039 40TorchArc YAML benchmarks replace original hardcoded network architectures across all benchmark categories.41 42- **TorchArc integration**: All algorithms (REINFORCE, SARSA, DQN, DDQN+PER, A2C, PPO, SAC) now use TorchArc YAML-defined networks instead of hardcoded PyTorch modules43- **Full benchmark validation**: Classic Control, Box2D, MuJoCo (11 envs), and Atari (54 games) re-benchmarked with TorchArc — results match or exceed original scores44- **SAC Atari**: New SAC Atari benchmarks (48 games) with discrete action support45- **Pre-commit hooks**: Conventional commit message validation via `.githooks/commit-msg`46 47---48 49# SLM-Lab v5.0.050 51Modernization release for the current RL ecosystem. Updates SLM-Lab from OpenAI Gym to Gymnasium, adds correct handling of episode termination (the `terminated`/`truncated` fix), and migrates to modern Python tooling.52 53**TL;DR:** Install with `uv sync`, run with `slm-lab run`. Specs are simpler (no more `body` section or array wrappers). Environment names changed (`CartPole-v1`, `ALE/Pong-v5`, `Hopper-v5`). Code structure preserved for book readers.54 55> **Book readers:** For exact code from *Foundations of Deep Reinforcement Learning*, use `git checkout v4.1.1`56 57---58 59## Why This Release60 61SLM-Lab was created as an educational framework for deep reinforcement learning, accompanying *Foundations of Deep Reinforcement Learning*. The code prioritizes clarity and correctness—it should help you understand RL algorithms, not just run them.62 63Since v4, the RL ecosystem changed significantly:64 65- **OpenAI Gym is deprecated.** The Farama Foundation forked it as [Gymnasium](https://gymnasium.farama.org/), now the standard. Gym's `done` flag conflated two concepts: true termination (agent failed/succeeded) and time-limit truncation. Gymnasium fixes this with separate `terminated` and `truncated` signals—important for correct value estimation (see [below](#the-gymnasium-api-change)).66 67- **Roboschool is abandoned.** MuJoCo became free in 2022, so roboschool is no longer maintained. Gymnasium includes native MuJoCo bindings.68 69- **Python tooling modernized.** `conda` + `setup.py` → `uv` + `pyproject.toml`. Python 3.12+, PyTorch 2.8+. [uv](https://docs.astral.sh/uv/) emerged as a fast, reliable Python package manager—no more conda environment headaches.70 71- **Old dependencies don't build anymore.** The v4 dependency stack (old PyTorch, atari-py, mujoco-py, etc.) won't compile on modern hardware, especially ARM machines (Apple Silicon, AWS Graviton). Many deprecated packages simply don't run. A full rebuild was necessary.72 73This release updates SLM-Lab to work with modern dependencies while preserving the educational code structure. If you've read the book, the code should still be recognizable.74 75### Critical: Atari v5 Sticky Actions76 77**SLM-Lab uses Gymnasium ALE v5 defaults.** v5 default `repeat_action_probability=0.25` (sticky actions) randomly repeats agent actions to simulate console stochasticity, making evaluation harder but more realistic than v4 default 0.0 used by most benchmarks (CleanRL, SB3, RL Zoo). This follows [Machado et al. (2018)](https://arxiv.org/abs/1709.06009) research best practices. See [ALE version history](https://ale.farama.org/environments/#version-history-and-naming-schemes).78 79### Summary80 81| v4 | v5 |82|----|----|83| `conda activate lab && python run_lab.py` | `slm-lab run` |84| `CartPole-v0`, `PongNoFrameskip-v4` | `CartPole-v1`, `ALE/Pong-v5` |85| `RoboschoolHopper-v1` | `Hopper-v5` |86| `agent: [{...}]`, `env: [{...}]`, `body: {...}` | `agent: {...}`, `env: {...}` |87| `body.state_dim`, `body.memory` | `agent.state_dim`, `agent.memory` |88 89---90 91## Migration from v492 93### 1. Install94 95```bash96uv sync97uv tool install --editable .98```99 100### 2. Update specs101 102Remove array brackets and `body` section:103 104```diff105 {106-  "agent": [{ "name": "PPO", ... }],107-  "env": [{ "name": "CartPole-v0", ... }],108-  "body": { "product": "outer", "num": 1 },109+  "agent": { "name": "PPO", ... },110+  "env": { "name": "CartPole-v1", ... },111   "meta": { ... }112 }113```114 115### 3. Update environment names116 117- Classic control: `v0`/`v1` → current version (`CartPole-v1`, `Pendulum-v1`, `LunarLander-v3`)118- Atari: `PongNoFrameskip-v4` → `ALE/Pong-v5`119- Roboschool → MuJoCo: see [Deprecations](#roboschool) for full mapping120 121### 4. Run122 123```bash124slm-lab run spec.json spec_name train125```126 127See `slm_lab/spec/benchmark/` for updated reference specs.128 129---130 131## The Gymnasium API Change132 133This matters for understanding the code, not just running it.134 135### The Problem136 137Gym's `done` flag was ambiguous—it meant "episode ended" but episodes end for two different reasons:138 1391. **Terminated:** True end state (CartPole fell, agent died, goal reached)1402. **Truncated:** Time limit hit (MuJoCo's 1000-step cap)141 142For value estimation, these need different treatment. Terminated means future returns are zero. Truncated means future returns exist but weren't observed—you should bootstrap from V(s').143 144### The Fix145 146Gymnasium separates the signals:147 148```python149# Gym150obs, reward, done, info = env.step(action)151 152# Gymnasium153obs, reward, terminated, truncated, info = env.step(action)154```155 156All SLM-Lab algorithms now use `terminated` for bootstrapping decisions:157 158```python159# Only zero out future returns on TRUE termination160q_targets = rewards + gamma * (1 - terminateds) * next_q_preds161```162 163This is why the code stores `terminateds` and `truncateds` separately in memory—algorithms need `terminated` for correct bootstrapping, `done` for episode boundaries.164 165This fix particularly matters for time-limited environments like MuJoCo (1000-step limit) where episodes frequently truncate during training. Using `done` instead of `terminated` there significantly hurts learning.166 167---168 169## Code Structure Changes170 171For book readers who want to trace through the code:172 173### Simplified Agent Design174 175The `Body` class was removed. Its responsibilities moved to more natural locations:176 177```python178# v4179state_dim = agent.body.state_dim180memory = agent.body.memory181env = agent.body.env182 183# v5184state_dim = agent.state_dim185memory = agent.memory186env = agent.env187```188 189Training metrics tracking is now in `MetricsTracker` (what `Body` was renamed to).190 191### Simplified Specs192 193Multi-agent configurations were rarely used. Specs are now flat:194 195```python196# v4: agent_spec = spec['agent'][0]197# v5: agent_spec = spec['agent']198```199 200### Architecture Preserved201 202The core design is unchanged:203 204```205Session → Agent → Algorithm → Network206              ↘ Memory207        → Env208```209 210---211 212## Algorithm Updates213 214**PPO:** New options for value target handling—`normalize_v_targets`, `symlog_transform` (from DreamerV3), `clip_vloss` (CleanRL-style).215 216**SAC:** Discrete action support uses exact expectation (Christodoulou 2019). Target entropy auto-calculated.217 218**Networks:** Optional `layer_norm` for MLP hidden layers. Custom optimizers (Lookahead, RAdam) removed—use native PyTorch `AdamW`.219 220All algorithms use `terminated` (not `done`) for correct bootstrapping.221 222---223 224## Benchmarks225 226All algorithms validated on Gymnasium. Full results in `docs/BENCHMARKS.md`.227 228| Category | REINFORCE | SARSA | DQN | DDQN+PER | A2C | PPO | SAC |229|----------|-----------|-------|-----|----------|-----|-----|-----|230| Classic Control | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |231| Box2D | — | — | ✅ | ✅ | ⚠️ | ✅ | ✅ |232| MuJoCo (11 envs) | — | — | — | — | ⚠️ | ✅ All | ✅ All |233| Atari (54 games) | — | — | — | — | ✅ | ✅ | — |234 235**Atari benchmarks** use ALE v5 with sticky actions (`repeat_action_probability=0.25`). PPO tested with lambda variants (0.95, 0.85, 0.70) to optimize per-game performance. A2C uses GAE with lambda 0.95.236 237**Note on scores:** Gymnasium environment versions differ from old Gym—some are harder (CartPole-v1 has stricter termination than v0), some have different reward scales (MuJoCo v5 vs roboschool). Targets reference [CleanRL](https://docs.cleanrl.dev/) and [Stable-Baselines3](https://stable-baselines3.readthedocs.io/) gymnasium benchmarks.238 239---240 241## New Features242 243**Hyperparameter search** now uses Ray Tune + Optuna + ASHA early stopping:244 245```bash246slm-lab run spec.json spec_name search    # Run search locally247```248 249Add `search_scheduler` to spec for ASHA early termination of poor trials. See `docs/BENCHMARKS.md` for search methodology.250 251---252 253## CLI Usage254 255The CLI uses [Typer](https://typer.tiangolo.com/). Use `--help` on any command for details:256 257```bash258slm-lab --help                           # List all commands259slm-lab run --help                       # Options for run command260 261# Installation262uv sync                                  # Install dependencies263uv tool install --editable .             # Install slm-lab command264 265# Basic usage266slm-lab run                              # PPO CartPole (default demo)267slm-lab run --render                     # With visualization268slm-lab run spec.json spec_name train    # Train from spec file269slm-lab run spec.json spec_name dev      # Dev mode (shorter run)270slm-lab run spec.json spec_name search   # Hyperparameter search271 272# Variable substitution (for template specs)273slm-lab run -s env=ALE/Breakout-v5 slm_lab/spec/benchmark/ppo/ppo_atari.json ppo_atari train274 275# Cloud training (dstack + HuggingFace)276slm-lab run-remote --gpu spec.json spec_name train   # Launch on cloud GPU277slm-lab list                                         # List experiments on HuggingFace278slm-lab pull spec_name                               # Download results locally279 280# Utilities281slm-lab run --stop-ray                   # Stop Ray processes282```283 284Modes: `dev` (quick test), `train` (full training), `search` (hyperparameter search), `enjoy` (evaluate saved model).285 286---287 288## Deprecations289 290### Multi-Agent / Multi-Environment291 292The v4 `body` spec section and array wrappers (`agent: [{...}]`) supported multi-agent and multi-environment configurations. These were rarely used and added complexity. v5 simplifies to single-agent single-env, which covers the vast majority of use cases and matches how most RL research is done.293 294### Unity ML-Agents and VizDoom295 296These integrations are removed from the core package. Both ecosystems have their own gymnasium-compatible wrappers now:297- Unity: [gymnasium-unity](https://gymnasium.farama.org/environments/third_party_environments/)298- VizDoom: [vizdoom gymnasium wrapper](https://gymnasium.farama.org/environments/third_party_environments/)299 300You can still use these environments with SLM-Lab by installing their wrappers and specifying the environment name in your spec.301 302### Roboschool303 304Roboschool is abandoned (MuJoCo became free in 2022). Use gymnasium's native MuJoCo environments instead:305- `RoboschoolHopper-v1` → `Hopper-v5`306- `RoboschoolHalfCheetah-v1` → `HalfCheetah-v5`307- `RoboschoolWalker2d-v1` → `Walker2d-v5`308- `RoboschoolAnt-v1` → `Ant-v5`309- `RoboschoolHumanoid-v1` → `Humanoid-v5`310