Compare commits

..

16 Commits

Author SHA1 Message Date
CarolinePascal 6cd4f1839b chore(format) 2026-07-30 11:22:56 +02:00
CarolinePascal dc0570586b fix(RGB only): remove the stereo module fallback when setting RGB/color parameters to avoid unexpected impacts on depth sensing 2026-07-30 11:18:24 +02:00
Steven Palma 36b8face98 fix(utils): validate precise_sleep spin/margin args (#4218)
* fix(utils): validate precise_sleep spin/margin args

Negative spin_threshold/sleep_margin make remaining arithmetic wrong
and can overshoot. Reject them early; cover the no-op path.

* test: drop flaky wall-clock assertion in no-op test

Per review: the 50ms wall-clock check can exceed its bound on a preempted
CI worker even when precise_sleep returns immediately. The direct calls
already exercise the non-positive no-op path, so the assertion is redundant.

* chore(tests): remove precise_sleep test negative values

---------

Co-authored-by: Bartok9 <danielrpike9@gmail.com>
2026-07-29 20:24:07 +02:00
Steven Palma cd8984cc0a fix(utils): allow any JSON payload in write_json - #3993 (#4217)
* fix(utils): allow any JSON payload in write_json

The dict-only type stub blocked lists/scalars callers already dump.
Accept Any, set utf-8 encoding, and cover list roundtrip.

* fix(utils): json type

---------

Co-authored-by: Bartok9 <danielrpike9@gmail.com>
2026-07-29 20:11:14 +02:00
Steven Palma b9ded9e761 fix(utils): mark Transition.complementary_info NotRequired (#4216)
* fix(utils): mark Transition.complementary_info NotRequired

TypedDict class-body ``= None`` does not make a key optional and confuses
type checkers. Use ``NotRequired[...]`` so transitions without metadata
are valid.

* refactor(utils): complete NotRequired

---------

Co-authored-by: Bartok9 <danielrpike9@gmail.com>
2026-07-29 19:55:39 +02:00
Steven Palma 185f3e1708 fix(utils): preserve exc_info/stack_info in init_logging formatter (#4215)
* fix(utils): preserve exc_info/stack_info in init_logging formatter

Replacing Formatter.format dropped logging.exception() tracebacks,
hurting HIL-SERL actor/learner crash diagnosis. Append formatted
exceptions and stack_info like the stdlib formatter.

Fixes #3978

* refactor(utils): format logging

---------

Co-authored-by: Bartok9 <danielrpike9@gmail.com>
2026-07-29 19:32:30 +02:00
Bartok e36783253a fix(utils): raise ValueError from get_safe_torch_device (#3992)
* fix(utils): raise ValueError from get_safe_torch_device

Bare asserts vanish under python -O and look like programmer bugs.
Convert unavailable CUDA/MPS/XPU requests into clear ValueErrors.

* style: combine nested with in device util tests (ruff)

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 19:21:50 +02:00
Bartok 289e577fc7 fix(utils): reject zero-norm / invalid quaternions in Rotation (#3988)
* fix(utils): reject zero-norm / invalid quaternions in Rotation

Zero or non-finite inputs previously slipped through and produced NaN
rotation matrices on later convert/apply. Validate shape and scept for
norm > 0 before normalizing.

* fix(teleop): degrade phone AR quat parse like missing pose

Address review on #3988: Rotation.from_quat now rejects zero/NaN
quaternions. Wrap HEBI iOS ARKit permission in ValueError and return the
existing (False, None, None, None) path so teleop does not die mid-session
before tracking is ready.

* style: ruff format long ValueError in rotation.py

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 19:13:40 +02:00
Anes Benmerzoug 9c32722eb9 fix(find-cameras): enforce sequential lifecycle and add configurable warmup (#3593)
* Connect, test and disconnected camera instances sequentially

* Add warmup-s cli argument to lerobot-find-cameras script

* Reduce default record time from 6 to 2 seconds in find_cameras

* Annotate return value of save_image function

* Initialize logging configuration in find_cameras
2026-07-29 19:01:44 +02:00
Kunal b49cb50e01 docs(agent-guide): prioritize uv over pip in §4.1 install block (#3799)
Co-authored-by: Altman <64389901+Altman-conquer@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 18:46:20 +02:00
Steven Palma dd08d4eb53 fix(robot): type FK-to-EE action features as ACTION not STATE (#4213)
* fix(robot): type FK-to-EE action features as ACTION not STATE

ForwardKinematicsJointsToEEAction.transform_features declared its
end-effector action features (ee.x/y/z/wx/wy/wz/gripper_pos) with
FeatureType.STATE, copied verbatim from the sibling
ForwardKinematicsJointsToEEObservation (where STATE is correct for
OBSERVATION features). Every other action-producing step in this file
(EEReferenceAndDelta, InverseKinematicsEEToJoints, InverseKinematicsRLStep)
types its ACTION-bucket features as FeatureType.ACTION.

The mismatch mis-classifies the converted EE actions as state, which
propagates a wrong feature schema to downstream consumers keyed on
FeatureType (e.g. normalization norm_map, policy input/output feature
classification).


* test(robot): FK-to-EE step feature-type contract (action vs observation)

Asserts ForwardKinematicsJointsToEEAction emits EE features in the ACTION
bucket typed FeatureType.ACTION, and ForwardKinematicsJointsToEEObservation
emits them in the OBSERVATION bucket typed FeatureType.STATE.


* chore: delete user file

* chore(processor): reduce verbosity

---------

Co-authored-by: Jaagat-P <jaagatp05@gmail.com>
2026-07-29 18:06:01 +02:00
Martino Russi 6e5f6df6e7 fix(evo1): re-pad normalizer stats when loading from checkpoint (#3945)
* fix(evo1): re-pad normalizer stats when loading from checkpoint

reconcile_evo1_processors did not re-pad the (un)normalizer stats to
max_state_dim/max_action_dim on the checkpoint-load path. When
lerobot-train loads a checkpoint (e.g. stage2 from a stage1 checkpoint)
it injects the raw dataset stats via processor overrides, so LIBERO's
8-dim state stats normalized a 24-dim padded state and crashed with
"size of tensor a (24) must match tensor b (8)".

Restore _refresh_evo1_normalization_steps (removed in the "remove legacy
codepaths" refactor) and call it from reconcile_evo1_processors so the
loaded stats/features are re-padded to EVO1's fixed widths. Padding is a
no-op when stats are already at the target width.

Co-authored-by: Cursor <cursoragent@cursor.com>

* test(evo1): cover reconcile re-padding of overridden normalizer stats

Regression test for the stage2-from-checkpoint crash: reloading a
checkpoint with raw (unpadded) dataset stats injected via processor
overrides must be re-padded to max_state_dim/max_action_dim by
reconcile_evo1_processors, otherwise normalizing the padded state
raises a shape mismatch.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Martino Russi <martino@huggingface.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 17:26:39 +02:00
Steven Palma 265abe6c79 chore(datasets): add typing to aggregate helpers (#4211)
* chore(datasets): add typing to aggregate helpers

Signed-off-by: nathon-lee <leejianwoo@gmail.com>

* chore(dataset): add more typing aggregate

* chore(test): remove panda test

---------

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Co-authored-by: nathon-lee <leejianwoo@gmail.com>
2026-07-29 17:07:34 +02:00
Old-Ding b4e2d0b610 docs: fix wording in guides (#3939)
Generated-by: OpenAI Codex

Signed-off-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 16:24:03 +02:00
Old-Ding 5594eba06a docs: fix repeated word in backward compatibility guide (#3938)
Generated-by: OpenAI Codex

Signed-off-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 16:23:01 +02:00
saime428 207183c2f8 docs: fix dataset split fraction example (#3936)
* docs: fix dataset split fraction example

* docs: preserve three-way dataset split example

---------

Co-authored-by: saime <2286263079@qq.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 16:21:12 +02:00
29 changed files with 465 additions and 820 deletions
+10 -6
View File
@@ -61,16 +61,20 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
**4.1 Install** **4.1 Install**
```bash ```bash
pip install 'lerobot[feetech]' # SO-100/SO-101 motor stack # uv (recommended — see AGENTS.md and CLAUDE.md)
# pip install 'lerobot[all]' # everything uv sync --locked --extra feetech # SO-100/SO-101 motor stack
# pip install 'lerobot[aloha,pusht]' # specific features # uv sync --locked --extra all # everything
# pip install 'lerobot[smolvla]' # add SmolVLA deps # uv sync --locked --extra smolvla # add SmolVLA deps
# pip (alternative, e.g. when not working from source)
# pip install 'lerobot[feetech]'
# pip install 'lerobot[all]'
# pip install 'lerobot[smolvla]'
git lfs install && git lfs pull git lfs install && git lfs pull
hf auth login # required to push datasets/policies hf auth login # required to push datasets/policies
``` ```
Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
**4.2 Find USB ports** — run once per arm, unplug when prompted. **4.2 Find USB ports** — run once per arm, unplug when prompted.
```bash ```bash
+3 -3
View File
@@ -58,7 +58,7 @@ final_action = postprocessor(action)
## Hardware API redesign ## Hardware API redesign
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request. PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is an overview of what changed and how you can continue to work with datasets created before this pull request.
### What changed? ### What changed?
@@ -129,8 +129,8 @@ python examples/backward_compatibility/replay.py \
Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied. Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.
To find these transformations, we recommend to first try and and replay an episode of the dataset your policy was trained on using the section above. To find these transformations, we recommend first replaying an episode of the dataset your policy was trained on using the section above.
Then, add these same transformations on your inference script (shown here in the `record.py` script): Then, add these same transformations to your inference script (shown here in the `record.py` script):
```diff ```diff
action_values = predict_action( action_values = predict_action(
+2 -2
View File
@@ -164,8 +164,8 @@ includes the range reported by the sensor. Requesting an unsupported control als
Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
require `use_rgb=True`. require `use_rgb=True`.
On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual Manual color controls require a dedicated RGB module. Cameras without one, such as the RealSense
exposure or gain also affects the depth stream. D405, do not support them and raise an error at connection time.
</hfoption> </hfoption>
</hfoptions> </hfoptions>
+4 -4
View File
@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
For each motor you want to update: For each motor you want to update:
1. **Select the motor** from the list by clicking on it 1. **Select the motor** from the list by clicking on it
2. **Click on Upgrade tab**: 2. **Click the Upgrade tab**:
3. **Click on Online button**: 3. **Click the Online button**:
- If an potential firmware update is found, it will be displayed in the box - If a potential firmware update is found, it will be displayed in the box
4. **Click on Upgrade button**: 4. **Click the Upgrade button**:
- The update progress will be displayed - The update progress will be displayed
## Step 6: Verify Update ## Step 6: Verify Update
+3 -3
View File
@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
## Three pipelines ## Three pipelines
We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match. We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline. Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets) 1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
2. Pipeline 2: Dataset action space → robot command space (EE targets → joints) 2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
@@ -74,7 +74,7 @@ In the phone to SO-100 follower examples we use the following adapters:
- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition. - `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
- `transition_to_robot_action`: transforms the pipeline transition to a robot action dict. - `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
- `observation_to_transition`: transforms the robot observation dict to a pipeline transition. - `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
- `transition_to_observation`: transforms the pipeline transition to a observation dict. - `transition_to_observation`: transforms the pipeline transition to an observation dict.
Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details. Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
@@ -82,7 +82,7 @@ Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/le
Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`. Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples: Below is an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
```python ```python
def transform_features( def transform_features(
+2 -2
View File
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda") policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
# Now use predict_action_chunk with RTC parameters # Now use predict_action_chunk with RTC parameters
inference_delay = 4 # How many steps of inference latency, this values should be calculated based on the inference latency of the policy inference_delay = 4 # How many steps of inference latency, this value should be calculated based on the inference latency of the policy
# Initialize the action queue # Initialize the action queue
action_queue = ActionQueue(policy_cfg.rtc_config) action_queue = ActionQueue(policy_cfg.rtc_config)
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
RTCConfig(execution_horizon=10) RTCConfig(execution_horizon=10)
``` ```
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value. **`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is an optimal value.
**`prefix_attention_schedule`**: How to weight consistency across the overlap region. **`prefix_attention_schedule`**: How to weight consistency across the overlap region.
+2 -2
View File
@@ -50,11 +50,11 @@ lerobot-edit-dataset \
Divide a dataset into multiple subsets. Divide a dataset into multiple subsets.
```bash ```bash
# Split by fractions (e.g. 80% train, 20% test, 20% val) # Split by fractions (e.g. 60% train, 20% val, 20% test)
lerobot-edit-dataset \ lerobot-edit-dataset \
--repo_id lerobot/pusht \ --repo_id lerobot/pusht \
--operation.type split \ --operation.type split \
--operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}' --operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'
# Split by specific episode indices # Split by specific episode indices
lerobot-edit-dataset \ lerobot-edit-dataset \
Binary file not shown.

Before

Width:  |  Height:  |  Size: 682 KiB

@@ -365,11 +365,12 @@ class RealSenseCamera(Camera):
return self._async_read(timeout_ms=10000, read_depth=read_depth) return self._async_read(timeout_ms=10000, read_depth=read_depth)
def _get_color_sensor(self) -> "rs.sensor": def _get_color_sensor(self) -> "rs.sensor":
"""Returns the sensor that controls the color stream. """Returns the dedicated "RGB Camera" sensor that controls the color stream.
Most RealSense cameras expose "RGB Camera" for color. The D405 has no Manual color controls are only applied to a dedicated RGB module. Cameras
separate RGB module — its color stream comes from "Stereo Module". without one (e.g. the D405, whose color stream comes from the shared
We try RGB Camera first, then fall back to Stereo Module. "Stereo Module") are unsupported, so we never fall back to another sensor
to avoid altering the depth stream.
""" """
if self.rs_profile is None: if self.rs_profile is None:
raise RuntimeError(f"{self}: rs_profile must be initialized before use.") raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
@@ -377,12 +378,14 @@ class RealSenseCamera(Camera):
device = self.rs_profile.get_device() device = self.rs_profile.get_device()
sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()} sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
for name in ("RGB Camera", "Stereo Module"): if "RGB Camera" in sensors:
if name in sensors: return sensors["RGB Camera"]
return sensors[name]
available = list(sensors.keys()) available = list(sensors.keys())
raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}") raise RuntimeError(
f"{self}: manual color controls require a dedicated 'RGB Camera' module, which this camera does not have. ",
f"Available sensors: {available}.",
)
def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None: def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
"""Sets a sensor option, re-raising range errors with actionable diagnostics.""" """Sets a sensor option, re-raising range errors with actionable diagnostics."""
+96 -40
View File
@@ -19,6 +19,7 @@ import copy
import logging import logging
import shutil import shutil
from pathlib import Path from pathlib import Path
from typing import Any, NotRequired, TypedDict
import datasets import datasets
import pandas as pd import pandas as pd
@@ -49,8 +50,32 @@ from .utils import (
) )
from .video_utils import concatenate_video_files, get_video_duration_in_s from .video_utils import concatenate_video_files, get_video_duration_in_s
logger = logging.getLogger(__name__)
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]: type FeatureDict = dict[str, dict[str, Any]]
type ChunkFile = tuple[int, int]
class IndexState(TypedDict):
chunk: int
file: int
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
class VideoIndex(TypedDict):
chunk: int
file: int
latest_duration: float
episode_duration: float
src_to_offset: NotRequired[dict[ChunkFile, float]]
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
dst_file_durations: NotRequired[dict[ChunkFile, float]]
type VideoIndexState = dict[str, VideoIndex]
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> FeatureDict:
"""Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged. """Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.
Args: Args:
@@ -59,14 +84,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
Returns: Returns:
dict: A dictionary of merged video feature info. dict: A dictionary of merged video feature info.
""" """
merged_info = copy.deepcopy(all_metadata[0].features) merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"] video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
for vk in video_keys: for vk in video_keys:
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata] video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
base_video_info = video_infos[0] base_video_info = video_infos[0]
merged_encoder_info: dict = {} merged_encoder_info: dict[str, Any] = {}
fallback_keys: list[str] = [] fallback_keys: list[str] = []
for info_key in VIDEO_ENCODER_INFO_KEYS: for info_key in VIDEO_ENCODER_INFO_KEYS:
values = [info.get(info_key, None) for info in video_infos] values = [info.get(info_key, None) for info in video_infos]
@@ -80,7 +105,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None
if fallback_keys: if fallback_keys:
logging.warning( logger.warning(
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. " f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
f"Setting these keys to null: {fallback_keys}.", f"Setting these keys to null: {fallback_keys}.",
) )
@@ -92,7 +117,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
return merged_info return merged_info
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]): def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
"""Validates that all dataset metadata have consistent properties. """Validates that all dataset metadata have consistent properties.
Ensures all datasets have the same fps, robot_type, and features to guarantee Ensures all datasets have the same fps, robot_type, and features to guarantee
@@ -129,7 +154,9 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
return fps, robot_type, features return fps, robot_type, features
def update_data_df(df, src_meta, dst_meta): def update_data_df(
df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
) -> pd.DataFrame:
"""Updates a data DataFrame with new indices and task mappings for aggregation. """Updates a data DataFrame with new indices and task mappings for aggregation.
Adjusts episode indices, frame indices, and task indices to account for Adjusts episode indices, frame indices, and task indices to account for
@@ -154,12 +181,12 @@ def update_data_df(df, src_meta, dst_meta):
def update_meta_data( def update_meta_data(
df, df: pd.DataFrame,
dst_meta, dst_meta: LeRobotDatasetMetadata,
meta_idx, meta_idx: IndexState,
data_idx, data_idx: IndexState,
videos_idx, videos_idx: VideoIndexState,
): ) -> pd.DataFrame:
"""Updates metadata DataFrame with new chunk, file, and timestamp indices. """Updates metadata DataFrame with new chunk, file, and timestamp indices.
Adjusts all indices and timestamps to account for previously aggregated Adjusts all indices and timestamps to account for previously aggregated
@@ -289,7 +316,7 @@ def aggregate_datasets(
chunk_size: int | None = None, chunk_size: int | None = None,
concatenate_videos: bool = True, concatenate_videos: bool = True,
concatenate_data: bool = True, concatenate_data: bool = True,
): ) -> None:
"""Aggregates multiple LeRobot datasets into a single unified dataset. """Aggregates multiple LeRobot datasets into a single unified dataset.
This is the main function that orchestrates the aggregation process by: This is the main function that orchestrates the aggregation process by:
@@ -309,7 +336,7 @@ def aggregate_datasets(
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards. concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
concatenate_data: When False, keep one parquet per source file instead of packing into shards. concatenate_data: When False, keep one parquet per source file instead of packing into shards.
""" """
logging.info("Start aggregate_datasets") logger.info("Start aggregate_datasets")
if data_files_size_in_mb is None: if data_files_size_in_mb is None:
data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
@@ -341,15 +368,15 @@ def aggregate_datasets(
video_files_size_in_mb=video_files_size_in_mb, video_files_size_in_mb=video_files_size_in_mb,
) )
logging.info("Find all tasks") logger.info("Find all tasks")
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique() unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
dst_meta.tasks = pd.DataFrame( dst_meta.tasks = pd.DataFrame(
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task") {"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
) )
meta_idx = {"chunk": 0, "file": 0} meta_idx: IndexState = {"chunk": 0, "file": 0}
data_idx = {"chunk": 0, "file": 0} data_idx: IndexState = {"chunk": 0, "file": 0}
videos_idx = { videos_idx: VideoIndexState = {
key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
} }
@@ -373,12 +400,17 @@ def aggregate_datasets(
dst_meta.info.total_frames += src_meta.total_frames dst_meta.info.total_frames += src_meta.total_frames
finalize_aggregation(dst_meta, all_metadata) finalize_aggregation(dst_meta, all_metadata)
logging.info("Aggregation complete.") logger.info("Aggregation complete.")
def aggregate_videos( def aggregate_videos(
src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True src_meta: LeRobotDatasetMetadata,
): dst_meta: LeRobotDatasetMetadata,
videos_idx: VideoIndexState,
video_files_size_in_mb: float,
chunk_size: int,
concatenate_videos: bool = True,
) -> VideoIndexState:
"""Aggregates video chunks from a source dataset into the destination dataset. """Aggregates video chunks from a source dataset into the destination dataset.
Handles video file concatenation and rotation based on file size limits. Handles video file concatenation and rotation based on file size limits.
@@ -406,7 +438,8 @@ def aggregate_videos(
videos_idx[key]["dst_file_durations"] = {} videos_idx[key]["dst_file_durations"] = {}
for key, video_idx in videos_idx.items(): for key, video_idx in videos_idx.items():
unique_chunk_file_pairs = { unique_chunk_file_pairs: list[ChunkFile] = sorted(
{
(chunk, file) (chunk, file)
for chunk, file in zip( for chunk, file in zip(
src_meta.episodes[f"videos/{key}/chunk_index"], src_meta.episodes[f"videos/{key}/chunk_index"],
@@ -414,7 +447,7 @@ def aggregate_videos(
strict=False, strict=False,
) )
} }
unique_chunk_file_pairs = sorted(unique_chunk_file_pairs) )
chunk_idx = video_idx["chunk"] chunk_idx = video_idx["chunk"]
file_idx = video_idx["file"] file_idx = video_idx["file"]
@@ -489,7 +522,14 @@ def aggregate_videos(
return videos_idx return videos_idx
def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True): def aggregate_data(
src_meta: LeRobotDatasetMetadata,
dst_meta: LeRobotDatasetMetadata,
data_idx: IndexState,
data_files_size_in_mb: float,
chunk_size: int,
concatenate_data: bool = True,
) -> IndexState:
"""Aggregates data chunks from a source dataset into the destination dataset. """Aggregates data chunks from a source dataset into the destination dataset.
Reads source data files, updates indices to match the aggregated dataset, Reads source data files, updates indices to match the aggregated dataset,
@@ -510,14 +550,16 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
Returns: Returns:
dict: Updated data_idx with current chunk and file indices. dict: Updated data_idx with current chunk and file indices.
""" """
unique_chunk_file_ids = { unique_chunk_file_ids: list[ChunkFile] = sorted(
{
(c, f) (c, f)
for c, f in zip( for c, f in zip(
src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False src_meta.episodes["data/chunk_index"],
src_meta.episodes["data/file_index"],
strict=False,
) )
} }
)
unique_chunk_file_ids = sorted(unique_chunk_file_ids)
contains_images = len(dst_meta.image_keys) > 0 contains_images = len(dst_meta.image_keys) > 0
# retrieve features schema for proper image typing in parquet # retrieve features schema for proper image typing in parquet
@@ -525,7 +567,7 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
# Track source to destination file mapping for metadata update # Track source to destination file mapping for metadata update
# This is critical for handling datasets that are already results of a merge # This is critical for handling datasets that are already results of a merge
src_to_dst: dict[tuple[int, int], tuple[int, int]] = {} src_to_dst: dict[ChunkFile, ChunkFile] = {}
for src_chunk_idx, src_file_idx in unique_chunk_file_ids: for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
src_path = src_meta.root / DEFAULT_DATA_PATH.format( src_path = src_meta.root / DEFAULT_DATA_PATH.format(
@@ -564,7 +606,13 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
return data_idx return data_idx
def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx): def aggregate_metadata(
src_meta: LeRobotDatasetMetadata,
dst_meta: LeRobotDatasetMetadata,
meta_idx: IndexState,
data_idx: IndexState,
videos_idx: VideoIndexState,
) -> IndexState:
"""Aggregates metadata from a source dataset into the destination dataset. """Aggregates metadata from a source dataset into the destination dataset.
Reads source metadata files, updates all indices and timestamps, Reads source metadata files, updates all indices and timestamps,
@@ -580,7 +628,8 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
Returns: Returns:
dict: Updated meta_idx with current chunk and file indices. dict: Updated meta_idx with current chunk and file indices.
""" """
chunk_file_ids = { chunk_file_ids: list[ChunkFile] = sorted(
{
(c, f) (c, f)
for c, f in zip( for c, f in zip(
src_meta.episodes["meta/episodes/chunk_index"], src_meta.episodes["meta/episodes/chunk_index"],
@@ -588,8 +637,7 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
strict=False, strict=False,
) )
} }
)
chunk_file_ids = sorted(chunk_file_ids)
for chunk_idx, file_idx in chunk_file_ids: for chunk_idx, file_idx in chunk_file_ids:
src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx) src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
df = pd.read_parquet(src_path) df = pd.read_parquet(src_path)
@@ -622,16 +670,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
def append_or_create_parquet_file( def append_or_create_parquet_file(
df: pd.DataFrame, df: pd.DataFrame,
src_path: Path, src_path: Path,
idx: dict[str, int], idx: IndexState,
max_mb: float, max_mb: float,
chunk_size: int, chunk_size: int,
default_path: str, default_path: str,
contains_images: bool = False, contains_images: bool = False,
aggr_root: Path = None, aggr_root: Path | None = None,
hf_features: datasets.Features | None = None, hf_features: datasets.Features | None = None,
concatenate: bool = True, concatenate: bool = True,
one_row_group_per_episode: bool = False, one_row_group_per_episode: bool = False,
) -> tuple[dict[str, int], tuple[int, int]]: ) -> tuple[IndexState, ChunkFile]:
"""Appends data to an existing parquet file or creates a new one based on size constraints. """Appends data to an existing parquet file or creates a new one based on size constraints.
Manages file rotation when size limits are exceeded to prevent individual files Manages file rotation when size limits are exceeded to prevent individual files
@@ -654,7 +702,13 @@ def append_or_create_parquet_file(
Returns: Returns:
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
and (dst_chunk, dst_file) is the actual destination file the data was written to. and (dst_chunk, dst_file) is the actual destination file the data was written to.
Raises:
ValueError: If aggr_root is not provided.
""" """
if aggr_root is None:
raise ValueError("aggr_root must be provided.")
dst_chunk, dst_file = idx["chunk"], idx["file"] dst_chunk, dst_file = idx["chunk"], idx["file"]
dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file) dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
@@ -698,7 +752,9 @@ def append_or_create_parquet_file(
return idx, (dst_chunk, dst_file) return idx, (dst_chunk, dst_file)
def finalize_aggregation(aggr_meta, all_metadata): def finalize_aggregation(
aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
) -> None:
"""Finalizes the dataset aggregation by writing summary files and statistics. """Finalizes the dataset aggregation by writing summary files and statistics.
Writes the tasks file, info file with total counts and splits, and Writes the tasks file, info file with total counts and splits, and
@@ -708,16 +764,16 @@ def finalize_aggregation(aggr_meta, all_metadata):
aggr_meta: Aggregated dataset metadata. aggr_meta: Aggregated dataset metadata.
all_metadata: List of all source dataset metadata objects. all_metadata: List of all source dataset metadata objects.
""" """
logging.info("write tasks") logger.info("write tasks")
write_tasks(aggr_meta.tasks, aggr_meta.root) write_tasks(aggr_meta.tasks, aggr_meta.root)
logging.info("write info") logger.info("write info")
aggr_meta.info.total_tasks = len(aggr_meta.tasks) aggr_meta.info.total_tasks = len(aggr_meta.tasks)
aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata) aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata)
aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata) aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata)
aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"} aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
write_info(aggr_meta.info, aggr_meta.root) write_info(aggr_meta.info, aggr_meta.root)
logging.info("write stats") logger.info("write stats")
aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata]) aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
write_stats(aggr_meta.stats, aggr_meta.root) write_stats(aggr_meta.stats, aggr_meta.root)
+35 -5
View File
@@ -302,6 +302,33 @@ def _pad_evo1_stats(
return padded_stats return padded_stats
def _refresh_evo1_normalization_steps(
config: Evo1Config,
preprocessor: PolicyProcessorPipeline,
postprocessor: PolicyProcessorPipeline,
) -> None:
"""Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
LIBERO stats). Padding is a no-op when stats are already at the target width.
"""
normalization_features = _evo1_normalization_features(config)
action_features = _evo1_action_features(config)
for step in preprocessor.steps:
if isinstance(step, NormalizerProcessorStep):
step.features = normalization_features
step.stats = _pad_evo1_stats(config, step.stats)
step.to(device=step.device, dtype=step.dtype)
for step in postprocessor.steps:
if isinstance(step, UnnormalizerProcessorStep):
step.features = action_features
step.stats = _pad_evo1_stats(config, step.stats)
step.to(device=step.device, dtype=step.dtype)
def reconcile_evo1_processors( def reconcile_evo1_processors(
config: Evo1Config, config: Evo1Config,
preprocessor: PolicyProcessorPipeline, preprocessor: PolicyProcessorPipeline,
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config. """Reconcile checkpoint-loaded pipelines with the current EVO1 config.
Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
(converters are plain functions and are never serialized), and eval-time CLI overrides of the (converters are plain functions and are never serialized), eval-time CLI overrides of the
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
restores the converter and rebuilds the action step from the current config so those overrides (un)normalizer stats/features when the generic override path injects raw, unpadded dataset
take effect. stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and
rebuilds the action step from the current config so those overrides take effect.
""" """
# Pipelines reloaded from a checkpoint come back with the default batch converter, which drops # Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1. # non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
preprocessor.to_transition = evo1_batch_to_transition preprocessor.to_transition = evo1_batch_to_transition
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
action_step = Evo1ActionProcessorStep( action_step = Evo1ActionProcessorStep(
action_dim=_evo1_action_dim(config), action_dim=_evo1_action_dim(config),
binarize_gripper=config.binarize_gripper, binarize_gripper=config.binarize_gripper,
+2 -2
View File
@@ -18,7 +18,7 @@ import functools
import threading import threading
from collections.abc import Callable, Sequence from collections.abc import Callable, Sequence
from contextlib import suppress from contextlib import suppress
from typing import TypedDict from typing import NotRequired, TypedDict
import torch import torch
import torch.nn.functional as F # noqa: N812 import torch.nn.functional as F # noqa: N812
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
next_state: dict[str, torch.Tensor] next_state: dict[str, torch.Tensor]
done: torch.Tensor done: torch.Tensor
truncated: torch.Tensor truncated: torch.Tensor
complementary_info: dict[str, torch.Tensor | float | int] | None = None complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor: def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
@@ -510,10 +510,10 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
# We only use the ee pose in the dataset, so we don't need the joint positions # We only use the ee pose in the dataset, so we don't need the joint positions
for n in self.motor_names: for n in self.motor_names:
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None) features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
# We specify the dataset features of this step that we want to be stored in the dataset # Store end-effector features as actions in the dataset schema
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]: for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature( features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
type=FeatureType.STATE, shape=(1,) type=FeatureType.ACTION, shape=(1,)
) )
return features return features
+28 -39
View File
@@ -28,7 +28,6 @@ lerobot-find-cameras
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful. # NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
import argparse import argparse
import concurrent.futures
import logging import logging
import time import time
from pathlib import Path from pathlib import Path
@@ -133,7 +132,7 @@ def save_image(
camera_identifier: str | int, camera_identifier: str | int,
images_dir: Path, images_dir: Path,
camera_type: str, camera_type: str,
): ) -> None:
""" """
Saves a single image to disk using Pillow. Handles color conversion if necessary. Saves a single image to disk using Pillow. Handles color conversion if necessary.
""" """
@@ -152,7 +151,7 @@ def save_image(
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}") logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None: def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None:
"""Create and connect to a camera instance based on metadata.""" """Create and connect to a camera instance based on metadata."""
cam_type = cam_meta.get("type") cam_type = cam_meta.get("type")
cam_id = cam_meta.get("id") cam_id = cam_meta.get("id")
@@ -165,12 +164,14 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
cv_config = OpenCVCameraConfig( cv_config = OpenCVCameraConfig(
index_or_path=cam_id, index_or_path=cam_id,
color_mode=ColorMode.RGB, color_mode=ColorMode.RGB,
warmup_s=warmup_s,
) )
instance = OpenCVCamera(cv_config) instance = OpenCVCamera(cv_config)
elif cam_type == "RealSense": elif cam_type == "RealSense":
rs_config = RealSenseCameraConfig( rs_config = RealSenseCameraConfig(
serial_number_or_name=cam_id, serial_number_or_name=cam_id,
color_mode=ColorMode.RGB, color_mode=ColorMode.RGB,
warmup_s=warmup_s,
) )
instance = RealSenseCamera(rs_config) instance = RealSenseCamera(rs_config)
else: else:
@@ -188,9 +189,7 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
return None return None
def process_camera_image( def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None:
cam_dict: dict[str, Any], output_dir: Path, current_time: float
) -> concurrent.futures.Future | None:
"""Capture and process an image from a single camera.""" """Capture and process an image from a single camera."""
cam = cam_dict["instance"] cam = cam_dict["instance"]
meta = cam_dict["meta"] meta = cam_dict["meta"]
@@ -200,7 +199,7 @@ def process_camera_image(
try: try:
image_data = cam.read() image_data = cam.read()
return save_image( save_image(
image_data, image_data,
cam_id_str, cam_id_str,
output_dir, output_dir,
@@ -215,10 +214,9 @@ def process_camera_image(
return None return None
def cleanup_cameras(cameras_to_use: list[dict[str, Any]]): def cleanup_camera(cam_dict: dict[str, Any]) -> None:
"""Disconnect all cameras.""" """Disconnect all cameras."""
logger.info(f"Disconnecting {len(cameras_to_use)} cameras...") logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...")
for cam_dict in cameras_to_use:
try: try:
if cam_dict["instance"] and cam_dict["instance"].is_connected: if cam_dict["instance"] and cam_dict["instance"].is_connected:
cam_dict["instance"].disconnect() cam_dict["instance"].disconnect()
@@ -230,6 +228,7 @@ def save_images_from_all_cameras(
output_dir: Path, output_dir: Path,
record_time_s: float = 2.0, record_time_s: float = 2.0,
camera_type: str | None = None, camera_type: str | None = None,
warmup_s: int = 1,
): ):
""" """
Connects to detected cameras (optionally filtered by type) and saves images from each. Connects to detected cameras (optionally filtered by type) and saves images from each.
@@ -240,6 +239,7 @@ def save_images_from_all_cameras(
record_time_s: Duration in seconds to record images. record_time_s: Duration in seconds to record images.
camera_type: Optional string to filter cameras ("realsense" or "opencv"). camera_type: Optional string to filter cameras ("realsense" or "opencv").
If None, uses all detected cameras. If None, uses all detected cameras.
warmup_s: Duration in seconds to warmup camera before recording images.
""" """
output_dir.mkdir(parents=True, exist_ok=True) output_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Saving images to {output_dir}") logger.info(f"Saving images to {output_dir}")
@@ -249,39 +249,23 @@ def save_images_from_all_cameras(
logger.warning("No cameras detected matching the criteria. Cannot save images.") logger.warning("No cameras detected matching the criteria. Cannot save images.")
return return
cameras_to_use = [] logger.info(
for cam_meta in all_camera_metadata: f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras."
camera_instance = create_camera_instance(cam_meta) )
if camera_instance:
cameras_to_use.append(camera_instance)
if not cameras_to_use:
logger.warning("No cameras could be connected. Aborting image save.")
return
logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
start_time = time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
try: try:
for cam_meta in all_camera_metadata:
cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s)
if cam_dict is None:
continue
start_time = time.perf_counter()
while time.perf_counter() - start_time < record_time_s: while time.perf_counter() - start_time < record_time_s:
futures = []
current_capture_time = time.perf_counter() current_capture_time = time.perf_counter()
process_camera_image(cam_dict, output_dir, current_capture_time)
for cam_dict in cameras_to_use: cleanup_camera(cam_dict)
future = process_camera_image(cam_dict, output_dir, current_capture_time)
if future:
futures.append(future)
if futures:
concurrent.futures.wait(futures)
except KeyboardInterrupt: except KeyboardInterrupt:
logger.info("Capture interrupted by user.") logger.info("Capture interrupted by user.")
finally: finally:
print("\nFinalizing image saving...")
executor.shutdown(wait=True)
cleanup_cameras(cameras_to_use)
print(f"Image capture finished. Images saved to {output_dir}") print(f"Image capture finished. Images saved to {output_dir}")
@@ -291,7 +275,6 @@ def main():
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="Unified camera utility script for listing cameras and capturing images." description="Unified camera utility script for listing cameras and capturing images."
) )
parser.add_argument( parser.add_argument(
"camera_type", "camera_type",
type=str, type=str,
@@ -309,8 +292,14 @@ def main():
parser.add_argument( parser.add_argument(
"--record-time-s", "--record-time-s",
type=float, type=float,
default=6.0, default=2.0,
help="Time duration to attempt capturing frames. Default: 6 seconds.", help="Time duration to attempt capturing frames. Default: 2 seconds.",
)
parser.add_argument(
"--warmup-s",
type=int,
default=1,
help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
) )
args = parser.parse_args() args = parser.parse_args()
save_images_from_all_cameras(**vars(args)) save_images_from_all_cameras(**vars(args))
@@ -171,7 +171,13 @@ class IOSPhone(BasePhone, Teleoperator):
# HEBI provides orientation in w, x, y, z format. # HEBI provides orientation in w, x, y, z format.
# Scipy's Rotation expects x, y, z, w. # Scipy's Rotation expects x, y, z, w.
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
# ARKit can emit zero/NaN quaternions before tracking is ready or on a
# dropped packet. Rotation.from_quat now rejects those; degrade the same
# way as a missing pose so teleop stays alive mid-session.
try:
rot = Rotation.from_quat(quat_xyzw) rot = Rotation.from_quat(quat_xyzw)
except ValueError:
return False, None, None, None
pos = ar_pos - rot.apply(self.config.camera_offset) pos = ar_pos - rot.apply(self.config.camera_offset)
return True, pos, rot, pose return True, pos, rot, pose
-16
View File
@@ -13,34 +13,18 @@
# limitations under the License. # limitations under the License.
from .transforms import ( from .transforms import (
CoarseDropout,
GammaCorrection,
GaussianNoise,
GaussianPatchBrightness,
ImageTransformConfig, ImageTransformConfig,
ImageTransforms, ImageTransforms,
ImageTransformsConfig, ImageTransformsConfig,
JPEGCompression,
MotionBlur,
PlanckianJitter,
RandomShadow,
RandomSubsetApply, RandomSubsetApply,
SharpnessJitter, SharpnessJitter,
make_transform_from_config, make_transform_from_config,
) )
__all__ = [ __all__ = [
"CoarseDropout",
"GammaCorrection",
"GaussianNoise",
"GaussianPatchBrightness",
"ImageTransformConfig", "ImageTransformConfig",
"ImageTransforms", "ImageTransforms",
"ImageTransformsConfig", "ImageTransformsConfig",
"JPEGCompression",
"MotionBlur",
"PlanckianJitter",
"RandomShadow",
"RandomSubsetApply", "RandomSubsetApply",
"SharpnessJitter", "SharpnessJitter",
"make_transform_from_config", "make_transform_from_config",
+3 -471
View File
@@ -14,13 +14,11 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import collections import collections
import math
from collections.abc import Callable, Sequence from collections.abc import Callable, Sequence
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any from typing import Any
import torch import torch
from torchvision.io import decode_image, encode_jpeg
from torchvision.transforms import v2 from torchvision.transforms import v2
from torchvision.transforms.v2 import ( from torchvision.transforms.v2 import (
Transform, Transform,
@@ -146,471 +144,6 @@ class SharpnessJitter(Transform):
return self._call_kernel(F.adjust_sharpness, inpt, sharpness_factor=sharpness_factor) return self._call_kernel(F.adjust_sharpness, inpt, sharpness_factor=sharpness_factor)
class GaussianNoise(Transform):
"""Add Gaussian noise to simulate camera sensor noise.
Models readout noise from ADC quantization, which increases in low-light conditions.
Common in real-robot setups where wrist cameras operate in suboptimal lighting.
Args:
std: Range (min, max) for noise standard deviation in pixel-value scale (0-255).
"""
def __init__(self, std: float | Sequence[float] = (5.0, 25.0)) -> None:
super().__init__()
if isinstance(std, (int, float)):
self.std = (0.0, float(std))
elif isinstance(std, Sequence) and len(std) == 2:
self.std = (float(std[0]), float(std[1]))
else:
raise TypeError("std must be a number or a sequence with length 2.")
if not 0.0 <= self.std[0] <= self.std[1]:
raise ValueError(f"std must satisfy 0 <= min <= max, but got {self.std}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
return {
"std": torch.empty(1).uniform_(self.std[0], self.std[1]).item(),
"seed": torch.randint(0, torch.iinfo(torch.int64).max, ()).item(),
}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if isinstance(inpt, torch.Tensor) and inpt.is_floating_point():
generator = torch.Generator(device=inpt.device).manual_seed(params["seed"])
noise = torch.randn(inpt.shape, device=inpt.device, dtype=inpt.dtype, generator=generator)
return (inpt + noise * (params["std"] / 255.0)).clamp(0.0, 1.0)
return inpt
class MotionBlur(Transform):
"""Apply directional motion blur to simulate fast robot or object movement.
Generates a 1D averaging kernel along a random direction, applied via depthwise convolution.
Args:
kernel_size: An odd kernel size or a range containing at least one odd kernel size.
"""
def __init__(self, kernel_size: int | Sequence[int] = (3, 11)) -> None:
super().__init__()
if isinstance(kernel_size, int):
self.kernel_size = (kernel_size, kernel_size)
elif isinstance(kernel_size, Sequence) and len(kernel_size) == 2:
self.kernel_size = (int(kernel_size[0]), int(kernel_size[1]))
else:
raise TypeError("kernel_size must be an int or a sequence with length 2.")
if not 1 <= self.kernel_size[0] <= self.kernel_size[1]:
raise ValueError(f"kernel_size must satisfy 1 <= min <= max, but got {self.kernel_size}.")
self._first_odd_kernel_size = self.kernel_size[0] + (self.kernel_size[0] + 1) % 2
if self._first_odd_kernel_size > self.kernel_size[1]:
raise ValueError(f"kernel_size range must contain an odd value, but got {self.kernel_size}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
num_odd_sizes = (self.kernel_size[1] - self._first_odd_kernel_size) // 2 + 1
size_index = int(torch.randint(0, num_odd_sizes, ()).item())
ks = self._first_odd_kernel_size + 2 * size_index
angle = torch.empty(1).uniform_(0, 360).item()
return {"kernel_size": ks, "angle": angle}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
if inpt.ndim < 3:
raise ValueError(f"MotionBlur expects [..., C, H, W] input, but got shape {inpt.shape}.")
kernel_size = params["kernel_size"]
radius = kernel_size // 2
angle = math.radians(params["angle"])
positions = torch.linspace(-radius, radius, kernel_size, device=inpt.device)
x_coords = (positions * math.cos(angle)).round().to(torch.long) + radius
y_coords = (positions * math.sin(angle)).round().to(torch.long) + radius
kernel = torch.zeros((kernel_size, kernel_size), device=inpt.device, dtype=inpt.dtype)
kernel[y_coords, x_coords] = 1
kernel /= kernel.sum()
channels, height, width = inpt.shape[-3:]
flat_input = inpt.reshape(-1, channels, height, width)
depthwise_kernel = kernel.expand(channels, 1, kernel_size, kernel_size)
padded = torch.nn.functional.pad(flat_input, (radius,) * 4, mode="replicate")
output = torch.nn.functional.conv2d(padded, depthwise_kernel, groups=channels)
return output.reshape(inpt.shape).clamp(0.0, 1.0)
class JPEGCompression(Transform):
"""Simulate JPEG compression artifacts (block artifacts, color banding).
Models quality degradation from video compression in network-streamed camera feeds.
Args:
quality: Range (min, max) for JPEG quality factor (lower = more artifacts).
"""
def __init__(self, quality: int | Sequence[int] = (15, 75)) -> None:
super().__init__()
if isinstance(quality, int):
self.quality = (quality, quality)
elif isinstance(quality, Sequence) and len(quality) == 2:
self.quality = (int(quality[0]), int(quality[1]))
else:
raise TypeError("quality must be an int or a sequence with length 2.")
if not 1 <= self.quality[0] <= self.quality[1] <= 100:
raise ValueError(f"quality must satisfy 1 <= min <= max <= 100, but got {self.quality}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
return {"quality": int(torch.randint(self.quality[0], self.quality[1] + 1, (1,)).item())}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
if inpt.ndim < 3:
raise ValueError(f"JPEGCompression expects [..., C, H, W] input, but got shape {inpt.shape}.")
channels, height, width = inpt.shape[-3:]
if channels not in (1, 3):
raise ValueError(f"JPEGCompression expects 1 or 3 channels, but got {channels}.")
flat_input = inpt.reshape(-1, channels, height, width)
flat_uint8 = (flat_input.clamp(0.0, 1.0) * 255).round().to(torch.uint8).cpu()
decoded_frames = [
decode_image(encode_jpeg(frame, quality=params["quality"])) for frame in flat_uint8.unbind()
]
output = torch.stack(decoded_frames).to(device=inpt.device, dtype=inpt.dtype) / 255.0
return output.reshape(inpt.shape)
class GaussianPatchBrightness(Transform):
"""Apply spatially-varying brightness with Gaussian patches.
Simulates uneven overhead lighting, spotlights, and shadow patches commonly
encountered in real robot workspaces with multiple light sources.
Args:
num_patches: Range (min, max) for number of brightness patches.
sigma_range: Range for Gaussian sigma as fraction of image size.
factor_range: Range for brightness factor (< 1 darkens, > 1 brightens).
"""
def __init__(
self,
num_patches: int | Sequence[int] = (1, 4),
sigma_range: Sequence[float] = (0.05, 0.25),
factor_range: Sequence[float] = (0.4, 1.6),
) -> None:
super().__init__()
if isinstance(num_patches, int):
self.num_patches = (num_patches, num_patches)
elif isinstance(num_patches, Sequence) and len(num_patches) == 2:
self.num_patches = (int(num_patches[0]), int(num_patches[1]))
else:
raise TypeError("num_patches must be an int or a sequence with length 2.")
if not 1 <= self.num_patches[0] <= self.num_patches[1]:
raise ValueError(f"num_patches must satisfy 1 <= min <= max, but got {self.num_patches}.")
if not isinstance(sigma_range, Sequence) or len(sigma_range) != 2:
raise TypeError("sigma_range must be a sequence with length 2.")
self.sigma_range = (float(sigma_range[0]), float(sigma_range[1]))
if not 0.0 < self.sigma_range[0] <= self.sigma_range[1]:
raise ValueError(f"sigma_range must satisfy 0 < min <= max, but got {self.sigma_range}.")
if not isinstance(factor_range, Sequence) or len(factor_range) != 2:
raise TypeError("factor_range must be a sequence with length 2.")
self.factor_range = (float(factor_range[0]), float(factor_range[1]))
if not 0.0 <= self.factor_range[0] <= self.factor_range[1]:
raise ValueError(f"factor_range must satisfy 0 <= min <= max, but got {self.factor_range}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
n = int(torch.randint(self.num_patches[0], self.num_patches[1] + 1, (1,)).item())
return {
"centers": torch.rand(n, 2).tolist(),
"sigmas": torch.empty(n).uniform_(self.sigma_range[0], self.sigma_range[1]).tolist(),
"factors": torch.empty(n).uniform_(self.factor_range[0], self.factor_range[1]).tolist(),
}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
h, w = inpt.shape[-2:]
mask = torch.ones(h, w, device=inpt.device, dtype=inpt.dtype)
grid_y = torch.linspace(0, 1, h, device=inpt.device, dtype=inpt.dtype)
grid_x = torch.linspace(0, 1, w, device=inpt.device, dtype=inpt.dtype)
yy, xx = torch.meshgrid(grid_y, grid_x, indexing="ij")
for (cy, cx), sigma, factor in zip(
params["centers"], params["sigmas"], params["factors"], strict=True
):
gauss = torch.exp(-((yy - cy) ** 2 + (xx - cx) ** 2) / (2 * sigma**2))
mask = mask * (1.0 + (factor - 1.0) * gauss)
broadcast_shape = (1,) * (inpt.ndim - 2) + (h, w)
return (inpt * mask.reshape(broadcast_shape)).clamp(0.0, 1.0)
class RandomShadow(Transform):
"""Add random vertical band shadow with smooth edges.
Simulates cast shadows from objects or people near the robot workspace.
Symmetric: randomly brightens or darkens to prevent BatchNorm stats shift.
Args:
opacity: Range (min, max) for shadow/highlight opacity.
"""
def __init__(self, opacity: float | Sequence[float] = (0.3, 0.6)) -> None:
super().__init__()
if isinstance(opacity, (int, float)):
self.opacity = (float(opacity), float(opacity))
elif isinstance(opacity, Sequence) and len(opacity) == 2:
self.opacity = (float(opacity[0]), float(opacity[1]))
else:
raise TypeError("opacity must be a number or a sequence with length 2.")
if not 0.0 <= self.opacity[0] <= self.opacity[1] <= 1.0:
raise ValueError(f"opacity must satisfy 0 <= min <= max <= 1, but got {self.opacity}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
return {
"opacity": torch.empty(1).uniform_(self.opacity[0], self.opacity[1]).item(),
"start": torch.rand(1).item(),
"width": torch.empty(1).uniform_(1 / 3, 2 / 3).item(),
"direction": -1.0 if torch.rand(1).item() < 0.5 else 1.0,
}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
if inpt.ndim < 3:
raise ValueError(f"RandomShadow expects [..., C, H, W] input, but got shape {inpt.shape}.")
h, w = inpt.shape[-2:]
band_width = max(1, min(w, round(params["width"] * w)))
x_start = round(params["start"] * (w - band_width))
x_end = x_start + band_width
mask = torch.ones(h, w, device=inpt.device, dtype=inpt.dtype)
mask[:, x_start:x_end] = 1.0 + params["direction"] * params["opacity"]
smoothing_size = min(8, h, w)
if smoothing_size > 1:
batched_mask = mask[None, None]
small = torch.nn.functional.avg_pool2d(batched_mask, smoothing_size, stride=smoothing_size)
mask = torch.nn.functional.interpolate(small, size=(h, w), mode="bilinear", align_corners=False)[
0, 0
]
broadcast_shape = (1,) * (inpt.ndim - 2) + (h, w)
return (inpt * mask.reshape(broadcast_shape)).clamp(0.0, 1.0)
class CoarseDropout(Transform):
"""Drop random rectangular patches to simulate partial occlusion.
Models objects, hands, or cables passing through the camera field of view
during robot manipulation.
Args:
max_holes: Maximum number of rectangular patches to drop.
max_height_frac: Maximum patch height as fraction of image height.
max_width_frac: Maximum patch width as fraction of image width.
fill_value: Value to fill dropped regions with.
"""
def __init__(
self,
max_holes: int = 8,
max_height_frac: float = 0.07,
max_width_frac: float = 0.07,
fill_value: float = 0.0,
) -> None:
super().__init__()
if not isinstance(max_holes, int):
raise TypeError("max_holes must be an int.")
if max_holes < 1:
raise ValueError(f"max_holes must be at least 1, but got {max_holes}.")
if not 0.0 < max_height_frac <= 1.0:
raise ValueError(f"max_height_frac must be in (0, 1], but got {max_height_frac}.")
if not 0.0 < max_width_frac <= 1.0:
raise ValueError(f"max_width_frac must be in (0, 1], but got {max_width_frac}.")
if not 0.0 <= fill_value <= 1.0:
raise ValueError(f"fill_value must be in [0, 1], but got {fill_value}.")
self.max_holes = max_holes
self.max_height_frac = max_height_frac
self.max_width_frac = max_width_frac
self.fill_value = fill_value
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
n = int(torch.randint(1, self.max_holes + 1, (1,)).item())
sizes = torch.rand(n, 2)
sizes[:, 0] *= self.max_height_frac
sizes[:, 1] *= self.max_width_frac
return {"sizes": sizes.tolist(), "positions": torch.rand(n, 2).tolist()}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
if inpt.ndim < 3:
raise ValueError(f"CoarseDropout expects [..., C, H, W] input, but got shape {inpt.shape}.")
h, w = inpt.shape[-2:]
result = inpt.clone()
for (height_frac, width_frac), (y_frac, x_frac) in zip(
params["sizes"], params["positions"], strict=True
):
hole_h = max(1, min(h, round(height_frac * h)))
hole_w = max(1, min(w, round(width_frac * w)))
y = round(y_frac * (h - hole_h))
x = round(x_frac * (w - hole_w))
result[..., y : y + hole_h, x : x + hole_w] = self.fill_value
return result
class GammaCorrection(Transform):
"""Apply random gamma correction to simulate exposure variation.
Models different camera auto-exposure settings and sensor response curves.
Uses log-symmetric sampling so brightening and darkening are equally likely,
preventing BatchNorm statistics shift.
Args:
gamma: Range (min, max) for gamma value. Values < 1 brighten, > 1 darken.
"""
def __init__(self, gamma: float | Sequence[float] = (0.5, 2.0)) -> None:
super().__init__()
if isinstance(gamma, (int, float)):
gamma = float(gamma)
if gamma <= 0:
raise ValueError(f"gamma must be positive, but got {gamma}.")
self.gamma = (min(gamma, 1.0 / gamma), max(gamma, 1.0 / gamma))
elif isinstance(gamma, Sequence) and len(gamma) == 2:
self.gamma = (float(gamma[0]), float(gamma[1]))
else:
raise TypeError("gamma must be a number or a sequence with length 2.")
if not 0.0 < self.gamma[0] <= self.gamma[1]:
raise ValueError(f"gamma must satisfy 0 < min <= max, but got {self.gamma}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
log_lo = math.log(self.gamma[0])
log_hi = math.log(self.gamma[1])
gamma = math.exp(torch.empty(1).uniform_(log_lo, log_hi).item())
return {"gamma": gamma}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if isinstance(inpt, torch.Tensor) and inpt.is_floating_point():
return inpt.pow(params["gamma"]).clamp(0.0, 1.0)
return inpt
# From the paper authors' MIT-licensed reference implementation:
# https://github.com/TheZino/PlanckianJitter
_PLANCKIAN_BLACKBODY_COEFFICIENTS = (
(0.6743, 0.4029, 0.0013),
(0.6281, 0.4241, 0.1665),
(0.5919, 0.4372, 0.2513),
(0.5623, 0.4457, 0.3154),
(0.5376, 0.4515, 0.3672),
(0.5163, 0.4555, 0.4103),
(0.4979, 0.4584, 0.4468),
(0.4816, 0.4604, 0.4782),
(0.4672, 0.4619, 0.5053),
(0.4542, 0.4630, 0.5289),
(0.4426, 0.4638, 0.5497),
(0.4320, 0.4644, 0.5681),
(0.4223, 0.4648, 0.5844),
(0.4135, 0.4651, 0.5990),
(0.4054, 0.4653, 0.6121),
(0.3980, 0.4654, 0.6239),
(0.3911, 0.4655, 0.6346),
(0.3847, 0.4656, 0.6444),
(0.3787, 0.4656, 0.6532),
(0.3732, 0.4656, 0.6613),
(0.3680, 0.4655, 0.6688),
(0.3632, 0.4655, 0.6756),
(0.3586, 0.4655, 0.6820),
(0.3544, 0.4654, 0.6878),
(0.3503, 0.4653, 0.6933),
)
_PLANCKIAN_MIN_TEMPERATURE = 3_000
_PLANCKIAN_MAX_TEMPERATURE = 15_000
_PLANCKIAN_TEMPERATURE_STEP = 500
class PlanckianJitter(Transform):
"""Simulate color temperature shift along the Planckian locus.
Samples one black-body temperature and applies the corresponding correlated red
and blue channel scaling while preserving the green channel. Coefficients between
the tabulated 500 K intervals are linearly interpolated.
Reference: Zini et al., "Planckian Jitter", CVPR 2022 Workshop.
Args:
temperature: A fixed color temperature or range in Kelvin. Supported values
are between 3000 K and 15000 K.
"""
def __init__(self, temperature: int | Sequence[int] = (3_000, 15_000)) -> None:
super().__init__()
if isinstance(temperature, int):
self.temperature = (temperature, temperature)
elif isinstance(temperature, Sequence) and len(temperature) == 2:
self.temperature = (int(temperature[0]), int(temperature[1]))
else:
raise TypeError("temperature must be an int or a sequence with length 2.")
if not (
_PLANCKIAN_MIN_TEMPERATURE
<= self.temperature[0]
<= self.temperature[1]
<= _PLANCKIAN_MAX_TEMPERATURE
):
raise ValueError(
"temperature must satisfy "
f"{_PLANCKIAN_MIN_TEMPERATURE} <= min <= max <= {_PLANCKIAN_MAX_TEMPERATURE}, "
f"but got {self.temperature}."
)
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
temperature = int(torch.randint(self.temperature[0], self.temperature[1] + 1, ()).item())
return {"temperature": temperature}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
if inpt.ndim < 3 or inpt.shape[-3] != 3:
raise ValueError(f"PlanckianJitter expects [..., 3, H, W] input, but got shape {inpt.shape}.")
table_position = (params["temperature"] - _PLANCKIAN_MIN_TEMPERATURE) / _PLANCKIAN_TEMPERATURE_STEP
left_index = math.floor(table_position)
right_index = min(left_index + 1, len(_PLANCKIAN_BLACKBODY_COEFFICIENTS) - 1)
interpolation_weight = table_position - left_index
left = torch.tensor(
_PLANCKIAN_BLACKBODY_COEFFICIENTS[left_index],
device=inpt.device,
dtype=inpt.dtype,
)
right = torch.tensor(
_PLANCKIAN_BLACKBODY_COEFFICIENTS[right_index],
device=inpt.device,
dtype=inpt.dtype,
)
coefficients = torch.lerp(left, right, interpolation_weight)
scale = torch.stack(
(
coefficients[0] / coefficients[1],
coefficients.new_tensor(1.0),
coefficients[2] / coefficients[1],
)
)
broadcast_shape = (1,) * (inpt.ndim - 3) + (3, 1, 1)
return (inpt * scale.reshape(broadcast_shape)).clamp(0.0, 1.0)
_CUSTOM_TRANSFORMS: dict[str, type[Transform]] = {
"SharpnessJitter": SharpnessJitter,
"GaussianNoise": GaussianNoise,
"MotionBlur": MotionBlur,
"JPEGCompression": JPEGCompression,
"GaussianPatchBrightness": GaussianPatchBrightness,
"RandomShadow": RandomShadow,
"CoarseDropout": CoarseDropout,
"GammaCorrection": GammaCorrection,
"PlanckianJitter": PlanckianJitter,
}
@dataclass @dataclass
class ImageTransformConfig: class ImageTransformConfig:
""" """
@@ -683,17 +216,16 @@ class ImageTransformsConfig:
def make_transform_from_config(cfg: ImageTransformConfig) -> Transform: def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
if cfg.type in _CUSTOM_TRANSFORMS: if cfg.type == "SharpnessJitter":
return _CUSTOM_TRANSFORMS[cfg.type](**cfg.kwargs) return SharpnessJitter(**cfg.kwargs)
transform_cls = getattr(v2, cfg.type, None) transform_cls = getattr(v2, cfg.type, None)
if isinstance(transform_cls, type) and issubclass(transform_cls, Transform): if isinstance(transform_cls, type) and issubclass(transform_cls, Transform):
return transform_cls(**cfg.kwargs) return transform_cls(**cfg.kwargs)
valid_custom = ", ".join(sorted(_CUSTOM_TRANSFORMS.keys()))
raise ValueError( raise ValueError(
f"Transform '{cfg.type}' is not valid. It must be a class in " f"Transform '{cfg.type}' is not valid. It must be a class in "
f"torchvision.transforms.v2 or one of: {valid_custom}." f"torchvision.transforms.v2 or 'SharpnessJitter'."
) )
+13 -4
View File
@@ -37,16 +37,25 @@ def auto_select_torch_device() -> torch.device:
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level # TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device: def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
"""Given a string, return a torch.device with checks on whether the device is available.""" """Given a string, return a torch.device with checks on whether the device is available.
Raises:
ValueError: If the requested device family is known but not available on
this machine (``AssertionError`` was previously used and is easy to
mistake for a programmer bug under ``python -O`` where asserts vanish).
"""
try_device = str(try_device) try_device = str(try_device)
if try_device.startswith("cuda"): if try_device.startswith("cuda"):
assert torch.cuda.is_available() if not torch.cuda.is_available():
raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
device = torch.device(try_device) device = torch.device(try_device)
elif try_device == "mps": elif try_device == "mps":
assert torch.backends.mps.is_available() if not torch.backends.mps.is_available():
raise ValueError("Requested device 'mps' but MPS is not available.")
device = torch.device("mps") device = torch.device("mps")
elif try_device == "xpu": elif try_device == "xpu":
assert torch.xpu.is_available() if not torch.xpu.is_available():
raise ValueError("Requested device 'xpu' but XPU is not available.")
device = torch.device("xpu") device = torch.device("xpu")
elif try_device == "cpu": elif try_device == "cpu":
device = torch.device("cpu") device = torch.device("cpu")
+5 -5
View File
@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
Returns: Returns:
Any: The data loaded from the JSON file. Any: The data loaded from the JSON file.
""" """
with open(fpath) as f: with open(fpath, encoding="utf-8") as f:
return json.load(f) return json.load(f)
def write_json(data: dict, fpath: Path) -> None: def write_json(data: JsonLike, fpath: Path) -> None:
"""Write data to a JSON file. """Write JSON-serializable data to a file.
Creates parent directories if they don't exist. Creates parent directories if they don't exist.
Args: Args:
data (dict): The dictionary to write. data: JSON-serializable data to write.
fpath (Path): The path to the output JSON file. fpath (Path): The path to the output JSON file.
""" """
fpath.parent.mkdir(exist_ok=True, parents=True) fpath.parent.mkdir(exist_ok=True, parents=True)
with open(fpath, "w") as f: with open(fpath, "w", encoding="utf-8") as f:
json.dump(data, f, indent=4, ensure_ascii=False) json.dump(data, f, indent=4, ensure_ascii=False)
+4
View File
@@ -30,6 +30,10 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
""" """
if seconds <= 0: if seconds <= 0:
return return
if spin_threshold < 0:
raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
if sleep_margin < 0:
raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")
system = platform.system() system = platform.system()
# On macOS and Windows the scheduler / sleep granularity can make # On macOS and Windows the scheduler / sleep granularity can make
+5 -2
View File
@@ -29,9 +29,12 @@ class Rotation:
def __init__(self, quat: np.ndarray) -> None: def __init__(self, quat: np.ndarray) -> None:
"""Initialize rotation from quaternion [x, y, z, w].""" """Initialize rotation from quaternion [x, y, z, w]."""
self._quat = np.asarray(quat, dtype=float) self._quat = np.asarray(quat, dtype=float)
# Normalize quaternion if self._quat.shape != (4,):
raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
# Normalize quaternion. Reject the zero vector — it has no orientation.
norm = np.linalg.norm(self._quat) norm = np.linalg.norm(self._quat)
if norm > 0: if norm <= 0.0 or not np.isfinite(norm):
raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})")
self._quat = self._quat / norm self._quat = self._quat / norm
@classmethod @classmethod
+2 -2
View File
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
from typing import TypedDict from typing import NotRequired, TypedDict
import torch import torch
@@ -28,7 +28,7 @@ class Transition(TypedDict):
next_state: dict[str, torch.Tensor] next_state: dict[str, torch.Tensor]
done: bool done: bool
truncated: bool truncated: bool
complementary_info: dict[str, torch.Tensor | float | int] | None = None complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition: def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
+9 -8
View File
@@ -24,7 +24,6 @@ import sys
import time import time
from collections.abc import Iterator from collections.abc import Iterator
from copy import copy, deepcopy from copy import copy, deepcopy
from datetime import datetime
from pathlib import Path from pathlib import Path
from statistics import mean from statistics import mean
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
@@ -61,14 +60,16 @@ def init_logging(
accelerator: Optional Accelerator instance (for multi-GPU detection) accelerator: Optional Accelerator instance (for multi-GPU detection)
""" """
def custom_format(record: logging.LogRecord) -> str: class LeRobotFormatter(logging.Formatter):
dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S") def format(self, record: logging.LogRecord) -> str:
fnameline = f"{record.pathname}:{record.lineno}" record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:]
pid_str = f"[PID: {os.getpid()}] " if display_pid else "" record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else ""
return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}" return super().format(record)
formatter = logging.Formatter() formatter = LeRobotFormatter(
formatter.format = custom_format "%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger() logger = logging.getLogger()
logger.setLevel(logging.NOTSET) logger.setLevel(logging.NOTSET)
+4 -3
View File
@@ -322,15 +322,16 @@ def test_get_color_sensor_prefers_rgb_camera():
assert camera._get_color_sensor() is rgb assert camera._get_color_sensor() is rgb
def test_get_color_sensor_falls_back_to_stereo_module(): def test_get_color_sensor_raises_without_dedicated_rgb_module():
"""D405 has no separate RGB module; color comes from Stereo Module.""" """D405 has no separate RGB module; we refuse to touch the shared Stereo Module."""
config = RealSenseCameraConfig(serial_number_or_name="042") config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config) camera = RealSenseCamera(config)
stereo = _make_mock_sensor("Stereo Module") stereo = _make_mock_sensor("Stereo Module")
_attach_mock_color_sensor(camera, stereo) _attach_mock_color_sensor(camera, stereo)
assert camera._get_color_sensor() is stereo with pytest.raises(RuntimeError, match="dedicated 'RGB Camera' module"):
camera._get_color_sensor()
def test_get_color_sensor_raises_with_available_sensors(): def test_get_color_sensor_raises_with_available_sensors():
-158
View File
@@ -28,17 +28,9 @@ from lerobot.scripts.lerobot_imgtransform_viz import (
save_each_transform, save_each_transform,
) )
from lerobot.transforms import ( from lerobot.transforms import (
CoarseDropout,
GammaCorrection,
GaussianNoise,
GaussianPatchBrightness,
ImageTransformConfig, ImageTransformConfig,
ImageTransforms, ImageTransforms,
ImageTransformsConfig, ImageTransformsConfig,
JPEGCompression,
MotionBlur,
PlanckianJitter,
RandomShadow,
RandomSubsetApply, RandomSubsetApply,
SharpnessJitter, SharpnessJitter,
make_transform_from_config, make_transform_from_config,
@@ -463,153 +455,3 @@ def test_save_each_transform(img_tensor_factory, tmp_path):
assert (transform_dir / file_name).exists(), ( assert (transform_dir / file_name).exists(), (
f"{file_name} was not found in {transform} directory." f"{file_name} was not found in {transform} directory."
) )
# --- Tests for robotics-relevant augmentations ---
ROBOTICS_TRANSFORMS = [
("GaussianNoise", GaussianNoise, {"std": (5.0, 25.0)}),
("MotionBlur", MotionBlur, {"kernel_size": (3, 11)}),
("JPEGCompression", JPEGCompression, {"quality": (15, 75)}),
("GaussianPatchBrightness", GaussianPatchBrightness, {}),
("RandomShadow", RandomShadow, {"opacity": (0.3, 0.6)}),
("CoarseDropout", CoarseDropout, {"max_holes": 8}),
("GammaCorrection", GammaCorrection, {"gamma": (0.5, 2.0)}),
("PlanckianJitter", PlanckianJitter, {"temperature": (3_000, 15_000)}),
]
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
def test_robotics_transform_shape_preserved(name, cls, kwargs, img_tensor_factory):
img = img_tensor_factory()
tf = cls(**kwargs)
out = tf(img)
assert out.shape == img.shape, f"{name} changed shape: {img.shape} -> {out.shape}"
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
def test_robotics_transform_output_range(name, cls, kwargs, img_tensor_factory):
img = img_tensor_factory()
tf = cls(**kwargs)
out = tf(img)
assert out.min() >= -0.01, f"{name} min below range: {out.min():.4f}"
assert out.max() <= 1.01, f"{name} max above range: {out.max():.4f}"
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
def test_robotics_transform_float_output(name, cls, kwargs, img_tensor_factory):
img = img_tensor_factory()
tf = cls(**kwargs)
out = tf(img)
assert out.is_floating_point(), f"{name} output dtype={out.dtype}"
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
def test_robotics_transform_non_float_passthrough(name, cls, kwargs):
int_img = torch.randint(0, 255, (3, 32, 32), dtype=torch.uint8)
tf = cls(**kwargs)
out = tf(int_img)
assert torch.equal(out, int_img), f"{name} modified non-float input"
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
def test_robotics_transform_via_config(name, cls, kwargs):
cfg = ImageTransformConfig(type=name, kwargs=kwargs)
tf = make_transform_from_config(cfg)
assert isinstance(tf, cls), f"Config produced {type(tf)}, expected {cls}"
def test_make_transform_error_message_includes_custom():
"""Error message should list all registered custom transforms."""
with pytest.raises(ValueError, match="GaussianNoise"):
make_transform_from_config(ImageTransformConfig(type="NonExistent"))
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
@pytest.mark.parametrize("shape", [(4, 3, 32, 32), (2, 4, 3, 16, 16)])
def test_robotics_transform_supports_temporal_batches(name, cls, kwargs, shape):
img = torch.rand(shape)
out = cls(**kwargs)(img)
assert out.shape == img.shape, f"{name} changed shape: {img.shape} -> {out.shape}"
assert out.min() >= 0
assert out.max() <= 1
@pytest.mark.parametrize(
"cls,kwargs",
[
(GaussianNoise, {"std": (25.0, 25.0)}),
(MotionBlur, {"kernel_size": 5}),
(JPEGCompression, {"quality": 10}),
(
GaussianPatchBrightness,
{"num_patches": 1, "sigma_range": (0.2, 0.2), "factor_range": (0.5, 0.5)},
),
(RandomShadow, {"opacity": 0.5}),
(CoarseDropout, {"max_holes": 1, "fill_value": 0.0}),
(GammaCorrection, {"gamma": (2.0, 2.0)}),
(PlanckianJitter, {"temperature": 3_000}),
],
)
def test_robotics_transform_is_not_silent_noop(cls, kwargs):
img = torch.rand(3, 32, 32)
out = cls(**kwargs)(img)
assert not torch.equal(out, img)
@pytest.mark.parametrize(
"transform",
[
GaussianNoise(std=25),
RandomShadow(opacity=0.5),
CoarseDropout(max_holes=4),
],
)
def test_robotics_transform_random_params_are_reused(transform):
img = torch.rand(3, 32, 32)
params = transform.make_params([img])
torch.testing.assert_close(transform.transform(img, params), transform.transform(img, params))
def test_motion_blur_kernel_size_stays_in_configured_range():
transform = MotionBlur(kernel_size=(4, 10))
sampled_sizes = {transform.make_params([])["kernel_size"] for _ in range(100)}
assert sampled_sizes <= {5, 7, 9}
assert sampled_sizes
def test_gamma_correction_scalar_below_one_defines_symmetric_range():
transform = GammaCorrection(gamma=0.5)
assert transform.gamma == (0.5, 2.0)
assert transform(torch.rand(3, 8, 8)).shape == (3, 8, 8)
def test_planckian_jitter_uses_correlated_temperature_coefficients():
img = torch.full((2, 3, 8, 8), 0.25)
out = PlanckianJitter(temperature=3_000)(img)
torch.testing.assert_close(out[:, 1], img[:, 1])
assert torch.all(out[:, 0] > out[:, 1])
assert torch.all(out[:, 2] < out[:, 1])
def test_random_shadow_supports_small_images():
img = torch.rand(3, 7, 7)
assert RandomShadow()(img).shape == img.shape
@pytest.mark.parametrize(
"cls,kwargs",
[
(GaussianNoise, {"std": (-1.0, 1.0)}),
(MotionBlur, {"kernel_size": 4}),
(JPEGCompression, {"quality": (0, 75)}),
(GaussianPatchBrightness, {"sigma_range": (0.0, 0.25)}),
(RandomShadow, {"opacity": (0.3, 1.1)}),
(CoarseDropout, {"max_holes": 0}),
(GammaCorrection, {"gamma": 0.0}),
(PlanckianJitter, {"temperature": (2_000, 6_500)}),
],
)
def test_robotics_transform_rejects_invalid_config(cls, kwargs):
with pytest.raises(ValueError):
cls(**kwargs)
+54
View File
@@ -496,6 +496,60 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
assert "embodiment_id" in processed assert "embodiment_id" in processed
def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path):
"""Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded.
Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer
overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats
are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to
max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the
stats or normalization crashes with a shape mismatch.
"""
config = make_config()
preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
preprocessor.save_pretrained(tmp_path)
postprocessor.save_pretrained(tmp_path)
# Reload with the generic override path injecting raw, unpadded dataset stats.
raw_stats = make_stats()
loaded_pre = PolicyProcessorPipeline.from_pretrained(
tmp_path,
config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
overrides={"normalizer_processor": {"stats": raw_stats}},
to_transition=batch_to_transition,
to_output=transition_to_batch,
)
loaded_post = PolicyProcessorPipeline.from_pretrained(
tmp_path,
config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
overrides={"unnormalizer_processor": {"stats": raw_stats}},
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
)
# Sanity: the override really injected unpadded stats before reconciliation.
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,)
loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post)
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep))
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
# Normalizing a padded state must not raise (this is the exact runtime path that crashed).
processed = loaded_pre(
{
"task": "pick the block",
OBS_STATE: torch.zeros(STATE_DIM),
f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
}
)
assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM)
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch): def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model) monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
policy = modeling_evo1.Evo1Policy(make_config()) policy = modeling_evo1.Evo1Policy(make_config())
@@ -0,0 +1,45 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.robots.so_follower.robot_kinematic_processor import (
ForwardKinematicsJointsToEEAction,
ForwardKinematicsJointsToEEObservation,
)
MOTOR_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
EE_KEYS = {f"ee.{k}" for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]}
def _joint_bucket(feature_type: FeatureType) -> dict[str, PolicyFeature]:
return {f"{n}.pos": PolicyFeature(type=feature_type, shape=(1,)) for n in MOTOR_NAMES}
@pytest.mark.parametrize(
("step_cls", "bucket", "feature_type"),
[
(ForwardKinematicsJointsToEEAction, PipelineFeatureType.ACTION, FeatureType.ACTION),
(ForwardKinematicsJointsToEEObservation, PipelineFeatureType.OBSERVATION, FeatureType.STATE),
],
)
def test_fk_feature_schema(step_cls, bucket, feature_type):
features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}}
features[bucket] = _joint_bucket(feature_type)
out = step_cls(kinematics=None, motor_names=MOTOR_NAMES).transform_features(features)[bucket]
assert set(out) == EE_KEYS
assert {feature.type for feature in out.values()} == {feature_type}
+36
View File
@@ -0,0 +1,36 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest.mock import patch
import pytest
import torch
from lerobot.utils.device_utils import get_safe_torch_device, is_torch_device_available
def test_cpu_always_available():
assert get_safe_torch_device("cpu") == torch.device("cpu")
assert is_torch_device_available("cpu")
def test_missing_cuda_raises_valueerror():
with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"):
get_safe_torch_device("cuda")
def test_missing_mps_raises_valueerror():
with patch("torch.backends.mps.is_available", return_value=False), pytest.raises(ValueError, match="MPS"):
get_safe_torch_device("mps")
+46
View File
@@ -0,0 +1,46 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import pytest
from lerobot.utils.rotation import Rotation
def test_zero_quaternion_rejected():
with pytest.raises(ValueError, match="non-zero"):
Rotation(np.zeros(4))
def test_non_finite_quaternion_rejected():
with pytest.raises(ValueError, match="non-zero|finite"):
Rotation(np.array([np.nan, 0.0, 0.0, 1.0]))
def test_wrong_shape_rejected():
with pytest.raises(ValueError, match="shape"):
Rotation(np.array([1.0, 0.0, 0.0]))
def test_identity_roundtrip():
r = Rotation.from_rotvec(np.zeros(3))
assert np.allclose(r.as_rotvec(), 0.0)
assert np.allclose(r.as_matrix(), np.eye(3))
def test_rotvec_roundtrip():
rotvec = np.array([0.1, -0.2, 0.3])
r = Rotation.from_rotvec(rotvec)
assert np.allclose(r.as_rotvec(), rotvec, atol=1e-6)