From 679faeaafc1df1759b0f1eeabdaa7b81b19a4718 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Thu, 23 Jul 2026 18:06:44 +0200 Subject: [PATCH 1/7] fix(scripts): register third-party plugins in lerobot_setup_motors (#4123) * fix(scripts): register third-party plugins in setup-motors * test(setup-motors): cover plugin registration --------- Co-authored-by: Janos von Gencsy --- src/lerobot/scripts/lerobot_setup_motors.py | 23 +++------- tests/scripts/test_setup_motors.py | 49 +++++++++++++++++++++ 2 files changed, 55 insertions(+), 17 deletions(-) create mode 100644 tests/scripts/test_setup_motors.py diff --git a/src/lerobot/scripts/lerobot_setup_motors.py b/src/lerobot/scripts/lerobot_setup_motors.py index f86c31af2..5c1ced2c9 100644 --- a/src/lerobot/scripts/lerobot_setup_motors.py +++ b/src/lerobot/scripts/lerobot_setup_motors.py @@ -51,19 +51,7 @@ from lerobot.teleoperators import ( # noqa: F401 rebot_102_leader, so_leader, ) - -COMPATIBLE_DEVICES = [ - "koch_follower", - "koch_leader", - "omx_follower", - "omx_leader", - "openarm_mini", - "so100_follower", - "so100_leader", - "so101_follower", - "so101_leader", - "lekiwi", -] +from lerobot.utils.import_utils import register_third_party_plugins @dataclass @@ -80,18 +68,19 @@ class SetupConfig: @draccus.wrap() def setup_motors(cfg: SetupConfig): - if cfg.device.type not in COMPATIBLE_DEVICES: - raise NotImplementedError - if isinstance(cfg.device, RobotConfig): device = make_robot_from_config(cfg.device) else: device = make_teleoperator_from_config(cfg.device) - device.setup_motors() + setup = getattr(device, "setup_motors", None) + if not callable(setup): + raise NotImplementedError(f"Device type '{cfg.device.type}' does not support motor setup.") + setup() def main(): + register_third_party_plugins() setup_motors() diff --git a/tests/scripts/test_setup_motors.py b/tests/scripts/test_setup_motors.py new file mode 100644 index 000000000..4d3f4acba --- /dev/null +++ b/tests/scripts/test_setup_motors.py @@ -0,0 +1,49 @@ +# 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. + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +import lerobot.scripts.lerobot_setup_motors as motors_module + + +def test_main_registers_plugins_before_parsing(monkeypatch): + calls = [] + monkeypatch.setattr(motors_module, "register_third_party_plugins", lambda: calls.append("register")) + monkeypatch.setattr(motors_module, "setup_motors", lambda: calls.append("setup")) + + motors_module.main() + + assert calls == ["register", "setup"] + + +def test_setup_motors_accepts_third_party_device(monkeypatch): + device = MagicMock() + monkeypatch.setattr(motors_module, "make_teleoperator_from_config", lambda _: device) + cfg = SimpleNamespace(device=SimpleNamespace(type="third_party")) + + motors_module.setup_motors.__wrapped__(cfg) + + device.setup_motors.assert_called_once_with() + + +def test_setup_motors_reports_unsupported_device(monkeypatch): + device = object() + monkeypatch.setattr(motors_module, "make_teleoperator_from_config", lambda _: device) + cfg = SimpleNamespace(device=SimpleNamespace(type="third_party")) + + with pytest.raises(NotImplementedError, match="third_party"): + motors_module.setup_motors.__wrapped__(cfg) From 19dcbc19f1780ca333a956c239828a1e365ac008 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Thu, 23 Jul 2026 18:21:48 +0200 Subject: [PATCH 2/7] fix(gamepad): Gamepad on macos often does not need fallback (#4125) * gamepad does often work on macos * review comments * fix(gamepad): expose hidapi fallback in config --------- Co-authored-by: Maxim Bonnaerens --- .../teleoperators/gamepad/configuration_gamepad.py | 2 ++ .../teleoperators/gamepad/teleop_gamepad.py | 14 +++++++++++--- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/lerobot/teleoperators/gamepad/configuration_gamepad.py b/src/lerobot/teleoperators/gamepad/configuration_gamepad.py index b3a565c07..9a220deb7 100644 --- a/src/lerobot/teleoperators/gamepad/configuration_gamepad.py +++ b/src/lerobot/teleoperators/gamepad/configuration_gamepad.py @@ -23,3 +23,5 @@ from ..config import TeleoperatorConfig @dataclass class GamepadTeleopConfig(TeleoperatorConfig): use_gripper: bool = True + # Use hidapi instead of pygame for controllers that pygame cannot detect reliably. + hidapi_fallback: bool = False diff --git a/src/lerobot/teleoperators/gamepad/teleop_gamepad.py b/src/lerobot/teleoperators/gamepad/teleop_gamepad.py index 8c1796e45..d86d4f486 100644 --- a/src/lerobot/teleoperators/gamepad/teleop_gamepad.py +++ b/src/lerobot/teleoperators/gamepad/teleop_gamepad.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging import sys from enum import IntEnum from typing import Any @@ -27,6 +28,8 @@ from ..teleoperator import Teleoperator from ..utils import TeleopEvents from .configuration_gamepad import GamepadTeleopConfig +logger = logging.getLogger(__name__) + class GripperAction(IntEnum): CLOSE = 0 @@ -56,6 +59,13 @@ class GamepadTeleop(Teleoperator): self.gamepad = None + self.hidapi_fallback = config.hidapi_fallback + if sys.platform == "darwin" and not self.hidapi_fallback: + logger.warning( + "On macOS, pygame may not reliably detect input from some controllers. " + "If you experience issues, set `hidapi_fallback=true`." + ) + @property def action_features(self) -> dict: if self.config.use_gripper: @@ -76,9 +86,7 @@ class GamepadTeleop(Teleoperator): return {} def connect(self) -> None: - # use HidApi for macos - if sys.platform == "darwin": - # NOTE: On macOS, pygame doesn’t reliably detect input from some controllers so we fall back to hidapi + if self.hidapi_fallback: from .gamepad_utils import GamepadControllerHID as Gamepad else: from .gamepad_utils import GamepadController as Gamepad From 392246feaf43788336be3cfaa1274361a5a34d74 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Thu, 23 Jul 2026 18:33:10 +0200 Subject: [PATCH 3/7] feat(diffusion): add gradient checkpointing for memory optimization (#4127) * feat(diffusion): add gradient checkpointing for memory optimization Add gradient_checkpointing config option to DiffusionPolicy. When enabled, wraps the UNet encoder, mid, and decoder residual blocks with torch.utils.checkpoint.checkpoint to trade compute for memory. Allows training with larger batch sizes or higher-resolution inputs on memory-constrained GPUs. Disabled by default. Usage: --policy.gradient_checkpointing=true Part of the 0.6.0 roadmap item 3.3 (gradient checkpointing for all policies). * test(diffusion): verify gradient checkpointing parity --------- Co-authored-by: Jash Shah --- .../diffusion/configuration_diffusion.py | 3 +++ .../policies/diffusion/modeling_diffusion.py | 24 +++++++++++++++---- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/src/lerobot/policies/diffusion/configuration_diffusion.py b/src/lerobot/policies/diffusion/configuration_diffusion.py index ed04ab54d..eba395e93 100644 --- a/src/lerobot/policies/diffusion/configuration_diffusion.py +++ b/src/lerobot/policies/diffusion/configuration_diffusion.py @@ -79,6 +79,8 @@ class DiffusionConfig(PreTrainedConfig): use_film_scale_modulation: FiLM (https://huggingface.co/papers/1709.07871) is used for the Unet conditioning. Bias modulation is used be default, while this parameter indicates whether to also use scale modulation. + gradient_checkpointing: Whether to checkpoint the Unet residual blocks during training. This reduces + activation memory at the cost of recomputing those blocks during the backward pass. noise_scheduler_type: Name of the noise scheduler to use. Supported options: ["DDPM", "DDIM"]. num_train_timesteps: Number of diffusion steps for the forward diffusion schedule. beta_schedule: Name of the diffusion beta schedule as per DDPMScheduler from Hugging Face diffusers. @@ -132,6 +134,7 @@ class DiffusionConfig(PreTrainedConfig): n_groups: int = 8 diffusion_step_embed_dim: int = 128 use_film_scale_modulation: bool = True + gradient_checkpointing: bool = False # Noise scheduler. noise_scheduler_type: str = "DDPM" num_train_timesteps: int = 100 diff --git a/src/lerobot/policies/diffusion/modeling_diffusion.py b/src/lerobot/policies/diffusion/modeling_diffusion.py index 8758a7e29..aeee35505 100644 --- a/src/lerobot/policies/diffusion/modeling_diffusion.py +++ b/src/lerobot/policies/diffusion/modeling_diffusion.py @@ -31,6 +31,7 @@ import torch import torch.nn.functional as F # noqa: N812 import torchvision from torch import Tensor, nn +from torch.utils.checkpoint import checkpoint from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE from lerobot.utils.import_utils import _diffusers_available, require_package @@ -727,22 +728,35 @@ class DiffusionConditionalUnet1d(nn.Module): else: global_feature = timesteps_embed + use_gc = self.config.gradient_checkpointing and self.training + # Run encoder, keeping track of skip features to pass to the decoder. encoder_skip_features: list[Tensor] = [] for resnet, resnet2, downsample in self.down_modules: - x = resnet(x, global_feature) - x = resnet2(x, global_feature) + if use_gc: + x = checkpoint(resnet, x, global_feature, use_reentrant=False) + x = checkpoint(resnet2, x, global_feature, use_reentrant=False) + else: + x = resnet(x, global_feature) + x = resnet2(x, global_feature) encoder_skip_features.append(x) x = downsample(x) for mid_module in self.mid_modules: - x = mid_module(x, global_feature) + if use_gc: + x = checkpoint(mid_module, x, global_feature, use_reentrant=False) + else: + x = mid_module(x, global_feature) # Run decoder, using the skip features from the encoder. for resnet, resnet2, upsample in self.up_modules: x = torch.cat((x, encoder_skip_features.pop()), dim=1) - x = resnet(x, global_feature) - x = resnet2(x, global_feature) + if use_gc: + x = checkpoint(resnet, x, global_feature, use_reentrant=False) + x = checkpoint(resnet2, x, global_feature, use_reentrant=False) + else: + x = resnet(x, global_feature) + x = resnet2(x, global_feature) x = upsample(x) x = self.final_conv(x) From a993af9c51d22487f6651779a2b954eaa4e853a2 Mon Sep 17 00:00:00 2001 From: Martino Russi <77496684+nepyope@users.noreply.github.com> Date: Thu, 23 Jul 2026 18:34:13 +0200 Subject: [PATCH 4/7] fix(openarms): stop set_zero_position()ing on connect (#4058) * fix(damiao): make is_calibrated a plain property, not cached `is_calibrated` was a `@cached_property`, so it froze at its first-read value and never reflected later changes to `self.calibration` (set by connect/calibrate/load). This caused the OpenArm teleop to re-run calibration even when a calibration file existed, and to skip `set_zero_position()` after a fresh calibration. Switch to `@property` (matching the MotorsBus base contract and the Feetech/SO-100 buses) and drop the now-unused `functools.cached_property` import. Co-authored-by: Cursor * don't set_zero_position() on connect --------- Co-authored-by: Cursor --- src/lerobot/motors/damiao/damiao.py | 3 +-- src/lerobot/robots/openarm_follower/openarm_follower.py | 3 --- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/src/lerobot/motors/damiao/damiao.py b/src/lerobot/motors/damiao/damiao.py index 572741cb4..afbec85d3 100644 --- a/src/lerobot/motors/damiao/damiao.py +++ b/src/lerobot/motors/damiao/damiao.py @@ -20,7 +20,6 @@ import logging import time from contextlib import contextmanager from copy import deepcopy -from functools import cached_property from typing import TYPE_CHECKING, Any, TypedDict from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected @@ -854,7 +853,7 @@ class DamiaoMotorsBus(MotorsBusBase): else: raise ValueError(f"Motor {motor_obj} doesn't have a valid recv_id (None).") - @cached_property + @property def is_calibrated(self) -> bool: """Check if motors are calibrated.""" return bool(self.calibration) diff --git a/src/lerobot/robots/openarm_follower/openarm_follower.py b/src/lerobot/robots/openarm_follower/openarm_follower.py index e2c7c8cf5..5f2286eb4 100644 --- a/src/lerobot/robots/openarm_follower/openarm_follower.py +++ b/src/lerobot/robots/openarm_follower/openarm_follower.py @@ -150,9 +150,6 @@ class OpenArmFollower(Robot): self.configure() - if self.is_calibrated: - self.bus.set_zero_position() - self.bus.enable_torque() logger.info(f"{self} connected.") From f59eae4e27a49aa044a1146c5fa73762045678ba Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Thu, 23 Jul 2026 18:41:48 +0200 Subject: [PATCH 5/7] fix(robots): add retries while recording motor ranges (#4126) * Add retries while recording motor ranges * fix(motors): throttle calibration reads consistently --------- Co-authored-by: tom-doerr --- src/lerobot/motors/motors_bus.py | 14 +++++++++----- tests/motors/test_dynamixel.py | 12 +++++++++--- tests/motors/test_feetech.py | 12 +++++++++--- 3 files changed, 27 insertions(+), 11 deletions(-) diff --git a/src/lerobot/motors/motors_bus.py b/src/lerobot/motors/motors_bus.py index 4688eaa7f..f4a747b44 100644 --- a/src/lerobot/motors/motors_bus.py +++ b/src/lerobot/motors/motors_bus.py @@ -23,6 +23,7 @@ from __future__ import annotations import abc import logging +import time from collections.abc import Sequence from contextlib import contextmanager from dataclasses import dataclass @@ -818,13 +819,13 @@ class SerialMotorsBus(MotorsBusBase): """ motor_names = self._get_motors_list(motors) - start_positions = self.sync_read("Present_Position", motor_names, normalize=False) + start_positions = self.sync_read("Present_Position", motor_names, normalize=False, num_retry=5) mins = start_positions.copy() maxes = start_positions.copy() user_pressed_enter = False while not user_pressed_enter: - positions = self.sync_read("Present_Position", motor_names, normalize=False) + positions = self.sync_read("Present_Position", motor_names, normalize=False, num_retry=5) mins = {motor: min(positions[motor], min_) for motor, min_ in mins.items()} maxes = {motor: max(positions[motor], max_) for motor, max_ in maxes.items()} @@ -837,9 +838,12 @@ class SerialMotorsBus(MotorsBusBase): if enter_pressed(): user_pressed_enter = True - if display_values and not user_pressed_enter: - # Move cursor up to overwrite the previous output - move_cursor_up(len(motor_names) + 3) + if not user_pressed_enter: + if display_values: + # Move cursor up to overwrite the previous output + move_cursor_up(len(motor_names) + 3) + # Throttle reads even when the live table is disabled. + time.sleep(0.02) same_min_max = [motor for motor in motor_names if mins[motor] == maxes[motor]] if same_min_max: diff --git a/tests/motors/test_dynamixel.py b/tests/motors/test_dynamixel.py index 8b02d4330..9e58ef5e0 100644 --- a/tests/motors/test_dynamixel.py +++ b/tests/motors/test_dynamixel.py @@ -405,12 +405,18 @@ def test_record_ranges_of_motion(mock_motors, dummy_motors): read_pos_stub = mock_motors.build_sequential_sync_read_stub( *X_SERIES_CONTROL_TABLE["Present_Position"], positions ) - with patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]): - bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors) - bus.connect(handshake=False) + bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors) + bus.connect(handshake=False) + with ( + patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]), + patch("lerobot.motors.motors_bus.time.sleep") as mock_sleep, + patch.object(bus, "sync_read", wraps=bus.sync_read) as mock_sync_read, + ): mins, maxes = bus.record_ranges_of_motion(display_values=False) assert mock_motors.stubs[read_pos_stub].calls == 3 + assert all(call.kwargs["num_retry"] == 5 for call in mock_sync_read.call_args_list) + mock_sleep.assert_called_once_with(0.02) assert mins == expected_mins assert maxes == expected_maxes diff --git a/tests/motors/test_feetech.py b/tests/motors/test_feetech.py index e8cdddbbd..6fc9d684e 100644 --- a/tests/motors/test_feetech.py +++ b/tests/motors/test_feetech.py @@ -509,12 +509,18 @@ def test_record_ranges_of_motion(mock_motors, dummy_motors): stub = mock_motors.build_sequential_sync_read_stub( *STS_SMS_SERIES_CONTROL_TABLE["Present_Position"], positions ) - with patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]): - bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors) - bus.connect(handshake=False) + bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors) + bus.connect(handshake=False) + with ( + patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]), + patch("lerobot.motors.motors_bus.time.sleep") as mock_sleep, + patch.object(bus, "sync_read", wraps=bus.sync_read) as mock_sync_read, + ): mins, maxes = bus.record_ranges_of_motion(display_values=False) assert mock_motors.stubs[stub].calls == 3 + assert all(call.kwargs["num_retry"] == 5 for call in mock_sync_read.call_args_list) + mock_sleep.assert_called_once_with(0.02) assert mins == expected_mins assert maxes == expected_maxes From cfd9ff969ca91acf22a68c9f85d3bc2e05c1bc92 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Thu, 23 Jul 2026 19:49:19 +0200 Subject: [PATCH 6/7] fix(envs): set LiberoEnvConfig.fps default to 20 to match robosuite (#4124) * fix(envs): set LiberoEnvConfig.fps default to 20 to match robosuite LiberoEnvConfig.fps was set to 30, but the underlying robosuite OffScreenRenderEnv always runs at its default control_freq of 20 Hz since fps is never passed through. This mismatch silently decouples the dataset/eval loop rate from the actual simulation step rate. Set the default to 20 to match the real sim rate and avoid the footgun. Fixes #3368 * fix(libero): apply configured control frequency --------- Co-authored-by: xinmotlanthua <275663218+xinmotlanthua@users.noreply.github.com> --- src/lerobot/envs/configs.py | 6 +++++- src/lerobot/envs/libero.py | 5 +++++ tests/envs/test_dispatch.py | 11 +++++++++++ 3 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index 3624357e2..3f6fd75f9 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -322,7 +322,7 @@ class HILSerlRobotEnvConfig(EnvConfig): class LiberoEnv(EnvConfig): task: str = "libero_10" # can also choose libero_spatial, libero_object, etc. task_ids: list[int] | None = None - fps: int = 30 + fps: int = 20 # Must match robosuite's default control_freq (20 Hz) episode_length: int | None = None obs_type: str = "pixels_agent_pos" render_mode: str = "rgb_array" @@ -354,6 +354,9 @@ class LiberoEnv(EnvConfig): control_mode: str = "relative" # or "absolute" def __post_init__(self): + if self.fps <= 0: + raise ValueError(f"fps must be positive, got {self.fps}") + if self.obs_type == "pixels": self.features[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature( type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3) @@ -412,6 +415,7 @@ class LiberoEnv(EnvConfig): "render_mode": self.render_mode, "observation_height": self.observation_height, "observation_width": self.observation_width, + "control_freq": self.fps, } if self.task_ids is not None: kwargs["task_ids"] = self.task_ids diff --git a/src/lerobot/envs/libero.py b/src/lerobot/envs/libero.py index 12be9e196..958723277 100644 --- a/src/lerobot/envs/libero.py +++ b/src/lerobot/envs/libero.py @@ -125,10 +125,13 @@ class LiberoEnv(gym.Env): n_envs: int = 1, camera_name_mapping: dict[str, str] | None = None, num_steps_wait: int = 10, + control_freq: int = 20, control_mode: str = "relative", is_libero_plus: bool = False, ): super().__init__() + if control_freq <= 0: + raise ValueError(f"control_freq must be positive, got {control_freq}") self.task_id = task_id self.is_libero_plus = is_libero_plus self.obs_type = obs_type @@ -154,6 +157,7 @@ class LiberoEnv(gym.Env): } self.camera_name_mapping = camera_name_mapping self.num_steps_wait = num_steps_wait + self.control_freq = control_freq self.episode_index = episode_index self.episode_length = episode_length # Load once and keep @@ -260,6 +264,7 @@ class LiberoEnv(gym.Env): bddl_file_name=self._task_bddl_file, camera_heights=self.observation_height, camera_widths=self.observation_width, + control_freq=self.control_freq, ) env.reset() self._env = env diff --git a/tests/envs/test_dispatch.py b/tests/envs/test_dispatch.py index 50f208422..9e23410dc 100644 --- a/tests/envs/test_dispatch.py +++ b/tests/envs/test_dispatch.py @@ -35,6 +35,17 @@ def test_unknown_type(): make_env_config("nonexistent") +def test_libero_fps_controls_simulator_frequency(): + cfg = LiberoEnv(fps=17) + + assert cfg.gym_kwargs["control_freq"] == 17 + + +def test_libero_rejects_nonpositive_fps(): + with pytest.raises(ValueError, match="fps must be positive"): + LiberoEnv(fps=0) + + def test_identity_processors(): """Base class get_env_processors() returns identity pipelines.""" cfg = make_env_config("aloha") From a0eb860d1e4d566157dd0eba11b5c86f9ea7d07a Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Thu, 23 Jul 2026 22:05:29 +0200 Subject: [PATCH 7/7] feat(dataset): add slice support to LeRobotDataset.__getitem__ (#4129) * feat(dataset): add efficient slice support * fix(dataset): handle empty dataset slices * refactor(dataset): reuse scalar path for slices --------- Co-authored-by: Francesco Capuano --- src/lerobot/datasets/lerobot_dataset.py | 14 +++++--- tests/datasets/test_datasets.py | 46 +++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 5 deletions(-) diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index aba86efe3..77b3032df 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -478,18 +478,19 @@ class LeRobotDataset(torch.utils.data.Dataset): """Return the number of frames in the selected episodes.""" return self.num_frames - def __getitem__(self, idx) -> dict: - """Return a single frame by index, with all transforms applied. + def __getitem__(self, idx: int | slice) -> dict | list[dict]: + """Return one frame or a slice of frames, with all transforms applied. Loads the frame from the underlying HF dataset, expands delta-timestamp windows, decodes video frames, and applies image transforms. Delegates - the core logic to :meth:`DatasetReader.get_item`. + the core logic to :class:`DatasetReader`. Args: - idx: Index into the (possibly episode-filtered) dataset. + idx: Integer index or slice into the possibly episode-filtered dataset. Returns: - Dict mapping feature names to their tensor values for this frame. + A frame dictionary for an integer index, or a list of frame + dictionaries for a slice. Raises: RuntimeError: If the dataset is currently being recorded and @@ -499,6 +500,9 @@ class LeRobotDataset(torch.utils.data.Dataset): raise RuntimeError( "Cannot read from a dataset that is being recorded. Call finalize() first, then access items." ) + if isinstance(idx, slice): + return [self[item_idx] for item_idx in range(*idx.indices(len(self)))] + reader = self._ensure_reader() if reader.hf_dataset is None: # One-shot load after finalize() diff --git a/tests/datasets/test_datasets.py b/tests/datasets/test_datasets.py index 97f799d9f..d76ed2e1d 100644 --- a/tests/datasets/test_datasets.py +++ b/tests/datasets/test_datasets.py @@ -114,6 +114,20 @@ def test_dataset_initialization(tmp_path, lerobot_dataset_factory): assert dataset.num_frames == len(dataset) +def test_dataset_slice(tmp_path, lerobot_dataset_factory): + dataset = lerobot_dataset_factory( + root=tmp_path / "test", total_episodes=3, total_frames=30, use_videos=False + ) + + assert len(dataset[:5]) == 5 + assert len(dataset[::2]) == (len(dataset) + 1) // 2 + assert [item["index"].item() for item in dataset[4::-1]] == [4, 3, 2, 1, 0] + assert [item["index"].item() for item in dataset[-3:]] == list(range(len(dataset) - 3, len(dataset))) + assert dataset[len(dataset) :] == [] + assert isinstance(dataset[0], dict) + assert dataset[:1][0].keys() == dataset[0].keys() + + # TODO(rcadene, aliberts): do not run LeRobotDataset.create, instead refactor LeRobotDatasetMetadata.create # and test the small resulting function that validates the features def test_dataset_feature_with_forward_slash_raises_error(): @@ -1741,6 +1755,38 @@ def test_delta_timestamps_query_returns_correct_values(tmp_path, empty_lerobot_d assert is_pad == [True, False], f"Expected [True, False], got {is_pad}" +def test_dataset_slice_with_delta_timestamps(tmp_path, empty_lerobot_dataset_factory): + features = { + "observation.state": {"dtype": "float32", "shape": (1,), "names": ["x"]}, + } + dataset = empty_lerobot_dataset_factory( + root=tmp_path / "test_slice_delta", features=features, use_videos=False, fps=10 + ) + + for frame_idx in range(5): + dataset.add_frame( + { + "observation.state": torch.tensor([frame_idx], dtype=torch.float32), + "task": "task_0", + } + ) + dataset.save_episode() + dataset.finalize() + + sliced_dataset = LeRobotDataset( + dataset.repo_id, + root=dataset.root, + delta_timestamps={"observation.state": [-0.1, 0.0]}, + tolerance_s=0.04, + ) + + items = sliced_dataset[:2] + + assert items[0]["observation.state"].tolist() == [0.0, 0.0] + assert items[0]["observation.state_is_pad"].tolist() == [True, False] + assert items[1]["observation.state"].tolist() == [0.0, 1.0] + + def test_episode_filter_filters_dataset(tmp_path, lerobot_dataset_factory): """episode_filter on LeRobotDataset narrows the loaded dataset to matching episodes.""" dataset = lerobot_dataset_factory(root=tmp_path / "test", total_episodes=8, total_frames=200)