SLM-Lab/benchmark

1# Phase 5: Brax PPO vs SLM-Lab PPO — Comprehensive Comparison2 3Source: `google/brax` (latest `main`) and `google-deepmind/mujoco_playground` (latest `main`).4All values extracted from actual code, not documentation.5 6---7 8## 1. Batch Collection Mechanics9 10### Brax11The training loop in `brax/training/agents/ppo/train.py` (line 586–591) collects data via nested `jax.lax.scan`:12 13```python14(state, _), data = jax.lax.scan(15    f, (state, key_generate_unroll), (),16    length=batch_size * num_minibatches // num_envs,17)18```19 20Each inner call does `generate_unroll(env, state, policy, key, unroll_length)` — a `jax.lax.scan` of `unroll_length` sequential env steps. The outer scan repeats this `batch_size * num_minibatches // num_envs` times **sequentially**, rolling the env state forward continuously.21 22**DM Control default**: `num_envs=2048`, `batch_size=1024`, `num_minibatches=32`, `unroll_length=30`.23- Outer scan length = `1024 * 32 / 2048 = 16` sequential unrolls.24- Each unroll = 30 steps.25- Total data per training step = 16 * 2048 * 30 = **983,040 transitions** reshaped to `(32768, 30)`.26- Then `num_updates_per_batch=16` SGD passes, each splitting into 32 minibatches.27- **Effective gradient steps per collect**: 16 * 32 = 512.28 29### SLM-Lab30`time_horizon=30`, `num_envs=2048` → collects `30 * 2048 = 61,440` transitions.31`training_epoch=16`, `minibatch_size=4096` → 15 minibatches per epoch → 16 * 15 = 240 gradient steps.32 33### Difference34**Brax collects 16x more data per training step** by doing 16 sequential unrolls before updating. SLM-Lab does 1 unroll. This means Brax's advantages are computed over much longer trajectories (480 steps vs 30 steps), providing much better value bootstrap targets.35 36Brax also shuffles the entire 983K-transition dataset into minibatches, enabling better gradient estimates.37 38**Classification: CRITICAL**39 40**Fix**: Increase `time_horizon` or implement multi-unroll collection. The simplest fix: increase `time_horizon` from 30 to 480 (= 30 * 16). This gives the same data-per-update ratio. However, this would require more memory. Alternative: keep `time_horizon=30` but change `training_epoch` to 1 and let the loop collect multiple horizons before training — requires architectural changes.41 42**Simplest spec-only fix**: Set `time_horizon=480` (or even 256 as a compromise). This is safe because GAE with `lam=0.95` naturally discounts old data. Risk: memory usage increases 16x for the batch buffer.43 44---45 46## 2. Reward Scaling47 48### Brax49`reward_scaling` is applied **inside the loss function** (`losses.py` line 212):50```python51rewards = data.reward * reward_scaling52```53This scales rewards just before GAE computation. It does NOT modify the environment rewards.54 55**DM Control default**: `reward_scaling=10.0`56**Locomotion default**: `reward_scaling=1.0`57**Manipulation default**: `reward_scaling=1.0` (except PandaPickCubeCartesian: 0.1)58 59### SLM-Lab60`reward_scale` is applied in the **environment wrapper** (`playground.py` line 149):61```python62rewards = np.asarray(self._state.reward) * self._reward_scale63```64 65**Current spec**: `reward_scale: 10.0` (DM Control)66 67### Difference68Functionally equivalent — both multiply rewards by a constant before GAE. The location (env vs loss) shouldn't matter for PPO since rewards are only used in GAE computation.69 70**Classification: MINOR** — Already matching for DM Control.71 72---73 74## 3. Observation Normalization75 76### Brax77Uses Welford's online algorithm to track per-feature running mean/std. Applied via `running_statistics.normalize()`:78```python79data = (data - mean) / std80```81Mean-centered AND divided by std. Updated **every training step** before SGD (line 614).82`normalize_observations=True` for all environments.83`std_eps=0.0` (default, no epsilon in std).84 85### SLM-Lab86Uses gymnasium's `VectorNormalizeObservation` (CPU) or `TorchNormalizeObservation` (GPU), which also uses Welford's algorithm with mean-centering and std division.87 88**Current spec**: `normalize_obs: true`89 90### Difference91Both use mean-centered running normalization. Brax updates normalizer params inside the training loop (not during rollout), while SLM-Lab updates during rollout (gymnasium wrapper). This is a subtle timing difference but functionally equivalent.92 93Brax uses `std_eps=0.0` by default, while gymnasium uses `epsilon=1e-8`. Minor numerical difference.94 95**Classification: MINOR** — Already matching.96 97---98 99## 4. Value Function100 101### Brax102- **Loss**: Unclipped MSE by default (`losses.py` line 252–263):103  ```python104  v_error = vs - baseline105  v_loss = jnp.mean(v_error * v_error) * 0.5 * vf_coefficient106  ```107- **vf_coefficient**: 0.5 (default in `train.py`)108- **Value clipping**: Only if `clipping_epsilon_value` is set (default `None` = no clipping)109- **No value target normalization** — raw GAE targets110- **Separate policy and value networks** (always separate in Brax's architecture)111- Value network: 5 hidden layers of 256 (DM Control default) with `swish` activation112- **Bootstrap on timeout**: Optional, default `False`113 114### SLM-Lab115- **Loss**: MSE with `val_loss_coef=0.5`116- **Value clipping**: Optional via `clip_vloss` (default False)117- **Value target normalization**: Optional via `normalize_v_targets: true` using `ReturnNormalizer`118- **Architecture**: `[256, 256, 256]` with SiLU (3 layers vs Brax's 5)119 120### Difference1211. **Value network depth**: Brax uses **5 layers of 256** for DM Control, SLM-Lab uses **3 layers of 256**. This is a meaningful capacity difference for the value function, which needs to accurately estimate returns.122 1232. **Value target normalization**: SLM-Lab has `normalize_v_targets: true` with a `ReturnNormalizer`. Brax does NOT normalize value targets. This could cause issues if the normalizer is poorly calibrated.124 1253. **Value network architecture (Loco)**: Brax uses `[256, 256, 256, 256, 256]` for loco too.126 127**Classification: IMPORTANT**128 129**Fix**:130- Consider increasing value network to 5 layers: `[256, 256, 256, 256, 256]` to match Brax.131- Consider disabling `normalize_v_targets` since Brax doesn't use it and `reward_scaling=10.0` already provides good gradient magnitudes.132- Risk of regressing: the return normalizer may be helping envs with high reward variance. Test with and without.133 134---135 136## 5. Advantage Computation (GAE)137 138### Brax139`compute_gae` in `losses.py` (line 38–100):140- Standard GAE with `lambda_=0.95`, `discount=0.995` (DM Control)141- Computed over each unroll of `unroll_length` timesteps142- Uses `truncation` mask to handle episode boundaries within an unroll143- `normalize_advantage=True` (default): `advs = (advs - mean) / (std + 1e-8)` over the **entire batch**144- GAE is computed **inside the loss function**, once per SGD pass (recomputed each time with current value estimates? No — computed once with data from rollout, including stored baseline values)145 146### SLM-Lab147- GAE computed in `calc_gae_advs_v_targets` using `math_util.calc_gaes`148- Computed once before training epochs149- Advantage normalization: per-minibatch standardization in `calc_policy_loss`:150  ```python151  advs = math_util.standardize(advs)  # per minibatch152  ```153 154### Difference1551. **GAE horizon**: Brax computes GAE over 30-step unrolls. SLM-Lab also uses 30-step horizon. **Match**.1562. **Advantage normalization scope**: Brax normalizes over the **entire batch** (983K transitions). SLM-Lab normalizes **per minibatch** (4096 transitions). Per-minibatch normalization has more variance. However, both approaches are standard — SB3 also normalizes per-minibatch.1573. **Truncation handling**: Brax explicitly handles truncation with `truncation_mask` in GAE. SLM-Lab uses `terminateds` from the env wrapper, with truncation handled by gymnasium's auto-reset. These should be functionally equivalent.158 159**Classification: MINOR** — Approaches are different but both standard.160 161---162 163## 6. Learning Rate Schedule164 165### Brax166Default: `learning_rate_schedule=None` → **no schedule** (constant LR).167Optional: `ADAPTIVE_KL` schedule that adjusts LR based on KL divergence.168Base LR: `1e-3` (DM Control), `3e-4` (Locomotion).169 170### SLM-Lab171Uses `LinearToMin` scheduler:172```yaml173lr_scheduler_spec:174  name: LinearToMin175  frame: "${max_frame}"176  min_factor: 0.033177```178This linearly decays LR from `1e-3` to `1e-3 * 0.033 = 3.3e-5` over training.179 180### Difference181**Brax uses constant LR. SLM-Lab decays LR by 30x over training.** This is a significant difference. Linear LR decay can help convergence in the final phase but can also hurt by reducing the LR too early for long training runs.182 183**Classification: IMPORTANT**184 185**Fix**: Consider removing or weakening the LR decay for playground envs:186- Option A: Set `min_factor: 1.0` (effectively constant LR) to match Brax187- Option B: Use a much gentler decay, e.g. `min_factor: 0.1` (10x instead of 30x)188- Risk: Some envs may benefit from the decay. Test both.189 190---191 192## 7. Entropy Coefficient193 194### Brax195**Fixed** (no decay):196- DM Control: `entropy_cost=1e-2`197- Locomotion: `entropy_cost=1e-2` (some overrides to `5e-3`)198- Manipulation: varies, typically `1e-2` or `2e-2`199 200### SLM-Lab201**Fixed** (no_decay):202```yaml203entropy_coef_spec:204  name: no_decay205  start_val: 0.01206```207 208### Difference209**Match**: Both use fixed `0.01`.210 211**Classification: MINOR** — Already matching.212 213---214 215## 8. Gradient Clipping216 217### Brax218`max_grad_norm` via `optax.clip_by_global_norm()`:219- DM Control default: **None** (no clipping!)220- Locomotion default: `1.0`221- Vision PPO and some manipulation: `1.0`222 223### SLM-Lab224`clip_grad_val: 1.0` — always clips gradients by global norm.225 226### Difference227**Brax does NOT clip gradients for DM Control by default.** SLM-Lab always clips at 1.0.228 229Gradient clipping can be overly conservative, preventing the optimizer from taking large useful steps when gradients are naturally large (e.g., early training with `reward_scaling=10.0`).230 231**Classification: IMPORTANT** — Could explain slow convergence on DM Control envs.232 233**Fix**: Remove gradient clipping for DM Control playground spec:234```yaml235clip_grad_val: null  # match Brax DM Control default236```237Keep `clip_grad_val: 1.0` for locomotion spec. Risk: gradient explosions without clipping, but Brax demonstrates it works for DM Control.238 239---240 241## 9. Action Distribution242 243### Brax244Default: `NormalTanhDistribution` — samples from `Normal(loc, scale)` then applies `tanh` postprocessing.245- `param_size = 2 * action_size` (network outputs both mean and log_scale)246- Scale: `scale = (softplus(raw_scale) + 0.001) * 1.0` (min_std=0.001, var_scale=1)247- **State-dependent std**: The scale is output by the policy network (not a separate parameter)248- Uses `tanh` bijector with log-det-jacobian correction249 250### SLM-Lab251Default: `Normal(loc, scale)` without tanh.252- `log_std_init` creates a **state-independent** `nn.Parameter` for log_std253- Scale: `scale = clamp(log_std, -5, 0.5).exp()` → std range [0.0067, 1.648]254- **State-independent std** (when `log_std_init` is set)255 256### Difference2571. **Tanh squashing**: Brax applies `tanh` to bound actions to [-1, 1]. SLM-Lab does NOT. This is a fundamental architectural difference:258   - With tanh: actions are bounded, log-prob includes jacobian correction259   - Without tanh: actions can exceed env bounds, relying on env clipping260 2612. **State-dependent vs independent std**: Brax uses state-dependent std (network outputs it), SLM-Lab uses state-independent learnable parameter.262 2633. **Std parameterization**: Brax uses `softplus + 0.001` (min_std=0.001), SLM-Lab uses `clamp(log_std, -5, 0.5).exp()` with max std of 1.648.264 2654. **Max std cap**: SLM-Lab caps at exp(0.5)=1.648. Brax has no explicit cap (softplus can grow unbounded). However, Brax's `tanh` squashing means even large std doesn't produce out-of-range actions.266 267**Classification: IMPORTANT**268 269**Note**: For MuJoCo Playground where actions are already in [-1, 1] and the env wrapper has `PlaygroundVecEnv` with action space `Box(-1, 1)`, the `tanh` squashing may not be critical since the env naturally clips. But the log-prob correction matters for policy gradient quality.270 271**Fix**:272- The state-independent log_std is a reasonable simplification (CleanRL also uses it). Keep.273- The `max=0.5` clamp may be too restrictive. Consider increasing to `max=2.0` (CleanRL default) or removing the upper clamp entirely.274- Consider implementing tanh squashing as an option for playground envs.275 276---277 278## 10. Network Initialization279 280### Brax281Default: `lecun_uniform` for all layers (policy and value).282Activation: `swish` (= SiLU).283No special output layer initialization by default.284 285### SLM-Lab286Default: `orthogonal_` initialization.287Activation: SiLU (same as swish).288 289### Difference290- Brax uses `lecun_uniform`, SLM-Lab uses `orthogonal_`. Both are reasonable for swish/SiLU activations.291- `orthogonal_` tends to preserve gradient magnitudes across layers, which can be beneficial for deeper networks.292 293**Classification: MINOR** — Both are standard choices. `orthogonal_` may actually be slightly better for the 3-layer SLM-Lab network.294 295---296 297## 11. Network Architecture298 299### Brax (DM Control defaults)300- **Policy**: `(32, 32, 32, 32)` — 4 layers of 32, swish activation301- **Value**: `(256, 256, 256, 256, 256)` — 5 layers of 256, swish activation302 303### Brax (Locomotion defaults)304- **Policy**: `(128, 128, 128, 128)` — 4 layers of 128305- **Value**: `(256, 256, 256, 256, 256)` — 5 layers of 256306 307### SLM-Lab (ppo_playground)308- **Policy**: `(64, 64)` — 2 layers of 64, SiLU309- **Value**: `(256, 256, 256)` — 3 layers of 256, SiLU310 311### Difference3121. **Policy width**: SLM-Lab uses wider layers (64) but fewer (2 vs 4). Total params: ~similar for DM Control (4*32*32=4096 vs 2*64*64=8192). SLM-Lab's policy is actually larger per layer but shallower.313 3142. **Value depth**: 3 vs 5 layers. This is significant — the value function benefits from more depth to accurately represent complex return landscapes, especially for long-horizon tasks.315 3163. **DM Control policy**: Brax uses very small 32-wide networks. SLM-Lab's 64-wide may be slightly over-parameterized but shouldn't hurt.317 318**Classification: IMPORTANT** (mainly the value network depth)319 320**Fix**: Consider increasing value network to 5 layers to match Brax:321```yaml322_value_body: &value_body323  modules:324    body:325      Sequential:326        - LazyLinear: {out_features: 256}327        - SiLU:328        - LazyLinear: {out_features: 256}329        - SiLU:330        - LazyLinear: {out_features: 256}331        - SiLU:332        - LazyLinear: {out_features: 256}333        - SiLU:334        - LazyLinear: {out_features: 256}335        - SiLU:336```337 338---339 340## 12. Clipping Epsilon341 342### Brax343Default: `clipping_epsilon=0.3` (in `train.py` line 206).344DM Control: not overridden → **0.3**.345Locomotion: some envs override to `0.2`.346 347### SLM-Lab348Default: `clip_eps=0.2` (in spec).349 350### Difference351Brax uses **0.3** while SLM-Lab uses **0.2**. This is notable — 0.3 allows larger policy updates per step, which can accelerate learning but risks instability. Given that Brax collects 16x more data per update (see #1), the larger clip epsilon is safe because the policy ratio variance is lower with more data.352 353**Classification: IMPORTANT** — Especially in combination with the batch size difference (#1).354 355**Fix**: Consider increasing to 0.3 for DM Control playground spec. However, this should only be done together with the batch size fix (#1), since larger clip epsilon with small batches risks instability.356 357---358 359## 13. Discount Factor360 361### Brax (DM Control)362Default: `discounting=0.995`363Overrides: BallInCup=0.95, FingerSpin=0.95364 365### Brax (Locomotion)366Default: `discounting=0.97`367Overrides: Go1Backflip=0.95368 369### SLM-Lab370DM Control: `gamma=0.995`371Locomotion: `gamma=0.97`372Overrides: FingerSpin=0.95373 374### Difference375**Match** for the main categories.376 377**Classification: MINOR** — Already matching.378 379---380 381## Summary: Priority-Ordered Fixes382 383### CRITICAL384 385| # | Issue | Brax Value | SLM-Lab Value | Fix |386|---|-------|-----------|--------------|-----|387| 1 | **Batch size (data per training step)** | 983K transitions (16 unrolls of 30) | 61K transitions (1 unroll of 30) | Increase `time_horizon` to 480, or implement multi-unroll collection |388 389### IMPORTANT390 391| # | Issue | Brax Value | SLM-Lab Value | Fix |392|---|-------|-----------|--------------|-----|393| 4 | **Value network depth** | 5 layers of 256 | 3 layers of 256 | Add 2 more hidden layers |394| 6 | **LR schedule** | Constant | Linear decay to 0.033x | Set `min_factor: 1.0` or weaken to 0.1 |395| 8 | **Gradient clipping (DM Control)** | None | 1.0 | Set `clip_grad_val: null` for DM Control |396| 9 | **Action std upper bound** | Softplus (unbounded) | exp(0.5)=1.65 | Increase max clamp from 0.5 to 2.0 |397| 11 | **Clipping epsilon** | 0.3 | 0.2 | Increase to 0.3 (only with larger batch) |398 399### MINOR (already matching or small effect)400 401| # | Issue | Status |402|---|-------|--------|403| 2 | Reward scaling | Match (10.0 for DM Control) |404| 3 | Obs normalization | Match (Welford running stats) |405| 5 | GAE computation | Match (lam=0.95, per-minibatch normalization) |406| 7 | Entropy coefficient | Match (0.01, fixed) |407| 10 | Network init | Minor difference (orthogonal vs lecun_uniform) |408| 13 | Discount factor | Match |409 410---411 412## Recommended Implementation Order413 414### Phase 1: Low-risk spec changes (test on CartpoleBalance/Swingup first)4151. Remove gradient clipping for DM Control: `clip_grad_val: null`4162. Weaken LR decay: `min_factor: 0.1` (or `1.0` for constant)4173. Increase log_std clamp from 0.5 to 2.0418 419### Phase 2: Architecture changes (test on several envs)4204. Increase value network to 5 layers of 2564215. Consider disabling `normalize_v_targets` since Brax doesn't use it422 423### Phase 3: Batch size alignment (largest expected impact, highest risk)4246. Increase `time_horizon` to 240 or 480 to match Brax's effective batch size4257. If time_horizon increase works, consider increasing `clipping_epsilon` to 0.3426 427### Risk Assessment428- **Safest changes**: #1 (no grad clip), #2 (weaker LR decay), #3 (wider std range)429- **Medium risk**: #4 (deeper value net — more compute, could slow training), #5 (removing normalization)430- **Highest risk/reward**: #6 (larger time_horizon — 16x more memory, biggest expected improvement)431 432### Envs Already Solved433Changes should be tested against already-solved envs (CartpoleBalance, CartpoleSwingup, etc.) to ensure no regression. The safest approach is to implement spec variants rather than modifying the default spec.434 435---436 437## Key Insight438 439The single largest difference is **data collection volume per training step**. Brax collects 16x more transitions before each update cycle. This provides:4401. Better advantage estimates (longer trajectory context)4412. More diverse minibatches (less overfitting per update)4423. Safety for larger clip epsilon and no gradient clipping443 444Without matching this, the other improvements will have diminished returns. The multi-unroll collection in Brax is fundamentally tied to its JAX/vectorized architecture — SLM-Lab's sequential PyTorch loop can approximate this by simply increasing `time_horizon`, at the cost of memory.445 446A practical compromise: increase `time_horizon` from 30 to 128 or 256 (4-8x, not full 16x) and adjust other hyperparameters accordingly.447