chore: trim training comments and obsolete rerun test

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).
+        log_metrics: Whether to synchronize and record GPU metrics this step.
-            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.
 
     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;
+    # Materialize GPU metrics only when logging to avoid synchronizing every step.
-    # 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.
     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
+        # Materialize detached loss components during the same logging synchronization.
-        # 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.
         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
+        # Static graphs restore DDP overlap when conditional parameter usage is stable.
-        # breaks DDP's gradient/backward overlap and bucket coalescing, which is
+        # Environment flags retain the existing defaults.
-        # 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.
         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
+        # Allow rank 0 enough time to index large datasets before other ranks leave the barrier.
-        # ``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.
         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
+    # Build PI052 processors from the current config so recipe and FAST labels are generated.
-    # 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.
     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
+    # Allow spawn/forkserver workers where forking large rank processes exhausts memory.
-    # 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).
     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
+        # Synchronize GPU metrics only for updates that will be logged.
-        # will be logged (mirrors the is_log_step gate below). Everything else stays async.
         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

@@ -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