admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-16 22:41:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: trim training comments and obsolete rerun test

Browse Source
This commit is contained in:
Pepijn
2026-07-15 16:25:25 +02:00
parent 6795b22b1e
commit d304d75ad7
2 changed files with 9 additions and 375 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/lerobot_train.py
+9 -34
View File
@@ -100,9 +100,7 @@ def update_policy(
lr_scheduler: An optional learning rate scheduler.
lock: An optional lock for thread-safe optimizer updates.
sample_weighter: Optional SampleWeighter instance for per-sample loss weighting.
log_metrics: When True, read loss/grad_norm/update_s off the GPU via `.item()` (a CUDA sync).
On non-logging steps set False so the step stays fully async — letting the next batch's
dataloading and enqueue overlap GPU compute instead of stalling on a per-step sync.
log_metrics: Whether to synchronize and record GPU metrics this step.
Returns:
A tuple containing:
@@ -173,17 +171,12 @@ def update_policy(
train_metrics.lr = optimizer.param_groups[0]["lr"]
if torch.cuda.is_available():
train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
# `loss.item()` / `grad_norm.item()` each block on a CUDA sync. Only pay that on logging steps;
# on the other steps the readouts are never consumed, so skipping them keeps the step async and
# lets CPU-side dataloading/enqueue overlap GPU compute. update_s is only accurate under that
# sync, so it too is recorded on logging steps only.
# Materialize GPU metrics only when logging to avoid synchronizing every step.
if log_metrics:
train_metrics.loss = loss.item()
train_metrics.grad_norm = grad_norm.item()
train_metrics.update_s = time.perf_counter() - start_time
# Policies may hand back loss components as detached tensors to keep the forward async on
# non-logging steps (e.g. pi052). Materialize them to python floats here, on logging steps only,
# so the per-step CUDA sync is paid 1-in-log_freq rather than every step.
# Materialize detached loss components during the same logging synchronization.
if output_dict:
output_dict = {
k: (v.item() if isinstance(v, torch.Tensor) else v) for k, v in output_dict.items()
@@ -228,25 +221,15 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
from accelerate.utils import InitProcessGroupKwargs
# find_unused_parameters=True is needed for conditional computation but
# breaks DDP's gradient/backward overlap and bucket coalescing, which is
# cheap intra-node (NVLink) but very costly across nodes (EFA). When the
# set of used params is stable, static_graph=True keeps unused-param
# support AND restores overlap. Env-gated; defaults preserve old behavior.
# Static graphs restore DDP overlap when conditional parameter usage is stable.
# Environment flags retain the existing defaults.
ddp_find_unused = os.environ.get("LEROBOT_DDP_FIND_UNUSED", "1") == "1"
ddp_static_graph = os.environ.get("LEROBOT_DDP_STATIC_GRAPH", "0") == "1"
ddp_kwargs = DistributedDataParallelKwargs(
find_unused_parameters=ddp_find_unused and not ddp_static_graph,
static_graph=ddp_static_graph,
)
# Bump the c10d store-get / barrier timeout so the rank-0-only
# ``make_dataset`` block below doesn't trigger a barrier crash on
# large datasets. Default is 10 min (``store->get`` 600 s); a
# 32 k-episode v3 dataset (e.g. ``robocasa_pretrain_human300_v4``)
# spends >13 min on rank 0 building the episode/frame index
# while ranks 1-N idle at ``wait_for_everyone()`` and crash with
# ``DistBackendError: ... wait timeout after 600000ms``. 2 h is
# plenty of headroom; fast paths are unaffected.
# Allow rank 0 enough time to index large datasets before other ranks leave the barrier.
ipg_kwargs = InitProcessGroupKwargs(timeout=timedelta(hours=2))
# Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting.
# Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training).
@@ -353,10 +336,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
active_cfg = cfg.trainable_config
processor_pretrained_path = active_cfg.pretrained_path
# pi052: even when loading pretrained weights, build the processors
# from the current pi052 config so the recipe text-label and FAST
# action-label steps are generated and not silently swapped for the
# checkpoint's older processor stack.
# Build PI052 processors from the current config so recipe and FAST labels are generated.
if cfg.policy.type == "pi052" and processor_pretrained_path is not None and not cfg.resume:
logging.warning(
"pi052 is loading pretrained weights from %s, but building processors from the current "
@@ -382,7 +362,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.is_reward_model_training:
processor_kwargs["dataset_meta"] = dataset.meta
# Policies that optionally fit dataset-specific processor artifacts need the repo id.
if cfg.policy.type in {"pi0_fast", "pi052"}:
processor_kwargs["dataset_repo_id"] = cfg.dataset.repo_id
@@ -530,10 +509,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
# declares language columns; otherwise stay on PyTorch's default
# collate so non-language training runs are unaffected.
collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None
# On multi-node EFA clusters, forking workers from a multi-GB rank process can
# fail with OSError(ENOMEM) because fork() reserve-charges the parent's full
# virtual footprint. Allow opting into "forkserver"/"spawn" so workers come
# from a clean process instead. Unset => default "fork" (unchanged behavior).
# Allow spawn/forkserver workers where forking large rank processes exhausts memory.
mp_context = os.environ.get("LEROBOT_DATALOADER_MP_CONTEXT") or None
dataloader = torch.utils.data.DataLoader(
dataset,
@@ -647,8 +623,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
batch = preprocessor(batch)
train_tracker.dataloading_s = time.perf_counter() - start_time
# This update produces step number `step + 1`; only sync metrics off the GPU when that step
# will be logged (mirrors the is_log_step gate below). Everything else stays async.
# Synchronize GPU metrics only for updates that will be logged.
log_metrics = cfg.log_freq > 0 and (step + 1) % cfg.log_freq == 0
train_tracker, output_dict = update_policy(
tests/utils/test_rerun_visualization.py
-341
View File
@@ -1,341 +0,0 @@
#!/usr/bin/env python
# Copyright 2025 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.
import importlib
import sys
from types import SimpleNamespace
import numpy as np
import pytest
pytest.importorskip("rerun", reason="rerun-sdk is required (install lerobot[viz])")
from lerobot.types import TransitionKey
from lerobot.utils.constants import OBS_STATE
@pytest.fixture
def mock_rerun(monkeypatch):
"""
Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't
depend on the real library. Also reload the module-under-test so it binds to
this mock `rr`.
"""
calls = []
blueprints = []
class DummyScalar:
def __init__(self, value):
# Scalars may be built from a single float or from a 1D array batch.
self.value = value
class DummyImage:
def __init__(self, arr):
self.arr = arr
def compress(self, *a, **k):
return self
class DummyDepthImage:
def __init__(self, arr, meter=None, colormap=None):
self.arr = arr
self.meter = meter
self.colormap = colormap
def dummy_log(key, obj=None, **kwargs):
# Accept either positional `obj` or keyword `entity` and record remaining kwargs.
if obj is None and "entity" in kwargs:
obj = kwargs.pop("entity")
calls.append((key, obj, kwargs))
def dummy_send_blueprint(blueprint, *a, **k):
blueprints.append(blueprint)
# Mock the `rerun.blueprint` submodule used to build the layout.
dummy_rrb = SimpleNamespace(
Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
kind="Spatial2DView", origin=origin, name=name
),
TimeSeriesView=lambda name=None, contents=None: SimpleNamespace(
kind="TimeSeriesView", name=name, contents=contents
),
Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)),
Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root),
)
dummy_rr = SimpleNamespace(
__name__="rerun",
__package__="rerun",
__spec__=SimpleNamespace(name="rerun", submodule_search_locations=None),
Scalars=DummyScalar,
Image=DummyImage,
DepthImage=DummyDepthImage,
components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")),
log=dummy_log,
send_blueprint=dummy_send_blueprint,
init=lambda *a, **k: None,
spawn=lambda *a, **k: None,
blueprint=dummy_rrb,
)
# Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`).
monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
# Now import and reload the module under test, to bind to our rerun mock
import lerobot.utils.rerun_visualization as rv
importlib.reload(rv)
# Expose the reloaded module, the call recorder and the captured blueprints
yield rv, calls, blueprints
def _keys(calls):
"""Helper to extract just the keys logged to rr.log"""
return [k for (k, _obj, _kw) in calls]
def _obj_for(calls, key):
"""Find the first object logged under a given key."""
for k, obj, _kw in calls:
if k == key:
return obj
raise KeyError(f"Key {key} not found in calls: {calls}")
def _kwargs_for(calls, key):
for k, _obj, kw in calls:
if k == key:
return kw
raise KeyError(f"Key {key} not found in calls: {calls}")
def _views_by_kind(blueprint, kind):
"""Return the views of a given kind from the (single) blueprint's grid."""
return [v for v in blueprint.root.views if v.kind == kind]
def test_init_rerun_can_serve_headless_web_viewer(mock_rerun, monkeypatch):
rv, _calls, _blueprints = mock_rerun
rr = sys.modules["rerun"]
served = {}
def serve_grpc(grpc_port):
served["grpc"] = grpc_port
return "rerun+http://localhost"
monkeypatch.setattr(
rr,
"serve_grpc",
serve_grpc,
raising=False,
)
monkeypatch.setattr(
rr,
"serve_web_viewer",
lambda **kwargs: served.setdefault("web", kwargs),
raising=False,
)
rv.init_rerun(session_name="runtime", port=8765, web_port=9091)
assert served["grpc"] == 8765
assert served["web"]["web_port"] == 9091
assert served["web"]["open_browser"] is False
assert served["web"]["connect_to"] == "rerun+http://localhost"
def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
rv, calls, blueprints = mock_rerun
# Build EnvTransition dict
obs = {
f"{OBS_STATE}.temperature": np.float32(25.0),
# CHW image should be converted to HWC for rr.Image
"observation.camera": np.zeros((3, 10, 20), dtype=np.uint8),
}
act = {
"action.throttle": 0.7,
# 1D array should be logged as a single Scalars batch under one entity path
"action.vector": np.array([1.0, 2.0], dtype=np.float32),
}
transition = {
TransitionKey.OBSERVATION: obs,
TransitionKey.ACTION: act,
}
# Extract observation and action data from transition like in the real call sites
obs_data = transition.get(TransitionKey.OBSERVATION, {})
action_data = transition.get(TransitionKey.ACTION, {})
rv.log_rerun_data(observation=obs_data, action=action_data)
# We expect:
# - observation.state.temperature -> Scalars
# - observation.camera -> Image (HWC) with static=True
# - action.throttle -> Scalars
# - action.vector -> single Scalars batch (no per-element suffix)
expected_keys = {
f"{OBS_STATE}.temperature",
"observation.camera",
"action.throttle",
"action.vector",
}
assert set(_keys(calls)) == expected_keys
# Check scalar types and values
temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature")
assert type(temp_obj).__name__ == "DummyScalar"
assert float(temp_obj.value) == pytest.approx(25.0)
throttle_obj = _obj_for(calls, "action.throttle")
assert type(throttle_obj).__name__ == "DummyScalar"
assert float(throttle_obj.value) == pytest.approx(0.7)
# 1D vector logged as a single batched Scalars under one entity path
vec = _obj_for(calls, "action.vector")
assert type(vec).__name__ == "DummyScalar"
np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0])
# Check image handling: CHW -> HWC
img_obj = _obj_for(calls, "observation.camera")
assert type(img_obj).__name__ == "DummyImage"
assert img_obj.arr.shape == (10, 20, 3) # transposed
assert _kwargs_for(calls, "observation.camera").get("static", False) is True # static=True for images
# A blueprint should have been built and sent exactly once, and cached on the function.
assert len(blueprints) == 1
assert rv.log_rerun_data.blueprint is blueprints[0]
bp = blueprints[0]
# One spatial view per image path
spatial_views = _views_by_kind(bp, "Spatial2DView")
assert {v.origin for v in spatial_views} == {"observation.camera"}
# One time-series view each for observation and action scalars
ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
assert set(ts_views) == {"observation", "action"}
assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"]
assert ts_views["action"].contents == ["action.throttle", "action.vector"]
def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
rv, calls, blueprints = mock_rerun
# First dict without prefixes treated as observation
# Second dict without prefixes treated as action
obs_plain = {
"temp": 1.5,
# Already HWC image => should stay as-is
"img": np.zeros((5, 6, 3), dtype=np.uint8),
"none": None, # should be skipped
}
act_plain = {
"throttle": 0.3,
"vec": np.array([9, 8, 7], dtype=np.float32),
}
# Extract observation and action data from list like the old function logic did
# First dict was treated as observation, second as action
rv.log_rerun_data(observation=obs_plain, action=act_plain)
# Expected keys with auto-prefixes. The 1D vector is a single batched Scalars.
expected = {
"observation.temp",
"observation.img",
"action.throttle",
"action.vec",
}
logged = set(_keys(calls))
assert logged == expected
# Scalars
t = _obj_for(calls, "observation.temp")
assert type(t).__name__ == "DummyScalar"
assert float(t.value) == pytest.approx(1.5)
throttle = _obj_for(calls, "action.throttle")
assert type(throttle).__name__ == "DummyScalar"
assert float(throttle.value) == pytest.approx(0.3)
# Image stays HWC
img = _obj_for(calls, "observation.img")
assert type(img).__name__ == "DummyImage"
assert img.arr.shape == (5, 6, 3)
assert _kwargs_for(calls, "observation.img").get("static", False) is True
# Vector logged as a single batched Scalars under one entity path
vec = _obj_for(calls, "action.vec")
assert type(vec).__name__ == "DummyScalar"
np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7])
# Blueprint sent once with the expected view layout
assert len(blueprints) == 1
bp = blueprints[0]
spatial_views = _views_by_kind(bp, "Spatial2DView")
assert {v.origin for v in spatial_views} == {"observation.img"}
ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
assert ts_views["observation"].contents == ["observation.temp"]
assert ts_views["action"].contents == ["action.throttle", "action.vec"]
def test_log_rerun_data_kwargs_only(mock_rerun):
rv, calls, blueprints = mock_rerun
rv.log_rerun_data(
observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)},
action={"action.a": 1.0},
)
keys = set(_keys(calls))
assert "observation.temp" in keys
assert "observation.gray" in keys
assert "action.a" in keys
temp = _obj_for(calls, "observation.temp")
assert type(temp).__name__ == "DummyScalar"
assert float(temp.value) == pytest.approx(10.0)
img = _obj_for(calls, "observation.gray")
assert type(img).__name__ == "DummyDepthImage" # single-channel -> DepthImage
assert img.arr.shape == (8, 8, 1) # remains HWC
assert _kwargs_for(calls, "observation.gray").get("static", False) is True
a = _obj_for(calls, "action.a")
assert type(a).__name__ == "DummyScalar"
assert float(a.value) == pytest.approx(1.0)
# Blueprint sent once, with a spatial view for the image and time-series views for scalars
assert len(blueprints) == 1
bp = blueprints[0]
assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"}
ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
assert ts_views["observation"].contents == ["observation.temp"]
assert ts_views["action"].contents == ["action.a"]
def test_log_rerun_data_blueprint_sent_only_once(mock_rerun):
"""The blueprint is built from the first call and not resent on subsequent calls."""
rv, calls, blueprints = mock_rerun
rv.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0})
assert len(blueprints) == 1
first_blueprint = rv.log_rerun_data.blueprint
rv.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0})
# Still only one blueprint, and the cached one is unchanged.
assert len(blueprints) == 1
assert rv.log_rerun_data.blueprint is first_blueprint