Merge branch 'main' into feat/bump_rerun

Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>

tests/utils/test_keyboard_input.py

@@ -0,0 +1,228 @@
+#!/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.
+
+"""Unit tests for the display-independent keyboard input helpers.
+
+These cover the parts most likely to regress: the environment-detection decision
+table (the heart of the Wayland/headless fix), the macOS trust probe, the control
+mapping, the terminal escape-sequence parsing, and backend selection. They require
+neither ``pynput`` nor a real terminal.
+"""
+
+import io
+import platform
+import sys
+
+import pytest
+
+import lerobot.utils.keyboard_input as ki
+from lerobot.utils.keyboard_input import (
+    TerminalKeyListener,
+    apply_recording_control,
+    create_key_listener,
+    init_keyboard_listener,
+    is_headless,
+    is_wayland,
+    pynput_can_capture,
+    pynput_listener_is_trusted,
+)
+
+
+@pytest.fixture(autouse=True)
+def _clear_detection_caches():
+    """The detection helpers are ``@cache``-decorated; clear around each test."""
+    for fn in (is_headless, is_wayland, pynput_can_capture):
+        fn.cache_clear()
+    yield
+    for fn in (is_headless, is_wayland, pynput_can_capture):
+        fn.cache_clear()
+
+
+def _set_platform(monkeypatch, name):
+    monkeypatch.setattr(platform, "system", lambda: name)
+
+
+def _set_tty(monkeypatch, is_tty):
+    stdin = io.StringIO("")
+    stdin.isatty = lambda: is_tty
+    monkeypatch.setattr(sys, "stdin", stdin)
+
+
+# --- Environment detection (the core of the fix) ---------------------------
+@pytest.mark.parametrize(
+    ("system", "env", "expected"),
+    [
+        ("Linux", {}, True),  # no display server
+        ("Linux", {"DISPLAY": ":0"}, False),  # X11
+        ("Linux", {"WAYLAND_DISPLAY": "wayland-0"}, False),  # Wayland
+        ("Darwin", {}, False),  # display always assumed present
+    ],
+)
+def test_is_headless(monkeypatch, system, env, expected):
+    _set_platform(monkeypatch, system)
+    monkeypatch.delenv("DISPLAY", raising=False)
+    monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
+    for key, value in env.items():
+        monkeypatch.setenv(key, value)
+    assert is_headless() is expected
+
+
+@pytest.mark.parametrize(
+    ("env", "expected"),
+    [
+        ({"XDG_SESSION_TYPE": "wayland"}, True),
+        ({"WAYLAND_DISPLAY": "wayland-0"}, True),
+        ({"XDG_SESSION_TYPE": "x11"}, False),
+        ({}, False),
+    ],
+)
+def test_is_wayland(monkeypatch, env, expected):
+    monkeypatch.delenv("XDG_SESSION_TYPE", raising=False)
+    monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
+    for key, value in env.items():
+        monkeypatch.setenv(key, value)
+    assert is_wayland() is expected
+
+
+@pytest.mark.parametrize(
+    ("system", "env", "pynput_available", "expected"),
+    [
+        ("Linux", {"DISPLAY": ":0"}, True, True),  # X11
+        ("Linux", {"DISPLAY": ":0", "WAYLAND_DISPLAY": "wayland-0"}, True, False),  # Wayland
+        ("Linux", {}, True, False),  # headless
+        ("Darwin", {}, True, True),
+        ("Linux", {"DISPLAY": ":0"}, False, False),  # pynput not installed
+    ],
+)
+def test_pynput_can_capture(monkeypatch, system, env, pynput_available, expected):
+    _set_platform(monkeypatch, system)
+    monkeypatch.setattr(ki, "_pynput_available", pynput_available)
+    for var in ("DISPLAY", "WAYLAND_DISPLAY", "XDG_SESSION_TYPE"):
+        monkeypatch.delenv(var, raising=False)
+    for key, value in env.items():
+        monkeypatch.setenv(key, value)
+    assert pynput_can_capture() is expected
+
+
+# --- macOS trust probe ------------------------------------------------------
+class _FakeListener:
+    def __init__(self, is_trusted):
+        self.IS_TRUSTED = is_trusted
+
+
+def test_pynput_listener_is_trusted(monkeypatch):
+    _set_platform(monkeypatch, "Linux")
+    assert pynput_listener_is_trusted(_FakeListener(False)) is True  # non-macOS: always assumed ok
+    _set_platform(monkeypatch, "Darwin")
+    assert pynput_listener_is_trusted(_FakeListener(False), timeout_s=0.05) is False
+
+
+# --- Control mapping --------------------------------------------------------
+def test_apply_recording_control():
+    events = {"exit_early": False, "rerecord_episode": False, "stop_recording": False}
+    apply_recording_control("left", events)
+    assert events == {"exit_early": True, "rerecord_episode": True, "stop_recording": False}
+    apply_recording_control("esc", events)
+    assert events["stop_recording"] is True
+    apply_recording_control("up", events)  # unknown control -> no-op (no error)
+
+
+# --- Terminal escape-sequence parsing (the tricky bit) ----------------------
+def _drive(listener, byte_seq):
+    """Run the listener's read loop over a scripted list of bytes (no real terminal)."""
+    script = list(byte_seq)
+
+    def fake_read(timeout):
+        if script:
+            return script.pop(0)
+        listener._running = False
+        return None
+
+    listener._read_char = fake_read
+    listener._running = True
+    listener._run()
+
+
+@pytest.mark.parametrize(
+    ("byte_seq", "expected"),
+    [
+        (["\x1b", "[", "C"], ["right"]),  # CSI arrow
+        (["\x1b", "O", "D"], ["left"]),  # SS3 arrow (e.g. over SSH/tmux)
+        (["\x1b"], ["esc"]),  # bare ESC
+        (["\x1b", "[", "A"], ["up"]),  # decoded even though the record handler ignores it
+        (["n"], ["n"]),  # letter passthrough
+    ],
+)
+def test_terminal_parsing(byte_seq, expected):
+    collected = []
+    _drive(TerminalKeyListener(collected.append), byte_seq)
+    assert collected == expected
+
+
+# --- Backend selection ------------------------------------------------------
+def test_init_selects_terminal_when_pynput_cannot_capture(monkeypatch):
+    monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
+    _set_tty(monkeypatch, is_tty=True)
+    monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None)  # avoid touching termios
+    listener, _ = init_keyboard_listener()
+    assert isinstance(listener, TerminalKeyListener)
+
+
+def test_init_returns_none_without_tty(monkeypatch):
+    monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
+    _set_tty(monkeypatch, is_tty=False)
+    listener, _ = init_keyboard_listener()
+    assert listener is None
+
+
+@pytest.mark.parametrize(
+    ("key", "flag"),
+    [("right", "exit_early"), ("r", "rerecord_episode"), ("q", "stop_recording")],
+)
+def test_init_terminal_key_routing(monkeypatch, key, flag):
+    """Arrows and their letter equivalents drive the same events (terminal backend)."""
+    monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
+    _set_tty(monkeypatch, is_tty=True)
+    monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None)
+    listener, events = init_keyboard_listener()
+    listener._on_key(key)
+    assert events[flag] is True
+
+
+# --- Shared factory + pynput key resolver -----------------------------------
+def test_resolve_pynput_key_char_fallback():
+    """Unmapped keys fall back to ``.char`` (and yield None when there is none)."""
+    assert ki._resolve_pynput_key(type("K", (), {"char": "s"})()) == "s"
+    assert ki._resolve_pynput_key(type("K", (), {"char": None})()) is None
+    assert ki._resolve_pynput_key(type("K", (), {"char": ""})()) is None  # empty char -> no key
+
+
+def test_create_key_listener_routes_to_dispatch(monkeypatch):
+    """The terminal backend forwards canonical key names straight to ``dispatch``."""
+    monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
+    _set_tty(monkeypatch, is_tty=True)
+    monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None)
+    seen = []
+    listener = create_key_listener(seen.append, controls_help="save='s'")
+    assert isinstance(listener, TerminalKeyListener)
+    listener._on_key("space")
+    assert seen == ["space"]
+
+
+def test_create_key_listener_none_without_tty(monkeypatch):
+    monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
+    _set_tty(monkeypatch, is_tty=False)
+    assert create_key_listener(lambda name: None) is None

tests/utils/test_logging_utils.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 
 import pytest
+import torch
 
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 
@@ -25,8 +26,16 @@ def mock_metrics():
 
 
 class MockAccelerator:
-    def __init__(self, num_processes: int):
+    def __init__(self, num_processes: int, reduce_fn=None):
         self.num_processes = num_processes
+        self.device = torch.device("cpu")
+        self._reduce_fn = reduce_fn
+
+    def reduce(self, tensor, reduction="mean"):
+        # In single-process tests we just want a deterministic stand-in for accelerate's reduce.
+        if self._reduce_fn is not None:
+            return self._reduce_fn(tensor, reduction)
+        return tensor
 
 
 def test_average_meter_initialization():
@@ -157,3 +166,70 @@ def test_metrics_tracker_reset_averages(mock_metrics):
     tracker.reset_averages()
     assert tracker.loss.avg == 0.0
     assert tracker.accuracy.avg == 0.0
+
+
+def test_average_meter_invalid_reduction():
+    with pytest.raises(ValueError):
+        AverageMeter("loss", reduction="median")
+
+
+def test_average_meter_reduction_stored():
+    meter = AverageMeter("updt_s", reduction="max")
+    assert meter.reduction == "max"
+
+
+def test_metrics_tracker_reduce_across_ranks_no_accelerator():
+    metrics = {"update_s": AverageMeter("update_s", reduction="max")}
+    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics)
+    tracker.update_s = 0.5
+    tracker.reduce_across_ranks()  # no-op without accelerator
+    assert tracker.update_s.avg == 0.5
+
+
+def test_metrics_tracker_reduce_across_ranks_single_process():
+    metrics = {"update_s": AverageMeter("update_s", reduction="max")}
+    tracker = MetricsTracker(
+        batch_size=32,
+        num_frames=1000,
+        num_episodes=50,
+        metrics=metrics,
+        accelerator=MockAccelerator(num_processes=1),
+    )
+    tracker.update_s = 0.5
+    tracker.reduce_across_ranks()  # no-op when world size is 1
+    assert tracker.update_s.avg == 0.5
+
+
+def test_metrics_tracker_reduce_across_ranks_invokes_reduce():
+    captured = {}
+
+    def fake_reduce(tensor, reduction):
+        captured["reduction"] = reduction
+        captured["values"] = tensor.clone()
+        # Pretend the slowest rank reported 0.9 instead of this rank's 0.4.
+        return torch.tensor([0.9], dtype=tensor.dtype, device=tensor.device)
+
+    metrics = {
+        "loss": AverageMeter("loss"),  # reduction="none" -> not touched
+        "update_s": AverageMeter("update_s", reduction="max"),
+    }
+    tracker = MetricsTracker(
+        batch_size=32,
+        num_frames=1000,
+        num_episodes=50,
+        metrics=metrics,
+        accelerator=MockAccelerator(num_processes=4, reduce_fn=fake_reduce),
+    )
+    tracker.loss = 1.0
+    tracker.update_s = 0.4
+    tracker.reduce_across_ranks()
+
+    assert captured["reduction"] == "max"
+    assert torch.allclose(captured["values"], torch.tensor([0.4]))
+    assert tracker.update_s.avg == pytest.approx(0.9)
+    # Metrics without a reduction stay untouched.
+    assert tracker.loss.avg == 1.0
+    # Invariant: avg == sum / count must hold after reduce, so subsequent .update() calls
+    # accumulate against the cluster view rather than the stale per-rank sum.
+    meter = tracker.update_s
+    assert meter.sum / meter.count == pytest.approx(meter.avg)

tests/utils/test_train_utils.py

@@ -20,6 +20,8 @@ from unittest.mock import Mock, patch
 from lerobot.common.train_utils import (
     get_step_checkpoint_dir,
     get_step_identifier,
+    load_training_batch_size,
+    load_training_num_processes,
     load_training_state,
     load_training_step,
     save_checkpoint,
@@ -63,6 +65,28 @@ def test_load_training_step(tmp_path):
     assert loaded_step == step
 
 
+def test_save_training_state_records_num_processes(tmp_path, optimizer, scheduler):
+    save_training_state(tmp_path, 10, optimizer, scheduler, num_processes=4)
+    assert load_training_num_processes(tmp_path) == 4
+
+
+def test_load_training_num_processes_absent_returns_none(tmp_path, optimizer, scheduler):
+    # Checkpoints written before the world size was recorded must still load (back-compat).
+    save_training_state(tmp_path, 10, optimizer, scheduler)
+    assert load_training_num_processes(tmp_path) is None
+
+
+def test_save_training_state_records_batch_size(tmp_path, optimizer, scheduler):
+    save_training_state(tmp_path, 10, optimizer, scheduler, batch_size=32)
+    assert load_training_batch_size(tmp_path) == 32
+
+
+def test_load_training_batch_size_absent_returns_none(tmp_path, optimizer, scheduler):
+    # Checkpoints written before the batch size was recorded must still load (back-compat).
+    save_training_state(tmp_path, 10, optimizer, scheduler)
+    assert load_training_batch_size(tmp_path) is None
+
+
 def test_update_last_checkpoint(tmp_path):
     checkpoint = tmp_path / "0005"
     checkpoint.mkdir()
@@ -112,3 +136,18 @@ def test_save_load_training_state(tmp_path, optimizer, scheduler):
     assert loaded_step == 10
     assert loaded_optimizer is optimizer
     assert loaded_scheduler is scheduler
+
+
+def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler):
+    # FSDP loads optimizer separately (after accelerator.prepare)
+    # load_training_state(load_optimizer=False) must restore step + scheduler but leave the
+    # optimizer untouched and never touch the on-disk optimizer state.
+    save_training_state(tmp_path, 10, optimizer, scheduler)
+    with patch("lerobot.common.train_utils.load_optimizer_state") as mock_load_optimizer_state:
+        loaded_step, loaded_optimizer, loaded_scheduler = load_training_state(
+            tmp_path, optimizer, scheduler, load_optimizer=False
+        )
+    mock_load_optimizer_state.assert_not_called()
+    assert loaded_step == 10
+    assert loaded_optimizer is optimizer
+    assert loaded_scheduler is scheduler

tests/utils/test_visualization_utils.py

@@ -48,6 +48,11 @@ def mock_rerun(monkeypatch):
 
         def compress(self, *a, **k):
             return self
+    
+    class DummyDepthImage:
+        def __init__(self, arr, colormap=None):
+            self.arr = arr
+            self.colormap = colormap
 
     def dummy_log(key, obj=None, **kwargs):
         # Accept either positional `obj` or keyword `entity` and record remaining kwargs.
@@ -76,6 +81,8 @@ def mock_rerun(monkeypatch):
         __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,
@@ -272,7 +279,7 @@ def test_log_rerun_data_kwargs_only(mock_rerun):
     assert float(temp.value) == pytest.approx(10.0)
 
     img = _obj_for(calls, "observation.gray")
-    assert type(img).__name__ == "DummyImage"
+    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