mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 35d40f353f |
@@ -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
|
||||
|
||||
@@ -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
@@ -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"]
|
||||
|
||||
@@ -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,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".
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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.
|
||||
"""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user