From d304d75ad70baa0b4a5d09e3a35cfcce4f2c27a5 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 15 Jul 2026 16:25:25 +0200 Subject: [PATCH] chore: trim training comments and obsolete rerun test --- src/lerobot/scripts/lerobot_train.py | 43 +-- tests/utils/test_rerun_visualization.py | 341 ------------------------ 2 files changed, 9 insertions(+), 375 deletions(-) delete mode 100644 tests/utils/test_rerun_visualization.py diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 55846c821..428f1ff8e 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -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( diff --git a/tests/utils/test_rerun_visualization.py b/tests/utils/test_rerun_visualization.py deleted file mode 100644 index a52ef9077..000000000 --- a/tests/utils/test_rerun_visualization.py +++ /dev/null @@ -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