fix(scripts): keep policy train/eval

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
+    # Preserve the mode for direct callers. eval_policy_all scopes the mode
-    # the caller (notably the training loop in lerobot_train.py) sees the policy
+    # around all tasks so parallel evaluations cannot race with each other.
-    # 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.
     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
+    policy.train(was_training)
-    # (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()
 
     return info
 
@@ -1026,6 +1015,12 @@ def eval_policy_all(
         recording_private=recording_private,
     )
 
+    # 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):
@@ -1060,6 +1055,8 @@ def eval_policy_all(
                         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):

tests/scripts/test_eval.py

@@ -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."
-    )