mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-27 19:56:09 +00:00
fix(scripts): keep policy train/eval
This commit is contained in:
Steven Palma
@@ -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,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):
|
||||
|
||||
@@ -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."
|
||||
)
|
||||
Reference in New Issue
Block a user