fix pi052 runtime and training safety

tests/policies/pi052/test_pi052_bucketed_ce.py

@@ -19,7 +19,71 @@ import torch
 
 pytest.importorskip("transformers")
 
-from lerobot.policies.pi052.modeling_pi052 import _lin_ce_flat
+from lerobot.policies.pi052.modeling_pi052 import _lin_ce_flat, _shifted_lin_ce
+
+
+def test_shifted_ce_none_retains_distinct_per_sample_losses():
+    hidden = torch.tensor(
+        [
+            [[8.0, 0.0], [0.0, 8.0], [0.0, 0.0]],
+            [[0.0, 8.0], [8.0, 0.0], [0.0, 0.0]],
+        ]
+    )
+    labels = torch.tensor([[0, 0, 1], [0, 0, 1]])
+    losses = _shifted_lin_ce(hidden, torch.eye(2), labels, reduction="none")
+
+    assert losses.shape == (2,)
+    assert losses[0] < losses[1]
+
+
+def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_path):
+    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+
+    checkpoint = tmp_path / "model.safetensors"
+    checkpoint.touch()
+    calls = []
+
+    def fake_cached_file(model_id, filename, **kwargs):
+        calls.append((model_id, filename, kwargs))
+        return None if filename.endswith("index.json") else str(checkpoint)
+
+    monkeypatch.setattr(modeling_pi052, "cached_file", fake_cached_file)
+    files = modeling_pi052._resolve_weight_files(
+        "org/model",
+        force_download=True,
+        resume_download=True,
+        proxies={"https": "proxy"},
+        token="secret",
+        cache_dir=tmp_path / "cache",
+        local_files_only=True,
+        revision="commit",
+    )
+
+    assert files == [checkpoint]
+    for _model_id, _filename, kwargs in calls:
+        assert kwargs["revision"] == "commit"
+        assert kwargs["cache_dir"] == tmp_path / "cache"
+        assert kwargs["force_download"] is True
+        assert kwargs["resume_download"] is True
+        assert kwargs["proxies"] == {"https": "proxy"}
+        assert kwargs["token"] == "secret"
+        assert kwargs["local_files_only"] is True
+
+
+def test_checkpoint_resolution_rejects_local_directory_without_weights(tmp_path):
+    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+
+    with pytest.raises(FileNotFoundError, match="model.safetensors"):
+        modeling_pi052._resolve_weight_files(
+            tmp_path,
+            force_download=False,
+            resume_download=None,
+            proxies=None,
+            token=None,
+            cache_dir=None,
+            local_files_only=False,
+            revision=None,
+        )
 
 
 @pytest.mark.parametrize("z_loss_weight", [0.0, 1e-4])

tests/policies/pi052/test_pi052_import.py

@@ -0,0 +1,27 @@
+# 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.
+
+import subprocess
+import sys
+
+
+def test_pi052_config_import_does_not_load_model_or_dataset_processor():
+    code = """
+import sys
+from lerobot.policies import PI052Config
+assert PI052Config.__name__ == "PI052Config"
+assert "lerobot.policies.pi052.modeling_pi052" not in sys.modules
+assert "lerobot.policies.pi052.processor_pi052" not in sys.modules
+"""
+    subprocess.run([sys.executable, "-c", code], check=True)

tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py

@@ -16,6 +16,8 @@
 
 from types import SimpleNamespace
 
+import pytest
+
 from lerobot.policies import factory
 from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
 from lerobot.policies.pi052 import fit_fast_tokenizer as fit_module
@@ -48,6 +50,27 @@ def test_pi0_fast_resolves_dataset_specific_tokenizer(monkeypatch, tmp_path):
     }
 
 
+def test_fast_fit_failure_is_not_silently_replaced(monkeypatch, tmp_path):
+    config = PI0FastConfig(auto_fit_fast_tokenizer=True, fast_tokenizer_cache_dir=str(tmp_path))
+    monkeypatch.setattr(
+        fit_module,
+        "fit_fast_tokenizer",
+        lambda **kwargs: (_ for _ in ()).throw(RuntimeError("fit failed")),
+    )
+
+    with pytest.raises(RuntimeError, match="fit failed"):
+        fit_module.resolve_fast_tokenizer(config, "user/dataset")
+
+
+def test_each_node_uses_its_local_rank_zero_as_fit_leader(monkeypatch):
+    monkeypatch.setenv("RANK", "8")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+    assert fit_module._is_local_leader()
+
+    monkeypatch.setenv("LOCAL_RANK", "1")
+    assert not fit_module._is_local_leader()
+
+
 def test_pretrained_pi0_fast_overrides_only_fitted_tokenizer(monkeypatch):
     config = PI0FastConfig(auto_fit_fast_tokenizer=True)
     calls = []

tests/runtime/test_language_runtime.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from lerobot.runtime import (
-    LanguageConditionedRuntime,
-)
+import threading
+import time
+
+from lerobot.runtime import LanguageConditionedRuntime, Tick
 
 
 class FakeAdapter:
@@ -69,3 +70,31 @@ def test_runtime_handles_user_interjection():
 
     assert "please say ok" in adapter.interjections
     assert runtime.state.language_context["plan"] == "new plan"
+
+
+def test_prompt_change_discards_in_flight_action_chunk():
+    started = threading.Event()
+    release = threading.Event()
+
+    class BlockingAdapter(FakeAdapter):
+        def select_action(self, observation, state):
+            started.set()
+            assert release.wait(timeout=2)
+            return ["stale"]
+
+    runtime = LanguageConditionedRuntime(
+        policy_adapter=BlockingAdapter(),
+        observation_provider=lambda: {"observation.state": 1},
+    )
+    runtime.set_task("old task")
+    runtime.state.tick = Tick(index=1, monotonic_seconds=time.monotonic())
+    inference = threading.Thread(target=runtime.maybe_enqueue_action_chunk, kwargs={"force": True})
+    inference.start()
+    assert started.wait(timeout=2)
+
+    runtime.set_task("new task")
+    release.set()
+    inference.join(timeout=2)
+
+    assert not inference.is_alive()
+    assert list(runtime.state.action_queue) == []

tests/runtime/test_sim_robocasa.py

@@ -17,6 +17,7 @@ from types import SimpleNamespace
 
 import numpy as np
 
+from lerobot.runtime.sim_robocasa import RoboCasaSimBackend
 from lerobot.utils.video_annotation import annotate_frame
 
 
@@ -59,3 +60,23 @@ def test_overlay_draws_each_label_once(monkeypatch):
     assert all(color == (255, 255, 255) and thickness == 1 for _, color, thickness in put_text_calls)
     assert len(rectangle_calls) == 1
     assert not np.shares_memory(annotated, frame)
+
+
+def test_capture_updates_live_frame_when_recording_is_disabled(monkeypatch):
+    backend = object.__new__(RoboCasaSimBackend)
+    frame = np.full((8, 8, 3), 42, dtype=np.uint8)
+    written = []
+    backend.record = False
+    backend.runtime_state = None
+    backend._multiview_frame = lambda: frame
+    backend._current_task = lambda: "task"
+    backend._subtask_getter = None
+    backend._memory_getter = None
+    backend._latest_frame = None
+    backend._write_live_frame = written.append
+    monkeypatch.setattr("lerobot.runtime.sim_robocasa.annotate_frame", lambda image, labels: image)
+
+    backend._capture_frame()
+
+    assert backend._latest_frame is frame
+    assert written == [frame]

tests/utils/test_logging_utils.py

@@ -37,6 +37,12 @@ class MockAccelerator:
             return self._reduce_fn(tensor, reduction)
         return tensor
 
+    def gather(self, tensor):
+        if self._reduce_fn is None:
+            return tensor.repeat(self.num_processes)
+        reduced = self._reduce_fn(tensor, "max")
+        return torch.cat([tensor.repeat(self.num_processes - 1), reduced])
+
 
 def test_average_meter_initialization():
     meter = AverageMeter("loss", ":.2f")
@@ -168,6 +174,18 @@ def test_metrics_tracker_reset_averages(mock_metrics):
     assert tracker.accuracy.avg == 0.0
 
 
+def test_metrics_tracker_materializes_full_tensor_window(mock_metrics):
+    tracker = MetricsTracker(batch_size=2, num_frames=10, num_episodes=2, metrics=mock_metrics)
+    tracker.accumulate_tensor("loss", torch.tensor(1.0))
+    tracker.accumulate_tensor("loss", torch.tensor(3.0))
+
+    assert tracker.loss.count == 0
+    tracker.materialize_tensors()
+
+    assert tracker.loss.avg == pytest.approx(2.0)
+    assert tracker.loss.count == 2
+
+
 def test_average_meter_invalid_reduction():
     with pytest.raises(ValueError):
         AverageMeter("loss", reduction="median")

tests/utils/test_rerun_visualization.py

@@ -0,0 +1,311 @@
+#!/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_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