mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 13:09:40 +00:00
Compare commits
99 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 4cd9fc73f8 | |||
| 57d33c295e | |||
| 6e23429105 | |||
| fffa42cc5e | |||
| f33089b027 | |||
| 9f32f57b59 | |||
| 0a9b9d9a93 | |||
| 16b915ede5 | |||
| 95d9029039 | |||
| bbfc4ff443 | |||
| 0c57cd03f2 | |||
| a6b06eac38 | |||
| 36b8face98 | |||
| cd8984cc0a | |||
| b9ded9e761 | |||
| 185f3e1708 | |||
| e36783253a | |||
| 289e577fc7 | |||
| 9c32722eb9 | |||
| b49cb50e01 | |||
| 962ed236af | |||
| 3ae036ee40 | |||
| dd08d4eb53 | |||
| 6e5f6df6e7 | |||
| 265abe6c79 | |||
| cdf5141688 | |||
| 9f0663e9e3 | |||
| b4e2d0b610 | |||
| 5594eba06a | |||
| 207183c2f8 | |||
| 503f3e57ae | |||
| 6d24f20eb4 | |||
| af163fd032 | |||
| 7d615acf9a | |||
| 09572babee | |||
| 35339d31e5 | |||
| f37be3edbe | |||
| 4d076845ac | |||
| 413972c812 | |||
| 0449aa02f6 | |||
| a05c0833e1 | |||
| 7b76d94c5b | |||
| ec2dbc1c98 | |||
| d526785e47 | |||
| 4af7c70664 | |||
| a855570097 | |||
| 167e22ba51 | |||
| 00c25c65c2 | |||
| 23f6d5dabd | |||
| 9b25b7fe0a | |||
| c1b6ea85d6 | |||
| ffe25afb8f | |||
| 77259f436e | |||
| 85f5c3606d | |||
| b587e81587 | |||
| 4658dada9b | |||
| 57ea6f4106 | |||
| 4209639f33 | |||
| fc7a0bc2fd | |||
| 5f6513551c | |||
| 70e157e00f | |||
| 1837be51bf | |||
| bedd56eed9 | |||
| c165e4df68 | |||
| 5e24da483a | |||
| 9c54665a76 | |||
| f6a845c30c | |||
| 45e8336854 | |||
| 5046e2df32 | |||
| 1c88e26c6d | |||
| 69a3edfa33 | |||
| 2492ce2c29 | |||
| c8e75da55f | |||
| 2eae31ea2b | |||
| c997abe739 | |||
| c73579055e | |||
| 4be438161b | |||
| 806d28a883 | |||
| 573b65ff6b | |||
| bc55713e7c | |||
| 4f53c42583 | |||
| bfced3d149 | |||
| 4969813d4e | |||
| 1c87ca31a3 | |||
| 4bcde762cc | |||
| 943ae78cfe | |||
| 3363688f1e | |||
| 0876629e72 | |||
| 305614b8c6 | |||
| 02d3202c4f | |||
| 3b6de2fdf8 | |||
| 744f3667c0 | |||
| fdde436776 | |||
| 5c683c65c6 | |||
| dfbc25c58f | |||
| 804c76bcc2 | |||
| e6afa69be9 | |||
| 31d1439e29 | |||
| 1c118c6359 |
+12
-8
@@ -61,15 +61,19 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
|
|||||||
**4.1 Install**
|
**4.1 Install**
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
pip install 'lerobot[feetech]' # SO-100/SO-101 motor stack
|
# uv (recommended — see AGENTS.md and CLAUDE.md)
|
||||||
# pip install 'lerobot[all]' # everything
|
uv sync --locked --extra feetech # SO-100/SO-101 motor stack
|
||||||
# pip install 'lerobot[aloha,pusht]' # specific features
|
# uv sync --locked --extra all # everything
|
||||||
# pip install 'lerobot[smolvla]' # add SmolVLA deps
|
# uv sync --locked --extra smolvla # add SmolVLA deps
|
||||||
git lfs install && git lfs pull
|
|
||||||
hf auth login # required to push datasets/policies
|
|
||||||
```
|
|
||||||
|
|
||||||
Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
|
# pip (alternative, e.g. when not working from source)
|
||||||
|
# pip install 'lerobot[feetech]'
|
||||||
|
# pip install 'lerobot[all]'
|
||||||
|
# pip install 'lerobot[smolvla]'
|
||||||
|
|
||||||
|
git lfs install && git lfs pull
|
||||||
|
hf auth login # required to push datasets/policies
|
||||||
|
```
|
||||||
|
|
||||||
**4.2 Find USB ports** — run once per arm, unplug when prompted.
|
**4.2 Find USB ports** — run once per arm, unplug when prompted.
|
||||||
|
|
||||||
|
|||||||
@@ -68,17 +68,16 @@ ENV HOME=/home/user_lerobot \
|
|||||||
# issues with MuJoCo and OpenGL drivers.
|
# issues with MuJoCo and OpenGL drivers.
|
||||||
RUN uv venv --python python${PYTHON_VERSION}
|
RUN uv venv --python python${PYTHON_VERSION}
|
||||||
|
|
||||||
# Install Python dependencies for caching
|
# Install third-party dependencies separately for layer caching
|
||||||
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
|
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
|
||||||
COPY --chown=user_lerobot:user_lerobot src/ src/
|
RUN uv sync --locked --extra all --no-install-project --no-cache
|
||||||
|
|
||||||
RUN uv sync --locked --extra all --no-cache
|
|
||||||
|
|
||||||
RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
|
RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
|
||||||
|
|
||||||
# Copy the rest of the application source code
|
# Copy the application source code and install the local project
|
||||||
# Make sure to have the git-LFS files for testing
|
# Make sure to have the git-LFS files for testing
|
||||||
COPY --chown=user_lerobot:user_lerobot . .
|
COPY --chown=user_lerobot:user_lerobot . .
|
||||||
|
RUN uv sync --locked --extra all --no-cache
|
||||||
|
|
||||||
# Set the default command
|
# Set the default command
|
||||||
CMD ["/bin/bash"]
|
CMD ["/bin/bash"]
|
||||||
|
|||||||
@@ -60,15 +60,14 @@ ENV HOME=/home/user_lerobot \
|
|||||||
# run other Python projects in the same container without dependency conflicts.
|
# run other Python projects in the same container without dependency conflicts.
|
||||||
RUN uv venv
|
RUN uv venv
|
||||||
|
|
||||||
# Install Python dependencies for caching
|
# Install third-party dependencies separately for layer caching
|
||||||
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
|
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
|
||||||
COPY --chown=user_lerobot:user_lerobot src/ src/
|
RUN uv sync --locked --extra all --no-install-project --no-cache
|
||||||
|
|
||||||
RUN uv sync --locked --extra all --no-cache
|
# Copy the application code and install the local project
|
||||||
|
|
||||||
# Copy the rest of the application code
|
|
||||||
# Make sure to have the git-LFS files for testing
|
# Make sure to have the git-LFS files for testing
|
||||||
COPY --chown=user_lerobot:user_lerobot . .
|
COPY --chown=user_lerobot:user_lerobot . .
|
||||||
|
RUN uv sync --locked --extra all --no-cache
|
||||||
|
|
||||||
# Set the default command
|
# Set the default command
|
||||||
CMD ["/bin/bash"]
|
CMD ["/bin/bash"]
|
||||||
|
|||||||
@@ -58,7 +58,7 @@ final_action = postprocessor(action)
|
|||||||
|
|
||||||
## Hardware API redesign
|
## Hardware API redesign
|
||||||
|
|
||||||
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request.
|
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is an overview of what changed and how you can continue to work with datasets created before this pull request.
|
||||||
|
|
||||||
### What changed?
|
### What changed?
|
||||||
|
|
||||||
@@ -129,8 +129,8 @@ python examples/backward_compatibility/replay.py \
|
|||||||
|
|
||||||
Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.
|
Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.
|
||||||
|
|
||||||
To find these transformations, we recommend to first try and and replay an episode of the dataset your policy was trained on using the section above.
|
To find these transformations, we recommend first replaying an episode of the dataset your policy was trained on using the section above.
|
||||||
Then, add these same transformations on your inference script (shown here in the `record.py` script):
|
Then, add these same transformations to your inference script (shown here in the `record.py` script):
|
||||||
|
|
||||||
```diff
|
```diff
|
||||||
action_values = predict_action(
|
action_values = predict_action(
|
||||||
|
|||||||
@@ -136,6 +136,10 @@ config = RealSenseCameraConfig(
|
|||||||
height=480,
|
height=480,
|
||||||
color_mode=ColorMode.RGB,
|
color_mode=ColorMode.RGB,
|
||||||
use_depth=True,
|
use_depth=True,
|
||||||
|
# Optional fixed color controls. Omit them to leave the current sensor settings unchanged.
|
||||||
|
exposure=120,
|
||||||
|
gain=64,
|
||||||
|
white_balance=4600,
|
||||||
rotation=Cv2Rotation.NO_ROTATION
|
rotation=Cv2Rotation.NO_ROTATION
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -154,6 +158,15 @@ finally:
|
|||||||
```
|
```
|
||||||
<!-- prettier-ignore-end -->
|
<!-- prettier-ignore-end -->
|
||||||
|
|
||||||
|
Manual color controls disable the corresponding automatic exposure or white-balance mode. Their
|
||||||
|
supported ranges vary by camera model; an invalid value raises an error at connection time that
|
||||||
|
includes the range reported by the sensor. Requesting an unsupported control also raises an error.
|
||||||
|
Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
|
||||||
|
require `use_rgb=True`.
|
||||||
|
|
||||||
|
On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
|
||||||
|
exposure or gain also affects the depth stream.
|
||||||
|
|
||||||
</hfoption>
|
</hfoption>
|
||||||
</hfoptions>
|
</hfoptions>
|
||||||
|
|
||||||
|
|||||||
@@ -88,20 +88,6 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
|
|||||||
|
|
||||||
The same policy can work with different environment processors, and the same environment processor can work with different policies:
|
The same policy can work with different environment processors, and the same environment processor can work with different policies:
|
||||||
|
|
||||||
````python
|
|
||||||
# Use SmolVLA policy with LIBERO environment
|
|
||||||
# Use SmolVLA policy with LIBERO environment
|
|
||||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
|
||||||
env_cfg=libero_cfg,
|
|
||||||
policy_cfg=smolvla_cfg,
|
|
||||||
)
|
|
||||||
smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
|
|
||||||
# Or use ACT policy with the same LIBERO environment
|
|
||||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
|
||||||
env_cfg=libero_cfg,
|
|
||||||
policy_cfg=act_cfg,
|
|
||||||
)
|
|
||||||
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
|
||||||
```python
|
```python
|
||||||
# Use SmolVLA policy with LIBERO environment
|
# Use SmolVLA policy with LIBERO environment
|
||||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||||
@@ -116,6 +102,7 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
|||||||
policy_cfg=act_cfg,
|
policy_cfg=act_cfg,
|
||||||
)
|
)
|
||||||
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
||||||
|
```
|
||||||
|
|
||||||
### 3. **Easier Experimentation**
|
### 3. **Easier Experimentation**
|
||||||
|
|
||||||
@@ -145,7 +132,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
|
|||||||
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
|
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
|
||||||
gripper_pos, gripper_vel], dim=-1) # 14D
|
gripper_pos, gripper_vel], dim=-1) # 14D
|
||||||
return state
|
return state
|
||||||
````
|
```
|
||||||
|
|
||||||
### 4. **Cleaner Environment Code**
|
### 4. **Cleaner Environment Code**
|
||||||
|
|
||||||
|
|||||||
@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
|
|||||||
For each motor you want to update:
|
For each motor you want to update:
|
||||||
|
|
||||||
1. **Select the motor** from the list by clicking on it
|
1. **Select the motor** from the list by clicking on it
|
||||||
2. **Click on Upgrade tab**:
|
2. **Click the Upgrade tab**:
|
||||||
3. **Click on Online button**:
|
3. **Click the Online button**:
|
||||||
- If an potential firmware update is found, it will be displayed in the box
|
- If a potential firmware update is found, it will be displayed in the box
|
||||||
4. **Click on Upgrade button**:
|
4. **Click the Upgrade button**:
|
||||||
- The update progress will be displayed
|
- The update progress will be displayed
|
||||||
|
|
||||||
## Step 6: Verify Update
|
## Step 6: Verify Update
|
||||||
|
|||||||
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
|
|||||||
|
|
||||||
### Record
|
### Record
|
||||||
|
|
||||||
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
|
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
lerobot-record \
|
lerobot-record \
|
||||||
|
|||||||
@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
|
|||||||
- [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
|
- [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
|
||||||
|
|
||||||
Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
|
Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
|
||||||
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
|
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
|
||||||
|
|
||||||
Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
|
Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
|
||||||
|
|
||||||
|
|||||||
@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
|
|||||||
pip install -e ".[lekiwi]"
|
pip install -e ".[lekiwi]"
|
||||||
```
|
```
|
||||||
|
|
||||||
Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
|
Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
|
||||||
Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
|
Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
|
||||||
|
|
||||||
# Step-by-Step Assembly Instructions
|
# Step-by-Step Assembly Instructions
|
||||||
|
|||||||
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
|
|||||||
|
|
||||||
## Reproducing π₀Fast results
|
## Reproducing π₀Fast results
|
||||||
|
|
||||||
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
|
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
|
||||||
|
|
||||||
The finetuned model can be found here:
|
The finetuned model can be found here:
|
||||||
|
|
||||||
|
|||||||
@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
|
|||||||
## Three pipelines
|
## Three pipelines
|
||||||
|
|
||||||
We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
|
We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
|
||||||
Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
|
Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
|
||||||
|
|
||||||
1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
|
1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
|
||||||
2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
|
2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
|
||||||
@@ -74,15 +74,15 @@ In the phone to SO-100 follower examples we use the following adapters:
|
|||||||
- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
|
- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
|
||||||
- `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
|
- `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
|
||||||
- `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
|
- `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
|
||||||
- `transition_to_observation`: transforms the pipeline transition to a observation dict.
|
- `transition_to_observation`: transforms the pipeline transition to an observation dict.
|
||||||
|
|
||||||
Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
|
Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
|
||||||
|
|
||||||
## Dataset feature contracts
|
## Dataset feature contracts
|
||||||
|
|
||||||
Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
|
Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
|
||||||
|
|
||||||
Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
|
Below is an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
|
||||||
|
|
||||||
```python
|
```python
|
||||||
def transform_features(
|
def transform_features(
|
||||||
|
|||||||
+2
-2
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
|
|||||||
policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
|
policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
|
||||||
|
|
||||||
# Now use predict_action_chunk with RTC parameters
|
# Now use predict_action_chunk with RTC parameters
|
||||||
inference_delay = 4 # How many steps of inference latency, this values should be calculated based on the inference latency of the policy
|
inference_delay = 4 # How many steps of inference latency, this value should be calculated based on the inference latency of the policy
|
||||||
|
|
||||||
# Initialize the action queue
|
# Initialize the action queue
|
||||||
action_queue = ActionQueue(policy_cfg.rtc_config)
|
action_queue = ActionQueue(policy_cfg.rtc_config)
|
||||||
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
|
|||||||
RTCConfig(execution_horizon=10)
|
RTCConfig(execution_horizon=10)
|
||||||
```
|
```
|
||||||
|
|
||||||
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value.
|
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is an optimal value.
|
||||||
|
|
||||||
**`prefix_attention_schedule`**: How to weight consistency across the overlap region.
|
**`prefix_attention_schedule`**: How to weight consistency across the overlap region.
|
||||||
|
|
||||||
|
|||||||
@@ -93,7 +93,7 @@ lerobot-train --help
|
|||||||
|
|
||||||
## Evaluate the finetuned model and run it in real-time
|
## Evaluate the finetuned model and run it in real-time
|
||||||
|
|
||||||
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
|
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
|
||||||
Once you are logged in, you can run inference in your setup by doing:
|
Once you are logged in, you can run inference in your setup by doing:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
|
|||||||
@@ -50,11 +50,11 @@ lerobot-edit-dataset \
|
|||||||
Divide a dataset into multiple subsets.
|
Divide a dataset into multiple subsets.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# Split by fractions (e.g. 80% train, 20% test, 20% val)
|
# Split by fractions (e.g. 60% train, 20% val, 20% test)
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht \
|
--repo_id lerobot/pusht \
|
||||||
--operation.type split \
|
--operation.type split \
|
||||||
--operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}'
|
--operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'
|
||||||
|
|
||||||
# Split by specific episode indices
|
# Split by specific episode indices
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
|
|||||||
@@ -494,6 +494,19 @@ ignore_errors = true
|
|||||||
module = "lerobot.envs.*"
|
module = "lerobot.envs.*"
|
||||||
ignore_errors = false
|
ignore_errors = false
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "lerobot.annotations.*"
|
||||||
|
ignore_errors = false
|
||||||
|
disallow_untyped_defs = true
|
||||||
|
disallow_incomplete_defs = true
|
||||||
|
check_untyped_defs = true
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "lerobot.transforms.*"
|
||||||
|
ignore_errors = false
|
||||||
|
disallow_untyped_defs = true
|
||||||
|
disallow_incomplete_defs = true
|
||||||
|
check_untyped_defs = true
|
||||||
|
|
||||||
# [[tool.mypy.overrides]]
|
# [[tool.mypy.overrides]]
|
||||||
# module = "lerobot.utils.*"
|
# module = "lerobot.utils.*"
|
||||||
|
|||||||
@@ -120,14 +120,22 @@ class OpenCVCamera(Camera):
|
|||||||
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
||||||
self.backend: int = config.backend
|
self.backend: int = config.backend
|
||||||
|
|
||||||
if self.height and self.width:
|
self.capture_width: int | None = None
|
||||||
self.capture_width, self.capture_height = self.width, self.height
|
self.capture_height: int | None = None
|
||||||
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
|
self._reset_connection_settings()
|
||||||
self.capture_width, self.capture_height = self.height, self.width
|
|
||||||
|
|
||||||
def __str__(self) -> str:
|
def __str__(self) -> str:
|
||||||
return f"{self.__class__.__name__}({self.index_or_path})"
|
return f"{self.__class__.__name__}({self.index_or_path})"
|
||||||
|
|
||||||
|
def _reset_connection_settings(self) -> None:
|
||||||
|
"""Restore settings that may have been auto-detected during a failed connection."""
|
||||||
|
self.fps = self.config.fps
|
||||||
|
self.width = self.config.width
|
||||||
|
self.height = self.config.height
|
||||||
|
self.capture_width, self.capture_height = self.width, self.height
|
||||||
|
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
|
||||||
|
self.capture_width, self.capture_height = self.height, self.width
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_connected(self) -> bool:
|
def is_connected(self) -> bool:
|
||||||
"""Checks if the camera is currently connected and opened."""
|
"""Checks if the camera is currently connected and opened."""
|
||||||
@@ -164,17 +172,25 @@ class OpenCVCamera(Camera):
|
|||||||
f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
|
f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
|
||||||
)
|
)
|
||||||
|
|
||||||
self._configure_capture_settings()
|
try:
|
||||||
self._start_read_thread()
|
self._configure_capture_settings()
|
||||||
|
self._start_read_thread()
|
||||||
|
|
||||||
if warmup and self.warmup_s > 0:
|
if warmup and self.warmup_s > 0:
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
while time.time() - start_time < self.warmup_s:
|
while time.time() - start_time < self.warmup_s:
|
||||||
self.async_read(timeout_ms=self.warmup_s * 1000)
|
self.async_read(timeout_ms=self.warmup_s * 1000)
|
||||||
time.sleep(0.1)
|
time.sleep(0.1)
|
||||||
with self.frame_lock:
|
with self.frame_lock:
|
||||||
if self.latest_frame is None:
|
if self.latest_frame is None:
|
||||||
raise ConnectionError(f"{self} failed to capture frames during warmup.")
|
raise ConnectionError(f"{self} failed to capture frames during warmup.")
|
||||||
|
except BaseException:
|
||||||
|
try:
|
||||||
|
self._cleanup_resources()
|
||||||
|
except Exception:
|
||||||
|
logger.exception(f"Failed to fully clean up {self} after connect() failed.")
|
||||||
|
self._reset_connection_settings()
|
||||||
|
raise
|
||||||
|
|
||||||
logger.info(f"{self} connected.")
|
logger.info(f"{self} connected.")
|
||||||
|
|
||||||
@@ -312,32 +328,36 @@ class OpenCVCamera(Camera):
|
|||||||
|
|
||||||
for target in targets_to_scan:
|
for target in targets_to_scan:
|
||||||
camera = cv2.VideoCapture(target)
|
camera = cv2.VideoCapture(target)
|
||||||
if camera.isOpened():
|
try:
|
||||||
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
|
if camera.isOpened():
|
||||||
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||||
default_fps = camera.get(cv2.CAP_PROP_FPS)
|
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||||
default_format = camera.get(cv2.CAP_PROP_FORMAT)
|
default_fps = camera.get(cv2.CAP_PROP_FPS)
|
||||||
|
default_format = camera.get(cv2.CAP_PROP_FORMAT)
|
||||||
|
|
||||||
# Get FOURCC code and convert to string
|
# Get FOURCC code and convert to string
|
||||||
default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
|
default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
|
||||||
default_fourcc_code_int = int(default_fourcc_code)
|
default_fourcc_code_int = int(default_fourcc_code)
|
||||||
default_fourcc = "".join([chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)])
|
default_fourcc = "".join(
|
||||||
|
[chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)]
|
||||||
|
)
|
||||||
|
|
||||||
camera_info = {
|
camera_info = {
|
||||||
"name": f"OpenCV Camera @ {target}",
|
"name": f"OpenCV Camera @ {target}",
|
||||||
"type": "OpenCV",
|
"type": "OpenCV",
|
||||||
"id": target,
|
"id": target,
|
||||||
"backend_api": camera.getBackendName(),
|
"backend_api": camera.getBackendName(),
|
||||||
"default_stream_profile": {
|
"default_stream_profile": {
|
||||||
"format": default_format,
|
"format": default_format,
|
||||||
"fourcc": default_fourcc,
|
"fourcc": default_fourcc,
|
||||||
"width": default_width,
|
"width": default_width,
|
||||||
"height": default_height,
|
"height": default_height,
|
||||||
"fps": default_fps,
|
"fps": default_fps,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
found_cameras_info.append(camera_info)
|
found_cameras_info.append(camera_info)
|
||||||
|
finally:
|
||||||
camera.release()
|
camera.release()
|
||||||
|
|
||||||
return found_cameras_info
|
return found_cameras_info
|
||||||
@@ -496,6 +516,26 @@ class OpenCVCamera(Camera):
|
|||||||
self.latest_timestamp = None
|
self.latest_timestamp = None
|
||||||
self.new_frame_event.clear()
|
self.new_frame_event.clear()
|
||||||
|
|
||||||
|
def _cleanup_resources(self) -> None:
|
||||||
|
"""Stop background reads and release the capture, including after partial setup."""
|
||||||
|
read_thread = self.thread
|
||||||
|
videocapture = self.videocapture
|
||||||
|
|
||||||
|
try:
|
||||||
|
self._stop_read_thread()
|
||||||
|
finally:
|
||||||
|
self.videocapture = None
|
||||||
|
try:
|
||||||
|
if videocapture is not None:
|
||||||
|
videocapture.release()
|
||||||
|
finally:
|
||||||
|
# Releasing the device may unblock a hardware read that outlived
|
||||||
|
# the first bounded join in _stop_read_thread().
|
||||||
|
if read_thread is not None and read_thread.is_alive():
|
||||||
|
read_thread.join(timeout=2.0)
|
||||||
|
if read_thread.is_alive(): # pragma: no cover
|
||||||
|
logger.warning(f"{self} read thread remained alive after releasing the capture.")
|
||||||
|
|
||||||
@check_if_not_connected
|
@check_if_not_connected
|
||||||
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
|
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
|
||||||
"""
|
"""
|
||||||
@@ -586,16 +626,6 @@ class OpenCVCamera(Camera):
|
|||||||
if not self.is_connected and self.thread is None:
|
if not self.is_connected and self.thread is None:
|
||||||
raise DeviceNotConnectedError(f"{self} not connected.")
|
raise DeviceNotConnectedError(f"{self} not connected.")
|
||||||
|
|
||||||
if self.thread is not None:
|
self._cleanup_resources()
|
||||||
self._stop_read_thread()
|
|
||||||
|
|
||||||
if self.videocapture is not None:
|
|
||||||
self.videocapture.release()
|
|
||||||
self.videocapture = None
|
|
||||||
|
|
||||||
with self.frame_lock:
|
|
||||||
self.latest_frame = None
|
|
||||||
self.latest_timestamp = None
|
|
||||||
self.new_frame_event.clear()
|
|
||||||
|
|
||||||
logger.info(f"{self} disconnected.")
|
logger.info(f"{self} disconnected.")
|
||||||
|
|||||||
@@ -121,6 +121,9 @@ class RealSenseCamera(Camera):
|
|||||||
|
|
||||||
self.config = config
|
self.config = config
|
||||||
|
|
||||||
|
self.width: int | None = config.width
|
||||||
|
self.height: int | None = config.height
|
||||||
|
|
||||||
if config.serial_number_or_name.isdigit():
|
if config.serial_number_or_name.isdigit():
|
||||||
self.serial_number = config.serial_number_or_name
|
self.serial_number = config.serial_number_or_name
|
||||||
else:
|
else:
|
||||||
@@ -131,6 +134,9 @@ class RealSenseCamera(Camera):
|
|||||||
self.use_rgb = config.use_rgb
|
self.use_rgb = config.use_rgb
|
||||||
self.use_depth = config.use_depth
|
self.use_depth = config.use_depth
|
||||||
self.warmup_s = config.warmup_s
|
self.warmup_s = config.warmup_s
|
||||||
|
self.exposure: int | None = config.exposure
|
||||||
|
self.gain: int | None = config.gain
|
||||||
|
self.white_balance: int | None = config.white_balance
|
||||||
|
|
||||||
self.rs_pipeline: rs.pipeline | None = None
|
self.rs_pipeline: rs.pipeline | None = None
|
||||||
self.rs_profile: rs.pipeline_profile | None = None
|
self.rs_profile: rs.pipeline_profile | None = None
|
||||||
@@ -145,14 +151,23 @@ class RealSenseCamera(Camera):
|
|||||||
|
|
||||||
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
||||||
|
|
||||||
if self.height and self.width:
|
self.capture_width: int | None = None
|
||||||
self.capture_width, self.capture_height = self.width, self.height
|
self.capture_height: int | None = None
|
||||||
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
|
self._reset_connection_settings()
|
||||||
self.capture_width, self.capture_height = self.height, self.width
|
|
||||||
|
|
||||||
def __str__(self) -> str:
|
def __str__(self) -> str:
|
||||||
return f"{self.__class__.__name__}({self.serial_number})"
|
return f"{self.__class__.__name__}({self.serial_number})"
|
||||||
|
|
||||||
|
def _reset_connection_settings(self) -> None:
|
||||||
|
"""Restore settings that may have been auto-detected during a failed connection."""
|
||||||
|
self.fps = self.config.fps
|
||||||
|
self.width = self.config.width
|
||||||
|
self.height = self.config.height
|
||||||
|
self.warmup_s = self.config.warmup_s
|
||||||
|
self.capture_width, self.capture_height = self.width, self.height
|
||||||
|
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
|
||||||
|
self.capture_width, self.capture_height = self.height, self.width
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_connected(self) -> bool:
|
def is_connected(self) -> bool:
|
||||||
"""Checks if the camera pipeline is started and streams are active."""
|
"""Checks if the camera pipeline is started and streams are active."""
|
||||||
@@ -172,7 +187,8 @@ class RealSenseCamera(Camera):
|
|||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
DeviceAlreadyConnectedError: If the camera is already connected.
|
DeviceAlreadyConnectedError: If the camera is already connected.
|
||||||
ValueError: If the configuration is invalid (e.g., missing serial/name, name not unique).
|
ValueError: If the configuration is invalid, a requested sensor option is unsupported,
|
||||||
|
or a requested sensor value is invalid.
|
||||||
ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all.
|
ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all.
|
||||||
RuntimeError: If the pipeline starts but fails to apply requested settings.
|
RuntimeError: If the pipeline starts but fails to apply requested settings.
|
||||||
"""
|
"""
|
||||||
@@ -190,22 +206,31 @@ class RealSenseCamera(Camera):
|
|||||||
f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
|
f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
|
||||||
) from e
|
) from e
|
||||||
|
|
||||||
self._configure_capture_settings()
|
try:
|
||||||
self._start_read_thread()
|
self._configure_capture_settings()
|
||||||
|
self._configure_sensor_options()
|
||||||
|
self._start_read_thread()
|
||||||
|
|
||||||
# NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
|
# NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
|
||||||
self.warmup_s = max(self.warmup_s, 1)
|
self.warmup_s = max(self.warmup_s, 1)
|
||||||
|
|
||||||
warmup_read = self.async_read if self.use_rgb else self.async_read_depth
|
warmup_read = self.async_read if self.use_rgb else self.async_read_depth
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
while time.time() - start_time < self.warmup_s:
|
while time.time() - start_time < self.warmup_s:
|
||||||
warmup_read(timeout_ms=self.warmup_s * 1000)
|
warmup_read(timeout_ms=self.warmup_s * 1000)
|
||||||
time.sleep(0.1)
|
time.sleep(0.1)
|
||||||
with self.frame_lock:
|
with self.frame_lock:
|
||||||
if (self.use_rgb and self.latest_color_frame is None) or (
|
if (self.use_rgb and self.latest_color_frame is None) or (
|
||||||
self.use_depth and self.latest_depth_frame is None
|
self.use_depth and self.latest_depth_frame is None
|
||||||
):
|
):
|
||||||
raise ConnectionError(f"{self} failed to capture frames during warmup.")
|
raise ConnectionError(f"{self} failed to capture frames during warmup.")
|
||||||
|
except BaseException:
|
||||||
|
try:
|
||||||
|
self._cleanup_resources()
|
||||||
|
except Exception:
|
||||||
|
logger.exception(f"Failed to fully clean up {self} after connect() failed.")
|
||||||
|
self._reset_connection_settings()
|
||||||
|
raise
|
||||||
|
|
||||||
logger.info(f"{self} connected.")
|
logger.info(f"{self} connected.")
|
||||||
|
|
||||||
@@ -339,6 +364,111 @@ class RealSenseCamera(Camera):
|
|||||||
self.new_frame_event.clear()
|
self.new_frame_event.clear()
|
||||||
return self._async_read(timeout_ms=10000, read_depth=read_depth)
|
return self._async_read(timeout_ms=10000, read_depth=read_depth)
|
||||||
|
|
||||||
|
def _get_color_sensor(self) -> "rs.sensor":
|
||||||
|
"""Returns the sensor that controls the color stream.
|
||||||
|
|
||||||
|
Most RealSense cameras expose "RGB Camera" for color. The D405 has no
|
||||||
|
separate RGB module — its color stream comes from "Stereo Module".
|
||||||
|
We try RGB Camera first, then fall back to Stereo Module.
|
||||||
|
"""
|
||||||
|
if self.rs_profile is None:
|
||||||
|
raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
|
||||||
|
|
||||||
|
device = self.rs_profile.get_device()
|
||||||
|
sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
|
||||||
|
|
||||||
|
for name in ("RGB Camera", "Stereo Module"):
|
||||||
|
if name in sensors:
|
||||||
|
return sensors[name]
|
||||||
|
|
||||||
|
available = list(sensors.keys())
|
||||||
|
raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
|
||||||
|
|
||||||
|
def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
|
||||||
|
"""Sets a sensor option, re-raising range errors with actionable diagnostics."""
|
||||||
|
try:
|
||||||
|
sensor.set_option(option, value)
|
||||||
|
except Exception as e:
|
||||||
|
range_info = ""
|
||||||
|
try:
|
||||||
|
option_range = sensor.get_option_range(option)
|
||||||
|
range_info = (
|
||||||
|
f" (supported range: min={option_range.min}, max={option_range.max}, "
|
||||||
|
f"step={option_range.step}, default={option_range.default})"
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
range_info = " (option range unavailable)"
|
||||||
|
raise ValueError(
|
||||||
|
f"{self}: failed to set {label} to {value}{range_info}. Original error: {e}"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
def _configure_sensor_options(self) -> None:
|
||||||
|
"""Applies manual sensor options (exposure, gain, white balance) to the color sensor.
|
||||||
|
|
||||||
|
When exposure or gain is set, auto-exposure is disabled first. When white_balance
|
||||||
|
is set, auto white balance is disabled first. An omitted option is left unchanged,
|
||||||
|
and configuration is skipped entirely if all options are omitted.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If the sensor does not support a requested option or a requested
|
||||||
|
value is invalid. Invalid-value errors include the option name, requested
|
||||||
|
value, and supported range when available.
|
||||||
|
"""
|
||||||
|
if self.exposure is None and self.gain is None and self.white_balance is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
color_sensor = self._get_color_sensor()
|
||||||
|
|
||||||
|
requested_options = (
|
||||||
|
(rs.option.exposure, self.exposure, "exposure"),
|
||||||
|
(rs.option.gain, self.gain, "gain"),
|
||||||
|
(rs.option.white_balance, self.white_balance, "white balance"),
|
||||||
|
)
|
||||||
|
unsupported_options = [
|
||||||
|
label
|
||||||
|
for option, value, label in requested_options
|
||||||
|
if value is not None and not color_sensor.supports(option)
|
||||||
|
]
|
||||||
|
if unsupported_options:
|
||||||
|
raise ValueError(
|
||||||
|
f"{self}: color sensor does not support requested manual options: {unsupported_options}."
|
||||||
|
)
|
||||||
|
|
||||||
|
manual_exposure_requested = self.exposure is not None or self.gain is not None
|
||||||
|
if manual_exposure_requested:
|
||||||
|
if color_sensor.supports(rs.option.enable_auto_exposure):
|
||||||
|
self._set_sensor_option(color_sensor, rs.option.enable_auto_exposure, 0, "auto-exposure")
|
||||||
|
logger.info(f"{self} auto-exposure disabled.")
|
||||||
|
else:
|
||||||
|
logger.warning(
|
||||||
|
f"{self} sensor does not support disabling auto-exposure; "
|
||||||
|
"applying manual exposure/gain directly."
|
||||||
|
)
|
||||||
|
|
||||||
|
if self.exposure is not None:
|
||||||
|
self._set_sensor_option(color_sensor, rs.option.exposure, self.exposure, "exposure")
|
||||||
|
logger.info(f"{self} exposure set to {self.exposure}.")
|
||||||
|
|
||||||
|
if self.gain is not None:
|
||||||
|
self._set_sensor_option(color_sensor, rs.option.gain, self.gain, "gain")
|
||||||
|
logger.info(f"{self} gain set to {self.gain}.")
|
||||||
|
|
||||||
|
if self.white_balance is not None:
|
||||||
|
if color_sensor.supports(rs.option.enable_auto_white_balance):
|
||||||
|
self._set_sensor_option(
|
||||||
|
color_sensor, rs.option.enable_auto_white_balance, 0, "auto white balance"
|
||||||
|
)
|
||||||
|
logger.info(f"{self} auto white balance disabled.")
|
||||||
|
else:
|
||||||
|
logger.warning(
|
||||||
|
f"{self} sensor does not support disabling auto white balance; "
|
||||||
|
"applying manual white balance directly."
|
||||||
|
)
|
||||||
|
self._set_sensor_option(
|
||||||
|
color_sensor, rs.option.white_balance, self.white_balance, "white balance"
|
||||||
|
)
|
||||||
|
logger.info(f"{self} white balance set to {self.white_balance}.")
|
||||||
|
|
||||||
@check_if_not_connected
|
@check_if_not_connected
|
||||||
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
|
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
|
||||||
"""
|
"""
|
||||||
@@ -541,6 +671,27 @@ class RealSenseCamera(Camera):
|
|||||||
self.latest_timestamp = None
|
self.latest_timestamp = None
|
||||||
self.new_frame_event.clear()
|
self.new_frame_event.clear()
|
||||||
|
|
||||||
|
def _cleanup_resources(self) -> None:
|
||||||
|
"""Stop background reads and stop the pipeline, including after partial setup."""
|
||||||
|
read_thread = self.thread
|
||||||
|
rs_pipeline = self.rs_pipeline
|
||||||
|
|
||||||
|
try:
|
||||||
|
self._stop_read_thread()
|
||||||
|
finally:
|
||||||
|
self.rs_pipeline = None
|
||||||
|
self.rs_profile = None
|
||||||
|
try:
|
||||||
|
if rs_pipeline is not None:
|
||||||
|
rs_pipeline.stop()
|
||||||
|
finally:
|
||||||
|
# Stopping the pipeline may unblock a hardware read that outlived
|
||||||
|
# the first bounded join in _stop_read_thread().
|
||||||
|
if read_thread is not None and read_thread.is_alive():
|
||||||
|
read_thread.join(timeout=2.0)
|
||||||
|
if read_thread.is_alive(): # pragma: no cover
|
||||||
|
logger.warning(f"{self} read thread remained alive after stopping the pipeline.")
|
||||||
|
|
||||||
def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]:
|
def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]:
|
||||||
"""Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame."""
|
"""Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame."""
|
||||||
if self.thread is None or not self.thread.is_alive():
|
if self.thread is None or not self.thread.is_alive():
|
||||||
@@ -684,18 +835,5 @@ class RealSenseCamera(Camera):
|
|||||||
f"Attempted to disconnect {self}, but it appears already disconnected."
|
f"Attempted to disconnect {self}, but it appears already disconnected."
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.thread is not None:
|
self._cleanup_resources()
|
||||||
self._stop_read_thread()
|
|
||||||
|
|
||||||
if self.rs_pipeline is not None:
|
|
||||||
self.rs_pipeline.stop()
|
|
||||||
self.rs_pipeline = None
|
|
||||||
self.rs_profile = None
|
|
||||||
|
|
||||||
with self.frame_lock:
|
|
||||||
self.latest_color_frame = None
|
|
||||||
self.latest_depth_frame = None
|
|
||||||
self.latest_timestamp = None
|
|
||||||
self.new_frame_event.clear()
|
|
||||||
|
|
||||||
logger.info(f"{self} disconnected.")
|
logger.info(f"{self} disconnected.")
|
||||||
|
|||||||
@@ -46,6 +46,17 @@ class RealSenseCameraConfig(CameraConfig):
|
|||||||
use_depth: Whether to enable depth stream. Defaults to False.
|
use_depth: Whether to enable depth stream. Defaults to False.
|
||||||
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
|
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
|
||||||
warmup_s: Time reading frames before returning from connect (in seconds)
|
warmup_s: Time reading frames before returning from connect (in seconds)
|
||||||
|
exposure: Manual exposure value for the color sensor. When set, auto-exposure is
|
||||||
|
disabled and this fixed value is used. Valid ranges are camera-model specific
|
||||||
|
and reported if the value is rejected. Defaults to None (leave unchanged).
|
||||||
|
gain: Manual gain value for the color sensor. When set, auto-exposure is disabled
|
||||||
|
and this fixed gain is used, which also freezes exposure at its current value
|
||||||
|
when no exposure is configured. Valid ranges are camera-model specific and
|
||||||
|
reported if the value is rejected. Defaults to None (leave unchanged).
|
||||||
|
white_balance: Manual white balance value for the color sensor. When set, auto
|
||||||
|
white balance is disabled and this fixed value is used. Valid ranges are
|
||||||
|
camera-model specific and reported if the value is rejected. Defaults to None
|
||||||
|
(leave unchanged).
|
||||||
|
|
||||||
Note:
|
Note:
|
||||||
- Either name or serial_number must be specified.
|
- Either name or serial_number must be specified.
|
||||||
@@ -61,6 +72,9 @@ class RealSenseCameraConfig(CameraConfig):
|
|||||||
use_depth: bool = False
|
use_depth: bool = False
|
||||||
rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
|
rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
|
||||||
warmup_s: int = 1
|
warmup_s: int = 1
|
||||||
|
exposure: int | None = None
|
||||||
|
gain: int | None = None
|
||||||
|
white_balance: int | None = None
|
||||||
|
|
||||||
def __post_init__(self) -> None:
|
def __post_init__(self) -> None:
|
||||||
self.color_mode = ColorMode(self.color_mode)
|
self.color_mode = ColorMode(self.color_mode)
|
||||||
@@ -69,6 +83,18 @@ class RealSenseCameraConfig(CameraConfig):
|
|||||||
if not self.use_rgb and not self.use_depth:
|
if not self.use_rgb and not self.use_depth:
|
||||||
raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")
|
raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")
|
||||||
|
|
||||||
|
manual_color_options = {
|
||||||
|
"exposure": self.exposure,
|
||||||
|
"gain": self.gain,
|
||||||
|
"white_balance": self.white_balance,
|
||||||
|
}
|
||||||
|
configured_color_options = [name for name, value in manual_color_options.items() if value is not None]
|
||||||
|
if configured_color_options and not self.use_rgb:
|
||||||
|
raise ValueError(
|
||||||
|
"Manual color sensor options require `use_rgb=True`. "
|
||||||
|
f"Configured options: {configured_color_options}."
|
||||||
|
)
|
||||||
|
|
||||||
values = (self.fps, self.width, self.height)
|
values = (self.fps, self.width, self.height)
|
||||||
if any(v is not None for v in values) and any(v is None for v in values):
|
if any(v is not None for v in values) and any(v is None for v in values):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
|
|||||||
@@ -71,13 +71,19 @@ class DatasetRecordConfig:
|
|||||||
# Number of threads per encoder instance. None = auto (codec default).
|
# Number of threads per encoder instance. None = auto (codec default).
|
||||||
# Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
|
# Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
|
||||||
encoder_threads: int | None = None
|
encoder_threads: int | None = None
|
||||||
|
# Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
|
||||||
|
# (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
|
||||||
|
no_stamp: bool = False
|
||||||
|
|
||||||
def stamp_repo_id(self) -> None:
|
def stamp_repo_id(self) -> None:
|
||||||
"""Append a date-time tag to ``repo_id`` so each recording session gets a unique name.
|
"""Append a date-time tag to ``repo_id`` so each recording session gets a unique name.
|
||||||
|
|
||||||
Must be called explicitly at dataset *creation* time — not on resume,
|
Must be called explicitly at dataset *creation* time — not on resume,
|
||||||
where the existing ``repo_id`` (already stamped) must be preserved.
|
where the existing ``repo_id`` (already stamped) must be preserved.
|
||||||
|
No-op when ``no_stamp`` is set, preserving a user-managed ``repo_id``.
|
||||||
"""
|
"""
|
||||||
|
if self.no_stamp:
|
||||||
|
return
|
||||||
if self.repo_id:
|
if self.repo_id:
|
||||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||||
self.repo_id = f"{self.repo_id}_{timestamp}"
|
self.repo_id = f"{self.repo_id}_{timestamp}"
|
||||||
|
|||||||
@@ -19,6 +19,7 @@ import copy
|
|||||||
import logging
|
import logging
|
||||||
import shutil
|
import shutil
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from typing import Any, NotRequired, TypedDict
|
||||||
|
|
||||||
import datasets
|
import datasets
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
@@ -49,8 +50,32 @@ from .utils import (
|
|||||||
)
|
)
|
||||||
from .video_utils import concatenate_video_files, get_video_duration_in_s
|
from .video_utils import concatenate_video_files, get_video_duration_in_s
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]:
|
type FeatureDict = dict[str, dict[str, Any]]
|
||||||
|
type ChunkFile = tuple[int, int]
|
||||||
|
|
||||||
|
|
||||||
|
class IndexState(TypedDict):
|
||||||
|
chunk: int
|
||||||
|
file: int
|
||||||
|
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
|
||||||
|
|
||||||
|
|
||||||
|
class VideoIndex(TypedDict):
|
||||||
|
chunk: int
|
||||||
|
file: int
|
||||||
|
latest_duration: float
|
||||||
|
episode_duration: float
|
||||||
|
src_to_offset: NotRequired[dict[ChunkFile, float]]
|
||||||
|
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
|
||||||
|
dst_file_durations: NotRequired[dict[ChunkFile, float]]
|
||||||
|
|
||||||
|
|
||||||
|
type VideoIndexState = dict[str, VideoIndex]
|
||||||
|
|
||||||
|
|
||||||
|
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> FeatureDict:
|
||||||
"""Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.
|
"""Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
@@ -59,14 +84,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
|||||||
Returns:
|
Returns:
|
||||||
dict: A dictionary of merged video feature info.
|
dict: A dictionary of merged video feature info.
|
||||||
"""
|
"""
|
||||||
merged_info = copy.deepcopy(all_metadata[0].features)
|
merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
|
||||||
video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
|
video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
|
||||||
|
|
||||||
for vk in video_keys:
|
for vk in video_keys:
|
||||||
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
|
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
|
||||||
base_video_info = video_infos[0]
|
base_video_info = video_infos[0]
|
||||||
|
|
||||||
merged_encoder_info: dict = {}
|
merged_encoder_info: dict[str, Any] = {}
|
||||||
fallback_keys: list[str] = []
|
fallback_keys: list[str] = []
|
||||||
for info_key in VIDEO_ENCODER_INFO_KEYS:
|
for info_key in VIDEO_ENCODER_INFO_KEYS:
|
||||||
values = [info.get(info_key, None) for info in video_infos]
|
values = [info.get(info_key, None) for info in video_infos]
|
||||||
@@ -80,7 +105,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
|||||||
merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None
|
merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None
|
||||||
|
|
||||||
if fallback_keys:
|
if fallback_keys:
|
||||||
logging.warning(
|
logger.warning(
|
||||||
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
|
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
|
||||||
f"Setting these keys to null: {fallback_keys}.",
|
f"Setting these keys to null: {fallback_keys}.",
|
||||||
)
|
)
|
||||||
@@ -92,7 +117,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
|||||||
return merged_info
|
return merged_info
|
||||||
|
|
||||||
|
|
||||||
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
|
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
|
||||||
"""Validates that all dataset metadata have consistent properties.
|
"""Validates that all dataset metadata have consistent properties.
|
||||||
|
|
||||||
Ensures all datasets have the same fps, robot_type, and features to guarantee
|
Ensures all datasets have the same fps, robot_type, and features to guarantee
|
||||||
@@ -129,7 +154,9 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
|
|||||||
return fps, robot_type, features
|
return fps, robot_type, features
|
||||||
|
|
||||||
|
|
||||||
def update_data_df(df, src_meta, dst_meta):
|
def update_data_df(
|
||||||
|
df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
|
||||||
|
) -> pd.DataFrame:
|
||||||
"""Updates a data DataFrame with new indices and task mappings for aggregation.
|
"""Updates a data DataFrame with new indices and task mappings for aggregation.
|
||||||
|
|
||||||
Adjusts episode indices, frame indices, and task indices to account for
|
Adjusts episode indices, frame indices, and task indices to account for
|
||||||
@@ -154,12 +181,12 @@ def update_data_df(df, src_meta, dst_meta):
|
|||||||
|
|
||||||
|
|
||||||
def update_meta_data(
|
def update_meta_data(
|
||||||
df,
|
df: pd.DataFrame,
|
||||||
dst_meta,
|
dst_meta: LeRobotDatasetMetadata,
|
||||||
meta_idx,
|
meta_idx: IndexState,
|
||||||
data_idx,
|
data_idx: IndexState,
|
||||||
videos_idx,
|
videos_idx: VideoIndexState,
|
||||||
):
|
) -> pd.DataFrame:
|
||||||
"""Updates metadata DataFrame with new chunk, file, and timestamp indices.
|
"""Updates metadata DataFrame with new chunk, file, and timestamp indices.
|
||||||
|
|
||||||
Adjusts all indices and timestamps to account for previously aggregated
|
Adjusts all indices and timestamps to account for previously aggregated
|
||||||
@@ -289,7 +316,7 @@ def aggregate_datasets(
|
|||||||
chunk_size: int | None = None,
|
chunk_size: int | None = None,
|
||||||
concatenate_videos: bool = True,
|
concatenate_videos: bool = True,
|
||||||
concatenate_data: bool = True,
|
concatenate_data: bool = True,
|
||||||
):
|
) -> None:
|
||||||
"""Aggregates multiple LeRobot datasets into a single unified dataset.
|
"""Aggregates multiple LeRobot datasets into a single unified dataset.
|
||||||
|
|
||||||
This is the main function that orchestrates the aggregation process by:
|
This is the main function that orchestrates the aggregation process by:
|
||||||
@@ -309,7 +336,7 @@ def aggregate_datasets(
|
|||||||
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
|
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
|
||||||
concatenate_data: When False, keep one parquet per source file instead of packing into shards.
|
concatenate_data: When False, keep one parquet per source file instead of packing into shards.
|
||||||
"""
|
"""
|
||||||
logging.info("Start aggregate_datasets")
|
logger.info("Start aggregate_datasets")
|
||||||
|
|
||||||
if data_files_size_in_mb is None:
|
if data_files_size_in_mb is None:
|
||||||
data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
|
data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
|
||||||
@@ -341,15 +368,15 @@ def aggregate_datasets(
|
|||||||
video_files_size_in_mb=video_files_size_in_mb,
|
video_files_size_in_mb=video_files_size_in_mb,
|
||||||
)
|
)
|
||||||
|
|
||||||
logging.info("Find all tasks")
|
logger.info("Find all tasks")
|
||||||
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
|
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
|
||||||
dst_meta.tasks = pd.DataFrame(
|
dst_meta.tasks = pd.DataFrame(
|
||||||
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
|
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
|
||||||
)
|
)
|
||||||
|
|
||||||
meta_idx = {"chunk": 0, "file": 0}
|
meta_idx: IndexState = {"chunk": 0, "file": 0}
|
||||||
data_idx = {"chunk": 0, "file": 0}
|
data_idx: IndexState = {"chunk": 0, "file": 0}
|
||||||
videos_idx = {
|
videos_idx: VideoIndexState = {
|
||||||
key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
|
key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -373,12 +400,17 @@ def aggregate_datasets(
|
|||||||
dst_meta.info.total_frames += src_meta.total_frames
|
dst_meta.info.total_frames += src_meta.total_frames
|
||||||
|
|
||||||
finalize_aggregation(dst_meta, all_metadata)
|
finalize_aggregation(dst_meta, all_metadata)
|
||||||
logging.info("Aggregation complete.")
|
logger.info("Aggregation complete.")
|
||||||
|
|
||||||
|
|
||||||
def aggregate_videos(
|
def aggregate_videos(
|
||||||
src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
|
src_meta: LeRobotDatasetMetadata,
|
||||||
):
|
dst_meta: LeRobotDatasetMetadata,
|
||||||
|
videos_idx: VideoIndexState,
|
||||||
|
video_files_size_in_mb: float,
|
||||||
|
chunk_size: int,
|
||||||
|
concatenate_videos: bool = True,
|
||||||
|
) -> VideoIndexState:
|
||||||
"""Aggregates video chunks from a source dataset into the destination dataset.
|
"""Aggregates video chunks from a source dataset into the destination dataset.
|
||||||
|
|
||||||
Handles video file concatenation and rotation based on file size limits.
|
Handles video file concatenation and rotation based on file size limits.
|
||||||
@@ -406,15 +438,16 @@ def aggregate_videos(
|
|||||||
videos_idx[key]["dst_file_durations"] = {}
|
videos_idx[key]["dst_file_durations"] = {}
|
||||||
|
|
||||||
for key, video_idx in videos_idx.items():
|
for key, video_idx in videos_idx.items():
|
||||||
unique_chunk_file_pairs = {
|
unique_chunk_file_pairs: list[ChunkFile] = sorted(
|
||||||
(chunk, file)
|
{
|
||||||
for chunk, file in zip(
|
(chunk, file)
|
||||||
src_meta.episodes[f"videos/{key}/chunk_index"],
|
for chunk, file in zip(
|
||||||
src_meta.episodes[f"videos/{key}/file_index"],
|
src_meta.episodes[f"videos/{key}/chunk_index"],
|
||||||
strict=False,
|
src_meta.episodes[f"videos/{key}/file_index"],
|
||||||
)
|
strict=False,
|
||||||
}
|
)
|
||||||
unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
|
}
|
||||||
|
)
|
||||||
|
|
||||||
chunk_idx = video_idx["chunk"]
|
chunk_idx = video_idx["chunk"]
|
||||||
file_idx = video_idx["file"]
|
file_idx = video_idx["file"]
|
||||||
@@ -489,7 +522,14 @@ def aggregate_videos(
|
|||||||
return videos_idx
|
return videos_idx
|
||||||
|
|
||||||
|
|
||||||
def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True):
|
def aggregate_data(
|
||||||
|
src_meta: LeRobotDatasetMetadata,
|
||||||
|
dst_meta: LeRobotDatasetMetadata,
|
||||||
|
data_idx: IndexState,
|
||||||
|
data_files_size_in_mb: float,
|
||||||
|
chunk_size: int,
|
||||||
|
concatenate_data: bool = True,
|
||||||
|
) -> IndexState:
|
||||||
"""Aggregates data chunks from a source dataset into the destination dataset.
|
"""Aggregates data chunks from a source dataset into the destination dataset.
|
||||||
|
|
||||||
Reads source data files, updates indices to match the aggregated dataset,
|
Reads source data files, updates indices to match the aggregated dataset,
|
||||||
@@ -510,14 +550,16 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
|
|||||||
Returns:
|
Returns:
|
||||||
dict: Updated data_idx with current chunk and file indices.
|
dict: Updated data_idx with current chunk and file indices.
|
||||||
"""
|
"""
|
||||||
unique_chunk_file_ids = {
|
unique_chunk_file_ids: list[ChunkFile] = sorted(
|
||||||
(c, f)
|
{
|
||||||
for c, f in zip(
|
(c, f)
|
||||||
src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False
|
for c, f in zip(
|
||||||
)
|
src_meta.episodes["data/chunk_index"],
|
||||||
}
|
src_meta.episodes["data/file_index"],
|
||||||
|
strict=False,
|
||||||
unique_chunk_file_ids = sorted(unique_chunk_file_ids)
|
)
|
||||||
|
}
|
||||||
|
)
|
||||||
contains_images = len(dst_meta.image_keys) > 0
|
contains_images = len(dst_meta.image_keys) > 0
|
||||||
|
|
||||||
# retrieve features schema for proper image typing in parquet
|
# retrieve features schema for proper image typing in parquet
|
||||||
@@ -525,7 +567,7 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
|
|||||||
|
|
||||||
# Track source to destination file mapping for metadata update
|
# Track source to destination file mapping for metadata update
|
||||||
# This is critical for handling datasets that are already results of a merge
|
# This is critical for handling datasets that are already results of a merge
|
||||||
src_to_dst: dict[tuple[int, int], tuple[int, int]] = {}
|
src_to_dst: dict[ChunkFile, ChunkFile] = {}
|
||||||
|
|
||||||
for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
|
for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
|
||||||
src_path = src_meta.root / DEFAULT_DATA_PATH.format(
|
src_path = src_meta.root / DEFAULT_DATA_PATH.format(
|
||||||
@@ -564,7 +606,13 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
|
|||||||
return data_idx
|
return data_idx
|
||||||
|
|
||||||
|
|
||||||
def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
|
def aggregate_metadata(
|
||||||
|
src_meta: LeRobotDatasetMetadata,
|
||||||
|
dst_meta: LeRobotDatasetMetadata,
|
||||||
|
meta_idx: IndexState,
|
||||||
|
data_idx: IndexState,
|
||||||
|
videos_idx: VideoIndexState,
|
||||||
|
) -> IndexState:
|
||||||
"""Aggregates metadata from a source dataset into the destination dataset.
|
"""Aggregates metadata from a source dataset into the destination dataset.
|
||||||
|
|
||||||
Reads source metadata files, updates all indices and timestamps,
|
Reads source metadata files, updates all indices and timestamps,
|
||||||
@@ -580,16 +628,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
|
|||||||
Returns:
|
Returns:
|
||||||
dict: Updated meta_idx with current chunk and file indices.
|
dict: Updated meta_idx with current chunk and file indices.
|
||||||
"""
|
"""
|
||||||
chunk_file_ids = {
|
chunk_file_ids: list[ChunkFile] = sorted(
|
||||||
(c, f)
|
{
|
||||||
for c, f in zip(
|
(c, f)
|
||||||
src_meta.episodes["meta/episodes/chunk_index"],
|
for c, f in zip(
|
||||||
src_meta.episodes["meta/episodes/file_index"],
|
src_meta.episodes["meta/episodes/chunk_index"],
|
||||||
strict=False,
|
src_meta.episodes["meta/episodes/file_index"],
|
||||||
)
|
strict=False,
|
||||||
}
|
)
|
||||||
|
}
|
||||||
chunk_file_ids = sorted(chunk_file_ids)
|
)
|
||||||
for chunk_idx, file_idx in chunk_file_ids:
|
for chunk_idx, file_idx in chunk_file_ids:
|
||||||
src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
|
src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
|
||||||
df = pd.read_parquet(src_path)
|
df = pd.read_parquet(src_path)
|
||||||
@@ -622,16 +670,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
|
|||||||
def append_or_create_parquet_file(
|
def append_or_create_parquet_file(
|
||||||
df: pd.DataFrame,
|
df: pd.DataFrame,
|
||||||
src_path: Path,
|
src_path: Path,
|
||||||
idx: dict[str, int],
|
idx: IndexState,
|
||||||
max_mb: float,
|
max_mb: float,
|
||||||
chunk_size: int,
|
chunk_size: int,
|
||||||
default_path: str,
|
default_path: str,
|
||||||
contains_images: bool = False,
|
contains_images: bool = False,
|
||||||
aggr_root: Path = None,
|
aggr_root: Path | None = None,
|
||||||
hf_features: datasets.Features | None = None,
|
hf_features: datasets.Features | None = None,
|
||||||
concatenate: bool = True,
|
concatenate: bool = True,
|
||||||
one_row_group_per_episode: bool = False,
|
one_row_group_per_episode: bool = False,
|
||||||
) -> tuple[dict[str, int], tuple[int, int]]:
|
) -> tuple[IndexState, ChunkFile]:
|
||||||
"""Appends data to an existing parquet file or creates a new one based on size constraints.
|
"""Appends data to an existing parquet file or creates a new one based on size constraints.
|
||||||
|
|
||||||
Manages file rotation when size limits are exceeded to prevent individual files
|
Manages file rotation when size limits are exceeded to prevent individual files
|
||||||
@@ -654,7 +702,13 @@ def append_or_create_parquet_file(
|
|||||||
Returns:
|
Returns:
|
||||||
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
|
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
|
||||||
and (dst_chunk, dst_file) is the actual destination file the data was written to.
|
and (dst_chunk, dst_file) is the actual destination file the data was written to.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If aggr_root is not provided.
|
||||||
"""
|
"""
|
||||||
|
if aggr_root is None:
|
||||||
|
raise ValueError("aggr_root must be provided.")
|
||||||
|
|
||||||
dst_chunk, dst_file = idx["chunk"], idx["file"]
|
dst_chunk, dst_file = idx["chunk"], idx["file"]
|
||||||
dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
|
dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
|
||||||
|
|
||||||
@@ -698,7 +752,9 @@ def append_or_create_parquet_file(
|
|||||||
return idx, (dst_chunk, dst_file)
|
return idx, (dst_chunk, dst_file)
|
||||||
|
|
||||||
|
|
||||||
def finalize_aggregation(aggr_meta, all_metadata):
|
def finalize_aggregation(
|
||||||
|
aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
|
||||||
|
) -> None:
|
||||||
"""Finalizes the dataset aggregation by writing summary files and statistics.
|
"""Finalizes the dataset aggregation by writing summary files and statistics.
|
||||||
|
|
||||||
Writes the tasks file, info file with total counts and splits, and
|
Writes the tasks file, info file with total counts and splits, and
|
||||||
@@ -708,16 +764,16 @@ def finalize_aggregation(aggr_meta, all_metadata):
|
|||||||
aggr_meta: Aggregated dataset metadata.
|
aggr_meta: Aggregated dataset metadata.
|
||||||
all_metadata: List of all source dataset metadata objects.
|
all_metadata: List of all source dataset metadata objects.
|
||||||
"""
|
"""
|
||||||
logging.info("write tasks")
|
logger.info("write tasks")
|
||||||
write_tasks(aggr_meta.tasks, aggr_meta.root)
|
write_tasks(aggr_meta.tasks, aggr_meta.root)
|
||||||
|
|
||||||
logging.info("write info")
|
logger.info("write info")
|
||||||
aggr_meta.info.total_tasks = len(aggr_meta.tasks)
|
aggr_meta.info.total_tasks = len(aggr_meta.tasks)
|
||||||
aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata)
|
aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata)
|
||||||
aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata)
|
aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata)
|
||||||
aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
|
aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
|
||||||
write_info(aggr_meta.info, aggr_meta.root)
|
write_info(aggr_meta.info, aggr_meta.root)
|
||||||
|
|
||||||
logging.info("write stats")
|
logger.info("write stats")
|
||||||
aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
|
aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
|
||||||
write_stats(aggr_meta.stats, aggr_meta.root)
|
write_stats(aggr_meta.stats, aggr_meta.root)
|
||||||
|
|||||||
@@ -188,8 +188,8 @@ class LeRobotDatasetMetadata:
|
|||||||
def _load_metadata(self):
|
def _load_metadata(self):
|
||||||
self.info = load_info(self.root)
|
self.info = load_info(self.root)
|
||||||
check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
|
check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
|
||||||
self.tasks = load_tasks(self.root)
|
self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
|
||||||
self.episodes = load_episodes(self.root)
|
self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
|
||||||
self.stats = load_stats(self.root)
|
self.stats = load_stats(self.root)
|
||||||
|
|
||||||
def ensure_readable(self) -> None:
|
def ensure_readable(self) -> None:
|
||||||
|
|||||||
@@ -384,7 +384,12 @@ class LiberoEnv(gym.Env):
|
|||||||
|
|
||||||
def close(self):
|
def close(self):
|
||||||
if self._env is not None:
|
if self._env is not None:
|
||||||
self._env.close()
|
try:
|
||||||
|
self._env.close()
|
||||||
|
finally:
|
||||||
|
# LIBERO deletes its inner env on close, so this wrapper must
|
||||||
|
# be recreated before the next reset.
|
||||||
|
self._env = None
|
||||||
|
|
||||||
|
|
||||||
def _make_env_fns(
|
def _make_env_fns(
|
||||||
|
|||||||
@@ -384,7 +384,9 @@ class RoboTwinEnv(gym.Env):
|
|||||||
|
|
||||||
self._env: Any | None = None # deferred — created on first reset() inside worker
|
self._env: Any | None = None # deferred — created on first reset() inside worker
|
||||||
self._step_count: int = 0
|
self._step_count: int = 0
|
||||||
self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
|
self._black_frame: np.ndarray = np.zeros(
|
||||||
|
(self.observation_height, self.observation_width, 3), dtype=np.uint8
|
||||||
|
)
|
||||||
|
|
||||||
image_spaces = {
|
image_spaces = {
|
||||||
cam: spaces.Box(
|
cam: spaces.Box(
|
||||||
|
|||||||
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
|
|||||||
|
|
||||||
if action.shape[0] != 7:
|
if action.shape[0] != 7:
|
||||||
# Unknown layout — fall back to zero-pad so the sim doesn't crash.
|
# Unknown layout — fall back to zero-pad so the sim doesn't crash.
|
||||||
padded = np.zeros(ctrl_dim, dtype=np.float64)
|
padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
|
||||||
padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
|
padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
|
||||||
return padded
|
return padded
|
||||||
|
|
||||||
|
|||||||
@@ -122,6 +122,9 @@ MODEL_ENCODING_TABLE = {
|
|||||||
"xm430-w350": X_SERIES_ENCODINGS_TABLE,
|
"xm430-w350": X_SERIES_ENCODINGS_TABLE,
|
||||||
"xm540-w270": X_SERIES_ENCODINGS_TABLE,
|
"xm540-w270": X_SERIES_ENCODINGS_TABLE,
|
||||||
"xc430-w150": X_SERIES_ENCODINGS_TABLE,
|
"xc430-w150": X_SERIES_ENCODINGS_TABLE,
|
||||||
|
"xh540-w150": X_SERIES_ENCODINGS_TABLE,
|
||||||
|
"xc330-t288": X_SERIES_ENCODINGS_TABLE,
|
||||||
|
"xc330-t181": X_SERIES_ENCODINGS_TABLE,
|
||||||
}
|
}
|
||||||
|
|
||||||
# {model: model_resolution}
|
# {model: model_resolution}
|
||||||
@@ -134,6 +137,9 @@ MODEL_RESOLUTION = {
|
|||||||
"xm430-w350": 4096,
|
"xm430-w350": 4096,
|
||||||
"xm540-w270": 4096,
|
"xm540-w270": 4096,
|
||||||
"xc430-w150": 4096,
|
"xc430-w150": 4096,
|
||||||
|
"xh540-w150": 4096,
|
||||||
|
"xc330-t288": 4096,
|
||||||
|
"xc330-t181": 4096,
|
||||||
}
|
}
|
||||||
|
|
||||||
# {model: model_number}
|
# {model: model_number}
|
||||||
@@ -145,6 +151,9 @@ MODEL_NUMBER_TABLE = {
|
|||||||
"xm430-w350": 1020,
|
"xm430-w350": 1020,
|
||||||
"xm540-w270": 1120,
|
"xm540-w270": 1120,
|
||||||
"xc430-w150": 1070,
|
"xc430-w150": 1070,
|
||||||
|
"xh540-w150": 1110,
|
||||||
|
"xc330-t288": 1220,
|
||||||
|
"xc330-t181": 1210,
|
||||||
}
|
}
|
||||||
|
|
||||||
# {model: available_operating_modes}
|
# {model: available_operating_modes}
|
||||||
@@ -156,6 +165,9 @@ MODEL_OPERATING_MODES = {
|
|||||||
"xm430-w350": [0, 1, 3, 4, 5, 16],
|
"xm430-w350": [0, 1, 3, 4, 5, 16],
|
||||||
"xm540-w270": [0, 1, 3, 4, 5, 16],
|
"xm540-w270": [0, 1, 3, 4, 5, 16],
|
||||||
"xc430-w150": [1, 3, 4, 16],
|
"xc430-w150": [1, 3, 4, 16],
|
||||||
|
"xh540-w150": [0, 1, 3, 4, 5, 16],
|
||||||
|
"xc330-t288": [0, 1, 3, 4, 5, 16],
|
||||||
|
"xc330-t181": [0, 1, 3, 4, 5, 16],
|
||||||
}
|
}
|
||||||
|
|
||||||
MODEL_CONTROL_TABLE = {
|
MODEL_CONTROL_TABLE = {
|
||||||
@@ -166,6 +178,9 @@ MODEL_CONTROL_TABLE = {
|
|||||||
"xm430-w350": X_SERIES_CONTROL_TABLE,
|
"xm430-w350": X_SERIES_CONTROL_TABLE,
|
||||||
"xm540-w270": X_SERIES_CONTROL_TABLE,
|
"xm540-w270": X_SERIES_CONTROL_TABLE,
|
||||||
"xc430-w150": X_SERIES_CONTROL_TABLE,
|
"xc430-w150": X_SERIES_CONTROL_TABLE,
|
||||||
|
"xh540-w150": X_SERIES_CONTROL_TABLE,
|
||||||
|
"xc330-t288": X_SERIES_CONTROL_TABLE,
|
||||||
|
"xc330-t181": X_SERIES_CONTROL_TABLE,
|
||||||
}
|
}
|
||||||
|
|
||||||
MODEL_BAUDRATE_TABLE = {
|
MODEL_BAUDRATE_TABLE = {
|
||||||
@@ -176,6 +191,9 @@ MODEL_BAUDRATE_TABLE = {
|
|||||||
"xm430-w350": X_SERIES_BAUDRATE_TABLE,
|
"xm430-w350": X_SERIES_BAUDRATE_TABLE,
|
||||||
"xm540-w270": X_SERIES_BAUDRATE_TABLE,
|
"xm540-w270": X_SERIES_BAUDRATE_TABLE,
|
||||||
"xc430-w150": X_SERIES_BAUDRATE_TABLE,
|
"xc430-w150": X_SERIES_BAUDRATE_TABLE,
|
||||||
|
"xh540-w150": X_SERIES_BAUDRATE_TABLE,
|
||||||
|
"xc330-t288": X_SERIES_BAUDRATE_TABLE,
|
||||||
|
"xc330-t181": X_SERIES_BAUDRATE_TABLE,
|
||||||
}
|
}
|
||||||
|
|
||||||
AVAILABLE_BAUDRATES = [
|
AVAILABLE_BAUDRATES = [
|
||||||
|
|||||||
@@ -302,6 +302,33 @@ def _pad_evo1_stats(
|
|||||||
return padded_stats
|
return padded_stats
|
||||||
|
|
||||||
|
|
||||||
|
def _refresh_evo1_normalization_steps(
|
||||||
|
config: Evo1Config,
|
||||||
|
preprocessor: PolicyProcessorPipeline,
|
||||||
|
postprocessor: PolicyProcessorPipeline,
|
||||||
|
) -> None:
|
||||||
|
"""Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
|
||||||
|
|
||||||
|
Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
|
||||||
|
into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
|
||||||
|
and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
|
||||||
|
normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
|
||||||
|
LIBERO stats). Padding is a no-op when stats are already at the target width.
|
||||||
|
"""
|
||||||
|
normalization_features = _evo1_normalization_features(config)
|
||||||
|
action_features = _evo1_action_features(config)
|
||||||
|
for step in preprocessor.steps:
|
||||||
|
if isinstance(step, NormalizerProcessorStep):
|
||||||
|
step.features = normalization_features
|
||||||
|
step.stats = _pad_evo1_stats(config, step.stats)
|
||||||
|
step.to(device=step.device, dtype=step.dtype)
|
||||||
|
for step in postprocessor.steps:
|
||||||
|
if isinstance(step, UnnormalizerProcessorStep):
|
||||||
|
step.features = action_features
|
||||||
|
step.stats = _pad_evo1_stats(config, step.stats)
|
||||||
|
step.to(device=step.device, dtype=step.dtype)
|
||||||
|
|
||||||
|
|
||||||
def reconcile_evo1_processors(
|
def reconcile_evo1_processors(
|
||||||
config: Evo1Config,
|
config: Evo1Config,
|
||||||
preprocessor: PolicyProcessorPipeline,
|
preprocessor: PolicyProcessorPipeline,
|
||||||
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
|
|||||||
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
|
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
|
||||||
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config.
|
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config.
|
||||||
|
|
||||||
Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
|
Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
|
||||||
(converters are plain functions and are never serialized), and eval-time CLI overrides of the
|
(converters are plain functions and are never serialized), eval-time CLI overrides of the
|
||||||
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
|
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
|
||||||
restores the converter and rebuilds the action step from the current config so those overrides
|
(un)normalizer stats/features when the generic override path injects raw, unpadded dataset
|
||||||
take effect.
|
stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and
|
||||||
|
rebuilds the action step from the current config so those overrides take effect.
|
||||||
"""
|
"""
|
||||||
# Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
|
# Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
|
||||||
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
|
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
|
||||||
preprocessor.to_transition = evo1_batch_to_transition
|
preprocessor.to_transition = evo1_batch_to_transition
|
||||||
|
|
||||||
|
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
|
||||||
|
|
||||||
action_step = Evo1ActionProcessorStep(
|
action_step = Evo1ActionProcessorStep(
|
||||||
action_dim=_evo1_action_dim(config),
|
action_dim=_evo1_action_dim(config),
|
||||||
binarize_gripper=config.binarize_gripper,
|
binarize_gripper=config.binarize_gripper,
|
||||||
|
|||||||
@@ -44,12 +44,19 @@ from lerobot.utils.constants import (
|
|||||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||||
)
|
)
|
||||||
from lerobot.utils.feature_utils import dataset_to_policy_features
|
from lerobot.utils.feature_utils import dataset_to_policy_features
|
||||||
|
from lerobot.utils.import_utils import _peft_available, require_package
|
||||||
|
|
||||||
from .evo1.configuration_evo1 import Evo1Config
|
from .evo1.configuration_evo1 import Evo1Config
|
||||||
from .groot.configuration_groot import GrootConfig
|
from .groot.configuration_groot import GrootConfig
|
||||||
from .pretrained import PreTrainedPolicy
|
from .pretrained import PreTrainedPolicy
|
||||||
from .utils import validate_visual_features_consistency
|
from .utils import validate_visual_features_consistency
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import PeftConfig, PeftModel
|
||||||
|
else:
|
||||||
|
PeftConfig = None
|
||||||
|
PeftModel = None
|
||||||
|
|
||||||
|
|
||||||
def _reconnect_relative_absolute_steps(
|
def _reconnect_relative_absolute_steps(
|
||||||
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
|
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
|
||||||
@@ -334,12 +341,15 @@ def make_policy(
|
|||||||
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
|
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
|
||||||
# of the adapter and the adapter's config contains the path to the base policy. So we need the
|
# of the adapter and the adapter's config contains the path to the base policy. So we need the
|
||||||
# adapter config first, then load the correct policy and then apply PEFT.
|
# adapter config first, then load the correct policy and then apply PEFT.
|
||||||
from peft import PeftConfig, PeftModel
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
logging.info("Loading policy's PEFT adapter.")
|
logging.info("Loading policy's PEFT adapter.")
|
||||||
|
|
||||||
peft_pretrained_path = str(cfg.pretrained_path)
|
peft_pretrained_path = str(cfg.pretrained_path)
|
||||||
peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
|
peft_config = PeftConfig.from_pretrained(
|
||||||
|
peft_pretrained_path,
|
||||||
|
revision=cfg.pretrained_revision,
|
||||||
|
)
|
||||||
|
|
||||||
kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
|
kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
|
||||||
if not kwargs["pretrained_name_or_path"]:
|
if not kwargs["pretrained_name_or_path"]:
|
||||||
@@ -350,9 +360,14 @@ def make_policy(
|
|||||||
"the adapter was trained."
|
"the adapter was trained."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
kwargs["revision"] = peft_config.revision
|
||||||
policy = policy_cls.from_pretrained(**kwargs)
|
policy = policy_cls.from_pretrained(**kwargs)
|
||||||
policy = PeftModel.from_pretrained(
|
policy = PeftModel.from_pretrained(
|
||||||
policy, peft_pretrained_path, config=peft_config, is_trainable=True
|
policy,
|
||||||
|
peft_pretrained_path,
|
||||||
|
config=peft_config,
|
||||||
|
revision=cfg.pretrained_revision,
|
||||||
|
is_trainable=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -37,13 +37,19 @@ def is_image_feature(key: str) -> bool:
|
|||||||
@dataclass
|
@dataclass
|
||||||
class ConcurrencyConfig:
|
class ConcurrencyConfig:
|
||||||
"""Configuration for the concurrency of the actor and learner.
|
"""Configuration for the concurrency of the actor and learner.
|
||||||
|
|
||||||
Possible values are:
|
Possible values are:
|
||||||
- "threads": Use threads for the actor and learner.
|
- "threads": Use threads for the actor and learner.
|
||||||
- "processes": Use processes for the actor and learner.
|
- "processes": Use processes for the actor and learner.
|
||||||
|
|
||||||
|
``multiprocessing_context`` selects the process-wide start method when
|
||||||
|
processes are used. Set it to ``None`` to preserve Python's default or a
|
||||||
|
method already selected by the embedding application.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
actor: str = "threads"
|
actor: str = "threads"
|
||||||
learner: str = "threads"
|
learner: str = "threads"
|
||||||
|
multiprocessing_context: str | None = "spawn"
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|||||||
@@ -43,11 +43,22 @@ from torch.distributions import Beta
|
|||||||
|
|
||||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||||
from lerobot.utils.constants import ACTION
|
from lerobot.utils.constants import ACTION
|
||||||
from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
|
from lerobot.utils.import_utils import (
|
||||||
|
_peft_available,
|
||||||
|
_scipy_available,
|
||||||
|
_transformers_available,
|
||||||
|
require_package,
|
||||||
|
)
|
||||||
|
|
||||||
from ..rtc.modeling_rtc import RTCProcessor
|
from ..rtc.modeling_rtc import RTCProcessor
|
||||||
from .configuration_molmoact2 import MolmoAct2Config
|
from .configuration_molmoact2 import MolmoAct2Config
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import LoraConfig, get_peft_model
|
||||||
|
else:
|
||||||
|
LoraConfig = None
|
||||||
|
get_peft_model = None
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -1731,13 +1742,11 @@ class MolmoAct2Policy(PreTrainedPolicy):
|
|||||||
|
|
||||||
def _build_inner_lora_config(self):
|
def _build_inner_lora_config(self):
|
||||||
require_package("peft", extra="molmoact2")
|
require_package("peft", extra="molmoact2")
|
||||||
from peft import LoraConfig
|
|
||||||
|
|
||||||
return LoraConfig(**self._get_inner_peft_targets())
|
return LoraConfig(**self._get_inner_peft_targets())
|
||||||
|
|
||||||
def _apply_lora_adapters(self) -> None:
|
def _apply_lora_adapters(self) -> None:
|
||||||
require_package("peft", extra="molmoact2")
|
require_package("peft", extra="molmoact2")
|
||||||
from peft import get_peft_model
|
|
||||||
|
|
||||||
peft_config = self._build_inner_lora_config()
|
peft_config = self._build_inner_lora_config()
|
||||||
self._validate_peft_config(peft_config)
|
self._validate_peft_config(peft_config)
|
||||||
|
|||||||
@@ -34,14 +34,22 @@ from lerobot.configs import PreTrainedConfig
|
|||||||
from lerobot.configs.train import TrainPipelineConfig
|
from lerobot.configs.train import TrainPipelineConfig
|
||||||
from lerobot.utils.device_utils import resolve_safetensors_device
|
from lerobot.utils.device_utils import resolve_safetensors_device
|
||||||
from lerobot.utils.hub import HubMixin
|
from lerobot.utils.hub import HubMixin
|
||||||
|
from lerobot.utils.import_utils import _peft_available, require_package
|
||||||
|
|
||||||
from .utils import log_model_loading_keys
|
from .utils import log_model_loading_keys
|
||||||
|
|
||||||
T = TypeVar("T", bound="PreTrainedPolicy")
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model
|
||||||
|
else:
|
||||||
|
PEFT_TYPE_TO_CONFIG_MAPPING = None
|
||||||
|
PeftType = None
|
||||||
|
get_peft_model = None
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||||
|
|
||||||
|
T = TypeVar("T", bound="PreTrainedPolicy")
|
||||||
|
|
||||||
|
|
||||||
def _build_card_context(
|
def _build_card_context(
|
||||||
cfg: TrainPipelineConfig | None,
|
cfg: TrainPipelineConfig | None,
|
||||||
@@ -384,7 +392,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
|||||||
peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
|
peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
|
||||||
These are merged with policy defaults to build the final config.
|
These are merged with policy defaults to build the final config.
|
||||||
"""
|
"""
|
||||||
from peft import get_peft_model
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
# If user provided a complete config, use it directly (with overrides)
|
# If user provided a complete config, use it directly (with overrides)
|
||||||
if peft_config is not None:
|
if peft_config is not None:
|
||||||
@@ -455,7 +463,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
|||||||
Returns:
|
Returns:
|
||||||
Preprocessed dict with renamed keys and init_type mapped to method-specific key.
|
Preprocessed dict with renamed keys and init_type mapped to method-specific key.
|
||||||
"""
|
"""
|
||||||
from peft import PeftType
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
cli_overrides = cli_overrides.copy()
|
cli_overrides = cli_overrides.copy()
|
||||||
|
|
||||||
@@ -480,7 +488,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
|||||||
|
|
||||||
def _build_peft_config(self, cli_overrides: dict):
|
def _build_peft_config(self, cli_overrides: dict):
|
||||||
"""Build a PEFT config from policy defaults and CLI overrides."""
|
"""Build a PEFT config from policy defaults and CLI overrides."""
|
||||||
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
# Determine PEFT method type (default to LORA)
|
# Determine PEFT method type (default to LORA)
|
||||||
method_type_str = cli_overrides.get("method_type") or "lora"
|
method_type_str = cli_overrides.get("method_type") or "lora"
|
||||||
@@ -507,7 +515,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
|||||||
|
|
||||||
def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
|
def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
|
||||||
"""Apply CLI overrides to an existing PEFT config."""
|
"""Apply CLI overrides to an existing PEFT config."""
|
||||||
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
# Get method type from existing config or CLI override
|
# Get method type from existing config or CLI override
|
||||||
method_type_str = cli_overrides.get("method_type")
|
method_type_str = cli_overrides.get("method_type")
|
||||||
|
|||||||
@@ -132,10 +132,20 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
|
|||||||
def transform_features(
|
def transform_features(
|
||||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
for axis in ["x", "y", "z", "gripper"]:
|
for axis in ["x", "y", "z"]:
|
||||||
features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
|
features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
|
||||||
|
features[PipelineFeatureType.ACTION].pop("gripper", None)
|
||||||
|
|
||||||
for feat in ["enabled", "target_x", "target_y", "target_z", "target_wx", "target_wy", "target_wz"]:
|
for feat in [
|
||||||
|
"enabled",
|
||||||
|
"target_x",
|
||||||
|
"target_y",
|
||||||
|
"target_z",
|
||||||
|
"target_wx",
|
||||||
|
"target_wy",
|
||||||
|
"target_wz",
|
||||||
|
"gripper_vel",
|
||||||
|
]:
|
||||||
features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
|
features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
|
||||||
type=FeatureType.ACTION, shape=(1,)
|
type=FeatureType.ACTION, shape=(1,)
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -91,7 +91,7 @@ from lerobot.robots import so_follower # noqa: F401
|
|||||||
from lerobot.teleoperators import gamepad, so_leader # noqa: F401
|
from lerobot.teleoperators import gamepad, so_leader # noqa: F401
|
||||||
from lerobot.teleoperators.utils import TeleopEvents
|
from lerobot.teleoperators.utils import TeleopEvents
|
||||||
from lerobot.utils.device_utils import get_safe_torch_device
|
from lerobot.utils.device_utils import get_safe_torch_device
|
||||||
from lerobot.utils.process import ProcessSignalHandler
|
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
|
||||||
from lerobot.utils.random_utils import set_seed
|
from lerobot.utils.random_utils import set_seed
|
||||||
from lerobot.utils.robot_utils import precise_sleep
|
from lerobot.utils.robot_utils import precise_sleep
|
||||||
from lerobot.utils.transition import (
|
from lerobot.utils.transition import (
|
||||||
@@ -124,9 +124,7 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
|
|||||||
cfg.validate()
|
cfg.validate()
|
||||||
display_pid = False
|
display_pid = False
|
||||||
if not use_threads(cfg):
|
if not use_threads(cfg):
|
||||||
import torch.multiprocessing as mp
|
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
|
||||||
|
|
||||||
mp.set_start_method("spawn")
|
|
||||||
display_pid = True
|
display_pid = True
|
||||||
|
|
||||||
# Create logs directory to ensure it exists
|
# Create logs directory to ensure it exists
|
||||||
|
|||||||
@@ -18,7 +18,7 @@ import functools
|
|||||||
import threading
|
import threading
|
||||||
from collections.abc import Callable, Sequence
|
from collections.abc import Callable, Sequence
|
||||||
from contextlib import suppress
|
from contextlib import suppress
|
||||||
from typing import TypedDict
|
from typing import NotRequired, TypedDict
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
import torch.nn.functional as F # noqa: N812
|
import torch.nn.functional as F # noqa: N812
|
||||||
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
|
|||||||
next_state: dict[str, torch.Tensor]
|
next_state: dict[str, torch.Tensor]
|
||||||
done: torch.Tensor
|
done: torch.Tensor
|
||||||
truncated: torch.Tensor
|
truncated: torch.Tensor
|
||||||
complementary_info: dict[str, torch.Tensor | float | int] | None = None
|
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
|
||||||
|
|
||||||
|
|
||||||
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
|
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
|
||||||
|
|||||||
@@ -102,7 +102,7 @@ from lerobot.utils.constants import (
|
|||||||
)
|
)
|
||||||
from lerobot.utils.device_utils import get_safe_torch_device
|
from lerobot.utils.device_utils import get_safe_torch_device
|
||||||
from lerobot.utils.io_utils import load_json, write_json
|
from lerobot.utils.io_utils import load_json, write_json
|
||||||
from lerobot.utils.process import ProcessSignalHandler
|
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
|
||||||
from lerobot.utils.random_utils import set_seed
|
from lerobot.utils.random_utils import set_seed
|
||||||
from lerobot.utils.utils import (
|
from lerobot.utils.utils import (
|
||||||
format_big_number,
|
format_big_number,
|
||||||
@@ -123,9 +123,7 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
|
|||||||
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
|
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
|
||||||
require_package("grpcio", extra="hilserl", import_name="grpc")
|
require_package("grpcio", extra="hilserl", import_name="grpc")
|
||||||
if not use_threads(cfg):
|
if not use_threads(cfg):
|
||||||
import torch.multiprocessing as mp
|
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
|
||||||
|
|
||||||
mp.set_start_method("spawn")
|
|
||||||
|
|
||||||
# Use the job_name from the config
|
# Use the job_name from the config
|
||||||
train(
|
train(
|
||||||
|
|||||||
@@ -46,6 +46,12 @@ class SOFollowerConfig:
|
|||||||
position_i_coefficient: int = 0
|
position_i_coefficient: int = 0
|
||||||
position_d_coefficient: int = 32
|
position_d_coefficient: int = 32
|
||||||
|
|
||||||
|
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
|
||||||
|
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
|
||||||
|
# at once, which otherwise aborts the control loop. Retries are immediate (no sleep) and only happen on
|
||||||
|
# failure, so the steady-state read cost is unchanged.
|
||||||
|
num_read_retries: int = 2
|
||||||
|
|
||||||
|
|
||||||
@RobotConfig.register_subclass("so101_follower")
|
@RobotConfig.register_subclass("so101_follower")
|
||||||
@RobotConfig.register_subclass("so100_follower")
|
@RobotConfig.register_subclass("so100_follower")
|
||||||
|
|||||||
@@ -510,10 +510,10 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
|
|||||||
# We only use the ee pose in the dataset, so we don't need the joint positions
|
# We only use the ee pose in the dataset, so we don't need the joint positions
|
||||||
for n in self.motor_names:
|
for n in self.motor_names:
|
||||||
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
|
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
|
||||||
# We specify the dataset features of this step that we want to be stored in the dataset
|
# Store end-effector features as actions in the dataset schema
|
||||||
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
|
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
|
||||||
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
|
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
|
||||||
type=FeatureType.STATE, shape=(1,)
|
type=FeatureType.ACTION, shape=(1,)
|
||||||
)
|
)
|
||||||
return features
|
return features
|
||||||
|
|
||||||
|
|||||||
@@ -180,7 +180,7 @@ class SOFollower(Robot):
|
|||||||
def get_observation(self) -> RobotObservation:
|
def get_observation(self) -> RobotObservation:
|
||||||
# Read arm position
|
# Read arm position
|
||||||
start = time.perf_counter()
|
start = time.perf_counter()
|
||||||
obs_dict = self.bus.sync_read("Present_Position")
|
obs_dict = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
|
||||||
obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
|
obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
|
||||||
dt_ms = (time.perf_counter() - start) * 1e3
|
dt_ms = (time.perf_counter() - start) * 1e3
|
||||||
logger.debug(f"{self} read state: {dt_ms:.1f}ms")
|
logger.debug(f"{self} read state: {dt_ms:.1f}ms")
|
||||||
@@ -221,7 +221,7 @@ class SOFollower(Robot):
|
|||||||
# Cap goal position when too far away from present position.
|
# Cap goal position when too far away from present position.
|
||||||
# /!\ Slower fps expected due to reading from the follower.
|
# /!\ Slower fps expected due to reading from the follower.
|
||||||
if self.config.max_relative_target is not None:
|
if self.config.max_relative_target is not None:
|
||||||
present_pos = self.bus.sync_read("Present_Position")
|
present_pos = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
|
||||||
goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
|
goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
|
||||||
goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
|
goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
|
||||||
|
|
||||||
|
|||||||
@@ -68,6 +68,10 @@ class UnitreeG1Config(RobotConfig):
|
|||||||
# Compensates for gravity on the unitree's arms using the arm ik solver
|
# Compensates for gravity on the unitree's arms using the arm ik solver
|
||||||
gravity_compensation: bool = False
|
gravity_compensation: bool = False
|
||||||
|
|
||||||
# Lower-body controller class name, e.g. "GrootLocomotionController" or
|
# Locomotion controller class name, e.g. "GrootLocomotionController",
|
||||||
# "HolosomaLocomotionController". None disables it.
|
# "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
|
||||||
|
# Selecting "SonicWholeBodyController" implicitly switches the robot to the 64-D
|
||||||
|
# latent-token action/observation interface (``motion_token.{i}.pos`` action and a
|
||||||
|
# ``motion_token_state.{i}.pos`` state echo) so ``lerobot-rollout`` can drive a
|
||||||
|
# policy trained on SONIC motion tokens (e.g. nepyope/sonic_walk).
|
||||||
controller: str | None = None
|
controller: str | None = None
|
||||||
|
|||||||
@@ -0,0 +1,27 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
|
||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
|
||||||
|
|
||||||
|
from .gr00t_locomotion import GrootLocomotionController
|
||||||
|
from .holosoma_locomotion import HolosomaLocomotionController
|
||||||
|
from .sonic_whole_body import SonicWholeBodyController
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"GrootLocomotionController",
|
||||||
|
"HolosomaLocomotionController",
|
||||||
|
"SonicWholeBodyController",
|
||||||
|
]
|
||||||
+1
-1
@@ -21,7 +21,7 @@ import numpy as np
|
|||||||
import onnxruntime as ort
|
import onnxruntime as ort
|
||||||
from huggingface_hub import hf_hub_download
|
from huggingface_hub import hf_hub_download
|
||||||
|
|
||||||
from .g1_utils import (
|
from ..g1_utils import (
|
||||||
REMOTE_AXES,
|
REMOTE_AXES,
|
||||||
REMOTE_BUTTONS,
|
REMOTE_BUTTONS,
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
+1
-1
@@ -22,7 +22,7 @@ import onnx
|
|||||||
import onnxruntime as ort
|
import onnxruntime as ort
|
||||||
from huggingface_hub import hf_hub_download
|
from huggingface_hub import hf_hub_download
|
||||||
|
|
||||||
from .g1_utils import (
|
from ..g1_utils import (
|
||||||
REMOTE_AXES,
|
REMOTE_AXES,
|
||||||
G1_29_JointArmIndex,
|
G1_29_JointArmIndex,
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
@@ -0,0 +1,360 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
|
||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""SONIC decoder whole-body controller for the Unitree G1 (token-only).
|
||||||
|
|
||||||
|
Pure-Python/ONNX re-implementation of the *decode* half of NVIDIA's SONIC deploy stack.
|
||||||
|
The encoder is intentionally absent: a token-output VLA (e.g. ``nepyope/sonic_walk``)
|
||||||
|
supplies the 64-D latent ``motion_token`` directly each tick, and the SONIC **decoder**
|
||||||
|
maps ``token + recent proprioception history`` to a residual action that is scaled and
|
||||||
|
added onto the standing pose (``default_angles``) to produce 50 Hz joint-position targets
|
||||||
|
for the robot's PD controller.
|
||||||
|
|
||||||
|
Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and
|
||||||
|
**MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils)
|
||||||
|
convert between them. Quaternions are scalar-first ``(w, x, y, z)``.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import onnx
|
||||||
|
import onnxruntime as ort
|
||||||
|
from huggingface_hub import hf_hub_download
|
||||||
|
|
||||||
|
from ..g1_utils import (
|
||||||
|
ISAACLAB_TO_MUJOCO,
|
||||||
|
MUJOCO_TO_ISAACLAB,
|
||||||
|
G1_29_JointIndex,
|
||||||
|
get_gravity_orientation,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# ── Constants (hardware-validated; see the NVIDIA SONIC deploy reference) ──────
|
||||||
|
CONTROL_DT = 0.02 # 50 Hz control period (s)
|
||||||
|
TOKEN_DIM = 64 # decoder latent size
|
||||||
|
|
||||||
|
# SONIC decoder checkpoint: NVIDIA's decoder ONNX re-packaged with its deploy constants
|
||||||
|
# (kp/kd PD gains, the standing pose default_angles, and the residual action_scale) embedded
|
||||||
|
# in the ONNX metadata; see upload_sonic_decoder.py for provisioning. The runtime loads the
|
||||||
|
# model *and* all of these straight from the checkpoint (the Holosoma convention), so no
|
||||||
|
# motor-physics math happens at deploy time.
|
||||||
|
DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
|
||||||
|
DECODER_FILENAME = "model_decoder.onnx"
|
||||||
|
DECODER_INPUT_DIM = 994 # token(64) + 10-frame proprio history + gravity
|
||||||
|
|
||||||
|
|
||||||
|
def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID):
|
||||||
|
"""Load the SONIC decoder ONNX and its baked-in deploy constants from the checkpoint.
|
||||||
|
|
||||||
|
Returns ``(decoder_session, kp, kd, default_angles, action_scale, neutral_token)``. The
|
||||||
|
gains/pose/scale are (29,) float32 in IsaacLab joint order and ``neutral_token`` is the
|
||||||
|
(64,) float32 idle latent -- all read from the ONNX ``metadata_props`` rather than
|
||||||
|
recomputed/hardcoded at deploy time (mirrors ``holosoma_locomotion.load_policy``).
|
||||||
|
"""
|
||||||
|
decoder_path = hf_hub_download(repo_id=repo_id, filename=DECODER_FILENAME)
|
||||||
|
so = ort.SessionOptions()
|
||||||
|
so.log_severity_level = 3 # quiet ORT logs
|
||||||
|
session = ort.InferenceSession(decoder_path, sess_options=so)
|
||||||
|
dec_dim = int(session.get_inputs()[0].shape[1])
|
||||||
|
if dec_dim != DECODER_INPUT_DIM:
|
||||||
|
raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected {DECODER_INPUT_DIM})")
|
||||||
|
|
||||||
|
meta = {p.key: p.value for p in onnx.load(decoder_path, load_external_data=False).metadata_props}
|
||||||
|
required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
|
||||||
|
missing = [k for k in required if k not in meta]
|
||||||
|
if missing:
|
||||||
|
raise ValueError(
|
||||||
|
f"SONIC decoder ONNX at {repo_id} is missing metadata {missing}; "
|
||||||
|
"re-run upload_sonic_decoder.py to (re)provision the checkpoint."
|
||||||
|
)
|
||||||
|
arr = {k: np.array(json.loads(meta[k]), dtype=np.float32) for k in required}
|
||||||
|
logger.info("Loaded SONIC deploy constants from %s (%d joints)", repo_id, len(arr["kp"]))
|
||||||
|
return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
|
||||||
|
|
||||||
|
|
||||||
|
# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
|
||||||
|
TOKEN_ACTION_PREFIX = "motion_token" # nosec B105 - feature-key prefix, not a secret
|
||||||
|
# Proprio-state prefix for the token interface: the robot echoes the last commanded token
|
||||||
|
# here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``.
|
||||||
|
TOKEN_STATE_PREFIX = "motion_token_state" # nosec B105 - feature-key prefix, not a secret
|
||||||
|
|
||||||
|
|
||||||
|
def token_action_key(i: int) -> str:
|
||||||
|
"""Action-dict key for the i-th component of the 64-D SONIC latent token.
|
||||||
|
|
||||||
|
The ``.pos`` suffix is required so the value flows through ``lerobot-rollout``, which
|
||||||
|
only routes ``.pos`` scalar features onto the policy action vector.
|
||||||
|
"""
|
||||||
|
return f"{TOKEN_ACTION_PREFIX}.{i}.pos"
|
||||||
|
|
||||||
|
|
||||||
|
def token_state_key(i: int) -> str:
|
||||||
|
"""Observation key for the i-th component of the 64-D SONIC latent token state."""
|
||||||
|
return f"{TOKEN_STATE_PREFIX}.{i}.pos"
|
||||||
|
|
||||||
|
|
||||||
|
# Startup blend duration: over the first control ticks, linearly interpolate every joint
|
||||||
|
# from the robot's initial measured pose into the policy's commanded target, so control
|
||||||
|
# eases in without a snap on the first command.
|
||||||
|
INIT_RAMP_S = 3.0
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
|
||||||
|
"""Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
|
||||||
|
|
||||||
|
The token-only interface: the caller supplies the 64-D encoder latent directly (e.g. a
|
||||||
|
token-output VLA's action), which the decoder consumes with the encoder bypassed.
|
||||||
|
Requires the full dense token; a partial one is ignored (returns None).
|
||||||
|
"""
|
||||||
|
if not action:
|
||||||
|
return None
|
||||||
|
keys = [token_action_key(i) for i in range(TOKEN_DIM)]
|
||||||
|
if any(key not in action for key in keys):
|
||||||
|
return None
|
||||||
|
return np.fromiter((float(action[key]) for key in keys), dtype=np.float32, count=TOKEN_DIM)
|
||||||
|
|
||||||
|
|
||||||
|
class SonicDecoder:
|
||||||
|
"""Runs the SONIC decoder ONNX model and owns the proprioception history.
|
||||||
|
|
||||||
|
Each tick it appends the latest robot state to 10-frame history buffers, then maps the
|
||||||
|
supplied 64-D ``token`` + that history to a residual action added onto ``default_angles``.
|
||||||
|
The encoder is bypassed entirely (token supplied by the policy). ``default_angles`` and
|
||||||
|
``action_scale`` are (29,) float32 in IsaacLab order, loaded from the checkpoint.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, decoder, default_angles, action_scale):
|
||||||
|
self.decoder = decoder
|
||||||
|
self.decoder_input = decoder.get_inputs()[0].name
|
||||||
|
self.default_angles = np.asarray(default_angles, np.float32)
|
||||||
|
self.action_scale = np.asarray(action_scale, np.float32)
|
||||||
|
self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB]
|
||||||
|
self.token = np.zeros(TOKEN_DIM, np.float32)
|
||||||
|
self.last_action_mj = np.zeros(29, np.float32)
|
||||||
|
self.h_q_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_ang = [np.zeros(3, np.float32)] * 10
|
||||||
|
self.h_act_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
|
||||||
|
|
||||||
|
def reset(self):
|
||||||
|
"""Clear the token and 10-frame proprioception history.
|
||||||
|
|
||||||
|
``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new episode
|
||||||
|
are not contaminated by the previous episode's state.
|
||||||
|
"""
|
||||||
|
self.token = np.zeros(TOKEN_DIM, np.float32)
|
||||||
|
self.last_action_mj = np.zeros(29, np.float32)
|
||||||
|
self.h_q_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_ang = [np.zeros(3, np.float32)] * 10
|
||||||
|
self.h_act_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
|
||||||
|
|
||||||
|
def update_history(self, q, dq, ang, quat):
|
||||||
|
"""Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
|
||||||
|
quat = quat / (np.linalg.norm(quat) + 1e-8)
|
||||||
|
# Reorder IsaacLab-order state into the MuJoCo order the decoder consumes. This
|
||||||
|
# permutation direction is validated against the deployed SONIC ONNX; don't flip it.
|
||||||
|
q_mj = q[MUJOCO_TO_ISAACLAB]
|
||||||
|
dq_mj = dq[MUJOCO_TO_ISAACLAB]
|
||||||
|
self.h_q_mj = [q_mj - self.default_angles_mj] + self.h_q_mj[:-1]
|
||||||
|
self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
|
||||||
|
self.h_ang = [ang.copy()] + self.h_ang[:-1]
|
||||||
|
self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
|
||||||
|
self.h_quat = [quat.copy()] + self.h_quat[:-1]
|
||||||
|
|
||||||
|
def build_decoder_obs(self):
|
||||||
|
"""Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
|
||||||
|
obs = np.zeros(994, np.float32)
|
||||||
|
off = 0
|
||||||
|
obs[off : off + 64] = self.token
|
||||||
|
off += 64
|
||||||
|
for h, sz in [
|
||||||
|
(list(reversed(self.h_ang)), 3),
|
||||||
|
(list(reversed(self.h_q_mj)), 29),
|
||||||
|
(list(reversed(self.h_dq_mj)), 29),
|
||||||
|
(list(reversed(self.h_act_mj)), 29),
|
||||||
|
]:
|
||||||
|
for f in range(10):
|
||||||
|
obs[off : off + sz] = h[f]
|
||||||
|
off += sz
|
||||||
|
for q in reversed(self.h_quat):
|
||||||
|
obs[off : off + 3] = get_gravity_orientation(q)
|
||||||
|
off += 3
|
||||||
|
assert off == 994, f"Decoder obs mismatch: {off}"
|
||||||
|
return obs
|
||||||
|
|
||||||
|
def step(self, robot_obs, token, debug=False):
|
||||||
|
"""One control tick: read robot obs, decode the supplied token -> joint targets.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
|
||||||
|
token: 64-D latent supplied by the policy (encoder bypassed).
|
||||||
|
debug: log action/delta norms.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
|
||||||
|
"""
|
||||||
|
self.token = np.asarray(token, np.float32)
|
||||||
|
jnames = [m.name for m in G1_29_JointIndex]
|
||||||
|
q = np.array(
|
||||||
|
[
|
||||||
|
robot_obs.get(f"{n}.q", self.default_angles[m.value])
|
||||||
|
for m, n in zip(G1_29_JointIndex, jnames, strict=False)
|
||||||
|
],
|
||||||
|
np.float32,
|
||||||
|
)
|
||||||
|
dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
|
||||||
|
quat = np.array(
|
||||||
|
[
|
||||||
|
robot_obs.get("imu.quat.w", 1),
|
||||||
|
robot_obs.get("imu.quat.x", 0),
|
||||||
|
robot_obs.get("imu.quat.y", 0),
|
||||||
|
robot_obs.get("imu.quat.z", 0),
|
||||||
|
],
|
||||||
|
np.float32,
|
||||||
|
)
|
||||||
|
ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
|
||||||
|
self.update_history(q, dq, ang, quat)
|
||||||
|
action_mj = (
|
||||||
|
self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
|
||||||
|
.squeeze()
|
||||||
|
.astype(np.float32)
|
||||||
|
)
|
||||||
|
self.last_action_mj = action_mj.copy()
|
||||||
|
target = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.action_scale
|
||||||
|
if debug:
|
||||||
|
delta = target - q
|
||||||
|
logger.debug(
|
||||||
|
"token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
|
||||||
|
np.linalg.norm(self.token),
|
||||||
|
np.linalg.norm(action_mj),
|
||||||
|
np.max(np.abs(delta)),
|
||||||
|
np.sqrt(np.mean(delta**2)),
|
||||||
|
)
|
||||||
|
return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
|
||||||
|
|
||||||
|
|
||||||
|
class SonicRuntime:
|
||||||
|
"""Loads the SONIC decoder ONNX model and owns the decode controller.
|
||||||
|
|
||||||
|
Token-only deploy: the encoder is bypassed; each tick the decoder consumes a 64-D
|
||||||
|
latent token supplied directly by the policy.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
decoder_sess, self.kp, self.kd, default_angles, action_scale, neutral_token = load_sonic_decoder()
|
||||||
|
self.default_angles = default_angles
|
||||||
|
self.neutral_token = neutral_token
|
||||||
|
self.controller = SonicDecoder(decoder_sess, default_angles, action_scale)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pipeline(self):
|
||||||
|
return self.controller
|
||||||
|
|
||||||
|
def reset(self):
|
||||||
|
self.controller.reset()
|
||||||
|
|
||||||
|
def shutdown(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SonicWholeBodyController:
|
||||||
|
"""Full-body SONIC controller for UnitreeG1's background controller thread."""
|
||||||
|
|
||||||
|
control_dt = CONTROL_DT
|
||||||
|
full_body = True
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
logger.info("Loading SONIC whole-body controller...")
|
||||||
|
self._runtime = SonicRuntime()
|
||||||
|
self.kp = self._runtime.kp
|
||||||
|
self.kd = self._runtime.kd
|
||||||
|
self.controller = self._runtime.controller
|
||||||
|
self._default_angles = self._runtime.default_angles
|
||||||
|
self._neutral_token = self._runtime.neutral_token
|
||||||
|
|
||||||
|
# Startup blend: ease from the robot's initial pose into the first commanded policy
|
||||||
|
# targets over INIT_RAMP_S (captured on the first control tick).
|
||||||
|
self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
|
||||||
|
self._init_step = 0
|
||||||
|
self._start_pose: dict[str, float] = {}
|
||||||
|
|
||||||
|
# Token-interface state. The controller holds a stable *neutral* token until the first
|
||||||
|
# real token arrives, and afterwards holds the *last* token received between ticks (the
|
||||||
|
# async controller runs ~50 Hz while a token VLA streams ~30 Hz).
|
||||||
|
self._last_token: np.ndarray | None = None
|
||||||
|
|
||||||
|
logger.info("SONIC ready (decoder, 64-D token command path)")
|
||||||
|
|
||||||
|
def _startup_blend(self, obs: dict, out: dict) -> dict:
|
||||||
|
"""Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
|
||||||
|
interpolate between the robot's pose captured on the first tick and the policy's
|
||||||
|
live commanded target, so the handoff has no snap.
|
||||||
|
|
||||||
|
``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
|
||||||
|
climbs 0->1 over the ramp, after which the raw policy target passes through.
|
||||||
|
"""
|
||||||
|
if self._init_step >= self._init_ramp_steps or not out:
|
||||||
|
return out
|
||||||
|
if self._init_step == 0:
|
||||||
|
# Capture the robot's actual pose as the interpolation start point.
|
||||||
|
self._start_pose = {
|
||||||
|
f"{m.name}.q": float(obs.get(f"{m.name}.q", self._default_angles[m.value]))
|
||||||
|
for m in G1_29_JointIndex
|
||||||
|
}
|
||||||
|
self._init_step += 1
|
||||||
|
ratio = min(1.0, self._init_step / self._init_ramp_steps)
|
||||||
|
blended = {
|
||||||
|
k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
|
||||||
|
for k, tgt in out.items()
|
||||||
|
}
|
||||||
|
if self._init_step >= self._init_ramp_steps:
|
||||||
|
logger.info("SONIC startup blend complete -> full policy control")
|
||||||
|
return blended
|
||||||
|
|
||||||
|
def run_step(self, action: dict, obs: dict) -> dict:
|
||||||
|
if not obs:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
# Token-only interface (token-output VLA): a dense 64-D ``motion_token.{i}`` command
|
||||||
|
# is decoded directly, encoder bypassed.
|
||||||
|
token = _extract_token_from_action(action)
|
||||||
|
if token is not None:
|
||||||
|
self._last_token = token
|
||||||
|
elif self._last_token is None:
|
||||||
|
# No token has arrived yet: hold the checkpoint's neutral token, which the decoder
|
||||||
|
# maps to a stable, natural standing pose.
|
||||||
|
self._last_token = self._neutral_token.copy()
|
||||||
|
# Either a fresh token this tick or the last one received (held between the ~30 Hz
|
||||||
|
# token stream and the ~50 Hz control loop).
|
||||||
|
return self._startup_blend(obs, self.controller.step(obs, self._last_token))
|
||||||
|
|
||||||
|
def reset(self):
|
||||||
|
self._runtime.reset()
|
||||||
|
self._init_step = 0 # re-run the startup blend after a reset
|
||||||
|
self._start_pose = {}
|
||||||
|
# Drop the held token so the neutral token is re-seeded after a reset.
|
||||||
|
self._last_token = None
|
||||||
|
|
||||||
|
def shutdown(self):
|
||||||
|
self._runtime.shutdown()
|
||||||
@@ -23,6 +23,47 @@ import numpy as np
|
|||||||
|
|
||||||
NUM_MOTORS = 29
|
NUM_MOTORS = 29
|
||||||
|
|
||||||
|
# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
|
||||||
|
# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
|
||||||
|
# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
|
||||||
|
ISAACLAB_TO_MUJOCO = np.array(
|
||||||
|
[
|
||||||
|
0,
|
||||||
|
3,
|
||||||
|
6,
|
||||||
|
9,
|
||||||
|
13,
|
||||||
|
17,
|
||||||
|
1,
|
||||||
|
4,
|
||||||
|
7,
|
||||||
|
10,
|
||||||
|
14,
|
||||||
|
18,
|
||||||
|
2,
|
||||||
|
5,
|
||||||
|
8,
|
||||||
|
11,
|
||||||
|
15,
|
||||||
|
19,
|
||||||
|
21,
|
||||||
|
23,
|
||||||
|
25,
|
||||||
|
27,
|
||||||
|
12,
|
||||||
|
16,
|
||||||
|
20,
|
||||||
|
22,
|
||||||
|
24,
|
||||||
|
26,
|
||||||
|
28,
|
||||||
|
],
|
||||||
|
dtype=np.int32,
|
||||||
|
)
|
||||||
|
# The two orderings are inverses of each other, so derive one from the other (argsort) to
|
||||||
|
# guarantee they can never drift out of sync.
|
||||||
|
MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32)
|
||||||
|
|
||||||
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
|
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
|
||||||
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
|
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
|
||||||
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
|
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
|
||||||
@@ -68,8 +109,9 @@ def make_locomotion_controller(name: str | None):
|
|||||||
if name is None:
|
if name is None:
|
||||||
return None
|
return None
|
||||||
controllers = {
|
controllers = {
|
||||||
"GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
|
"GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
|
||||||
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
|
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
|
||||||
|
"SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
|
||||||
}
|
}
|
||||||
module_path = controllers.get(name)
|
module_path = controllers.get(name)
|
||||||
if module_path is None:
|
if module_path is None:
|
||||||
|
|||||||
@@ -34,7 +34,6 @@ from .config_unitree_g1 import UnitreeG1Config
|
|||||||
from .g1_kinematics import G1_29_ArmIK
|
from .g1_kinematics import G1_29_ArmIK
|
||||||
from .g1_utils import (
|
from .g1_utils import (
|
||||||
REMOTE_AXES,
|
REMOTE_AXES,
|
||||||
REMOTE_KEYS,
|
|
||||||
G1_29_JointArmIndex,
|
G1_29_JointArmIndex,
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
default_remote_input,
|
default_remote_input,
|
||||||
@@ -148,22 +147,54 @@ class UnitreeG1(Robot):
|
|||||||
|
|
||||||
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
|
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
|
||||||
|
|
||||||
# Lower-body controller loaded dynamically
|
# Lower-body / whole-body controller loaded dynamically
|
||||||
self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
|
self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
|
||||||
|
|
||||||
# Controller thread state
|
# Controller thread state
|
||||||
self._controller_thread = None
|
self._controller_thread = None
|
||||||
|
# When set, the controller loop stops publishing low commands so reset() can
|
||||||
|
# drive the joints directly without two publishers fighting (single-publisher).
|
||||||
|
self._controller_paused = threading.Event()
|
||||||
self._controller_action_lock = threading.Lock()
|
self._controller_action_lock = threading.Lock()
|
||||||
self.controller_input = default_remote_input()
|
self.controller_input = default_remote_input()
|
||||||
self.controller_output = {}
|
self.controller_output = {}
|
||||||
|
|
||||||
|
# Token-mode state: last 64-D SONIC latent token commanded by the policy,
|
||||||
|
# echoed back as ``observation.state`` so a token-output VLA closes the loop
|
||||||
|
# on its own previous token. Implicit whenever the SONIC whole-body controller
|
||||||
|
# is active. Seeded to zeros; the controller's startup blend eases joints in.
|
||||||
|
self._last_token: np.ndarray | None = None
|
||||||
|
if self._sonic_token:
|
||||||
|
from .controllers.sonic_whole_body import TOKEN_DIM
|
||||||
|
|
||||||
|
self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def _sonic_token(self) -> bool:
|
||||||
|
"""Whether the SONIC whole-body decoder is active.
|
||||||
|
|
||||||
|
A SONIC controller consumes a 64-D latent motion token as its action and echoes
|
||||||
|
the last commanded token as ``observation.state``. Keyed purely off the selected
|
||||||
|
controller so the token interface is implicit -- no separate config flag.
|
||||||
|
"""
|
||||||
|
return self.config.controller == "SonicWholeBodyController"
|
||||||
|
|
||||||
def _subscribe_lowstate(self): # polls robot state @ 250Hz
|
def _subscribe_lowstate(self): # polls robot state @ 250Hz
|
||||||
while not self._shutdown_event.is_set():
|
while not self._shutdown_event.is_set():
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
|
|
||||||
# Step simulation if in simulation mode
|
# Step simulation if in simulation mode
|
||||||
if self.config.is_simulation and self.sim_env is not None:
|
if self.config.is_simulation and self.sim_env is not None:
|
||||||
self.sim_env.step()
|
try:
|
||||||
|
self.sim_env.step()
|
||||||
|
except ValueError as e:
|
||||||
|
# Startup race: the sim thread can step once before reset() has
|
||||||
|
# written a valid base pose, giving a zero-norm pelvis quaternion
|
||||||
|
# (scipy>=1.11 raises instead of normalizing). Skip and retry so
|
||||||
|
# the thread survives instead of dying and freezing the sim.
|
||||||
|
if "zero norm" not in str(e).lower():
|
||||||
|
raise
|
||||||
|
time.sleep(self.control_dt)
|
||||||
|
continue
|
||||||
|
|
||||||
msg = self.lowstate_subscriber.Read()
|
msg = self.lowstate_subscriber.Read()
|
||||||
if msg is not None:
|
if msg is not None:
|
||||||
@@ -231,15 +262,38 @@ class UnitreeG1(Robot):
|
|||||||
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
||||||
return features
|
return features
|
||||||
|
|
||||||
|
@property
|
||||||
|
def _token_state_ft(self) -> dict[str, type]:
|
||||||
|
"""64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
|
||||||
|
|
||||||
|
Exposed only when a SONIC whole-body controller is active; aggregated by the
|
||||||
|
rollout into a 64-D ``observation.state`` (the last token the policy commanded).
|
||||||
|
"""
|
||||||
|
if not self._sonic_token:
|
||||||
|
return {}
|
||||||
|
from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
|
||||||
|
|
||||||
|
return {token_state_key(i): float for i in range(TOKEN_DIM)}
|
||||||
|
|
||||||
@cached_property
|
@cached_property
|
||||||
def observation_features(self) -> dict[str, type | tuple]:
|
def observation_features(self) -> dict[str, type | tuple]:
|
||||||
return {**self._motors_ft, **self._cameras_ft}
|
return {**self._motors_ft, **self._token_state_ft, **self._cameras_ft}
|
||||||
|
|
||||||
@cached_property
|
@cached_property
|
||||||
def action_features(self) -> dict[str, type]:
|
def action_features(self) -> dict[str, type]:
|
||||||
|
# No controller configured at all: raw 29-DoF joint teleop.
|
||||||
if self.controller is None:
|
if self.controller is None:
|
||||||
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
|
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
|
||||||
|
|
||||||
|
# Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
|
||||||
|
# (``motion_token.{i}.pos``) so ``lerobot-rollout`` maps a 64-D policy output
|
||||||
|
# straight onto the decoder, bypassing the encoder.
|
||||||
|
if self._sonic_token:
|
||||||
|
from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key
|
||||||
|
|
||||||
|
return {token_action_key(i): float for i in range(TOKEN_DIM)}
|
||||||
|
|
||||||
|
# Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
|
||||||
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
|
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
|
||||||
remote_features = dict.fromkeys(REMOTE_AXES, float)
|
remote_features = dict.fromkeys(REMOTE_AXES, float)
|
||||||
return {**arm_features, **remote_features}
|
return {**arm_features, **remote_features}
|
||||||
@@ -255,6 +309,11 @@ class UnitreeG1(Robot):
|
|||||||
while not self._shutdown_event.is_set():
|
while not self._shutdown_event.is_set():
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
|
|
||||||
|
# Paused during reset() so the reset routine is the sole low-cmd publisher.
|
||||||
|
if self._controller_paused.is_set():
|
||||||
|
time.sleep(control_dt)
|
||||||
|
continue
|
||||||
|
|
||||||
with self._lowstate_lock:
|
with self._lowstate_lock:
|
||||||
lowstate = self._lowstate
|
lowstate = self._lowstate
|
||||||
|
|
||||||
@@ -271,8 +330,12 @@ class UnitreeG1(Robot):
|
|||||||
with self._controller_action_lock:
|
with self._controller_action_lock:
|
||||||
controller_input = dict(self.controller_input)
|
controller_input = dict(self.controller_input)
|
||||||
|
|
||||||
# Run controller step
|
# Full-body controllers (SONIC) consume the full observation dict; others
|
||||||
controller_action = self.controller.run_step(controller_input, lowstate)
|
# take the raw lowstate. get_observation() is the single lowstate -> obs builder.
|
||||||
|
controller_state = (
|
||||||
|
self.get_observation() if getattr(self.controller, "full_body", False) else lowstate
|
||||||
|
)
|
||||||
|
controller_action = self.controller.run_step(controller_input, controller_state)
|
||||||
|
|
||||||
# Write controller output snapshot
|
# Write controller output snapshot
|
||||||
with self._controller_action_lock:
|
with self._controller_action_lock:
|
||||||
@@ -343,6 +406,9 @@ class UnitreeG1(Robot):
|
|||||||
|
|
||||||
self.kp = np.array(self.config.kp, dtype=np.float32)
|
self.kp = np.array(self.config.kp, dtype=np.float32)
|
||||||
self.kd = np.array(self.config.kd, dtype=np.float32)
|
self.kd = np.array(self.config.kd, dtype=np.float32)
|
||||||
|
if self.controller is not None and hasattr(self.controller, "kp"):
|
||||||
|
self.kp = np.array(self.controller.kp, dtype=np.float32)
|
||||||
|
self.kd = np.array(self.controller.kd, dtype=np.float32)
|
||||||
|
|
||||||
for joint in G1_29_JointIndex:
|
for joint in G1_29_JointIndex:
|
||||||
self.msg.motor_cmd[joint].mode = 1
|
self.msg.motor_cmd[joint].mode = 1
|
||||||
@@ -391,6 +457,10 @@ class UnitreeG1(Robot):
|
|||||||
if self._controller_thread.is_alive():
|
if self._controller_thread.is_alive():
|
||||||
logger.warning("Controller thread did not stop cleanly")
|
logger.warning("Controller thread did not stop cleanly")
|
||||||
|
|
||||||
|
# Release controller resources (e.g. SONIC decoder sessions).
|
||||||
|
if self.controller is not None and hasattr(self.controller, "shutdown"):
|
||||||
|
self.controller.shutdown()
|
||||||
|
|
||||||
# Close simulation environment
|
# Close simulation environment
|
||||||
if self.config.is_simulation and self.sim_env is not None:
|
if self.config.is_simulation and self.sim_env is not None:
|
||||||
try:
|
try:
|
||||||
@@ -461,6 +531,15 @@ class UnitreeG1(Robot):
|
|||||||
if lowstate.wireless_remote:
|
if lowstate.wireless_remote:
|
||||||
obs["wireless_remote"] = lowstate.wireless_remote
|
obs["wireless_remote"] = lowstate.wireless_remote
|
||||||
|
|
||||||
|
# Token mode: echo the last commanded latent token as observation.state so a
|
||||||
|
# token-output VLA closes the loop on its own previous token.
|
||||||
|
if self._sonic_token:
|
||||||
|
from .controllers.sonic_whole_body import token_state_key
|
||||||
|
|
||||||
|
token = self._last_token if self._last_token is not None else []
|
||||||
|
for i, v in enumerate(token):
|
||||||
|
obs[token_state_key(i)] = float(v)
|
||||||
|
|
||||||
# Cameras - read images from ZMQ cameras
|
# Cameras - read images from ZMQ cameras
|
||||||
for cam_name, cam in self._cameras.items():
|
for cam_name, cam in self._cameras.items():
|
||||||
if getattr(cam, "use_rgb", True):
|
if getattr(cam, "use_rgb", True):
|
||||||
@@ -473,9 +552,22 @@ class UnitreeG1(Robot):
|
|||||||
def send_action(self, action: RobotAction) -> RobotAction:
|
def send_action(self, action: RobotAction) -> RobotAction:
|
||||||
action_to_publish = action
|
action_to_publish = action
|
||||||
if self.controller is not None:
|
if self.controller is not None:
|
||||||
|
# SONIC decoder: pull the 64-D latent token out of the action and remember it
|
||||||
|
# for the observation.state echo. The controller thread reads it back from
|
||||||
|
# controller_input (populated below) and decodes it into a 29-DoF command.
|
||||||
|
if self._sonic_token:
|
||||||
|
from .controllers.sonic_whole_body import _extract_token_from_action
|
||||||
|
|
||||||
|
token = _extract_token_from_action(action)
|
||||||
|
if token is not None:
|
||||||
|
self._last_token = token
|
||||||
|
self._update_controller_action(action)
|
||||||
|
# Full-body controllers (SONIC) own the whole 29-DoF command; nothing to
|
||||||
|
# publish here (the controller thread is the sole publisher).
|
||||||
|
if getattr(self.controller, "full_body", False):
|
||||||
|
return action
|
||||||
# Controller thread owns legs/waist. Here we only update joystick inputs
|
# Controller thread owns legs/waist. Here we only update joystick inputs
|
||||||
# and publish arm targets from the teleoperator.
|
# and publish arm targets from the teleoperator.
|
||||||
self._update_controller_action(action)
|
|
||||||
arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
|
arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
|
||||||
action_to_publish = {
|
action_to_publish = {
|
||||||
key: value
|
key: value
|
||||||
@@ -503,11 +595,17 @@ class UnitreeG1(Robot):
|
|||||||
return action
|
return action
|
||||||
|
|
||||||
def _update_controller_action(self, action: RobotAction) -> None:
|
def _update_controller_action(self, action: RobotAction) -> None:
|
||||||
"""Update controller input state from incoming teleop action."""
|
"""Update controller input state from an incoming teleop action.
|
||||||
|
|
||||||
|
Controller-agnostic: every value-carrying key (locomotion ``remote.*`` axes or
|
||||||
|
SONIC ``motion_token.*`` values) is forwarded verbatim into ``controller_input``
|
||||||
|
and each controller extracts only the keys it understands. The robot deliberately
|
||||||
|
does not enumerate any controller's key schema here.
|
||||||
|
"""
|
||||||
with self._controller_action_lock:
|
with self._controller_action_lock:
|
||||||
for key in REMOTE_KEYS:
|
for key, value in action.items():
|
||||||
if key in action:
|
if isinstance(key, str) and value is not None:
|
||||||
self.controller_input[key] = action[key]
|
self.controller_input[key] = value
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_calibrated(self) -> bool:
|
def is_calibrated(self) -> bool:
|
||||||
@@ -537,43 +635,64 @@ class UnitreeG1(Robot):
|
|||||||
if default_positions is None:
|
if default_positions is None:
|
||||||
default_positions = np.array(self.config.default_positions, dtype=np.float32)
|
default_positions = np.array(self.config.default_positions, dtype=np.float32)
|
||||||
|
|
||||||
if self.config.is_simulation and self.sim_env is not None:
|
# Full-body controllers (SONIC) own the whole 29-DoF command and ignore
|
||||||
self.sim_env.reset()
|
# ``<joint>.q`` in send_action(), so reset() must publish the default pose
|
||||||
self.publish_lowcmd(
|
# directly. Pause the background controller first so the two aren't both writing
|
||||||
{f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
|
# low commands while the robot moves to the default pose.
|
||||||
)
|
full_body = getattr(self.controller, "full_body", False)
|
||||||
else:
|
paused = False
|
||||||
total_time = 3.0
|
if full_body and self._controller_thread is not None:
|
||||||
num_steps = int(total_time / control_dt)
|
self._controller_paused.set()
|
||||||
|
paused = True
|
||||||
|
time.sleep(control_dt) # let any in-flight controller tick settle
|
||||||
|
|
||||||
# get current state
|
try:
|
||||||
obs = self.get_observation()
|
if self.config.is_simulation and self.sim_env is not None:
|
||||||
|
self.sim_env.reset()
|
||||||
|
self.publish_lowcmd(
|
||||||
|
{f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
total_time = 3.0
|
||||||
|
num_steps = int(total_time / control_dt)
|
||||||
|
|
||||||
# record current positions
|
# get current state
|
||||||
init_dof_pos = np.zeros(29, dtype=np.float32)
|
obs = self.get_observation()
|
||||||
for motor in G1_29_JointIndex:
|
|
||||||
init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
|
|
||||||
|
|
||||||
# Interpolate to default position
|
# record current positions
|
||||||
for step in range(num_steps):
|
init_dof_pos = np.zeros(29, dtype=np.float32)
|
||||||
start_time = time.time()
|
|
||||||
|
|
||||||
alpha = step / num_steps
|
|
||||||
action_dict = {}
|
|
||||||
for motor in G1_29_JointIndex:
|
for motor in G1_29_JointIndex:
|
||||||
target_pos = default_positions[motor.value]
|
init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
|
||||||
interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
|
|
||||||
action_dict[f"{motor.name}.q"] = float(interp_pos)
|
|
||||||
|
|
||||||
self.send_action(action_dict)
|
# Interpolate to default position
|
||||||
|
for step in range(num_steps):
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
# Maintain constant control rate
|
alpha = step / num_steps
|
||||||
elapsed = time.time() - start_time
|
action_dict = {}
|
||||||
sleep_time = max(0, control_dt - elapsed)
|
for motor in G1_29_JointIndex:
|
||||||
time.sleep(sleep_time)
|
target_pos = default_positions[motor.value]
|
||||||
|
interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
|
||||||
|
action_dict[f"{motor.name}.q"] = float(interp_pos)
|
||||||
|
|
||||||
# Reset controller internal state (gait phase, obs history, etc.)
|
# Full-body controllers no-op in send_action(); publish the pose
|
||||||
if self.controller is not None and hasattr(self.controller, "reset"):
|
# directly (arm-only controllers keep the send_action() path).
|
||||||
self.controller.reset()
|
if full_body:
|
||||||
|
self.publish_lowcmd(action_dict)
|
||||||
|
else:
|
||||||
|
self.send_action(action_dict)
|
||||||
|
|
||||||
|
# Maintain constant control rate
|
||||||
|
elapsed = time.time() - start_time
|
||||||
|
sleep_time = max(0, control_dt - elapsed)
|
||||||
|
time.sleep(sleep_time)
|
||||||
|
|
||||||
|
# Reset controller internal state (gait phase, obs history, etc.) before
|
||||||
|
# resuming so its buffers reflect the post-reset pose.
|
||||||
|
if self.controller is not None and hasattr(self.controller, "reset"):
|
||||||
|
self.controller.reset()
|
||||||
|
finally:
|
||||||
|
if paused:
|
||||||
|
self._controller_paused.clear()
|
||||||
|
|
||||||
logger.info("Reached default position")
|
logger.info("Reached default position")
|
||||||
|
|||||||
@@ -24,6 +24,7 @@ from __future__ import annotations
|
|||||||
import logging
|
import logging
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from threading import Event
|
from threading import Event
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
@@ -47,6 +48,7 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor
|
|||||||
from lerobot.robots import make_robot_from_config
|
from lerobot.robots import make_robot_from_config
|
||||||
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
|
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
|
||||||
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
|
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
|
||||||
|
from lerobot.utils.import_utils import _peft_available, require_package
|
||||||
|
|
||||||
from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
|
from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
|
||||||
from .inference import (
|
from .inference import (
|
||||||
@@ -57,6 +59,12 @@ from .inference import (
|
|||||||
)
|
)
|
||||||
from .robot_wrapper import ThreadSafeRobot
|
from .robot_wrapper import ThreadSafeRobot
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import PeftConfig, PeftModel
|
||||||
|
else:
|
||||||
|
PeftConfig = None
|
||||||
|
PeftModel = None
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -171,7 +179,7 @@ def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy
|
|||||||
revision=pretrained_revision,
|
revision=pretrained_revision,
|
||||||
)
|
)
|
||||||
|
|
||||||
from peft import PeftConfig, PeftModel
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
peft_path = policy_config.pretrained_path
|
peft_path = policy_config.pretrained_path
|
||||||
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
|
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
|
||||||
@@ -294,12 +302,22 @@ def build_rollout_context(
|
|||||||
# ``observation_features`` values are either a tuple (camera shape) or the
|
# ``observation_features`` values are either a tuple (camera shape) or the
|
||||||
# ``float`` type itself used as a sentinel for scalar motor features —
|
# ``float`` type itself used as a sentinel for scalar motor features —
|
||||||
# see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
|
# see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
|
||||||
|
# Keep cameras (tuple) plus both joint-position (.pos) and base-velocity (.vel)
|
||||||
|
# scalar state features. LeKiwi's observation.state is 9-dim (6 arm .pos +
|
||||||
|
# x/y/theta.vel) and the policy was trained/normalized on all 9; the old .pos-only
|
||||||
|
# filter fed a 6-dim state into a 9-dim normalizer → RuntimeError (size 6 vs 9).
|
||||||
|
# Pure-arm robots have no .vel state keys, so this is a no-op for them.
|
||||||
observation_features_hw = {
|
observation_features_hw = {
|
||||||
k: v
|
k: v
|
||||||
for k, v in all_obs_features.items()
|
for k, v in all_obs_features.items()
|
||||||
if isinstance(v, tuple) or (v is float and k.endswith(".pos"))
|
if isinstance(v, tuple) or (v is float and k.endswith((".pos", ".vel")))
|
||||||
}
|
}
|
||||||
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith(".pos")}
|
# Keep both joint-position (.pos) and base-velocity (.vel) action features so
|
||||||
|
# mobile manipulators command the base too (e.g. LeKiwi: 6 arm .pos +
|
||||||
|
# x/y/theta.vel = 9-dim action). Pure-arm robots have no .vel keys, so this is
|
||||||
|
# a no-op for them. Without the .vel keys the base velocities are silently
|
||||||
|
# dropped from dataset_features[ACTION]/ordered_action_keys and the base never moves.
|
||||||
|
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith((".pos", ".vel"))}
|
||||||
|
|
||||||
# The action side is always needed: sync inference reads action names from
|
# The action side is always needed: sync inference reads action names from
|
||||||
# ``dataset_features[ACTION]`` to map policy tensors back to robot actions.
|
# ``dataset_features[ACTION]`` to map policy tensors back to robot actions.
|
||||||
|
|||||||
@@ -36,6 +36,7 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
|
|||||||
import argparse
|
import argparse
|
||||||
import concurrent.futures
|
import concurrent.futures
|
||||||
import logging
|
import logging
|
||||||
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -52,6 +53,7 @@ from lerobot.datasets import (
|
|||||||
get_feature_stats,
|
get_feature_stats,
|
||||||
write_stats,
|
write_stats,
|
||||||
)
|
)
|
||||||
|
from lerobot.datasets.compute_stats import sample_indices
|
||||||
from lerobot.utils.utils import init_logging
|
from lerobot.utils.utils import init_logging
|
||||||
|
|
||||||
|
|
||||||
@@ -77,12 +79,14 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
|
def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
|
||||||
"""Process a single episode and return its statistics.
|
"""Process a single episode and return its statistics.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
dataset: The LeRobot dataset
|
dataset: The LeRobot dataset
|
||||||
episode_idx: Index of the episode to process
|
episode_idx: Index of the episode to process
|
||||||
|
use_sampling: If True, sub-sample image/video frames per episode to bound
|
||||||
|
memory. If False, use every frame (exact, higher memory).
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Dictionary containing episode statistics
|
Dictionary containing episode statistics
|
||||||
@@ -92,16 +96,31 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
|
|||||||
start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
|
start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
|
||||||
end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
|
end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
|
||||||
|
|
||||||
collected_data: dict[str, list] = {}
|
episode_len = end_idx - start_idx
|
||||||
for idx in range(start_idx, end_idx):
|
|
||||||
item = dataset[idx]
|
|
||||||
for key, value in item.items():
|
|
||||||
if key not in dataset.features:
|
|
||||||
continue
|
|
||||||
|
|
||||||
if key not in collected_data:
|
# Images/video are the memory hog, so sub-sample those frames per episode;
|
||||||
collected_data[key] = []
|
# numeric columns are cheap, so read them in full (exact).
|
||||||
collected_data[key].append(value)
|
image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")]
|
||||||
|
numeric_keys = [
|
||||||
|
k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
|
||||||
|
]
|
||||||
|
|
||||||
|
collected_data: dict[str, list] = {}
|
||||||
|
|
||||||
|
# Numeric features: every frame, read directly from the underlying table.
|
||||||
|
if numeric_keys:
|
||||||
|
numeric_cols = dataset.hf_dataset.select_columns(numeric_keys)[start_idx:end_idx]
|
||||||
|
for key in numeric_keys:
|
||||||
|
collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]]
|
||||||
|
|
||||||
|
# Image/video features: decode only a sampled subset of frames.
|
||||||
|
if image_keys:
|
||||||
|
sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len))
|
||||||
|
for offset in sampled_offsets:
|
||||||
|
item = dataset[start_idx + offset]
|
||||||
|
for key in image_keys:
|
||||||
|
if key in item:
|
||||||
|
collected_data.setdefault(key, []).append(item[key])
|
||||||
|
|
||||||
ep_stats = {}
|
ep_stats = {}
|
||||||
for key, data_list in collected_data.items():
|
for key, data_list in collected_data.items():
|
||||||
@@ -131,11 +150,13 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
|
|||||||
return ep_stats
|
return ep_stats
|
||||||
|
|
||||||
|
|
||||||
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]:
|
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
|
||||||
"""Compute quantile statistics for all episodes in the dataset.
|
"""Compute quantile statistics for all episodes in the dataset.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
dataset: The LeRobot dataset to compute statistics for
|
dataset: The LeRobot dataset to compute statistics for
|
||||||
|
use_sampling: If True, sub-sample image/video frames per episode to bound
|
||||||
|
memory. If False, use every frame (exact, higher memory).
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Dictionary containing aggregated statistics with quantiles
|
Dictionary containing aggregated statistics with quantiles
|
||||||
@@ -153,15 +174,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dic
|
|||||||
if has_videos:
|
if has_videos:
|
||||||
logging.info("Dataset contains video keys - using sequential processing for thread safety")
|
logging.info("Dataset contains video keys - using sequential processing for thread safety")
|
||||||
for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
|
for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
|
||||||
ep_stats = process_single_episode(dataset, episode_idx)
|
ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
|
||||||
episode_stats_list.append(ep_stats)
|
episode_stats_list.append(ep_stats)
|
||||||
else:
|
else:
|
||||||
logging.info("Dataset has no video keys - using parallel processing for better performance")
|
logging.info("Dataset has no video keys - using parallel processing for better performance")
|
||||||
max_workers = min(dataset.num_episodes, 16)
|
max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))
|
||||||
|
|
||||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||||
future_to_episode = {
|
future_to_episode = {
|
||||||
executor.submit(process_single_episode, dataset, episode_idx): episode_idx
|
executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
|
||||||
for episode_idx in range(dataset.num_episodes)
|
for episode_idx in range(dataset.num_episodes)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -188,6 +209,7 @@ def augment_dataset_with_quantile_stats(
|
|||||||
repo_id: str,
|
repo_id: str,
|
||||||
root: str | Path | None = None,
|
root: str | Path | None = None,
|
||||||
overwrite: bool = False,
|
overwrite: bool = False,
|
||||||
|
use_sampling: bool = True,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Augment a dataset with quantile statistics if they are missing.
|
"""Augment a dataset with quantile statistics if they are missing.
|
||||||
|
|
||||||
@@ -195,6 +217,8 @@ def augment_dataset_with_quantile_stats(
|
|||||||
repo_id: Repository ID of the dataset
|
repo_id: Repository ID of the dataset
|
||||||
root: Local root directory for the dataset
|
root: Local root directory for the dataset
|
||||||
overwrite: Overwrite existing quantile statistics if they already exist
|
overwrite: Overwrite existing quantile statistics if they already exist
|
||||||
|
use_sampling: If True, sub-sample image/video frames per episode to bound
|
||||||
|
memory. If False, use every frame (exact, higher memory).
|
||||||
"""
|
"""
|
||||||
logging.info(f"Loading dataset: {repo_id}")
|
logging.info(f"Loading dataset: {repo_id}")
|
||||||
dataset = LeRobotDataset(
|
dataset = LeRobotDataset(
|
||||||
@@ -208,7 +232,7 @@ def augment_dataset_with_quantile_stats(
|
|||||||
|
|
||||||
logging.info("Dataset does not contain quantile statistics. Computing them now...")
|
logging.info("Dataset does not contain quantile statistics. Computing them now...")
|
||||||
|
|
||||||
new_stats = compute_quantile_stats_for_dataset(dataset)
|
new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
|
||||||
|
|
||||||
logging.info("Updating dataset metadata with new quantile statistics")
|
logging.info("Updating dataset metadata with new quantile statistics")
|
||||||
dataset.meta.stats = new_stats
|
dataset.meta.stats = new_stats
|
||||||
@@ -248,6 +272,14 @@ def main():
|
|||||||
action="store_true",
|
action="store_true",
|
||||||
help="Overwrite existing quantile statistics if they already exist",
|
help="Overwrite existing quantile statistics if they already exist",
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--no-sampling",
|
||||||
|
action="store_true",
|
||||||
|
help=(
|
||||||
|
"Compute stats over every frame (exact, higher memory). By default, "
|
||||||
|
"image/video frames are sub-sampled per episode to bound memory."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
root = Path(args.root) if args.root else None
|
root = Path(args.root) if args.root else None
|
||||||
@@ -258,6 +290,7 @@ def main():
|
|||||||
repo_id=args.repo_id,
|
repo_id=args.repo_id,
|
||||||
root=root,
|
root=root,
|
||||||
overwrite=args.overwrite,
|
overwrite=args.overwrite,
|
||||||
|
use_sampling=not args.no_sampling,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -94,6 +94,8 @@ from lerobot.datasets.video_utils import concatenate_video_files, get_video_dura
|
|||||||
from lerobot.utils.constants import HF_LEROBOT_HOME
|
from lerobot.utils.constants import HF_LEROBOT_HOME
|
||||||
from lerobot.utils.utils import flatten_dict, init_logging
|
from lerobot.utils.utils import flatten_dict, init_logging
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
V21 = "v2.1"
|
V21 = "v2.1"
|
||||||
V30 = "v3.0"
|
V30 = "v3.0"
|
||||||
|
|
||||||
@@ -476,11 +478,11 @@ def convert_dataset(
|
|||||||
# First check if the dataset already has a v3.0 version
|
# First check if the dataset already has a v3.0 version
|
||||||
if root is None and not force_conversion:
|
if root is None and not force_conversion:
|
||||||
try:
|
try:
|
||||||
print("Trying to download v3.0 version of the dataset from the hub...")
|
logger.info("Trying to download v3.0 version of the dataset from the hub...")
|
||||||
snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
|
snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
|
||||||
return
|
return
|
||||||
except Exception:
|
except Exception:
|
||||||
print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
|
logger.info("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
|
||||||
|
|
||||||
# Set root based on whether local dataset path is provided
|
# Set root based on whether local dataset path is provided
|
||||||
use_local_dataset = False
|
use_local_dataset = False
|
||||||
@@ -488,7 +490,7 @@ def convert_dataset(
|
|||||||
if root.exists():
|
if root.exists():
|
||||||
validate_local_dataset_version(root)
|
validate_local_dataset_version(root)
|
||||||
use_local_dataset = True
|
use_local_dataset = True
|
||||||
print(f"Using local dataset at {root}")
|
logger.info(f"Using local dataset at {root}")
|
||||||
|
|
||||||
old_root = root.parent / f"{root.name}_old"
|
old_root = root.parent / f"{root.name}_old"
|
||||||
new_root = root.parent / f"{root.name}_v30"
|
new_root = root.parent / f"{root.name}_v30"
|
||||||
@@ -523,7 +525,7 @@ def convert_dataset(
|
|||||||
try:
|
try:
|
||||||
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
|
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
|
||||||
except (HTTPError, RevisionNotFoundError) as e:
|
except (HTTPError, RevisionNotFoundError) as e:
|
||||||
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
|
logger.warning(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
|
||||||
pass
|
pass
|
||||||
hub_api.delete_files(
|
hub_api.delete_files(
|
||||||
delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
|
delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
|
||||||
|
|||||||
@@ -154,14 +154,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
|||||||
repo_id = cfg.new_repo_id or cfg.repo_id
|
repo_id = cfg.new_repo_id or cfg.repo_id
|
||||||
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
|
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
|
||||||
api = HfApi()
|
api = HfApi()
|
||||||
print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True)
|
logger.info(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...")
|
||||||
api.create_repo(
|
api.create_repo(
|
||||||
repo_id=repo_id,
|
repo_id=repo_id,
|
||||||
repo_type="dataset",
|
repo_type="dataset",
|
||||||
private=cfg.push_private,
|
private=cfg.push_private,
|
||||||
exist_ok=True,
|
exist_ok=True,
|
||||||
)
|
)
|
||||||
print(f"[lerobot-annotate] uploading {root} -> {repo_id}...", flush=True)
|
logger.info(f"[lerobot-annotate] uploading {root} -> {repo_id}...")
|
||||||
commit_info = api.upload_folder(
|
commit_info = api.upload_folder(
|
||||||
folder_path=str(root),
|
folder_path=str(root),
|
||||||
repo_id=repo_id,
|
repo_id=repo_id,
|
||||||
@@ -172,7 +172,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
|||||||
# at the source dataset; a fresh card is generated below instead.
|
# at the source dataset; a fresh card is generated below instead.
|
||||||
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
|
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
|
||||||
)
|
)
|
||||||
print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
|
logger.info(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}")
|
||||||
|
|
||||||
dataset_info = load_info(root)
|
dataset_info = load_info(root)
|
||||||
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
|
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
|
||||||
@@ -200,14 +200,13 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
|||||||
with suppress(RevisionNotFoundError):
|
with suppress(RevisionNotFoundError):
|
||||||
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
|
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
|
||||||
api.create_tag(**tag_kwargs)
|
api.create_tag(**tag_kwargs)
|
||||||
print(f"[lerobot-annotate] tagged {repo_id} as {version_tag}", flush=True)
|
logger.info(f"[lerobot-annotate] tagged {repo_id} as {version_tag}")
|
||||||
except Exception as exc: # noqa: BLE001
|
except Exception as exc: # noqa: BLE001
|
||||||
print(
|
logger.warning(
|
||||||
f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
|
f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
|
||||||
"Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
|
"Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
|
||||||
"Run: from huggingface_hub import HfApi; "
|
"Run: from huggingface_hub import HfApi; "
|
||||||
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)",
|
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)"
|
||||||
flush=True,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -89,6 +89,8 @@ from lerobot.datasets import LeRobotDataset
|
|||||||
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
|
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
|
||||||
from lerobot.utils.utils import init_logging
|
from lerobot.utils.utils import init_logging
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
DEFAULT_FOXGLOVE_PORT = 8765
|
DEFAULT_FOXGLOVE_PORT = 8765
|
||||||
DEFAULT_RERUN_PORT = 9090
|
DEFAULT_RERUN_PORT = 9090
|
||||||
|
|
||||||
@@ -299,7 +301,7 @@ def visualize_dataset(
|
|||||||
while True:
|
while True:
|
||||||
time.sleep(1)
|
time.sleep(1)
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
print("Ctrl-C received. Exiting.")
|
logger.info("Ctrl-C received. Exiting.")
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
|
|||||||
@@ -62,7 +62,7 @@ from dataclasses import asdict
|
|||||||
from functools import partial
|
from functools import partial
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from pprint import pformat
|
from pprint import pformat
|
||||||
from typing import Any, TypedDict
|
from typing import TYPE_CHECKING, Any, TypedDict
|
||||||
|
|
||||||
import einops
|
import einops
|
||||||
import gymnasium as gym
|
import gymnasium as gym
|
||||||
@@ -87,7 +87,7 @@ from lerobot.processor import PolicyProcessorPipeline
|
|||||||
from lerobot.types import PolicyAction
|
from lerobot.types import PolicyAction
|
||||||
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
|
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
|
||||||
from lerobot.utils.device_utils import get_safe_torch_device
|
from lerobot.utils.device_utils import get_safe_torch_device
|
||||||
from lerobot.utils.import_utils import register_third_party_plugins
|
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
|
||||||
from lerobot.utils.io_utils import write_video
|
from lerobot.utils.io_utils import write_video
|
||||||
from lerobot.utils.random_utils import set_seed
|
from lerobot.utils.random_utils import set_seed
|
||||||
from lerobot.utils.utils import (
|
from lerobot.utils.utils import (
|
||||||
@@ -95,6 +95,14 @@ from lerobot.utils.utils import (
|
|||||||
inside_slurm,
|
inside_slurm,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import PeftModel
|
||||||
|
else:
|
||||||
|
PeftModel = None
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
def _env_features_to_dataset_features(env_features: dict) -> dict:
|
def _env_features_to_dataset_features(env_features: dict) -> dict:
|
||||||
"""Convert EnvConfig.features to the dict format expected by LeRobotDataset.create()."""
|
"""Convert EnvConfig.features to the dict format expected by LeRobotDataset.create()."""
|
||||||
@@ -444,13 +452,11 @@ def eval_policy(
|
|||||||
exc = ValueError(
|
exc = ValueError(
|
||||||
f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
|
f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
|
||||||
)
|
)
|
||||||
try:
|
if not _peft_available:
|
||||||
from peft import PeftModel
|
raise exc
|
||||||
|
require_package("peft", extra="peft")
|
||||||
if not isinstance(policy, PeftModel):
|
if not isinstance(policy, PeftModel):
|
||||||
raise exc
|
raise exc
|
||||||
except ImportError:
|
|
||||||
raise exc from None
|
|
||||||
|
|
||||||
start = time.time()
|
start = time.time()
|
||||||
# Preserve the mode for direct callers. eval_policy_all scopes the mode
|
# Preserve the mode for direct callers. eval_policy_all scopes the mode
|
||||||
@@ -558,7 +564,7 @@ def eval_policy(
|
|||||||
if seeds:
|
if seeds:
|
||||||
all_seeds.extend(seeds)
|
all_seeds.extend(seeds)
|
||||||
else:
|
else:
|
||||||
all_seeds.append(None)
|
all_seeds.extend([None] * env.num_envs)
|
||||||
|
|
||||||
# FIXME: episode_data is either None or it doesn't exist
|
# FIXME: episode_data is either None or it doesn't exist
|
||||||
if return_episode_data:
|
if return_episode_data:
|
||||||
@@ -796,13 +802,13 @@ def eval_main(cfg: EvalPipelineConfig):
|
|||||||
recording_repo_id=cfg.eval.recording_repo_id,
|
recording_repo_id=cfg.eval.recording_repo_id,
|
||||||
recording_private=cfg.eval.recording_private,
|
recording_private=cfg.eval.recording_private,
|
||||||
)
|
)
|
||||||
print("Overall Aggregated Metrics:")
|
logger.info("Overall Aggregated Metrics:")
|
||||||
print(info["overall"])
|
logger.info(info["overall"])
|
||||||
|
|
||||||
# Print per-suite stats
|
# Print per-suite stats
|
||||||
for task_group, task_group_info in info.items():
|
for task_group, task_group_info in info.items():
|
||||||
print(f"\nAggregated Metrics for {task_group}:")
|
logger.info(f"\nAggregated Metrics for {task_group}:")
|
||||||
print(task_group_info)
|
logger.info(task_group_info)
|
||||||
# Close all vec envs
|
# Close all vec envs
|
||||||
close_envs(envs)
|
close_envs(envs)
|
||||||
|
|
||||||
|
|||||||
@@ -28,7 +28,6 @@ lerobot-find-cameras
|
|||||||
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
|
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import concurrent.futures
|
|
||||||
import logging
|
import logging
|
||||||
import time
|
import time
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -40,6 +39,7 @@ from PIL import Image
|
|||||||
from lerobot.cameras import ColorMode
|
from lerobot.cameras import ColorMode
|
||||||
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
|
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
|
||||||
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
|
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
|
||||||
|
from lerobot.utils.utils import init_logging
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -132,7 +132,7 @@ def save_image(
|
|||||||
camera_identifier: str | int,
|
camera_identifier: str | int,
|
||||||
images_dir: Path,
|
images_dir: Path,
|
||||||
camera_type: str,
|
camera_type: str,
|
||||||
):
|
) -> None:
|
||||||
"""
|
"""
|
||||||
Saves a single image to disk using Pillow. Handles color conversion if necessary.
|
Saves a single image to disk using Pillow. Handles color conversion if necessary.
|
||||||
"""
|
"""
|
||||||
@@ -151,7 +151,7 @@ def save_image(
|
|||||||
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
|
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
|
||||||
|
|
||||||
|
|
||||||
def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
|
def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None:
|
||||||
"""Create and connect to a camera instance based on metadata."""
|
"""Create and connect to a camera instance based on metadata."""
|
||||||
cam_type = cam_meta.get("type")
|
cam_type = cam_meta.get("type")
|
||||||
cam_id = cam_meta.get("id")
|
cam_id = cam_meta.get("id")
|
||||||
@@ -164,12 +164,14 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
|
|||||||
cv_config = OpenCVCameraConfig(
|
cv_config = OpenCVCameraConfig(
|
||||||
index_or_path=cam_id,
|
index_or_path=cam_id,
|
||||||
color_mode=ColorMode.RGB,
|
color_mode=ColorMode.RGB,
|
||||||
|
warmup_s=warmup_s,
|
||||||
)
|
)
|
||||||
instance = OpenCVCamera(cv_config)
|
instance = OpenCVCamera(cv_config)
|
||||||
elif cam_type == "RealSense":
|
elif cam_type == "RealSense":
|
||||||
rs_config = RealSenseCameraConfig(
|
rs_config = RealSenseCameraConfig(
|
||||||
serial_number_or_name=cam_id,
|
serial_number_or_name=cam_id,
|
||||||
color_mode=ColorMode.RGB,
|
color_mode=ColorMode.RGB,
|
||||||
|
warmup_s=warmup_s,
|
||||||
)
|
)
|
||||||
instance = RealSenseCamera(rs_config)
|
instance = RealSenseCamera(rs_config)
|
||||||
else:
|
else:
|
||||||
@@ -187,9 +189,7 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
def process_camera_image(
|
def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None:
|
||||||
cam_dict: dict[str, Any], output_dir: Path, current_time: float
|
|
||||||
) -> concurrent.futures.Future | None:
|
|
||||||
"""Capture and process an image from a single camera."""
|
"""Capture and process an image from a single camera."""
|
||||||
cam = cam_dict["instance"]
|
cam = cam_dict["instance"]
|
||||||
meta = cam_dict["meta"]
|
meta = cam_dict["meta"]
|
||||||
@@ -199,7 +199,7 @@ def process_camera_image(
|
|||||||
try:
|
try:
|
||||||
image_data = cam.read()
|
image_data = cam.read()
|
||||||
|
|
||||||
return save_image(
|
save_image(
|
||||||
image_data,
|
image_data,
|
||||||
cam_id_str,
|
cam_id_str,
|
||||||
output_dir,
|
output_dir,
|
||||||
@@ -214,21 +214,21 @@ def process_camera_image(
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
def cleanup_cameras(cameras_to_use: list[dict[str, Any]]):
|
def cleanup_camera(cam_dict: dict[str, Any]) -> None:
|
||||||
"""Disconnect all cameras."""
|
"""Disconnect all cameras."""
|
||||||
logger.info(f"Disconnecting {len(cameras_to_use)} cameras...")
|
logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...")
|
||||||
for cam_dict in cameras_to_use:
|
try:
|
||||||
try:
|
if cam_dict["instance"] and cam_dict["instance"].is_connected:
|
||||||
if cam_dict["instance"] and cam_dict["instance"].is_connected:
|
cam_dict["instance"].disconnect()
|
||||||
cam_dict["instance"].disconnect()
|
except Exception as e:
|
||||||
except Exception as e:
|
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
|
||||||
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
|
|
||||||
|
|
||||||
|
|
||||||
def save_images_from_all_cameras(
|
def save_images_from_all_cameras(
|
||||||
output_dir: Path,
|
output_dir: Path,
|
||||||
record_time_s: float = 2.0,
|
record_time_s: float = 2.0,
|
||||||
camera_type: str | None = None,
|
camera_type: str | None = None,
|
||||||
|
warmup_s: int = 1,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Connects to detected cameras (optionally filtered by type) and saves images from each.
|
Connects to detected cameras (optionally filtered by type) and saves images from each.
|
||||||
@@ -239,6 +239,7 @@ def save_images_from_all_cameras(
|
|||||||
record_time_s: Duration in seconds to record images.
|
record_time_s: Duration in seconds to record images.
|
||||||
camera_type: Optional string to filter cameras ("realsense" or "opencv").
|
camera_type: Optional string to filter cameras ("realsense" or "opencv").
|
||||||
If None, uses all detected cameras.
|
If None, uses all detected cameras.
|
||||||
|
warmup_s: Duration in seconds to warmup camera before recording images.
|
||||||
"""
|
"""
|
||||||
output_dir.mkdir(parents=True, exist_ok=True)
|
output_dir.mkdir(parents=True, exist_ok=True)
|
||||||
logger.info(f"Saving images to {output_dir}")
|
logger.info(f"Saving images to {output_dir}")
|
||||||
@@ -248,47 +249,32 @@ def save_images_from_all_cameras(
|
|||||||
logger.warning("No cameras detected matching the criteria. Cannot save images.")
|
logger.warning("No cameras detected matching the criteria. Cannot save images.")
|
||||||
return
|
return
|
||||||
|
|
||||||
cameras_to_use = []
|
logger.info(
|
||||||
for cam_meta in all_camera_metadata:
|
f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras."
|
||||||
camera_instance = create_camera_instance(cam_meta)
|
)
|
||||||
if camera_instance:
|
|
||||||
cameras_to_use.append(camera_instance)
|
|
||||||
|
|
||||||
if not cameras_to_use:
|
try:
|
||||||
logger.warning("No cameras could be connected. Aborting image save.")
|
for cam_meta in all_camera_metadata:
|
||||||
return
|
cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s)
|
||||||
|
if cam_dict is None:
|
||||||
logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
|
continue
|
||||||
start_time = time.perf_counter()
|
start_time = time.perf_counter()
|
||||||
|
|
||||||
with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
|
|
||||||
try:
|
|
||||||
while time.perf_counter() - start_time < record_time_s:
|
while time.perf_counter() - start_time < record_time_s:
|
||||||
futures = []
|
|
||||||
current_capture_time = time.perf_counter()
|
current_capture_time = time.perf_counter()
|
||||||
|
process_camera_image(cam_dict, output_dir, current_capture_time)
|
||||||
for cam_dict in cameras_to_use:
|
cleanup_camera(cam_dict)
|
||||||
future = process_camera_image(cam_dict, output_dir, current_capture_time)
|
except KeyboardInterrupt:
|
||||||
if future:
|
logger.info("Capture interrupted by user.")
|
||||||
futures.append(future)
|
finally:
|
||||||
|
print(f"Image capture finished. Images saved to {output_dir}")
|
||||||
if futures:
|
|
||||||
concurrent.futures.wait(futures)
|
|
||||||
|
|
||||||
except KeyboardInterrupt:
|
|
||||||
logger.info("Capture interrupted by user.")
|
|
||||||
finally:
|
|
||||||
print("\nFinalizing image saving...")
|
|
||||||
executor.shutdown(wait=True)
|
|
||||||
cleanup_cameras(cameras_to_use)
|
|
||||||
print(f"Image capture finished. Images saved to {output_dir}")
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
|
init_logging()
|
||||||
|
|
||||||
parser = argparse.ArgumentParser(
|
parser = argparse.ArgumentParser(
|
||||||
description="Unified camera utility script for listing cameras and capturing images."
|
description="Unified camera utility script for listing cameras and capturing images."
|
||||||
)
|
)
|
||||||
|
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"camera_type",
|
"camera_type",
|
||||||
type=str,
|
type=str,
|
||||||
@@ -306,8 +292,14 @@ def main():
|
|||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--record-time-s",
|
"--record-time-s",
|
||||||
type=float,
|
type=float,
|
||||||
default=6.0,
|
default=2.0,
|
||||||
help="Time duration to attempt capturing frames. Default: 6 seconds.",
|
help="Time duration to attempt capturing frames. Default: 2 seconds.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--warmup-s",
|
||||||
|
type=int,
|
||||||
|
default=1,
|
||||||
|
help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
|
||||||
)
|
)
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
save_images_from_all_cameras(**vars(args))
|
save_images_from_all_cameras(**vars(args))
|
||||||
|
|||||||
@@ -165,6 +165,7 @@ from lerobot.robots import ( # noqa: F401
|
|||||||
earthrover_mini_plus,
|
earthrover_mini_plus,
|
||||||
hope_jr,
|
hope_jr,
|
||||||
koch_follower,
|
koch_follower,
|
||||||
|
lekiwi,
|
||||||
omx_follower,
|
omx_follower,
|
||||||
openarm_follower,
|
openarm_follower,
|
||||||
reachy2,
|
reachy2,
|
||||||
|
|||||||
@@ -22,7 +22,8 @@ import dataclasses
|
|||||||
import logging
|
import logging
|
||||||
import sys
|
import sys
|
||||||
import time
|
import time
|
||||||
from contextlib import nullcontext
|
from collections.abc import Iterator
|
||||||
|
from contextlib import contextmanager, nullcontext
|
||||||
from pprint import pformat
|
from pprint import pformat
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
@@ -57,7 +58,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
|
|||||||
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
|
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
|
||||||
from lerobot.rewards import make_reward_pre_post_processors
|
from lerobot.rewards import make_reward_pre_post_processors
|
||||||
from lerobot.utils.collate import lerobot_collate_fn
|
from lerobot.utils.collate import lerobot_collate_fn
|
||||||
from lerobot.utils.import_utils import register_third_party_plugins
|
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
|
||||||
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
|
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
|
||||||
from lerobot.utils.random_utils import set_seed
|
from lerobot.utils.random_utils import set_seed
|
||||||
from lerobot.utils.utils import (
|
from lerobot.utils.utils import (
|
||||||
@@ -68,9 +69,28 @@ from lerobot.utils.utils import (
|
|||||||
inside_slurm,
|
inside_slurm,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import PeftModel
|
||||||
|
else:
|
||||||
|
PeftModel = None
|
||||||
|
|
||||||
from .lerobot_eval import eval_policy_all
|
from .lerobot_eval import eval_policy_all
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
|
||||||
|
"""Create evaluation environments for one run and always dispose of them."""
|
||||||
|
envs = make_env(
|
||||||
|
cfg.env,
|
||||||
|
n_envs=cfg.eval.batch_size,
|
||||||
|
use_async_envs=cfg.eval.use_async_envs,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
yield envs
|
||||||
|
finally:
|
||||||
|
close_envs(envs)
|
||||||
|
|
||||||
|
|
||||||
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
|
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
|
||||||
"""Return worker-only DataLoader options, disabling them for single-process loading."""
|
"""Return worker-only DataLoader options, disabling them for single-process loading."""
|
||||||
workers_enabled = cfg.num_workers > 0
|
workers_enabled = cfg.num_workers > 0
|
||||||
@@ -207,8 +227,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
if cfg.job.is_remote:
|
if cfg.job.is_remote:
|
||||||
return submit_to_hf(cfg)
|
return submit_to_hf(cfg)
|
||||||
|
|
||||||
from lerobot.utils.import_utils import require_package
|
|
||||||
|
|
||||||
require_package("accelerate", extra="training")
|
require_package("accelerate", extra="training")
|
||||||
from accelerate import Accelerator
|
from accelerate import Accelerator
|
||||||
from accelerate.utils import DistributedDataParallelKwargs, DistributedType
|
from accelerate.utils import DistributedDataParallelKwargs, DistributedType
|
||||||
@@ -277,14 +295,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
if not is_main_process:
|
if not is_main_process:
|
||||||
dataset, eval_dataset = make_train_eval_datasets(cfg)
|
dataset, eval_dataset = make_train_eval_datasets(cfg)
|
||||||
|
|
||||||
# Create environment used for evaluating checkpoints during training on simulation data.
|
|
||||||
# On real-world data, no need to create an environment as evaluations are done outside train.py,
|
|
||||||
# using the eval.py instead, with gym_dora environment and dora-rs.
|
|
||||||
eval_env = None
|
|
||||||
if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process:
|
|
||||||
logging.info("Creating env")
|
|
||||||
eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
|
|
||||||
|
|
||||||
if cfg.is_reward_model_training:
|
if cfg.is_reward_model_training:
|
||||||
if is_main_process:
|
if is_main_process:
|
||||||
logging.info("Creating reward model")
|
logging.info("Creating reward model")
|
||||||
@@ -312,7 +322,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
if cfg.peft is not None:
|
if cfg.peft is not None:
|
||||||
if cfg.is_reward_model_training:
|
if cfg.is_reward_model_training:
|
||||||
raise ValueError("PEFT is only supported for policy training. ")
|
raise ValueError("PEFT is only supported for policy training. ")
|
||||||
from peft import PeftModel
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
if isinstance(policy, PeftModel):
|
if isinstance(policy, PeftModel):
|
||||||
logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
|
logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
|
||||||
@@ -692,7 +702,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
if is_main_process:
|
if is_main_process:
|
||||||
step_id = get_step_identifier(step, cfg.steps)
|
step_id = get_step_identifier(step, cfg.steps)
|
||||||
logging.info(f"Eval policy at step {step}")
|
logging.info(f"Eval policy at step {step}")
|
||||||
with torch.no_grad(), accelerator.autocast():
|
with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
|
||||||
eval_info = eval_policy_all(
|
eval_info = eval_policy_all(
|
||||||
envs=eval_env, # dict[suite][task_id] -> vec_env
|
envs=eval_env, # dict[suite][task_id] -> vec_env
|
||||||
policy=accelerator.unwrap_model(policy),
|
policy=accelerator.unwrap_model(policy),
|
||||||
@@ -740,9 +750,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
if is_main_process:
|
if is_main_process:
|
||||||
progbar.close()
|
progbar.close()
|
||||||
|
|
||||||
if eval_env:
|
|
||||||
close_envs(eval_env)
|
|
||||||
|
|
||||||
is_fsdp = accelerator.distributed_type == DistributedType.FSDP
|
is_fsdp = accelerator.distributed_type == DistributedType.FSDP
|
||||||
model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
|
model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
|
||||||
if is_main_process:
|
if is_main_process:
|
||||||
|
|||||||
@@ -45,6 +45,7 @@ lerobot-train-tokenizer \
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
import logging
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
@@ -63,6 +64,9 @@ else:
|
|||||||
from lerobot.configs import NormalizationMode, parser
|
from lerobot.configs import NormalizationMode, parser
|
||||||
from lerobot.datasets import LeRobotDataset
|
from lerobot.datasets import LeRobotDataset
|
||||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||||
|
from lerobot.utils.utils import init_logging
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
@@ -274,11 +278,8 @@ def process_episode(args):
|
|||||||
|
|
||||||
return action_chunks
|
return action_chunks
|
||||||
|
|
||||||
except Exception as e:
|
except Exception:
|
||||||
print(f"Error processing episode {ep_idx}: {e}")
|
logger.exception("Error processing episode %s", ep_idx)
|
||||||
import traceback
|
|
||||||
|
|
||||||
traceback.print_exc()
|
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
@@ -300,10 +301,10 @@ def train_fast_tokenizer(
|
|||||||
Returns:
|
Returns:
|
||||||
Trained FAST tokenizer
|
Trained FAST tokenizer
|
||||||
"""
|
"""
|
||||||
print(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
|
logger.info(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
|
||||||
print(f"Action chunk shape: {action_chunks.shape}")
|
logger.info(f"Action chunk shape: {action_chunks.shape}")
|
||||||
print(f"Vocab size: {vocab_size}")
|
logger.info(f"Vocab size: {vocab_size}")
|
||||||
print(f"DCT scale: {scale}")
|
logger.info(f"DCT scale: {scale}")
|
||||||
|
|
||||||
# download the tokenizer source code (not pretrained weights)
|
# download the tokenizer source code (not pretrained weights)
|
||||||
# we'll train a new tokenizer on our own data
|
# we'll train a new tokenizer on our own data
|
||||||
@@ -314,7 +315,7 @@ def train_fast_tokenizer(
|
|||||||
|
|
||||||
# train the new tokenizer on our action data using .fit()
|
# train the new tokenizer on our action data using .fit()
|
||||||
# this trains the BPE tokenizer on DCT coefficients
|
# this trains the BPE tokenizer on DCT coefficients
|
||||||
print("Training new tokenizer (this may take a few minutes)...")
|
logger.info("Training new tokenizer (this may take a few minutes)...")
|
||||||
tokenizer = base_tokenizer.fit(
|
tokenizer = base_tokenizer.fit(
|
||||||
action_data_list,
|
action_data_list,
|
||||||
scale=scale,
|
scale=scale,
|
||||||
@@ -322,21 +323,21 @@ def train_fast_tokenizer(
|
|||||||
time_horizon=action_chunks.shape[1], # action_horizon
|
time_horizon=action_chunks.shape[1], # action_horizon
|
||||||
action_dim=action_chunks.shape[2], # encoded dimensions
|
action_dim=action_chunks.shape[2], # encoded dimensions
|
||||||
)
|
)
|
||||||
print("✓ Tokenizer training complete!")
|
logger.info("✓ Tokenizer training complete!")
|
||||||
|
|
||||||
# validate it works
|
# validate it works
|
||||||
sample_chunk = action_chunks[0]
|
sample_chunk = action_chunks[0]
|
||||||
encoded = tokenizer(sample_chunk[None])[0]
|
encoded = tokenizer(sample_chunk[None])[0]
|
||||||
if isinstance(encoded, list):
|
if isinstance(encoded, list):
|
||||||
encoded = np.array(encoded)
|
encoded = np.array(encoded)
|
||||||
print(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
|
logger.info(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
|
||||||
|
|
||||||
return tokenizer
|
return tokenizer
|
||||||
|
|
||||||
|
|
||||||
def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
|
def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
|
||||||
"""Compute compression statistics."""
|
"""Compute compression statistics."""
|
||||||
print("\nComputing compression statistics...")
|
logger.info("\nComputing compression statistics...")
|
||||||
|
|
||||||
# sample for stats (use max 1000 chunks for speed)
|
# sample for stats (use max 1000 chunks for speed)
|
||||||
sample_size = min(1000, len(action_chunks))
|
sample_size = min(1000, len(action_chunks))
|
||||||
@@ -366,12 +367,12 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
|
|||||||
"max_token_length": float(np.max(token_lengths)),
|
"max_token_length": float(np.max(token_lengths)),
|
||||||
}
|
}
|
||||||
|
|
||||||
print("Compression Statistics:")
|
logger.info("Compression Statistics:")
|
||||||
print(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
|
logger.info(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
|
||||||
print(f" Mean token length: {stats['mean_token_length']:.1f}")
|
logger.info(f" Mean token length: {stats['mean_token_length']:.1f}")
|
||||||
print(f" P99 token length: {stats['p99_token_length']:.0f}")
|
logger.info(f" P99 token length: {stats['p99_token_length']:.0f}")
|
||||||
print(f" Min token length: {stats['min_token_length']:.0f}")
|
logger.info(f" Min token length: {stats['min_token_length']:.0f}")
|
||||||
print(f" Max token length: {stats['max_token_length']:.0f}")
|
logger.info(f" Max token length: {stats['max_token_length']:.0f}")
|
||||||
|
|
||||||
return stats
|
return stats
|
||||||
|
|
||||||
@@ -385,9 +386,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
cfg: TokenizerTrainingConfig dataclass with all configuration parameters
|
cfg: TokenizerTrainingConfig dataclass with all configuration parameters
|
||||||
"""
|
"""
|
||||||
# load dataset
|
# load dataset
|
||||||
print(f"Loading dataset: {cfg.repo_id}")
|
logger.info(f"Loading dataset: {cfg.repo_id}")
|
||||||
dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
|
||||||
print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
|
logger.info(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
|
||||||
|
|
||||||
# parse normalization mode
|
# parse normalization mode
|
||||||
try:
|
try:
|
||||||
@@ -397,7 +398,7 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
f"Invalid normalization_mode: {cfg.normalization_mode}. "
|
f"Invalid normalization_mode: {cfg.normalization_mode}. "
|
||||||
f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
|
f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
|
||||||
) from err
|
) from err
|
||||||
print(f"Normalization mode: {norm_mode.value}")
|
logger.info(f"Normalization mode: {norm_mode.value}")
|
||||||
|
|
||||||
# parse encoded dimensions
|
# parse encoded dimensions
|
||||||
encoded_dim_ranges = []
|
encoded_dim_ranges = []
|
||||||
@@ -406,38 +407,38 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
encoded_dim_ranges.append((start, end))
|
encoded_dim_ranges.append((start, end))
|
||||||
|
|
||||||
total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
|
total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
|
||||||
print(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
|
logger.info(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
|
||||||
|
|
||||||
# parse relative dimensions
|
# parse relative dimensions
|
||||||
relative_dim_list = None
|
relative_dim_list = None
|
||||||
if cfg.relative_dims is not None and cfg.relative_dims.strip():
|
if cfg.relative_dims is not None and cfg.relative_dims.strip():
|
||||||
relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
|
relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
|
||||||
print(f"Relative dimensions: {relative_dim_list}")
|
logger.info(f"Relative dimensions: {relative_dim_list}")
|
||||||
else:
|
else:
|
||||||
print("No relative dimensions specified")
|
logger.info("No relative dimensions specified")
|
||||||
|
|
||||||
print(f"Use relative transform: {cfg.use_relative_transform}")
|
logger.info(f"Use relative transform: {cfg.use_relative_transform}")
|
||||||
if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
|
if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
|
||||||
print(
|
logger.warning(
|
||||||
"Warning: use_relative_transform=True but no relative_dims specified. "
|
"Warning: use_relative_transform=True but no relative_dims specified. "
|
||||||
"No relative transform will be applied."
|
"No relative transform will be applied."
|
||||||
)
|
)
|
||||||
|
|
||||||
print(f"Action horizon: {cfg.action_horizon}")
|
logger.info(f"Action horizon: {cfg.action_horizon}")
|
||||||
print(f"State key: {cfg.state_key}")
|
logger.info(f"State key: {cfg.state_key}")
|
||||||
|
|
||||||
# determine episodes to process
|
# determine episodes to process
|
||||||
num_episodes = dataset.num_episodes
|
num_episodes = dataset.num_episodes
|
||||||
if cfg.max_episodes is not None:
|
if cfg.max_episodes is not None:
|
||||||
num_episodes = min(cfg.max_episodes, num_episodes)
|
num_episodes = min(cfg.max_episodes, num_episodes)
|
||||||
|
|
||||||
print(f"Processing {num_episodes} episodes...")
|
logger.info(f"Processing {num_episodes} episodes...")
|
||||||
|
|
||||||
# process episodes sequentially (to avoid pickling issues with dataset)
|
# process episodes sequentially (to avoid pickling issues with dataset)
|
||||||
all_chunks = []
|
all_chunks = []
|
||||||
for ep_idx in range(num_episodes):
|
for ep_idx in range(num_episodes):
|
||||||
if ep_idx % 10 == 0:
|
if ep_idx % 10 == 0:
|
||||||
print(f" Processing episode {ep_idx}/{num_episodes}...")
|
logger.info(f" Processing episode {ep_idx}/{num_episodes}...")
|
||||||
|
|
||||||
chunks = process_episode(
|
chunks = process_episode(
|
||||||
(
|
(
|
||||||
@@ -455,19 +456,19 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
|
|
||||||
# concatenate all chunks
|
# concatenate all chunks
|
||||||
all_chunks = np.concatenate(all_chunks, axis=0)
|
all_chunks = np.concatenate(all_chunks, axis=0)
|
||||||
print(f"Collected {len(all_chunks)} action chunks")
|
logger.info(f"Collected {len(all_chunks)} action chunks")
|
||||||
|
|
||||||
# extract only encoded dimensions FIRST (before normalization)
|
# extract only encoded dimensions FIRST (before normalization)
|
||||||
encoded_chunks = []
|
encoded_chunks = []
|
||||||
for start, end in encoded_dim_ranges:
|
for start, end in encoded_dim_ranges:
|
||||||
encoded_chunks.append(all_chunks[:, :, start:end])
|
encoded_chunks.append(all_chunks[:, :, start:end])
|
||||||
encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded]
|
encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded]
|
||||||
print(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
|
logger.info(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
|
||||||
|
|
||||||
# apply normalization to encoded dimensions
|
# apply normalization to encoded dimensions
|
||||||
print("\nBefore normalization - overall stats:")
|
logger.info("\nBefore normalization - overall stats:")
|
||||||
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
|
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
|
||||||
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
|
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
|
||||||
|
|
||||||
# get normalization stats from dataset
|
# get normalization stats from dataset
|
||||||
norm_stats = dataset.meta.stats
|
norm_stats = dataset.meta.stats
|
||||||
@@ -489,9 +490,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
encoded_stats[stat_name] = stat_array[encoded_dim_indices]
|
encoded_stats[stat_name] = stat_array[encoded_dim_indices]
|
||||||
|
|
||||||
if encoded_stats:
|
if encoded_stats:
|
||||||
print(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
|
logger.info(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
|
||||||
for stat_name, stat_values in encoded_stats.items():
|
for stat_name, stat_values in encoded_stats.items():
|
||||||
print(
|
logger.info(
|
||||||
f" {stat_name}: shape={stat_values.shape}, "
|
f" {stat_name}: shape={stat_values.shape}, "
|
||||||
f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
|
f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
|
||||||
)
|
)
|
||||||
@@ -499,27 +500,27 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
# apply normalization based on mode
|
# apply normalization based on mode
|
||||||
try:
|
try:
|
||||||
encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
|
encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
|
||||||
print(f"\nApplied {norm_mode.value} normalization")
|
logger.info(f"\nApplied {norm_mode.value} normalization")
|
||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
print(f"Warning: {e}. Using raw actions without normalization.")
|
logger.warning(f"Warning: {e}. Using raw actions without normalization.")
|
||||||
|
|
||||||
print("\nAfter normalization - overall stats:")
|
logger.info("\nAfter normalization - overall stats:")
|
||||||
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
|
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
|
||||||
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
|
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
|
||||||
|
|
||||||
print("\nPer-dimension stats (after normalization):")
|
logger.info("\nPer-dimension stats (after normalization):")
|
||||||
for d in range(encoded_chunks.shape[-1]):
|
for d in range(encoded_chunks.shape[-1]):
|
||||||
dim_data = encoded_chunks[:, :, d]
|
dim_data = encoded_chunks[:, :, d]
|
||||||
print(
|
logger.info(
|
||||||
f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
|
f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
|
||||||
f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
|
f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
print("Warning: Could not extract stats for encoded dimensions, using raw actions")
|
logger.warning("Warning: Could not extract stats for encoded dimensions, using raw actions")
|
||||||
else:
|
else:
|
||||||
print("Warning: No normalization stats found in dataset, using raw actions")
|
logger.warning("Warning: No normalization stats found in dataset, using raw actions")
|
||||||
|
|
||||||
print(f"Encoded chunks shape: {encoded_chunks.shape}")
|
logger.info(f"Encoded chunks shape: {encoded_chunks.shape}")
|
||||||
|
|
||||||
# train FAST tokenizer
|
# train FAST tokenizer
|
||||||
tokenizer = train_fast_tokenizer(
|
tokenizer = train_fast_tokenizer(
|
||||||
@@ -561,8 +562,8 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
with open(output_path / "metadata.json", "w") as f:
|
with open(output_path / "metadata.json", "w") as f:
|
||||||
json.dump(metadata, f, indent=2)
|
json.dump(metadata, f, indent=2)
|
||||||
|
|
||||||
print(f"\nSaved FAST tokenizer to {output_path}")
|
logger.info(f"\nSaved FAST tokenizer to {output_path}")
|
||||||
print(f"Metadata: {json.dumps(metadata, indent=2)}")
|
logger.info(f"Metadata: {json.dumps(metadata, indent=2)}")
|
||||||
|
|
||||||
# push to Hugging Face Hub if requested
|
# push to Hugging Face Hub if requested
|
||||||
if cfg.push_to_hub:
|
if cfg.push_to_hub:
|
||||||
@@ -570,10 +571,10 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
hub_repo_id = cfg.hub_repo_id
|
hub_repo_id = cfg.hub_repo_id
|
||||||
if hub_repo_id is None:
|
if hub_repo_id is None:
|
||||||
hub_repo_id = output_path.name
|
hub_repo_id = output_path.name
|
||||||
print(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
|
logger.info(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
|
||||||
|
|
||||||
print(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
|
logger.info(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
|
||||||
print(f" Private: {cfg.hub_private}")
|
logger.info(f" Private: {cfg.hub_private}")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# use the tokenizer's push_to_hub method
|
# use the tokenizer's push_to_hub method
|
||||||
@@ -593,14 +594,15 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
commit_message="Upload tokenizer metadata",
|
commit_message="Upload tokenizer metadata",
|
||||||
)
|
)
|
||||||
|
|
||||||
print(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
|
logger.info(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"Error pushing to hub: {e}")
|
logger.error(f"Error pushing to hub: {e}")
|
||||||
print(" Make sure you're logged in with `huggingface-cli login`")
|
logger.error(" Make sure you're logged in with `huggingface-cli login`")
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
"""CLI entry point that parses arguments and runs the tokenizer training."""
|
"""CLI entry point that parses arguments and runs the tokenizer training."""
|
||||||
|
init_logging()
|
||||||
train_tokenizer()
|
train_tokenizer()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -171,7 +171,13 @@ class IOSPhone(BasePhone, Teleoperator):
|
|||||||
# HEBI provides orientation in w, x, y, z format.
|
# HEBI provides orientation in w, x, y, z format.
|
||||||
# Scipy's Rotation expects x, y, z, w.
|
# Scipy's Rotation expects x, y, z, w.
|
||||||
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
|
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
|
||||||
rot = Rotation.from_quat(quat_xyzw)
|
# ARKit can emit zero/NaN quaternions before tracking is ready or on a
|
||||||
|
# dropped packet. Rotation.from_quat now rejects those; degrade the same
|
||||||
|
# way as a missing pose so teleop stays alive mid-session.
|
||||||
|
try:
|
||||||
|
rot = Rotation.from_quat(quat_xyzw)
|
||||||
|
except ValueError:
|
||||||
|
return False, None, None, None
|
||||||
pos = ar_pos - rot.apply(self.config.camera_offset)
|
pos = ar_pos - rot.apply(self.config.camera_offset)
|
||||||
return True, pos, rot, pose
|
return True, pos, rot, pose
|
||||||
|
|
||||||
|
|||||||
@@ -29,6 +29,12 @@ class SOLeaderConfig:
|
|||||||
# Whether to use degrees for angles
|
# Whether to use degrees for angles
|
||||||
use_degrees: bool = True
|
use_degrees: bool = True
|
||||||
|
|
||||||
|
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
|
||||||
|
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
|
||||||
|
# at once, which otherwise aborts the teleoperation loop. Retries are immediate (no sleep) and only
|
||||||
|
# happen on failure, so the steady-state read cost is unchanged.
|
||||||
|
num_read_retries: int = 2
|
||||||
|
|
||||||
|
|
||||||
@TeleoperatorConfig.register_subclass("so101_leader")
|
@TeleoperatorConfig.register_subclass("so101_leader")
|
||||||
@TeleoperatorConfig.register_subclass("so100_leader")
|
@TeleoperatorConfig.register_subclass("so100_leader")
|
||||||
|
|||||||
@@ -145,7 +145,7 @@ class SOLeader(Teleoperator):
|
|||||||
@check_if_not_connected
|
@check_if_not_connected
|
||||||
def get_action(self) -> dict[str, float]:
|
def get_action(self) -> dict[str, float]:
|
||||||
start = time.perf_counter()
|
start = time.perf_counter()
|
||||||
action = self.bus.sync_read("Present_Position")
|
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()}
|
action = {f"{motor}.pos": val for motor, val in action.items()}
|
||||||
dt_ms = (time.perf_counter() - start) * 1e3
|
dt_ms = (time.perf_counter() - start) * 1e3
|
||||||
logger.debug(f"{self} read action: {dt_ms:.1f}ms")
|
logger.debug(f"{self} read action: {dt_ms:.1f}ms")
|
||||||
|
|||||||
@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):
|
|||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
transforms: Sequence[Callable],
|
transforms: Sequence[Callable[..., Any]],
|
||||||
p: list[float] | None = None,
|
p: list[float] | None = None,
|
||||||
n_subset: int | None = None,
|
n_subset: int | None = None,
|
||||||
random_order: bool = False,
|
random_order: bool = False,
|
||||||
@@ -50,7 +50,7 @@ class RandomSubsetApply(Transform):
|
|||||||
if not isinstance(transforms, Sequence):
|
if not isinstance(transforms, Sequence):
|
||||||
raise TypeError("Argument transforms should be a sequence of callables")
|
raise TypeError("Argument transforms should be a sequence of callables")
|
||||||
if p is None:
|
if p is None:
|
||||||
p = [1] * len(transforms)
|
p = [1.0] * len(transforms)
|
||||||
elif len(p) != len(transforms):
|
elif len(p) != len(transforms):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
|
f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
|
||||||
@@ -69,7 +69,7 @@ class RandomSubsetApply(Transform):
|
|||||||
self.n_subset = n_subset
|
self.n_subset = n_subset
|
||||||
self.random_order = random_order
|
self.random_order = random_order
|
||||||
|
|
||||||
self.selected_transforms = None
|
self.selected_transforms: list[Callable[..., Any]] = []
|
||||||
|
|
||||||
def forward(self, *inputs: Any) -> Any:
|
def forward(self, *inputs: Any) -> Any:
|
||||||
needs_unpacking = len(inputs) > 1
|
needs_unpacking = len(inputs) > 1
|
||||||
@@ -119,7 +119,7 @@ class SharpnessJitter(Transform):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
self.sharpness = self._check_input(sharpness)
|
self.sharpness = self._check_input(sharpness)
|
||||||
|
|
||||||
def _check_input(self, sharpness):
|
def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]:
|
||||||
if isinstance(sharpness, (int | float)):
|
if isinstance(sharpness, (int | float)):
|
||||||
if sharpness < 0:
|
if sharpness < 0:
|
||||||
raise ValueError("If sharpness is a single number, it must be non negative.")
|
raise ValueError("If sharpness is a single number, it must be non negative.")
|
||||||
@@ -215,7 +215,7 @@ class ImageTransformsConfig:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def make_transform_from_config(cfg: ImageTransformConfig):
|
def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
|
||||||
if cfg.type == "SharpnessJitter":
|
if cfg.type == "SharpnessJitter":
|
||||||
return SharpnessJitter(**cfg.kwargs)
|
return SharpnessJitter(**cfg.kwargs)
|
||||||
|
|
||||||
@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
self._cfg = cfg
|
self._cfg = cfg
|
||||||
|
|
||||||
self.weights = []
|
self.weights: list[float] = []
|
||||||
self.transforms = {}
|
self.transforms: dict[str, Transform] = {}
|
||||||
for tf_name, tf_cfg in cfg.tfs.items():
|
for tf_name, tf_cfg in cfg.tfs.items():
|
||||||
if tf_cfg.weight <= 0.0:
|
if tf_cfg.weight <= 0.0:
|
||||||
continue
|
continue
|
||||||
|
|||||||
@@ -37,16 +37,25 @@ def auto_select_torch_device() -> torch.device:
|
|||||||
|
|
||||||
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
|
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
|
||||||
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
|
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
|
||||||
"""Given a string, return a torch.device with checks on whether the device is available."""
|
"""Given a string, return a torch.device with checks on whether the device is available.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If the requested device family is known but not available on
|
||||||
|
this machine (``AssertionError`` was previously used and is easy to
|
||||||
|
mistake for a programmer bug under ``python -O`` where asserts vanish).
|
||||||
|
"""
|
||||||
try_device = str(try_device)
|
try_device = str(try_device)
|
||||||
if try_device.startswith("cuda"):
|
if try_device.startswith("cuda"):
|
||||||
assert torch.cuda.is_available()
|
if not torch.cuda.is_available():
|
||||||
|
raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
|
||||||
device = torch.device(try_device)
|
device = torch.device(try_device)
|
||||||
elif try_device == "mps":
|
elif try_device == "mps":
|
||||||
assert torch.backends.mps.is_available()
|
if not torch.backends.mps.is_available():
|
||||||
|
raise ValueError("Requested device 'mps' but MPS is not available.")
|
||||||
device = torch.device("mps")
|
device = torch.device("mps")
|
||||||
elif try_device == "xpu":
|
elif try_device == "xpu":
|
||||||
assert torch.xpu.is_available()
|
if not torch.xpu.is_available():
|
||||||
|
raise ValueError("Requested device 'xpu' but XPU is not available.")
|
||||||
device = torch.device("xpu")
|
device = torch.device("xpu")
|
||||||
elif try_device == "cpu":
|
elif try_device == "cpu":
|
||||||
device = torch.device("cpu")
|
device = torch.device("cpu")
|
||||||
|
|||||||
@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
|
|||||||
Returns:
|
Returns:
|
||||||
Any: The data loaded from the JSON file.
|
Any: The data loaded from the JSON file.
|
||||||
"""
|
"""
|
||||||
with open(fpath) as f:
|
with open(fpath, encoding="utf-8") as f:
|
||||||
return json.load(f)
|
return json.load(f)
|
||||||
|
|
||||||
|
|
||||||
def write_json(data: dict, fpath: Path) -> None:
|
def write_json(data: JsonLike, fpath: Path) -> None:
|
||||||
"""Write data to a JSON file.
|
"""Write JSON-serializable data to a file.
|
||||||
|
|
||||||
Creates parent directories if they don't exist.
|
Creates parent directories if they don't exist.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
data (dict): The dictionary to write.
|
data: JSON-serializable data to write.
|
||||||
fpath (Path): The path to the output JSON file.
|
fpath (Path): The path to the output JSON file.
|
||||||
"""
|
"""
|
||||||
fpath.parent.mkdir(exist_ok=True, parents=True)
|
fpath.parent.mkdir(exist_ok=True, parents=True)
|
||||||
with open(fpath, "w") as f:
|
with open(fpath, "w", encoding="utf-8") as f:
|
||||||
json.dump(data, f, indent=4, ensure_ascii=False)
|
json.dump(data, f, indent=4, ensure_ascii=False)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -16,11 +16,39 @@
|
|||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
import multiprocessing
|
||||||
import os
|
import os
|
||||||
import signal
|
import signal
|
||||||
import sys
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_multiprocessing_start_method(start_method: str | None) -> None:
|
||||||
|
"""Set a multiprocessing start method once, or verify the existing method matches.
|
||||||
|
|
||||||
|
Passing ``None`` leaves Python's process-wide default untouched. This is useful
|
||||||
|
when LeRobot is embedded in an application that owns multiprocessing setup.
|
||||||
|
"""
|
||||||
|
if start_method is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
available_methods = multiprocessing.get_all_start_methods()
|
||||||
|
if start_method not in available_methods:
|
||||||
|
raise ValueError(
|
||||||
|
f"Multiprocessing start method must be one of {available_methods} on this platform, "
|
||||||
|
f"got {start_method!r}."
|
||||||
|
)
|
||||||
|
|
||||||
|
current_method = multiprocessing.get_start_method(allow_none=True)
|
||||||
|
if current_method is None:
|
||||||
|
multiprocessing.set_start_method(start_method)
|
||||||
|
elif current_method != start_method:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Multiprocessing start method is already {current_method!r}; cannot change it to "
|
||||||
|
f"{start_method!r}. Set the configured multiprocessing context to null to keep the "
|
||||||
|
"application's existing method, or launch LeRobot in a fresh process."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class ProcessSignalHandler:
|
class ProcessSignalHandler:
|
||||||
"""Utility class to attach graceful shutdown signal handlers.
|
"""Utility class to attach graceful shutdown signal handlers.
|
||||||
|
|
||||||
|
|||||||
@@ -30,6 +30,10 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
|
|||||||
"""
|
"""
|
||||||
if seconds <= 0:
|
if seconds <= 0:
|
||||||
return
|
return
|
||||||
|
if spin_threshold < 0:
|
||||||
|
raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
|
||||||
|
if sleep_margin < 0:
|
||||||
|
raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")
|
||||||
|
|
||||||
system = platform.system()
|
system = platform.system()
|
||||||
# On macOS and Windows the scheduler / sleep granularity can make
|
# On macOS and Windows the scheduler / sleep granularity can make
|
||||||
|
|||||||
@@ -29,10 +29,13 @@ class Rotation:
|
|||||||
def __init__(self, quat: np.ndarray) -> None:
|
def __init__(self, quat: np.ndarray) -> None:
|
||||||
"""Initialize rotation from quaternion [x, y, z, w]."""
|
"""Initialize rotation from quaternion [x, y, z, w]."""
|
||||||
self._quat = np.asarray(quat, dtype=float)
|
self._quat = np.asarray(quat, dtype=float)
|
||||||
# Normalize quaternion
|
if self._quat.shape != (4,):
|
||||||
|
raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
|
||||||
|
# Normalize quaternion. Reject the zero vector — it has no orientation.
|
||||||
norm = np.linalg.norm(self._quat)
|
norm = np.linalg.norm(self._quat)
|
||||||
if norm > 0:
|
if norm <= 0.0 or not np.isfinite(norm):
|
||||||
self._quat = self._quat / norm
|
raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})")
|
||||||
|
self._quat = self._quat / norm
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":
|
def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":
|
||||||
|
|||||||
@@ -14,7 +14,7 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
from typing import TypedDict
|
from typing import NotRequired, TypedDict
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
@@ -28,7 +28,7 @@ class Transition(TypedDict):
|
|||||||
next_state: dict[str, torch.Tensor]
|
next_state: dict[str, torch.Tensor]
|
||||||
done: bool
|
done: bool
|
||||||
truncated: bool
|
truncated: bool
|
||||||
complementary_info: dict[str, torch.Tensor | float | int] | None = None
|
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
|
||||||
|
|
||||||
|
|
||||||
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
|
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
|
||||||
|
|||||||
+16
-12
@@ -24,7 +24,6 @@ import sys
|
|||||||
import time
|
import time
|
||||||
from collections.abc import Iterator
|
from collections.abc import Iterator
|
||||||
from copy import copy, deepcopy
|
from copy import copy, deepcopy
|
||||||
from datetime import datetime
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from statistics import mean
|
from statistics import mean
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
@@ -61,14 +60,16 @@ def init_logging(
|
|||||||
accelerator: Optional Accelerator instance (for multi-GPU detection)
|
accelerator: Optional Accelerator instance (for multi-GPU detection)
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def custom_format(record: logging.LogRecord) -> str:
|
class LeRobotFormatter(logging.Formatter):
|
||||||
dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
def format(self, record: logging.LogRecord) -> str:
|
||||||
fnameline = f"{record.pathname}:{record.lineno}"
|
record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:]
|
||||||
pid_str = f"[PID: {os.getpid()}] " if display_pid else ""
|
record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else ""
|
||||||
return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}"
|
return super().format(record)
|
||||||
|
|
||||||
formatter = logging.Formatter()
|
formatter = LeRobotFormatter(
|
||||||
formatter.format = custom_format
|
"%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s",
|
||||||
|
datefmt="%Y-%m-%d %H:%M:%S",
|
||||||
|
)
|
||||||
|
|
||||||
logger = logging.getLogger()
|
logger = logging.getLogger()
|
||||||
logger.setLevel(logging.NOTSET)
|
logger.setLevel(logging.NOTSET)
|
||||||
@@ -133,10 +134,13 @@ def say(text: str, blocking: bool = False):
|
|||||||
else:
|
else:
|
||||||
raise RuntimeError("Unsupported operating system for text-to-speech.")
|
raise RuntimeError("Unsupported operating system for text-to-speech.")
|
||||||
|
|
||||||
if blocking:
|
try:
|
||||||
subprocess.run(cmd, check=True)
|
if blocking:
|
||||||
else:
|
subprocess.run(cmd, check=True, timeout=5)
|
||||||
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
|
else:
|
||||||
|
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
|
||||||
|
except (FileNotFoundError, subprocess.TimeoutExpired) as e:
|
||||||
|
logging.warning("Text-to-speech command failed: %s | Error: %s", cmd, e)
|
||||||
|
|
||||||
|
|
||||||
def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
|
def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
|
||||||
|
|||||||
@@ -20,7 +20,7 @@
|
|||||||
# ```
|
# ```
|
||||||
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from unittest.mock import patch
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -123,6 +123,73 @@ def test_invalid_width_connect():
|
|||||||
camera.connect(warmup=False)
|
camera.connect(warmup=False)
|
||||||
|
|
||||||
|
|
||||||
|
def test_connect_cleans_up_after_settings_failure_and_allows_retry():
|
||||||
|
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
|
||||||
|
camera = OpenCVCamera(config)
|
||||||
|
opened_captures = []
|
||||||
|
|
||||||
|
def fail_settings():
|
||||||
|
opened_captures.append(camera.videocapture)
|
||||||
|
raise RuntimeError("settings failed")
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(camera, "_configure_capture_settings", side_effect=fail_settings),
|
||||||
|
pytest.raises(RuntimeError, match="settings failed"),
|
||||||
|
):
|
||||||
|
camera.connect(warmup=False)
|
||||||
|
|
||||||
|
assert camera.videocapture is None
|
||||||
|
assert camera.thread is None
|
||||||
|
assert not camera.is_connected
|
||||||
|
assert opened_captures[0] is not None
|
||||||
|
assert not opened_captures[0].isOpened()
|
||||||
|
|
||||||
|
camera.connect(warmup=False)
|
||||||
|
assert camera.is_connected
|
||||||
|
camera.disconnect()
|
||||||
|
|
||||||
|
|
||||||
|
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
|
||||||
|
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=1)
|
||||||
|
camera = OpenCVCamera(config)
|
||||||
|
read_threads = []
|
||||||
|
|
||||||
|
def fail_warmup(*_args, **_kwargs):
|
||||||
|
read_threads.append(camera.thread)
|
||||||
|
raise TimeoutError("no frame")
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(camera, "async_read", side_effect=fail_warmup),
|
||||||
|
pytest.raises(TimeoutError, match="no frame"),
|
||||||
|
):
|
||||||
|
camera.connect()
|
||||||
|
|
||||||
|
assert camera.videocapture is None
|
||||||
|
assert camera.thread is None
|
||||||
|
assert not camera.is_connected
|
||||||
|
assert read_threads[0] is not None
|
||||||
|
assert not read_threads[0].is_alive()
|
||||||
|
|
||||||
|
camera.connect(warmup=False)
|
||||||
|
assert camera.is_connected
|
||||||
|
camera.disconnect()
|
||||||
|
|
||||||
|
|
||||||
|
def test_find_cameras_releases_unopened_handles():
|
||||||
|
module_path = OpenCVCamera.__module__
|
||||||
|
unopened_capture = MagicMock()
|
||||||
|
unopened_capture.isOpened.return_value = False
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(f"{module_path}.platform.system", return_value="Darwin"),
|
||||||
|
patch(f"{module_path}.MAX_OPENCV_INDEX", 1),
|
||||||
|
patch(f"{module_path}.cv2.VideoCapture", return_value=unopened_capture),
|
||||||
|
):
|
||||||
|
assert OpenCVCamera.find_cameras() == []
|
||||||
|
|
||||||
|
unopened_capture.release.assert_called_once_with()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
|
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
|
||||||
def test_read(index_or_path):
|
def test_read(index_or_path):
|
||||||
config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
|
config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
|
||||||
|
|||||||
@@ -20,7 +20,7 @@
|
|||||||
# ```
|
# ```
|
||||||
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from unittest.mock import patch
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pytest
|
import pytest
|
||||||
@@ -30,6 +30,8 @@ from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnected
|
|||||||
|
|
||||||
pytest.importorskip("pyrealsense2")
|
pytest.importorskip("pyrealsense2")
|
||||||
|
|
||||||
|
import pyrealsense2 as rs
|
||||||
|
|
||||||
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
|
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
|
||||||
|
|
||||||
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
|
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
|
||||||
@@ -61,6 +63,17 @@ def test_abc_implementation():
|
|||||||
_ = RealSenseCamera(config)
|
_ = RealSenseCamera(config)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("option", ["exposure", "gain", "white_balance"])
|
||||||
|
def test_manual_color_option_requires_rgb(option):
|
||||||
|
with pytest.raises(ValueError, match="use_rgb=True"):
|
||||||
|
RealSenseCameraConfig(
|
||||||
|
serial_number_or_name="042",
|
||||||
|
use_rgb=False,
|
||||||
|
use_depth=True,
|
||||||
|
**{option: 100},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def test_connect():
|
def test_connect():
|
||||||
config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
|
config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
|
||||||
|
|
||||||
@@ -83,6 +96,27 @@ def test_connect_invalid_camera_path(patch_realsense):
|
|||||||
camera.connect(warmup=False)
|
camera.connect(warmup=False)
|
||||||
|
|
||||||
|
|
||||||
|
def test_connect_cleans_up_when_sensor_configuration_fails():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
pipeline = MagicMock()
|
||||||
|
pipeline.start.return_value = MagicMock()
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("lerobot.cameras.realsense.camera_realsense.rs.pipeline", return_value=pipeline),
|
||||||
|
patch.object(camera, "_configure_rs_pipeline_config"),
|
||||||
|
patch.object(camera, "_configure_capture_settings"),
|
||||||
|
patch.object(camera, "_configure_sensor_options", side_effect=ValueError("invalid exposure")),
|
||||||
|
pytest.raises(ValueError, match="invalid exposure"),
|
||||||
|
):
|
||||||
|
camera.connect(warmup=False)
|
||||||
|
|
||||||
|
pipeline.stop.assert_called_once_with()
|
||||||
|
assert camera.rs_pipeline is None
|
||||||
|
assert camera.rs_profile is None
|
||||||
|
assert not camera.is_connected
|
||||||
|
|
||||||
|
|
||||||
def test_invalid_width_connect():
|
def test_invalid_width_connect():
|
||||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
|
config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
|
||||||
camera = RealSenseCamera(config)
|
camera = RealSenseCamera(config)
|
||||||
@@ -91,6 +125,33 @@ def test_invalid_width_connect():
|
|||||||
camera.connect(warmup=False)
|
camera.connect(warmup=False)
|
||||||
|
|
||||||
|
|
||||||
|
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
read_threads = []
|
||||||
|
|
||||||
|
def fail_warmup(*_args, **_kwargs):
|
||||||
|
read_threads.append(camera.thread)
|
||||||
|
raise TimeoutError("no frame")
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(camera, "async_read", side_effect=fail_warmup),
|
||||||
|
pytest.raises(TimeoutError, match="no frame"),
|
||||||
|
):
|
||||||
|
camera.connect()
|
||||||
|
|
||||||
|
assert camera.rs_pipeline is None
|
||||||
|
assert camera.rs_profile is None
|
||||||
|
assert camera.thread is None
|
||||||
|
assert not camera.is_connected
|
||||||
|
assert read_threads[0] is not None
|
||||||
|
assert not read_threads[0].is_alive()
|
||||||
|
|
||||||
|
camera.connect(warmup=False)
|
||||||
|
assert camera.is_connected
|
||||||
|
camera.disconnect()
|
||||||
|
|
||||||
|
|
||||||
def test_read():
|
def test_read():
|
||||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
|
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
|
||||||
with RealSenseCamera(config) as camera:
|
with RealSenseCamera(config) as camera:
|
||||||
@@ -228,6 +289,203 @@ def test_read_latest_too_old():
|
|||||||
_ = camera.read_latest(max_age_ms=0) # immediately too old
|
_ = camera.read_latest(max_age_ms=0) # immediately too old
|
||||||
|
|
||||||
|
|
||||||
|
def _make_mock_sensor(name: str, supported_options: set | None = None) -> MagicMock:
|
||||||
|
"""Build a fake rs.sensor that reports a name and a configurable supported-options set."""
|
||||||
|
supported = supported_options if supported_options is not None else set()
|
||||||
|
sensor = MagicMock()
|
||||||
|
sensor.get_info.return_value = name
|
||||||
|
sensor.supports.side_effect = lambda opt: opt in supported
|
||||||
|
return sensor
|
||||||
|
|
||||||
|
|
||||||
|
def _attach_mock_color_sensor(camera: RealSenseCamera, sensor: MagicMock) -> None:
|
||||||
|
"""Wire camera.rs_profile so _get_color_sensor finds the given sensor."""
|
||||||
|
profile = MagicMock()
|
||||||
|
device = MagicMock()
|
||||||
|
device.query_sensors.return_value = [sensor]
|
||||||
|
profile.get_device.return_value = device
|
||||||
|
camera.rs_profile = profile
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_color_sensor_prefers_rgb_camera():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
rgb = _make_mock_sensor("RGB Camera")
|
||||||
|
stereo = _make_mock_sensor("Stereo Module")
|
||||||
|
profile = MagicMock()
|
||||||
|
device = MagicMock()
|
||||||
|
device.query_sensors.return_value = [stereo, rgb]
|
||||||
|
profile.get_device.return_value = device
|
||||||
|
camera.rs_profile = profile
|
||||||
|
|
||||||
|
assert camera._get_color_sensor() is rgb
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_color_sensor_falls_back_to_stereo_module():
|
||||||
|
"""D405 has no separate RGB module; color comes from Stereo Module."""
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
stereo = _make_mock_sensor("Stereo Module")
|
||||||
|
_attach_mock_color_sensor(camera, stereo)
|
||||||
|
|
||||||
|
assert camera._get_color_sensor() is stereo
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_color_sensor_raises_with_available_sensors():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
other = _make_mock_sensor("Motion Module")
|
||||||
|
_attach_mock_color_sensor(camera, other)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Motion Module"):
|
||||||
|
camera._get_color_sensor()
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_sensor_options_skipped_when_none():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
with patch.object(RealSenseCamera, "_get_color_sensor") as mock_get:
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
mock_get.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_sensor_options_applies_all_values():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120, gain=64, white_balance=4600)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor(
|
||||||
|
"RGB Camera",
|
||||||
|
supported_options={
|
||||||
|
rs.option.enable_auto_exposure,
|
||||||
|
rs.option.exposure,
|
||||||
|
rs.option.gain,
|
||||||
|
rs.option.enable_auto_white_balance,
|
||||||
|
rs.option.white_balance,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
sensor.set_option.assert_any_call(rs.option.enable_auto_exposure, 0)
|
||||||
|
sensor.set_option.assert_any_call(rs.option.exposure, 120)
|
||||||
|
sensor.set_option.assert_any_call(rs.option.gain, 64)
|
||||||
|
sensor.set_option.assert_any_call(rs.option.enable_auto_white_balance, 0)
|
||||||
|
sensor.set_option.assert_any_call(rs.option.white_balance, 4600)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("config_field", "option", "label"),
|
||||||
|
[
|
||||||
|
("exposure", rs.option.exposure, "exposure"),
|
||||||
|
("gain", rs.option.gain, "gain"),
|
||||||
|
("white_balance", rs.option.white_balance, "white balance"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_configure_sensor_options_raises_when_requested_option_is_unsupported(config_field, option, label):
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: 100})
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor("RGB Camera", supported_options=set())
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match=label):
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
sensor.supports.assert_any_call(option)
|
||||||
|
sensor.set_option.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("config_field", "option", "value"),
|
||||||
|
[
|
||||||
|
("exposure", rs.option.exposure, 120),
|
||||||
|
("gain", rs.option.gain, 64),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_configure_sensor_options_exposure_or_gain_disables_auto_exposure(config_field, option, value):
|
||||||
|
"""white_balance=None should not touch auto white balance."""
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: value})
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor(
|
||||||
|
"RGB Camera",
|
||||||
|
supported_options={rs.option.enable_auto_exposure, option},
|
||||||
|
)
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
calls = [call.args for call in sensor.set_option.call_args_list]
|
||||||
|
assert (rs.option.enable_auto_exposure, 0) in calls
|
||||||
|
assert (option, value) in calls
|
||||||
|
for opt, _ in calls:
|
||||||
|
assert opt != rs.option.enable_auto_white_balance
|
||||||
|
assert opt != rs.option.white_balance
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_sensor_options_warns_when_auto_exposure_control_is_unsupported(caplog):
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.exposure})
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
with caplog.at_level("WARNING"):
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
sensor.set_option.assert_called_once_with(rs.option.exposure, 120)
|
||||||
|
assert "does not support disabling auto-exposure" in caplog.text
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_sensor_options_warns_when_auto_white_balance_control_is_unsupported(caplog):
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", white_balance=4600)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.white_balance})
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
with caplog.at_level("WARNING"):
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
sensor.set_option.assert_called_once_with(rs.option.white_balance, 4600)
|
||||||
|
assert "does not support disabling auto white balance" in caplog.text
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_sensor_options_out_of_range_raises_value_error():
|
||||||
|
"""set_option errors should be re-raised as ValueError with range diagnostics."""
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=999999)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor(
|
||||||
|
"RGB Camera",
|
||||||
|
supported_options={rs.option.enable_auto_exposure, rs.option.exposure},
|
||||||
|
)
|
||||||
|
|
||||||
|
def fake_set_option(option, value):
|
||||||
|
if option == rs.option.exposure:
|
||||||
|
raise RuntimeError("value out of range")
|
||||||
|
|
||||||
|
sensor.set_option.side_effect = fake_set_option
|
||||||
|
|
||||||
|
option_range = MagicMock(min=1, max=10000, step=1, default=156)
|
||||||
|
sensor.get_option_range.return_value = option_range
|
||||||
|
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="exposure") as exc_info:
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
msg = str(exc_info.value)
|
||||||
|
assert "999999" in msg
|
||||||
|
assert "min=1" in msg
|
||||||
|
assert "max=10000" in msg
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"rotation",
|
"rotation",
|
||||||
[
|
[
|
||||||
|
|||||||
@@ -0,0 +1,104 @@
|
|||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||||
|
|
||||||
|
from lerobot.scripts.augment_dataset_quantile_stats import (
|
||||||
|
compute_quantile_stats_for_dataset,
|
||||||
|
has_quantile_stats,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _numeric_keys(dataset):
|
||||||
|
return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
|
||||||
|
|
||||||
|
|
||||||
|
def _image_keys(dataset):
|
||||||
|
return [k for k, v in dataset.features.items() if v["dtype"] in ("image", "video")]
|
||||||
|
|
||||||
|
|
||||||
|
def test_numeric_stats_are_unaffected_by_sampling(tmp_path, lerobot_dataset_factory):
|
||||||
|
"""Sampling only touches image/video frames; numeric features are read in
|
||||||
|
full either way, so their stats must be identical with and without sampling."""
|
||||||
|
dataset = lerobot_dataset_factory(
|
||||||
|
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
|
||||||
|
)
|
||||||
|
|
||||||
|
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
|
||||||
|
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
|
||||||
|
|
||||||
|
numeric_keys = _numeric_keys(dataset)
|
||||||
|
assert numeric_keys, "fixture should expose numeric features"
|
||||||
|
for key in numeric_keys:
|
||||||
|
if key not in exact:
|
||||||
|
continue
|
||||||
|
for stat in ("mean", "std", "q01", "q50", "q99"):
|
||||||
|
if stat in exact[key]:
|
||||||
|
np.testing.assert_allclose(
|
||||||
|
sampled[key][stat],
|
||||||
|
exact[key][stat],
|
||||||
|
rtol=1e-6,
|
||||||
|
atol=1e-6,
|
||||||
|
err_msg=f"numeric feature '{key}' stat '{stat}' changed under sampling",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_image_sampling_reduces_data_but_keeps_stats_close(tmp_path, lerobot_dataset_factory):
|
||||||
|
"""For images, sampling should reduce the number of samples considered while
|
||||||
|
keeping the resulting statistics close to the exact ones."""
|
||||||
|
dataset = lerobot_dataset_factory(
|
||||||
|
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
|
||||||
|
)
|
||||||
|
|
||||||
|
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
|
||||||
|
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
|
||||||
|
|
||||||
|
image_keys = _image_keys(dataset)
|
||||||
|
assert image_keys, "fixture should expose at least one image feature"
|
||||||
|
for key in image_keys:
|
||||||
|
# sampling actually looked at fewer pixels
|
||||||
|
assert sampled[key]["count"][0] < exact[key]["count"][0]
|
||||||
|
# but per-channel mean stays close
|
||||||
|
np.testing.assert_allclose(
|
||||||
|
sampled[key]["mean"],
|
||||||
|
exact[key]["mean"],
|
||||||
|
rtol=0.15,
|
||||||
|
err_msg=f"image feature '{key}' mean drifted too far under sampling",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_short_episodes_use_all_frames(tmp_path, lerobot_dataset_factory):
|
||||||
|
"""With episodes shorter than the sampling floor, sampling is a no-op and
|
||||||
|
must produce exactly the same stats as the exact path."""
|
||||||
|
dataset = lerobot_dataset_factory(
|
||||||
|
root=tmp_path / "ds", total_episodes=2, total_frames=40, use_videos=False
|
||||||
|
)
|
||||||
|
|
||||||
|
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
|
||||||
|
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
|
||||||
|
|
||||||
|
for key in _image_keys(dataset):
|
||||||
|
assert sampled[key]["count"][0] == exact[key]["count"][0]
|
||||||
|
|
||||||
|
|
||||||
|
def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory):
|
||||||
|
"""The computed stats should contain quantile keys for the dataset."""
|
||||||
|
dataset = lerobot_dataset_factory(
|
||||||
|
root=tmp_path / "ds", total_episodes=2, total_frames=200, use_videos=False
|
||||||
|
)
|
||||||
|
stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
|
||||||
|
assert has_quantile_stats(stats)
|
||||||
@@ -482,6 +482,20 @@ def test_add_frame_works_in_write_mode(tmp_path):
|
|||||||
# ── Resume mode ──────────────────────────────────────────────────────
|
# ── Resume mode ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def test_resume_freshly_created_empty_dataset(tmp_path):
|
||||||
|
"""resume() accepts a local dataset created before any episode was recorded."""
|
||||||
|
root = tmp_path / "resume_empty_ds"
|
||||||
|
LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root)
|
||||||
|
|
||||||
|
resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
|
||||||
|
|
||||||
|
assert isinstance(resumed.writer, DatasetWriter)
|
||||||
|
assert resumed.meta.total_episodes == 0
|
||||||
|
assert resumed.meta.total_frames == 0
|
||||||
|
assert resumed.meta.tasks is None
|
||||||
|
assert resumed.meta.episodes is None
|
||||||
|
|
||||||
|
|
||||||
def test_resume_creates_writer(tmp_path):
|
def test_resume_creates_writer(tmp_path):
|
||||||
"""After resume(), writer is a DatasetWriter."""
|
"""After resume(), writer is a DatasetWriter."""
|
||||||
root = tmp_path / "resume_ds"
|
root = tmp_path / "resume_ds"
|
||||||
|
|||||||
@@ -294,6 +294,19 @@ def test__sync_read(addr, length, ids_values, mock_motors, dummy_motors):
|
|||||||
assert read_values == ids_values
|
assert read_values == ids_values
|
||||||
|
|
||||||
|
|
||||||
|
def test__sync_read_retries_after_transient_failure(mock_motors, dummy_motors):
|
||||||
|
addr, length, ids_values = (10, 4, {1: 1337})
|
||||||
|
stub = mock_motors.build_sync_read_stub(addr, length, ids_values, num_invalid_try=1)
|
||||||
|
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||||
|
bus.connect(handshake=False)
|
||||||
|
|
||||||
|
read_values, read_comm = bus._sync_read(addr, length, list(ids_values), num_retry=1)
|
||||||
|
|
||||||
|
assert read_comm == scs.COMM_SUCCESS
|
||||||
|
assert read_values == ids_values
|
||||||
|
assert mock_motors.stubs[stub].calls == 2
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||||
def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
|
def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
|
||||||
addr, length, ids_values = (10, 4, {1: 1337})
|
addr, length, ids_values = (10, 4, {1: 1337})
|
||||||
|
|||||||
@@ -496,6 +496,60 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
|
|||||||
assert "embodiment_id" in processed
|
assert "embodiment_id" in processed
|
||||||
|
|
||||||
|
|
||||||
|
def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path):
|
||||||
|
"""Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded.
|
||||||
|
|
||||||
|
Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer
|
||||||
|
overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats
|
||||||
|
are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to
|
||||||
|
max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the
|
||||||
|
stats or normalization crashes with a shape mismatch.
|
||||||
|
"""
|
||||||
|
config = make_config()
|
||||||
|
preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
|
||||||
|
preprocessor.save_pretrained(tmp_path)
|
||||||
|
postprocessor.save_pretrained(tmp_path)
|
||||||
|
|
||||||
|
# Reload with the generic override path injecting raw, unpadded dataset stats.
|
||||||
|
raw_stats = make_stats()
|
||||||
|
loaded_pre = PolicyProcessorPipeline.from_pretrained(
|
||||||
|
tmp_path,
|
||||||
|
config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
|
||||||
|
overrides={"normalizer_processor": {"stats": raw_stats}},
|
||||||
|
to_transition=batch_to_transition,
|
||||||
|
to_output=transition_to_batch,
|
||||||
|
)
|
||||||
|
loaded_post = PolicyProcessorPipeline.from_pretrained(
|
||||||
|
tmp_path,
|
||||||
|
config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
|
||||||
|
overrides={"unnormalizer_processor": {"stats": raw_stats}},
|
||||||
|
to_transition=policy_action_to_transition,
|
||||||
|
to_output=transition_to_policy_action,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Sanity: the override really injected unpadded stats before reconciliation.
|
||||||
|
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
|
||||||
|
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,)
|
||||||
|
|
||||||
|
loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post)
|
||||||
|
|
||||||
|
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
|
||||||
|
unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep))
|
||||||
|
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
|
||||||
|
assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
|
||||||
|
assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
|
||||||
|
|
||||||
|
# Normalizing a padded state must not raise (this is the exact runtime path that crashed).
|
||||||
|
processed = loaded_pre(
|
||||||
|
{
|
||||||
|
"task": "pick the block",
|
||||||
|
OBS_STATE: torch.zeros(STATE_DIM),
|
||||||
|
f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM)
|
||||||
|
|
||||||
|
|
||||||
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
|
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
|
||||||
monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
|
monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
|
||||||
policy = modeling_evo1.Evo1Policy(make_config())
|
policy = modeling_evo1.Evo1Policy(make_config())
|
||||||
|
|||||||
@@ -0,0 +1,83 @@
|
|||||||
|
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
import lerobot.policies.factory as policy_factory
|
||||||
|
|
||||||
|
|
||||||
|
def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch):
|
||||||
|
cfg = SimpleNamespace(
|
||||||
|
type="mock",
|
||||||
|
device="cpu",
|
||||||
|
pretrained_path="user/adapter",
|
||||||
|
pretrained_revision="adapter-sha",
|
||||||
|
use_peft=True,
|
||||||
|
input_features={},
|
||||||
|
output_features={},
|
||||||
|
)
|
||||||
|
dataset_meta = SimpleNamespace(features={}, stats={})
|
||||||
|
|
||||||
|
base_policy = torch.nn.Linear(1, 1)
|
||||||
|
policy_from_pretrained = MagicMock(return_value=base_policy)
|
||||||
|
policy_class = SimpleNamespace(from_pretrained=policy_from_pretrained)
|
||||||
|
monkeypatch.setattr(policy_factory, "get_policy_class", lambda _: policy_class)
|
||||||
|
monkeypatch.setattr(policy_factory, "dataset_to_policy_features", lambda _: {})
|
||||||
|
monkeypatch.setattr(policy_factory, "validate_visual_features_consistency", lambda *args: None)
|
||||||
|
|
||||||
|
peft_config = SimpleNamespace(
|
||||||
|
base_model_name_or_path="user/base-policy",
|
||||||
|
revision="base-sha",
|
||||||
|
)
|
||||||
|
peft_config_from_pretrained = MagicMock(return_value=peft_config)
|
||||||
|
adapted_policy = torch.nn.Linear(1, 1)
|
||||||
|
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
|
||||||
|
require_package = MagicMock()
|
||||||
|
monkeypatch.setattr(policy_factory, "require_package", require_package)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
policy_factory,
|
||||||
|
"PeftConfig",
|
||||||
|
SimpleNamespace(from_pretrained=peft_config_from_pretrained),
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
policy_factory,
|
||||||
|
"PeftModel",
|
||||||
|
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
|
||||||
|
)
|
||||||
|
|
||||||
|
policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)
|
||||||
|
|
||||||
|
assert policy is adapted_policy
|
||||||
|
require_package.assert_called_once_with("peft", extra="peft")
|
||||||
|
peft_config_from_pretrained.assert_called_once_with(
|
||||||
|
"user/adapter",
|
||||||
|
revision="adapter-sha",
|
||||||
|
)
|
||||||
|
policy_from_pretrained.assert_called_once_with(
|
||||||
|
config=cfg,
|
||||||
|
dataset_stats=dataset_meta.stats,
|
||||||
|
dataset_meta=dataset_meta,
|
||||||
|
pretrained_name_or_path="user/base-policy",
|
||||||
|
revision="base-sha",
|
||||||
|
)
|
||||||
|
peft_model_from_pretrained.assert_called_once_with(
|
||||||
|
base_policy,
|
||||||
|
"user/adapter",
|
||||||
|
config=peft_config,
|
||||||
|
revision="adapter-sha",
|
||||||
|
is_trainable=True,
|
||||||
|
)
|
||||||
@@ -113,6 +113,7 @@ def test_gaussian_actor_config_default_initialization():
|
|||||||
# Concurrency configuration
|
# Concurrency configuration
|
||||||
assert config.concurrency.actor == "threads"
|
assert config.concurrency.actor == "threads"
|
||||||
assert config.concurrency.learner == "threads"
|
assert config.concurrency.learner == "threads"
|
||||||
|
assert config.concurrency.multiprocessing_context == "spawn"
|
||||||
|
|
||||||
assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
|
assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
|
||||||
assert isinstance(config.policy_kwargs, PolicyConfig)
|
assert isinstance(config.policy_kwargs, PolicyConfig)
|
||||||
@@ -152,6 +153,7 @@ def test_concurrency_config():
|
|||||||
config = ConcurrencyConfig()
|
config = ConcurrencyConfig()
|
||||||
assert config.actor == "threads"
|
assert config.actor == "threads"
|
||||||
assert config.learner == "threads"
|
assert config.learner == "threads"
|
||||||
|
assert config.multiprocessing_context == "spawn"
|
||||||
|
|
||||||
|
|
||||||
def test_gaussian_actor_config_custom_initialization():
|
def test_gaussian_actor_config_custom_initialization():
|
||||||
|
|||||||
@@ -0,0 +1,45 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
|
||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||||
|
from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||||
|
ForwardKinematicsJointsToEEAction,
|
||||||
|
ForwardKinematicsJointsToEEObservation,
|
||||||
|
)
|
||||||
|
|
||||||
|
MOTOR_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
|
||||||
|
EE_KEYS = {f"ee.{k}" for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]}
|
||||||
|
|
||||||
|
|
||||||
|
def _joint_bucket(feature_type: FeatureType) -> dict[str, PolicyFeature]:
|
||||||
|
return {f"{n}.pos": PolicyFeature(type=feature_type, shape=(1,)) for n in MOTOR_NAMES}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("step_cls", "bucket", "feature_type"),
|
||||||
|
[
|
||||||
|
(ForwardKinematicsJointsToEEAction, PipelineFeatureType.ACTION, FeatureType.ACTION),
|
||||||
|
(ForwardKinematicsJointsToEEObservation, PipelineFeatureType.OBSERVATION, FeatureType.STATE),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_fk_feature_schema(step_cls, bucket, feature_type):
|
||||||
|
features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}}
|
||||||
|
features[bucket] = _joint_bucket(feature_type)
|
||||||
|
out = step_cls(kinematics=None, motor_names=MOTOR_NAMES).transform_features(features)[bucket]
|
||||||
|
assert set(out) == EE_KEYS
|
||||||
|
assert {feature.type for feature in out.values()} == {feature_type}
|
||||||
@@ -49,7 +49,7 @@ def _make_bus_mock() -> MagicMock:
|
|||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def follower():
|
def follower(tmp_path):
|
||||||
bus_mock = _make_bus_mock()
|
bus_mock = _make_bus_mock()
|
||||||
|
|
||||||
def _bus_side_effect(*_args, **kwargs):
|
def _bus_side_effect(*_args, **kwargs):
|
||||||
@@ -71,7 +71,7 @@ def follower():
|
|||||||
),
|
),
|
||||||
patch.object(SO100Follower, "configure", lambda self: None),
|
patch.object(SO100Follower, "configure", lambda self: None),
|
||||||
):
|
):
|
||||||
cfg = SO100FollowerConfig(port="/dev/null")
|
cfg = SO100FollowerConfig(port="/dev/null", calibration_dir=tmp_path)
|
||||||
robot = SO100Follower(cfg)
|
robot = SO100Follower(cfg)
|
||||||
yield robot
|
yield robot
|
||||||
if robot.is_connected:
|
if robot.is_connected:
|
||||||
@@ -99,6 +99,27 @@ def test_get_observation(follower):
|
|||||||
assert obs[f"{motor}.pos"] == idx
|
assert obs[f"{motor}.pos"] == idx
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_observation_uses_read_retries(follower):
|
||||||
|
# Feetech buses can intermittently fail a sync_read; the follower should forward the configured
|
||||||
|
# retry count so transient failures don't abort the control loop (see #3131).
|
||||||
|
follower.config.num_read_retries = 7
|
||||||
|
follower.connect()
|
||||||
|
follower.get_observation()
|
||||||
|
|
||||||
|
follower.bus.sync_read.assert_called_once_with("Present_Position", num_retry=7)
|
||||||
|
|
||||||
|
|
||||||
|
def test_send_action_uses_read_retries(follower):
|
||||||
|
follower.config.max_relative_target = 10.0
|
||||||
|
follower.config.num_read_retries = 7
|
||||||
|
follower.connect()
|
||||||
|
|
||||||
|
action = {f"{motor}.pos": value * 10 for value, motor in enumerate(follower.bus.motors, 1)}
|
||||||
|
follower.send_action(action)
|
||||||
|
|
||||||
|
follower.bus.sync_read.assert_called_once_with("Present_Position", num_retry=7)
|
||||||
|
|
||||||
|
|
||||||
def test_send_action(follower):
|
def test_send_action(follower):
|
||||||
follower.connect()
|
follower.connect()
|
||||||
|
|
||||||
|
|||||||
+14
-7
@@ -185,18 +185,25 @@ def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(m
|
|||||||
peft_config_from_pretrained = MagicMock(return_value=peft_config)
|
peft_config_from_pretrained = MagicMock(return_value=peft_config)
|
||||||
adapted_policy = MagicMock()
|
adapted_policy = MagicMock()
|
||||||
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
|
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
|
||||||
monkeypatch.setitem(
|
require_package = MagicMock()
|
||||||
sys.modules,
|
monkeypatch.setattr(rollout_context, "require_package", require_package)
|
||||||
"peft",
|
monkeypatch.setattr(
|
||||||
SimpleNamespace(
|
rollout_context,
|
||||||
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
|
"PeftConfig",
|
||||||
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
|
SimpleNamespace(from_pretrained=peft_config_from_pretrained),
|
||||||
),
|
raising=False,
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
rollout_context,
|
||||||
|
"PeftModel",
|
||||||
|
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
|
||||||
|
raising=False,
|
||||||
)
|
)
|
||||||
|
|
||||||
policy = rollout_context._load_pretrained_policy(policy_config)
|
policy = rollout_context._load_pretrained_policy(policy_config)
|
||||||
|
|
||||||
assert policy is adapted_policy
|
assert policy is adapted_policy
|
||||||
|
require_package.assert_called_once_with("peft", extra="peft")
|
||||||
peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
|
peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
|
||||||
policy_class.from_pretrained.assert_called_once_with(
|
policy_class.from_pretrained.assert_called_once_with(
|
||||||
pretrained_name_or_path="user/base-policy",
|
pretrained_name_or_path="user/base-policy",
|
||||||
|
|||||||
@@ -0,0 +1,36 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from lerobot.utils.device_utils import get_safe_torch_device, is_torch_device_available
|
||||||
|
|
||||||
|
|
||||||
|
def test_cpu_always_available():
|
||||||
|
assert get_safe_torch_device("cpu") == torch.device("cpu")
|
||||||
|
assert is_torch_device_available("cpu")
|
||||||
|
|
||||||
|
|
||||||
|
def test_missing_cuda_raises_valueerror():
|
||||||
|
with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"):
|
||||||
|
get_safe_torch_device("cuda")
|
||||||
|
|
||||||
|
|
||||||
|
def test_missing_mps_raises_valueerror():
|
||||||
|
with patch("torch.backends.mps.is_available", return_value=False), pytest.raises(ValueError, match="MPS"):
|
||||||
|
get_safe_torch_device("mps")
|
||||||
@@ -0,0 +1,46 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from lerobot.utils.rotation import Rotation
|
||||||
|
|
||||||
|
|
||||||
|
def test_zero_quaternion_rejected():
|
||||||
|
with pytest.raises(ValueError, match="non-zero"):
|
||||||
|
Rotation(np.zeros(4))
|
||||||
|
|
||||||
|
|
||||||
|
def test_non_finite_quaternion_rejected():
|
||||||
|
with pytest.raises(ValueError, match="non-zero|finite"):
|
||||||
|
Rotation(np.array([np.nan, 0.0, 0.0, 1.0]))
|
||||||
|
|
||||||
|
|
||||||
|
def test_wrong_shape_rejected():
|
||||||
|
with pytest.raises(ValueError, match="shape"):
|
||||||
|
Rotation(np.array([1.0, 0.0, 0.0]))
|
||||||
|
|
||||||
|
|
||||||
|
def test_identity_roundtrip():
|
||||||
|
r = Rotation.from_rotvec(np.zeros(3))
|
||||||
|
assert np.allclose(r.as_rotvec(), 0.0)
|
||||||
|
assert np.allclose(r.as_matrix(), np.eye(3))
|
||||||
|
|
||||||
|
|
||||||
|
def test_rotvec_roundtrip():
|
||||||
|
rotvec = np.array([0.1, -0.2, 0.3])
|
||||||
|
r = Rotation.from_rotvec(rotvec)
|
||||||
|
assert np.allclose(r.as_rotvec(), rotvec, atol=1e-6)
|
||||||
Reference in New Issue
Block a user