mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-27 11:46:04 +00:00
ModeEric
545522d3c1
`eval_policy` calls `policy.eval()` before the rollout but never restores
the prior mode on return. When called from the training loop
(`lerobot_train.py`'s `eval_policy_all -> run_one -> eval_one ->
eval_policy` chain), the policy is left in eval mode for every subsequent
training step, which silently:
* disables Dropout (no regularisation),
* freezes BatchNorm running stats (no further EMA updates).
Under DDP only `is_main_process` runs eval (lerobot_train.py:527), so the
main rank ends up in eval mode while workers stay in train mode — the
all-reduced gradients then combine forward passes computed with different
dropout masks and different BN behaviour, a real DDP-correctness issue.
Scope of impact:
* Affects every policy with Dropout in its forward path. In-tree, that
includes the default ACT (6 Dropout layers at p=0.1), Diffusion (vision
backbone), VQ-BeT, Multi-Task DiT, X-VLA, plus all VLA policies that
inherit Dropout from their pretrained HF backbone (PI0/PI0.5/PI0-FAST,
SmolVLA, GR00T-N1.5, EO1, Wall-X).
* Triggers from the first eval onward. On the default config
(steps=100k, eval_freq=20k) that's 80% of training; on the LIBERO /
RoboCasa / VLABench example commands in docs/ (eval_freq=1k–5k)
it's 95–99% of training.
* Policies using only LayerNorm/GroupNorm and no Dropout (TDMPC, RTC)
are unaffected. Policies using `FrozenBatchNorm2d` (ACT's ResNet
backbone) are immune to the BN-stat half; the Dropout half still bites.
Fix:
* Snapshot `policy.training` on entry to `eval_policy`.
* Restore it on normal return.
* Save-and-restore is a strict no-op for callers that pass an
already-eval-mode policy (e.g. the standalone `lerobot-eval` script
loading a frozen checkpoint).
* Restoration is placed before the normal return only, not in a
try/finally — exception paths leave the policy in eval mode, same as
today. A try/finally upgrade would require re-indenting ~165 lines and
can land as a separate cleanup if desired.
Tests (tests/scripts/test_eval.py, 7 tests total, ~1.6s):
* Regression gates on the lerobot_eval fix itself: training-mode
preservation, eval-mode preservation, dropout-active behavioural
check, non-crash for both entry modes.
* Quantitative mechanism demonstration
(`test_missing_mode_restoration_hurts_generalisation`): trains a tiny
Dropout+BatchNorm MLP under both the bug pattern and the fix pattern
on identical data and seed, then asserts the buggy variant generalises
at least 5% worse on a held-out val set. In repeated runs we see
10-25% deltas on this toy problem; real policies (more layers, more
Dropout, longer training) generally see larger gaps. Lives alongside
the regression tests so the empirical proof is reproducible from the
repo without adding a separate benchmarks/ directory.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>