Compare commits

...

4 Commits

Author SHA1 Message Date
Maxime Ellerbach 2de042690e fix: pre-commit auto-fix (prettier markdown table formatting) 2026-07-27 09:40:20 +00:00
griffinaddison 124d03608c feat(rollout): add smooth_handover flag to DAgger strategy config
The DAgger phase transitions run blocking smooth handovers: on pause the
leader is driven to the follower (~2 s), and on correction start the
follower is slid to the teleop pose (~1 s), both inside the record loop.

For clutch-style teleoperators (e.g. VR controllers) that re-reference
their command frame at the current robot pose on engage, the handover is
already continuous — the interpolation only delays the start of the
correction and eats its first frames.

Add --strategy.smooth_handover (default true, existing behavior
unchanged) to let such setups skip it, mirroring the episodic strategy's
smooth_leader_to_follower_handover flag.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-27 09:38:22 +00:00
Steven Palma 0d383d09f2 feat(dataset): accept token argument for private HF Hub datasets (#4136) 2026-07-24 18:51:35 +02:00
Caroline Pascal ab2b5b04dd (depth image processing): excluding depth frames from the RGB to BGR image processing (#4135)
* (depth image processing): excluding depth frames from the RGB to BGR image processing

* test(update): updating tests to include RGB/BGR conversion checks
2026-07-24 17:43:17 +02:00
17 changed files with 315 additions and 42 deletions
+8 -7
View File
@@ -149,13 +149,14 @@ lerobot-rollout \
Foot pedal input is also supported via `--strategy.input_device=pedal`. Configure pedal codes with `--strategy.pedal.*` flags.
| Flag | Description |
| ------------------------------------ | ------------------------------------------------------- |
| `--strategy.num_episodes` | Number of correction episodes to record (default: 10) |
| `--strategy.record_autonomous` | Record autonomous frames too (default: false) |
| `--strategy.upload_every_n_episodes` | Push to Hub every N episodes (default: 5) |
| `--strategy.input_device` | Input device: `keyboard` or `pedal` (default: keyboard) |
| `--teleop.type` | **Required.** Teleoperator type |
| Flag | Description |
| ------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `--strategy.num_episodes` | Number of correction episodes to record (default: 10) |
| `--strategy.record_autonomous` | Record autonomous frames too (default: false) |
| `--strategy.upload_every_n_episodes` | Push to Hub every N episodes (default: 5) |
| `--strategy.input_device` | Input device: `keyboard` or `pedal` (default: keyboard) |
| `--strategy.smooth_handover` | Smoothly hand control over at pause / correction start (default: true). Disable for clutch-style teleops that re-reference at the current robot pose on engage |
| `--teleop.type` | **Required.** Teleoperator type |
### Episodic (`--strategy.type=episodic`)
+4
View File
@@ -252,6 +252,10 @@ lerobot-dataset-viz \
--episode-index 0
```
For a private or gated dataset, authenticate first with `hf auth login`, or set the
`HF_TOKEN` environment variable. The Hub client then discovers the credential
automatically; no token argument is needed.
**From a local folder:**
Add the `--root` option and set `--mode local`. For example, to search in `./my_local_data_dir/lerobot/pusht`:
@@ -173,7 +173,8 @@ class Reachy2Camera(Camera):
raise ValueError(
f"Invalid color mode '{self.color_mode}'. Expected {ColorMode.RGB} or {ColorMode.BGR}."
)
if self.color_mode == ColorMode.RGB:
is_depth_frame = self.config.name == "depth" and self.config.image_type == "depth"
if not is_depth_frame and self.color_mode == ColorMode.RGB:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
self.latest_frame = frame
@@ -453,7 +453,7 @@ class RealSenseCamera(Camera):
)
processed_image = image
if self.color_mode == ColorMode.BGR:
if not depth_frame and self.color_mode == ColorMode.BGR:
processed_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE, cv2.ROTATE_180]:
+16 -2
View File
@@ -73,6 +73,8 @@ class LeRobotDatasetMetadata:
revision: str | None = None,
force_cache_sync: bool = False,
metadata_buffer_size: int = 10,
*,
token: str | bool | None = None,
):
"""Load or download metadata for an existing LeRobot dataset.
@@ -94,6 +96,10 @@ class LeRobotDatasetMetadata:
even when local files exist.
metadata_buffer_size: Number of episode metadata records to buffer
in memory before flushing to parquet.
token: Authentication token used for Hub requests. Pass a string
token, ``True`` to require the locally stored token, ``False``
to disable authentication, or ``None`` to use the Hugging Face
Hub default.
"""
self.repo_id = repo_id
self.revision = revision if revision else CODEBASE_VERSION
@@ -113,9 +119,12 @@ class LeRobotDatasetMetadata:
self._load_metadata()
except (FileNotFoundError, NotADirectoryError):
if is_valid_version(self.revision):
self.revision = get_safe_version(self.repo_id, self.revision)
if token is None:
self.revision = get_safe_version(self.repo_id, self.revision)
else:
self.revision = get_safe_version(self.repo_id, self.revision, token=token)
self._pull_from_repo(allow_patterns="meta/")
self._pull_from_repo(allow_patterns="meta/", token=token)
self._load_metadata()
def _flush_metadata_buffer(self) -> None:
@@ -220,7 +229,10 @@ class LeRobotDatasetMetadata:
self,
allow_patterns: list[str] | str | None = None,
ignore_patterns: list[str] | str | None = None,
*,
token: str | bool | None = None,
) -> None:
token_kwargs = {} if token is None else {"token": token}
if self._requested_root is None:
self.root = Path(
snapshot_download(
@@ -230,6 +242,7 @@ class LeRobotDatasetMetadata:
cache_dir=HF_LEROBOT_HUB_CACHE,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
**token_kwargs,
)
)
return
@@ -242,6 +255,7 @@ class LeRobotDatasetMetadata:
local_dir=self._requested_root,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
**token_kwargs,
)
self.root = self._requested_root
+30 -5
View File
@@ -65,6 +65,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
encoder_threads: int | None = None,
streaming_encoding: bool = False,
encoder_queue_maxsize: int = 30,
*,
token: str | bool | None = None,
):
"""
2 modes are available for instantiating this class, depending on 2 different use cases:
@@ -197,6 +199,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
instead of writing PNG images first. This makes save_episode() near-instant. Defaults to False.
encoder_queue_maxsize (int, optional): Maximum number of frames to buffer per camera when using
streaming encoding. Defaults to 30 (~1s at 30fps).
token: Authentication token used while downloading this dataset
from the Hub. Pass a string token, ``True`` to require the
locally stored token, ``False`` to disable authentication, or
``None`` to use the Hugging Face Hub default. The token is not
retained on the dataset instance after initialization.
Note:
Write-mode parameters (``streaming_encoding``, ``batch_encoding_size``) passed to
@@ -220,7 +227,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Load metadata (sets self.root once from the resolved metadata root)
self.meta = LeRobotDatasetMetadata(
self.repo_id, self._requested_root, self.revision, force_cache_sync=force_cache_sync
self.repo_id,
self._requested_root,
self.revision,
force_cache_sync=force_cache_sync,
token=token,
)
self.root = self.meta.root
self.revision = self.meta.revision
@@ -260,8 +271,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Load actual data
if force_cache_sync or not self.reader.try_load():
if is_valid_version(self.revision):
self.revision = get_safe_version(self.repo_id, self.revision)
self._download(download_videos)
if token is None:
self.revision = get_safe_version(self.repo_id, self.revision)
else:
self.revision = get_safe_version(self.repo_id, self.revision, token=token)
self._download(download_videos, token=token)
self.reader.load_and_activate()
# Detect write-mode params for backward compatibility
@@ -626,10 +640,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
hub_api.delete_tag(self.repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
hub_api.create_tag(self.repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
def _download(self, download_videos: bool = True) -> None:
def _download(self, download_videos: bool = True, *, token: str | bool | None = None) -> None:
"""Downloads the dataset from the given 'repo_id' at the provided version."""
ignore_patterns = None if download_videos else "videos/"
files = None
token_kwargs = {} if token is None else {"token": token}
if self.episodes is not None:
# Reader is guaranteed to exist here (created in __init__ before _download)
files = self.reader.get_episodes_file_paths()
@@ -643,6 +658,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
cache_dir=HF_LEROBOT_HUB_CACHE,
allow_patterns=files,
ignore_patterns=ignore_patterns,
**token_kwargs,
)
)
else:
@@ -654,6 +670,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
local_dir=self._requested_root,
allow_patterns=files,
ignore_patterns=ignore_patterns,
**token_kwargs,
)
self.meta.root = self._requested_root
@@ -793,6 +810,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
image_writer_threads: int = 0,
streaming_encoding: bool = False,
encoder_queue_maxsize: int = 30,
*,
token: str | bool | None = None,
) -> "LeRobotDataset":
"""Resume recording on an existing dataset.
@@ -826,6 +845,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
streaming_encoding: If ``True``, encode video in real-time during
capture.
encoder_queue_maxsize: Max buffered frames per camera for streaming.
token: Authentication token used if metadata must be downloaded
from the Hub. The token is not retained on the dataset instance.
Returns:
A :class:`LeRobotDataset` in write mode, ready to append episodes.
@@ -854,7 +875,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Load metadata (revision-safe when root is not provided)
obj.meta = LeRobotDatasetMetadata(
obj.repo_id, obj._requested_root, obj.revision, force_cache_sync=force_cache_sync
obj.repo_id,
obj._requested_root,
obj.revision,
force_cache_sync=force_cache_sync,
token=token,
)
obj._encoder_threads = encoder_threads
+3
View File
@@ -48,6 +48,8 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
tolerances_s: dict | None = None,
download_videos: bool = True,
video_backend: str | None = None,
*,
token: str | bool | None = None,
):
super().__init__()
self.repo_ids = repo_ids
@@ -65,6 +67,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
tolerance_s=self.tolerances_s[repo_id],
download_videos=download_videos,
video_backend=video_backend,
token=token,
)
for repo_id in repo_ids
]
+14 -1
View File
@@ -256,6 +256,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
shuffle: bool = True,
return_uint8: bool = False,
depth_output_unit: str = DEFAULT_DEPTH_UNIT,
*,
token: str | bool | None = None,
):
"""Initialize a StreamingLeRobotDataset.
@@ -278,6 +280,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm").
Defaults to "mm".
token: Authentication token used while streaming this dataset from
the Hub. Pass a string token, ``True`` to require the locally
stored token, ``False`` to disable authentication, or ``None``
to use the Hugging Face Hub default. The token is not retained
on the dataset instance after initialization.
"""
super().__init__()
self.repo_id = repo_id
@@ -306,7 +313,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
# Load metadata
self.meta = LeRobotDatasetMetadata(
self.repo_id, self._requested_root, self.revision, force_cache_sync=force_cache_sync
self.repo_id,
self._requested_root,
self.revision,
force_cache_sync=force_cache_sync,
token=token,
)
self.root = self.meta.root
self.revision = self.meta.revision
@@ -334,12 +345,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
self.delta_timestamps = delta_timestamps
self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)
token_kwargs = {} if token is None or self.streaming_from_local else {"token": token}
self.hf_dataset: datasets.IterableDataset = load_dataset(
self.repo_id if not self.streaming_from_local else str(self.root),
split="train",
streaming=self.streaming,
data_files="data/*/*.parquet",
revision=self.revision,
**token_kwargs,
)
self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
+13 -4
View File
@@ -325,16 +325,19 @@ def check_version_compatibility(
logging.warning(FUTURE_MESSAGE.format(repo_id=repo_id, version=v_check))
def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
def get_repo_versions(repo_id: str, *, token: str | bool | None = None) -> list[packaging.version.Version]:
"""Return available valid versions (branches and tags) on a given Hub repo.
Args:
repo_id (str): The repository ID on the Hugging Face Hub.
token: Authentication token used for Hub requests. Pass a string token,
``True`` to require the locally stored token, ``False`` to disable
authentication, or ``None`` to use the Hugging Face Hub default.
Returns:
list[packaging.version.Version]: A list of valid versions found.
"""
api = HfApi()
api = HfApi() if token is None else HfApi(token=token)
repo_refs = api.list_repo_refs(repo_id, repo_type="dataset")
repo_refs = [b.name for b in repo_refs.branches + repo_refs.tags]
repo_versions = []
@@ -345,7 +348,12 @@ def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
return repo_versions
def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> str:
def get_safe_version(
repo_id: str,
version: str | packaging.version.Version,
*,
token: str | bool | None = None,
) -> str:
"""Return the specified version if available on repo, or the latest compatible one.
If the exact version is not found, it looks for the latest version with the
@@ -354,6 +362,7 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
Args:
repo_id (str): The repository ID on the Hugging Face Hub.
version (str | packaging.version.Version): The target version.
token: Authentication token forwarded to the Hub version lookup.
Returns:
str: The safe version string (e.g., "v1.2.3") to use as a revision.
@@ -366,7 +375,7 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
target_version = (
packaging.version.parse(version) if not isinstance(version, packaging.version.Version) else version
)
hub_versions = get_repo_versions(repo_id)
hub_versions = get_repo_versions(repo_id) if token is None else get_repo_versions(repo_id, token=token)
if not hub_versions:
raise RevisionNotFoundError(
+8
View File
@@ -180,6 +180,14 @@ class DAggerStrategyConfig(RolloutStrategyConfig):
# Target video file size in MB for episode rotation (record_autonomous
# mode only). Defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB when None.
target_video_file_size_mb: int | None = None
# Whether to turn on or off the smooth handover behavior at phase transitions:
# the leader is driven to the follower position on pause (teleops with
# `send_feedback` capability), and the follower is slid to the teleop pose when
# a correction starts (non-actuated teleops). Disable for clutch-style
# teleoperators (e.g. VR controllers) that re-reference at the current robot
# pose on engage: the handover is already continuous there, and the blocking
# interpolation only delays the start of the correction.
smooth_handover: bool = True
input_device: str = "keyboard"
keyboard: DAggerKeyboardConfig = field(default_factory=DAggerKeyboardConfig)
pedal: DAggerPedalConfig = field(default_factory=DAggerPedalConfig)
+11 -3
View File
@@ -623,8 +623,8 @@ class DAggerStrategy(RolloutStrategy):
# State-machine transition side-effects
# ------------------------------------------------------------------
@staticmethod
def _apply_transition(
self,
old_phase: DAggerPhase,
new_phase: DAggerPhase,
engine,
@@ -634,6 +634,10 @@ class DAggerStrategy(RolloutStrategy):
) -> None:
"""Execute side-effects for a validated phase transition, including smooth handovers.
The smooth handovers below can be disabled with
``--strategy.smooth_handover=false`` (useful for clutch-style teleops
that re-reference at the current robot pose on engage).
AUTONOMOUS -> PAUSED (actuated teleop):
Pause the engine, then drive the leader arm to the follower's last
commanded position so the operator takes over without a jerk.
@@ -657,7 +661,7 @@ class DAggerStrategy(RolloutStrategy):
logger.info("Pausing engine - robot holds position")
engine.pause()
if teleop_supports_feedback(teleop) and prev_action is not None:
if self.config.smooth_handover and teleop_supports_feedback(teleop) and prev_action is not None:
# TODO(Maxime): prev_action is in robot action key space (output of robot_action_processor).
# send_feedback expects teleop feedback key space. For homogeneous setups (e.g. SO-101
# leader + SO-101 follower) the keys are identical so this works. If the processor pipeline
@@ -668,7 +672,11 @@ class DAggerStrategy(RolloutStrategy):
elif old_phase == DAggerPhase.PAUSED and new_phase == DAggerPhase.CORRECTING:
logger.info("Entering correction mode - human teleop control")
if not teleop_supports_feedback(teleop) and prev_action is not None:
if (
self.config.smooth_handover
and not teleop_supports_feedback(teleop)
and prev_action is not None
):
logger.info("Smooth handover: sliding follower to teleop position")
obs = robot.get_observation()
teleop_action = teleop.get_action()
+23 -1
View File
@@ -26,7 +26,7 @@ import cv2
import numpy as np
import pytest
from lerobot.cameras.configs import Cv2Rotation
from lerobot.cameras.configs import ColorMode, Cv2Rotation
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
@@ -132,6 +132,28 @@ def test_read(index_or_path):
assert isinstance(img, np.ndarray)
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
def test_color_mode_conversion(index_or_path):
"""RGB and BGR reads of the same frame must differ only by a channel-axis reversal."""
rgb_config = OpenCVCameraConfig(index_or_path=index_or_path, color_mode=ColorMode.RGB, warmup_s=0)
bgr_config = OpenCVCameraConfig(index_or_path=index_or_path, color_mode=ColorMode.BGR, warmup_s=0)
with OpenCVCamera(rgb_config) as rgb_cam:
rgb = rgb_cam.read()
with OpenCVCamera(bgr_config) as bgr_cam:
bgr = bgr_cam.read()
assert rgb.shape == bgr.shape
np.testing.assert_array_equal(rgb, bgr[..., ::-1])
def test_postprocess_invalid_color_mode():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
camera = OpenCVCamera(config)
camera.color_mode = "invalid"
with pytest.raises(ValueError):
camera._postprocess_image(np.zeros((120, 160, 3), dtype=np.uint8))
def test_read_before_connect():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
+44 -15
View File
@@ -22,6 +22,7 @@ import pytest
pytest.importorskip("reachy2_sdk")
from lerobot.cameras.configs import ColorMode
from lerobot.cameras.reachy2_camera import Reachy2Camera, Reachy2CameraConfig
from lerobot.utils.errors import DeviceNotConnectedError
@@ -33,28 +34,19 @@ PARAMS = [
]
def _make_cam_manager_mock():
def _make_cam_manager_mock(color_frame, depth_frame=None):
c = MagicMock(name="CameraManagerMock")
teleop = MagicMock(name="TeleopCam")
teleop.width = 640
teleop.height = 480
teleop.get_frame = MagicMock(
side_effect=lambda *_, **__: (
np.zeros((480, 640, 3), dtype=np.uint8),
time.time(),
)
)
teleop.get_frame = MagicMock(side_effect=lambda *_, **__: (color_frame, time.time()))
depth = MagicMock(name="DepthCam")
depth.width = 640
depth.height = 480
depth.get_frame = MagicMock(
side_effect=lambda *_, **__: (
np.zeros((480, 640, 3), dtype=np.uint8),
time.time(),
)
)
depth.get_frame = MagicMock(side_effect=lambda *_, **__: (color_frame, time.time()))
depth.get_depth_frame = MagicMock(side_effect=lambda *_, **__: (depth_frame, time.time()))
c.is_connected.return_value = True
c.teleop = teleop
@@ -84,12 +76,14 @@ def _make_cam_manager_mock():
# ids=["teleop-left", "teleop-right", "torso-rgb", "torso-depth"],
ids=["teleop-left", "teleop-right", "torso-rgb"],
)
def camera(request):
def camera(request, img_array_factory):
name, image_type = request.param
color_frame = img_array_factory(height=480, width=640)
depth_frame = img_array_factory(height=480, width=640, channels=1, dtype=np.uint16)[..., 0]
with (
patch(
"lerobot.cameras.reachy2_camera.reachy2_camera.CameraManager",
side_effect=lambda *a, **k: _make_cam_manager_mock(),
side_effect=lambda *a, **k: _make_cam_manager_mock(color_frame, depth_frame),
),
):
config = Reachy2CameraConfig(name=name, image_type=image_type)
@@ -188,6 +182,41 @@ def test_read_latest_too_old(camera):
_ = camera.read_latest(max_age_ms=0) # immediately too old
def test_color_mode_conversion(img_array_factory):
"""teleop frames are native BGR: RGB reverses the channel axis, BGR is passed through."""
frame = img_array_factory(height=8, width=8)
outputs = {}
for color_mode in (ColorMode.RGB, ColorMode.BGR):
with patch(
"lerobot.cameras.reachy2_camera.reachy2_camera.CameraManager",
side_effect=lambda *a, **k: _make_cam_manager_mock(frame),
):
cam = Reachy2Camera(Reachy2CameraConfig(name="teleop", image_type="left", color_mode=color_mode))
cam.connect()
outputs[color_mode] = cam.read()
cam.disconnect()
np.testing.assert_array_equal(outputs[ColorMode.BGR], frame)
np.testing.assert_array_equal(outputs[ColorMode.RGB], frame[..., ::-1])
def test_depth_frame_not_color_converted(img_array_factory):
"""A depth/depth frame must be returned as-is, without BGR<->RGB conversion."""
color_frame = img_array_factory(height=8, width=8)
depth = img_array_factory(height=8, width=8, channels=1, dtype=np.uint16)[..., 0]
with patch(
"lerobot.cameras.reachy2_camera.reachy2_camera.CameraManager",
side_effect=lambda *a, **k: _make_cam_manager_mock(color_frame, depth_frame=depth),
):
cam = Reachy2Camera(Reachy2CameraConfig(name="depth", image_type="depth"))
cam.connect()
out = cam.read()
cam.disconnect()
np.testing.assert_array_equal(out, depth)
def test_wrong_camera_name():
with pytest.raises(ValueError):
_ = Reachy2CameraConfig(name="wrong-name", image_type="left")
+27 -1
View File
@@ -25,7 +25,7 @@ from unittest.mock import patch
import numpy as np
import pytest
from lerobot.cameras.configs import Cv2Rotation
from lerobot.cameras.configs import ColorMode, Cv2Rotation
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
pytest.importorskip("pyrealsense2")
@@ -109,6 +109,32 @@ def test_read_depth():
assert isinstance(img, np.ndarray)
# These exercise _postprocess_image directly rather than read(): the bag playback returns
# non-deterministic frames we can't compare against, and the depth read() path is skipped
# (see test_read_depth) with the current pyrealsense2 version.
def test_color_mode_conversion(img_array_factory):
"""RGB (native for RealSense) is passed through; BGR reverses the channel axis."""
color = img_array_factory(height=3, width=4)
outputs = {}
for color_mode in (ColorMode.RGB, ColorMode.BGR):
camera = RealSenseCamera(RealSenseCameraConfig(serial_number_or_name="042", color_mode=color_mode))
camera.capture_height, camera.capture_width = color.shape[:2]
outputs[color_mode] = camera._postprocess_image(color)
np.testing.assert_array_equal(outputs[ColorMode.RGB], color)
np.testing.assert_array_equal(outputs[ColorMode.BGR], color[..., ::-1])
def test_depth_frame_not_color_converted(img_array_factory):
"""Depth frames must bypass color conversion, even when a BGR color_mode is set."""
camera = RealSenseCamera(RealSenseCameraConfig(serial_number_or_name="042", color_mode=ColorMode.BGR))
depth = img_array_factory(height=3, width=4, channels=1, dtype=np.uint16)[..., 0]
camera.capture_height, camera.capture_width = depth.shape
np.testing.assert_array_equal(camera._postprocess_image(depth, depth_frame=True), depth)
def test_read_before_connect():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
+30 -1
View File
@@ -14,16 +14,21 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
from packaging.version import Version
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from datasets import Dataset # noqa: E402
from huggingface_hub import DatasetCard
import lerobot.datasets.utils as dataset_utils
from lerobot.datasets.io_utils import hf_transform_to_torch
from lerobot.datasets.utils import create_lerobot_dataset_card
from lerobot.datasets.utils import create_lerobot_dataset_card, get_repo_versions, get_safe_version
from lerobot.utils.constants import ACTION, OBS_IMAGES
from lerobot.utils.feature_utils import combine_feature_dicts
@@ -57,6 +62,30 @@ def test_default_parameters():
]
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
def test_get_repo_versions_forwards_token(monkeypatch, token):
api = Mock()
api.list_repo_refs.return_value = SimpleNamespace(
branches=[SimpleNamespace(name="v3.0")],
tags=[],
)
hf_api = Mock(return_value=api)
monkeypatch.setattr(dataset_utils, "HfApi", hf_api)
assert get_repo_versions("private/repo", token=token) == [Version("3.0")]
hf_api.assert_called_once_with(token=token)
api.list_repo_refs.assert_called_once_with("private/repo", repo_type="dataset")
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
def test_get_safe_version_forwards_token(monkeypatch, token):
get_versions = Mock(return_value=[Version("3.0")])
monkeypatch.setattr(dataset_utils, "get_repo_versions", get_versions)
assert get_safe_version("private/repo", "v3.0", token=token) == "v3.0"
get_versions.assert_called_once_with("private/repo", token=token)
def test_with_tags():
tags = ["tag1", "tag2"]
card = create_lerobot_dataset_card(tags=tags)
+43
View File
@@ -20,6 +20,7 @@ property delegation, and the full create-record-finalize-read lifecycle.
"""
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
@@ -191,6 +192,48 @@ def test_metadata_without_root_uses_hub_cache_snapshot_download(
}
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
def test_metadata_download_forwards_token(tmp_path, monkeypatch, token):
snapshot_root = tmp_path / "snapshot"
snapshot_download = Mock(return_value=str(snapshot_root))
get_safe_version = Mock(return_value="v3.0")
load_metadata = Mock(side_effect=[FileNotFoundError, None])
monkeypatch.setattr(dataset_metadata_module, "snapshot_download", snapshot_download)
monkeypatch.setattr(dataset_metadata_module, "get_safe_version", get_safe_version)
monkeypatch.setattr(LeRobotDatasetMetadata, "_load_metadata", load_metadata)
meta = LeRobotDatasetMetadata(
repo_id=DUMMY_REPO_ID,
revision="v3.0",
token=token,
)
assert meta.root == snapshot_root
assert not hasattr(meta, "_token")
get_safe_version.assert_called_once_with(DUMMY_REPO_ID, "v3.0", token=token)
assert snapshot_download.call_args.kwargs["token"] is token
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
def test_data_download_forwards_token(tmp_path, monkeypatch, token):
snapshot_root = tmp_path / "snapshot"
snapshot_download = Mock(return_value=str(snapshot_root))
monkeypatch.setattr(lerobot_dataset_module, "snapshot_download", snapshot_download)
dataset = LeRobotDataset.__new__(LeRobotDataset)
dataset.repo_id = DUMMY_REPO_ID
dataset.revision = "main"
dataset.episodes = None
dataset._requested_root = None
dataset.meta = SimpleNamespace(root=None)
dataset.reader = SimpleNamespace(root=None)
dataset._download(token=token)
assert dataset.root == snapshot_root
assert snapshot_download.call_args.kwargs["token"] is token
def test_without_root_reads_different_revisions_from_distinct_snapshot_roots(
tmp_path,
info_factory,
+38
View File
@@ -13,12 +13,16 @@
# 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 Mock
import numpy as np
import pytest
import torch
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
import lerobot.datasets.streaming_dataset as streaming_dataset_module
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
from lerobot.datasets.utils import safe_shard
from lerobot.utils.constants import ACTION
@@ -71,6 +75,40 @@ def get_frames_expected_order(streaming_ds: StreamingLeRobotDataset) -> list[int
return expected_indices
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
@pytest.mark.parametrize("from_local", [False, True])
def test_streaming_dataset_forwards_hub_token_only_for_remote_data(tmp_path, monkeypatch, token, from_local):
requested_root = tmp_path / "local" if from_local else None
metadata = SimpleNamespace(
root=requested_root or tmp_path / "snapshot",
revision=streaming_dataset_module.CODEBASE_VERSION,
_version=streaming_dataset_module.CODEBASE_VERSION,
features={},
depth_keys=[],
image_keys=[],
rescale_depth_stats=Mock(),
)
metadata_cls = Mock(return_value=metadata)
load_dataset = Mock(return_value=SimpleNamespace(num_shards=1))
monkeypatch.setattr(streaming_dataset_module, "LeRobotDatasetMetadata", metadata_cls)
monkeypatch.setattr(streaming_dataset_module, "load_dataset", load_dataset)
dataset = StreamingLeRobotDataset(DUMMY_REPO_ID, root=requested_root, token=token)
metadata_cls.assert_called_once_with(
DUMMY_REPO_ID,
requested_root,
streaming_dataset_module.CODEBASE_VERSION,
force_cache_sync=False,
token=token,
)
if from_local:
assert "token" not in load_dataset.call_args.kwargs
else:
assert load_dataset.call_args.kwargs["token"] is token
assert not hasattr(dataset, "_token")
def test_single_frame_consistency(tmp_path, lerobot_dataset_factory):
"""Test if are correctly accessed"""
ds_num_frames = 400