lerobot/tests/scripts/test_eval.py

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