Compare commits

..

1 Commits

Author SHA1 Message Date
CarolinePascal 35d40f353f docs(teleoperators): write the API reference docstrings
Completes Wave 1. Takes src/lerobot/teleoperators/ (excluding teleoperator.py, off-limits) to 100% public
docstring coverage across all 16 hardware families. Fixes a real check_docstrings.py-breaking bug in
ExoskeletonIKHelper's docstring format. Several other real bugs (missing @property, undefined attribute
reference, wrong parameter name in an existing docstring) were found and documented accurately but left
unfixed per the docstrings-only scope, detailed in the PR description.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-07 01:25:26 +02:00
64 changed files with 2770 additions and 948 deletions
-34
View File
@@ -21,37 +21,3 @@ See [Using LeRobotDataset](../lerobot-dataset-v3) for the format and the common
## StreamingLeRobotDataset
[[autodoc]] lerobot.datasets.StreamingLeRobotDataset
## EpisodeAwareSampler
[[autodoc]] lerobot.datasets.sampler.EpisodeAwareSampler
## Editing a dataset
Functions in `lerobot.datasets.dataset_tools` for editing an existing `LeRobotDataset` on disk: adding,
removing, or modifying features; splitting, merging, or deleting episodes; re-encoding video; and
recomputing statistics. Each returns a new dataset rather than mutating the source in place.
[[autodoc]] lerobot.datasets.dataset_tools.add_features
[[autodoc]] lerobot.datasets.dataset_tools.remove_feature
[[autodoc]] lerobot.datasets.dataset_tools.modify_features
[[autodoc]] lerobot.datasets.dataset_tools.modify_tasks
[[autodoc]] lerobot.datasets.dataset_tools.delete_episodes
[[autodoc]] lerobot.datasets.dataset_tools.split_dataset
[[autodoc]] lerobot.datasets.dataset_tools.merge_datasets
[[autodoc]] lerobot.datasets.dataset_tools.recompute_stats
[[autodoc]] lerobot.datasets.dataset_tools.reencode_dataset
[[autodoc]] lerobot.datasets.dataset_tools.convert_image_to_video_dataset
## Aggregating datasets
[[autodoc]] lerobot.datasets.aggregate.aggregate_datasets
+228
View File
@@ -28,3 +28,231 @@ See [Phone teleoperation](../phone_teleop) and [Isaac Teleop](../isaac_teleop) f
## make_teleoperator_from_config
[[autodoc]] lerobot.teleoperators.make_teleoperator_from_config
## SO-100 and SO-101 leaders
`SO100Leader` and `SO101Leader` are aliases of the same `SOLeader` class; the two arms differ in their
configuration, not their control code. `SO100LeaderConfig` and `SO101LeaderConfig` are likewise aliases of
`SOLeaderTeleopConfig`.
[[autodoc]] lerobot.teleoperators.so_leader.SOLeader
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.so_leader.SOLeaderTeleopConfig
## KochLeader
[[autodoc]] lerobot.teleoperators.koch_leader.KochLeader
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.koch_leader.KochLeaderConfig
## OmxLeader
[[autodoc]] lerobot.teleoperators.omx_leader.OmxLeader
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.omx_leader.OmxLeaderConfig
## OpenArmLeader
CAN-based leader arm using Damiao motors.
[[autodoc]] lerobot.teleoperators.openarm_leader.OpenArmLeader
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.openarm_leader.OpenArmLeaderConfig
## BiOpenArmLeader
A bimanual pair of `OpenArmLeader` arms.
[[autodoc]] lerobot.teleoperators.bi_openarm_leader.BiOpenArmLeader
- all
[[autodoc]] lerobot.teleoperators.bi_openarm_leader.BiOpenArmLeaderConfig
## OpenArmMini
CAN-based leader arm using Damiao motors, a smaller/simpler OpenArm variant.
[[autodoc]] lerobot.teleoperators.openarm_mini.OpenArmMini
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.openarm_mini.OpenArmMiniConfig
## BiOpenArmMini
A bimanual pair of `OpenArmMini` arms.
[[autodoc]] lerobot.teleoperators.bi_openarm_mini.BiOpenArmMini
- all
[[autodoc]] lerobot.teleoperators.bi_openarm_mini.BiOpenArmMiniConfig
## HomunculusArm
A wearable exoskeleton arm read over a serial link.
[[autodoc]] lerobot.teleoperators.homunculus.HomunculusArm
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.homunculus.HomunculusArmConfig
## HomunculusGlove
A wearable exoskeleton glove read over a serial link, remapped to HopeJR hand joints via
`homunculus_glove_to_hope_jr_hand`.
[[autodoc]] lerobot.teleoperators.homunculus.HomunculusGlove
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.homunculus.HomunculusGloveConfig
[[autodoc]] lerobot.teleoperators.homunculus.homunculus_glove_to_hope_jr_hand
## RebotArm102Leader
[[autodoc]] lerobot.teleoperators.rebot_102_leader.RebotArm102Leader
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.rebot_102_leader.RebotArm102LeaderTeleopConfig
## BiRebot102Leader
A bimanual pair of `RebotArm102Leader` arms.
[[autodoc]] lerobot.teleoperators.bi_rebot_102_leader.BiRebot102Leader
- all
[[autodoc]] lerobot.teleoperators.bi_rebot_102_leader.BiRebot102LeaderConfig
## BiSOLeader
A bimanual pair of `SOLeader` arms.
[[autodoc]] lerobot.teleoperators.bi_so_leader.BiSOLeader
- all
[[autodoc]] lerobot.teleoperators.bi_so_leader.BiSOLeaderConfig
## Phone
Reads pose and touch input from a phone app (iOS or Android).
[[autodoc]] lerobot.teleoperators.phone.Phone
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.phone.PhoneConfig
## Keyboard
`KeyboardTeleop`, `KeyboardEndEffectorTeleop`, and `KeyboardRoverTeleop` read key-press events for manual
control, targeting joint-space, end-effector, or mobile-base actions respectively.
[[autodoc]] lerobot.teleoperators.keyboard.KeyboardTeleop
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.keyboard.KeyboardTeleopConfig
[[autodoc]] lerobot.teleoperators.keyboard.KeyboardEndEffectorTeleop
- all
- action_features
[[autodoc]] lerobot.teleoperators.keyboard.KeyboardEndEffectorTeleopConfig
[[autodoc]] lerobot.teleoperators.keyboard.KeyboardRoverTeleop
- all
- action_features
- is_calibrated
[[autodoc]] lerobot.teleoperators.keyboard.KeyboardRoverTeleopConfig
## GamepadTeleop
Reads joystick/button input from a gamepad via pygame.
[[autodoc]] lerobot.teleoperators.gamepad.GamepadTeleop
- all
- action_features
- feedback_features
- is_connected
[[autodoc]] lerobot.teleoperators.gamepad.GamepadTeleopConfig
## UnitreeG1Teleoperator
A wearable exoskeleton for teleoperating the Unitree G1 humanoid's arms, mapping exoskeleton joint angles to
G1 end-effector poses via forward/inverse kinematics.
[[autodoc]] lerobot.teleoperators.unitree_g1.UnitreeG1Teleoperator
- all
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.unitree_g1.UnitreeG1TeleoperatorConfig
[[autodoc]] lerobot.teleoperators.unitree_g1.ExoskeletonArm
- all
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.unitree_g1.ExoskeletonArmPortConfig
[[autodoc]] lerobot.teleoperators.unitree_g1.ExoskeletonIKHelper
- all
[[autodoc]] lerobot.teleoperators.unitree_g1.ExoskeletonCalibration
[[autodoc]] lerobot.teleoperators.unitree_g1.ExoskeletonJointCalibration
## Reachy2Teleoperator
[[autodoc]] lerobot.teleoperators.reachy2_teleoperator.Reachy2Teleoperator
- all
- action_features
- feedback_features
- is_connected
- is_calibrated
[[autodoc]] lerobot.teleoperators.reachy2_teleoperator.Reachy2TeleoperatorConfig
+2 -1
View File
@@ -441,6 +441,7 @@ ignore = [
"src/lerobot/common/**" = ["D"]
"src/lerobot/configs/**" = ["D"]
"src/lerobot/data_processing/**" = ["D"]
"src/lerobot/datasets/**" = ["D"]
"src/lerobot/distributed/**" = ["D"]
"src/lerobot/envs/**" = ["D"]
"src/lerobot/jobs/**" = ["D"]
@@ -453,7 +454,7 @@ ignore = [
"src/lerobot/rl/**" = ["D"]
"src/lerobot/rollout/**" = ["D"]
"src/lerobot/scripts/**" = ["D"]
"src/lerobot/teleoperators/**" = ["D"]
"src/lerobot/teleoperators/teleoperator.py" = ["D"]
"src/lerobot/transforms/**" = ["D"]
"src/lerobot/transport/**" = ["D"]
"src/lerobot/utils/**" = ["D"]
+53 -89
View File
@@ -58,38 +58,12 @@ type ChunkFile = tuple[int, int]
class IndexState(TypedDict):
"""The current write cursor for a non-video (parquet) output stream during aggregation.
**Attributes**:
- **chunk** (`int`) -- The chunk index currently being written to.
- **file** (`int`) -- The file index, within `chunk`, currently being written to.
- **src_to_dst** (`dict[ChunkFile, ChunkFile]`, *optional*) -- Maps each source dataset's
`(chunk, file)` to the destination `(chunk, file)` its rows were merged into.
"""
chunk: int
file: int
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
class VideoIndex(TypedDict):
"""The current write cursor for a video output stream during aggregation.
**Attributes**:
- **chunk** (`int`) -- The chunk index currently being written to.
- **file** (`int`) -- The file index, within `chunk`, currently being written to.
- **latest_duration** (`float`) -- The duration, in seconds, appended to the current destination
file so far.
- **episode_duration** (`float`) -- The duration, in seconds, of the episode currently being
concatenated.
- **src_to_offset** (`dict[ChunkFile, float]`, *optional*) -- Maps each source `(chunk, file)` to
the time offset, in seconds, at which it was appended into its destination file.
- **src_to_dst** (`dict[ChunkFile, ChunkFile]`, *optional*) -- Maps each source `(chunk, file)` to
the destination `(chunk, file)` its video was concatenated into.
- **dst_file_durations** (`dict[ChunkFile, float]`, *optional*) -- The final total duration, in
seconds, of each completed destination file.
"""
chunk: int
file: int
latest_duration: float
@@ -106,7 +80,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
"""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:
all_metadata (`list`): List of `LeRobotDatasetMetadata` objects to merge.
all_metadata: List of LeRobotDatasetMetadata objects to merge.
Returns:
dict: A dictionary of merged video feature info.
@@ -152,7 +126,7 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[i
Video encoder info is not considered for validation but is merged during aggregation in ``merge_video_feature_info_for_aggregate``.
Args:
all_metadata (`list`): List of `LeRobotDatasetMetadata` objects to validate.
all_metadata: List of LeRobotDatasetMetadata objects to validate.
Returns:
tuple: A tuple containing (fps, robot_type, features) from the first metadata.
@@ -161,6 +135,7 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[i
ValueError: If any metadata has different fps, robot_type, or features
than the first metadata in the list.
"""
fps = all_metadata[0].fps
robot_type = all_metadata[0].robot_type
features = all_metadata[0].features
@@ -189,13 +164,14 @@ def update_data_df(
previously aggregated data in the destination dataset.
Args:
df (`DataFrame`): DataFrame containing the data to be updated.
src_meta (`LeRobotDatasetMetadata`): Source dataset metadata.
dst_meta (`LeRobotDatasetMetadata`): Destination dataset metadata.
df: DataFrame containing the data to be updated.
src_meta: Source dataset metadata.
dst_meta: Destination dataset metadata.
Returns:
pd.DataFrame: Updated DataFrame with adjusted indices.
"""
df["episode_index"] = df["episode_index"] + dst_meta.info.total_episodes
df["index"] = df["index"] + dst_meta.info.total_frames
@@ -221,15 +197,16 @@ def update_meta_data(
to correctly map source file indices to their destination locations.
Args:
df (`DataFrame`): DataFrame containing the metadata to be updated.
dst_meta (`LeRobotDatasetMetadata`): Destination dataset metadata.
meta_idx (`IndexState`): Dictionary containing current metadata chunk and file indices.
data_idx (`IndexState`): Dictionary containing current data chunk and file indices.
videos_idx (`VideoIndexState`): Dictionary containing current video indices and timestamps.
df: DataFrame containing the metadata to be updated.
dst_meta: Destination dataset metadata.
meta_idx: Dictionary containing current metadata chunk and file indices.
data_idx: Dictionary containing current data chunk and file indices.
videos_idx: Dictionary containing current video indices and timestamps.
Returns:
pd.DataFrame: Updated DataFrame with adjusted indices and timestamps.
"""
df["meta/episodes/chunk_index"] = df["meta/episodes/chunk_index"] + meta_idx["chunk"]
df["meta/episodes/file_index"] = df["meta/episodes/file_index"] + meta_idx["file"]
@@ -390,21 +367,15 @@ def aggregate_datasets(
4. Finalizing the aggregated dataset with proper statistics
Args:
repo_ids (`list`): List of repository IDs for the datasets to aggregate.
aggr_repo_id (`str`): Repository ID for the aggregated output dataset.
roots (`list[pathlib.Path] | None`, *optional*): List of root paths for the source
datasets.
aggr_root (`pathlib.Path | None`, *optional*): Root path for the aggregated dataset.
data_files_size_in_mb (`int | None`, *optional*): Maximum size for data files in MB. Falls
back to `DEFAULT_DATA_FILE_SIZE_IN_MB` when not set.
video_files_size_in_mb (`int | None`, *optional*): Maximum size for video files in MB. Falls
back to `DEFAULT_VIDEO_FILE_SIZE_IN_MB` when not set.
chunk_size (`int | None`, *optional*): Maximum number of files per chunk. Falls back to
`DEFAULT_CHUNK_SIZE` when not set.
concatenate_videos (`bool`, *optional*, defaults to `True`): When `False`, keep one mp4 per
source file instead of packing into shards.
concatenate_data (`bool`, *optional*, defaults to `True`): When `False`, keep one parquet per
source file instead of packing into shards.
repo_ids: List of repository IDs for the datasets to aggregate.
aggr_repo_id: Repository ID for the aggregated output dataset.
roots: Optional list of root paths for the source datasets.
aggr_root: Optional root path for the aggregated dataset.
data_files_size_in_mb: Maximum size for data files in MB (defaults to DEFAULT_DATA_FILE_SIZE_IN_MB)
video_files_size_in_mb: Maximum size for video files in MB (defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB)
chunk_size: Maximum number of files per chunk (defaults to DEFAULT_CHUNK_SIZE)
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.
"""
logger.info("Start aggregate_datasets")
@@ -487,14 +458,12 @@ def aggregate_videos(
Creates new video files when size limits are exceeded.
Args:
src_meta (`LeRobotDatasetMetadata`): Source dataset metadata.
dst_meta (`LeRobotDatasetMetadata`): Destination dataset metadata.
videos_idx (`VideoIndexState`): Dictionary tracking video chunk and file indices.
video_files_size_in_mb (`float`): Maximum size for video files in MB.
chunk_size (`int`): Maximum number of files per chunk.
concatenate_videos (`bool`, *optional*, defaults to `True`): When `False`, keep one mp4 per
source file instead of packing into shards.
src_meta: Source dataset metadata.
dst_meta: Destination dataset metadata.
videos_idx: Dictionary tracking video chunk and file indices.
video_files_size_in_mb: Maximum size for video files in MB (defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB)
chunk_size: Maximum number of files per chunk (defaults to DEFAULT_CHUNK_SIZE)
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
Returns:
dict: Updated videos_idx with current chunk and file indices.
"""
@@ -612,13 +581,12 @@ def aggregate_data(
have multiple data files (e.g., from a previous merge operation).
Args:
src_meta (`LeRobotDatasetMetadata`): Source dataset metadata.
dst_meta (`LeRobotDatasetMetadata`): Destination dataset metadata.
data_idx (`IndexState`): Dictionary tracking data chunk and file indices.
data_files_size_in_mb (`float`): Maximum size for data files in MB.
chunk_size (`int`): Maximum number of files per chunk.
concatenate_data (`bool`, *optional*, defaults to `True`): When `False`, keep one parquet per
source file instead of packing into shards.
src_meta: Source dataset metadata.
dst_meta: Destination dataset metadata.
data_idx: Dictionary tracking data chunk and file indices.
data_files_size_in_mb: Maximum size for data files in MB.
chunk_size: Maximum number of files per chunk.
concatenate_data: When False, keep one parquet per source file instead of packing into shards.
Returns:
dict: Updated data_idx with current chunk and file indices.
@@ -692,11 +660,11 @@ def aggregate_metadata(
and writes them to the destination with proper file rotation.
Args:
src_meta (`LeRobotDatasetMetadata`): Source dataset metadata.
dst_meta (`LeRobotDatasetMetadata`): Destination dataset metadata.
meta_idx (`IndexState`): Dictionary tracking metadata chunk and file indices.
data_idx (`IndexState`): Dictionary tracking data chunk and file indices.
videos_idx (`VideoIndexState`): Dictionary tracking video indices and timestamps.
src_meta: Source dataset metadata.
dst_meta: Destination dataset metadata.
meta_idx: Dictionary tracking metadata chunk and file indices.
data_idx: Dictionary tracking data chunk and file indices.
videos_idx: Dictionary tracking video indices and timestamps.
Returns:
dict: Updated meta_idx with current chunk and file indices.
@@ -759,22 +727,18 @@ def append_or_create_parquet_file(
from becoming too large. Handles both regular parquet files and those containing images.
Args:
df (`DataFrame`): DataFrame to write to the parquet file.
src_path (`Path`): Path to the source file, used for size estimation.
idx (`IndexState`): Dictionary containing current `chunk` and `file` indices.
max_mb (`float`): Maximum allowed file size in MB before rotation.
chunk_size (`int`): Maximum number of files per chunk before incrementing the chunk index.
default_path (`str`): Format string for generating file paths.
contains_images (`bool`, *optional*, defaults to `False`): Whether the data contains images
requiring special handling.
aggr_root (`pathlib.Path | None`, *optional*): Root path for the aggregated dataset.
hf_features (`datasets.features.features.Features | None`, *optional*): HuggingFace Features
schema used for proper image typing.
concatenate (`bool`, *optional*, defaults to `True`): When `False`, always rotate to a new
file instead of appending to the current one.
one_row_group_per_episode (`bool`, *optional*, defaults to `False`): Whether to emit one
parquet row group per episode. Set to `True` for data parquet files; left `False` for the
episodes-metadata parquet, which already has one row per episode.
df: DataFrame to write to the parquet file.
src_path: Path to the source file (used for size estimation).
idx: Dictionary containing current 'chunk' and 'file' indices.
max_mb: Maximum allowed file size in MB before rotation.
chunk_size: Maximum number of files per chunk before incrementing chunk index.
default_path: Format string for generating file paths.
contains_images: Whether the data contains images requiring special handling.
aggr_root: Root path for the aggregated dataset.
hf_features: Optional HuggingFace Features schema for proper image typing.
concatenate: When False, always rotate to a new file instead of appending to the current one.
one_row_group_per_episode: True for DATA parquet (emit one row group per episode); False for
the episodes-metadata parquet (already one row per episode).
Returns:
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
@@ -838,8 +802,8 @@ def finalize_aggregation(
aggregated statistics from all source datasets.
Args:
aggr_meta (`LeRobotDatasetMetadata`): Aggregated dataset metadata.
all_metadata (`list`): List of all source dataset metadata objects.
aggr_meta: Aggregated dataset metadata.
all_metadata: List of all source dataset metadata objects.
"""
logger.info("write tasks")
write_tasks(aggr_meta.tasks, aggr_meta.root)
+28 -62
View File
@@ -28,22 +28,16 @@ DEFAULT_QUANTILES = [0.01, 0.10, 0.50, 0.90, 0.99]
class RunningQuantileStats:
"""Maintains running statistics for batches of vectors.
"""
Maintains running statistics for batches of vectors, including mean,
standard deviation, min, max, and approximate quantiles.
Includes mean, standard deviation, min, max, and approximate quantiles. Statistics are computed per
feature dimension and updated incrementally
Statistics are computed per feature dimension and updated incrementally
as new batches are observed. Quantiles are estimated using histograms,
which adapt dynamically if the observed data range expands.
"""
def __init__(self, quantile_list: list[float] | None = None, num_quantile_bins: int = 5000):
"""Initialize empty running statistics.
Args:
quantile_list: Quantiles to track (e.g. `0.01` for the 1st percentile). Defaults to
`DEFAULT_QUANTILES` (1st, 10th, 50th, 90th, 99th percentiles).
num_quantile_bins: Number of histogram bins used to estimate quantiles.
"""
self._count = 0
self._mean = None
self._mean_of_squares = None
@@ -210,9 +204,8 @@ def estimate_num_samples(
dataset_len: int, min_num_samples: int = 100, max_num_samples: int = 10_000, power: float = 0.75
) -> int:
"""Heuristic to estimate the number of samples based on dataset size.
The power controls the sample growth relative to dataset size. Lower the power for less number of
samples.
The power controls the sample growth relative to dataset size.
Lower the power for less number of samples.
For default arguments, we have:
- from 1 to ~500, num_samples=100
@@ -228,24 +221,11 @@ def estimate_num_samples(
def sample_indices(data_len: int) -> list[int]:
"""Return evenly-spaced indices into a sequence of length `data_len`, sized by `estimate_num_samples`."""
num_samples = estimate_num_samples(data_len)
return np.round(np.linspace(0, data_len - 1, num_samples)).astype(int).tolist()
def auto_downsample_height_width(img: np.ndarray, target_size: int = 150, max_size_threshold: int = 300):
"""Downsample a `(C, H, W)` image by integer striding if either dimension exceeds `max_size_threshold`.
Args:
img (`np.ndarray`): Input image in `(C, H, W)` layout to potentially downsample.
target_size (`int`, *optional*, defaults to 150): Approximate size, in pixels, that the
largest side should be reduced to.
max_size_threshold (`int`, *optional*, defaults to 300): Size, in pixels, above which the
largest side of `img` triggers downsampling.
Returns:
`img` unchanged, or strided down so its largest side is roughly `target_size`.
"""
_, height, width = img.shape
if max(width, height) < max_size_threshold:
@@ -257,16 +237,6 @@ def auto_downsample_height_width(img: np.ndarray, target_size: int = 150, max_si
def sample_images(image_paths: list[str]) -> np.ndarray:
"""Load and downsample a sampled subset of `image_paths` into a single `uint8` array.
Args:
image_paths (`list[str]`): Paths of all images for the episode/feature, from which a
subset is sampled (see `sample_indices`).
Returns:
A `(N, C, H, W)` `uint8` array of the sampled, downsampled images (see
`auto_downsample_height_width`), where `N` is chosen by `sample_indices`.
"""
sampled_indices = sample_indices(len(image_paths))
images = None
@@ -439,7 +409,6 @@ def _compute_basic_stats(
Args:
array: Reshaped array ready for statistics computation
sample_count: Number of samples represented in the data
quantile_list: Quantiles to fill with the mean value. Defaults to `DEFAULT_QUANTILES`.
Returns:
Dictionary with basic statistics and quantiles set to mean values
@@ -478,14 +447,13 @@ def get_feature_stats(
- Global: axis=None computes statistics over entire array
Args:
array (`np.ndarray`): Input data array with a shape appropriate for the specified `axis`.
axis (`int | tuple[int, ...] | None`): Axis or axes along which to compute statistics:
`(0, 2, 3)` for image data (batch, channels, height, width), `0` or `(0,)` for
vector/tabular data (samples, features), `(1,)` to compute across features, or
`None` for global statistics over the entire array.
keepdims (`bool`): If `True`, reduced axes are kept as dimensions of size 1.
quantile_list (`list[float] | None`, *optional*): Quantiles to compute (e.g. `0.01` for
the 1st percentile). Defaults to `DEFAULT_QUANTILES` when not provided.
array: Input data array with shape appropriate for the specified axis
axis: Axis or axes along which to compute statistics
- (0, 2, 3): For image data (batch, channels, height, width)
- 0 or (0,): For vector/tabular data (samples, features)
- (1,): For computing across features
- None: For global statistics over entire array
keepdims: If True, reduced axes are kept as dimensions with size 1
Returns:
Dictionary containing:
@@ -528,13 +496,10 @@ def compute_episode_stats(
- Strings: Skipped (no statistics computed)
Args:
episode_data (`dict[str, list[str] | np.ndarray]`): Mapping from feature name to its data
for the episode: a list of file paths for `image`/`video` features, or a numpy array
for numerical features.
features (`dict`): Dataset feature metadata, keyed by feature name, describing each
feature's `dtype` and shape.
quantile_list (`list[float] | None`, *optional*): Quantiles to compute (e.g. `0.01` for
the 1st percentile). Defaults to `DEFAULT_QUANTILES` when not provided.
episode_data: Dictionary mapping feature names to data
- For images/videos: list of file paths
- For numerical data: numpy arrays
features: Dictionary describing each feature's dtype and shape
Returns:
Dictionary mapping feature names to their statistics dictionaries.
@@ -665,6 +630,7 @@ def aggregate_stats(stats_list: list[dict[str, dict]]) -> dict[str, dict[str, np
- new_mean = (mean of all data, weighted by counts)
- new_std = (std of all data)
"""
_assert_type_and_shape(stats_list)
data_keys = {key for stats in stats_list for key in stats}
@@ -724,16 +690,16 @@ def compute_relative_action_stats(
statistics suitable for normalization.
Args:
hf_dataset (`datasets.Dataset`): The underlying HuggingFace dataset, must expose
`"action"`, `"observation.state"`, and `"episode_index"` columns.
features (`dict`): Dataset feature metadata; must contain `"action"` with a `"shape"`
entry and optionally `"names"`.
chunk_size (`int`): Number of consecutive frames per action chunk.
exclude_joints (`list[str] | None`, *optional*): Joint names whose dimensions should
remain absolute instead of being converted to relative actions.
num_workers (`int`, *optional*, defaults to 0): Number of parallel threads used for
computation. Values `<= 1` run single-threaded; NumPy releases the GIL so threads
give real parallelism here.
hf_dataset: The underlying HuggingFace dataset with "action",
"observation.state", and "episode_index" columns.
features: Dataset feature metadata (must contain "action" with "shape"
and optionally "names").
chunk_size: Number of consecutive frames per action chunk.
exclude_joints: Joint names whose dimensions should remain absolute
(not converted to relative actions).
num_workers: Number of parallel threads for computation. Values ≤1
mean single-threaded. Numpy releases the GIL so threads give
real parallelism here.
Returns:
Statistics dict with keys "mean", "std", "min", "max", "q01", …, "q99".
+3 -4
View File
@@ -529,9 +529,9 @@ class LeRobotDatasetMetadata:
return self.info.video_files_size_in_mb
def get_task_index(self, task: str) -> int | None:
"""Given a task in natural language, returns its task_index if the task already exists in the dataset.
Otherwise return None.
"""
Given a task in natural language, returns its task_index if the task already exists in the dataset,
otherwise return None.
"""
if task in self.tasks.index:
return int(self.tasks.loc[task].task_index)
@@ -774,7 +774,6 @@ class LeRobotDatasetMetadata:
}
def __repr__(self):
"""A short summary: repo ID, total episode/frame counts, and feature keys."""
feature_keys = list(self.features)
return (
f"{self.__class__.__name__}({{\n"
+1 -3
View File
@@ -293,9 +293,7 @@ class DatasetReader:
return result
def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict[str, torch.Tensor]:
"""Decode the requested per-camera frame timestamps from `ep_idx`'s videos.
Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
"""Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
in the main process (e.g. by using a second Dataloader with num_workers=0). It will result in a
Segmentation Fault.
"""
+77 -86
View File
@@ -118,11 +118,10 @@ def delete_episodes(
consistent with its own metadata.
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset.
episode_indices (`list`): List of episode indices to delete.
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the edited dataset
will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/repo_id`.
repo_id (`str | None`, *optional*): Identifier for the edited dataset.
dataset: The source LeRobotDataset.
episode_indices: List of episode indices to delete.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
"""
if not episode_indices:
raise ValueError("No episodes to delete")
@@ -186,11 +185,10 @@ def split_dataset(
output split stays consistent with its own metadata.
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset to split.
splits (`dict`): Either a dict mapping split names to episode indices, or a dict mapping
split names to fractions (must sum to <= 1.0).
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the split
datasets will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/repo_id`.
dataset: The source LeRobotDataset to split.
splits: Either a dict mapping split names to episode indices, or a dict mapping
split names to fractions (must sum to <= 1.0).
output_dir: Root directory where the split datasets will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id.
Examples:
Split by specific episodes
@@ -282,14 +280,11 @@ def merge_datasets(
This is a wrapper around the aggregate_datasets functionality with a cleaner API.
Args:
datasets (`list`): List of LeRobotDatasets to merge.
output_repo_id (`str`): Identifier for the merged dataset.
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the merged dataset
will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/output_repo_id`.
concatenate_videos (`bool`, *optional*, defaults to `True`): When `False`, keep one mp4 per
source file instead of packing them into shards.
concatenate_data (`bool`, *optional*, defaults to `True`): When `False`, keep one parquet file
per source file instead of packing them into shards.
datasets: List of LeRobotDatasets to merge.
output_repo_id: Merged dataset identifier.
output_dir: Root directory where the merged dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/output_repo_id.
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.
"""
if not datasets:
raise ValueError("No datasets to merge")
@@ -332,14 +327,11 @@ def modify_features(
regardless of how many features are being added or removed.
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset.
add_features (`dict[str, tuple[numpy.ndarray | torch.Tensor | collections.abc.Callable, dict]] | None`, *optional*):
Dict mapping feature names to `(feature_values, feature_info)` tuples.
remove_features (`str | list[str] | None`, *optional*): Feature name(s) to remove. Can be a
single string or a list.
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the edited dataset
will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/repo_id`.
repo_id (`str | None`, *optional*): Identifier for the edited dataset.
dataset: The source LeRobotDataset.
add_features: Optional dict mapping feature names to (feature_values, feature_info) tuples.
remove_features: Optional feature name(s) to remove. Can be a single string or list.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
Returns:
New dataset with features modified.
@@ -438,11 +430,10 @@ def add_features(
copies the dataset once regardless of how many features are being added.
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset.
features (`dict`): Dictionary mapping feature names to `(feature_values, feature_info)` tuples.
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the edited dataset
will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/repo_id`.
repo_id (`str | None`, *optional*): Identifier for the edited dataset.
dataset: The source LeRobotDataset.
features: Dictionary mapping feature names to (feature_values, feature_info) tuples.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
Returns:
New dataset with all features added.
@@ -476,12 +467,10 @@ def remove_feature(
"""Remove features from a LeRobotDataset.
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset.
feature_names (`str | list[str]`): Name(s) of features to remove. Can be a single string or
a list.
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the edited dataset
will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/repo_id`.
repo_id (`str | None`, *optional*): Identifier for the edited dataset.
dataset: The source LeRobotDataset.
feature_names: Name(s) of features to remove. Can be a single string or list.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
Returns:
New dataset with features removed.
@@ -960,7 +949,7 @@ def _copy_and_reindex_episodes_metadata(
def _write_parquet(df: pd.DataFrame, path: Path, meta: LeRobotDatasetMetadata) -> None:
"""Write DataFrame to parquet.
"""Write DataFrame to parquet
This ensures images are properly embedded and the file can be loaded correctly by HF datasets.
"""
@@ -1468,14 +1457,13 @@ def modify_tasks(
- meta/info.json (total_tasks)
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset to modify.
new_task (`str | None`, *optional*): Default task applied to any episode not covered by
`episode_tasks` or a matching `task_replacements` entry.
episode_tasks (`dict[int, str] | None`, *optional*): Dict mapping episode indices to task
strings. Takes precedence over both `task_replacements` and `new_task`.
task_replacements (`dict[str, str] | None`, *optional*): Dict mapping existing task strings to
new ones. Applied to episodes whose current task matches a key. Every key must be an
existing task.
dataset: The source LeRobotDataset to modify.
new_task: Default task applied to any episode not covered by `episode_tasks` or a
matching `task_replacements` entry.
episode_tasks: Optional dict mapping episode indices to task strings. Takes precedence
over both `task_replacements` and `new_task`.
task_replacements: Optional dict mapping existing task strings to new ones. Applied to
episodes whose current task matches a key. Every key must be an existing task.
At least one of `new_task`, `episode_tasks`, or `task_replacements` must be provided.
@@ -1606,19 +1594,19 @@ def recompute_stats(
"""Recompute stats.json from scratch by iterating all episodes.
Args:
dataset (`LeRobotDataset`): The LeRobotDataset to recompute stats for.
skip_image_video (`bool`, *optional*, defaults to `True`): If `True`, only recompute stats for
numeric features (action, state, etc.) and keep existing image/video stats unchanged.
relative_action (`bool`, *optional*, defaults to `False`): If `True`, compute action stats in
relative space by iterating all valid action chunks and subtracting the current state.
This matches the normalization distribution the model sees during training with
`use_relative_actions=True`.
relative_exclude_joints (`list[str] | None`, *optional*): Joint names to exclude from relative
conversion when `relative_action=True`. These dims keep absolute stats.
chunk_size (`int`, *optional*, defaults to 50): Action chunk size used for relative stats
computation. Should match `policy.chunk_size`. Only used when `relative_action=True`.
num_workers (`int`, *optional*, defaults to 0): Number of parallel threads for relative action
stats computation. Values <=1 mean single-threaded. Only used when `relative_action=True`.
dataset: The LeRobotDataset to recompute stats for.
skip_image_video: If True (default), only recompute stats for numeric features
(action, state, etc.) and keep existing image/video stats unchanged.
relative_action: If True, compute action stats in relative space by
iterating all valid action chunks and subtracting the current state.
This matches the normalization distribution the model sees during
training with ``use_relative_actions=True``.
relative_exclude_joints: Joint names to exclude from relative conversion when
relative_action=True. These dims keep absolute stats.
chunk_size: Action chunk size used for relative stats computation. Should match
``policy.chunk_size``. Only used when ``relative_action=True``.
num_workers: Number of parallel threads for relative action stats computation.
Values 1 mean single-threaded. Only used when ``relative_action=True``.
Returns:
The same dataset with updated stats.
@@ -1721,22 +1709,24 @@ def convert_image_to_video_dataset(
LeRobot dataset structure with videos stored in chunked MP4 files.
Args:
dataset (`LeRobotDataset`): The source LeRobot dataset with images.
output_dir (`pathlib.Path | None`, *optional*): Root directory where the converted dataset will
be stored. When `None`, defaults to `$HF_LEROBOT_HOME/repo_id`.
repo_id (`str | None`, *optional*): Identifier for the converted dataset.
rgb_encoder (`lerobot.configs.video.RGBEncoderConfig | None`, *optional*): Video encoder settings
applied to RGB cameras. When `None`, `rgb_encoder_defaults` is used.
depth_encoder (`lerobot.configs.video.DepthEncoderConfig | None`, *optional*): Video encoder
settings applied to depth-map cameras, including the quantization parameters persisted to
the dataset metadata. When `None`, `depth_encoder_defaults` is used.
episode_indices (`list[int] | None`, *optional*): Episode indices to convert. When `None`, all
episodes are converted.
num_workers (`int`, *optional*, defaults to 4): Number of threads for parallel processing.
max_episodes_per_batch (`int | None`, *optional*): Maximum episodes per video batch, to bound
memory use. `None` means no limit.
max_frames_per_batch (`int | None`, *optional*): Maximum frames per video batch, to bound memory
use. `None` means no limit.
dataset: The source LeRobot dataset with images.
output_dir: Root directory where the converted dataset will be stored. When
``None``, defaults to ``$HF_LEROBOT_HOME/repo_id``. Equivalent to
``new_root`` in ``EditDatasetConfig``.
repo_id: Converted dataset identifier. Equivalent to ``new_repo_id`` in
``EditDatasetConfig``.
rgb_encoder: Video encoder settings applied to RGB cameras. When ``None``,
:func:`~lerobot.configs.video.rgb_encoder_defaults` is used.
depth_encoder: Video encoder settings applied to depth-map cameras, including
the quantization parameters persisted to the dataset metadata. When
``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used.
episode_indices: Episode indices to convert. When ``None``, all episodes are
converted.
num_workers: Number of threads for parallel processing.
max_episodes_per_batch: Maximum episodes per video batch, to bound memory use.
``None`` means no limit.
max_frames_per_batch: Maximum frames per video batch, to bound memory use.
``None`` means no limit.
Returns:
A new :class:`LeRobotDataset` with images encoded as videos.
@@ -1976,17 +1966,18 @@ def reencode_dataset(
Videos are re-encoded in-place and the video information in ``info.json`` is refreshed.
Args:
dataset (`LeRobotDataset`): An existing :class:`LeRobotDataset` whose videos will be re-encoded.
rgb_encoder (`lerobot.configs.video.RGBEncoderConfig | None`, *optional*): Target encoder
configuration applied to every RGB video file. If `None`, re-encoding is skipped for RGB
videos.
depth_encoder (`lerobot.configs.video.DepthEncoderConfig | None`, *optional*): Target encoder
configuration applied to every depth video file. If `None`, re-encoding is skipped for depth
videos. Quantization parameters will not override the ones in the current dataset.
encoder_threads (`int | None`, *optional*): Per-encoder thread count forwarded to
`reencode_video`. `None` lets the codec decide.
num_workers (`int | None`, *optional*): Number of parallel processes. `None` or `0` means
sequential (no multiprocessing); `1+` spawns a `ProcessPoolExecutor`.
dataset: An existing :class:`LeRobotDataset` whose videos will be
re-encoded.
rgb_encoder: Target encoder configuration applied to every RGB video
file. If ``None``, re-encoding is skipped for RGB videos.
depth_encoder: Target encoder configuration applied to every depth video
file. If ``None``, re-encoding is skipped for depth videos.
Quantization parameters will not override the ones in the current dataset.
encoder_threads: Per-encoder thread count forwarded to
:func:`reencode_video`. ``None`` lets the codec decide.
num_workers: Number of parallel processes. ``None`` or ``0`` means
sequential (no multiprocessing); ``1+`` spawns a
:class:`~concurrent.futures.ProcessPoolExecutor`.
Returns:
The same :class:`LeRobotDataset` instance with its metadata updated
+2 -1
View File
@@ -200,7 +200,8 @@ class DatasetWriter:
self.image_writer.save_image(image=image, fpath=fpath, compress_level=compress_level)
def add_frame(self, frame: dict) -> None:
"""Add a single frame to the current episode buffer.
"""
Add a single frame to the current episode buffer.
Apart from images written to a temporary directory, nothing is written to disk
until ``save_episode()`` is called.
+19 -41
View File
@@ -13,7 +13,9 @@
# 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.
"""Depth encoding/decoding helpers for :class:`DepthEncoderConfig`."""
"""
Depth encoding/decoding helpers for :class:`DepthEncoderConfig`.
"""
import math
from typing import Literal
@@ -90,24 +92,13 @@ def quantize_depth(
``depth_min``, ``depth_max``, and ``shift`` are always in **metres**.
Args:
depth (`numpy.ndarray[tuple[int, ...], numpy.dtype[numpy.uint16]] | numpy.ndarray[tuple[int, ...], numpy.dtype[numpy.float32]] | torch.Tensor`): Depth
map to quantize. A `torch.Tensor` is moved to CPU before conversion.
depth_min (`float`, *optional*, defaults to 0.01): Depth, in metres, mapped to quantum
`0`.
depth_max (`float`, *optional*, defaults to 10.0): Depth, in metres, mapped to quantum
`DEPTH_QMAX`.
shift (`float`, *optional*, defaults to 3.5): Depth shift, in metres, used in log mode.
Must satisfy `depth_min + shift > 0`.
use_log (`bool`, *optional*, defaults to `True`): If `True`, quantize in log space, which
allocates more quanta to near-range depth.
pix_fmt (`str`, *optional*, defaults to `"gray12le"`): Pixel format used to build the
`av.VideoFrame` when `video_backend="pyav"`.
video_backend (`str | None`, *optional*, defaults to `"pyav"`): Video backend used for
encoding. When `"pyav"`, returns an `av.VideoFrame`; otherwise returns the raw
`uint16` array.
input_unit (`Literal`, *optional*, defaults to `"auto"`): Input unit policy: `"auto"`
infers the unit from `depth`'s dtype, while `"mm"` or `"m"` force millimetres or
metres respectively.
depth: Depth map; ``torch.Tensor`` is moved to CPU for conversion.
depth_min: Depth (metres) at quantum ``0``.
depth_max: Depth (metres) at quantum :data:`DEPTH_QMAX`.
shift: Depth shift (metres); used in log mode. Must satisfy ``depth_min + shift > 0``.
use_log: If ``True`` (default), quantize in log space.
video_backend: Video backend to use for encoding. Defaults to "pyav".
input_unit: Input unit policy (``"auto"``, ``"mm"``, ``"m"``).
Returns:
``numpy.ndarray``, ``dtype=uint16``, same shape as ``depth``, values in
@@ -183,28 +174,15 @@ def dequantize_depth(
Output layout is determined by ``output_channel_last``.
Args:
quantized (`numpy.ndarray[tuple[int, ...], numpy.dtype[numpy.uint16]] | av.video.frame.VideoFrame | torch.Tensor`): 12-bit
codes in `[0, DEPTH_QMAX]`, as a numpy array, `av.VideoFrame`, or `torch.Tensor`
(any integer or float dtype).
depth_min (`float`, *optional*, defaults to 0.01): Depth, in metres, mapped to quantum
`0`. Must match the value passed to `quantize_depth`.
depth_max (`float`, *optional*, defaults to 10.0): Depth, in metres, mapped to quantum
`DEPTH_QMAX`. Must match the value passed to `quantize_depth`.
shift (`float`, *optional*, defaults to 3.5): Depth shift, in metres, used in log mode.
Must match the value passed to `quantize_depth`.
use_log (`bool`, *optional*, defaults to `True`): If `True`, invert the log-space mapping
used by `quantize_depth`. Must match the encoding call.
pix_fmt (`str`, *optional*, defaults to `"gray12le"`): Pixel format used to extract the
plane data when `quantized` is an `av.VideoFrame`.
output_unit (`Literal`, *optional*, defaults to `"mm"`): `"mm"` returns `uint16`
millimetres, clipped to `[0, 65535]`, when returning a numpy array, or `float32`
millimetres when `output_tensor=True`. `"m"` returns `float32` metres in
`[depth_min, depth_max]`.
output_tensor (`bool`, *optional*, defaults to `True`): If `True`, return a
`torch.Tensor` instead of a numpy array.
output_channel_last (`bool`, *optional*, defaults to `False`): If `True`, add the
restored singleton channel dimension as the last axis instead of the third-to-last
axis.
quantized: 12-bit codes in ``[0, DEPTH_QMAX]``. ``np.ndarray``,
``av.VideoFrame``, or ``torch.Tensor`` (any integer or float dtype).
depth_min, depth_max, shift, use_log: Same as :func:`quantize_depth` (metres).
pix_fmt: Pixel format used to extract the plane from an ``av.VideoFrame``.
output_unit: ``"mm"`` returns ``uint16`` millimetres (rint, clip
``[0, 65535]``) when returning a numpy array, or ``float32`` mm when
``output_tensor=True``. ``"m"`` returns ``float32`` metres in
``[depth_min, depth_max]``.
output_tensor: If True, return a ``torch.Tensor`` instead of a numpy array.
Returns:
Depth map in the requested unit and dtype.
+8 -21
View File
@@ -101,10 +101,10 @@ def create_empty_dataset_info(
fps (int): The frames per second of the data.
features (dict): The LeRobot features dictionary for the dataset.
use_videos (bool): Whether the dataset will store videos.
robot_type (str | None, *optional*): The type of robot used, if any.
chunks_size (int | None, *optional*): Max files per chunk directory. Defaults to ``DEFAULT_CHUNK_SIZE``.
data_files_size_in_mb (int | None, *optional*): Max parquet file size in MB. Defaults to ``DEFAULT_DATA_FILE_SIZE_IN_MB``.
video_files_size_in_mb (int | None, *optional*): Max video file size in MB. Defaults to ``DEFAULT_VIDEO_FILE_SIZE_IN_MB``.
robot_type (str | None): The type of robot used, if any.
chunks_size (int | None): Max files per chunk directory. Defaults to ``DEFAULT_CHUNK_SIZE``.
data_files_size_in_mb (int | None): Max parquet file size in MB. Defaults to ``DEFAULT_DATA_FILE_SIZE_IN_MB``.
video_files_size_in_mb (int | None): Max video file size in MB. Defaults to ``DEFAULT_VIDEO_FILE_SIZE_IN_MB``.
Returns:
DatasetInfo: A typed dataset information object with initial metadata.
@@ -170,7 +170,7 @@ def check_delta_timestamps(
deltas in seconds.
fps (int): The frames per second of the dataset.
tolerance_s (float): The allowed tolerance in seconds.
raise_value_error (bool, *optional*, defaults to `True`): If True, raises an error on failure.
raise_value_error (bool): If True, raises an error on failure.
Returns:
bool: True if all deltas are valid, False otherwise.
@@ -219,18 +219,6 @@ def get_delta_indices(delta_timestamps: dict[str, list[float]], fps: int) -> dic
def validate_frame(frame: dict, features: dict) -> None:
"""Check that `frame` has a `"task"` key and matches `features` (minus auto-populated defaults).
Args:
frame (`dict`): The frame to validate, mapping feature names to their values, as passed by the
caller to `add_frame`.
features (`dict`): The dataset's feature specification, mapping feature names to their dtype and
shape metadata.
Raises:
ValueError: If `frame` is missing `"task"`, or has missing/extra features, or a feature's dtype
or shape doesn't match its definition in `features`.
"""
# DEFAULT_FEATURES (timestamp, frame_index, episode_index, index, task_index) are
# auto-populated by the recording pipeline (add_frame / save_episode) and must not
# be supplied by the caller. Excluding them here means any frame dict that contains
@@ -287,7 +275,7 @@ def validate_feature_dtype_and_shape(
Args:
name (str): The name of the feature.
feature (dict): The feature specification from the LeRobot features dictionary.
value (`numpy.ndarray | PIL.Image.Image | str`): The value of the feature to validate.
value: The value of the feature to validate.
Returns:
str: An error message if validation fails, otherwise an empty string.
@@ -349,7 +337,7 @@ def validate_feature_image_or_video(
Args:
name (str): The name of the feature.
expected_shape (list[str]): The expected shape, e.g. (C, H, W) or (H, W, C).
value (`numpy.ndarray | PIL.Image.Image`): The image or video frame data to validate.
value: The image data to validate.
Returns:
str: An error message if validation fails, otherwise an empty string.
@@ -395,8 +383,7 @@ def validate_feature_language(name: str, value) -> str:
Args:
name (str): The name of the feature.
value (`Any`): The value supplied for the language feature. Only checked for being `None`; any
other value is dropped with a warning.
value: The value to validate.
Returns:
str: Always an empty string — language values are non-fatal.
+6 -22
View File
@@ -27,10 +27,7 @@ logger = logging.getLogger(__name__)
def safe_stop_image_writer(func):
"""Decorator: on an exception from `func`, stop the `dataset` kwarg's image writer before re-raising."""
def wrapper(*args, **kwargs):
"""Call `func`; on any exception, stop `kwargs["dataset"].writer.image_writer` before re-raising."""
try:
return func(*args, **kwargs)
except BaseException:
@@ -129,7 +126,8 @@ def save_kwargs_for_path(fpath: Path, compress_level: int) -> dict:
def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1):
"""Saves a NumPy array or PIL Image to a file.
"""
Saves a NumPy array or PIL Image to a file.
This function handles both NumPy arrays and PIL Image objects, converting
the former to a PIL Image before saving. It includes error handling for
@@ -140,7 +138,7 @@ def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level
Args:
image (np.ndarray | PIL.Image.Image): The image data to save.
fpath (Path): The destination file path for the image.
compress_level (int, optional, *optional*, defaults to 1): The compression level for the saved
compress_level (int, optional): The compression level for the saved
image, as used by PIL.Image.save(). Defaults to 1.
Refer to: https://github.com/huggingface/lerobot/pull/2135
for more details on the default value rationale.
@@ -165,7 +163,6 @@ def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level
def worker_thread_loop(queue: queue.Queue):
"""Pop `(image_array, fpath, compress_level)` items from `queue` and write each until a `None` sentinel."""
while True:
item = queue.get()
if item is None:
@@ -177,7 +174,6 @@ def worker_thread_loop(queue: queue.Queue):
def worker_process(queue: queue.Queue, num_threads: int):
"""Run `num_threads` `worker_thread_loop` threads against `queue` and block until they all exit."""
threads = []
for _ in range(num_threads):
t = threading.Thread(target=worker_thread_loop, args=(queue,))
@@ -189,9 +185,9 @@ def worker_process(queue: queue.Queue, num_threads: int):
class AsyncImageWriter:
"""This class abstracts away the initialisation of processes or/and threads.
It saves images on disk asynchronously, which is critical to control a robot and record data
"""
This class abstract away the initialisation of processes or/and threads to
save images on disk asynchronously, which is critical to control a robot and record data
at a high frame rate.
When `num_processes=0`, it creates a threads pool of size `num_threads`.
@@ -204,15 +200,6 @@ class AsyncImageWriter:
"""
def __init__(self, num_processes: int = 0, num_threads: int = 1):
"""Start the thread or process pool.
Args:
num_processes: Number of worker subprocesses. `0` uses threads only (in this process).
num_threads: Number of writer threads per process (or in this process, if `num_processes=0`).
Raises:
ValueError: If both `num_threads` and `num_processes` are non-positive.
"""
self.num_processes = num_processes
self.num_threads = num_threads
self.queue = None
@@ -243,18 +230,15 @@ class AsyncImageWriter:
def save_image(
self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
):
"""Enqueue `image` to be written to `fpath` asynchronously; returns immediately."""
if isinstance(image, torch.Tensor):
# Convert tensor to numpy array to minimize main process time
image = image.cpu().numpy()
self.queue.put((image, fpath, compress_level))
def wait_until_done(self):
"""Block until every enqueued image has been written to disk."""
self.queue.join()
def stop(self):
"""Signal all worker threads/processes to exit and wait for them to join. No-op if already stopped."""
if self._stopped:
return
+10 -21
View File
@@ -46,7 +46,6 @@ from .utils import (
def get_parquet_file_size_in_mb(parquet_path: str | Path) -> float:
"""Return the uncompressed size, in megabytes, of a parquet file's column data (from its metadata)."""
metadata = pq.read_metadata(parquet_path)
total_uncompressed_size = 0
for row_group in range(metadata.num_row_groups):
@@ -58,24 +57,20 @@ def get_parquet_file_size_in_mb(parquet_path: str | Path) -> float:
def get_hf_dataset_size_in_mb(hf_ds: Dataset) -> int:
"""Return the in-memory (Arrow buffer) size of a Hugging Face `Dataset`, in megabytes."""
return hf_ds.data.nbytes // (1024**2)
def load_nested_dataset(
pq_dir: Path, features: datasets.Features | None = None, episodes: list[int] | None = None
) -> Dataset:
"""Find parquet files in provided directory {pq_dir}/chunk-xxx/file-xxx.parquet.
Convert parquet files to pyarrow memory mapped in a cache folder for efficient RAM usage, then
concatenate all pyarrow references to return HF Dataset format.
"""Find parquet files in provided directory {pq_dir}/chunk-xxx/file-xxx.parquet
Convert parquet files to pyarrow memory mapped in a cache folder for efficient RAM usage
Concatenate all pyarrow references to return HF Dataset format
Args:
pq_dir (`Path`): Directory containing parquet files.
features (`datasets.features.features.Features | None`, *optional*): Features schema used to ensure
consistent loading of complex types like images.
episodes (`list[int] | None`, *optional*): List of episode indices to filter. Uses PyArrow
predicate pushdown for efficiency.
pq_dir: Directory containing parquet files
features: Optional features schema to ensure consistent loading of complex types like images
episodes: Optional list of episode indices to filter. Uses PyArrow predicate pushdown for efficiency.
"""
paths = sorted(pq_dir.glob("*/*.parquet"))
if len(paths) == 0:
@@ -88,7 +83,6 @@ def load_nested_dataset(
def get_parquet_num_frames(parquet_path: str | Path) -> int:
"""Return the number of rows in a parquet file, read from its metadata (no data is loaded)."""
metadata = pq.read_metadata(parquet_path)
return metadata.num_rows
@@ -124,7 +118,6 @@ def embed_images(dataset: datasets.Dataset) -> datasets.Dataset:
def write_info(info: DatasetInfo, local_dir: Path) -> None:
"""Write dataset info metadata to its standard file path (the inverse of `load_info`)."""
write_json(info.to_dict(), local_dir / INFO_PATH)
@@ -183,14 +176,12 @@ def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]] | None:
def write_tasks(tasks: pandas.DataFrame, local_dir: Path) -> None:
"""Write the task-prompt table to its standard parquet file path (the inverse of `load_tasks`)."""
path = local_dir / DEFAULT_TASKS_PATH
path.parent.mkdir(parents=True, exist_ok=True)
tasks.to_parquet(path)
def load_tasks(local_dir: Path) -> pandas.DataFrame:
"""Load the task-prompt table from its standard file path, indexed by task string."""
tasks = pd.read_parquet(local_dir / DEFAULT_TASKS_PATH)
tasks.index.name = "task"
return tasks
@@ -198,13 +189,12 @@ def load_tasks(local_dir: Path) -> pandas.DataFrame:
def write_episodes(episodes: Dataset, local_dir: Path) -> None:
"""Write episode metadata to a parquet file in the LeRobot v3.0 format.
This function writes episode-level metadata to a single parquet file.
Used primarily during dataset conversion (v2.1 → v3.0) and in test fixtures.
Args:
episodes (`Dataset`): Hugging Face `Dataset` containing the episode metadata.
local_dir (`Path`): Root directory where the dataset is stored.
episodes: HuggingFace Dataset containing episode metadata
local_dir: Root directory where the dataset will be stored
"""
episode_size_mb = get_hf_dataset_size_in_mb(episodes)
if episode_size_mb > DEFAULT_DATA_FILE_SIZE_IN_MB:
@@ -220,7 +210,6 @@ def write_episodes(episodes: Dataset, local_dir: Path) -> None:
def load_episodes(local_dir: Path) -> datasets.Dataset:
"""Load episode metadata, excluding per-episode `stats/*` columns (for faster access to the rest)."""
episodes = load_nested_dataset(local_dir / EPISODES_DIR)
# Select episode features/columns containing references to episode data and videos
# (e.g. tasks, dataset_from_index, dataset_to_index, data/chunk_index, data/file_index, etc.)
@@ -236,9 +225,9 @@ def load_image_as_numpy(
Args:
fpath (str | Path): Path to the image file.
dtype (np.dtype, *optional*, defaults to `float32`): The desired data type of the output array. If floating,
dtype (np.dtype): The desired data type of the output array. If floating,
pixels are scaled to [0, 1]. Only used for RGB images.
channel_first (bool, *optional*, defaults to `True`): If True, converts the image to (C, H, W) format.
channel_first (bool): If True, converts the image to (C, H, W) format.
Otherwise, it remains in (H, W, C) format.
Returns:
+4 -23
View File
@@ -44,14 +44,6 @@ logger = logging.getLogger(__name__)
class LeRobotDataset(torch.utils.data.Dataset):
"""A PyTorch `Dataset` over episodic robot data: per-frame state/action tensors, optional videos.
Backed by parquet files (`data/`) for tabular observation/action/reward data, optional video files
(`videos/`) for image observations, and a `meta/` directory holding `info.json` (shapes, keys, fps),
`stats.json` (normalization statistics), and per-episode metadata. See `__init__`'s docstring for the
on-disk layout, and `create()` for building a new (empty) dataset from scratch.
"""
def __init__(
self,
repo_id: str,
@@ -76,7 +68,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
*,
token: str | bool | None = None,
):
"""2 modes are available for instantiating this class, depending on 2 different use cases.
"""
2 modes are available for instantiating this class, depending on 2 different use cases:
1. Your dataset already exists:
- On your local disk in the 'root' folder. This is typically the case when you recorded your
@@ -175,9 +168,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
conversion. This works for both image-backed and video-backed observations and can later be
updated with `set_image_transforms()` or cleared with `clear_image_transforms()`.
Defaults to None.
delta_timestamps (dict[list[float]] | None, optional): Per-feature timestamp offsets (in
seconds, relative to a frame's own timestamp) of additional frames to return alongside it.
Defaults to None.
delta_timestamps (dict[list[float]] | None, optional): _description_. Defaults to None.
tolerance_s (float, optional): Tolerance in seconds used to ensure data timestamps are actually in
sync with the fps value. It is used at the init of the dataset to make sure that each
timestamps is separated to the next by 1/fps +/- tolerance_s. This also applies to frames
@@ -194,11 +185,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
True.
video_backend (str | None, optional): Video backend to use for decoding videos. Defaults to torchcodec when available int the platform; otherwise, defaults to 'pyav'.
You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision.
return_uint8 (bool, optional): For RGB videos, whether to return raw uint8 frames instead of
the default float32 frames normalized to [0, 1]. Defaults to False.
depth_output_unit (str, optional): Physical unit depth maps are dequantized to at load time:
"mm" (millimeters) or "m" (metres). Has no effect on datasets without depth cameras.
Defaults to "mm".
batch_encoding_size (int, optional): Number of episodes to accumulate before batch encoding videos.
Set to 1 for immediate encoding (default), or higher for batched encoding. Defaults to 1.
rgb_encoder (RGBEncoderConfig | None, optional): Video encoder settings for cameras
@@ -409,7 +395,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
@property
def hf_dataset(self) -> datasets.Dataset:
"""The underlying Hugging Face Dataset object."""
"""The underlying Hugging Face Dataset object"""
self.reader = self._ensure_reader()
if self.reader.hf_dataset is None:
self.reader.load_and_activate()
@@ -554,7 +540,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
return self.hf_dataset[idx]
def __repr__(self):
"""A short summary: repo ID, selected episode/sample counts, and feature keys."""
feature_keys = list(self.features)
return (
f"{self.__class__.__name__}({{\n"
@@ -753,10 +738,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
during capture instead of writing images first.
encoder_queue_maxsize: Max buffered frames per camera when using
streaming encoding.
video_files_size_in_mb: Max video file size in MB. Defaults to
``DEFAULT_VIDEO_FILE_SIZE_IN_MB``.
data_files_size_in_mb: Max parquet file size in MB. Defaults to
``DEFAULT_DATA_FILE_SIZE_IN_MB``.
Returns:
A new :class:`LeRobotDataset` in write mode.
+3 -28
View File
@@ -51,20 +51,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
*,
token: str | bool | None = None,
):
"""Construct a `LeRobotDataset` for each `repo_id` and concatenate them.
Args:
repo_ids: The Hub repo IDs (or local dataset names, if `root` is set) to load.
root: Root directory containing the underlying datasets. Defaults to `$HF_LEROBOT_HOME`.
episodes: Optional mapping from `repo_id` to the episode indices to load from it.
image_transforms: Transform applied to visual observations in each underlying dataset.
delta_timestamps: Passed through to each underlying `LeRobotDataset`.
tolerances_s: Optional mapping from `repo_id` to its timestamp tolerance, in seconds. Defaults
to `1e-4` for every dataset.
download_videos: Whether to download video files for each underlying dataset.
video_backend: The video decoding backend to use.
token: Hugging Face Hub authentication token.
"""
super().__init__()
self.repo_ids = repo_ids
self.root = Path(root) if root else HF_LEROBOT_HOME
@@ -154,7 +140,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
@property
def features(self) -> datasets.Features:
"""The union of all underlying datasets' features (minus `disabled_features`)."""
features = {}
for dataset in self._datasets:
features.update(
@@ -201,26 +186,17 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
@property
def tolerance_s(self) -> float:
"""Tolerance in seconds used to discard loaded frames when their timestamps aren't close enough.
Only used when `delta_timestamps` is provided or when loading video frames from mp4 files.
"""Tolerance in seconds used to discard loaded frames when their timestamps
are not close enough from the requested frames. It is only used when `delta_timestamps`
is provided or when loading video frames from mp4 files.
"""
# 1e-4 to account for possible numerical error
return 1 / self.fps - 1e-4
def __len__(self):
"""The total number of frames across all underlying datasets."""
return self.num_frames
def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
"""Return the frame at `idx`, resolved to the underlying dataset it falls in.
Adds a `"dataset_index"` key identifying which underlying dataset the frame came from, and drops
any `disabled_features` keys.
Raises:
IndexError: If `idx` is out of bounds.
"""
if idx >= len(self):
raise IndexError(f"Index {idx} out of bounds.")
# Determine which dataset to get an item from based on the index.
@@ -243,7 +219,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
return item
def __repr__(self):
"""A summary: repo IDs, sample/episode counts, media type, fps, camera keys, and transforms."""
return (
f"{self.__class__.__name__}(\n"
f" Repository IDs: '{self.repo_ids}',\n"
+12 -17
View File
@@ -26,13 +26,12 @@ from lerobot.utils.feature_utils import hw_to_dataset_features
def create_initial_features(
action: RobotAction | None = None, observation: RobotObservation | None = None
) -> dict[PipelineFeatureType, dict[str, Any]]:
"""Creates the initial features dict for the dataset from action and observation specs.
"""
Creates the initial features dict for the dataset from action and observation specs.
Args:
action (`dict[str, typing.Any] | None`, *optional*): A dictionary of action feature names to their
types/shapes.
observation (`dict[str, typing.Any] | None`, *optional*): A dictionary of observation feature names
to their types/shapes.
action: A dictionary of action feature names to their types/shapes.
observation: A dictionary of observation feature names to their types/shapes.
Returns:
The initial features dictionary structured by PipelineFeatureType.
@@ -47,14 +46,12 @@ def create_initial_features(
# Helper to filter state/action keys based on compiled regex patterns.
def should_keep(key: str, patterns: tuple[re.Pattern] | None) -> bool:
"""Return `True` if `patterns` is `None` or any pattern in it matches `key`."""
if patterns is None:
return True
return any(pat.search(key) for pat in patterns)
def strip_prefix(key: str, prefixes_to_strip: tuple[str]) -> str:
"""Remove the first prefix in `prefixes_to_strip` that `key` starts with, if any."""
for prefix in prefixes_to_strip:
if key.startswith(prefix):
return key[len(prefix) :]
@@ -76,22 +73,20 @@ def aggregate_pipeline_dataset_features(
exclude_images: bool = False,
patterns: Sequence[str] | None = None,
) -> dict[str, dict]:
"""Aggregates and filters pipeline features to create a dataset-ready features dictionary.
"""
Aggregates and filters pipeline features to create a dataset-ready features dictionary.
This function transforms initial features using the pipeline, categorizes them as action or observations
(image or state), filters them based on `exclude_images` and `patterns`, and finally
formats them for use with a Hugging Face LeRobot Dataset.
Args:
pipeline (`DataProcessorPipeline`): The processor pipeline to apply to `initial_features`.
initial_features (`dict`): A dictionary of raw feature specs for actions and observations, keyed by
`PipelineFeatureType`.
use_videos (`bool`, *optional*, defaults to `True`): Controls the storage dtype for image features.
If `True`, images are stored as `"video"`; if `False`, they are stored as `"image"`.
exclude_images (`bool`, *optional*, defaults to `False`): If `True`, image features are dropped
entirely from the output.
patterns (`collections.abc.Sequence[str] | None`, *optional*): A sequence of regex patterns used to
filter action and state features.
pipeline: The DataProcessorPipeline to apply.
initial_features: A dictionary of raw feature specs for actions and observations.
use_videos: Controls the storage dtype for image features. If True, images are stored as "video"; if False, they are stored as "image".
exclude_images: If True, image features are dropped entirely from the output.
patterns: A sequence of regex patterns to filter action and state features.
Image features are not affected by this filter.
Returns:
A dictionary of features formatted for a Hugging Face LeRobot Dataset.
+4 -5
View File
@@ -41,11 +41,10 @@ def write_u16_plane(plane: av.video.plane.VideoPlane, src: np.ndarray, fill_valu
leave the padding untouched.
Args:
plane (`VideoPlane`): Destination 16-bit plane to copy into.
src (`ndarray`): Source image, shape `(height, width)`, dtype `uint16`.
fill_value (`int | None`, *optional*): If given, every pixel of the plane
(including the row padding) is set to this value first, so the padding
holds clean data instead of garbage.
plane: Destination 16-bit plane.
src: Source image, shape ``(height, width)``, dtype ``uint16``.
fill_value: If given, every pixel (padding included) is set to this first, so the
padding holds clean data instead of garbage.
"""
height, width = src.shape
stride_u16 = plane.line_size // np.dtype(np.uint16).itemsize
+1 -17
View File
@@ -55,8 +55,7 @@ class EpisodeAwareSampler:
seed: int = 0,
absolute_to_relative_idx: dict[int, int] | None = None,
):
"""Build the sampler from per-episode `[from, to)` frame-index boundaries.
"""
Args:
dataset_from_indices: Start index of each episode in the dataset.
dataset_to_indices: End index of each episode in the dataset.
@@ -65,13 +64,6 @@ class EpisodeAwareSampler:
drop_n_last_frames: Frames to drop from the end of each episode.
shuffle: Whether to shuffle the indices.
seed: Seed the permutation is derived from (together with the epoch).
absolute_to_relative_idx: Optional mapping from absolute dataset frame index to the relative
index actually yielded (e.g. when the sampler is used over a filtered subset).
Raises:
ValueError: If `drop_n_first_frames`/`drop_n_last_frames` is negative, if
`dataset_from_indices`/`dataset_to_indices` have different lengths, or if no episode has
any frames remaining after dropping.
"""
if drop_n_first_frames < 0:
raise ValueError(f"drop_n_first_frames must be >= 0, got {drop_n_first_frames}")
@@ -124,15 +116,12 @@ class EpisodeAwareSampler:
return [self._frame_index(k) for k in range(self._num_frames)]
def set_epoch(self, epoch: int) -> None:
"""Set the epoch the next `__iter__` call will use, without consuming an auto-advance."""
self._epoch = epoch
def state_dict(self) -> dict:
"""Return `{"epoch": ..., "start_index": ...}`, enough to resume mid-epoch sample-exactly."""
return {"epoch": self._epoch, "start_index": self._start_index}
def load_state_dict(self, state: dict) -> None:
"""Restore the epoch and within-epoch offset from a `state_dict()`-produced dict."""
self._epoch = state["epoch"]
self._start_index = state["start_index"]
@@ -151,10 +140,6 @@ class EpisodeAwareSampler:
return absolute_idx
def __iter__(self) -> Iterator[int]:
"""Yield frame indices for the current epoch (from `set_epoch`/`load_state_dict`), then advance it.
Shuffled if `self.shuffle`, using a permutation seeded from `(seed, epoch)`.
"""
# Advance epoch state eagerly, not on first consumption of the generator.
epoch, start = self._epoch, self._start_index
self._epoch += 1
@@ -171,7 +156,6 @@ class EpisodeAwareSampler:
yield self._frame_index(k)
def __len__(self) -> int:
"""The total number of frames across the sampled episodes (full length, even mid-resume)."""
return self._num_frames
+33 -49
View File
@@ -44,13 +44,17 @@ from .video_utils import (
class LookBackError(Exception):
"""Exception raised when trying to look back in the history of a Backtrackable object."""
"""
Exception raised when trying to look back in the history of a Backtrackable object.
"""
pass
class LookAheadError(Exception):
"""Exception raised when trying to look ahead in the future of a Backtrackable object."""
"""
Exception raised when trying to look ahead in the future of a Backtrackable object.
"""
pass
@@ -60,10 +64,11 @@ class _ShardExhaustedError(Exception):
class Backtrackable[T]:
"""Wrap any iterator/iterable so you can step back up to `history` items and look ahead.
"""
Wrap any iterator/iterable so you can step back up to `history` items
and look ahead up to `lookahead` items.
Looking ahead is bounded by `lookahead` items. This is useful for streaming datasets where you need
to access previous and future items
This is useful for streaming datasets where you need to access previous and future items
but can't load the entire dataset into memory.
Example:
@@ -93,16 +98,6 @@ class Backtrackable[T]:
__slots__ = ("_source", "_back_buf", "_ahead_buf", "_cursor", "_history", "_lookahead")
def __init__(self, iterable: Iterable[T], *, history: int = 1, lookahead: int = 0):
"""Wrap `iterable`, buffering up to `history` past items and `lookahead` future items.
Args:
iterable: The iterable to wrap.
history: How many past items `prev()`/`peek_back()` can reach. Must be `>= 1`.
lookahead: How many future items `peek_ahead()` can reach. Must be `> 0`.
Raises:
ValueError: If `history < 1` or `lookahead <= 0`.
"""
if history < 1:
raise ValueError("history must be >= 1")
if lookahead <= 0:
@@ -116,11 +111,9 @@ class Backtrackable[T]:
self._lookahead = lookahead
def __iter__(self) -> "Backtrackable[T]":
"""Return `self`; `Backtrackable` is its own iterator."""
return self
def __next__(self) -> T:
"""Return the next item, consuming from the back buffer first if `prev()` stepped back."""
# If we've stepped back, consume from back buffer first
if self._cursor < 0: # -1 means "last item", etc.
self._cursor += 1
@@ -135,9 +128,9 @@ class Backtrackable[T]:
return item
def prev(self) -> T:
"""Step one item back in history and return it.
Raises `LookBackError` if already at the oldest buffered item.
"""
Step one item back in history and return it.
Raises IndexError if already at the oldest buffered item.
"""
if len(self._back_buf) + self._cursor <= 1:
raise LookBackError("At start of history")
@@ -146,15 +139,17 @@ class Backtrackable[T]:
return self._back_buf[self._cursor]
def peek_back(self, n: int = 1) -> T:
"""Look `n` items back (n=1 == previous item) without moving the cursor."""
"""
Look `n` items back (n=1 == previous item) without moving the cursor.
"""
if n < 0 or n + 1 > len(self._back_buf) + self._cursor:
raise LookBackError("peek_back distance out of range")
return self._back_buf[self._cursor - (n + 1)]
def peek_ahead(self, n: int = 1) -> T:
"""Look `n` items ahead (n=1 == next item) without moving the cursor.
"""
Look `n` items ahead (n=1 == next item) without moving the cursor.
Fills the ahead buffer if necessary.
"""
if n < 1:
@@ -174,9 +169,9 @@ class Backtrackable[T]:
return self._ahead_buf[n - 1]
def history(self) -> list[T]:
"""Return a copy of the buffered history (most recent last).
The list length is at most the `history` argument passed at construction.
"""
Return a copy of the buffered history (most recent last).
The list length `history` argument passed at construction.
"""
if self._cursor == 0:
return list(self._back_buf)
@@ -185,12 +180,14 @@ class Backtrackable[T]:
return list(self._back_buf)[: self._cursor or None]
def can_peek_back(self, steps: int = 1) -> bool:
"""Check if we can go back `steps` items without raising a `LookBackError`."""
"""
Check if we can go back `steps` items without raising an IndexError.
"""
return steps < len(self._back_buf) + self._cursor
def can_peek_ahead(self, steps: int = 1) -> bool:
"""Check if we can peek ahead `steps` items.
"""
Check if we can peek ahead `steps` items.
This may involve trying to fill the ahead buffer.
"""
if self._lookahead > 0 and steps > self._lookahead:
@@ -278,8 +275,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
episodes (list[int] | None, optional): If specified, this will only load episodes specified by
their episode_index in this list.
image_transforms (Callable | None, optional): Transform to apply to image data.
delta_timestamps (dict[list[float]] | None, optional): Per-feature timestamp offsets (in
seconds, relative to a frame's own timestamp) of additional frames to return alongside it.
tolerance_s (float, optional): Tolerance in seconds for timestamp matching.
revision (str, optional): Git revision id (branch name, tag, or commit hash).
force_cache_sync (bool, optional): Flag to sync and refresh local files first.
@@ -289,8 +284,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
seed (int, optional): Reproducibility random seed.
rng (np.random.Generator | None, optional): Random number generator.
shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
return_uint8 (bool, optional): For RGB videos, whether to return raw uint8 frames instead of
the default float32 frames normalized to [0, 1].
depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm").
Defaults to "mm".
repo_type: "dataset" (default) or "bucket" to stream from an HF Storage Bucket
@@ -390,17 +383,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
@property
def num_frames(self):
"""The total number of frames in the dataset."""
return self.meta.total_frames
@property
def num_episodes(self):
"""The total number of episodes in the dataset."""
return self.meta.total_episodes
@property
def fps(self):
"""The dataset's recording frame rate."""
return self.meta.fps
@property
@@ -425,11 +415,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
# could be used with a ThreadPoolExecutor to run `make_frame` (especially video decoding)
# in parallel, feeding a queue from which this iterator will yield processed items.
def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
"""Yield frames via reservoir-buffered random sampling across shards, streaming indefinitely.
Samples a random shard, then a random frame from a fixed-size buffer refilled from that shard, so
no full shuffle or shard is ever fully materialized in memory.
"""
if self.video_decoder_cache is None:
self.video_decoder_cache = VideoDecoderCache()
@@ -507,7 +492,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
return dict.fromkeys(self.meta.video_keys, [start_ts])
def _make_padding_camera_frame(self, camera_key: str):
"""Variable-shape padding frame for the given camera key, shaped (H, W, C)."""
"""Variable-shape padding frame for given camera keys, given in (H, W, C)"""
return torch.zeros(self.meta.info.features[camera_key]["shape"]).permute(-1, 0, 1)
def _get_video_frame_padding_mask(
@@ -538,7 +523,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
return padding_mask
def make_frame(self, dataset_iterator: Backtrackable) -> Generator:
"""Makes a frame starting from a dataset iterator."""
"""Makes a frame starting from a dataset iterator"""
try:
item = next(dataset_iterator)
except StopIteration as e:
@@ -631,13 +616,12 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
return query_timestamps
def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict:
"""Decode the requested per-camera frame timestamps from `ep_idx`'s videos.
Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
"""Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
in the main process (e.g. by using a second Dataloader with num_workers=0). It will result in a
Segmentation Fault. This probably happens because a memory reference to the video loader is created in
the main process and a subprocess fails to access it.
"""
item = {}
for video_key, query_ts in query_timestamps.items():
root = self.meta.url_root if self.streaming and not self.streaming_from_local else self.root
@@ -680,9 +664,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
"""Get frames with delta offsets using the backtrackable iterator.
Args:
dataset_iterator (Backtrackable): The backtrackable iterator to peek/step through for delta
frames.
current_item (dict): Current item from the iterator.
ep_idx (int): Episode index.
Returns:
tuple: (query_result, padding) - frames at delta offsets and padding info.
@@ -787,7 +770,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
return query_result, padding
def _validate_delta_timestamp_keys(self, delta_timestamps: dict[list[float]]) -> None:
"""Validate that all keys in delta_timestamps correspond to actual features in the dataset.
"""
Validate that all keys in delta_timestamps correspond to actual features in the dataset.
Raises:
ValueError: If any delta timestamp key doesn't correspond to a dataset feature.
+13 -41
View File
@@ -63,21 +63,11 @@ hub_api.create_tag("{repo_id}", tag="_version_", repo_type="dataset")
"""
class CompatibilityError(Exception):
"""Base class for errors raised when a dataset's `codebase_version` doesn't match this install."""
...
class CompatibilityError(Exception): ...
class BackwardCompatibilityError(CompatibilityError):
"""Raised when a dataset was saved with an older, unsupported `codebase_version`."""
def __init__(self, repo_id: str, version: packaging.version.Version):
"""Build the error message pointing the user at the v2.1-to-v3.0 conversion script.
Raises:
NotImplementedError: If `version` isn't the one supported legacy version (2.1).
"""
if version.major == 2 and version.minor == 1:
message = V30_MESSAGE.format(repo_id=repo_id, version=version)
else:
@@ -88,10 +78,7 @@ class BackwardCompatibilityError(CompatibilityError):
class ForwardCompatibilityError(CompatibilityError):
"""Raised when a dataset was saved with a newer `codebase_version` than this install supports."""
def __init__(self, repo_id: str, version: packaging.version.Version):
"""Build the error message pointing the user at upgrading their `lerobot` install."""
message = FUTURE_MESSAGE.format(repo_id=repo_id, version=version)
super().__init__(message)
@@ -202,12 +189,6 @@ class DatasetInfo:
tools: list[dict] | None = None
def __post_init__(self) -> None:
"""Coerce feature shapes from list to tuple, and validate `fps`/`chunks_size`/file-size fields.
Raises:
ValueError: If `fps`, `chunks_size`, `data_files_size_in_mb`, or `video_files_size_in_mb` isn't
positive.
"""
# Coerce feature shapes from list to tuple — JSON deserialisation
# returns lists, but the rest of the codebase expects tuples.
for ft in self.features.values():
@@ -258,11 +239,6 @@ class DatasetInfo:
# Once all callers have been migrated to attribute access, remove these.
# ---------------------------------------------------------------------------
def __getitem__(self, key: str):
"""Deprecated dict-style read; use attribute access instead.
Raises:
KeyError: If `key` isn't a field on this class.
"""
import warnings
warnings.warn(
@@ -277,7 +253,6 @@ class DatasetInfo:
raise KeyError(key) from err
def __setitem__(self, key: str, value) -> None:
"""Deprecated dict-style write; use attribute assignment instead."""
import warnings
warnings.warn(
@@ -315,7 +290,6 @@ def has_legacy_hub_download_metadata(root: Path) -> bool:
def update_chunk_file_indices(chunk_idx: int, file_idx: int, chunks_size: int) -> tuple[int, int]:
"""Advance to the next `(chunk_idx, file_idx)`, rolling over to a new chunk once `chunks_size` is hit."""
if file_idx == chunks_size - 1:
file_idx = 0
chunk_idx += 1
@@ -381,7 +355,7 @@ def check_version_compatibility(
repo_id (str): The repository ID for logging purposes.
version_to_check (str | packaging.version.Version): The version of the dataset.
current_version (str | packaging.version.Version): The current version of the codebase.
enforce_breaking_major (bool, *optional*, defaults to `True`): If True, raise an error on major version mismatch.
enforce_breaking_major (bool): If True, raise an error on major version mismatch.
Raises:
BackwardCompatibilityError: If the dataset version is from a newer, incompatible
@@ -408,9 +382,9 @@ def get_repo_versions(repo_id: str, *, token: str | bool | None = None) -> list[
Args:
repo_id (str): The repository ID on the Hugging Face Hub.
token (`str | bool | None`, *optional*): Authentication token used for Hub requests. Pass a string
token, `True` to require the locally stored token, `False` to disable authentication, or `None`
to use the Hugging Face Hub default.
token: Authentication token used for Hub requests. Pass a string token,
``True`` to require the locally stored token, ``False`` to disable
authentication, or ``None`` to use the Hugging Face Hub default.
Returns:
list[packaging.version.Version]: A list of valid versions found.
@@ -440,7 +414,7 @@ def get_safe_version(
Args:
repo_id (str): The repository ID on the Hugging Face Hub.
version (str | packaging.version.Version): The target version.
token (`str | bool | None`, *optional*): Authentication token forwarded to the Hub version lookup.
token: Authentication token forwarded to the Hub version lookup.
Returns:
str: The safe version string (e.g., "v1.2.3") to use as a revision.
@@ -487,7 +461,7 @@ def create_branch(repo_id: str, *, branch: str, repo_type: str | None = None) ->
Args:
repo_id (str): The ID of the repository.
branch (str): The name of the branch to create.
repo_type (str | None, *optional*): The type of the repository (e.g., "dataset").
repo_type (str | None): The type of the repository (e.g., "dataset").
"""
api = HfApi()
@@ -512,12 +486,10 @@ def create_lerobot_dataset_card(
https://huggingface.co/docs/hub/repositories-licenses.
Args:
tags (list | None, *optional*): A list of tags to add to the dataset card.
dataset_info (DatasetInfo | None, *optional*): The dataset's info object, which will
tags (list | None): A list of tags to add to the dataset card.
dataset_info (DatasetInfo | None): The dataset's info object, which will
be displayed on the card.
kwargs (`Any`, *optional*): Values used to replace placeholders in the card template, e.g. `license`, which
must be a valid license identifier from
https://huggingface.co/docs/hub/repositories-licenses.
**kwargs: Additional keyword arguments to populate the card template.
Returns:
DatasetCard: The generated dataset card object.
@@ -552,12 +524,10 @@ def create_lerobot_dataset_card(
def is_float_in_list(target, float_list, threshold=1e-6):
"""Return `True` if `float_list` contains a value within `threshold` of `target`."""
return any(abs(target - x) <= threshold for x in float_list)
def find_float_index(target, float_list, threshold=1e-6):
"""Return the index of the first value in `float_list` within `threshold` of `target`, or -1."""
for i, x in enumerate(float_list):
if abs(target - x) <= threshold:
return i
@@ -565,7 +535,9 @@ def find_float_index(target, float_list, threshold=1e-6):
def safe_shard(dataset: datasets.IterableDataset, index: int, num_shards: int) -> datasets.Dataset:
"""Safe shards the dataset."""
"""
Safe shards the dataset.
"""
shard_idx = min(dataset.num_shards, index + 1) - 1
return dataset.shard(num_shards, index=shard_idx)
+70 -124
View File
@@ -61,18 +61,19 @@ def decode_video_frames(
return_uint8: bool = False,
is_depth: bool = False,
) -> torch.Tensor:
"""Decodes video frames using the specified backend.
"""
Decodes video frames using the specified backend.
Args:
video_path (Path): Path to the video file.
timestamps (list[float]): List of timestamps to extract frames.
tolerance_s (float): Allowed deviation in seconds for frame retrieval.
backend (str, optional, *optional*): Backend to use for decoding. Defaults to "torchcodec" when available
backend (str, optional): Backend to use for decoding. Defaults to "torchcodec" when available
in the platform; otherwise, defaults to "pyav". The legacy value "video_reader" is
accepted for one release as an alias for "pyav" and will be removed in a future version.
return_uint8 (bool, *optional*, defaults to `False`): For RGB videos, if True return raw uint8 frames without float32 normalization.
return_uint8 (bool): For RGB videos, if True return raw uint8 frames without float32 normalization.
This reduces memory for DataLoader IPC; normalization can be done on GPU afterward.
is_depth (bool, *optional*, defaults to `False`): Set to True if the video is a depth map (1 channel, uint12).
is_depth (bool): Set to True if the video is a depth map (1 channel, uint12).
Returns:
torch.Tensor: Decoded frames (RGB: float32 in [0,1] by default, or uint8 if return_uint8=True, Depth: uint12).
@@ -123,16 +124,14 @@ def decode_video_frames_pyav(
video can be adjusted at encoding time to trade off decoding speed against file size.
Args:
video_path (`pathlib.Path | str`): Path to the video file.
timestamps (`list`): List of timestamps, in seconds, to extract frames for.
tolerance_s (`float`): Allowed deviation in seconds between a queried timestamp
and the closest decoded frame.
log_loaded_timestamps (`bool`, *optional*, defaults to `False`): Whether to log
every decoded frame's timestamp at INFO level.
return_uint8 (`bool`, *optional*, defaults to `False`): For RGB videos, whether to
return raw uint8 frames instead of the default float32 frames normalized to [0, 1].
is_depth (`bool`, *optional*, defaults to `False`): Whether the video is a depth map
(1 channel, uint12).
video_path: Path to the video file.
timestamps: List of timestamps (in seconds) to extract frames for.
tolerance_s: Allowed deviation in seconds between a queried timestamp and the closest
decoded frame.
log_loaded_timestamps: When True, log every decoded frame's timestamp at INFO level.
return_uint8: For RGB videos, if True return raw uint8 frames (C, H, W).
Otherwise, return float32 in [0, 1] range.
is_depth: Set to True if the video is a depth map (1 channel, uint12).
Returns:
torch.Tensor of shape (len(timestamps), C, H, W).
@@ -266,31 +265,15 @@ class VideoDecoderCache:
ever opened until the process exits).
Args:
max_size (`int | None | object`, *optional*, defaults to `<unset>`): Maximum
number of decoders to retain. `None` disables eviction and restores legacy unbounded
behaviour. The sentinel default defers to the value of `LEROBOT_VIDEO_DECODER_CACHE_SIZE`
if set, otherwise `DEFAULT_DECODER_CACHE_SIZE`.
max_size: Maximum number of decoders to retain. ``None`` disables
eviction and restores legacy unbounded behaviour. Defaults to the
value of ``LEROBOT_VIDEO_DECODER_CACHE_SIZE`` if set, otherwise
:data:`DEFAULT_DECODER_CACHE_SIZE`.
"""
class _UnsetSentinel:
"""Singleton marker distinguishing "not passed" from an explicit `None` `max_size`.
Has a fixed `__repr__` (unlike a bare `object()`) so it renders identically across
processes, which keeps the class docstring's `defaults to` clause stable.
"""
def __repr__(self) -> str:
"""Return `"<unset>"`, a stable placeholder for docstrings/logging."""
return "<unset>"
_SENTINEL: ClassVar[object] = _UnsetSentinel()
_SENTINEL: ClassVar[object] = object()
def __init__(self, max_size: int | None | object = _SENTINEL):
"""Create the cache. See the class docstring for `max_size`.
Raises:
ValueError: If `max_size` is neither `None` nor a positive integer.
"""
if max_size is VideoDecoderCache._SENTINEL:
max_size = _default_max_cache_size()
if max_size is not None and max_size <= 0:
@@ -300,7 +283,6 @@ class VideoDecoderCache:
self._lock = Lock()
def __contains__(self, video_path: object) -> bool:
"""Return `True` if `video_path` (as `str`) has a cached decoder."""
with self._lock:
return str(video_path) in self._cache
@@ -356,7 +338,7 @@ class VideoDecoderCache:
class FrameTimestampError(ValueError):
"""Helper error to indicate the retrieved timestamps exceed the queried ones."""
"""Helper error to indicate the retrieved timestamps exceed the queried ones"""
pass
@@ -375,16 +357,11 @@ def decode_video_frames_torchcodec(
"""Loads frames associated with the requested timestamps of a video using torchcodec.
Args:
video_path (`pathlib.Path | str`): Path to the video file.
timestamps (`list`): List of timestamps, in seconds, to extract frames for.
tolerance_s (`float`): Allowed deviation in seconds between a queried timestamp
and the closest decoded frame.
log_loaded_timestamps (`bool`, *optional*, defaults to `False`): Whether to log
every decoded frame's timestamp at INFO level.
decoder_cache (`lerobot.datasets.video_utils.VideoDecoderCache | None`, *optional*): Decoder
cache to fetch the `VideoDecoder` from. Uses the module-level default cache if `None`.
return_uint8 (`bool`, *optional*, defaults to `False`): For RGB videos, whether to
return raw uint8 frames instead of the default float32 frames normalized to [0, 1].
video_path: Path to the video file.
timestamps: List of timestamps to extract frames.
tolerance_s: Allowed deviation in seconds for frame retrieval.
log_loaded_timestamps: Whether to log loaded timestamps.
decoder_cache: Optional decoder cache instance. Uses default if None.
Note: Setting device="cuda" outside the main process, e.g. in data loader workers, will lead to CUDA initialization errors.
@@ -474,19 +451,19 @@ def encode_video_frames(
RGB frames are encoded directly.
Args:
imgs_dir (`pathlib.Path | str`): Directory containing the frames to encode, named
`frame-000000` onwards (`.png` for RGB, `.tiff` for depth).
video_path (`pathlib.Path | str`): Output path for the encoded `.mp4` file.
fps (`int`): Frame rate of the output video.
video_encoder (`lerobot.configs.video.VideoEncoderConfig | None`, *optional*): Encoder settings
(codec, pixel format, quality, ...). When `None`, `rgb_encoder_defaults` is used. Pass a
`DepthEncoderConfig` to encode depth frames.
encoder_threads (`int | None`, *optional*): Per-encoder thread count forwarded to the codec.
`None` lets the codec decide.
log_level (`int | None`, *optional*, defaults to 24): libav log level to set while encoding,
or `None` to leave the current logging configuration unchanged.
overwrite (`bool`, *optional*, defaults to `False`): When `False` and `video_path` already
exists, skip encoding and log a warning. When `True`, re-encode and replace the existing file.
imgs_dir: Directory containing the frames to encode, named ``frame-000000``
onwards (``.png`` for RGB, ``.tiff`` for depth).
video_path: Output path for the encoded ``.mp4`` file.
fps: Frame rate of the output video.
video_encoder: Encoder settings (codec, pixel format, quality, ...). When
``None``, :func:`rgb_encoder_defaults` is used. Pass a
:class:`~lerobot.configs.video.DepthEncoderConfig` to encode depth frames.
encoder_threads: Per-encoder thread count forwarded to the codec. ``None``
lets the codec decide.
log_level: libav log level to set while encoding, or ``None`` to leave the
current logging configuration unchanged.
overwrite: When ``False`` and ``video_path`` already exists, skip encoding and
log a warning. When ``True``, re-encode and replace the existing file.
"""
if video_encoder is None:
video_encoder = rgb_encoder_defaults()
@@ -575,21 +552,16 @@ def reencode_video(
"""Re-encode a video file, optionally trimming it to ``[start_time_s, end_time_s)``.
Args:
input_video_path (`pathlib.Path | str`): Existing video file to read.
output_video_path (`pathlib.Path | str`): Path for the re-encoded file.
video_encoder (`lerobot.configs.video.VideoEncoderConfig | None`, *optional*): Encoder
configuration. Defaults to `rgb_encoder_defaults`.
encoder_threads (`int | None`, *optional*): Optional thread count forwarded to
`VideoEncoderConfig.get_codec_options`.
log_level (`int | None`, *optional*, defaults to 24): libav log level while encoding,
or `None` to leave logging unchanged.
overwrite (`bool`, *optional*, defaults to `False`): When `False` and `output_video_path`
already exists, skip and log a warning.
start_time_s (`float | None`, *optional*): When set, trim the output to start at this
timestamp, in seconds.
end_time_s (`float | None`, *optional*): When set, trim the output to end at this
timestamp, in seconds, exclusive.
input_video_path: Existing video file to read.
output_video_path: Path for the re-encoded file.
video_encoder: Encoder configuration. Defaults to :func:`rgb_encoder_defaults`.
encoder_threads: Optional thread count forwarded to :meth:`VideoEncoderConfig.get_codec_options`.
log_level: libav log level while encoding, or ``None`` to leave logging unchanged. Defaults to WARNING.
overwrite: When ``False`` and ``output_video_path`` already exists, skip and log a warning.
start_time_s: When set, trim the output to start at this timestamp (seconds).
end_time_s: When set, trim the output to end at this timestamp (seconds, exclusive).
"""
video_encoder = video_encoder or rgb_encoder_defaults()
if (start_time_s is not None and start_time_s < 0) or (end_time_s is not None and end_time_s < 0):
@@ -679,26 +651,25 @@ def concatenate_video_files(
overwrite: bool = True,
compatibility_check: bool = False,
):
"""Concatenate multiple video files into a single video file using pyav.
"""
Concatenate multiple video files into a single video file using pyav.
This function takes a list of video input file paths and concatenates them into a single
output video file. It uses ffmpeg's concat demuxer with stream copy mode for fast
concatenation without re-encoding.
Args:
input_video_paths (`list`): Ordered list of input video file paths to concatenate.
output_video_path (`Path`): Path to the output video file.
overwrite (`bool`, *optional*, defaults to `True`): Whether to overwrite the output
video file if it already exists.
compatibility_check (`bool`, *optional*, defaults to `False`): Whether to check that
the input videos share the same height, width, fps, codec, and pixel format
before concatenating.
input_video_paths: Ordered list of input video file paths to concatenate.
output_video_path: Path to the output video file.
overwrite: Whether to overwrite the output video file if it already exists. Default is True.
compatibility_check: Whether to check if the input videos are compatible. Default is False.
Note:
- Creates a temporary directory for intermediate files that is cleaned up after use.
- Uses ffmpeg's concat demuxer which requires all input videos to have the same
codec, resolution, and frame rate for proper concatenation.
"""
output_video_path = Path(output_video_path)
if output_video_path.exists() and not overwrite:
@@ -796,17 +767,6 @@ class _CameraEncoderThread(threading.Thread):
stop_event: threading.Event,
encoder_threads: int | None = None,
):
"""Set up the thread; frames are only consumed once `start()` is called.
Args:
video_path: Output MP4 path.
fps: Output frame rate.
video_encoder: Codec/quality settings; `DepthEncoderConfig` selects depth-map encoding.
frame_queue: Queue this thread reads `(frame, ...)` items from.
result_queue: Queue the final stats are pushed to once encoding finishes.
stop_event: Set by the caller to signal this thread to stop early.
encoder_threads: Number of threads passed to the codec, if it supports one.
"""
super().__init__(daemon=True)
self.video_path = video_path
self.fps = fps
@@ -818,10 +778,6 @@ class _CameraEncoderThread(threading.Thread):
self.encoder_threads = encoder_threads
def run(self) -> None:
"""Encode frames from `frame_queue` to `video_path` until a stop sentinel or `stop_event`.
Pushes the accumulated `RunningQuantileStats` to `result_queue` once encoding finishes.
"""
from .compute_stats import RunningQuantileStats, auto_downsample_height_width
container = None
@@ -942,8 +898,7 @@ class StreamingVideoEncoder:
queue_maxsize: int = 30,
encoder_threads: int | None = None,
):
"""Create the manager; per-camera encoder threads are started lazily on first frame.
"""
Args:
fps: Frames per second for the output videos.
rgb_encoder: Video encoder settings applied to all RGB cameras.
@@ -1150,9 +1105,11 @@ class StreamingVideoEncoder:
@dataclass
class VideoFrame:
# TODO(rcadene, lhoestq): move to Hugging Face `datasets` repo
"""Provides a type for a dataset containing video frames.
"""
Provides a type for a dataset containing video frames.
Example:
```python
data_dict = [{"image": {"path": "videos/episode_0.mp4", "timestamp": 0.3}}]
features = {"image": VideoFrame()}
@@ -1164,7 +1121,6 @@ class VideoFrame:
_type: str = field(default="VideoFrame", init=False, repr=False)
def __call__(self):
"""Return the pyarrow struct type backing this feature, as required by `datasets.Features`."""
return self.pa_type
@@ -1179,11 +1135,6 @@ with warnings.catch_warnings():
def get_audio_info(video_path: Path | str) -> dict:
"""Read audio-stream metadata (channels, codec, bit rate, sample rate, etc.) from a video file.
Returns:
A dict of `"audio.*"` keys, or `{"has_audio": False}` if `video_path` has no audio stream.
"""
# Set logging level
logging.getLogger("libav").setLevel(av.logging.WARNING)
@@ -1222,13 +1173,13 @@ def get_video_info(
"""Build the ``video.*`` / ``audio.*`` info dict persisted in ``info.json``.
Args:
video_path (`pathlib.Path | str`): Path to the encoded video file to probe.
video_encoder (`lerobot.configs.video.VideoEncoderConfig | None`, *optional*): If provided,
record the exact encoder settings used to encode this video. Stream-derived values take
precedence — encoder fields are only written for keys not already populated from the
video file itself. When a `DepthEncoderConfig` is passed, the depth quantization
parameters (`depth_min` / `depth_max` / `shift` / `use_log`) are recorded so frames can
be dequantized on read.
video_path: Path to the encoded video file to probe.
video_encoder: If provided, record the exact encoder settings used to encode this
video. Stream-derived values take precedence — encoder fields are only written for keys
not already populated from the video file itself. When a
:class:`~lerobot.configs.video.DepthEncoderConfig` is passed, the depth
quantization parameters (``depth_min`` / ``depth_max`` / ``shift`` /
``use_log``) are recorded so frames can be dequantized on read.
Returns:
The ``video.*`` / ``audio.*`` info dict, including ``is_depth_map`` which is
@@ -1276,10 +1227,11 @@ def get_video_info(
def get_video_duration_in_s(video_path: Path | str) -> float:
"""Get the duration of a video file in seconds using PyAV.
"""
Get the duration of a video file in seconds using PyAV.
Args:
video_path (`pathlib.Path | str`): Path to the video file.
video_path: Path to the video file.
Returns:
Duration of the video in seconds.
@@ -1297,7 +1249,8 @@ def get_video_duration_in_s(video_path: Path | str) -> float:
class VideoEncodingManager:
"""Context manager that ensures proper video encoding and data cleanup even if exceptions occur.
"""
Context manager that ensures proper video encoding and data cleanup even if exceptions occur.
This manager handles:
- Batch encoding for any remaining episodes when recording interrupted
@@ -1305,23 +1258,16 @@ class VideoEncodingManager:
- Removing empty image directories
Args:
dataset (`LeRobotDataset`): The LeRobotDataset instance.
dataset: The LeRobotDataset instance
"""
def __init__(self, dataset):
"""Store the `LeRobotDataset` this manager will finalize/clean up on exit."""
self.dataset = dataset
def __enter__(self):
"""Return `self`; no setup is needed on entry."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Finalize the dataset, cancelling pending videos and cleaning up interrupted-episode files.
Runs unconditionally (even if `exc_type` is set), so partial/interrupted recordings still leave
a consistent dataset on disk.
"""
writer = self.dataset.writer
if writer is not None:
if exc_type is not None and writer._streaming_encoder is not None:
@@ -29,14 +29,19 @@ logger = logging.getLogger(__name__)
class BiOpenArmLeader(BimanualMixin, Teleoperator):
"""
Bimanual OpenArm Leader Arms
"""
"""A bimanual pair of [`~teleoperators.openarm_leader.OpenArmLeader`] arms."""
config_class = BiOpenArmLeaderConfig
name = "bi_openarm_leader"
def __init__(self, config: BiOpenArmLeaderConfig):
"""Build the teleoperator from its configuration.
Args:
config (`BiOpenArmLeaderConfig`):
The teleoperator's configuration. Its `left_arm_config` and `right_arm_config` determine
what is connected on each side.
"""
super().__init__(config)
self.config = config
@@ -75,6 +80,10 @@ class BiOpenArmLeader(BimanualMixin, Teleoperator):
@cached_property
def action_features(self) -> dict[str, type]:
"""See [`~teleoperators.Teleoperator.action_features`].
Merges both arms' features, each key prefixed with `left_` or `right_`.
"""
left_arm_features = self.left_arm.action_features
right_arm_features = self.right_arm.action_features
@@ -85,15 +94,31 @@ class BiOpenArmLeader(BimanualMixin, Teleoperator):
@cached_property
def feedback_features(self) -> dict[str, type]:
"""See [`~teleoperators.Teleoperator.feedback_features`].
Always empty: feedback is not implemented for the OpenArm leader.
"""
return {}
def setup_motors(self) -> None:
"""Not supported: raises `NotImplementedError`.
Motor ID configuration for CAN motors is typically done via manufacturer tools rather than through
LeRobot.
Raises:
NotImplementedError: Always.
"""
raise NotImplementedError(
"Motor ID configuration is typically done via manufacturer tools for CAN motors."
)
@check_if_not_connected
def get_action(self) -> RobotAction:
"""See [`~teleoperators.Teleoperator.get_action`].
Merges both arms' actions, each key prefixed with `left_` or `right_`.
"""
action_dict = {}
# Add "left_" prefix
@@ -107,5 +132,14 @@ class BiOpenArmLeader(BimanualMixin, Teleoperator):
return action_dict
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not supported: raises `NotImplementedError`.
Args:
feedback (`dict[str, float]`):
Unused.
Raises:
NotImplementedError: Always.
"""
# TODO: Implement force feedback
raise NotImplementedError
@@ -23,7 +23,23 @@ from ..openarm_leader import OpenArmLeaderConfigBase
@TeleoperatorConfig.register_subclass("bi_openarm_leader")
@dataclass
class BiOpenArmLeaderConfig(TeleoperatorConfig):
"""Configuration class for Bi OpenArm Leader teleoperators."""
"""Configuration for a bimanual pair of OpenArm leader arms.
The two arms are configured independently, then driven as one teleoperator: action keys from each arm
are prefixed with `left_` and `right_`.
Calibration is per arm, taken from each arm config's own `id` and `calibration_dir`.
Args:
left_arm_config (`OpenArmLeaderConfigBase`):
Configuration for the left arm, including its own `port` and `motor_config`.
right_arm_config (`OpenArmLeaderConfigBase`):
Configuration for the right arm, including its own `port` and `motor_config`.
id (`str`, *optional*):
Identifier for the pair as a whole.
calibration_dir (`Path`, *optional*):
Unused at this level; each arm calibrates through its own config.
"""
left_arm_config: OpenArmLeaderConfigBase
right_arm_config: OpenArmLeaderConfigBase
@@ -40,6 +40,14 @@ class BiOpenArmMini(BimanualMixin, Teleoperator):
name = "bi_openarm_mini"
def __init__(self, config: BiOpenArmMiniConfig):
"""Build the teleoperator from its configuration.
Args:
config (`BiOpenArmMiniConfig`):
The teleoperator's configuration. Its `left_arm_config` and `right_arm_config` determine
what is connected on each side; each arm's `side` is forced to `"left"`/`"right"`
regardless of what was set on the per-arm config.
"""
super().__init__(config)
self.config = config
@@ -66,6 +74,10 @@ class BiOpenArmMini(BimanualMixin, Teleoperator):
@cached_property
def action_features(self) -> dict[str, type]:
"""See [`~teleoperators.Teleoperator.action_features`].
Merges both arms' features, each key prefixed with `left_` or `right_`.
"""
return {
**{f"left_{k}": v for k, v in self.left_arm.action_features.items()},
**{f"right_{k}": v for k, v in self.right_arm.action_features.items()},
@@ -73,17 +85,30 @@ class BiOpenArmMini(BimanualMixin, Teleoperator):
@cached_property
def feedback_features(self) -> dict[str, type]:
"""See [`~teleoperators.Teleoperator.feedback_features`].
Merges both arms' features, each key prefixed with `left_` or `right_`.
"""
return {
**{f"left_{k}": v for k, v in self.left_arm.feedback_features.items()},
**{f"right_{k}": v for k, v in self.right_arm.feedback_features.items()},
}
def setup_motors(self) -> None:
"""Assign each motor its bus ID, one arm at a time.
Run this once when building the teleoperator. Interactive: prompts you to connect the controller
board to a single motor at a time, left arm first.
"""
self.left_arm.setup_motors()
self.right_arm.setup_motors()
@check_if_not_connected
def get_action(self) -> RobotAction:
"""See [`~teleoperators.Teleoperator.get_action`].
Merges both arms' actions, each key prefixed with `left_` or `right_`.
"""
action: RobotAction = {}
for k, v in self.left_arm.get_action().items():
action[f"left_{k}"] = v
@@ -93,6 +118,14 @@ class BiOpenArmMini(BimanualMixin, Teleoperator):
@check_if_not_connected
def send_feedback(self, feedback: dict[str, float]) -> None:
"""See [`~teleoperators.Teleoperator.send_feedback`].
Args:
feedback (`dict[str, float]`):
Feedback values keyed with `left_`/`right_` prefixes, as produced by
[`~teleoperators.bi_openarm_mini.BiOpenArmMini.get_action`]. Each arm only receives the entries for its
own side.
"""
left_fb = {k.removeprefix("left_"): v for k, v in feedback.items() if k.startswith("left_")}
right_fb = {k.removeprefix("right_"): v for k, v in feedback.items() if k.startswith("right_")}
if left_fb:
@@ -23,7 +23,18 @@ from ..openarm_mini import OpenArmMiniConfigBase
@TeleoperatorConfig.register_subclass("bi_openarm_mini")
@dataclass
class BiOpenArmMiniConfig(TeleoperatorConfig):
"""Configuration class for Bi OpenArm Mini teleoperators."""
"""Configuration for a bimanual pair of OpenArm Mini leader arms.
Args:
left_arm_config (`OpenArmMiniConfigBase`):
Configuration for the left arm.
right_arm_config (`OpenArmMiniConfigBase`):
Configuration for the right arm.
id (`str`, *optional*):
Identifier for this particular unit; also names its calibration file.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
"""
left_arm_config: OpenArmMiniConfigBase
right_arm_config: OpenArmMiniConfigBase
@@ -40,6 +40,12 @@ class BiRebot102Leader(BimanualMixin, Teleoperator):
name = "bi_rebot_102_leader"
def __init__(self, config: BiRebot102LeaderConfig):
"""Build the two underlying [`~teleoperators.rebot_102_leader.RebotArm102Leader`] arms.
Args:
config (`BiRebot102LeaderConfig`):
The teleoperator's configuration.
"""
super().__init__(config)
self.config = config
@@ -68,6 +74,11 @@ class BiRebot102Leader(BimanualMixin, Teleoperator):
@cached_property
def action_features(self) -> dict[str, type]:
"""The union of both arms' action features, each key prefixed `left_` / `right_`.
Returns:
`dict[str, type]`: See [`~teleoperators.rebot_102_leader.RebotArm102Leader.action_features`].
"""
return {
**{f"left_{k}": v for k, v in self.left_arm.action_features.items()},
**{f"right_{k}": v for k, v in self.right_arm.action_features.items()},
@@ -75,14 +86,29 @@ class BiRebot102Leader(BimanualMixin, Teleoperator):
@cached_property
def feedback_features(self) -> dict[str, type]:
"""Neither arm accepts feedback.
Returns:
`dict[str, type]`: Always empty.
"""
return {}
@check_if_not_connected
def get_action(self) -> RobotAction:
"""Read both arms' actions and merge them under `left_` / `right_` prefixed keys.
Returns:
`dict[str, float]`: See [`~teleoperators.rebot_102_leader.RebotArm102Leader.get_action`].
"""
action_dict = {}
action_dict.update({f"left_{k}": v for k, v in self.left_arm.get_action().items()})
action_dict.update({f"right_{k}": v for k, v in self.right_arm.get_action().items()})
return action_dict
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not supported: neither arm has actuators to receive feedback.
Raises:
NotImplementedError: Always.
"""
raise NotImplementedError("Feedback is not implemented for the reBot Arm 102 leader.")
@@ -23,7 +23,23 @@ from ..rebot_102_leader import RebotArm102LeaderConfig
@TeleoperatorConfig.register_subclass("bi_rebot_102_leader")
@dataclass
class BiRebot102LeaderConfig(TeleoperatorConfig):
"""Configuration class for the bimanual reBot Arm 102 leader teleoperator."""
"""Configuration class for the bimanual reBot Arm 102 leader teleoperator.
Args:
left_arm_config (`RebotArm102LeaderConfig`):
Configuration of the left [`~teleoperators.rebot_102_leader.RebotArm102Leader`] arm. Its
`id` and `calibration_dir` are ignored; the bimanual `id` and `calibration_dir` below are
used for both arms instead.
right_arm_config (`RebotArm102LeaderConfig`):
Configuration of the right [`~teleoperators.rebot_102_leader.RebotArm102Leader`] arm. Same
caveat as `left_arm_config`.
id (`str`, *optional*):
Identifier for this particular unit; also names the calibration files for both arms
(suffixed `_left` / `_right`).
calibration_dir (`Path`, *optional*):
Where to read and write both arms' calibration files. Defaults to the LeRobot calibration
home.
"""
left_arm_config: RebotArm102LeaderConfig
right_arm_config: RebotArm102LeaderConfig
@@ -29,14 +29,19 @@ logger = logging.getLogger(__name__)
class BiSOLeader(BimanualMixin, Teleoperator):
"""
[Bimanual SO Leader Arms](https://github.com/TheRobotStudio/SO-ARM100) designed by TheRobotStudio
"""
"""A bimanual pair of [SO leader arms](https://github.com/TheRobotStudio/SO-ARM100) by TheRobotStudio."""
config_class = BiSOLeaderConfig
name = "bi_so_leader"
def __init__(self, config: BiSOLeaderConfig):
"""Build the teleoperator from its configuration.
Args:
config (`BiSOLeaderConfig`):
The teleoperator's configuration. Its `left_arm_config` and `right_arm_config` determine
what is connected.
"""
super().__init__(config)
self.config = config
@@ -61,6 +66,12 @@ class BiSOLeader(BimanualMixin, Teleoperator):
@cached_property
def action_features(self) -> dict[str, type]:
"""The values this teleoperator produces, and their types.
Returns:
`dict[str, type]`: Each arm's [`~teleoperators.so_leader.SOLeader.action_features`] keys,
prefixed with `left_` or `right_`.
"""
left_arm_features = self.left_arm.action_features
right_arm_features = self.right_arm.action_features
@@ -71,6 +82,12 @@ class BiSOLeader(BimanualMixin, Teleoperator):
@cached_property
def feedback_features(self) -> dict[str, type]:
"""The values this teleoperator accepts as feedback, and their types.
Returns:
`dict[str, type]`: Each arm's [`~teleoperators.so_leader.SOLeader.feedback_features`] keys,
prefixed with `left_` or `right_`.
"""
# Bimanual teleop has feedback (can be actuated for handover).
# Return the same structure as action_features for consistency with left/right arms.
left_arm_features = self.left_arm.feedback_features
@@ -82,11 +99,25 @@ class BiSOLeader(BimanualMixin, Teleoperator):
}
def setup_motors(self) -> None:
"""Assign each motor its bus ID on both arms, one at a time.
Run this once when building the teleoperator. Interactive: prompts you to connect the controller
board to a single motor at a time, left arm first.
"""
self.left_arm.setup_motors()
self.right_arm.setup_motors()
@check_if_not_connected
def get_action(self) -> RobotAction:
"""Retrieve the current action from both leader arms.
Returns:
`dict[str, Any]`: Each arm's action, keyed as described by
[`~teleoperators.bi_so_leader.BiSOLeader.action_features`].
Raises:
DeviceNotConnectedError: If [`~teleoperators.Teleoperator.connect`] has not been called.
"""
action_dict = {}
# Add "left_" prefix
@@ -23,7 +23,18 @@ from ..so_leader import SOLeaderConfig
@TeleoperatorConfig.register_subclass("bi_so_leader")
@dataclass
class BiSOLeaderConfig(TeleoperatorConfig):
"""Configuration class for Bi SO Leader teleoperators."""
"""Configuration for a bimanual pair of SO-family leader arms.
Args:
left_arm_config (`SOLeaderConfig`):
Configuration for the left arm.
right_arm_config (`SOLeaderConfig`):
Configuration for the right arm.
id (`str`, *optional*):
Identifier for this particular unit; also names its calibration file.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
"""
left_arm_config: SOLeaderConfig
right_arm_config: SOLeaderConfig
+20
View File
@@ -21,6 +21,21 @@ import draccus
@dataclass(kw_only=True)
class TeleoperatorConfig(draccus.ChoiceRegistry, abc.ABC):
"""Base configuration shared by every teleoperator.
Concrete teleoperators subclass this and register themselves with
`@TeleoperatorConfig.register_subclass("name")`, which is what makes `--teleop.type=name` work on the
command line. Subclasses inherit the two fields below and must document them alongside their own.
Args:
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
"""
# Allows to distinguish between different teleoperators of the same type
id: str | None = None
# Directory to store calibration file
@@ -28,4 +43,9 @@ class TeleoperatorConfig(draccus.ChoiceRegistry, abc.ABC):
@property
def type(self) -> str:
"""Return the registered name this config was registered under.
Returns:
`str`: The name passed to `@TeleoperatorConfig.register_subclass`, e.g. `"so101_leader"`.
"""
return self.get_choice_name(self.__class__)
@@ -22,6 +22,22 @@ from ..config import TeleoperatorConfig
@TeleoperatorConfig.register_subclass("gamepad")
@dataclass
class GamepadTeleopConfig(TeleoperatorConfig):
"""Configuration for the gamepad teleoperator.
Args:
use_gripper (`bool`, *optional*, defaults to `True`):
Whether to include a `gripper` entry in the produced actions.
hidapi_fallback (`bool`, *optional*, defaults to `False`):
Read the gamepad through `hidapi` instead of `pygame`. Set this on macOS if `pygame` does not
reliably detect input from your controller.
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
"""
use_gripper: bool = True
# Use hidapi instead of pygame for controllers that pygame cannot detect reliably.
hidapi_fallback: bool = False
@@ -34,68 +34,88 @@ else:
class InputController:
"""Base class for input controllers that generate motion deltas."""
"""Base class for input controllers that generate motion deltas for gamepad-style teleoperation.
Subclasses override `start`, `stop`, `update`, and `get_deltas` to read an actual device; this base
class returns inert defaults.
"""
def __init__(self, x_step_size=1.0, y_step_size=1.0, z_step_size=1.0):
"""
Initialize the controller.
"""Instantiate the controller's step sizes and reset its state.
Args:
x_step_size: Base movement step size in meters
y_step_size: Base movement step size in meters
z_step_size: Base movement step size in meters
x_step_size (`float`, *optional*, defaults to 1.0):
Movement step size along X, in meters.
y_step_size (`float`, *optional*, defaults to 1.0):
Movement step size along Y, in meters.
z_step_size (`float`, *optional*, defaults to 1.0):
Movement step size along Z, in meters.
"""
self.x_step_size = x_step_size
self.y_step_size = y_step_size
self.z_step_size = z_step_size
self.running = True
self.episode_end_status = None # None, "success", or "failure"
self.episode_end_status = None # None, or a TeleopEvents member (SUCCESS, FAILURE, RERECORD_EPISODE)
self.intervention_flag = False
self.open_gripper_command = False
self.close_gripper_command = False
def start(self):
"""Start the controller and initialize resources."""
"""Start the controller and initialize resources. Subclasses open the actual device here."""
pass
def stop(self):
"""Stop the controller and release resources."""
"""Stop the controller and release resources. Subclasses close the actual device here."""
pass
def get_deltas(self):
"""Get the current movement deltas (dx, dy, dz) in meters."""
"""Get the current movement deltas.
Returns:
`tuple[float, float, float]`: `(dx, dy, dz)` in meters. Always `(0.0, 0.0, 0.0)` on the base
class.
"""
return 0.0, 0.0, 0.0
def update(self):
"""Update controller state - call this once per frame."""
"""Refresh the controller's internal state. Call this once per frame before reading deltas or events."""
pass
def __enter__(self):
"""Support for use in 'with' statements."""
"""Support for use in `with` statements. Calls `start`."""
self.start()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Ensure resources are released when exiting 'with' block."""
"""Ensure resources are released when exiting a `with` block, even on error."""
self.stop()
def get_episode_end_status(self):
"""
Get the current episode end status.
"""Read and clear the current episode end status.
Returns:
None if episode should continue, "success" or "failure" otherwise
`TeleopEvents | None`: `None` if the episode should continue, otherwise whichever
[`~teleoperators.TeleopEvents`] member (e.g. `SUCCESS`, `FAILURE`, `RERECORD_EPISODE`) a
subclass most recently recorded.
"""
status = self.episode_end_status
self.episode_end_status = None # Reset after reading
return status
def should_intervene(self):
"""Return True if intervention flag was set."""
"""Whether the intervention flag is currently set.
Returns:
`bool`: `True` if a human is currently intervening.
"""
return self.intervention_flag
def gripper_command(self):
"""Return the current gripper command."""
"""Derive a gripper command from the open/close button flags.
Returns:
`str`: `"open"` or `"close"` if exactly one of the flags is set, `"stay"` otherwise.
"""
if self.open_gripper_command == self.close_gripper_command:
return "stay"
elif self.open_gripper_command:
@@ -105,9 +125,14 @@ class InputController:
class KeyboardController(InputController):
"""Generate motion deltas from keyboard input."""
"""Generate motion deltas from keyboard input via `pynput`, as an alternative to a physical gamepad.
Arrow keys drive X/Y, shift/shift_r drive Z, `enter`/`backspace` end the episode with success/failure,
and `esc` stops the listener.
"""
def __init__(self, x_step_size=1.0, y_step_size=1.0, z_step_size=1.0):
"""See `InputController.__init__`; the step sizes have the same meaning here."""
super().__init__(x_step_size, y_step_size, z_step_size)
self.key_states = {
"forward_x": False,
@@ -123,7 +148,7 @@ class KeyboardController(InputController):
self.listener = None
def start(self):
"""Start the keyboard listener."""
"""Start the `pynput` keyboard listener, if the current session can capture key events."""
if not pynput_can_capture():
logging.warning(
"Keyboard control is unavailable in this environment. pynput cannot capture keys "
@@ -136,6 +161,7 @@ class KeyboardController(InputController):
from pynput import keyboard
def on_press(key):
"""Update key/episode state for a key-down event."""
try:
if key == keyboard.Key.up:
self.key_states["forward_x"] = True
@@ -163,6 +189,7 @@ class KeyboardController(InputController):
pass
def on_release(key):
"""Update key state for a key-up event."""
try:
if key == keyboard.Key.up:
self.key_states["forward_x"] = False
@@ -194,12 +221,16 @@ class KeyboardController(InputController):
print(" ESC: Exit")
def stop(self):
"""Stop the keyboard listener."""
"""Stop the `pynput` keyboard listener."""
if self.listener and self.listener.is_alive():
self.listener.stop()
def get_deltas(self):
"""Get the current movement deltas from keyboard state."""
"""Get the current movement deltas from held-down arrow/shift keys.
Returns:
`tuple[float, float, float]`: `(dx, dy, dz)` in meters.
"""
delta_x = delta_y = delta_z = 0.0
if self.key_states["forward_x"]:
@@ -219,9 +250,29 @@ class KeyboardController(InputController):
class GamepadController(InputController):
"""Generate motion deltas from gamepad input."""
"""Generate motion deltas from gamepad input via `pygame`.
Left stick drives X/Y, the right stick's vertical axis drives Z. Y/Triangle, A/Cross, and X/Square
end the episode with success, failure, or rerecord respectively; RB/LT open and close the gripper;
holding RB also sets the intervention flag.
"""
def __init__(self, x_step_size=1.0, y_step_size=1.0, z_step_size=1.0, deadzone=0.1):
"""Instantiate the controller.
Args:
x_step_size (`float`, *optional*, defaults to 1.0):
Movement step size along X, in meters.
y_step_size (`float`, *optional*, defaults to 1.0):
Movement step size along Y, in meters.
z_step_size (`float`, *optional*, defaults to 1.0):
Movement step size along Z, in meters.
deadzone (`float`, *optional*, defaults to 0.1):
Minimum absolute stick reading before it is treated as input, to filter out drift.
Raises:
ImportError: If `pygame` is not installed.
"""
require_package("pygame", extra="gamepad")
super().__init__(x_step_size, y_step_size, z_step_size)
self.deadzone = deadzone
@@ -229,7 +280,7 @@ class GamepadController(InputController):
self.intervention_flag = False
def start(self):
"""Initialize pygame and the gamepad."""
"""Initialize `pygame` and connect to the first detected joystick."""
pygame.init()
pygame.joystick.init()
@@ -251,7 +302,7 @@ class GamepadController(InputController):
print(" X/Square button: Rerecord episode")
def stop(self):
"""Clean up pygame resources."""
"""Clean up `pygame` joystick and display resources."""
if pygame.joystick.get_init():
if self.joystick:
self.joystick.quit()
@@ -259,7 +310,7 @@ class GamepadController(InputController):
pygame.quit()
def update(self):
"""Process pygame events to get fresh gamepad readings."""
"""Drain pending `pygame` events to refresh button, episode, and intervention state."""
for event in pygame.event.get():
if event.type == pygame.JOYBUTTONDOWN:
if event.button == 3:
@@ -297,7 +348,12 @@ class GamepadController(InputController):
self.intervention_flag = False
def get_deltas(self):
"""Get the current movement deltas from gamepad state."""
"""Get the current movement deltas from the joystick axes, after applying the deadzone.
Returns:
`tuple[float, float, float]`: `(dx, dy, dz)` in meters. `(0.0, 0.0, 0.0)` if reading the
joystick raises `pygame.error` (e.g. the controller was disconnected).
"""
try:
# Read joystick axes
# Left stick X and Y (typically axes 0 and 1)
@@ -325,7 +381,12 @@ class GamepadController(InputController):
class GamepadControllerHID(InputController):
"""Generate motion deltas from gamepad input using HIDAPI."""
"""Generate motion deltas from gamepad input by reading raw HID reports via `hidapi`.
An alternative to `GamepadController` for controllers `pygame` does not reliably detect (notably on
macOS). Byte offsets in `update` are tuned for the Logitech RumblePad 2 and may need adjusting for
other controllers.
"""
def __init__(
self,
@@ -334,13 +395,20 @@ class GamepadControllerHID(InputController):
z_step_size=1.0,
deadzone=0.1,
):
"""
Initialize the HID gamepad controller.
"""Instantiate the controller.
Args:
step_size: Base movement step size in meters
z_scale: Scaling factor for Z-axis movement
deadzone: Joystick deadzone to prevent drift
x_step_size (`float`, *optional*, defaults to 1.0):
Movement step size along X, in meters.
y_step_size (`float`, *optional*, defaults to 1.0):
Movement step size along Y, in meters.
z_step_size (`float`, *optional*, defaults to 1.0):
Movement step size along Z, in meters.
deadzone (`float`, *optional*, defaults to 0.1):
Minimum absolute stick reading before it is treated as input, to filter out drift.
Raises:
ImportError: If `hidapi` is not installed.
"""
require_package("hidapi", extra="gamepad", import_name="hid")
super().__init__(x_step_size, y_step_size, z_step_size)
@@ -358,7 +426,14 @@ class GamepadControllerHID(InputController):
self.buttons = {}
def find_device(self):
"""Look for the gamepad device by vendor and product ID."""
"""Look for a supported gamepad among enumerated HID devices.
Matches the first device whose product string contains `"Logitech"`, `"Xbox"`, `"PS4"`, or
`"PS5"`.
Returns:
`dict | None`: The `hidapi` device info dict, or `None` if no matching device was found.
"""
devices = hid.enumerate()
for device in devices:
device_name = device["product_string"]
@@ -371,7 +446,7 @@ class GamepadControllerHID(InputController):
return None
def start(self):
"""Connect to the gamepad using HIDAPI."""
"""Find and open the gamepad's HID device in non-blocking mode."""
self.device_info = self.find_device()
if not self.device_info:
self.running = False
@@ -406,9 +481,9 @@ class GamepadControllerHID(InputController):
self.device = None
def update(self):
"""
Read and process the latest gamepad data.
Due to an issue with the HIDAPI, we need to read the read the device several times in order to get a stable reading
"""Read and process the latest gamepad HID report.
Reads the device 10 times in a row, since a single `hidapi` read can otherwise return stale data.
"""
for _ in range(10):
self._update()
@@ -464,7 +539,11 @@ class GamepadControllerHID(InputController):
logging.error(f"Error reading from gamepad: {e}")
def get_deltas(self):
"""Get the current movement deltas from gamepad state."""
"""Get the current movement deltas from the last-read HID report.
Returns:
`tuple[float, float, float]`: `(dx, dy, dz)` in meters.
"""
# Calculate deltas - invert as needed based on controller orientation
delta_x = -self.left_x * self.x_step_size # Forward/backward
delta_y = -self.left_y * self.y_step_size # Left/right
@@ -32,6 +32,14 @@ logger = logging.getLogger(__name__)
class GripperAction(IntEnum):
"""Gripper command levels produced by a gamepad's gripper buttons.
**Attributes**:
- **CLOSE** (`int`) -- Close the gripper.
- **STAY** (`int`) -- Leave the gripper where it is.
- **OPEN** (`int`) -- Open the gripper.
"""
CLOSE = 0
STAY = 1
OPEN = 2
@@ -45,14 +53,24 @@ gripper_action_map = {
class GamepadTeleop(Teleoperator):
"""
Teleop class to use gamepad inputs for control.
"""Teleoperator that reads a gamepad's analog sticks and buttons via `pygame` (or `hidapi`).
[`~teleoperators.Teleoperator.get_action`] reports the left stick as `delta_x`/`delta_y` and the
right stick's vertical axis as `delta_z`, plus an optional gripper command. See `gamepad_utils.py`'s
`GamepadController` (`pygame`) and `GamepadControllerHID` (`hidapi`) for the exact axis/button
mapping.
"""
config_class = GamepadTeleopConfig
name = "gamepad"
def __init__(self, config: GamepadTeleopConfig):
"""Instantiate the teleoperator.
Args:
config (`GamepadTeleopConfig`):
Configuration for this gamepad teleoperator.
"""
super().__init__(config)
self.config = config
self.robot_type = config.type
@@ -68,6 +86,12 @@ class GamepadTeleop(Teleoperator):
@property
def action_features(self) -> dict:
"""See [`~teleoperators.Teleoperator.action_features`].
Returns:
`dict`: A 3-element (or 4-element if `config.use_gripper` is `True`) `float32` vector named
`delta_x`, `delta_y`, `delta_z`, and optionally `gripper`.
"""
if self.config.use_gripper:
return {
"dtype": "float32",
@@ -83,9 +107,15 @@ class GamepadTeleop(Teleoperator):
@property
def feedback_features(self) -> dict:
"""See [`~teleoperators.Teleoperator.feedback_features`]. `GamepadTeleop` accepts no feedback."""
return {}
def connect(self) -> None:
"""See [`~teleoperators.Teleoperator.connect`].
Starts a `GamepadControllerHID` if `config.hidapi_fallback` is `True`, otherwise a
`GamepadController`.
"""
if self.hidapi_fallback:
from .gamepad_utils import GamepadControllerHID as Gamepad
else:
@@ -96,6 +126,18 @@ class GamepadTeleop(Teleoperator):
@check_if_not_connected
def get_action(self) -> RobotAction:
"""Read the gamepad's current stick positions and gripper button state.
The left analog stick drives `delta_x`/`delta_y`; the right stick's vertical axis drives
`delta_z`. When `config.use_gripper` is `True`, the gripper buttons additionally produce a
`gripper` entry (one of `GripperAction.CLOSE`, `STAY`, or `OPEN`).
Returns:
`dict[str, Any]`: `delta_x`, `delta_y`, `delta_z`, and, if enabled, `gripper`.
Raises:
DeviceNotConnectedError: If [`~teleoperators.Teleoperator.connect`] has not been called.
"""
# Update the controller to get fresh inputs
self.gamepad.update()
@@ -121,16 +163,15 @@ class GamepadTeleop(Teleoperator):
return action_dict
def get_teleop_events(self) -> dict[str, Any]:
"""
Get extra control events from the gamepad such as intervention status,
episode termination, success indicators, etc.
"""Read auxiliary gamepad events used to drive episode control during recording.
Holding the intervention button counts as an active intervention; the success/failure/rerecord
buttons are read once as one-shot signals, then cleared.
Returns:
Dictionary containing:
- is_intervention: bool - Whether human is currently intervening
- terminate_episode: bool - Whether to terminate the current episode
- success: bool - Whether the episode was successful
- rerecord_episode: bool - Whether to rerecord the episode
`dict[TeleopEvents, bool]`: Values for the [`~teleoperators.TeleopEvents`] keys
`IS_INTERVENTION`, `TERMINATE_EPISODE`, `SUCCESS`, and `RERECORD_EPISODE`. All `False` if
[`~teleoperators.Teleoperator.connect`] has not been called yet.
"""
if self.gamepad is None:
return {
@@ -163,32 +204,32 @@ class GamepadTeleop(Teleoperator):
}
def disconnect(self) -> None:
"""Disconnect from the gamepad."""
"""See [`~teleoperators.Teleoperator.disconnect`]. Stops and releases the underlying controller."""
if self.gamepad is not None:
self.gamepad.stop()
self.gamepad = None
@property
def is_connected(self) -> bool:
"""Check if gamepad is connected."""
"""See [`~teleoperators.Teleoperator.is_connected`]."""
return self.gamepad is not None
def calibrate(self) -> None:
"""Calibrate the gamepad."""
"""See [`~teleoperators.Teleoperator.calibrate`]. No-op: the gamepad does not require calibration."""
# No calibration needed for gamepad
pass
def is_calibrated(self) -> bool:
"""Check if gamepad is calibrated."""
"""See [`~teleoperators.Teleoperator.is_calibrated`]. Always `True`: no calibration is required."""
# Gamepad doesn't require calibration
return True
def configure(self) -> None:
"""Configure the gamepad."""
"""See [`~teleoperators.Teleoperator.configure`]. No-op: the gamepad needs no configuration."""
# No additional configuration needed
pass
def send_feedback(self, feedback: dict) -> None:
"""Send feedback to the gamepad."""
"""See [`~teleoperators.Teleoperator.send_feedback`]. No-op: `GamepadTeleop` accepts no feedback."""
# Gamepad doesn't support feedback
pass
@@ -22,11 +22,34 @@ from ..config import TeleoperatorConfig
@TeleoperatorConfig.register_subclass("homunculus_glove")
@dataclass
class HomunculusGloveConfig(TeleoperatorConfig):
"""Configuration for the Homunculus Glove teleoperator.
Args:
port (`str`):
Serial port the glove is connected to, e.g. `/dev/ttyACM0`.
side (`str`):
Which hand the glove is worn on, `"left"` or `"right"`. Selects which joints get their drive
mode inverted so the produced action matches the HopeJR hand convention.
baud_rate (`int`, *optional*, defaults to 115200):
Serial communication speed in bauds.
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
"""
port: str # Port to connect to the glove
side: str # "left" / "right"
baud_rate: int = 115_200
def __post_init__(self):
"""Validate that `side` is one of `"left"` or `"right"`.
Raises:
ValueError: If `side` is neither `"left"` nor `"right"`.
"""
if self.side not in ["right", "left"]:
raise ValueError(self.side)
@@ -34,5 +57,20 @@ class HomunculusGloveConfig(TeleoperatorConfig):
@TeleoperatorConfig.register_subclass("homunculus_arm")
@dataclass
class HomunculusArmConfig(TeleoperatorConfig):
"""Configuration for the Homunculus Arm teleoperator.
Args:
port (`str`):
Serial port the arm is connected to, e.g. `/dev/ttyACM0`.
baud_rate (`int`, *optional*, defaults to 115200):
Serial communication speed in bauds.
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
"""
port: str # Port to connect to the arm
baud_rate: int = 115_200
@@ -37,14 +37,25 @@ logger = logging.getLogger(__name__)
class HomunculusArm(Teleoperator):
"""
Homunculus Arm designed by Hugging Face.
"""Homunculus Arm designed by Hugging Face: a wearable exoskeleton arm read over a serial link.
The arm streams raw encoder values for each joint continuously over a background thread; readings are
smoothed with an exponential moving average before being normalized and returned as an action. It only
produces actions and accepts no feedback.
See [`~teleoperators.Teleoperator`] for the contract every method here implements.
"""
config_class = HomunculusArmConfig
name = "homunculus_arm"
def __init__(self, config: HomunculusArmConfig):
"""Open the serial connection and set up the background reader thread.
Args:
config (`HomunculusArmConfig`):
The teleoperator's configuration. Its `port` determines what is connected.
"""
require_package("pyserial", extra="pyserial-dep", import_name="serial")
super().__init__(config)
self.config = config
@@ -88,19 +99,43 @@ class HomunculusArm(Teleoperator):
@property
def action_features(self) -> dict:
"""The arm's joint positions.
Returns:
`dict`: `"<joint>.pos"` keys mapped to `float`, one per entry in `self.joints`.
"""
return {f"{joint}.pos": float for joint in self.joints}
@property
def feedback_features(self) -> dict:
"""This arm accepts no feedback.
Returns:
`dict`: Always empty.
"""
return {}
@property
def is_connected(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_connected`].
The serial port is open and the background reader thread is alive.
"""
with self.serial_lock:
return self.serial.is_open and self.thread.is_alive()
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
"""Open the serial port, start the background reader thread, and wait for the first reading.
Args:
calibrate (`bool`, *optional*, defaults to `True`):
Whether to run calibration when no calibration file exists yet. Calibration is
interactive and prompts on stdin.
Raises:
TimeoutError: If no state is received from the arm within 2 seconds of starting.
"""
if not self.serial.is_open:
self.serial.open()
self.thread.start()
@@ -116,9 +151,19 @@ class HomunculusArm(Teleoperator):
@property
def is_calibrated(self) -> bool:
"""Whether a calibration file has been saved for this arm.
Returns:
`bool`: `True` if the calibration file exists on disk.
"""
return self.calibration_fpath.is_file()
def calibrate(self) -> None:
"""Interactively record each joint's range of motion and save it as the arm's calibration.
Prompts the operator to move every joint through its full range, then persists the observed
min/max encoder values to the calibration file.
"""
print(
"\nMove all joints through their entire range of motion."
"\nRecording positions. Press ENTER to stop..."
@@ -197,6 +242,7 @@ class HomunculusArm(Teleoperator):
return mins, maxes
def configure(self) -> None:
"""No-op: the arm requires no runtime configuration beyond calibration."""
pass
# TODO(Steven): This function is copy/paste from the `HomunculusGlove` class. Consider moving it to an utility to reduce duplicated code.
@@ -239,9 +285,9 @@ class HomunculusArm(Teleoperator):
def _read(
self, joints: list[str] | None = None, normalize: bool = True, timeout: float = 1
) -> dict[str, int | float]:
"""
Return the most recent (single) values from self.last_d,
optionally applying calibration.
"""Return the most recent values from the reader thread.
Optionally applies calibration.
"""
if not self.new_state_event.wait(timeout=timeout):
raise TimeoutError(f"{self}: Timed out waiting for state after {timeout}s.")
@@ -265,9 +311,9 @@ class HomunculusArm(Teleoperator):
return state
def _read_loop(self):
"""
Continuously read from the serial buffer in its own thread and sends values to the main thread through
a queue.
"""Continuously read from the serial buffer in its own thread.
Sends values to the main thread through a queue.
"""
while not self.stop_event.is_set():
try:
@@ -305,14 +351,28 @@ class HomunculusArm(Teleoperator):
@check_if_not_connected
def get_action(self) -> dict[str, float]:
"""Read the most recent EMA-smoothed, normalized joint positions.
Returns:
`dict[str, float]`: `"<joint>.pos"` keys mapped to their normalized position.
Raises:
TimeoutError: If no new reading arrives from the background thread within 1 second.
"""
joint_positions = self._read()
return {f"{joint}.pos": pos for joint, pos in joint_positions.items()}
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not supported: the arm has no actuators to receive feedback.
Raises:
NotImplementedError: Always.
"""
raise NotImplementedError
@check_if_not_connected
def disconnect(self) -> None:
"""Stop the background reader thread and close the serial port."""
self.stop_event.set()
self.thread.join(timeout=1)
self.serial.close()
@@ -63,14 +63,27 @@ RIGHT_HAND_INVERSIONS = [
class HomunculusGlove(Teleoperator):
"""
Homunculus Glove designed by NepYope & Hugging Face.
"""Homunculus Glove designed by NepYope & Hugging Face: a wearable exoskeleton glove read over a serial link.
The glove streams raw encoder values for each finger joint continuously over a background thread;
readings are smoothed with an exponential moving average, normalized, then remapped from glove joint
names to HopeJR hand joint names via [`~teleoperators.homunculus.homunculus_glove_to_hope_jr_hand`]. It
only produces actions and accepts no feedback.
See [`~teleoperators.Teleoperator`] for the contract every method here implements.
"""
config_class = HomunculusGloveConfig
name = "homunculus_glove"
def __init__(self, config: HomunculusGloveConfig):
"""Open the serial connection and set up the background reader thread.
Args:
config (`HomunculusGloveConfig`):
The teleoperator's configuration. Its `port` determines what is connected and `side`
selects which joints are inverted for the left vs. right hand.
"""
require_package("pyserial", extra="pyserial-dep", import_name="serial")
super().__init__(config)
self.config = config
@@ -114,19 +127,43 @@ class HomunculusGlove(Teleoperator):
@property
def action_features(self) -> dict:
"""The glove's raw per-joint positions, before remapping to HopeJR hand joint names.
Returns:
`dict`: `"<joint>.pos"` keys mapped to `float`, one per entry in `self.joints`.
"""
return {f"{joint}.pos": float for joint in self.joints}
@property
def feedback_features(self) -> dict:
"""This glove accepts no feedback.
Returns:
`dict`: Always empty.
"""
return {}
@property
def is_connected(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_connected`].
The serial port is open and the background reader thread is alive.
"""
with self.serial_lock:
return self.serial.is_open and self.thread.is_alive()
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
"""Open the serial port, start the background reader thread, and wait for the first reading.
Args:
calibrate (`bool`, *optional*, defaults to `True`):
Whether to run calibration when no calibration file exists yet. Calibration is
interactive and prompts on stdin.
Raises:
TimeoutError: If no state is received from the glove within 2 seconds of starting.
"""
if not self.serial.is_open:
self.serial.open()
self.thread.start()
@@ -142,9 +179,19 @@ class HomunculusGlove(Teleoperator):
@property
def is_calibrated(self) -> bool:
"""Whether a calibration file has been saved for this glove.
Returns:
`bool`: `True` if the calibration file exists on disk.
"""
return self.calibration_fpath.is_file()
def calibrate(self) -> None:
"""Interactively record each finger's range of motion and save it as the glove's calibration.
Prompts the operator to move each finger through its full range, one finger at a time, then
persists the observed min/max encoder values to the calibration file.
"""
range_mins, range_maxes = {}, {}
for finger in ["thumb", "index", "middle", "ring", "pinky"]:
print(
@@ -228,6 +275,7 @@ class HomunculusGlove(Teleoperator):
return mins, maxes
def configure(self) -> None:
"""No-op: the glove requires no runtime configuration beyond calibration."""
pass
# TODO(Steven): This function is copy/paste from the `HomunculusArm` class. Consider moving it to an utility to reduce duplicated code.
@@ -271,9 +319,9 @@ class HomunculusGlove(Teleoperator):
def _read(
self, joints: list[str] | None = None, normalize: bool = True, timeout: float = 1
) -> dict[str, int | float]:
"""
Return the most recent (single) values from self.last_d,
optionally applying calibration.
"""Return the most recent values from the reader thread.
Optionally applies calibration.
"""
if not self.new_state_event.wait(timeout=timeout):
raise TimeoutError(f"{self}: Timed out waiting for state after {timeout}s.")
@@ -299,9 +347,9 @@ class HomunculusGlove(Teleoperator):
return state
def _read_loop(self):
"""
Continuously read from the serial buffer in its own thread and sends values to the main thread through
a queue.
"""Continuously read from the serial buffer in its own thread.
Sends values to the main thread through a queue.
"""
while not self.stop_event.is_set():
try:
@@ -331,16 +379,32 @@ class HomunculusGlove(Teleoperator):
@check_if_not_connected
def get_action(self) -> dict[str, float]:
"""Read the most recent EMA-smoothed, normalized joint positions, remapped to HopeJR hand joints.
Returns:
`dict[str, float]`: `"<joint>.pos"` keys, named after the HopeJR hand's joints, mapped to
their normalized position. See
[`~teleoperators.homunculus.homunculus_glove_to_hope_jr_hand`] for the remapping.
Raises:
TimeoutError: If no new reading arrives from the background thread within 1 second.
"""
joint_positions = self._read()
return homunculus_glove_to_hope_jr_hand(
{f"{joint}.pos": pos for joint, pos in joint_positions.items()}
)
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not supported: the glove has no actuators to receive feedback.
Raises:
NotImplementedError: Always.
"""
raise NotImplementedError
@check_if_not_connected
def disconnect(self) -> None:
"""Stop the background reader thread and close the serial port."""
self.stop_event.set()
self.thread.join(timeout=1)
self.serial.close()
@@ -19,14 +19,67 @@ PINKY_SPLAY = 0.5
def get_ulnar_flexion(flexion: float, abduction: float, splay: float):
"""Derive the ulnar-side tendon command for a HopeJR finger from its glove-sensed MCP angles.
The HopeJR hand flexes a finger with a pair of opposing tendons (radial and ulnar) rather than
independent flexion and abduction joints. This blends the glove's flexion and abduction readings for
one MCP joint into the ulnar tendon's share of the motion: an abduction toward the ulnar side pulls
this tendon further, while `splay` sets how much of the abduction reading leaks into it versus pure
flexion.
Args:
flexion (`float`):
MCP flexion reading for the finger, as reported by the glove.
abduction (`float`):
MCP abduction reading for the finger, as reported by the glove. Positive values pull toward
the radial side and are subtracted here.
splay (`float`):
Fraction, in `[0, 1]`, of the tendon command driven by abduction rather than flexion.
Returns:
`float`: The ulnar tendon's target position.
"""
return -abduction * splay + flexion * (1 - splay)
def get_radial_flexion(flexion: float, abduction: float, splay: float):
"""Derive the radial-side tendon command for a HopeJR finger from its glove-sensed MCP angles.
The counterpart to [`get_ulnar_flexion`]: same blend of flexion and abduction, but abduction toward
the radial side adds to this tendon's target instead of subtracting from it.
Args:
flexion (`float`):
MCP flexion reading for the finger, as reported by the glove.
abduction (`float`):
MCP abduction reading for the finger, as reported by the glove. Positive values pull toward
the radial side and are added here.
splay (`float`):
Fraction, in `[0, 1]`, of the tendon command driven by abduction rather than flexion.
Returns:
`float`: The radial tendon's target position.
"""
return abduction * splay + flexion * (1 - splay)
def homunculus_glove_to_hope_jr_hand(glove_action: dict[str, float]) -> dict[str, float]:
"""Translate a Homunculus Glove action into a HopeJR hand action.
The glove reports one flexion and one abduction value per finger's MCP joint, plus a DIP/PIP reading,
while the HopeJR hand is driven by a pair of tendons (radial and ulnar flexors) per finger and a
coupled PIP/DIP joint. This remaps and blends the glove's per-joint keys into the hand's per-tendon
keys via [`get_radial_flexion`] and [`get_ulnar_flexion`]; the thumb, whose joints map one-to-one, is
passed through unchanged.
Args:
glove_action (`dict[str, float]`):
Action produced by [`~teleoperators.homunculus.HomunculusGlove.get_action`], keyed by glove
joint name.
Returns:
`dict[str, float]`: The equivalent action keyed by HopeJR hand joint name.
"""
return {
"thumb_cmc.pos": glove_action["thumb_cmc.pos"],
"thumb_mcp.pos": glove_action["thumb_mcp.pos"],
@@ -23,7 +23,16 @@ from ..config import TeleoperatorConfig
@TeleoperatorConfig.register_subclass("keyboard")
@dataclass
class KeyboardTeleopConfig(TeleoperatorConfig):
"""KeyboardTeleopConfig"""
"""Configuration for the plain keyboard teleoperator.
Args:
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
"""
# TODO(Steven): Consider setting in here the keys that we want to capture/listen
@@ -31,12 +40,17 @@ class KeyboardTeleopConfig(TeleoperatorConfig):
@TeleoperatorConfig.register_subclass("keyboard_ee")
@dataclass
class KeyboardEndEffectorTeleopConfig(KeyboardTeleopConfig):
"""Configuration for keyboard end-effector teleoperator.
"""Configuration for controlling a robot end-effector with keyboard inputs.
Used for controlling robot end-effectors with keyboard inputs.
**Attributes**:
- **use_gripper** (`bool`) -- Whether to include gripper control in actions
Args:
use_gripper (`bool`, *optional*, defaults to `True`):
Whether to include a `gripper` entry in the produced actions.
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
"""
use_gripper: bool = True
@@ -45,18 +59,29 @@ class KeyboardEndEffectorTeleopConfig(KeyboardTeleopConfig):
@TeleoperatorConfig.register_subclass("keyboard_rover")
@dataclass
class KeyboardRoverTeleopConfig(TeleoperatorConfig):
"""Configuration for keyboard rover teleoperator.
"""Configuration for the WASD-style keyboard teleoperator for mobile robots like EarthRover Mini Plus.
Used for controlling mobile robots like EarthRover Mini Plus with WASD controls.
**Attributes**:
- **linear_speed** (`float`) -- Default linear velocity magnitude (-1 to 1 range for SDK robots)
- **angular_speed** (`float`) -- Default angular velocity magnitude (-1 to 1 range for SDK robots)
- **speed_increment** (`float`) -- Amount to increase/decrease speed with +/- keys
- **turn_assist_ratio** (`float`) -- Forward motion multiplier when turning with A/D keys (0.0-1.0)
- **angular_speed_ratio** (`float`) -- Ratio of angular to linear speed for synchronized adjustments
- **min_linear_speed** (`float`) -- Minimum linear speed when decreasing (prevents zero speed)
- **min_angular_speed** (`float`) -- Minimum angular speed when decreasing (prevents zero speed)
Args:
linear_speed (`float`, *optional*, defaults to 1.0):
Initial linear velocity magnitude (-1 to 1 range for SDK robots).
angular_speed (`float`, *optional*, defaults to 1.0):
Initial angular velocity magnitude (-1 to 1 range for SDK robots).
speed_increment (`float`, *optional*, defaults to 0.1):
Amount `current_linear_speed` changes by on each `+`/`-` key press.
turn_assist_ratio (`float`, *optional*, defaults to 0.3):
Forward-motion multiplier applied when turning with `a`/`d` while otherwise stationary.
angular_speed_ratio (`float`, *optional*, defaults to 0.6):
Ratio of angular to linear speed increment, so both scale together on `+`/`-`.
min_linear_speed (`float`, *optional*, defaults to 0.1):
Floor for `current_linear_speed` when decreasing it.
min_angular_speed (`float`, *optional*, defaults to 0.05):
Floor for `current_angular_speed` when decreasing it.
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
"""
linear_speed: float = 1.0
@@ -43,14 +43,28 @@ if PYNPUT_AVAILABLE:
class KeyboardTeleop(Teleoperator):
"""
Teleop class to use keyboard inputs for control.
"""Teleoperator that reads raw keyboard key states via `pynput` for manual control.
[`~teleoperators.Teleoperator.get_action`] reports every key currently held down. Requires an
interactive desktop session capable of capturing global key events — an X11 session (Linux), a
Windows desktop, or macOS with Accessibility / Input Monitoring permission granted. On Wayland or a
headless machine, [`~teleoperators.Teleoperator.connect`] logs a warning and the teleoperator produces
no actions.
"""
config_class = KeyboardTeleopConfig
name = "keyboard"
def __init__(self, config: KeyboardTeleopConfig):
"""Instantiate the teleoperator.
Args:
config (`KeyboardTeleopConfig`):
Configuration for this keyboard teleoperator.
Raises:
ImportError: If `pynput` is not installed.
"""
require_package("pynput", extra="pynput-dep")
super().__init__(config)
self.config = config
@@ -63,6 +77,11 @@ class KeyboardTeleop(Teleoperator):
@property
def action_features(self) -> dict:
"""See [`~teleoperators.Teleoperator.action_features`].
Returns:
`dict`: Motor count and names taken from `self.arm`.
"""
return {
"dtype": "float32",
"shape": (len(self.arm),),
@@ -71,18 +90,26 @@ class KeyboardTeleop(Teleoperator):
@property
def feedback_features(self) -> dict:
"""See [`~teleoperators.Teleoperator.feedback_features`]. `KeyboardTeleop` accepts no feedback."""
return {}
@property
def is_connected(self) -> bool:
"""See [`~teleoperators.Teleoperator.is_connected`]."""
return PYNPUT_AVAILABLE and isinstance(self.listener, keyboard.Listener) and self.listener.is_alive()
@property
def is_calibrated(self) -> bool:
"""See [`~teleoperators.Teleoperator.is_calibrated`]. Keyboard input does not require calibration."""
pass
@check_if_already_connected
def connect(self) -> None:
"""See [`~teleoperators.Teleoperator.connect`].
Starts a `pynput` keyboard listener if the current session can capture key events; otherwise logs
a warning and leaves the teleoperator producing no actions.
"""
if PYNPUT_AVAILABLE and pynput_can_capture():
logging.info("pynput is available - enabling local keyboard listener.")
self.listener = keyboard.Listener(
@@ -101,6 +128,7 @@ class KeyboardTeleop(Teleoperator):
self.listener = None
def calibrate(self) -> None:
"""See [`~teleoperators.Teleoperator.calibrate`]. No-op: keyboard input does not require calibration."""
pass
def _on_press(self, key):
@@ -123,10 +151,20 @@ class KeyboardTeleop(Teleoperator):
self.current_pressed[key_char] = is_pressed
def configure(self):
"""See [`~teleoperators.Teleoperator.configure`]. No-op: keyboard input needs no configuration."""
pass
@check_if_not_connected
def get_action(self) -> RobotAction:
"""Read the keys currently held down.
Returns:
`dict[str, Any]`: One entry per key character currently pressed, each mapped to `None`. An
empty dict means no key is currently held.
Raises:
DeviceNotConnectedError: If [`~teleoperators.Teleoperator.connect`] has not been called.
"""
before_read_t = time.perf_counter()
self._drain_pressed_keys()
@@ -138,30 +176,45 @@ class KeyboardTeleop(Teleoperator):
return dict.fromkeys(action, None)
def send_feedback(self, feedback: dict[str, Any]) -> None:
"""See [`~teleoperators.Teleoperator.send_feedback`]. No-op: `KeyboardTeleop` accepts no feedback."""
pass
@check_if_not_connected
def disconnect(self) -> None:
"""See [`~teleoperators.Teleoperator.disconnect`]. Stops the keyboard listener, if one is running."""
if self.listener is not None:
self.listener.stop()
class KeyboardEndEffectorTeleop(KeyboardTeleop):
"""
Teleop class to use keyboard inputs for end effector control.
Designed to be used with the `So100FollowerEndEffector` robot.
"""Keyboard teleoperator for end-effector (Cartesian delta) control.
Arrow keys and shift map to `delta_x`/`delta_y`/`delta_z`; `ctrl_l`/`ctrl_r` map to the gripper.
Designed for use with the `So100FollowerEndEffector` robot.
"""
config_class = KeyboardEndEffectorTeleopConfig
name = "keyboard_ee"
def __init__(self, config: KeyboardEndEffectorTeleopConfig):
"""Instantiate the teleoperator.
Args:
config (`KeyboardEndEffectorTeleopConfig`):
Configuration for this keyboard end-effector teleoperator.
"""
super().__init__(config)
self.config = config
self.misc_keys_queue = Queue()
@property
def action_features(self) -> dict:
"""See [`~teleoperators.Teleoperator.action_features`].
Returns:
`dict`: A 3-element (or 4-element if `config.use_gripper` is `True`) `float32` vector named
`delta_x`, `delta_y`, `delta_z`, and optionally `gripper`.
"""
if self.config.use_gripper:
return {
"dtype": "float32",
@@ -177,6 +230,19 @@ class KeyboardEndEffectorTeleop(KeyboardTeleop):
@check_if_not_connected
def get_action(self) -> RobotAction:
"""Translate held-down keys into an end-effector Cartesian delta.
Arrow keys drive `delta_x`/`delta_y`; `shift`/`shift_r` drive `delta_z`. `ctrl_r` opens the
gripper and `ctrl_l` closes it (only present when `config.use_gripper` is `True`); any other
pressed key is queued for [`~teleoperators.keyboard.KeyboardEndEffectorTeleop.get_teleop_events`]
instead of affecting the action.
Returns:
`dict[str, Any]`: `delta_x`, `delta_y`, `delta_z`, and, if enabled, `gripper`.
Raises:
DeviceNotConnectedError: If [`~teleoperators.Teleoperator.connect`] has not been called.
"""
self._drain_pressed_keys()
delta_x = 0.0
delta_y = 0.0
@@ -220,22 +286,15 @@ class KeyboardEndEffectorTeleop(KeyboardTeleop):
return action_dict
def get_teleop_events(self) -> dict[str, Any]:
"""
Get extra control events from the keyboard such as intervention status,
episode termination, success indicators, etc.
"""Read auxiliary keyboard events used to drive episode control during recording.
Keyboard mappings:
- Any movement keys pressed = intervention active
- 's' key = success (terminate episode successfully)
- 'r' key = rerecord episode (terminate and rerecord)
- 'q' key = quit episode (terminate without success)
Any of the movement/gripper keys held down counts as an active intervention. `s`, `r`, and `q`
are read once as one-shot signals for success, rerecord, and quit respectively; reading this
method clears the currently tracked key state.
Returns:
Dictionary containing:
- is_intervention: bool - Whether human is currently intervening
- terminate_episode: bool - Whether to terminate the current episode
- success: bool - Whether the episode was successful
- rerecord_episode: bool - Whether to rerecord the episode
`dict[TeleopEvents, bool]`: Values for the [`~teleoperators.TeleopEvents`] keys
`IS_INTERVENTION`, `TERMINATE_EPISODE`, `SUCCESS`, and `RERECORD_EPISODE`.
"""
if not self.is_connected:
return {
@@ -286,49 +345,24 @@ class KeyboardEndEffectorTeleop(KeyboardTeleop):
class KeyboardRoverTeleop(KeyboardTeleop):
"""
Keyboard teleoperator for mobile robots like EarthRover Mini Plus.
"""Keyboard teleoperator for mobile robots such as EarthRover Mini Plus.
Provides intuitive WASD-style controls for driving a mobile robot:
- Linear movement (forward/backward)
- Angular movement (turning/rotation)
- Speed adjustment
- Emergency stop
Keyboard Controls:
Movement:
- W: Move forward
- S: Move backward
- A: Turn left (with forward motion)
- D: Turn right (with forward motion)
- Q: Rotate left in place
- E: Rotate right in place
- X: Emergency stop
Speed Control:
- +/=: Increase speed
- -: Decrease speed
System:
- ESC: Disconnect teleoperator
Provides WASD-style driving controls: `w`/`s` drive forward/backward, `a`/`d` turn (with a forward
motion assist), `q`/`e` rotate in place, `x` is an emergency stop, and `+`/`-` adjust speed. `ESC`
disconnects the teleoperator.
**Attributes**:
- **config** -- Teleoperator configuration
- **current_linear_speed** -- Current linear velocity magnitude
- **current_angular_speed** -- Current angular velocity magnitude
- **current_linear_speed** (`float`) -- Current linear velocity magnitude, adjustable at runtime
with `+`/`-`.
- **current_angular_speed** (`float`) -- Current angular velocity magnitude, adjustable at
runtime with `+`/`-`.
Example:
```python
from lerobot.teleoperators.keyboard import KeyboardRoverTeleop, KeyboardRoverTeleopConfig
teleop = KeyboardRoverTeleop(
KeyboardRoverTeleopConfig(linear_speed=1.0, angular_speed=1.0, speed_increment=0.1)
)
teleop.connect()
while teleop.is_connected:
action = teleop.get_action()
robot.send_action(action)
>>> from lerobot.teleoperators.keyboard import KeyboardRoverTeleop, KeyboardRoverTeleopConfig
>>> teleop = KeyboardRoverTeleop(KeyboardRoverTeleopConfig(linear_speed=1.0)) # doctest: +SKIP
>>> teleop.connect() # doctest: +SKIP
>>> teleop.get_action() # doctest: +SKIP
```
"""
@@ -336,6 +370,12 @@ class KeyboardRoverTeleop(KeyboardTeleop):
name = "keyboard_rover"
def __init__(self, config: KeyboardRoverTeleopConfig):
"""Instantiate the teleoperator.
Args:
config (`KeyboardRoverTeleopConfig`):
Configuration for this keyboard rover teleoperator.
"""
super().__init__(config)
# Add rover-specific speed settings
self.current_linear_speed = config.linear_speed
@@ -343,7 +383,11 @@ class KeyboardRoverTeleop(KeyboardTeleop):
@property
def action_features(self) -> dict:
"""Return action format for rover (linear and angular velocities)."""
"""See [`~teleoperators.Teleoperator.action_features`].
Returns:
`dict`: `linear_velocity` and `angular_velocity`, each mapped to `float`.
"""
return {
"linear_velocity": float,
"angular_velocity": float,
@@ -351,11 +395,11 @@ class KeyboardRoverTeleop(KeyboardTeleop):
@property
def is_calibrated(self) -> bool:
"""Rover teleop doesn't require calibration."""
"""See [`~teleoperators.Teleoperator.is_calibrated`]. Rover teleop does not require calibration."""
return True
def _drain_pressed_keys(self):
"""Update current_pressed state from event queue without clearing held keys"""
"""Update current_pressed state from event queue without clearing held keys."""
while not self.event_queue.empty():
key_char, is_pressed = self.event_queue.get_nowait()
if is_pressed:
@@ -366,11 +410,18 @@ class KeyboardRoverTeleop(KeyboardTeleop):
@check_if_not_connected
def get_action(self) -> RobotAction:
"""
Get the current action based on pressed keys.
"""Translate held-down WASD-style keys into linear and angular rover velocities.
`w`/`s` set the linear velocity; `a`/`d` turn while adding a forward-motion assist
(`config.turn_assist_ratio`) when not already moving; `q`/`e` rotate in place; `x` stops both
axes. `+`/`-` adjust `current_linear_speed` and `current_angular_speed` in place, clamped to
`config.min_linear_speed` / `config.min_angular_speed`.
Returns:
RobotAction with 'linear_velocity' and 'angular_velocity' keys.
`dict[str, float]`: `linear_velocity` and `angular_velocity`.
Raises:
DeviceNotConnectedError: If [`~teleoperators.Teleoperator.connect`] has not been called.
"""
before_read_t = time.perf_counter()
@@ -22,6 +22,28 @@ from ..config import TeleoperatorConfig
@TeleoperatorConfig.register_subclass("koch_leader")
@dataclass
class KochLeaderConfig(TeleoperatorConfig):
"""Configuration for the Koch leader arm.
Args:
port (`str`):
Serial port the arm is connected to, e.g. `/dev/ttyACM0` on Linux or `COM3` on Windows. Run
`lerobot-find-port` to identify it.
gripper_open_pos (`float`, *optional*, defaults to 50.0):
Goal position written to the gripper motor, held under current-based position control so the
gripper springs back to this position when released, letting it be used as a physical trigger.
id (`str`, *optional*):
Identifier for this particular arm; also names its calibration file.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
Example:
```python
>>> from lerobot.teleoperators.koch_leader import KochLeader, KochLeaderConfig
>>> config = KochLeaderConfig(port="/dev/ttyACM0") # doctest: +SKIP
>>> teleop = KochLeader(config) # doctest: +SKIP
```
"""
# Port to connect to the arm
port: str
@@ -32,16 +32,34 @@ logger = logging.getLogger(__name__)
class KochLeader(Teleoperator):
"""
"""The Koch leader arm, in either of its two revisions, held by an operator to teleoperate a follower arm.
- [Koch v1.0](https://github.com/AlexanderKoch-Koch/low_cost_robot), with and without the wrist-to-elbow
expansion, developed by Alexander Koch from [Tau Robotics](https://tau-robotics.com)
- [Koch v1.1](https://github.com/jess-moss/koch-v1-1) developed by Jess Moss
expansion, developed by Alexander Koch from [Tau Robotics](https://tau-robotics.com).
- [Koch v1.1](https://github.com/jess-moss/koch-v1-1), developed by Jess Moss.
Actions are keyed `"<motor>.pos"`. See [`~teleoperators.Teleoperator`] for the contract every method
here implements.
Example:
```python
>>> from lerobot.teleoperators.koch_leader import KochLeader, KochLeaderConfig
>>> teleop = KochLeader(KochLeaderConfig(port="/dev/ttyACM0")) # doctest: +SKIP
>>> with teleop: # doctest: +SKIP
... action = teleop.get_action()
```
"""
config_class = KochLeaderConfig
name = "koch_leader"
def __init__(self, config: KochLeaderConfig):
"""Build the teleoperator from its configuration.
Args:
config (`KochLeaderConfig`):
The teleoperator's configuration. Its `port` determines what is connected.
"""
super().__init__(config)
self.config = config
self.bus = DynamixelMotorsBus(
@@ -59,18 +77,42 @@ class KochLeader(Teleoperator):
@property
def action_features(self) -> dict[str, type]:
"""The arm's joint positions.
Returns:
`dict[str, type]`: `"<motor>.pos"` keys mapped to `float`.
"""
return {f"{motor}.pos": float for motor in self.bus.motors}
@property
def feedback_features(self) -> dict[str, type]:
"""Same as [`~teleoperators.Teleoperator.feedback_features`].
This arm does not support feedback; [`~teleoperators.koch_leader.KochLeader.send_feedback`] always
raises `NotImplementedError`.
Returns:
`dict[str, type]`: Always empty.
"""
return {}
@property
def is_connected(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_connected`]."""
return self.bus.is_connected
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
"""Connect the motor bus, calibrating and configuring the arm.
Args:
calibrate (`bool`, *optional*, defaults to `True`):
Whether to run calibration when the motors disagree with the calibration file, or no file
exists yet. Calibration is interactive and prompts on stdin.
Raises:
DeviceAlreadyConnectedError: If the teleoperator is already connected.
"""
self.bus.connect()
if not self.is_calibrated and calibrate:
logger.info(
@@ -83,9 +125,16 @@ class KochLeader(Teleoperator):
@property
def is_calibrated(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_calibrated`]."""
return self.bus.is_calibrated
def calibrate(self) -> None:
"""Calibrate the arm, writing the result to the motors and the calibration file.
This is interactive: it prompts on stdin to reuse an existing calibration file, and otherwise asks
you to move the arm to its middle position and then through each joint's full range. The
`elbow_flex` motor is inverted, and `shoulder_pan` and `wrist_roll` are treated as full-turn joints.
"""
self.bus.disable_torque()
if self.calibration:
# Calibration file exists, ask user whether to use it or run new calibration
@@ -132,6 +181,12 @@ class KochLeader(Teleoperator):
logger.info(f"Calibration saved to {self.calibration_fpath}")
def configure(self) -> None:
"""Write the operating modes to every motor, including the gripper's spring-back trigger behavior.
All motors except the gripper are set to extended position mode. The gripper is set to
current-based position control and driven to `gripper_open_pos`, with torque enabled, so it springs
back to that position when released and can be used as a physical trigger.
"""
self.bus.disable_torque()
self.bus.configure_motors()
for motor in self.bus.motors:
@@ -154,6 +209,11 @@ class KochLeader(Teleoperator):
self.bus.write("Goal_Position", "gripper", self.config.gripper_open_pos)
def setup_motors(self) -> None:
"""Assign each motor its bus ID, one at a time.
Run this once when building an arm. It is interactive: it prompts you to connect the controller
board to a single motor at a time, working from the gripper back to the base.
"""
for motor in reversed(self.bus.motors):
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
self.bus.setup_motor(motor)
@@ -161,6 +221,14 @@ class KochLeader(Teleoperator):
@check_if_not_connected
def get_action(self) -> dict[str, float]:
"""Same as [`~teleoperators.Teleoperator.get_action`].
Returns:
`dict[str, float]`: `"<motor>.pos"` keys mapped to the arm's current joint positions.
Raises:
DeviceNotConnectedError: If the teleoperator is not connected.
"""
start = time.perf_counter()
action = self.bus.sync_read("Present_Position")
action = {f"{motor}.pos": val for motor, val in action.items()}
@@ -169,10 +237,20 @@ class KochLeader(Teleoperator):
return action
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not implemented for this arm.
Raises:
NotImplementedError: Always. This arm does not support force feedback.
"""
# TODO(rcadene, aliberts): Implement force feedback
raise NotImplementedError
@check_if_not_connected
def disconnect(self) -> None:
"""Same as [`~teleoperators.Teleoperator.disconnect`].
Raises:
DeviceNotConnectedError: If the teleoperator is not connected.
"""
self.bus.disconnect()
logger.info(f"{self} disconnected.")
@@ -22,6 +22,28 @@ from ..config import TeleoperatorConfig
@TeleoperatorConfig.register_subclass("omx_leader")
@dataclass
class OmxLeaderConfig(TeleoperatorConfig):
"""Configuration for the OMX leader arm.
Args:
port (`str`):
Serial port the arm is connected to, e.g. `/dev/ttyACM0` on Linux or `COM3` on Windows. Run
`lerobot-find-port` to identify it.
gripper_open_pos (`float`, *optional*, defaults to 60.0):
Goal position written to the gripper motor, held under current-based position control so the
gripper springs back to this position when released, letting it be used as a physical trigger.
id (`str`, *optional*):
Identifier for this particular arm; also names its calibration file.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
Example:
```python
>>> from lerobot.teleoperators.omx_leader import OmxLeader, OmxLeaderConfig
>>> config = OmxLeaderConfig(port="/dev/ttyACM0") # doctest: +SKIP
>>> teleop = OmxLeader(config) # doctest: +SKIP
```
"""
# Port to connect to the arm
port: str
@@ -32,15 +32,33 @@ logger = logging.getLogger(__name__)
class OmxLeader(Teleoperator):
"""
- [OMX](https://github.com/ROBOTIS-GIT/open_manipulator),
expansion, developed by Woojin Wie and Junha Cha from [ROBOTIS](https://ai.robotis.com/)
"""The OMX leader arm, held by an operator to teleoperate a follower arm.
[OMX](https://github.com/ROBOTIS-GIT/open_manipulator), developed by Woojin Wie and Junha Cha from
[ROBOTIS](https://ai.robotis.com/).
Actions are keyed `"<motor>.pos"`. See [`~teleoperators.Teleoperator`] for the contract every method
here implements.
Example:
```python
>>> from lerobot.teleoperators.omx_leader import OmxLeader, OmxLeaderConfig
>>> teleop = OmxLeader(OmxLeaderConfig(port="/dev/ttyACM0")) # doctest: +SKIP
>>> with teleop: # doctest: +SKIP
... action = teleop.get_action()
```
"""
config_class = OmxLeaderConfig
name = "omx_leader"
def __init__(self, config: OmxLeaderConfig):
"""Build the teleoperator from its configuration.
Args:
config (`OmxLeaderConfig`):
The teleoperator's configuration. Its `port` determines what is connected.
"""
super().__init__(config)
self.config = config
self.bus = DynamixelMotorsBus(
@@ -58,18 +76,42 @@ class OmxLeader(Teleoperator):
@property
def action_features(self) -> dict[str, type]:
"""The arm's joint positions.
Returns:
`dict[str, type]`: `"<motor>.pos"` keys mapped to `float`.
"""
return {f"{motor}.pos": float for motor in self.bus.motors}
@property
def feedback_features(self) -> dict[str, type]:
"""Same as [`~teleoperators.Teleoperator.feedback_features`].
This arm does not support feedback; [`~teleoperators.omx_leader.OmxLeader.send_feedback`] always
raises `NotImplementedError`.
Returns:
`dict[str, type]`: Always empty.
"""
return {}
@property
def is_connected(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_connected`]."""
return self.bus.is_connected
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
"""Connect the motor bus, calibrating and configuring the arm.
Args:
calibrate (`bool`, *optional*, defaults to `True`):
Whether to write the factory default calibration when the motors disagree with the
calibration file, or no file exists yet.
Raises:
DeviceAlreadyConnectedError: If the teleoperator is already connected.
"""
self.bus.connect()
if not self.is_calibrated and calibrate:
logger.info(
@@ -82,9 +124,15 @@ class OmxLeader(Teleoperator):
@property
def is_calibrated(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_calibrated`]."""
return self.bus.is_calibrated
def calibrate(self) -> None:
"""Write the factory default calibration to the motors and the calibration file.
Unlike other SO/Koch-family arms, this is not interactive: the OMX arm's homing offsets and ranges
of motion are fixed factory defaults, so no manual positioning is required.
"""
self.bus.disable_torque()
logger.info(f"\nUsing factory default calibration values for {self}")
logger.info(f"\nWriting default configuration of {self} to the motors")
@@ -113,6 +161,13 @@ class OmxLeader(Teleoperator):
logger.info(f"Calibration saved to {self.calibration_fpath}")
def configure(self) -> None:
"""Write the operating and drive modes to every motor, including the gripper's spring-back trigger.
All motors except the gripper are set to extended position mode with a non-inverted drive mode. The
gripper's drive mode is inverted, and it is set to current-based position control with a reduced
current limit and driven to `gripper_open_pos`, with torque enabled, so it springs back to that
position when released and can be used as a physical trigger.
"""
self.bus.disable_torque()
self.bus.configure_motors()
for motor in self.bus.motors:
@@ -143,6 +198,11 @@ class OmxLeader(Teleoperator):
self.bus.write("Goal_Position", "gripper", self.config.gripper_open_pos)
def setup_motors(self) -> None:
"""Assign each motor its bus ID, one at a time.
Run this once when building an arm. It is interactive: it prompts you to connect the controller
board to a single motor at a time, working from the gripper back to the base.
"""
for motor in reversed(self.bus.motors):
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
self.bus.setup_motor(motor)
@@ -150,6 +210,14 @@ class OmxLeader(Teleoperator):
@check_if_not_connected
def get_action(self) -> dict[str, float]:
"""Same as [`~teleoperators.Teleoperator.get_action`].
Returns:
`dict[str, float]`: `"<motor>.pos"` keys mapped to the arm's current joint positions.
Raises:
DeviceNotConnectedError: If the teleoperator is not connected.
"""
start = time.perf_counter()
action = self.bus.sync_read("Present_Position")
action = {f"{motor}.pos": val for motor, val in action.items()}
@@ -158,10 +226,20 @@ class OmxLeader(Teleoperator):
return action
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not implemented for this arm.
Raises:
NotImplementedError: Always. This arm does not support force feedback.
"""
# TODO(rcadene, aliberts): Implement force feedback
raise NotImplementedError
@check_if_not_connected
def disconnect(self) -> None:
"""Same as [`~teleoperators.Teleoperator.disconnect`].
Raises:
DeviceNotConnectedError: If the teleoperator is not connected.
"""
self.bus.disconnect()
logger.info(f"{self} disconnected.")
@@ -76,4 +76,41 @@ class OpenArmLeaderConfigBase:
@TeleoperatorConfig.register_subclass("openarm_leader")
@dataclass
class OpenArmLeaderConfig(TeleoperatorConfig, OpenArmLeaderConfigBase):
"""Configuration for the OpenArm leader/teleoperator arm (CAN bus, Damiao motors).
Args:
port (`str`):
CAN interface the arm is connected to, e.g. `"can0"` on Linux.
can_interface (`str`, *optional*, defaults to `"socketcan"`):
CAN backend type: `"socketcan"` (Linux), `"slcan"` (serial), or `"auto"` (auto-detect).
use_can_fd (`bool`, *optional*, defaults to `True`):
Whether to use CAN FD, which OpenArm uses by default.
can_bitrate (`int`, *optional*, defaults to 1000000):
Nominal CAN bus bitrate, in bits per second.
can_data_bitrate (`int`, *optional*, defaults to 5000000):
CAN FD data-phase bitrate, in bits per second. Only used when `use_can_fd` is `True`.
motor_config (`dict[str, tuple[int, int, str]]`, *optional*):
Maps motor name to `(send_can_id, recv_can_id, motor_type)`. Defaults to the standard 7-DOF
plus gripper OpenArm layout, using DM8009 (shoulder), DM4340 (shoulder rotation, elbow), and
DM4310 (wrist, gripper) Damiao motors.
manual_control (`bool`, *optional*, defaults to `True`):
Whether motors have torque disabled for manual movement. Required for a leader arm that is
moved by hand.
use_velocity_and_torque (`bool`, *optional*, defaults to `False`):
Whether to expose `.vel` and `.torque` per motor in [`~teleoperators.Teleoperator.action_features`],
in addition to `.pos`.
position_kp (`list[float]`, *optional*):
Per-joint position gain, used for MIT torque control when `manual_control` is `False`.
Defaults to the standard 8-value OpenArm gain set (one value per joint, plus gripper).
position_kd (`list[float]`, *optional*):
Per-joint velocity gain, used for MIT torque control when `manual_control` is `False`.
Defaults to the standard 8-value OpenArm damping set (one value per joint, plus gripper).
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
"""
pass
@@ -30,17 +30,33 @@ logger = logging.getLogger(__name__)
class OpenArmLeader(Teleoperator):
"""
OpenArm Leader/Teleoperator Arm with Damiao motors.
"""OpenArm Leader/Teleoperator Arm with Damiao motors.
This teleoperator uses CAN bus communication to read positions from
Damiao motors that are manually moved (torque disabled).
This teleoperator uses CAN bus communication to read positions from Damiao motors that are manually
moved (torque disabled). For the bimanual setup, see [`~teleoperators.bi_openarm_leader.BiOpenArmLeader`], which composes
two of these.
Example:
```python
>>> from lerobot.teleoperators.openarm_leader import OpenArmLeader, OpenArmLeaderConfig
>>> config = OpenArmLeaderConfig(port="can0")
>>> leader = OpenArmLeader(config) # doctest: +SKIP
>>> leader.connect() # doctest: +SKIP
>>> action = leader.get_action() # doctest: +SKIP
```
"""
config_class = OpenArmLeaderConfig
name = "openarm_leader"
def __init__(self, config: OpenArmLeaderConfig):
"""Build the teleoperator from its configuration.
Args:
config (`OpenArmLeaderConfig`):
The teleoperator's configuration. Its `port` and `motor_config` determine what is
connected and how the CAN bus is laid out.
"""
super().__init__(config)
self.config = config
@@ -66,7 +82,11 @@ class OpenArmLeader(Teleoperator):
@property
def action_features(self) -> dict[str, type]:
"""Features produced by this teleoperator."""
"""See [`~teleoperators.Teleoperator.action_features`].
Always includes `.pos` per motor; also includes `.vel` and `.torque` per motor when
`config.use_velocity_and_torque` is `True`.
"""
features: dict[str, type] = {}
for motor in self.bus.motors:
features[f"{motor}.pos"] = float
@@ -77,23 +97,23 @@ class OpenArmLeader(Teleoperator):
@property
def feedback_features(self) -> dict[str, type]:
"""Feedback features (not implemented for OpenArms)."""
"""See [`~teleoperators.Teleoperator.feedback_features`].
Always empty: feedback is not implemented for the OpenArm leader.
"""
return {}
@property
def is_connected(self) -> bool:
"""Check if teleoperator is connected."""
"""See [`~teleoperators.Teleoperator.is_connected`]."""
return self.bus.is_connected
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
"""
Connect to the teleoperator.
"""See [`~teleoperators.Teleoperator.connect`].
For manual control, we disable torque after connecting so the
arm can be moved by hand.
For manual control, torque is disabled after connecting so the arm can be moved by hand.
"""
# Connect to CAN bus
logger.info(f"Connecting arm on {self.config.port}...")
self.bus.connect()
@@ -114,12 +134,11 @@ class OpenArmLeader(Teleoperator):
@property
def is_calibrated(self) -> bool:
"""Check if teleoperator is calibrated."""
"""See [`~teleoperators.Teleoperator.is_calibrated`]."""
return self.bus.is_calibrated
def calibrate(self) -> None:
"""
Run calibration procedure for OpenArms leader.
"""See [`~teleoperators.Teleoperator.calibrate`].
The calibration procedure:
1. Disable torque (if not already disabled)
@@ -170,26 +189,29 @@ class OpenArmLeader(Teleoperator):
print(f"Calibration saved to {self.calibration_fpath}")
def configure(self) -> None:
"""
Configure motors for manual teleoperation.
"""See [`~teleoperators.Teleoperator.configure`].
For manual control, we disable torque so the arm can be moved by hand.
For manual control, torque is disabled so the arm can be moved by hand; otherwise the motors are
configured for MIT torque control.
"""
return self.bus.disable_torque() if self.config.manual_control else self.bus.configure_motors()
def setup_motors(self) -> None:
"""Not supported: raises `NotImplementedError`.
Motor ID configuration for CAN motors is typically done via manufacturer tools rather than through
LeRobot.
Raises:
NotImplementedError: Always.
"""
raise NotImplementedError(
"Motor ID configuration is typically done via manufacturer tools for CAN motors."
)
@check_if_not_connected
def get_action(self) -> RobotAction:
"""
Get current action from the leader arm.
This is the main method for teleoperators - it reads the current state
of the leader arm and returns it as an action that can be sent to a follower.
"""See [`~teleoperators.Teleoperator.get_action`].
Reads all motor states (pos/vel/torque) in one CAN refresh cycle.
"""
@@ -212,12 +234,20 @@ class OpenArmLeader(Teleoperator):
return action_dict
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not supported: raises `NotImplementedError`.
Args:
feedback (`dict[str, float]`):
Unused.
Raises:
NotImplementedError: Always.
"""
raise NotImplementedError("Feedback is not yet implemented for OpenArm leader.")
@check_if_not_connected
def disconnect(self) -> None:
"""Disconnect from teleoperator."""
"""See [`~teleoperators.Teleoperator.disconnect`]."""
# Disconnect CAN bus
# For manual control, ensure torque is disabled before disconnecting
self.bus.disconnect(disable_torque=self.config.manual_control)
@@ -36,4 +36,22 @@ class OpenArmMiniConfigBase:
@TeleoperatorConfig.register_subclass("openarm_mini")
@dataclass
class OpenArmMiniConfig(TeleoperatorConfig, OpenArmMiniConfigBase):
"""Configuration for the OpenArm Mini teleoperator (Feetech STS3215, 7DOF + gripper).
Args:
port (`str`):
Serial port the Feetech bus is connected to, e.g. `/dev/ttyUSB0`.
side (`str`, *optional*):
Which side of a bimanual pair this arm is: `"left"` or `"right"`. Controls per-joint
direction flips applied during readout. `None` disables flipping.
use_degrees (`bool`, *optional*, defaults to `True`):
Keep `True` for backward compatibility with existing policies and datasets.
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
"""
pass
@@ -46,13 +46,32 @@ GRIPPER_TELEOP_TO_DEGREES = -0.65
class OpenArmMini(Teleoperator):
"""OpenArm Mini single-arm teleoperator (Feetech STS3215, 7DOF + gripper).
For the bimanual setup, see :class:`BiOpenArmMini` which composes two of these.
For the bimanual setup, see [`~teleoperators.bi_openarm_mini.BiOpenArmMini`], which composes two of these.
Example:
```python
>>> from lerobot.teleoperators.openarm_mini import OpenArmMini, OpenArmMiniConfig
>>> config = OpenArmMiniConfig(port="/dev/ttyUSB0")
>>> teleop = OpenArmMini(config) # doctest: +SKIP
>>> teleop.connect() # doctest: +SKIP
>>> action = teleop.get_action() # doctest: +SKIP
```
"""
config_class = OpenArmMiniConfig
name = "openarm_mini"
def __init__(self, config: OpenArmMiniConfig):
"""Build the teleoperator from its configuration.
Args:
config (`OpenArmMiniConfig`):
The teleoperator's configuration. Its `port` and `side` determine what is connected and
which per-joint direction flips are applied.
Raises:
ValueError: If `config.side` is not `"left"`, `"right"`, or `None`.
"""
super().__init__(config)
self.config = config
@@ -80,18 +99,25 @@ class OpenArmMini(Teleoperator):
@property
def action_features(self) -> dict[str, type]:
"""See [`~teleoperators.Teleoperator.action_features`]. One `.pos` entry per motor."""
return {f"{motor}.pos": float for motor in self.bus.motors}
@property
def feedback_features(self) -> dict[str, type]:
"""See [`~teleoperators.Teleoperator.feedback_features`].
Same shape as [`~teleoperators.Teleoperator.action_features`]: one `.pos` entry per motor.
"""
return self.action_features
@property
def is_connected(self) -> bool:
"""See [`~teleoperators.Teleoperator.is_connected`]."""
return self.bus.is_connected
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
"""See [`~teleoperators.Teleoperator.connect`]."""
logger.info(f"Connecting arm on {self.config.port}...")
self.bus.connect()
@@ -103,11 +129,11 @@ class OpenArmMini(Teleoperator):
@property
def is_calibrated(self) -> bool:
"""See [`~teleoperators.Teleoperator.is_calibrated`]."""
return self.bus.is_calibrated
def calibrate(self) -> None:
"""
Run calibration procedure for a single OpenArm Mini arm.
"""See [`~teleoperators.Teleoperator.calibrate`].
1. Disable torque
2. Ask user to position arm in hanging position with gripper closed
@@ -201,12 +227,23 @@ class OpenArmMini(Teleoperator):
print(f"\nCalibration complete and saved to {self.calibration_fpath}")
def configure(self) -> None:
"""See [`~teleoperators.Teleoperator.configure`].
Disables torque, applies bus-level motor configuration, then sets every motor to position
operating mode.
"""
self.bus.disable_torque()
self.bus.configure_motors()
for motor in self.bus.motors:
self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
def setup_motors(self) -> None:
"""Assign each motor its bus ID, one at a time.
Run this once when building the teleoperator. Interactive: prompts you to connect the controller
board to a single motor at a time, in reverse order so downstream motors on the daisy chain don't
interfere.
"""
for motor in reversed(self.bus.motors):
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
self.bus.setup_motor(motor)
@@ -214,7 +251,11 @@ class OpenArmMini(Teleoperator):
@check_if_not_connected
def get_action(self) -> RobotAction:
"""Get current action (read positions from all motors)."""
"""See [`~teleoperators.Teleoperator.get_action`].
Applies the `joint_6`/`joint_7` remap, the per-side direction flip configured by `config.side`,
and the gripper teleop-to-degrees conversion before returning.
"""
start = time.perf_counter()
positions = self.bus.sync_read("Present_Position")
@@ -235,13 +276,24 @@ class OpenArmMini(Teleoperator):
return action
def enable_torque(self) -> None:
"""Enable torque on all motors, e.g. to hold position instead of being freely moved by hand."""
self.bus.enable_torque()
def disable_torque(self) -> None:
"""Disable torque on all motors so the arm can be moved by hand."""
self.bus.disable_torque()
def write_goal_positions(self, positions: dict[str, float]) -> None:
"""Write goal positions to motors (inverse of get_action flip/gripper/remap logic)."""
"""Write goal positions to the motors.
Applies the inverse of [`~teleoperators.openarm_mini.OpenArmMini.get_action`]'s remap, direction flip, and
gripper unit conversion before writing.
Args:
positions (`dict[str, float]`):
Target positions keyed by `{motor}.pos`, in the same units [`~teleoperators.openarm_mini.OpenArmMini.get_action`]
returns.
"""
goals: dict[str, float] = {}
for key, val in positions.items():
if not key.endswith(".pos"):
@@ -261,9 +313,15 @@ class OpenArmMini(Teleoperator):
@check_if_not_connected
def send_feedback(self, feedback: dict[str, float]) -> None:
"""See [`~teleoperators.Teleoperator.send_feedback`].
Delegates to [`~teleoperators.openarm_mini.OpenArmMini.write_goal_positions`], moving the arm's motors to the
given positions.
"""
self.write_goal_positions(feedback)
@check_if_not_connected
def disconnect(self) -> None:
"""See [`~teleoperators.Teleoperator.disconnect`]."""
self.bus.disconnect()
logger.info(f"{self} disconnected.")
@@ -23,6 +23,14 @@ from ..config import TeleoperatorConfig
class PhoneOS(Enum):
"""Which phone platform a `Phone` teleoperator talks to, selecting its backend implementation.
**Attributes**:
- **ANDROID** (`str`) -- WebXR-based backend (`AndroidPhone`), driven through the `teleop` Python
package.
- **IOS** (`str`) -- ARKit-based backend (`IOSPhone`), driven through the HEBI Mobile I/O app.
"""
ANDROID = "android"
IOS = "ios"
@@ -30,6 +38,35 @@ class PhoneOS(Enum):
@TeleoperatorConfig.register_subclass("phone")
@dataclass
class PhoneConfig(TeleoperatorConfig):
"""Configuration for the [`~teleoperators.phone.Phone`] teleoperator.
Args:
phone_os (`PhoneOS`, *optional*, defaults to `PhoneOS.IOS`):
Which phone platform and backend to use. `PhoneOS.IOS` talks to the HEBI Mobile I/O app over
ARKit; `PhoneOS.ANDROID` talks to a browser WebXR session over the `teleop` package.
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
Note:
`camera_offset` is a fixed class attribute, not a constructor argument, so it currently cannot be
overridden per instance or from the command line. It defaults to the offset between an iPhone 14
Pro's camera and the phone's physical center (2cm lateral, 4cm vertical) and is applied to
translate the ARKit/WebXR camera pose into the phone's own frame.
Example:
```python
>>> from lerobot.teleoperators.phone import PhoneConfig
>>> from lerobot.teleoperators.phone.config_phone import PhoneOS
>>> config = PhoneConfig(phone_os=PhoneOS.ANDROID)
>>> config.phone_os
<PhoneOS.ANDROID: 'android'>
```
"""
phone_os: PhoneOS = PhoneOS.IOS
camera_offset = np.array(
[0.0, -0.02, 0.04]
@@ -26,8 +26,7 @@ from .config_phone import PhoneOS
@ProcessorStepRegistry.register("map_phone_action_to_robot_action")
@dataclass
class MapPhoneActionToRobotAction(RobotActionProcessorStep):
"""
Maps calibrated phone pose actions to standardized robot action inputs.
"""Maps calibrated phone pose actions to standardized robot action inputs.
This processor step acts as a bridge between the phone teleoperator's output
and the robot's expected action format. It remaps the phone's 6-DoF pose
@@ -45,17 +44,22 @@ class MapPhoneActionToRobotAction(RobotActionProcessorStep):
_enabled_prev: bool = field(default=False, init=False, repr=False)
def action(self, action: RobotAction) -> RobotAction:
"""
Processes the phone action dictionary to create a robot action dictionary.
"""Processes the phone action dictionary to create a robot action dictionary.
Args:
act: The input action dictionary from the phone teleoperator.
action (`RobotAction`):
The input action dictionary from the phone teleoperator, keyed `"phone.pos"`,
`"phone.rot"`, `"phone.raw_inputs"`, and `"phone.enabled"`.
Returns:
A new action dictionary formatted for the robot controller.
`RobotAction`: A new action dictionary formatted for the robot controller, keyed
`"enabled"`, `"target_x"`/`"target_y"`/`"target_z"`, `"target_wx"`/`"target_wy"`/`"target_wz"`,
and `"gripper_vel"`.
Raises:
ValueError: If 'pos' or 'rot' keys are missing from the input action.
KeyError: If `"phone.pos"`, `"phone.rot"`, `"phone.raw_inputs"`, or `"phone.enabled"` is
missing from `action`.
ValueError: If `"phone.pos"` or `"phone.rot"` is `None`.
"""
# Pop them from the action
enabled = bool(action.pop("phone.enabled"))
@@ -92,6 +96,20 @@ class MapPhoneActionToRobotAction(RobotActionProcessorStep):
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
"""Replace the `phone.*` action feature entries with the robot action features `action` produces.
Drops the `"phone.enabled"`, `"phone.pos"`, `"phone.rot"`, and `"phone.raw_inputs"` feature
entries, and adds one scalar (`shape=(1,)`) entry for each of `"enabled"`, `"target_x"`,
`"target_y"`, `"target_z"`, `"target_wx"`, `"target_wy"`, `"target_wz"`, and `"gripper_vel"`.
Args:
features (`dict[PipelineFeatureType, dict[str, PolicyFeature]]`):
The pipeline's feature dictionary, keyed by pipeline feature type and then feature name.
Returns:
`dict[PipelineFeatureType, dict[str, PolicyFeature]]`: The same dictionary, with the action
feature entries updated in place.
"""
for feat in ["enabled", "pos", "rot", "raw_inputs"]:
features[PipelineFeatureType.ACTION].pop(f"phone.{feat}", None)
+249 -15
View File
@@ -46,6 +46,14 @@ logger = logging.getLogger(__name__)
class BasePhone:
"""Shared calibration state and `Teleoperator` interface parts common to both phone backends.
`IOSPhone` and `AndroidPhone` mix this in alongside `Teleoperator` so that the action/feedback feature
schemas, calibration status, and the no-op configuration step only need to be written once. Each
backend implements the parts that genuinely differ: connecting, reading the raw pose, and capturing a
calibration reference.
"""
_enabled: bool = False
_calib_pos: np.ndarray | None = None
_calib_rot_inv: Rotation | None = None
@@ -55,10 +63,24 @@ class BasePhone:
@property
def is_calibrated(self) -> bool:
"""Whether a calibration reference pose has been captured.
Returns:
`bool`: `True` once both a reference position and inverse rotation have been recorded by
`calibrate`.
"""
return (self._calib_pos is not None) and (self._calib_rot_inv is not None)
@property
def action_features(self) -> dict[str, type]:
"""Describe the action dictionary returned by `get_action`.
Returns:
`dict[str, type]`: Maps `"phone.pos"` (3D position, shape `(3,)`), `"phone.rot"` (orientation,
a `scipy.spatial.transform.Rotation`), `"phone.raw_inputs"` (device-specific analog/button or
WebXR values), and `"phone.enabled"` (whether the teleoperation trigger is currently held) to
their value types.
"""
return {
"phone.pos": np.ndarray, # shape (3,)
"phone.rot": Rotation, # scipy.spatial.transform.Rotation
@@ -68,22 +90,60 @@ class BasePhone:
@property
def feedback_features(self) -> dict[str, type]:
"""Feedback schema accepted by `send_feedback`.
No haptic or other feedback channel is implemented for phone teleoperators yet.
Returns:
`dict[str, type]`: Currently always `None`, since `feedback_features` has no implementation
yet; this deviates from the declared return type and should not be relied on.
"""
# No haptic or other feedback implemented yet
pass
def configure(self) -> None:
"""No-op. Phone teleoperators require no runtime configuration.
See [`~teleoperators.Teleoperator.configure`] for the base contract.
"""
# No additional configuration required for phone teleop
pass
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not implemented. Phone teleoperators do not support feedback yet.
Args:
feedback (`dict[str, float]`):
Feedback values; see [`~teleoperators.Teleoperator.send_feedback`] for the base contract.
Raises:
NotImplementedError: Always. Haptic feedback (phone vibration) is not implemented yet.
"""
# We could add haptic feedback (vibrations) here, but it's not implemented yet
raise NotImplementedError
class IOSPhone(BasePhone, Teleoperator):
"""ARKit-based teleoperator backend for iOS, driven through the HEBI Mobile I/O app.
Reads the phone's 6-DoF pose (position and orientation) captured by ARKit and relayed over the HEBI
SDK, along with the app's 8 analog (`a1`-`a8`) and 8 digital (`b1`-`b8`) inputs. `Phone` instantiates
this internally when `PhoneConfig.phone_os` is `PhoneOS.IOS`; use `Phone` directly rather than this
class.
"""
name = "ios_phone"
def __init__(self, config: PhoneConfig):
"""Check for the optional dependencies this backend needs and store the configuration.
Args:
config (`PhoneConfig`):
Configuration shared with the parent `Phone` teleoperator.
Raises:
ImportError: If the `hebi-py` or `teleop` packages are not installed.
"""
require_package("hebi-py", extra="phone", import_name="hebi")
require_package("teleop", extra="phone")
super().__init__(config)
@@ -92,10 +152,26 @@ class IOSPhone(BasePhone, Teleoperator):
@property
def is_connected(self) -> bool:
"""See [`~teleoperators.Teleoperator.is_connected`].
Returns:
`bool`: `True` once a HEBI feedback group has been acquired by `connect`.
"""
return self._group is not None
@check_if_already_connected
def connect(self) -> None:
"""Look up the HEBI Mobile I/O group over the network, then calibrate.
Waits briefly for the HEBI lookup service to discover the phone running the Mobile I/O app under
the `"HEBI"` family / `"mobileIO"` name, then immediately runs `calibrate`, which blocks until the
user captures a reference pose in the app. Unlike
[`~teleoperators.Teleoperator.connect`], this method always calibrates; there is no way to skip it.
Raises:
DeviceAlreadyConnectedError: If already connected.
RuntimeError: If no matching Mobile I/O group is found on the network.
"""
logger.info("Connecting to IPhone, make sure to open the HEBI Mobile I/O app.")
lookup = hebi.Lookup()
time.sleep(2.0)
@@ -108,6 +184,13 @@ class IOSPhone(BasePhone, Teleoperator):
self.calibrate()
def calibrate(self) -> None:
"""Block until the user captures a reference pose via the HEBI Mobile I/O app.
Prompts the user to hold the phone so its top edge points along the robot's +x axis and its
screen faces the robot's +z axis, then to press and hold button `B1` in the app to capture that
pose as the calibration reference. See [`~teleoperators.Teleoperator.calibrate`] for the base
contract.
"""
print(
"Hold the phone so that: top edge points forward in same direction as the robot (robot +x) and screen points up (robot +z)"
)
@@ -119,8 +202,7 @@ class IOSPhone(BasePhone, Teleoperator):
print("Calibration done\n")
def _wait_for_capture_trigger(self) -> tuple[np.ndarray, Rotation]:
"""
Blocks execution until the calibration trigger is detected from the iOS device.
"""Blocks execution until the calibration trigger is detected from the iOS device.
This method enters a loop, continuously reading the phone's state. It waits for the user to press
and hold the 'B1' button in the HEBI Mobile I/O app. Once B1 is pressed, the loop breaks and
@@ -147,8 +229,7 @@ class IOSPhone(BasePhone, Teleoperator):
time.sleep(0.01)
def _read_current_pose(self) -> tuple[bool, np.ndarray | None, Rotation | None, object | None]:
"""
Reads the instantaneous 6-DoF pose from the connected iOS device via the HEBI SDK.
"""Reads the instantaneous 6-DoF pose from the connected iOS device via the HEBI SDK.
This method fetches the latest feedback packet from the HEBI group, extracts the ARKit
position and orientation, and converts them into a standard format. It also applies a
@@ -183,6 +264,20 @@ class IOSPhone(BasePhone, Teleoperator):
@check_if_not_connected
def get_action(self) -> dict:
"""Read the phone's current calibrated pose and raw HEBI inputs.
Applies the calibration captured by `calibrate` to the raw ARKit pose, and re-anchors the
reference position on the rising edge of the `b1` "enable" button so that moving the phone while
disabled does not cause a jump once teleoperation resumes.
Returns:
`dict`: Matches `action_features`: `"phone.pos"`, `"phone.rot"`, `"phone.raw_inputs"` (the
app's analog/digital channel values, keyed e.g. `"a1"`, `"b1"`), and `"phone.enabled"`. An
empty `dict` if no pose has been received yet or the teleoperator has not been calibrated.
Raises:
DeviceNotConnectedError: If `connect` has not been called.
"""
has_pose, raw_position, raw_rotation, fb_pose = self._read_current_pose()
if not has_pose or not self.is_calibrated:
return {}
@@ -224,13 +319,34 @@ class IOSPhone(BasePhone, Teleoperator):
@check_if_not_connected
def disconnect(self) -> None:
"""See [`~teleoperators.Teleoperator.disconnect`].
Raises:
DeviceNotConnectedError: If `connect` has not been called.
"""
self._group = None
class AndroidPhone(BasePhone, Teleoperator):
"""WebXR-based teleoperator backend for Android, driven through the `teleop` Python package.
Runs the `teleop` package's local WebXR server on a background thread and reads the pose and touch
events posted by the phone's browser session. `Phone` instantiates this internally when
`PhoneConfig.phone_os` is `PhoneOS.ANDROID`; use `Phone` directly rather than this class.
"""
name = "android_phone"
def __init__(self, config: PhoneConfig):
"""Check for the optional dependencies this backend needs and store the configuration.
Args:
config (`PhoneConfig`):
Configuration shared with the parent `Phone` teleoperator.
Raises:
ImportError: If the `hebi-py` or `teleop` packages are not installed.
"""
require_package("hebi-py", extra="phone", import_name="hebi")
require_package("teleop", extra="phone")
super().__init__(config)
@@ -243,10 +359,26 @@ class AndroidPhone(BasePhone, Teleoperator):
@property
def is_connected(self) -> bool:
"""See [`~teleoperators.Teleoperator.is_connected`].
Returns:
`bool`: `True` once the `teleop` background thread has been started by `connect`.
"""
return self._teleop is not None
@check_if_already_connected
def connect(self) -> None:
"""Start the `teleop` WebXR server on a background thread, then calibrate.
Subscribes to pose/message updates from the `teleop` package and starts its server loop on a
daemon thread, then immediately runs `calibrate`, which blocks until the user captures a reference
pose from the phone's browser session. Unlike
[`~teleoperators.Teleoperator.connect`], this method always calibrates; there is no way to skip
it.
Raises:
DeviceAlreadyConnectedError: If already connected.
"""
logger.info("Starting teleop stream for Android...")
self._teleop = Teleop()
self._teleop.subscribe(self._android_callback)
@@ -257,6 +389,13 @@ class AndroidPhone(BasePhone, Teleoperator):
self.calibrate()
def calibrate(self) -> None:
"""Block until the user captures a reference pose via touch on the WebXR page.
Prompts the user to hold the phone so its top edge points along the robot's +x axis and its
screen faces the robot's +z axis, then to touch and move a finger on the WebXR page to capture
that pose as the calibration reference. See [`~teleoperators.Teleoperator.calibrate`] for the base
contract.
"""
print(
"Hold the phone so that: top edge points forward in same direction as the robot (robot +x) and screen points up (robot +z)"
)
@@ -269,8 +408,7 @@ class AndroidPhone(BasePhone, Teleoperator):
print("Calibration done\n")
def _wait_for_capture_trigger(self) -> tuple[np.ndarray, Rotation]:
"""
Blocks execution until the calibration trigger is detected from the Android device.
"""Blocks execution until the calibration trigger is detected from the Android device.
This method enters a loop, continuously checking the latest message received from the WebXR
session. It waits for the user to touch and move their finger on the screen, which generates
@@ -293,8 +431,7 @@ class AndroidPhone(BasePhone, Teleoperator):
time.sleep(0.01)
def _read_current_pose(self) -> tuple[bool, np.ndarray | None, Rotation | None, object | None]:
"""
Reads the latest 6-DoF pose received from the Android device's WebXR session.
"""Reads the latest 6-DoF pose received from the Android device's WebXR session.
This method accesses the most recent pose data stored by the `_android_callback`. It uses a
thread lock to safely read the shared `_latest_pose` variable. The pose, a 4x4 matrix, is
@@ -317,8 +454,7 @@ class AndroidPhone(BasePhone, Teleoperator):
return True, pos, rot, pose
def _android_callback(self, pose: np.ndarray, message: dict) -> None:
"""
Callback function to handle incoming data from the Android teleop stream.
"""Callback function to handle incoming data from the Android teleop stream.
This method is executed by the `teleop` package's subscriber thread whenever a new
pose and message are received from the WebXR session on the Android phone. It updates
@@ -336,6 +472,20 @@ class AndroidPhone(BasePhone, Teleoperator):
@check_if_not_connected
def get_action(self) -> dict:
"""Read the phone's current calibrated pose and raw touch/button state.
Applies the calibration captured by `calibrate` to the latest pose received from the `teleop`
background thread, and re-anchors the reference position on the rising edge of the `"move"` touch
event so that moving the phone while disabled does not cause a jump once teleoperation resumes.
Returns:
`dict`: Matches `action_features`: `"phone.pos"`, `"phone.rot"`, `"phone.raw_inputs"`
(`"move"`, `"scale"`, `"reservedButtonA"`, `"reservedButtonB"`), and `"phone.enabled"`. An
empty `dict` if no pose has been received yet or the teleoperator has not been calibrated.
Raises:
DeviceNotConnectedError: If `connect` has not been called.
"""
ok, raw_pos, raw_rot, pose = self._read_current_pose()
if not ok or not self.is_calibrated:
return {}
@@ -369,6 +519,11 @@ class AndroidPhone(BasePhone, Teleoperator):
@check_if_not_connected
def disconnect(self) -> None:
"""Stop the `teleop` background thread.
Raises:
DeviceNotConnectedError: If `connect` has not been called.
"""
self._teleop = None
if self._teleop_thread and self._teleop_thread.is_alive():
self._teleop_thread.join(timeout=1.0)
@@ -377,18 +532,42 @@ class AndroidPhone(BasePhone, Teleoperator):
class Phone(Teleoperator):
"""
Phone-based teleoperator using ARKit (iOS via HEBI Mobile I/O App) or the teleop Python package (Android via WebXR API).
For HEBI Mobile I/O we also expose 8 analog (a1-a8) and 8 digital (b1-b8) inputs.
"""Phone-based teleoperator: iOS via ARKit and the HEBI Mobile I/O app, Android via WebXR.
Press and hold **B1** to enable teleoperation. While enabled, the first B1 press
captures a reference pose and rotation, when disabled and pressed again the position is reapplied.
Reads the phone's 6-DoF pose and, for the HEBI Mobile I/O app, 8 analog (`a1`-`a8`) and 8 digital
(`b1`-`b8`) inputs. Which backend is used is picked at construction time from
`config.phone_os` and delegated to internally: [`~teleoperators.Teleoperator`] method calls on `Phone`
forward to either an `IOSPhone` or an `AndroidPhone` instance.
Press and hold **B1** (iOS) or touch and move on the WebXR page (Android) to enable teleoperation.
The first press/touch while enabled captures a reference pose; releasing and re-triggering re-anchors
the reference position to wherever the phone currently is, so motion is always relative to where
teleoperation was last resumed.
Example:
```python
>>> from lerobot.teleoperators.phone import Phone, PhoneConfig
>>> teleop = Phone(PhoneConfig()) # doctest: +SKIP
>>> teleop.connect() # doctest: +SKIP
>>> teleop.get_action() # doctest: +SKIP
```
"""
config_class = PhoneConfig
name = "phone"
def __init__(self, config: PhoneConfig):
"""Pick and construct the backend matching `config.phone_os`.
Args:
config (`PhoneConfig`):
Configuration selecting the phone platform (`config.phone_os`) and forwarded to the
chosen backend.
Raises:
ValueError: If `config.phone_os` is not a valid `PhoneOS` member.
ImportError: If the `hebi-py` or `teleop` packages are not installed.
"""
super().__init__(config)
self.config = config
@@ -403,34 +582,89 @@ class Phone(Teleoperator):
@property
def is_connected(self) -> bool:
"""See [`~teleoperators.Teleoperator.is_connected`].
Returns:
`bool`: `True` if the underlying `IOSPhone` or `AndroidPhone` backend is connected.
"""
return self._phone_impl.is_connected
def connect(self) -> None:
"""Connect and calibrate through the underlying backend.
Unlike [`~teleoperators.Teleoperator.connect`], this always calibrates; there is no `calibrate`
argument to opt out.
Raises:
DeviceAlreadyConnectedError: If already connected.
RuntimeError: If the iOS backend cannot find the Mobile I/O group on the network.
"""
return self._phone_impl.connect()
def calibrate(self) -> None:
"""See [`~teleoperators.Teleoperator.calibrate`]. Delegates to the underlying backend."""
return self._phone_impl.calibrate()
@property
def is_calibrated(self) -> bool:
"""See [`~teleoperators.Teleoperator.is_calibrated`].
Returns:
`bool`: `True` once a calibration reference pose has been captured.
"""
return self._phone_impl.is_calibrated
@property
def action_features(self) -> dict[str, type]:
"""See [`~teleoperators.Teleoperator.action_features`].
Returns:
`dict[str, type]`: `"phone.pos"`, `"phone.rot"`, `"phone.raw_inputs"`, and `"phone.enabled"`
mapped to their value types; see `get_action` for what each holds.
"""
return self._phone_impl.action_features
@property
def feedback_features(self) -> dict[str, type]:
"""See [`~teleoperators.Teleoperator.feedback_features`].
Returns:
`dict[str, type]`: Currently always `None`, since no feedback channel is implemented yet.
"""
return self._phone_impl.feedback_features
def configure(self) -> None:
"""No-op. See [`~teleoperators.Teleoperator.configure`]."""
return self._phone_impl.configure()
def get_action(self) -> dict:
"""Read the phone's current calibrated pose and raw inputs from the underlying backend.
Returns:
`dict`: Matches `action_features`. An empty `dict` if no pose has been received yet or the
teleoperator has not been calibrated.
Raises:
DeviceNotConnectedError: If `connect` has not been called.
"""
return self._phone_impl.get_action()
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not implemented. See [`~teleoperators.Teleoperator.send_feedback`].
Args:
feedback (`dict[str, float]`):
Feedback values; unused.
Raises:
NotImplementedError: Always. Haptic feedback is not implemented yet.
"""
return self._phone_impl.send_feedback(feedback)
def disconnect(self) -> None:
"""See [`~teleoperators.Teleoperator.disconnect`]. Delegates to the underlying backend.
Raises:
DeviceNotConnectedError: If `connect` has not been called.
"""
return self._phone_impl.disconnect()
@@ -22,6 +22,37 @@ from ..config import TeleoperatorConfig
@TeleoperatorConfig.register_subclass("reachy2_teleoperator")
@dataclass
class Reachy2TeleoperatorConfig(TeleoperatorConfig):
"""Configuration for reading teleoperation actions from a Reachy 2.
Reachy 2 can act as its own teleoperator: instead of a leader arm, another Reachy 2 (or the same one in
a different mode) reports its joint positions over the network as the action. There is no LeRobot
calibration file; Reachy 2 manages its own calibration.
Which joints are reported is selected by the `with_*` flags: turning a part off removes its joints
entirely. At least one part must stay enabled.
Args:
ip_address (`str`, *optional*, defaults to `"localhost"`):
Address of the Reachy 2 robot to read actions from.
use_present_position (`bool`, *optional*, defaults to `False`):
Whether to report each joint's present position as the action. If `False`, the joint's goal
position is reported instead.
with_mobile_base (`bool`, *optional*, defaults to `True`):
Whether to include the mobile base's velocity in actions.
with_l_arm (`bool`, *optional*, defaults to `True`):
Whether to include the left arm's joints.
with_r_arm (`bool`, *optional*, defaults to `True`):
Whether to include the right arm's joints.
with_neck (`bool`, *optional*, defaults to `True`):
Whether to include the neck's joints.
with_antennas (`bool`, *optional*, defaults to `True`):
Whether to include the antennas' joints.
id (`str`, *optional*):
Identifier for this particular teleoperator.
calibration_dir (`Path`, *optional*):
Unused: Reachy 2 manages its own calibration.
"""
# IP address of the Reachy 2 robot used as teleoperator
ip_address: str | None = "localhost"
@@ -37,6 +68,11 @@ class Reachy2TeleoperatorConfig(TeleoperatorConfig):
with_antennas: bool = True
def __post_init__(self):
"""Validate that at least one robot part is enabled.
Raises:
ValueError: If every robot part is disabled, which would leave no joints to report.
"""
if not (
self.with_mobile_base
or self.with_l_arm
@@ -76,14 +76,19 @@ REACHY2_VEL = {
class Reachy2Teleoperator(Teleoperator):
"""
[Reachy 2](https://www.pollen-robotics.com/reachy/), by Pollen Robotics.
"""
"""[Reachy 2](https://www.pollen-robotics.com/reachy/), by Pollen Robotics."""
config_class = Reachy2TeleoperatorConfig
name = "reachy2_specific"
def __init__(self, config: Reachy2TeleoperatorConfig):
"""Build the teleoperator from its configuration.
Args:
config (`Reachy2TeleoperatorConfig`):
The teleoperator's configuration. Its `ip_address` and `with_*` flags determine what is
read.
"""
require_package("reachy2_sdk", extra="reachy2")
super().__init__(config)
@@ -106,6 +111,13 @@ class Reachy2Teleoperator(Teleoperator):
@property
def action_features(self) -> dict[str, type]:
"""The joint positions (and mobile base velocity, if enabled) read from Reachy 2.
Returns:
`dict[str, type]`: `"<joint>.pos"` keys for each enabled part mapped to `float`, plus
`"mobile_base.vx"`, `"mobile_base.vy"`, and `"mobile_base.vtheta"` when
`config.with_mobile_base` is `True`.
"""
if self.config.with_mobile_base:
return {
**dict.fromkeys(
@@ -122,14 +134,32 @@ class Reachy2Teleoperator(Teleoperator):
@property
def feedback_features(self) -> dict[str, type]:
"""Always empty: this teleoperator does not accept feedback.
Returns:
`dict[str, type]`: An empty dictionary.
"""
return {}
@property
def is_connected(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_connected`]."""
return self.reachy.is_connected() if self.reachy is not None else False
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
"""Open the gRPC connection to Reachy 2's teleoperation interface.
The `calibrate` argument is accepted for interface compatibility but has no effect: Reachy 2
manages its own calibration.
Args:
calibrate (`bool`, *optional*, defaults to `True`):
Unused.
Raises:
DeviceNotConnectedError: If the connection could not be established.
"""
self.reachy = ReachySDK(self.config.ip_address)
if not self.is_connected:
@@ -138,16 +168,32 @@ class Reachy2Teleoperator(Teleoperator):
@property
def is_calibrated(self) -> bool:
"""Always `True`: Reachy 2 manages its own calibration.
Returns:
`bool`: Always `True`.
"""
return True
def calibrate(self) -> None:
"""No-op: Reachy 2 manages its own calibration."""
pass
def configure(self) -> None:
"""No-op: Reachy 2 requires no additional configuration."""
pass
@check_if_not_connected
def get_action(self) -> dict[str, float]:
"""Read the current (or goal) joint positions and mobile base velocity from Reachy 2.
Returns:
`dict[str, float]`: Values keyed as described by
[`~teleoperators.Teleoperator.action_features`].
Raises:
DeviceNotConnectedError: If [`~teleoperators.Teleoperator.connect`] has not been called.
"""
start = time.perf_counter()
joint_action: dict[str, float] = {}
@@ -170,8 +216,14 @@ class Reachy2Teleoperator(Teleoperator):
return {**joint_action, **vel_action}
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not supported.
Raises:
NotImplementedError: Always. This teleoperator does not accept feedback.
"""
raise NotImplementedError
def disconnect(self) -> None:
"""Close the gRPC connection to Reachy 2, if it is open."""
if self.is_connected:
self.reachy.disconnect()
@@ -21,10 +21,14 @@ from ..config import TeleoperatorConfig
@dataclass
class RebotArm102LeaderConfig:
"""Base configuration class for the Seeed Studio StarArm102 / reBot Arm 102 leader.
"""Field definitions shared by the reBot Arm 102 leader.
The reBot Arm 102 is a 7-joint (incl. gripper) leader arm driven by FashionStar
UART smart servos. Servo communication goes through ``motorbridge-smart-servo``.
The reBot Arm 102 is a 7-joint (incl. gripper) leader arm driven by FashionStar UART smart servos.
Servo communication goes through ``motorbridge-smart-servo``.
This class only carries the fields. The registered configuration users instantiate is
[`RebotArm102LeaderTeleopConfig`], which combines these with [`~teleoperators.TeleoperatorConfig`] and
documents them all in one place doc-builder renders only a class's own docstring, never its bases'.
"""
# USB-to-UART device the leader arm is connected to (e.g. "/dev/ttyUSB0").
@@ -78,6 +82,28 @@ class RebotArm102LeaderConfig:
@TeleoperatorConfig.register_subclass("rebot_102_leader")
@dataclass
class RebotArm102LeaderTeleopConfig(TeleoperatorConfig, RebotArm102LeaderConfig):
"""Registered configuration for the reBot Arm 102 leader teleoperator."""
"""Registered configuration for the reBot Arm 102 leader teleoperator.
Args:
port (`str`):
USB-to-UART device the leader arm is connected to, e.g. `/dev/ttyUSB0`.
baudrate (`int`, *optional*, defaults to 1000000):
Baud rate of the UART link to the FashionStar smart servos.
joint_ids (`dict[str, int]`, *optional*):
Servo id of each joint on the UART bus. Defaults to the reBot Arm 102's standard 7-joint
layout (`shoulder_pan`, `shoulder_lift`, `elbow_flex`, `wrist_flex`, `wrist_yaw`,
`wrist_roll`, `gripper`).
joint_directions (`dict[str, int]`, *optional*):
Per-joint sign applied to raw servo angles so the leader matches the follower convention. The
gripper additionally carries a scale (e.g. `-6`) to widen its range to the reBot B601
follower's gripper travel.
joint_ranges (`dict[str, list[int]]`, *optional*):
Per-joint `[min, max]` output range in degrees. Defaults to ranges matching the reBot B601
follower's joint limits so leader actions can drive the follower key-for-key.
id (`str`, *optional*):
Identifier for this particular arm; also names its calibration file.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
"""
pass
@@ -49,6 +49,12 @@ class RebotArm102Leader(Teleoperator):
name = "rebot_102_leader"
def __init__(self, config: RebotArm102LeaderTeleopConfig):
"""Build the teleoperator from its configuration.
Args:
config (`RebotArm102LeaderTeleopConfig`):
The teleoperator's configuration. Its `port` determines what is connected.
"""
require_package("motorbridge-smart-servo", extra="rebot", import_name="motorbridge_smart_servo")
super().__init__(config)
self.config = config
@@ -58,18 +64,39 @@ class RebotArm102Leader(Teleoperator):
@property
def action_features(self) -> dict[str, type]:
"""The arm's joint positions, in degrees.
Returns:
`dict[str, type]`: `"<motor>.pos"` keys mapped to `float`.
"""
return {f"{motor}.pos": float for motor in self.motor_names}
@property
def feedback_features(self) -> dict[str, type]:
"""This arm accepts no feedback.
Returns:
`dict[str, type]`: Always empty.
"""
return {}
@property
def is_connected(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_connected`]: the servo bus has been opened."""
return self.bus is not None
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
"""Open the UART servo bus, ping every configured joint, then calibrate and configure the arm.
Args:
calibrate (`bool`, *optional*, defaults to `True`):
Whether to run calibration when the arm is not already calibrated. Calibration is
interactive and prompts on stdin.
Raises:
RuntimeError: If a configured servo does not respond to a ping.
"""
logger.info(f"Connecting {self} on {self.config.port}...")
bus = FashionStarServo(self.config.port, baudrate=self.config.baudrate)
try:
@@ -95,9 +122,20 @@ class RebotArm102Leader(Teleoperator):
@property
def is_calibrated(self) -> bool:
"""Whether every configured joint has a saved calibration entry.
Returns:
`bool`: `True` if `self.calibration` has an entry for each of `self.motor_names`.
"""
return bool(self.calibration) and set(self.calibration) == set(self.motor_names)
def calibrate(self) -> None:
"""Set the zero position of every joint from the arm's current pose.
If a calibration file already exists, prompts the operator to reuse it or to redo calibration. To
redo it, the operator manually moves the arm to its zero pose (gripper closed); each servo's
origin point is then reset to that pose and the result is saved to the calibration file.
"""
if self.calibration:
user_input = input(
f"Press ENTER to use provided calibration file associated with the id {self.id}, "
@@ -132,6 +170,10 @@ class RebotArm102Leader(Teleoperator):
logger.info(f"Calibration saved to {self.calibration_fpath}")
def configure(self) -> None:
"""Unlock every servo's torque and reset each one's multi-turn counter.
Run once after connecting so subsequent readings start from a known turn count.
"""
for motor_id in self.config.joint_ids.values():
self.bus.unlock(motor_id)
time.sleep(_SETTLE_SEC)
@@ -165,6 +207,16 @@ class RebotArm102Leader(Teleoperator):
@check_if_not_connected
def get_action(self) -> RobotAction:
"""Read, unwrap, and sign-correct the current joint positions.
Each joint's raw multi-turn angle is unwrapped into its configured range (see
`_round_to_valid_range`), then flipped and clipped according to `joint_directions` and
`joint_ranges` so the result matches the follower's convention. If reading the servos fails, the
last successfully read positions are reused and the caller is expected to stop teleoperation.
Returns:
`dict[str, float]`: `"<motor>.pos"` keys mapped to the joint's position in degrees.
"""
start = time.perf_counter()
try:
raw_positions = self._read_raw_positions()
@@ -198,10 +250,16 @@ class RebotArm102Leader(Teleoperator):
return action_dict
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Not supported: the leader arm has no actuators to receive feedback.
Raises:
NotImplementedError: Always.
"""
raise NotImplementedError("Feedback is not implemented for the reBot Arm 102 leader.")
@check_if_not_connected
def disconnect(self) -> None:
"""Close the UART servo bus."""
self.bus.close()
self.bus = None
logger.info(f"{self} disconnected.")
@@ -21,7 +21,12 @@ from ..config import TeleoperatorConfig
@dataclass
class SOLeaderConfig:
"""Base configuration class for SO Leader teleoperators."""
"""Field definitions shared by the SO-family leader arms.
This class only carries the fields. The registered configuration users instantiate is
[`SOLeaderTeleopConfig`], which combines these with [`~teleoperators.TeleoperatorConfig`] and documents
them all in one place doc-builder renders only a class's own docstring, never its bases'.
"""
# Port to connect to the arm
port: str
@@ -40,6 +45,36 @@ class SOLeaderConfig:
@TeleoperatorConfig.register_subclass("so100_leader")
@dataclass
class SOLeaderTeleopConfig(TeleoperatorConfig, SOLeaderConfig):
"""Configuration for the SO-100 and SO-101 leader arms.
Both arms share this class; `SO100LeaderConfig` and `SO101LeaderConfig` are aliases for it. They differ
in their calibration and gearing, not in their control code.
Args:
port (`str`):
Serial port the arm is connected to, e.g. `/dev/ttyACM0` on Linux or `COM3` on Windows. Run
`lerobot-find-port` to identify it.
use_degrees (`bool`, *optional*, defaults to `True`):
Whether to report joint positions in degrees. Keep `True` for compatibility with existing
policies and datasets.
num_read_retries (`int`, *optional*, defaults to 2):
Extra attempts when a `sync_read` fails. Feetech buses occasionally return a corrupted status
packet, especially when several joints move at once, which would otherwise abort the
teleoperation loop. Retries are immediate and only happen on failure, so steady-state read cost
is unchanged.
id (`str`, *optional*):
Identifier for this particular arm; also names its calibration file.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
Example:
```python
>>> from lerobot.teleoperators.so_leader import SO101Leader, SO101LeaderConfig
>>> config = SO101LeaderConfig(port="/dev/ttyACM0") # doctest: +SKIP
>>> teleop = SO101Leader(config) # doctest: +SKIP
```
"""
pass
@@ -31,12 +31,34 @@ logger = logging.getLogger(__name__)
class SOLeader(Teleoperator):
"""Generic SO leader base for SO-100/101/10X teleoperators."""
"""The SO-family leader arm: a 5-DOF arm plus gripper on a Feetech bus, held to teleoperate a follower arm.
`SO100Leader` and `SO101Leader` are aliases of this class. The two arms differ in calibration and
gearing, not control code, so both are driven through the same implementation with a different
`config_class` and `name`.
Actions are keyed `"<motor>.pos"`. See [`~teleoperators.Teleoperator`] for the contract every method
here implements.
Example:
```python
>>> from lerobot.teleoperators.so_leader import SO101Leader, SO101LeaderConfig
>>> teleop = SO101Leader(SO101LeaderConfig(port="/dev/ttyACM0")) # doctest: +SKIP
>>> with teleop: # doctest: +SKIP
... action = teleop.get_action()
```
"""
config_class = SOLeaderTeleopConfig
name = "so_leader"
def __init__(self, config: SOLeaderTeleopConfig):
"""Build the teleoperator from its configuration.
Args:
config (`SOLeaderTeleopConfig`):
The teleoperator's configuration. Its `port` determines what is connected.
"""
super().__init__(config)
self.config = config
norm_mode_body = MotorNormMode.DEGREES if config.use_degrees else MotorNormMode.RANGE_M100_100
@@ -55,18 +77,42 @@ class SOLeader(Teleoperator):
@property
def action_features(self) -> dict[str, type]:
"""The arm's joint positions.
Returns:
`dict[str, type]`: `"<motor>.pos"` keys mapped to `float`.
"""
return {f"{motor}.pos": float for motor in self.bus.motors}
@property
def feedback_features(self) -> dict[str, type]:
"""The arm's target joint positions, used to sync this leader arm to another pose.
Shares the same keys as [`~teleoperators.Teleoperator.action_features`], since feedback for this
arm is a goal position written to each motor.
Returns:
`dict[str, type]`: `"<motor>.pos"` keys mapped to `float`.
"""
return self.action_features
@property
def is_connected(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_connected`]."""
return self.bus.is_connected
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
"""Connect the motor bus, calibrating and configuring the arm.
Args:
calibrate (`bool`, *optional*, defaults to `True`):
Whether to run calibration when the motors disagree with the calibration file, or no file
exists yet. Calibration is interactive and prompts on stdin.
Raises:
DeviceAlreadyConnectedError: If the teleoperator is already connected.
"""
self.bus.connect()
if not self.is_calibrated and calibrate:
logger.info(
@@ -79,9 +125,15 @@ class SOLeader(Teleoperator):
@property
def is_calibrated(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_calibrated`]."""
return self.bus.is_calibrated
def calibrate(self) -> None:
"""Calibrate the arm, writing the result to the motors and the calibration file.
This is interactive: it prompts on stdin to reuse an existing calibration file, and otherwise asks
you to move the arm to its middle position and then through each joint's full range.
"""
if self.calibration:
# Calibration file exists, ask user whether to use it or run new calibration
user_input = input(
@@ -125,18 +177,34 @@ class SOLeader(Teleoperator):
print(f"Calibration saved to {self.calibration_fpath}")
def configure(self) -> None:
"""Disable torque and write the position-mode operating mode to every motor.
Torque is left disabled so the arm can be moved freely by hand while teleoperating.
"""
self.bus.disable_torque()
self.bus.configure_motors()
for motor in self.bus.motors:
self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
def enable_torque(self) -> None:
"""Enable torque on every motor.
Useful to briefly drive the arm to a position (e.g. via
[`~teleoperators.so_leader.SOLeader.send_feedback`]) before releasing it back to free movement with
[`~teleoperators.so_leader.SOLeader.disable_torque`].
"""
self.bus.enable_torque()
def disable_torque(self) -> None:
"""Disable torque on every motor, letting the arm be moved freely by hand."""
self.bus.disable_torque()
def setup_motors(self) -> None:
"""Assign each motor its bus ID, one at a time.
Run this once when building an arm. It is interactive: it prompts you to connect the controller
board to a single motor at a time, working from the gripper back to the base.
"""
for motor in reversed(self.bus.motors):
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
self.bus.setup_motor(motor)
@@ -144,6 +212,14 @@ class SOLeader(Teleoperator):
@check_if_not_connected
def get_action(self) -> dict[str, float]:
"""Same as [`~teleoperators.Teleoperator.get_action`].
Returns:
`dict[str, float]`: `"<motor>.pos"` keys mapped to the arm's current joint positions.
Raises:
DeviceNotConnectedError: If the teleoperator is not connected.
"""
start = time.perf_counter()
action = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
action = {f"{motor}.pos": val for motor, val in action.items()}
@@ -153,12 +229,29 @@ class SOLeader(Teleoperator):
@check_if_not_connected
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Write goal positions to the arm's motors, e.g. to sync it to a follower's current pose.
Torque must be enabled (see [`~teleoperators.so_leader.SOLeader.enable_torque`]) for the arm to
actually move to the written positions.
Args:
feedback (`dict[str, float]`):
`"<motor>.pos"` keys mapped to target positions. Keys not ending in `.pos` are ignored.
Raises:
DeviceNotConnectedError: If the teleoperator is not connected.
"""
goals = {k.removesuffix(".pos"): v for k, v in feedback.items() if k.endswith(".pos")}
if goals:
self.bus.sync_write("Goal_Position", goals)
@check_if_not_connected
def disconnect(self) -> None:
"""Same as [`~teleoperators.Teleoperator.disconnect`].
Raises:
DeviceNotConnectedError: If the teleoperator is not connected.
"""
self.bus.disconnect()
logger.info(f"{self} disconnected.")
@@ -21,7 +21,15 @@ from ..config import TeleoperatorConfig
@dataclass
class ExoskeletonArmPortConfig:
"""Serial port configuration for individual exoskeleton arm."""
"""Serial port configuration for one exoskeleton arm.
Args:
port (`str`, *optional*, defaults to `""`):
Serial port the exoskeleton arm's sensor board is connected to, e.g. `/dev/ttyUSB0`. An empty
string disables exoskeleton control for that arm.
baud_rate (`int`, *optional*, defaults to 115200):
Baud rate for the serial connection.
"""
port: str = ""
baud_rate: int = 115200
@@ -30,6 +38,26 @@ class ExoskeletonArmPortConfig:
@TeleoperatorConfig.register_subclass("unitree_g1")
@dataclass
class UnitreeG1TeleoperatorConfig(TeleoperatorConfig):
"""Configuration for the Unitree G1 bimanual exoskeleton teleoperator.
Args:
left_arm_config (`ExoskeletonArmPortConfig`, *optional*):
Serial port settings for the left exoskeleton arm. Leave `port` empty to run without exoskeleton
control on this side.
right_arm_config (`ExoskeletonArmPortConfig`, *optional*):
Serial port settings for the right exoskeleton arm. Leave `port` empty to run without
exoskeleton control on this side.
frozen_joints (`str`, *optional*, defaults to `""`):
Comma-separated G1 arm joint names to exclude from the exoskeleton-driven inverse kinematics.
These joints are held at their neutral pose instead of being tracked.
id (`str`, *optional*):
Identifier for this particular unit, used to tell apart several teleoperators of the same
type. It also names the calibration file, so keep it stable for a given piece of hardware.
calibration_dir (`Path`, *optional*):
Where to read and write the calibration file. Defaults to a per-teleoperator directory under
the LeRobot calibration home.
"""
left_arm_config: ExoskeletonArmPortConfig = field(default_factory=ExoskeletonArmPortConfig)
right_arm_config: ExoskeletonArmPortConfig = field(default_factory=ExoskeletonArmPortConfig)
@@ -14,8 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This module handles calibration of hall effect sensors used in the exoskeleton.
"""This module handles calibration of hall effect sensors used in the exoskeleton.
Each joint has a pair of ADC channels outputting sin and cos values that trace an ellipse
as the joint rotates due to imprecision in magnet/sensor placement. We fit this ellipse to a unit circle,
and calculate arctan2 of the unit circle to get the joint angle.
@@ -59,6 +59,21 @@ JOINTS = {
@dataclass
class ExoskeletonJointCalibration:
"""Per-joint calibration mapping raw sin/cos ADC pairs to an angle in radians.
Args:
name (`str`):
Joint name, matching a key in `JOINTS`.
center_fit (`list[float]`):
The `[x, y]` center of the ellipse fitted to this joint's raw sin/cos ADC readings.
T (`list[list[float]]`):
2x2 transformation matrix mapping a centered raw reading onto the unit circle, correcting for
the fitted ellipse's scale and rotation.
zero_offset (`float`, *optional*, defaults to 0.0):
Angle, in radians, measured while the joint was held at its neutral pose. Subtracted from the
raw angle so the neutral pose reads as zero.
"""
name: str # joint name
center_fit: list[float] # center of the ellipse
T: list[list[float]] # 2x2 transformation matrix
@@ -75,6 +90,11 @@ class ExoskeletonCalibration:
joints: list[ExoskeletonJointCalibration] = field(default_factory=list)
def to_dict(self) -> dict:
"""Serialize this calibration to a plain dict suitable for JSON storage.
Returns:
`dict`: The calibration with nested joint calibrations flattened to plain dicts.
"""
return {
"version": self.version,
"side": self.side,
@@ -92,6 +112,15 @@ class ExoskeletonCalibration:
@classmethod
def from_dict(cls, data: dict) -> ExoskeletonCalibration:
"""Reconstruct a calibration from the dict produced by `to_dict`.
Args:
data (`dict`):
Parsed JSON calibration data. Missing optional keys fall back to their defaults.
Returns:
`ExoskeletonCalibration`: The reconstructed calibration.
"""
joints = [
ExoskeletonJointCalibration(
name=j["name"],
@@ -111,6 +140,32 @@ class ExoskeletonCalibration:
@dataclass(frozen=True)
class CalibParams:
"""Tuning knobs for the interactive ellipse-fitting calibration UI.
Args:
fit_every (`float`, *optional*, defaults to 0.15):
Minimum time, in seconds, between successive ellipse re-fits while mapping a joint's range.
min_fit_points (`int`, *optional*, defaults to 60):
Minimum number of buffered samples required before attempting an ellipse fit.
fit_window (`int`, *optional*, defaults to 900):
Number of most recent raw samples considered for each ellipse fit.
max_fit_points (`int`, *optional*, defaults to 300):
Maximum number of points passed to the ellipse fitter; the fit window is downsampled evenly
above this count.
trim_low (`float`, *optional*, defaults to 0.05):
Lower radius quantile below which points are treated as outliers and discarded before fitting.
trim_high (`float`, *optional*, defaults to 0.95):
Upper radius quantile above which points are treated as outliers and discarded before fitting.
median_window (`int`, *optional*, defaults to 5):
Number of raw samples averaged (median) to smooth each sin/cos reading before it is buffered.
history (`int`, *optional*, defaults to 3500):
Maximum number of samples retained per plot, for visualization only.
draw_hz (`float`, *optional*, defaults to 120.0):
Maximum refresh rate of the calibration plot.
sample_count (`int`, *optional*, defaults to 50):
Number of samples averaged to compute a joint's zero-pose offset.
"""
fit_every: float = 0.15
min_fit_points: int = 60
fit_window: int = 900
@@ -129,9 +184,7 @@ def normalize_angle(angle: float) -> float:
def joint_z_and_angle(raw16: list[int], j: ExoskeletonJointCalibration) -> tuple[np.ndarray, float]:
"""
Applies calibration to each joint: raw centered ellipse-to-circle angle.
"""
"""Applies calibration to each joint: raw → centered → ellipse-to-circle → angle."""
pair = JOINTS[j.name]
s, c = raw16[pair[0]], raw16[pair[1]] # get sin and cos
p = np.array([float(c) - ADC_HALF, float(s) - ADC_HALF]) # center the raw values
@@ -153,9 +206,7 @@ def run_exo_calibration(
save_path: Path,
params: CalibParams | None = None,
) -> ExoskeletonCalibration:
"""
Run interactive calibration for an exoskeleton arm.
"""
"""Run interactive calibration for an exoskeleton arm."""
require_package("pyserial", extra="unitree_g1", import_name="serial")
try:
import cv2
@@ -173,9 +224,11 @@ def run_exo_calibration(
logger.info(f"Starting calibration for {side} exoskeleton arm")
def running_median(win: deque) -> float:
"""Return the median of a buffered window of raw ADC samples, used to smooth sensor noise."""
return float(np.median(np.fromiter(win, dtype=float)))
def read_joint_point(raw16: list[int], pair: tuple[int, int]):
"""Extract one joint's centered (x, y) sin/cos point, plus its raw sin/cos values."""
s, c = raw16[pair[0]], raw16[pair[1]]
return float(c) - ADC_HALF, float(s) - ADC_HALF, float(s), float(c)
@@ -259,6 +312,7 @@ def run_exo_calibration(
zero_samples = []
def on_key(event):
"""Matplotlib key-press handler that requests advancing to the calibration's next phase."""
nonlocal advance_requested
if event.key in ("n", "N", "enter", " "):
advance_requested = True
@@ -266,6 +320,7 @@ def run_exo_calibration(
fig.canvas.mpl_connect("key_press_event", on_key)
def reset_state():
"""Build a fresh mutable state dict for tracking one joint's in-progress ellipse fit."""
return {
"xs": deque(maxlen=params.history),
"ys": deque(maxlen=params.history),
+82 -17
View File
@@ -14,9 +14,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""
IK helper for exoskeleton-to-G1 teleoperation. We map Exoskeleton joint angles to end-effector pose in world frame,
visualizing the result in meshcat after calibration.
"""IK helper for exoskeleton-to-G1 teleoperation.
We map Exoskeleton joint angles to end-effector pose in world frame, visualizing the result in meshcat
after calibration.
"""
import logging
@@ -43,6 +44,24 @@ def _frame_id(model, name: str) -> int | None:
@dataclass
class ArmCfg:
"""Static per-arm configuration linking an exoskeleton URDF to its G1 counterpart.
Args:
side (`str`):
Which arm this describes, `"left"` or `"right"`.
urdf (`str`):
Path to the exoskeleton arm's URDF file.
root (`str`):
Name of the exoskeleton's root node in the meshcat scene tree.
g1_ee (`str`):
Name of the corresponding end-effector frame on the G1 URDF model.
offset (`np.ndarray`):
World-frame translation applied to the exoskeleton and its IK target, so the exoskeleton's
visualization does not overlap the G1's.
marker_prefix (`str`):
Prefix used to namespace this arm's meshcat marker paths.
"""
side: str # "left" | "right"
urdf: str # exo_left.urdf / exo_right.urdf
root: str # "exo_left" / "exo_right"
@@ -52,12 +71,28 @@ class ArmCfg:
class Markers:
"""Creates meshcat visualization primitives, showing end-effector frames of exoskeleton and G1"""
"""Creates meshcat visualization primitives, showing end-effector frames of exoskeleton and G1."""
def __init__(self, viewer):
"""Store the meshcat viewer (or scene-tree node) markers will be attached under.
Args:
viewer:
The meshcat viewer markers are added to.
"""
self.v = viewer
def sphere(self, path: str, r: float, rgba: tuple[float, float, float, float]):
"""Add a colored sphere marker to the meshcat scene.
Args:
path (`str`):
Meshcat scene-tree path for this marker, e.g. `"markers/left_exo_ee"`.
r (`float`):
Sphere radius, in meters.
rgba (`tuple[float, float, float, float]`):
Red, green, and blue components (each 0-1) followed by opacity (0-1).
"""
import meshcat.geometry as mg
c = (int(rgba[0] * 255) << 16) | (int(rgba[1] * 255) << 8) | int(rgba[2] * 255)
@@ -67,6 +102,16 @@ class Markers:
)
def axes(self, path: str, axis_len: float = 0.1, axis_w: int = 6):
"""Add a red/green/blue XYZ axis-triad marker to the meshcat scene.
Args:
path (`str`):
Meshcat scene-tree path for this marker.
axis_len (`float`, *optional*, defaults to 0.1):
Length of each axis line, in meters.
axis_w (`int`, *optional*, defaults to 6):
Line width, in pixels.
"""
import meshcat.geometry as mg
pts = np.array(
@@ -85,21 +130,37 @@ class Markers:
)
def tf(self, path: str, mat: np.ndarray):
"""Update the transform of an existing marker.
Args:
path (`str`):
Meshcat scene-tree path of the marker to move.
mat (`np.ndarray`):
New 4x4 homogeneous transform for the marker, in world frame.
"""
self.v[path].set_transform(mat)
class ExoskeletonIKHelper:
"""
- Loads G1 robot and exoskeleton URDF models via Pinocchio
- Computes forward kinematics on exoskeleton to get end-effector poses
- Solves inverse kinematics on G1 to match those poses
- Provides meshcat visualization showing both robots and targets
"""Maps exoskeleton joint angles to G1 arm joint angles via forward and inverse kinematics.
Loads the G1 robot and both exoskeleton arm URDF models via Pinocchio, computes forward kinematics on
the exoskeleton to obtain end-effector poses in the world frame, then solves inverse kinematics on the
G1 model to find joint angles reproducing those poses. Also provides an optional meshcat
visualization showing both robots alongside their IK targets.
Args:
frozen_joints: List of G1 joint names to exclude from IK (kept at neutral).
frozen_joints (`list[str] | None`, *optional*):
G1 joint names to exclude from IK; these are held at their current pose instead of being
solved for.
"""
def __init__(self, frozen_joints: list[str] | None = None):
"""Load the G1 and exoskeleton Pinocchio models and precompute frozen-joint indices.
Raises:
ImportError: If `pinocchio` is not installed.
"""
try:
import pinocchio as pin
except ImportError as e:
@@ -188,9 +249,9 @@ class ExoskeletonIKHelper:
logger.info(f"loaded {a.side} exo urdf: {a.urdf}")
def init_visualization(self):
"""
Creates a browser-based visualization of exoskeleton and G1 robot,
highlighting end-effector frames and target positions.
"""Creates a browser-based visualization of exoskeleton and G1 robot.
Highlights end-effector frames and target positions.
"""
try:
from pinocchio.visualize import MeshcatVisualizer
@@ -237,7 +298,7 @@ class ExoskeletonIKHelper:
print(f"\nmeshcat url: {self.viewer.url()}\n")
def _fk_target_world(self, side: str, angles: dict[str, float]) -> np.ndarray | None:
"""returns wrist frame target to be used for G1 IK in 4x4 homogeneous transform. Takes offset into account."""
"""Returns wrist frame target to be used for G1 IK in 4x4 homogeneous transform. Takes offset into account."""
if side not in self.exo or not angles:
return None
@@ -263,6 +324,10 @@ class ExoskeletonIKHelper:
return target
def update_visualization(self):
"""Refresh the meshcat scene with the G1's and both exoskeletons' current poses and IK targets.
No-op if `init_visualization` has not been called yet.
"""
if self.viewer is None or self.markers is None:
return
@@ -311,9 +376,9 @@ class ExoskeletonIKHelper:
left_angles: dict[str, float],
right_angles: dict[str, float],
) -> dict[str, float]:
"""
Performs FK on exoskeleton to get end-effector poses in world frame,
after which it solves IK on G1 to return joint angles matching those poses in G1 motor order.
"""Performs FK on exoskeleton to get end-effector poses in world frame.
Solves IK on G1 to return joint angles matching those poses in G1 motor order.
"""
pin = self.pin
@@ -35,6 +35,17 @@ logger = logging.getLogger(__name__)
def parse_raw16(line: bytes) -> list[int] | None:
"""Parse one line of exoskeleton telemetry into 16 raw ADC channel readings.
Args:
line (`bytes`):
One raw line read from the exoskeleton's serial port, expected to contain 16
whitespace-separated integers (sin/cos pairs for each sensed joint, plus joystick channels).
Returns:
`list[int] | None`: The 16 raw ADC values in channel order, or `None` if the line is malformed or
has fewer than 16 values.
"""
try:
parts = line.decode("utf-8", errors="ignore").split()
if len(parts) < 16:
@@ -45,7 +56,18 @@ def parse_raw16(line: bytes) -> list[int] | None:
def read_raw_from_serial(ser) -> list[int] | None:
"""Read latest sample from serial; if buffer is backed up, keep only the newest."""
"""Read the latest sample from serial; if the input buffer is backed up, keep only the newest.
Draining the buffer down to the newest line keeps teleoperation responsive to the exoskeleton's
current pose instead of replaying a queue of stale samples.
Args:
ser (`serial.Serial`):
Open serial connection to the exoskeleton's sensor board.
Returns:
`list[int] | None`: The most recently parsed sample, or `None` if no valid line was available.
"""
try:
last = None
while ser.in_waiting > 0:
@@ -67,6 +89,27 @@ def read_raw_from_serial(ser) -> list[int] | None:
@dataclass
class ExoskeletonArm:
"""Serial link and calibration state for one exoskeleton arm (left or right).
Wraps the raw serial connection to the arm's sensor board and converts its hall-effect sensor readings
into calibrated joint angles via `get_angles`, once a calibration has been loaded or produced by
`calibrate`.
Args:
port (`str`):
Serial port the arm's sensor board is connected to, e.g. `/dev/ttyUSB0`.
calibration_fpath (`Path`):
Path to the JSON file used to load and save this arm's calibration.
side (`str`):
Which arm this is, `"left"` or `"right"`. Used to label saved calibration data and log
messages.
baud_rate (`int`, *optional*, defaults to 115200):
Baud rate for the serial connection.
calibration (`ExoskeletonCalibration | None`, *optional*):
Calibration data for this arm. Loaded automatically from `calibration_fpath` if that file
exists; otherwise populated by calling `calibrate`.
"""
port: str
calibration_fpath: Path
side: str
@@ -76,19 +119,39 @@ class ExoskeletonArm:
calibration: ExoskeletonCalibration | None = None
def __post_init__(self):
"""Check that `pyserial` is installed and load an existing calibration file, if any."""
require_package("pyserial", extra="unitree_g1", import_name="serial")
if self.calibration_fpath.is_file():
self._load_calibration()
@property
def is_connected(self) -> bool:
"""Whether the serial connection to the arm's sensor board is open.
Returns:
`bool`: `True` if the serial port has been opened and not yet closed.
"""
return self._ser is not None and getattr(self._ser, "is_open", False)
@property
def is_calibrated(self) -> bool:
"""Whether calibration data is available for this arm.
Returns:
`bool`: `True` if a calibration has been loaded from disk or produced by `calibrate`.
"""
return self.calibration is not None
def connect(self, calibrate: bool = True) -> None:
"""Open the serial connection to the arm's sensor board.
Args:
calibrate (`bool`, *optional*, defaults to `True`):
Whether to run `calibrate` automatically after connecting if no calibration is loaded yet.
Raises:
ConnectionError: If the serial port cannot be opened.
"""
if self.is_connected:
return
try:
@@ -102,6 +165,7 @@ class ExoskeletonArm:
self.calibrate()
def disconnect(self) -> None:
"""Close the serial connection to the arm's sensor board, if open."""
if self._ser:
try:
self._ser.close()
@@ -117,17 +181,41 @@ class ExoskeletonArm:
logger.warning(f"failed to load calibration: {e}")
def read_raw(self) -> list[int] | None:
"""Read the arm's latest raw ADC sample.
Returns:
`list[int] | None`: The 16 raw ADC channel values, or `None` if the arm is not connected or no
valid sample was available.
"""
if not self._ser:
return None
return read_raw_from_serial(self._ser)
def get_angles(self) -> dict[str, float]:
"""Read the arm's current sensor sample and convert it to calibrated joint angles.
Returns:
`dict[str, float]`: Joint name to angle in radians, or an empty dict if no sample was
available on the serial link.
Raises:
RuntimeError: If the arm has not been calibrated yet.
"""
if not self.calibration:
raise RuntimeError("exoskeleton not calibrated")
raw = self.read_raw()
return {} if raw is None else exo_raw_to_angles(raw, self.calibration)
def calibrate(self) -> None:
"""Run the interactive per-joint calibration procedure and store its result.
Delegates to `run_exo_calibration`, which walks the operator through moving each joint through
its range and holding a zero pose, then saves the resulting ellipse fits and zero offsets to
`calibration_fpath`.
Raises:
RuntimeError: If the arm is not connected.
"""
if not self.is_connected:
raise RuntimeError("Cannot calibrate: exoskeleton not connected")
self.calibration = run_exo_calibration(self._ser, self.side, self.calibration_fpath)
@@ -28,7 +28,18 @@ if TYPE_CHECKING or _unitree_sdk_available:
else:
class Joystick:
"""Placeholder used when `unitree_sdk2py` is not installed.
Raises `ImportError` on instantiation instead of on import, so the module can still be imported
(and its non-hardware members inspected) without the SDK present.
"""
def __init__(self):
"""Raise `ImportError` because `unitree_sdk2py` is required and not installed.
Raises:
ImportError: Always.
"""
raise ImportError(
"unitree_sdk2py is required for RemoteController. Install with: pip install unitree_sdk2py"
)
@@ -74,6 +85,7 @@ class RemoteController:
]
def __init__(self):
"""Initialize joystick axes, button state, and joystick-center calibration to their defaults."""
self.lx = 0.0
self.ly = 0.0
self.rx = 0.0
@@ -102,6 +114,19 @@ class RemoteController:
self.remote_action.update(zip(REMOTE_AXES, (self.lx, self.ly, self.rx, self.ry), strict=True))
def calibrate_center(self, raw16: list[int] | None, side: str) -> None:
"""Detect and record the center position of one side's exoskeleton-mounted joystick.
Meant to be called once at connect time. If the joystick's button ADC channel reads above
half-scale, an exoskeleton joystick is assumed present on that side, and its current X/Y ADC
reading is stored as the neutral center used by `set_from_exo`.
Args:
raw16 (`list[int] | None`):
The 16 raw ADC channel values read from the exoskeleton's sensor board, or `None` if no
sample was available.
side (`str`):
Which joystick to calibrate, `"left"` or `"right"`.
"""
if raw16 is None or len(raw16) < 16:
logger.info(f"{side.capitalize()} exo joystick: no data available")
return
@@ -123,6 +148,17 @@ class RemoteController:
logger.info(f"{side.capitalize()} exo joystick enabled, center: x={x}, y={y}")
def set_from_exo(self, raw16: list[int] | None, side: str) -> None:
"""Update one side's joystick axes and button from the exoskeleton-mounted joystick, if calibrated.
No-op if `calibrate_center` did not detect an exoskeleton joystick on that side.
Args:
raw16 (`list[int] | None`):
The 16 raw ADC channel values read from the exoskeleton's sensor board, or `None` if no
sample was available.
side (`str`):
Which joystick to update, `"left"` or `"right"`.
"""
if raw16 is None or len(raw16) < 16:
return
@@ -157,17 +193,39 @@ class RemoteController:
class UnitreeG1Teleoperator(Teleoperator):
"""
Bimanual exoskeleton arms teleoperator for Unitree G1 arms.
"""Bimanual exoskeleton-arm teleoperator for the Unitree G1 humanoid, plus its wireless remote.
Uses inverse kinematics: exoskeleton FK computes end-effector pose,
G1 IK solves for joint angles.
Two exoskeleton arms worn by the operator report joint angles, which are converted to a G1 arm action
via forward kinematics on the exoskeleton followed by inverse kinematics on the G1 (see
[`~teleoperators.unitree_g1.exo_ik.ExoskeletonIKHelper`]). A Unitree wireless remote (or an
exoskeleton-mounted joystick, when the remote is idle) supplies additional axes, typically used for
locomotion. If neither exoskeleton arm has a configured serial port, the teleoperator falls back to
remote-controller-only mode and reports no arm joint actions.
Example:
```python
>>> from lerobot.teleoperators.unitree_g1 import UnitreeG1Teleoperator, UnitreeG1TeleoperatorConfig
>>> teleop = UnitreeG1Teleoperator(UnitreeG1TeleoperatorConfig()) # doctest: +SKIP
>>> with teleop: # doctest: +SKIP
... action = teleop.get_action()
```
"""
config_class = UnitreeG1TeleoperatorConfig
name = "unitree_g1"
def __init__(self, config: UnitreeG1TeleoperatorConfig):
"""Build the teleoperator from its configuration.
Args:
config (`UnitreeG1TeleoperatorConfig`):
The teleoperator's configuration. Exoskeleton arm control is enabled only if both
`left_arm_config.port` and `right_arm_config.port` are set; leaving both empty runs in
remote-controller-only mode.
Raises:
ValueError: If exactly one of the two arm ports is configured.
"""
super().__init__(config)
self.config = config
left_exo_enabled = bool(config.left_arm_config.port.strip())
@@ -208,6 +266,15 @@ class UnitreeG1Teleoperator(Teleoperator):
@cached_property
def action_features(self) -> dict[str, type]:
"""Keys the teleoperator's actions are reported under.
Includes one `"<joint>.q"` key per G1 arm joint (radians) when both exoskeleton arms are
configured, plus the remote controller's stick and button axes. See
[`~teleoperators.Teleoperator.action_features`].
Returns:
`dict[str, type]`: Action names mapped to `float`.
"""
remote_features = dict.fromkeys(self.remote_controller.remote_action, float)
if not self._arm_control_enabled:
return remote_features
@@ -216,21 +283,48 @@ class UnitreeG1Teleoperator(Teleoperator):
@cached_property
def feedback_features(self) -> dict[str, type]:
"""Same as [`~teleoperators.Teleoperator.feedback_features`].
Returns:
`dict[str, type]`: A single `"wireless_remote"` key mapped to `bytes`, the raw Unitree
wireless remote packet to be parsed into joystick and button state.
"""
return {"wireless_remote": bytes}
@property
def is_connected(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_connected`].
Returns:
`bool`: `True` if exoskeleton arm control is disabled (remote-only mode), or if both
exoskeleton arms are connected.
"""
if not self._arm_control_enabled:
return True
return self.left_arm.is_connected and self.right_arm.is_connected
@property
def is_calibrated(self) -> bool:
"""Same as [`~teleoperators.Teleoperator.is_calibrated`].
Returns:
`bool`: `True` if exoskeleton arm control is disabled (remote-only mode), or if both
exoskeleton arms are calibrated.
"""
if not self._arm_control_enabled:
return True
return self.left_arm.is_calibrated and self.right_arm.is_calibrated
def connect(self, calibrate: bool = True) -> None:
"""Connect both exoskeleton arms, build the IK helper, and calibrate the remote's joystick centers.
If neither exoskeleton arm has a configured serial port, this is a no-op and the teleoperator
falls back to reporting only remote-controller actions.
Args:
calibrate (`bool`, *optional*, defaults to `True`):
Whether to calibrate each exoskeleton arm that is not yet calibrated.
"""
if not self._arm_control_enabled:
logger.warning("Exo ports not fully configured; teleop will send joystick only (no arm actions)")
return
@@ -250,6 +344,12 @@ class UnitreeG1Teleoperator(Teleoperator):
self.remote_controller.calibrate_center(right_raw, "right")
def calibrate(self) -> None:
"""Calibrate each exoskeleton arm that is not already calibrated, then verify tracking visually.
See [`~teleoperators.Teleoperator.calibrate`]. After both arms are calibrated, this opens the
interactive meshcat visualization (see `run_visualization_loop`) so the operator can confirm the
G1 arms track the exoskeleton before recording data.
"""
if not self.left_arm.is_calibrated:
logger.info("Starting calibration for left arm...")
self.left_arm.calibrate()
@@ -266,9 +366,27 @@ class UnitreeG1Teleoperator(Teleoperator):
self.run_visualization_loop()
def configure(self) -> None:
"""No-op: the exoskeleton arms require no runtime configuration beyond calibration.
See [`~teleoperators.Teleoperator.configure`].
"""
pass
def get_action(self) -> dict[str, float]:
"""Read both exoskeleton arms and the remote controller, and combine them into one action.
Exoskeleton joint angles are converted to G1 arm joint angles by forward kinematics on the
exoskeleton followed by inverse kinematics on the G1, via
[`~teleoperators.unitree_g1.exo_ik.ExoskeletonIKHelper.compute_g1_joints_from_exo`]. The wireless
remote takes priority over the exoskeleton-mounted joystick for stick/button axes whenever it
reports a non-zero stick or a pressed button; otherwise the exoskeleton-mounted joystick (if
calibrated) is used instead.
Returns:
`dict[str, float]`: G1 arm joint angles (`"<joint>.q"`, radians) when exoskeleton control is
enabled, merged with the remote controller's stick and button axes. Matches
[`~teleoperators.Teleoperator.action_features`].
"""
joint_action = {}
left_raw = None
right_raw = None
@@ -293,11 +411,19 @@ class UnitreeG1Teleoperator(Teleoperator):
return {**joint_action, **rc.remote_action}
def send_feedback(self, feedback: dict[str, Any]) -> None:
"""Update the remote controller's parsed state from a raw wireless remote packet.
Args:
feedback (`dict[str, Any]`):
Feedback dict; only the `"wireless_remote"` key (raw bytes) is used, matching
[`~teleoperators.Teleoperator.feedback_features`]. Ignored if the key is absent.
"""
wireless_remote = feedback.get("wireless_remote")
if wireless_remote is not None:
self.remote_controller.set_from_wireless(wireless_remote)
def disconnect(self) -> None:
"""Disconnect both exoskeleton arms. See [`~teleoperators.Teleoperator.disconnect`]."""
self.left_arm.disconnect()
self.right_arm.disconnect()
+13
View File
@@ -34,6 +34,19 @@ class TeleopEvents(Enum):
def make_teleoperator_from_config(config: TeleoperatorConfig) -> "Teleoperator":
"""Instantiate the [`~teleoperators.Teleoperator`] matching a config's registered [`~teleoperators.TeleoperatorConfig.type`].
Args:
config (`TeleoperatorConfig`):
Configuration of the teleoperator to build.
Returns:
`Teleoperator`: The instantiated teleoperator, not yet connected.
Raises:
ValueError: If the config's type is not a known teleoperator and building it via the generic
device factory also fails.
"""
# TODO(Steven): Consider just using the make_device_from_device_class for all types
if config.type == "keyboard":
from .keyboard import KeyboardTeleop
+1 -1
View File
@@ -60,7 +60,7 @@ PATH_TO_LEROBOT = PATH_TO_REPO / "src" / "lerobot"
# Modules whose public objects are checked. Add a module here once its docstrings follow the standard.
MODULES_TO_CHECK = [
"lerobot.robots",
"lerobot.datasets",
"lerobot.teleoperators",
]
# Objects that do not yet follow the standard, so the check can be green from day one. Removing an entry
+2
View File
@@ -15,4 +15,6 @@ src/lerobot/robots/robot.py
src/lerobot/robots/so_follower/config_so_follower.py
src/lerobot/robots/so_follower/so_follower.py
src/lerobot/robots/utils.py
src/lerobot/teleoperators/phone/config_phone.py
src/lerobot/teleoperators/teleoperator.py
src/lerobot/teleoperators/unitree_g1/unitree_g1.py