From 801346e18cea12999f095be5b1caa9c5a6842607 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Mon, 27 Jul 2026 14:08:11 +0200 Subject: [PATCH] fix(scripts): restore policy training mode after eval_policy() in lerobot-eval (#4162) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(scripts/eval): restore policy training mode after eval_policy() `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. * fix(scripts): keep policy train/eval --------- Co-authored-by: ModeEric --- src/lerobot/scripts/lerobot_eval.py | 69 +++++++++++++++++------------ 1 file changed, 41 insertions(+), 28 deletions(-) diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 722763d6e..c4ed35145 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -453,6 +453,9 @@ def eval_policy( raise exc from None start = time.time() + # Preserve the mode for direct callers. eval_policy_all scopes the mode + # around all tasks so parallel evaluations cannot race with each other. + was_training = policy.training policy.eval() # Determine how many batched rollouts we need to get n_episodes. Note that if n_episodes is not evenly @@ -674,6 +677,8 @@ def eval_policy( if save_predicted_video: info["predicted_video_paths"] = predicted_video_paths + policy.train(was_training) + return info @@ -1010,40 +1015,48 @@ def eval_policy_all( recording_private=recording_private, ) - if max_parallel_tasks <= 1: - prefetch_thread: threading.Thread | None = None - for i, (task_group, task_id, env) in enumerate(tasks): - if prefetch_thread is not None: - prefetch_thread.join() - prefetch_thread = None + # Set the shared policy's mode before launching any workers. Restoring it + # inside individual tasks would let one task enable training mode while + # another task is still evaluating. + was_training = policy.training + policy.eval() + try: + if max_parallel_tasks <= 1: + prefetch_thread: threading.Thread | None = None + for i, (task_group, task_id, env) in enumerate(tasks): + if prefetch_thread is not None: + prefetch_thread.join() + prefetch_thread = None - try: - tg, tid, metrics = task_runner(task_group, task_id, env) - _accumulate_to(tg, metrics) - per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics}) - finally: - env.close() - # Prefetch next task's workers *after* closing current env to prevent - # GPU memory overlap between consecutive tasks. - if i + 1 < len(tasks): - next_env = tasks[i + 1][2] - if hasattr(next_env, "_ensure"): - prefetch_thread = threading.Thread(target=next_env._ensure, daemon=True) - prefetch_thread.start() - else: - with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor: - fut2meta = {} - for task_group, task_id, env in tasks: - fut = executor.submit(task_runner, task_group, task_id, env) - fut2meta[fut] = (task_group, task_id, env) - for fut in cf.as_completed(fut2meta): - tg, tid, env = fut2meta[fut] try: - tg, tid, metrics = fut.result() + tg, tid, metrics = task_runner(task_group, task_id, env) _accumulate_to(tg, metrics) per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics}) finally: env.close() + # Prefetch next task's workers *after* closing current env to prevent + # GPU memory overlap between consecutive tasks. + if i + 1 < len(tasks): + next_env = tasks[i + 1][2] + if hasattr(next_env, "_ensure"): + prefetch_thread = threading.Thread(target=next_env._ensure, daemon=True) + prefetch_thread.start() + else: + with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor: + fut2meta = {} + for task_group, task_id, env in tasks: + fut = executor.submit(task_runner, task_group, task_id, env) + fut2meta[fut] = (task_group, task_id, env) + for fut in cf.as_completed(fut2meta): + tg, tid, env = fut2meta[fut] + try: + tg, tid, metrics = fut.result() + _accumulate_to(tg, metrics) + per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics}) + finally: + env.close() + finally: + policy.train(was_training) # compute aggregated metrics helper (robust to lists/scalars) def _agg_from_list(xs):