mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-27 11:46:04 +00:00
Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 3011e18c06 | |||
| 0d383d09f2 | |||
| ab2b5b04dd |
@@ -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]:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -177,6 +177,7 @@ def make_pre_post_processors(
|
||||
return make_groot_pre_post_processors_from_pretrained(
|
||||
config=policy_cfg,
|
||||
pretrained_path=pretrained_path,
|
||||
revision=pretrained_revision,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
preprocessor_overrides=kwargs.get("preprocessor_overrides"),
|
||||
|
||||
@@ -475,6 +475,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
||||
config: GrootConfig,
|
||||
pretrained_path: str,
|
||||
*,
|
||||
revision: str | None = None,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
dataset_meta: Any | None = None,
|
||||
preprocessor_overrides: dict[str, Any] | None = None,
|
||||
@@ -511,6 +512,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
||||
|
||||
preprocessor, postprocessor = _load_groot_processor_pipelines(
|
||||
pretrained_path,
|
||||
revision=revision,
|
||||
preprocessor_overrides=preprocessor_overrides,
|
||||
postprocessor_overrides=postprocessor_overrides,
|
||||
preprocessor_config_filename=preprocessor_config_filename,
|
||||
@@ -526,6 +528,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
||||
def _load_groot_processor_pipelines(
|
||||
pretrained_path: str,
|
||||
*,
|
||||
revision: str | None,
|
||||
preprocessor_overrides: dict[str, Any],
|
||||
postprocessor_overrides: dict[str, Any],
|
||||
preprocessor_config_filename: str,
|
||||
@@ -540,6 +543,7 @@ def _load_groot_processor_pipelines(
|
||||
preprocessor = PolicyProcessorPipeline.from_pretrained(
|
||||
pretrained_model_name_or_path=pretrained_path,
|
||||
config_filename=preprocessor_config_filename,
|
||||
revision=revision,
|
||||
overrides=preprocessor_overrides,
|
||||
to_transition=batch_to_transition,
|
||||
to_output=transition_to_batch,
|
||||
@@ -547,6 +551,7 @@ def _load_groot_processor_pipelines(
|
||||
postprocessor = PolicyProcessorPipeline.from_pretrained(
|
||||
pretrained_model_name_or_path=pretrained_path,
|
||||
config_filename=postprocessor_config_filename,
|
||||
revision=revision,
|
||||
overrides=postprocessor_overrides,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
|
||||
@@ -326,8 +326,17 @@ class RolloutConfig:
|
||||
|
||||
policy_path = parser.get_path_arg("policy")
|
||||
if policy_path:
|
||||
cli_overrides = parser.get_cli_overrides("policy")
|
||||
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
|
||||
yaml_overrides = parser.get_yaml_overrides("policy")
|
||||
cli_overrides = parser.get_cli_overrides("policy") or []
|
||||
policy_overrides = yaml_overrides + cli_overrides
|
||||
pretrained_revision = parser.parse_arg("pretrained_revision", cli_overrides)
|
||||
if pretrained_revision is None:
|
||||
pretrained_revision = parser.parse_arg("pretrained_revision", yaml_overrides)
|
||||
self.policy = PreTrainedConfig.from_pretrained(
|
||||
policy_path,
|
||||
revision=pretrained_revision,
|
||||
cli_overrides=policy_overrides,
|
||||
)
|
||||
self.policy.pretrained_path = policy_path
|
||||
if self.policy is None:
|
||||
raise ValueError("--policy.path is required for rollout")
|
||||
|
||||
@@ -27,7 +27,7 @@ from threading import Event
|
||||
|
||||
import torch
|
||||
|
||||
from lerobot.configs import FeatureType
|
||||
from lerobot.configs import FeatureType, PreTrainedConfig
|
||||
from lerobot.datasets import (
|
||||
LeRobotDataset,
|
||||
aggregate_pipeline_dataset_features,
|
||||
@@ -159,6 +159,35 @@ class RolloutContext:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy:
|
||||
"""Load policy weights, keeping adapter and base-model revisions independent."""
|
||||
pretrained_revision = policy_config.pretrained_revision
|
||||
policy_class = get_policy_class(policy_config.type)
|
||||
|
||||
if not policy_config.use_peft:
|
||||
return policy_class.from_pretrained(
|
||||
policy_config.pretrained_path,
|
||||
config=policy_config,
|
||||
revision=pretrained_revision,
|
||||
)
|
||||
|
||||
from peft import PeftConfig, PeftModel
|
||||
|
||||
peft_path = policy_config.pretrained_path
|
||||
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
|
||||
policy = policy_class.from_pretrained(
|
||||
pretrained_name_or_path=peft_config.base_model_name_or_path,
|
||||
config=policy_config,
|
||||
revision=peft_config.revision,
|
||||
)
|
||||
return PeftModel.from_pretrained(
|
||||
policy,
|
||||
peft_path,
|
||||
config=peft_config,
|
||||
revision=pretrained_revision,
|
||||
)
|
||||
|
||||
|
||||
def build_rollout_context(
|
||||
cfg: RolloutConfig,
|
||||
shutdown_event: Event,
|
||||
@@ -176,7 +205,6 @@ def build_rollout_context(
|
||||
# --- 1. Policy (heavy I/O, but no hardware yet) -------------------
|
||||
logger.info("Loading policy from '%s'...", cfg.policy.pretrained_path)
|
||||
policy_config = cfg.policy
|
||||
policy_class = get_policy_class(policy_config.type)
|
||||
|
||||
if hasattr(policy_config, "compile_model"):
|
||||
policy_config.compile_model = cfg.use_torch_compile
|
||||
@@ -187,17 +215,7 @@ def build_rollout_context(
|
||||
"Please use `cpu` or `cuda` backend."
|
||||
)
|
||||
|
||||
if policy_config.use_peft:
|
||||
from peft import PeftConfig, PeftModel
|
||||
|
||||
peft_path = policy_config.pretrained_path
|
||||
peft_config = PeftConfig.from_pretrained(peft_path)
|
||||
policy = policy_class.from_pretrained(
|
||||
pretrained_name_or_path=peft_config.base_model_name_or_path, config=policy_config
|
||||
)
|
||||
policy = PeftModel.from_pretrained(policy, peft_path, config=peft_config)
|
||||
else:
|
||||
policy = policy_class.from_pretrained(policy_config.pretrained_path, config=policy_config)
|
||||
policy = _load_pretrained_policy(policy_config)
|
||||
|
||||
if is_rtc:
|
||||
policy.config.rtc_config = cfg.inference.rtc
|
||||
@@ -392,6 +410,7 @@ def build_rollout_context(
|
||||
preprocessor, postprocessor = make_pre_post_processors(
|
||||
policy_cfg=policy_config,
|
||||
pretrained_path=cfg.policy.pretrained_path,
|
||||
pretrained_revision=policy_config.pretrained_revision,
|
||||
dataset_stats=dataset_stats,
|
||||
preprocessor_overrides={
|
||||
"device_processor": {"device": cfg.device},
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -17,6 +17,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
@@ -106,6 +108,109 @@ def test_sentry_config_defaults():
|
||||
assert cfg.target_video_file_size_mb is None
|
||||
|
||||
|
||||
def test_rollout_config_passes_policy_pretrained_revision(monkeypatch):
|
||||
from lerobot.configs import PreTrainedConfig, parser
|
||||
from lerobot.rollout import RolloutConfig
|
||||
from tests.mocks.mock_robot import MockRobotConfig
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_from_pretrained(cls, pretrained_name_or_path, **kwargs):
|
||||
captured["pretrained_name_or_path"] = pretrained_name_or_path
|
||||
captured.update(kwargs)
|
||||
return SimpleNamespace(device="cpu", pretrained_revision=kwargs["revision"])
|
||||
|
||||
monkeypatch.setattr(parser, "get_yaml_overrides", lambda _: ["--pretrained_revision=yaml-sha"])
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
["lerobot-rollout", "--policy.path=user/policy", "--policy.pretrained_revision=cli-sha"],
|
||||
)
|
||||
monkeypatch.setattr(PreTrainedConfig, "from_pretrained", classmethod(fake_from_pretrained))
|
||||
|
||||
cfg = RolloutConfig(robot=MockRobotConfig())
|
||||
|
||||
assert captured["pretrained_name_or_path"] == "user/policy"
|
||||
assert captured["revision"] == "cli-sha"
|
||||
assert captured["cli_overrides"] == [
|
||||
"--pretrained_revision=yaml-sha",
|
||||
"--pretrained_revision=cli-sha",
|
||||
]
|
||||
assert cfg.policy.pretrained_path == "user/policy"
|
||||
assert cfg.policy.pretrained_revision == "cli-sha"
|
||||
|
||||
|
||||
def test_load_pretrained_policy_passes_revision(monkeypatch):
|
||||
import lerobot.rollout.context as rollout_context
|
||||
|
||||
policy_config = SimpleNamespace(
|
||||
type="mock",
|
||||
use_peft=False,
|
||||
pretrained_path="user/policy",
|
||||
pretrained_revision="policy-sha",
|
||||
)
|
||||
policy_class = MagicMock()
|
||||
loaded_policy = MagicMock()
|
||||
policy_class.from_pretrained.return_value = loaded_policy
|
||||
monkeypatch.setattr(rollout_context, "get_policy_class", lambda _: policy_class)
|
||||
|
||||
policy = rollout_context._load_pretrained_policy(policy_config)
|
||||
|
||||
assert policy is loaded_policy
|
||||
policy_class.from_pretrained.assert_called_once_with(
|
||||
"user/policy",
|
||||
config=policy_config,
|
||||
revision="policy-sha",
|
||||
)
|
||||
|
||||
|
||||
def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(monkeypatch):
|
||||
import lerobot.rollout.context as rollout_context
|
||||
|
||||
policy_config = SimpleNamespace(
|
||||
type="mock",
|
||||
use_peft=True,
|
||||
pretrained_path="user/adapter",
|
||||
pretrained_revision="adapter-sha",
|
||||
)
|
||||
policy_class = MagicMock()
|
||||
base_policy = MagicMock()
|
||||
policy_class.from_pretrained.return_value = base_policy
|
||||
monkeypatch.setattr(rollout_context, "get_policy_class", lambda _: policy_class)
|
||||
|
||||
peft_config = SimpleNamespace(
|
||||
base_model_name_or_path="user/base-policy",
|
||||
revision="base-sha",
|
||||
)
|
||||
peft_config_from_pretrained = MagicMock(return_value=peft_config)
|
||||
adapted_policy = MagicMock()
|
||||
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"peft",
|
||||
SimpleNamespace(
|
||||
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
|
||||
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
|
||||
),
|
||||
)
|
||||
|
||||
policy = rollout_context._load_pretrained_policy(policy_config)
|
||||
|
||||
assert policy is adapted_policy
|
||||
peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
|
||||
policy_class.from_pretrained.assert_called_once_with(
|
||||
pretrained_name_or_path="user/base-policy",
|
||||
config=policy_config,
|
||||
revision="base-sha",
|
||||
)
|
||||
peft_model_from_pretrained.assert_called_once_with(
|
||||
base_policy,
|
||||
"user/adapter",
|
||||
config=peft_config,
|
||||
revision="adapter-sha",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RolloutRingBuffer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user