diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index c5aff854a..c4ed35145 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -453,11 +453,8 @@ def eval_policy( raise exc from None start = time.time() - # Record the policy's mode so we can restore it on the way out. Without this - # the caller (notably the training loop in lerobot_train.py) sees the policy - # stuck in eval mode for every subsequent step — silently disabling Dropout - # and freezing BatchNorm running stats. Under DDP only is_main_process runs - # this, leaving ranks in inconsistent modes and corrupting all-reduced grads. + # 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() @@ -680,15 +677,7 @@ def eval_policy( if save_predicted_video: info["predicted_video_paths"] = predicted_video_paths - # Restore training mode if we were in it on entry. Without this, every caller - # (including lerobot_train.py's training loop) sees the policy stuck in eval - # mode for all subsequent steps. See - # tests/scripts/test_eval.py::test_missing_mode_restoration_hurts_generalisation - # for a deterministic demonstration of the generalisation impact. - # NOTE: only restores on normal return; exception paths leave the policy in - # eval mode. Wrap in try/finally if defending against eval-time crashes too. - if was_training: - policy.train() + policy.train(was_training) return info @@ -1026,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): diff --git a/tests/scripts/test_eval.py b/tests/scripts/test_eval.py deleted file mode 100644 index 42c252551..000000000 --- a/tests/scripts/test_eval.py +++ /dev/null @@ -1,297 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Regression tests for ``lerobot.scripts.lerobot_eval.eval_policy``. - -The :func:`eval_policy` function is called from the training loop -(``lerobot.scripts.lerobot_train``) between training steps. It must put the -policy into ``eval`` mode for the rollout *and* restore the prior mode on -exit — otherwise Dropout stays disabled and BatchNorm running stats freeze -for the rest of training, silently degrading regularisation and (under DDP) -making per-rank forward passes inconsistent. -""" - -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -import torch -import torch.nn as nn -import torch.nn.functional as F # noqa: N812 - -from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.scripts import lerobot_eval - -# ── Test fixtures ──────────────────────────────────────────────────── - - -class _DuckPolicy(nn.Module): - """nn.Module with mode-sensitive layers, posing as a ``PreTrainedPolicy``. - - Inheriting from ``PreTrainedPolicy`` directly triggers ``__init_subclass__`` - requirements (``config_class``, ``name``) and pulls in 5 abstract methods. - We only need the isinstance check inside ``eval_policy`` to pass — registering - this class as a *virtual* subclass via ``abc.ABCMeta.register`` is the - minimum-overhead way to do that. - """ - - def __init__(self): - super().__init__() - # Two layers whose behaviour differs between train and eval mode: - # Dropout (active in train, no-op in eval) and BatchNorm (updates - # running stats in train, uses them in eval). - self.dropout = nn.Dropout(0.5) - self.bn = nn.BatchNorm1d(16) - - def reset(self) -> None: # called by rollout(); no-op suffices - pass - - -PreTrainedPolicy.register(_DuckPolicy) - - -def _canned_rollout_data() -> dict: - """Minimal rollout-output shape: one env, one done step.""" - return { - "done": torch.tensor([[True]], dtype=torch.bool), - "reward": torch.tensor([[1.0]], dtype=torch.float32), - "success": torch.tensor([[True]], dtype=torch.bool), - } - - -def _invoke_eval_policy(policy: nn.Module) -> dict: - """Call the real ``eval_policy`` with rollout mocked out.""" - env = MagicMock() - env.num_envs = 1 - with patch.object(lerobot_eval, "rollout", return_value=_canned_rollout_data()): - return lerobot_eval.eval_policy( - env=env, - policy=policy, - env_preprocessor=None, - env_postprocessor=None, - preprocessor=None, - postprocessor=None, - n_episodes=1, - max_episodes_rendered=0, - videos_dir=None, - return_episode_data=False, - start_seed=None, - ) - - -# ── Tests ──────────────────────────────────────────────────────────── - - -class TestEvalPolicyPreservesTrainingMode: - """The bug this guards against: ``eval_policy`` calling ``policy.eval()`` - without ever restoring ``policy.train()`` on return. - - Without restoration, the training loop's next forward pass runs with the - policy in eval mode — Dropout disabled, BN frozen — for the rest of - training. Under DDP, only ``is_main_process`` runs the eval, so the - main rank and workers end up in inconsistent modes. - """ - - def test_training_mode_preserved(self): - """A policy that entered eval_policy in train mode should leave in train mode.""" - policy = _DuckPolicy() - policy.train() - assert policy.training, "sanity: policy should start in train mode" - - _invoke_eval_policy(policy) - - assert policy.training, ( - "eval_policy must restore policy.training=True on return. " - "Without this, lerobot_train's training loop runs subsequent steps " - "with Dropout disabled and BatchNorm running-stats frozen." - ) - - def test_eval_mode_preserved(self): - """A policy that entered in eval mode should still be in eval mode after. - - Common for the standalone ``lerobot-eval`` script which loads a frozen - checkpoint. We must not accidentally flip it into train mode. - """ - policy = _DuckPolicy() - policy.eval() - assert not policy.training, "sanity: policy should start in eval mode" - - _invoke_eval_policy(policy) - - assert not policy.training, "eval_policy must not put a frozen policy into train mode" - - def test_dropout_remains_active_after_eval_policy(self): - """End-to-end behavioural check: dropout should still drop after the call. - - We measure dropout's zero-fraction on a continuous-valued input - (so any zeros come from dropout, not relu). In train mode this is - ≈ p = 0.5; in eval mode it's 0.0. - """ - policy = _DuckPolicy() - policy.train() - - _invoke_eval_policy(policy) - - torch.manual_seed(0) - x = torch.randn(256, 16) - with torch.no_grad(): - out = policy.dropout(x) - zero_frac = (out == 0).float().mean().item() - - # Loose bound: dropout's expected zero-frac is 0.5 but Bernoulli sampling - # has variance ~ p(1-p)/n. 0.3 gives plenty of room and still distinguishes - # the buggy case (zero_frac == 0) from the fixed case. - assert zero_frac > 0.3, ( - f"dropout zero-frac after eval_policy = {zero_frac:.3f}; " - "expected ~0.5 if the policy is in train mode. Dropout silently " - "disabled — the eval-mode restoration fix is missing or reverted." - ) - - -@pytest.mark.parametrize("starting_mode", ["train", "eval"]) -def test_eval_policy_does_not_crash_with_canned_rollout(starting_mode): - """Sanity guard: eval_policy itself completes and returns sensible aggregated - metrics from the canned rollout, regardless of entry mode.""" - policy = _DuckPolicy() - if starting_mode == "train": - policy.train() - else: - policy.eval() - - info = _invoke_eval_policy(policy) - assert info["aggregated"]["pc_success"] == 100.0 - assert "eval_s" in info["aggregated"] - - -# ── Quantitative impact demonstration ──────────────────────────────── -# The tests above gate against regression of the lerobot_eval fix itself. -# The test below quantifies *why the fix matters*: when the bug pattern is -# replayed on a tiny Dropout+BatchNorm MLP, the missing mode restoration -# causes measurably worse generalisation. The deterministic head-to-head -# below runs in ~3s and asserts the gap is real — a runnable, reproducible -# proof living alongside the regression test instead of in a separate -# benchmarks/ directory. - - -class _TinyMLP(nn.Module): - """MLP exercising both mode-sensitive layers: Dropout and BatchNorm.""" - - def __init__(self, in_dim: int = 16, hidden: int = 64, out_dim: int = 4): - super().__init__() - self.fc1 = nn.Linear(in_dim, hidden) - self.bn1 = nn.BatchNorm1d(hidden) - self.drop1 = nn.Dropout(0.5) - self.fc2 = nn.Linear(hidden, hidden) - self.bn2 = nn.BatchNorm1d(hidden) - self.drop2 = nn.Dropout(0.5) - self.head = nn.Linear(hidden, out_dim) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.drop1(F.relu(self.bn1(self.fc1(x)))) - x = self.drop2(F.relu(self.bn2(self.fc2(x)))) - return self.head(x) - - -def _make_data(n: int, in_dim: int, out_dim: int, noise: float, seed: int): - """Mildly nonlinear regression so the model has something to overfit to.""" - rng = np.random.default_rng(seed) - X = rng.standard_normal((n, in_dim)).astype(np.float32) # noqa: N806 - W1 = (rng.standard_normal((in_dim, in_dim)) / np.sqrt(in_dim)).astype(np.float32) # noqa: N806 - W2 = (rng.standard_normal((in_dim, out_dim)) / np.sqrt(in_dim)).astype(np.float32) # noqa: N806 - # Edge values in float32 matmul produce harmless RuntimeWarnings; suppress - # locally rather than relying on module-level filters (pytest captures - # warnings out of band and a module-level filterwarnings call would not - # take effect by the time the test runs). - with np.errstate(over="ignore", invalid="ignore", divide="ignore"): - y = np.tanh(X @ W1) @ W2 + noise * rng.standard_normal((n, out_dim)).astype(np.float32) - return torch.from_numpy(X), torch.from_numpy(y.astype(np.float32)) - - -def _run_synthetic_training( - apply_fix: bool, n_steps: int = 500, eval_every: int = 100, seed: int = 0 -) -> float: - """Replay the lerobot_train loop pattern with a stand-in for eval_policy. - - The stand-in sets ``policy.eval()`` exactly like the real chain at - ``lerobot_eval.py:307`` does. If ``apply_fix=True`` we restore - ``policy.train()`` after the call (mirroring this PR's patch). Returns the - final held-out val loss (mean of last 50 steps). - """ - torch.manual_seed(seed) - np.random.seed(seed) - - X_train, y_train = _make_data(n=2048, in_dim=16, out_dim=4, noise=0.3, seed=seed) # noqa: N806 - X_val, y_val = _make_data(n=512, in_dim=16, out_dim=4, noise=0.3, seed=seed + 100) # noqa: N806 - - policy = _TinyMLP() - optim = torch.optim.Adam(policy.parameters(), lr=1e-3) - policy.train() # mirrors lerobot_train.py:429 - - val_losses: list[float] = [] - for step in range(n_steps): - idx = torch.randint(0, X_train.shape[0], (64,)) - out = policy(X_train[idx]) - loss = F.mse_loss(out, y_train[idx]) - optim.zero_grad() - loss.backward() - optim.step() - - # Held-out val loss in eval mode (faithful generalisation metric) - prev = policy.training - policy.eval() - with torch.no_grad(): - val_losses.append(F.mse_loss(policy(X_val), y_val).item()) - policy.train(prev) - - # Periodic eval — replicates the lerobot_train.py:526-573 pattern - if (step + 1) % eval_every == 0: - policy.eval() # mirrors the bug at lerobot_eval.py:307 - with torch.no_grad(): - policy(X_val[:32]) # a stand-in for the eval rollout - if apply_fix: - policy.train() # the one-line fix in this PR - - return float(np.mean(val_losses[-50:])) - - -@pytest.mark.parametrize("seed", [0, 1]) -def test_missing_mode_restoration_hurts_generalisation(seed): - """Quantitative demonstration that the bug pattern degrades generalisation. - - Trains the same tiny Dropout+BatchNorm MLP twice with identical data and - seed: once mirroring the bug (``policy.eval()`` mid-training, never - restored), once with this PR's fix. Asserts the buggy variant generalises - measurably worse — confirming the mechanism that motivates the fix. - - This is a mechanism-demo test, not a regression gate on - ``lerobot_eval.eval_policy`` itself; the gates for that are the tests - above. Multiple seeds guard against single-seed flukes. - - The 5% margin is generous — in repeated runs we see ≈10–25% val-loss - deltas on this toy problem. Real policies (with more layers, more - Dropout, longer training) generally see larger gaps. - """ - bug_val_loss = _run_synthetic_training(apply_fix=False, seed=seed) - fix_val_loss = _run_synthetic_training(apply_fix=True, seed=seed) - - assert bug_val_loss > fix_val_loss * 1.05, ( - f"Expected the buggy training pattern (no mode restoration after eval) " - f"to generalise ≥5% worse than the fixed pattern. Got " - f"BUG val_loss={bug_val_loss:.4f}, FIXED val_loss={fix_val_loss:.4f} " - f"(delta={100 * (bug_val_loss - fix_val_loss) / fix_val_loss:+.1f}%). " - f"If this fails the mechanism may have changed — investigate before " - f"adjusting the threshold." - )