#!/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." )