mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 21:19:40 +00:00
Compare commits
44 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 21321d7a27 | |||
| fdf6014214 | |||
| 4db03ed535 | |||
| 79cfb52a71 | |||
| c0e9b0bbff | |||
| b5a13e43ce | |||
| 1e7e0b6de5 | |||
| a3edab661b | |||
| 564ba6395a | |||
| 85086fed7a | |||
| ee06d1005f | |||
| 1c47809bf6 | |||
| 3d70b21aac | |||
| b85620657f | |||
| fbfc861cf2 | |||
| 06aa6a0425 | |||
| be64ded80f | |||
| 88843ed675 | |||
| f2b5c4a47b | |||
| 9202fcea96 | |||
| ef47c35178 | |||
| 6d6c82eb8c | |||
| 9201be92cb | |||
| 0064a06205 | |||
| 710171ccac | |||
| 0f8257443c | |||
| 3a09d0c48a | |||
| 03fc5e3ea9 | |||
| 28c3e095bf | |||
| 5bfb749a9b | |||
| 51c023a7a1 | |||
| 51ea18cb7a | |||
| 04ab43b8d2 | |||
| cdfe192491 | |||
| 3451e53452 | |||
| 30849ce74f | |||
| 7d6907c444 | |||
| d99e1fe89d | |||
| 7fcde61b69 | |||
| bdfe8f8ce9 | |||
| 34d0495d03 | |||
| 834c282631 | |||
| f132885cbc | |||
| d0686be2f5 |
+6
-10
@@ -61,20 +61,16 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
|
||||
**4.1 Install**
|
||||
|
||||
```bash
|
||||
# uv (recommended — see AGENTS.md and CLAUDE.md)
|
||||
uv sync --locked --extra feetech # SO-100/SO-101 motor stack
|
||||
# uv sync --locked --extra all # everything
|
||||
# uv sync --locked --extra smolvla # add SmolVLA deps
|
||||
|
||||
# pip (alternative, e.g. when not working from source)
|
||||
# pip install 'lerobot[feetech]'
|
||||
# pip install 'lerobot[all]'
|
||||
# pip install 'lerobot[smolvla]'
|
||||
|
||||
pip install 'lerobot[feetech]' # SO-100/SO-101 motor stack
|
||||
# pip install 'lerobot[all]' # everything
|
||||
# pip install 'lerobot[aloha,pusht]' # specific features
|
||||
# pip install 'lerobot[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`).
|
||||
|
||||
**4.2 Find USB ports** — run once per arm, unplug when prompted.
|
||||
|
||||
```bash
|
||||
|
||||
@@ -68,16 +68,17 @@ ENV HOME=/home/user_lerobot \
|
||||
# issues with MuJoCo and OpenGL drivers.
|
||||
RUN uv venv --python python${PYTHON_VERSION}
|
||||
|
||||
# Install third-party dependencies separately for layer caching
|
||||
# Install Python dependencies for caching
|
||||
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
|
||||
RUN uv sync --locked --extra all --no-install-project --no-cache
|
||||
COPY --chown=user_lerobot:user_lerobot src/ src/
|
||||
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
|
||||
RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
|
||||
|
||||
# Copy the application source code and install the local project
|
||||
# Copy the rest of the application source code
|
||||
# Make sure to have the git-LFS files for testing
|
||||
COPY --chown=user_lerobot:user_lerobot . .
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
|
||||
# Set the default command
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
@@ -60,14 +60,15 @@ ENV HOME=/home/user_lerobot \
|
||||
# run other Python projects in the same container without dependency conflicts.
|
||||
RUN uv venv
|
||||
|
||||
# Install third-party dependencies separately for layer caching
|
||||
# Install Python dependencies for caching
|
||||
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
|
||||
RUN uv sync --locked --extra all --no-install-project --no-cache
|
||||
COPY --chown=user_lerobot:user_lerobot src/ src/
|
||||
|
||||
# Copy the application code and install the local project
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
|
||||
# Copy the rest of the application code
|
||||
# Make sure to have the git-LFS files for testing
|
||||
COPY --chown=user_lerobot:user_lerobot . .
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
|
||||
# Set the default command
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
@@ -37,6 +37,8 @@
|
||||
- sections:
|
||||
- local: lerobot-dataset-v3
|
||||
title: Using LeRobotDataset
|
||||
- local: training_dataset_streaming
|
||||
title: Training Dataset Streaming
|
||||
- local: porting_datasets_v3
|
||||
title: Porting Large Datasets
|
||||
- local: using_dataset_tools
|
||||
|
||||
@@ -58,7 +58,7 @@ final_action = postprocessor(action)
|
||||
|
||||
## Hardware API redesign
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
### 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.
|
||||
|
||||
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 to your inference script (shown here in the `record.py` script):
|
||||
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.
|
||||
Then, add these same transformations on your inference script (shown here in the `record.py` script):
|
||||
|
||||
```diff
|
||||
action_values = predict_action(
|
||||
|
||||
@@ -136,10 +136,6 @@ config = RealSenseCameraConfig(
|
||||
height=480,
|
||||
color_mode=ColorMode.RGB,
|
||||
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
|
||||
)
|
||||
|
||||
@@ -158,15 +154,6 @@ finally:
|
||||
```
|
||||
<!-- 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`.
|
||||
|
||||
Manual color controls require a dedicated RGB module. Cameras without one, such as the RealSense
|
||||
D405, do not support them and raise an error at connection time.
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
|
||||
@@ -88,6 +88,20 @@ 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:
|
||||
|
||||
````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
|
||||
# Use SmolVLA policy with LIBERO environment
|
||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||
@@ -102,7 +116,6 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||
policy_cfg=act_cfg,
|
||||
)
|
||||
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
||||
```
|
||||
|
||||
### 3. **Easier Experimentation**
|
||||
|
||||
@@ -132,7 +145,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
|
||||
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
|
||||
gripper_pos, gripper_vel], dim=-1) # 14D
|
||||
return state
|
||||
```
|
||||
````
|
||||
|
||||
### 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:
|
||||
|
||||
1. **Select the motor** from the list by clicking on it
|
||||
2. **Click the Upgrade tab**:
|
||||
3. **Click the Online button**:
|
||||
- If a potential firmware update is found, it will be displayed in the box
|
||||
4. **Click the Upgrade button**:
|
||||
2. **Click on Upgrade tab**:
|
||||
3. **Click on Online button**:
|
||||
- If an potential firmware update is found, it will be displayed in the box
|
||||
4. **Click on Upgrade button**:
|
||||
- The update progress will be displayed
|
||||
|
||||
## Step 6: Verify Update
|
||||
|
||||
@@ -59,7 +59,6 @@ The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with
|
||||
|
||||
- `bi_openarm_mini` - Bimanual OpenArm Mini
|
||||
- `so_leader` - SO100 / SO101 leader arm
|
||||
- `bi_so_leader` - Bimanual SO100 / SO101 leader arms
|
||||
|
||||
> [!IMPORTANT]
|
||||
> The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`.
|
||||
|
||||
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
|
||||
|
||||
### 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).
|
||||
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).
|
||||
|
||||
```bash
|
||||
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
|
||||
|
||||
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/so_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/so101_follower/so101_follower.py)
|
||||
|
||||
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]"
|
||||
```
|
||||
|
||||
Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
|
||||
Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
|
||||
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
|
||||
|
||||
@@ -131,15 +131,54 @@ for batch in data_loader:
|
||||
# model.forward(batch)
|
||||
```
|
||||
|
||||
## Stream a dataset (no downloads)
|
||||
## Stream a dataset during training
|
||||
|
||||
Use `StreamingLeRobotDataset` to iterate directly from the Hub without local copies. This allows to stream large datasets without the need to downloading them onto disk or loading them onto memory, and is a key feature of the new dataset format.
|
||||
Enable training-time streaming with the public `--dataset.streaming=true` flag:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=yaak-ai/L2D-v3 \
|
||||
--dataset.streaming=true \
|
||||
--policy.type=act \
|
||||
--output_dir=outputs/train/act_streaming
|
||||
```
|
||||
|
||||
This is separate from `--dataset.streaming_encoding=true`, which controls video encoding while
|
||||
recording. Training-time streaming leaves the map-style `LeRobotDataset` path unchanged.
|
||||
|
||||
`StreamingLeRobotDataset` assigns complete episodes disjointly across distributed ranks and
|
||||
DataLoader workers, reads only the selected episode rows from Parquet, and keeps a bounded pool of
|
||||
episode video bytes. Every selected frame is visited exactly once per streaming epoch.
|
||||
|
||||
On first use, LeRobot looks for a revision-matched MP4 index sidecar. If none is published with the
|
||||
dataset, it builds one in the revision-keyed local LeRobot cache under a process lock and installs it
|
||||
atomically. Training never uploads this sidecar. Dataset maintainers can build and publish one
|
||||
explicitly with `scripts/build_mp4_sidecar.py --push`.
|
||||
|
||||
The default memory cap is 8 GiB per DataLoader worker. It can be adjusted along with episode mixing
|
||||
and prefetch:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=yaak-ai/L2D-v3 \
|
||||
--dataset.streaming=true \
|
||||
--dataset.streaming_byte_budget_gb=4 \
|
||||
--dataset.streaming_episode_pool_size=16 \
|
||||
--dataset.streaming_prefetch_episodes=4 \
|
||||
--policy.type=act \
|
||||
--output_dir=outputs/train/act_streaming
|
||||
```
|
||||
|
||||
Use `--dataset.streaming_data_root=hf://buckets/OWNER/BUCKET/PREFIX` when metadata lives in a
|
||||
dataset repository but Parquet and MP4 data are mirrored in an HF Bucket.
|
||||
|
||||
The Python API uses the same implementation:
|
||||
|
||||
```python
|
||||
from lerobot.datasets import StreamingLeRobotDataset
|
||||
|
||||
repo_id = "yaak-ai/L2D-v3"
|
||||
dataset = StreamingLeRobotDataset(repo_id) # streams directly from the Hub
|
||||
dataset = StreamingLeRobotDataset(repo_id)
|
||||
```
|
||||
|
||||
<div style="display:flex; justify-content:center; gap:12px; flex-wrap:wrap;">
|
||||
|
||||
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
|
||||
|
||||
## 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 40k 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 40kk 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:
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
|
||||
## Three pipelines
|
||||
|
||||
We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
|
||||
Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
|
||||
Each of these pipelines handle 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)
|
||||
2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
|
||||
@@ -74,7 +74,7 @@ In the phone to SO-100 follower examples we use the following adapters:
|
||||
- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
|
||||
- `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.
|
||||
- `transition_to_observation`: transforms the pipeline transition to an observation dict.
|
||||
- `transition_to_observation`: transforms the pipeline transition to a observation dict.
|
||||
|
||||
Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
|
||||
|
||||
@@ -82,7 +82,7 @@ Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/l
|
||||
|
||||
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 an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
|
||||
Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
|
||||
|
||||
```python
|
||||
def transform_features(
|
||||
|
||||
@@ -82,8 +82,6 @@ By default the env samples objects only from the `lightwheel` registry (what `--
|
||||
|
||||
All eval snippets below mirror the CI command (see `.github/workflows/benchmark_tests.yml`). The `--rename_map` argument maps RoboCasa's native camera keys (`robot0_agentview_left` / `robot0_eye_in_hand` / `robot0_agentview_right`) onto the three-camera (`camera1` / `camera2` / `camera3`) input layout the released `smolvla_robocasa` policy was trained on.
|
||||
|
||||
By default, each task uses the rollout horizon registered by RoboCasa. Set `--env.episode_length=<steps>` to apply the same explicit horizon to every selected task.
|
||||
|
||||
### Single-task evaluation (recommended for quick iteration)
|
||||
|
||||
```bash
|
||||
|
||||
+2
-2
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
|
||||
policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
|
||||
|
||||
# Now use predict_action_chunk with RTC parameters
|
||||
inference_delay = 4 # How many steps of inference latency, this value should be calculated based on the inference latency of the policy
|
||||
inference_delay = 4 # How many steps of inference latency, this values should be calculated based on the inference latency of the policy
|
||||
|
||||
# Initialize the action queue
|
||||
action_queue = ActionQueue(policy_cfg.rtc_config)
|
||||
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
|
||||
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 an 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 a optimal value.
|
||||
|
||||
**`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
|
||||
|
||||
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).
|
||||
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).
|
||||
Once you are logged in, you can run inference in your setup by doing:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -338,7 +338,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
|
||||
<hfoption id="Leader">
|
||||
|
||||
- Mount the leader holder onto the wrist and secure it with 4 M3x6mm screws.
|
||||
- Attach the handle to the leader holder using 1 M2x6mm screw.
|
||||
- Attach the handle to motor 5 using 1 M2x6mm screw.
|
||||
- Insert the gripper motor, secure it with 2 M2x6mm screws on each side, attach a motor horn using a M3x6mm horn screw.
|
||||
- Attach the follower trigger with 4 M3x6mm screws.
|
||||
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
# Training Dataset Streaming
|
||||
|
||||
Training-time dataset streaming lets `lerobot-train` consume a LeRobotDataset v3 without first
|
||||
downloading its complete Parquet and video payload. Enable it through the existing public switch:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=OWNER/DATASET \
|
||||
--dataset.streaming=true \
|
||||
--policy.type=act \
|
||||
--output_dir=outputs/train/act_streaming
|
||||
```
|
||||
|
||||
This feature is independent of `--dataset.streaming_encoding=true`. `streaming_encoding` controls
|
||||
how videos are written while recording; `dataset.streaming` controls how an existing dataset is read
|
||||
during training. Recording and rollout encoders are not used by this training path.
|
||||
|
||||
## How an epoch is read
|
||||
|
||||
Each distributed rank owns a deterministic, frame-balanced set of complete episodes. One logical
|
||||
exact-coverage pool per rank mixes those episodes while visiting every selected frame once per
|
||||
rank-local coverage epoch. Parquet columns and compressed MP4 byte ranges are prefetched from the
|
||||
same byte-aware admission frontier. Temporal history and future windows are resolved inside the
|
||||
complete episode, including the same boundary padding masks as map-style loading.
|
||||
|
||||
When `--num_workers` is nonzero, training uses one dedicated DataLoader process per rank. Its
|
||||
bounded result queue holds decoded batches while the policy trains. The configured worker count is
|
||||
instead used as internal Parquet and byte-range fetch concurrency, so increasing it does not create
|
||||
independent samplers or multiply the cache budget.
|
||||
|
||||
The default map-style `LeRobotDataset` behavior is unchanged when `--dataset.streaming=false`.
|
||||
|
||||
## MP4 sidecars and the first run
|
||||
|
||||
Video streaming uses a small MP4 index sidecar. Dataset initialization first checks the
|
||||
revision-keyed local cache, then looks for a valid published sidecar. If neither is available,
|
||||
LeRobot builds the sidecar locally while holding a process lock and installs it atomically. A
|
||||
failed or interrupted build does not replace the previous valid file.
|
||||
|
||||
Training is read-only: it never uploads a sidecar or modifies the dataset repository. On a cluster
|
||||
with node-local caches, the first job may build once per node. A shared LeRobot cache avoids that
|
||||
duplication.
|
||||
|
||||
Dataset maintainers can build a sidecar ahead of time:
|
||||
|
||||
```bash
|
||||
uv run python scripts/build_mp4_sidecar.py \
|
||||
--repo-id=OWNER/DATASET \
|
||||
--revision=COMMIT_SHA \
|
||||
--data-root=hf://datasets/OWNER/DATASET@COMMIT_SHA \
|
||||
--output=/tmp/dataset-mp4-sidecar.npz
|
||||
```
|
||||
|
||||
Publication is always explicit. Add `--push` only after validating the complete-dataset sidecar.
|
||||
Subset sidecars cannot be published.
|
||||
|
||||
## Configuration
|
||||
|
||||
The production defaults are:
|
||||
|
||||
| Option | Default | Meaning |
|
||||
| ----------------------------------- | -------------: | --------------------------------------------------- |
|
||||
| `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank |
|
||||
| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool |
|
||||
| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank |
|
||||
| `streaming_decode_threads` | 2 | Parallel sample assembly and video decode workers |
|
||||
| `streaming_decoded_queue_size` | 8 | Decoded samples buffered ahead, in planner order |
|
||||
| `streaming_max_open_decoders` | pool × cameras | Independent open-decoder LRU cap per rank |
|
||||
| `streaming_native_http_connections` | unset | Native HTTP connection cap per rank |
|
||||
| `streaming_native_http_subranges` | 1 | Concurrent subranges per native HTTP range read |
|
||||
| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root |
|
||||
|
||||
The active episode set is capped by both episode count and the exact synthesized mini-MP4 sizes
|
||||
computed from the sidecar. An episode larger than the complete rank budget fails before training
|
||||
fetches its payload. The cache, decoder LRU, and decoded-batch queue remain independently bounded.
|
||||
Start with a smaller pool or budget on memory-constrained hosts:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=OWNER/DATASET \
|
||||
--dataset.streaming=true \
|
||||
--dataset.streaming_episode_pool_size=16 \
|
||||
--dataset.streaming_prefetch_episodes=4 \
|
||||
--dataset.streaming_byte_budget_gb=4 \
|
||||
--dataset.streaming_decode_threads=2 \
|
||||
--dataset.streaming_decoded_queue_size=8 \
|
||||
--num_workers=4 \
|
||||
--policy.type=act \
|
||||
--output_dir=outputs/train/act_streaming
|
||||
```
|
||||
|
||||
If metadata remains in a dataset repository while payload files are mirrored elsewhere, set
|
||||
`--dataset.streaming_data_root`. Supported values include local paths, revision-qualified
|
||||
`hf://datasets/...` roots, `hf://buckets/...` roots, and fsspec URLs.
|
||||
|
||||
## Resume and shuffle migration
|
||||
|
||||
The earlier streaming reader used a bounded row shuffle buffer. The episode reader instead has
|
||||
deterministic exact-coverage ordering derived from the seed and epoch. Checkpoint resume restores the
|
||||
per-rank sample offset using the checkpoint batch size. Changing distributed world size or
|
||||
batch size changes ownership or batch boundaries. For sample-exact comparisons, resume with the
|
||||
same world size and batch size. Keep the same internal fetch concurrency when comparing performance.
|
||||
|
||||
Streaming checkpoints created by the earlier multi-worker sampler record a different ownership
|
||||
topology and are rejected for sample-exact resume. Start a new run from the saved policy weights
|
||||
rather than claiming that its dataset stream resumes exactly.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- **The first batch takes a long time:** check logs for a sidecar build. Reuse a shared cache or
|
||||
publish a validated complete sidecar explicitly.
|
||||
- **A sidecar lock times out:** another process may still be indexing the same revision. Confirm it
|
||||
is healthy before removing a stale lock.
|
||||
- **A rank owns no data:** reduce the number of ranks so every rank owns at least one selected
|
||||
episode.
|
||||
- **The byte budget is exceeded:** lower the episode pool, increase the per-rank byte budget, or
|
||||
use an explicitly prepared payload layout with smaller episode ranges.
|
||||
- **Remote reads cannot authenticate:** verify the normal Hugging Face token or fsspec credentials
|
||||
are available in every worker environment. Credentials are never embedded in the sidecar.
|
||||
- **Refill stalls are high:** compare p95/p99 batch wait, reduce network contention, raise prefetch
|
||||
gradually, and verify that the dataset sidecar matches the exact revision.
|
||||
@@ -50,11 +50,11 @@ lerobot-edit-dataset \
|
||||
Divide a dataset into multiple subsets.
|
||||
|
||||
```bash
|
||||
# Split by fractions (e.g. 60% train, 20% val, 20% test)
|
||||
# Split by fractions (e.g. 80% train, 20% test, 20% val)
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht \
|
||||
--operation.type split \
|
||||
--operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'
|
||||
--operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}'
|
||||
|
||||
# Split by specific episode indices
|
||||
lerobot-edit-dataset \
|
||||
|
||||
@@ -44,7 +44,6 @@ from typing import Protocol
|
||||
|
||||
import numpy as np
|
||||
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
RobotProcessorPipeline,
|
||||
@@ -57,6 +56,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
EEBoundsAndSafety,
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, HF_LEROBOT_HOME, TELEOPERATORS
|
||||
from lerobot.utils.robot_utils import precise_sleep
|
||||
|
||||
|
||||
@@ -38,7 +38,7 @@ from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
|
||||
from lerobot.lerobot_types import RobotAction
|
||||
from lerobot.types import RobotAction
|
||||
|
||||
from .base import _GRIPPER_MOTOR_SCALE, IsaacTeleopTeleoperator, _isaacteleop_available
|
||||
from .config_isaac_teleop import SO101LeaderArmConfig
|
||||
|
||||
@@ -32,7 +32,7 @@ from typing import TYPE_CHECKING, Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from lerobot.lerobot_types import RobotAction
|
||||
from lerobot.types import RobotAction
|
||||
|
||||
from .base import IsaacTeleopTeleoperator, _isaacteleop_available
|
||||
from .config_isaac_teleop import XRControllerConfig
|
||||
|
||||
@@ -26,8 +26,8 @@ from __future__ import annotations
|
||||
from dataclasses import dataclass
|
||||
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import RobotAction
|
||||
from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep
|
||||
from lerobot.types import RobotAction
|
||||
from lerobot.utils.rotation import Rotation
|
||||
|
||||
from .base import _GRIPPER_MOTOR_SCALE
|
||||
|
||||
@@ -21,7 +21,6 @@ from lerobot.cameras.opencv import OpenCVCameraConfig
|
||||
from lerobot.common.control_utils import predict_action
|
||||
from lerobot.configs import FeatureType, PolicyFeature
|
||||
from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.policies import make_pre_post_processors
|
||||
from lerobot.policies.act import ACTPolicy
|
||||
@@ -39,6 +38,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
ForwardKinematicsJointsToEE,
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import ACTION, OBS_STR
|
||||
from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
|
||||
from lerobot.utils.keyboard_input import init_keyboard_listener
|
||||
|
||||
@@ -16,7 +16,6 @@
|
||||
|
||||
from lerobot.cameras.opencv import OpenCVCameraConfig
|
||||
from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
RobotProcessorPipeline,
|
||||
@@ -37,6 +36,7 @@ from lerobot.scripts.lerobot_record import record_loop
|
||||
from lerobot.teleoperators.phone import Phone, PhoneConfig
|
||||
from lerobot.teleoperators.phone.config_phone import PhoneOS
|
||||
from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.feature_utils import combine_feature_dicts
|
||||
from lerobot.utils.keyboard_input import init_keyboard_listener
|
||||
from lerobot.utils.utils import log_say
|
||||
|
||||
@@ -17,7 +17,6 @@
|
||||
import time
|
||||
|
||||
from lerobot.datasets import LeRobotDataset
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
RobotProcessorPipeline,
|
||||
@@ -28,6 +27,7 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
|
||||
from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import ACTION
|
||||
from lerobot.utils.robot_utils import precise_sleep
|
||||
from lerobot.utils.utils import log_say
|
||||
|
||||
@@ -27,7 +27,6 @@ Highlight, or DAgger via ``lerobot-rollout --strategy.type=...``.
|
||||
|
||||
from lerobot.cameras.opencv import OpenCVCameraConfig
|
||||
from lerobot.configs import PreTrainedConfig
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
RobotProcessorPipeline,
|
||||
@@ -44,6 +43,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
|
||||
from lerobot.rollout.inference import SyncInferenceConfig
|
||||
from lerobot.rollout.strategies import BaseStrategy
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.process import ProcessSignalHandler
|
||||
from lerobot.utils.utils import init_logging
|
||||
|
||||
|
||||
@@ -15,7 +15,6 @@
|
||||
|
||||
import time
|
||||
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
RobotProcessorPipeline,
|
||||
@@ -32,6 +31,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
from lerobot.teleoperators.phone import Phone, PhoneConfig
|
||||
from lerobot.teleoperators.phone.config_phone import PhoneOS
|
||||
from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.robot_utils import precise_sleep
|
||||
from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
|
||||
|
||||
|
||||
@@ -21,7 +21,6 @@ from lerobot.cameras.opencv import OpenCVCameraConfig
|
||||
from lerobot.common.control_utils import predict_action
|
||||
from lerobot.configs import FeatureType, PolicyFeature
|
||||
from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.policies import make_pre_post_processors
|
||||
from lerobot.policies.act import ACTPolicy
|
||||
@@ -39,6 +38,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
ForwardKinematicsJointsToEE,
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import ACTION, OBS_STR
|
||||
from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
|
||||
from lerobot.utils.keyboard_input import init_keyboard_listener
|
||||
|
||||
@@ -17,7 +17,6 @@
|
||||
|
||||
from lerobot.cameras.opencv import OpenCVCameraConfig
|
||||
from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
RobotProcessorPipeline,
|
||||
@@ -34,6 +33,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
)
|
||||
from lerobot.scripts.lerobot_record import record_loop
|
||||
from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.feature_utils import combine_feature_dicts
|
||||
from lerobot.utils.keyboard_input import init_keyboard_listener
|
||||
from lerobot.utils.utils import log_say
|
||||
|
||||
@@ -18,7 +18,6 @@
|
||||
import time
|
||||
|
||||
from lerobot.datasets import LeRobotDataset
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
RobotProcessorPipeline,
|
||||
@@ -29,6 +28,7 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
|
||||
from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import ACTION
|
||||
from lerobot.utils.robot_utils import precise_sleep
|
||||
from lerobot.utils.utils import log_say
|
||||
|
||||
@@ -25,7 +25,6 @@ forward/inverse kinematics.
|
||||
|
||||
from lerobot.cameras.opencv import OpenCVCameraConfig
|
||||
from lerobot.configs import PreTrainedConfig
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
RobotProcessorPipeline,
|
||||
@@ -42,6 +41,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
|
||||
from lerobot.rollout.inference import SyncInferenceConfig
|
||||
from lerobot.rollout.strategies import BaseStrategy
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.process import ProcessSignalHandler
|
||||
from lerobot.utils.utils import init_logging
|
||||
|
||||
|
||||
@@ -16,7 +16,6 @@
|
||||
|
||||
import time
|
||||
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
RobotProcessorPipeline,
|
||||
@@ -31,6 +30,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.robot_utils import precise_sleep
|
||||
from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
|
||||
|
||||
|
||||
@@ -12,8 +12,12 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""This script demonstrates how to train a Diffusion Policy on the PushT environment,
|
||||
using a dataset processed in streaming mode."""
|
||||
"""Train directly from episode-scoped Parquet and MP4 streams.
|
||||
|
||||
For normal training, prefer ``lerobot-train --dataset.streaming=true`` so distributed sharding,
|
||||
checkpoint resume, and device placement are configured by the training pipeline. This lower-level
|
||||
example shows the underlying Python API.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
@@ -64,18 +68,28 @@ def main():
|
||||
ACTION: [t / dataset_metadata.fps for t in range(cfg.n_action_steps)],
|
||||
}
|
||||
|
||||
# Instantiating the training dataset in streaming mode allows to not consume up memory as the data is fetched
|
||||
# iteratively rather than being load into memory all at once. Retrieved frames are shuffled across epochs
|
||||
dataset = StreamingLeRobotDataset(dataset_id, delta_timestamps=delta_timestamps, tolerance_s=1e-3)
|
||||
# The first run resolves or locally builds a revision-safe MP4 index sidecar. It is never
|
||||
# uploaded implicitly. Episode rows and video byte ranges are then prefetched together.
|
||||
dataset = StreamingLeRobotDataset(
|
||||
dataset_id,
|
||||
delta_timestamps=delta_timestamps,
|
||||
tolerance_s=1e-3,
|
||||
episode_pool_size=16,
|
||||
prefetch_episodes=4,
|
||||
byte_budget_gb=4,
|
||||
repeat=True,
|
||||
)
|
||||
|
||||
optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4)
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
dataset,
|
||||
num_workers=4,
|
||||
# One worker owns the rank-level pool. Internal fetch concurrency is configured through
|
||||
# StreamingLeRobotDataset.max_num_shards (the lerobot-train CLI derives it from num_workers).
|
||||
num_workers=1,
|
||||
batch_size=16,
|
||||
pin_memory=device.type != "cpu",
|
||||
drop_last=True,
|
||||
prefetch_factor=2, # loads batches with multiprocessing while policy trains
|
||||
prefetch_factor=2, # bounded decoded-batch queue while the policy trains
|
||||
persistent_workers=True,
|
||||
)
|
||||
|
||||
# Run training loop.
|
||||
|
||||
+5
-14
@@ -67,8 +67,11 @@ dependencies = [
|
||||
"einops>=0.8.0,<0.9.0",
|
||||
|
||||
# Config & Hub
|
||||
"draccus>=0.11.6,<0.12.0",
|
||||
"draccus==0.10.0", # TODO: Relax version constraint
|
||||
"huggingface-hub>=1.0.0,<2.0.0",
|
||||
"filelock>=3.12.0,<4.0.0",
|
||||
"fsspec>=2023.5.0,<2027.0.0",
|
||||
"httpx>=0.27.0,<1.0.0",
|
||||
"requests>=2.32.0,<3.0.0",
|
||||
|
||||
# Environments
|
||||
@@ -436,6 +439,7 @@ exclude_dirs = [
|
||||
skips = ["B101", "B311", "B404", "B603", "B615"]
|
||||
|
||||
[tool.typos]
|
||||
default.extend-words = { trak = "trak" }
|
||||
default.extend-ignore-re = [
|
||||
"(?Rm)^.*(#|//)\\s*spellchecker:disable-line$", # spellchecker:disable-line
|
||||
"(?s)(#|//)\\s*spellchecker:off.*?\\n\\s*(#|//)\\s*spellchecker:on", # spellchecker:<on|off>
|
||||
@@ -494,19 +498,6 @@ ignore_errors = true
|
||||
module = "lerobot.envs.*"
|
||||
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]]
|
||||
# module = "lerobot.utils.*"
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
#!/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
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import fsspec
|
||||
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
from lerobot.datasets.streaming_sidecar import (
|
||||
build_mp4_sidecar,
|
||||
make_sidecar_spec,
|
||||
published_sidecar_url,
|
||||
range_backend_for_root,
|
||||
)
|
||||
from lerobot.streaming.sidecar import SidecarSpec
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Build a reusable MP4 byte-index sidecar for streaming.")
|
||||
parser.add_argument("--repo-id", required=True)
|
||||
parser.add_argument("--revision", default=None)
|
||||
parser.add_argument("--data-root", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--episodes", type=int, default=None)
|
||||
parser.add_argument("--workers", type=int, default=8)
|
||||
parser.add_argument("--range-backend", choices=("fsspec", "native-http"), default=None)
|
||||
parser.add_argument("--max-probe-mb", type=int, default=64)
|
||||
parser.add_argument("--push", action="store_true", help="Explicitly publish the sidecar to data_root.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def push_sidecar(local_path: str, spec: SidecarSpec) -> list[str]:
|
||||
if not spec.data_root.startswith("hf://"):
|
||||
raise ValueError("--push currently supports only hf:// data roots")
|
||||
|
||||
fs = fsspec.filesystem("hf")
|
||||
remote = published_sidecar_url(spec)
|
||||
fs.put(str(Path(local_path)), remote)
|
||||
return [remote]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
meta = LeRobotDatasetMetadata(args.repo_id, revision=args.revision)
|
||||
meta.ensure_readable()
|
||||
total = (
|
||||
int(meta.total_episodes) if args.episodes is None else min(args.episodes, int(meta.total_episodes))
|
||||
)
|
||||
spec = make_sidecar_spec(meta, args.data_root)
|
||||
if total != int(meta.total_episodes):
|
||||
selected_paths = {
|
||||
str(meta.get_video_file_path(ep_idx, key)) for ep_idx in range(total) for key in meta.video_keys
|
||||
}
|
||||
spec = SidecarSpec(
|
||||
repo_id=spec.repo_id,
|
||||
revision=spec.revision,
|
||||
data_root=spec.data_root,
|
||||
source_files=tuple(item for item in spec.source_files if item[0] in selected_paths),
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
build_mp4_sidecar(
|
||||
args.output,
|
||||
spec,
|
||||
range_backend=args.range_backend or range_backend_for_root(args.data_root),
|
||||
workers=args.workers,
|
||||
max_probe_bytes=args.max_probe_mb * 1024 * 1024,
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
print(f"wrote {args.output}")
|
||||
print(f"episodes={total} files={len(spec.source_files)} elapsed_s={elapsed:.2f}")
|
||||
if args.push:
|
||||
if total != int(meta.total_episodes):
|
||||
raise ValueError("Only a complete dataset sidecar can be published")
|
||||
pushed = push_sidecar(args.output, spec)
|
||||
for remote in pushed:
|
||||
print(f"pushed {remote}")
|
||||
else:
|
||||
print("push_skipped: pass --push for explicit publication")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -38,7 +38,6 @@ import draccus
|
||||
import grpc
|
||||
import torch
|
||||
|
||||
from lerobot.lerobot_types import PolicyAction
|
||||
from lerobot.policies import get_policy_class, make_pre_post_processors
|
||||
from lerobot.processor import PolicyProcessorPipeline
|
||||
from lerobot.transport import (
|
||||
@@ -46,6 +45,7 @@ from lerobot.transport import (
|
||||
services_pb2_grpc, # type: ignore
|
||||
)
|
||||
from lerobot.transport.utils import receive_bytes_in_chunks
|
||||
from lerobot.types import PolicyAction
|
||||
|
||||
from .configs import PolicyServerConfig
|
||||
from .constants import SUPPORTED_POLICIES
|
||||
|
||||
@@ -120,22 +120,14 @@ class OpenCVCamera(Camera):
|
||||
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
||||
self.backend: int = config.backend
|
||||
|
||||
self.capture_width: int | None = None
|
||||
self.capture_height: int | None = None
|
||||
self._reset_connection_settings()
|
||||
|
||||
def __str__(self) -> str:
|
||||
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
|
||||
if self.height and self.width:
|
||||
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
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}({self.index_or_path})"
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Checks if the camera is currently connected and opened."""
|
||||
@@ -172,7 +164,6 @@ class OpenCVCamera(Camera):
|
||||
f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
|
||||
)
|
||||
|
||||
try:
|
||||
self._configure_capture_settings()
|
||||
self._start_read_thread()
|
||||
|
||||
@@ -184,13 +175,6 @@ class OpenCVCamera(Camera):
|
||||
with self.frame_lock:
|
||||
if self.latest_frame is None:
|
||||
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.")
|
||||
|
||||
@@ -328,7 +312,6 @@ class OpenCVCamera(Camera):
|
||||
|
||||
for target in targets_to_scan:
|
||||
camera = cv2.VideoCapture(target)
|
||||
try:
|
||||
if camera.isOpened():
|
||||
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
@@ -338,9 +321,7 @@ class OpenCVCamera(Camera):
|
||||
# Get FOURCC code and convert to string
|
||||
default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
|
||||
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 = {
|
||||
"name": f"OpenCV Camera @ {target}",
|
||||
@@ -357,7 +338,6 @@ class OpenCVCamera(Camera):
|
||||
}
|
||||
|
||||
found_cameras_info.append(camera_info)
|
||||
finally:
|
||||
camera.release()
|
||||
|
||||
return found_cameras_info
|
||||
@@ -516,26 +496,6 @@ class OpenCVCamera(Camera):
|
||||
self.latest_timestamp = None
|
||||
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
|
||||
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
|
||||
"""
|
||||
@@ -626,6 +586,16 @@ class OpenCVCamera(Camera):
|
||||
if not self.is_connected and self.thread is None:
|
||||
raise DeviceNotConnectedError(f"{self} not connected.")
|
||||
|
||||
self._cleanup_resources()
|
||||
if self.thread is not None:
|
||||
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.")
|
||||
|
||||
@@ -121,9 +121,6 @@ class RealSenseCamera(Camera):
|
||||
|
||||
self.config = config
|
||||
|
||||
self.width: int | None = config.width
|
||||
self.height: int | None = config.height
|
||||
|
||||
if config.serial_number_or_name.isdigit():
|
||||
self.serial_number = config.serial_number_or_name
|
||||
else:
|
||||
@@ -134,9 +131,6 @@ class RealSenseCamera(Camera):
|
||||
self.use_rgb = config.use_rgb
|
||||
self.use_depth = config.use_depth
|
||||
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_profile: rs.pipeline_profile | None = None
|
||||
@@ -151,23 +145,14 @@ class RealSenseCamera(Camera):
|
||||
|
||||
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
||||
|
||||
self.capture_width: int | None = None
|
||||
self.capture_height: int | None = None
|
||||
self._reset_connection_settings()
|
||||
|
||||
def __str__(self) -> str:
|
||||
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
|
||||
if self.height and self.width:
|
||||
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
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}({self.serial_number})"
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Checks if the camera pipeline is started and streams are active."""
|
||||
@@ -187,8 +172,7 @@ class RealSenseCamera(Camera):
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the camera is already connected.
|
||||
ValueError: If the configuration is invalid, a requested sensor option is unsupported,
|
||||
or a requested sensor value is invalid.
|
||||
ValueError: If the configuration is invalid (e.g., missing serial/name, name not unique).
|
||||
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.
|
||||
"""
|
||||
@@ -206,9 +190,7 @@ class RealSenseCamera(Camera):
|
||||
f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
|
||||
) from e
|
||||
|
||||
try:
|
||||
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.
|
||||
@@ -224,13 +206,6 @@ class RealSenseCamera(Camera):
|
||||
self.use_depth and self.latest_depth_frame is None
|
||||
):
|
||||
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.")
|
||||
|
||||
@@ -364,114 +339,6 @@ class RealSenseCamera(Camera):
|
||||
self.new_frame_event.clear()
|
||||
return self._async_read(timeout_ms=10000, read_depth=read_depth)
|
||||
|
||||
def _get_color_sensor(self) -> "rs.sensor":
|
||||
"""Returns the dedicated "RGB Camera" sensor that controls the color stream.
|
||||
|
||||
Manual color controls are only applied to a dedicated RGB module. Cameras
|
||||
without one (e.g. the D405, whose color stream comes from the shared
|
||||
"Stereo Module") are unsupported, so we never fall back to another sensor
|
||||
to avoid altering the depth stream.
|
||||
"""
|
||||
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()}
|
||||
|
||||
if "RGB Camera" in sensors:
|
||||
return sensors["RGB Camera"]
|
||||
|
||||
available = list(sensors.keys())
|
||||
raise RuntimeError(
|
||||
f"{self}: manual color controls require a dedicated 'RGB Camera' module, which this camera does not have. ",
|
||||
f"Available sensors: {available}.",
|
||||
)
|
||||
|
||||
def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
|
||||
"""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
|
||||
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
|
||||
"""
|
||||
@@ -674,27 +541,6 @@ class RealSenseCamera(Camera):
|
||||
self.latest_timestamp = None
|
||||
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]:
|
||||
"""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():
|
||||
@@ -838,5 +684,18 @@ class RealSenseCamera(Camera):
|
||||
f"Attempted to disconnect {self}, but it appears already disconnected."
|
||||
)
|
||||
|
||||
self._cleanup_resources()
|
||||
if self.thread is not None:
|
||||
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.")
|
||||
|
||||
@@ -46,17 +46,6 @@ class RealSenseCameraConfig(CameraConfig):
|
||||
use_depth: Whether to enable depth stream. Defaults to False.
|
||||
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
|
||||
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:
|
||||
- Either name or serial_number must be specified.
|
||||
@@ -72,9 +61,6 @@ class RealSenseCameraConfig(CameraConfig):
|
||||
use_depth: bool = False
|
||||
rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
|
||||
warmup_s: int = 1
|
||||
exposure: int | None = None
|
||||
gain: int | None = None
|
||||
white_balance: int | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.color_mode = ColorMode(self.color_mode)
|
||||
@@ -83,18 +69,6 @@ class RealSenseCameraConfig(CameraConfig):
|
||||
if not self.use_rgb and not self.use_depth:
|
||||
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)
|
||||
if any(v is not None for v in values) and any(v is None for v in values):
|
||||
raise ValueError(
|
||||
|
||||
@@ -35,9 +35,9 @@ else:
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.datasets import LeRobotDataset
|
||||
from lerobot.lerobot_types import PolicyAction
|
||||
from lerobot.processor import PolicyProcessorPipeline
|
||||
from lerobot.robots import Robot
|
||||
from lerobot.types import PolicyAction
|
||||
|
||||
|
||||
def predict_action(
|
||||
|
||||
@@ -53,7 +53,11 @@ def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Pa
|
||||
|
||||
|
||||
def save_training_step(
|
||||
step: int, save_dir: Path, num_processes: int | None = None, batch_size: int | None = None
|
||||
step: int,
|
||||
save_dir: Path,
|
||||
num_processes: int | None = None,
|
||||
batch_size: int | None = None,
|
||||
num_workers: int | None = None,
|
||||
) -> None:
|
||||
state: dict = {"step": step}
|
||||
# num_processes and batch_size are recorded so a resumed run can detect a changed world size or
|
||||
@@ -64,6 +68,8 @@ def save_training_step(
|
||||
state["num_processes"] = num_processes
|
||||
if batch_size is not None:
|
||||
state["batch_size"] = batch_size
|
||||
if num_workers is not None:
|
||||
state["num_workers"] = num_workers
|
||||
write_json(state, save_dir / TRAINING_STEP)
|
||||
|
||||
|
||||
@@ -82,6 +88,11 @@ def load_training_batch_size(checkpoint_dir: Path) -> int | None:
|
||||
return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("batch_size")
|
||||
|
||||
|
||||
def load_training_num_workers(checkpoint_dir: Path) -> int | None:
|
||||
"""DataLoader worker count recorded at checkpoint time, or None for older checkpoints."""
|
||||
return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("num_workers")
|
||||
|
||||
|
||||
def update_last_checkpoint(checkpoint_dir: Path) -> Path:
|
||||
last_checkpoint_dir = checkpoint_dir.parent / LAST_CHECKPOINT_LINK
|
||||
if last_checkpoint_dir.is_symlink():
|
||||
@@ -101,6 +112,7 @@ def save_checkpoint(
|
||||
postprocessor: PolicyProcessorPipeline | None = None,
|
||||
num_processes: int | None = None,
|
||||
batch_size: int | None = None,
|
||||
num_workers: int | None = None,
|
||||
model_state_dict: dict | None = None,
|
||||
optim_state_dict: dict | None = None,
|
||||
) -> None:
|
||||
@@ -132,6 +144,8 @@ def save_checkpoint(
|
||||
resume. Defaults to None (not recorded).
|
||||
batch_size (int | None, optional): Per-process batch size to record for sample-exact
|
||||
resume. Defaults to None (not recorded).
|
||||
num_workers (int | None, optional): Per-process DataLoader worker count to record for
|
||||
sample-exact streaming resume. Defaults to None (not recorded).
|
||||
model_state_dict: Pre-gathered full (unsharded) model state dict. Required under FSDP,
|
||||
where `policy.state_dict()` would return sharded tensors; the caller gathers it via a
|
||||
cross-rank collective and passes it here so rank 0 can write it directly. It holds
|
||||
@@ -160,6 +174,7 @@ def save_checkpoint(
|
||||
scheduler,
|
||||
num_processes=num_processes,
|
||||
batch_size=batch_size,
|
||||
num_workers=num_workers,
|
||||
optim_state_dict=optim_state_dict,
|
||||
)
|
||||
|
||||
@@ -171,6 +186,7 @@ def save_training_state(
|
||||
scheduler: LRScheduler | None = None,
|
||||
num_processes: int | None = None,
|
||||
batch_size: int | None = None,
|
||||
num_workers: int | None = None,
|
||||
optim_state_dict: dict | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
@@ -185,12 +201,20 @@ def save_training_state(
|
||||
Defaults to None.
|
||||
num_processes (int | None, optional): Distributed world size to record. Defaults to None.
|
||||
batch_size (int | None, optional): Per-process batch size to record. Defaults to None.
|
||||
num_workers (int | None, optional): Per-process DataLoader worker count to record.
|
||||
Defaults to None.
|
||||
optim_state_dict: Pre-gathered full optimizer state dict (for FSDP). Saved instead of
|
||||
`optimizer.state_dict()` when provided. Defaults to None.
|
||||
"""
|
||||
save_dir = checkpoint_dir / TRAINING_STATE_DIR
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
save_training_step(train_step, save_dir, num_processes=num_processes, batch_size=batch_size)
|
||||
save_training_step(
|
||||
train_step,
|
||||
save_dir,
|
||||
num_processes=num_processes,
|
||||
batch_size=batch_size,
|
||||
num_workers=num_workers,
|
||||
)
|
||||
save_rng_state(save_dir)
|
||||
if optimizer is not None:
|
||||
save_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict)
|
||||
|
||||
@@ -71,19 +71,13 @@ class DatasetRecordConfig:
|
||||
# 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..
|
||||
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:
|
||||
"""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,
|
||||
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:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
self.repo_id = f"{self.repo_id}_{timestamp}"
|
||||
|
||||
@@ -44,6 +44,22 @@ class DatasetConfig:
|
||||
# Has no effect on datasets without depth cameras.
|
||||
depth_output_unit: str = DEFAULT_DEPTH_UNIT
|
||||
streaming: bool = False
|
||||
# Optional data-plane root used by training-time streaming. Metadata still resolves from repo_id/root.
|
||||
streaming_data_root: str | None = None
|
||||
# Number of complete episodes mixed by the rank-level exact-coverage sampler.
|
||||
streaming_episode_pool_size: int = 32
|
||||
# Complete episodes fetched ahead of the current admission frontier.
|
||||
streaming_prefetch_episodes: int = 8
|
||||
# Hard per-rank cap for synthesized episode-video bytes.
|
||||
streaming_byte_budget_gb: float = 8.0
|
||||
# Parallel sample assembly/decode workers and their bounded in-order result queue.
|
||||
streaming_decode_threads: int = 2
|
||||
streaming_decoded_queue_size: int = 8
|
||||
# Independent decoder-state cap. None covers every camera in the configured episode pool.
|
||||
streaming_max_open_decoders: int | None = None
|
||||
# Per-rank native HTTP limits. None preserves the fetcher's worker-derived default.
|
||||
streaming_native_http_connections: int | None = None
|
||||
streaming_native_http_subranges: int = 1
|
||||
# Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
|
||||
eval_split: float = 0.0
|
||||
|
||||
@@ -54,6 +70,22 @@ class DatasetConfig:
|
||||
)
|
||||
if not (0.0 <= self.eval_split < 1.0):
|
||||
raise ValueError(f"eval_split must be in [0.0, 1.0), got {self.eval_split}")
|
||||
if self.streaming_episode_pool_size <= 0:
|
||||
raise ValueError("streaming_episode_pool_size must be positive")
|
||||
if self.streaming_prefetch_episodes < 0:
|
||||
raise ValueError("streaming_prefetch_episodes must be non-negative")
|
||||
if self.streaming_byte_budget_gb <= 0:
|
||||
raise ValueError("streaming_byte_budget_gb must be positive")
|
||||
if self.streaming_decode_threads <= 0:
|
||||
raise ValueError("streaming_decode_threads must be positive")
|
||||
if self.streaming_decoded_queue_size <= 0:
|
||||
raise ValueError("streaming_decoded_queue_size must be positive")
|
||||
if self.streaming_max_open_decoders is not None and self.streaming_max_open_decoders <= 0:
|
||||
raise ValueError("streaming_max_open_decoders must be positive")
|
||||
if self.streaming_native_http_connections is not None and self.streaming_native_http_connections <= 0:
|
||||
raise ValueError("streaming_native_http_connections must be positive")
|
||||
if self.streaming_native_http_subranges <= 0:
|
||||
raise ValueError("streaming_native_http_subranges must be positive")
|
||||
if self.episodes is not None:
|
||||
if any(ep < 0 for ep in self.episodes):
|
||||
raise ValueError(
|
||||
|
||||
@@ -163,10 +163,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
|
||||
return None
|
||||
|
||||
def _save_pretrained(self, save_directory: Path) -> None:
|
||||
# Encode against the base class so draccus includes the choice "type" key,
|
||||
# which `from_pretrained` needs to resolve the concrete subclass.
|
||||
with open(save_directory / CONFIG_NAME, "w") as f:
|
||||
json.dump(draccus.encode(self, PreTrainedConfig), f, indent=4)
|
||||
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
|
||||
draccus.dump(self, f, indent=4)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
|
||||
@@ -103,10 +103,8 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
|
||||
pass
|
||||
|
||||
def _save_pretrained(self, save_directory: Path) -> None:
|
||||
# Encode against the base class so draccus includes the choice "type" key,
|
||||
# which `from_pretrained` needs to resolve the concrete subclass.
|
||||
with open(save_directory / CONFIG_NAME, "w") as f:
|
||||
json.dump(draccus.encode(self, RewardModelConfig), f, indent=4)
|
||||
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
|
||||
draccus.dump(self, f, indent=4)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
|
||||
@@ -194,11 +194,7 @@ class TrainPipelineConfig(HubMixin):
|
||||
)
|
||||
|
||||
if Path(config_path).resolve().exists():
|
||||
# `config_path` may point at the checkpoint's train_config.json or at its
|
||||
# pretrained_model/ directory (both documented above) — resolve either to
|
||||
# the pretrained_model/ directory.
|
||||
config_path_obj = Path(config_path)
|
||||
policy_dir = config_path_obj.parent if config_path_obj.is_file() else config_path_obj
|
||||
policy_dir = Path(config_path).parent
|
||||
self.checkpoint_path = policy_dir.parent
|
||||
elif self.job.is_remote:
|
||||
return
|
||||
|
||||
@@ -19,7 +19,6 @@ import copy
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Any, NotRequired, TypedDict
|
||||
|
||||
import datasets
|
||||
import pandas as pd
|
||||
@@ -50,32 +49,8 @@ from .utils import (
|
||||
)
|
||||
from .video_utils import concatenate_video_files, get_video_duration_in_s
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
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:
|
||||
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]:
|
||||
"""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:
|
||||
@@ -84,14 +59,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
||||
Returns:
|
||||
dict: A dictionary of merged video feature info.
|
||||
"""
|
||||
merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
|
||||
merged_info = copy.deepcopy(all_metadata[0].features)
|
||||
video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
|
||||
|
||||
for vk in video_keys:
|
||||
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
|
||||
base_video_info = video_infos[0]
|
||||
|
||||
merged_encoder_info: dict[str, Any] = {}
|
||||
merged_encoder_info: dict = {}
|
||||
fallback_keys: list[str] = []
|
||||
for info_key in VIDEO_ENCODER_INFO_KEYS:
|
||||
values = [info.get(info_key, None) for info in video_infos]
|
||||
@@ -105,7 +80,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
|
||||
|
||||
if fallback_keys:
|
||||
logger.warning(
|
||||
logging.warning(
|
||||
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
|
||||
f"Setting these keys to null: {fallback_keys}.",
|
||||
)
|
||||
@@ -117,7 +92,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
||||
return merged_info
|
||||
|
||||
|
||||
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
|
||||
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
|
||||
"""Validates that all dataset metadata have consistent properties.
|
||||
|
||||
Ensures all datasets have the same fps, robot_type, and features to guarantee
|
||||
@@ -154,9 +129,7 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[i
|
||||
return fps, robot_type, features
|
||||
|
||||
|
||||
def update_data_df(
|
||||
df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
|
||||
) -> pd.DataFrame:
|
||||
def update_data_df(df, src_meta, dst_meta):
|
||||
"""Updates a data DataFrame with new indices and task mappings for aggregation.
|
||||
|
||||
Adjusts episode indices, frame indices, and task indices to account for
|
||||
@@ -181,12 +154,12 @@ def update_data_df(
|
||||
|
||||
|
||||
def update_meta_data(
|
||||
df: pd.DataFrame,
|
||||
dst_meta: LeRobotDatasetMetadata,
|
||||
meta_idx: IndexState,
|
||||
data_idx: IndexState,
|
||||
videos_idx: VideoIndexState,
|
||||
) -> pd.DataFrame:
|
||||
df,
|
||||
dst_meta,
|
||||
meta_idx,
|
||||
data_idx,
|
||||
videos_idx,
|
||||
):
|
||||
"""Updates metadata DataFrame with new chunk, file, and timestamp indices.
|
||||
|
||||
Adjusts all indices and timestamps to account for previously aggregated
|
||||
@@ -316,7 +289,7 @@ def aggregate_datasets(
|
||||
chunk_size: int | None = None,
|
||||
concatenate_videos: bool = True,
|
||||
concatenate_data: bool = True,
|
||||
) -> None:
|
||||
):
|
||||
"""Aggregates multiple LeRobot datasets into a single unified dataset.
|
||||
|
||||
This is the main function that orchestrates the aggregation process by:
|
||||
@@ -336,7 +309,7 @@ def aggregate_datasets(
|
||||
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
|
||||
concatenate_data: When False, keep one parquet per source file instead of packing into shards.
|
||||
"""
|
||||
logger.info("Start aggregate_datasets")
|
||||
logging.info("Start aggregate_datasets")
|
||||
|
||||
if data_files_size_in_mb is None:
|
||||
data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
|
||||
@@ -368,15 +341,15 @@ def aggregate_datasets(
|
||||
video_files_size_in_mb=video_files_size_in_mb,
|
||||
)
|
||||
|
||||
logger.info("Find all tasks")
|
||||
logging.info("Find all tasks")
|
||||
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
|
||||
dst_meta.tasks = pd.DataFrame(
|
||||
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
|
||||
)
|
||||
|
||||
meta_idx: IndexState = {"chunk": 0, "file": 0}
|
||||
data_idx: IndexState = {"chunk": 0, "file": 0}
|
||||
videos_idx: VideoIndexState = {
|
||||
meta_idx = {"chunk": 0, "file": 0}
|
||||
data_idx = {"chunk": 0, "file": 0}
|
||||
videos_idx = {
|
||||
key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
|
||||
}
|
||||
|
||||
@@ -400,17 +373,12 @@ def aggregate_datasets(
|
||||
dst_meta.info.total_frames += src_meta.total_frames
|
||||
|
||||
finalize_aggregation(dst_meta, all_metadata)
|
||||
logger.info("Aggregation complete.")
|
||||
logging.info("Aggregation complete.")
|
||||
|
||||
|
||||
def aggregate_videos(
|
||||
src_meta: LeRobotDatasetMetadata,
|
||||
dst_meta: LeRobotDatasetMetadata,
|
||||
videos_idx: VideoIndexState,
|
||||
video_files_size_in_mb: float,
|
||||
chunk_size: int,
|
||||
concatenate_videos: bool = True,
|
||||
) -> VideoIndexState:
|
||||
src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
|
||||
):
|
||||
"""Aggregates video chunks from a source dataset into the destination dataset.
|
||||
|
||||
Handles video file concatenation and rotation based on file size limits.
|
||||
@@ -438,8 +406,7 @@ def aggregate_videos(
|
||||
videos_idx[key]["dst_file_durations"] = {}
|
||||
|
||||
for key, video_idx in videos_idx.items():
|
||||
unique_chunk_file_pairs: list[ChunkFile] = sorted(
|
||||
{
|
||||
unique_chunk_file_pairs = {
|
||||
(chunk, file)
|
||||
for chunk, file in zip(
|
||||
src_meta.episodes[f"videos/{key}/chunk_index"],
|
||||
@@ -447,7 +414,7 @@ def aggregate_videos(
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
)
|
||||
unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
|
||||
|
||||
chunk_idx = video_idx["chunk"]
|
||||
file_idx = video_idx["file"]
|
||||
@@ -522,14 +489,7 @@ def aggregate_videos(
|
||||
return videos_idx
|
||||
|
||||
|
||||
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:
|
||||
def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True):
|
||||
"""Aggregates data chunks from a source dataset into the destination dataset.
|
||||
|
||||
Reads source data files, updates indices to match the aggregated dataset,
|
||||
@@ -550,16 +510,14 @@ def aggregate_data(
|
||||
Returns:
|
||||
dict: Updated data_idx with current chunk and file indices.
|
||||
"""
|
||||
unique_chunk_file_ids: list[ChunkFile] = sorted(
|
||||
{
|
||||
unique_chunk_file_ids = {
|
||||
(c, f)
|
||||
for c, f in zip(
|
||||
src_meta.episodes["data/chunk_index"],
|
||||
src_meta.episodes["data/file_index"],
|
||||
strict=False,
|
||||
src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
unique_chunk_file_ids = sorted(unique_chunk_file_ids)
|
||||
contains_images = len(dst_meta.image_keys) > 0
|
||||
|
||||
# retrieve features schema for proper image typing in parquet
|
||||
@@ -567,7 +525,7 @@ def aggregate_data(
|
||||
|
||||
# Track source to destination file mapping for metadata update
|
||||
# This is critical for handling datasets that are already results of a merge
|
||||
src_to_dst: dict[ChunkFile, ChunkFile] = {}
|
||||
src_to_dst: dict[tuple[int, int], tuple[int, int]] = {}
|
||||
|
||||
for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
|
||||
src_path = src_meta.root / DEFAULT_DATA_PATH.format(
|
||||
@@ -606,13 +564,7 @@ def aggregate_data(
|
||||
return data_idx
|
||||
|
||||
|
||||
def aggregate_metadata(
|
||||
src_meta: LeRobotDatasetMetadata,
|
||||
dst_meta: LeRobotDatasetMetadata,
|
||||
meta_idx: IndexState,
|
||||
data_idx: IndexState,
|
||||
videos_idx: VideoIndexState,
|
||||
) -> IndexState:
|
||||
def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
|
||||
"""Aggregates metadata from a source dataset into the destination dataset.
|
||||
|
||||
Reads source metadata files, updates all indices and timestamps,
|
||||
@@ -628,8 +580,7 @@ def aggregate_metadata(
|
||||
Returns:
|
||||
dict: Updated meta_idx with current chunk and file indices.
|
||||
"""
|
||||
chunk_file_ids: list[ChunkFile] = sorted(
|
||||
{
|
||||
chunk_file_ids = {
|
||||
(c, f)
|
||||
for c, f in zip(
|
||||
src_meta.episodes["meta/episodes/chunk_index"],
|
||||
@@ -637,7 +588,8 @@ def aggregate_metadata(
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
chunk_file_ids = sorted(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)
|
||||
df = pd.read_parquet(src_path)
|
||||
@@ -670,16 +622,16 @@ def aggregate_metadata(
|
||||
def append_or_create_parquet_file(
|
||||
df: pd.DataFrame,
|
||||
src_path: Path,
|
||||
idx: IndexState,
|
||||
idx: dict[str, int],
|
||||
max_mb: float,
|
||||
chunk_size: int,
|
||||
default_path: str,
|
||||
contains_images: bool = False,
|
||||
aggr_root: Path | None = None,
|
||||
aggr_root: Path = None,
|
||||
hf_features: datasets.Features | None = None,
|
||||
concatenate: bool = True,
|
||||
one_row_group_per_episode: bool = False,
|
||||
) -> tuple[IndexState, ChunkFile]:
|
||||
) -> tuple[dict[str, int], tuple[int, int]]:
|
||||
"""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
|
||||
@@ -702,13 +654,7 @@ def append_or_create_parquet_file(
|
||||
Returns:
|
||||
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.
|
||||
|
||||
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_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
|
||||
|
||||
@@ -752,9 +698,7 @@ def append_or_create_parquet_file(
|
||||
return idx, (dst_chunk, dst_file)
|
||||
|
||||
|
||||
def finalize_aggregation(
|
||||
aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
|
||||
) -> None:
|
||||
def finalize_aggregation(aggr_meta, all_metadata):
|
||||
"""Finalizes the dataset aggregation by writing summary files and statistics.
|
||||
|
||||
Writes the tasks file, info file with total counts and splits, and
|
||||
@@ -764,16 +708,16 @@ def finalize_aggregation(
|
||||
aggr_meta: Aggregated dataset metadata.
|
||||
all_metadata: List of all source dataset metadata objects.
|
||||
"""
|
||||
logger.info("write tasks")
|
||||
logging.info("write tasks")
|
||||
write_tasks(aggr_meta.tasks, aggr_meta.root)
|
||||
|
||||
logger.info("write info")
|
||||
logging.info("write info")
|
||||
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_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)}"}
|
||||
write_info(aggr_meta.info, aggr_meta.root)
|
||||
|
||||
logger.info("write stats")
|
||||
logging.info("write stats")
|
||||
aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
|
||||
write_stats(aggr_meta.stats, aggr_meta.root)
|
||||
|
||||
@@ -188,8 +188,8 @@ class LeRobotDatasetMetadata:
|
||||
def _load_metadata(self):
|
||||
self.info = load_info(self.root)
|
||||
check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
|
||||
self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
|
||||
self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
|
||||
self.tasks = load_tasks(self.root)
|
||||
self.episodes = load_episodes(self.root)
|
||||
self.stats = load_stats(self.root)
|
||||
|
||||
def ensure_readable(self) -> None:
|
||||
|
||||
@@ -1045,12 +1045,10 @@ def _copy_data_with_feature_changes(
|
||||
df[feature_name] = feature_values
|
||||
else:
|
||||
feature_slice = values[frame_idx:end_idx]
|
||||
if feature_slice.ndim == 1:
|
||||
df[feature_name] = feature_slice
|
||||
elif feature_slice.ndim == 2 and feature_slice.shape[1] == 1:
|
||||
if len(feature_slice.shape) > 1 and feature_slice.shape[1] == 1:
|
||||
df[feature_name] = feature_slice.flatten()
|
||||
else:
|
||||
df[feature_name] = list(feature_slice)
|
||||
df[feature_name] = feature_slice
|
||||
frame_idx = end_idx
|
||||
|
||||
# Write using the same chunk/file structure as source
|
||||
|
||||
@@ -66,7 +66,9 @@ def resolve_delta_timestamps(
|
||||
return delta_timestamps
|
||||
|
||||
|
||||
def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
|
||||
def make_dataset(
|
||||
cfg: TrainPipelineConfig,
|
||||
) -> LeRobotDataset | StreamingLeRobotDataset | MultiLeRobotDataset:
|
||||
"""Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
|
||||
|
||||
Args:
|
||||
@@ -108,9 +110,21 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
|
||||
delta_timestamps=delta_timestamps,
|
||||
image_transforms=image_transforms,
|
||||
revision=cfg.dataset.revision,
|
||||
max_num_shards=cfg.num_workers,
|
||||
max_num_shards=max(1, cfg.num_workers),
|
||||
tolerance_s=cfg.tolerance_s,
|
||||
return_uint8=True,
|
||||
depth_output_unit=cfg.dataset.depth_output_unit,
|
||||
video_backend=cfg.dataset.video_backend,
|
||||
data_root=cfg.dataset.streaming_data_root,
|
||||
episode_pool_size=cfg.dataset.streaming_episode_pool_size,
|
||||
prefetch_episodes=cfg.dataset.streaming_prefetch_episodes,
|
||||
byte_budget_gb=cfg.dataset.streaming_byte_budget_gb,
|
||||
decode_threads=cfg.dataset.streaming_decode_threads,
|
||||
decoded_queue_size=cfg.dataset.streaming_decoded_queue_size,
|
||||
max_open_decoders=cfg.dataset.streaming_max_open_decoders,
|
||||
native_http_connections=cfg.dataset.streaming_native_http_connections,
|
||||
native_http_subranges=cfg.dataset.streaming_native_http_subranges,
|
||||
repeat=True,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
|
||||
@@ -138,7 +152,10 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
|
||||
|
||||
def make_train_eval_datasets(
|
||||
cfg: TrainPipelineConfig,
|
||||
) -> tuple[LeRobotDataset | MultiLeRobotDataset, LeRobotDataset | None]:
|
||||
) -> tuple[
|
||||
LeRobotDataset | StreamingLeRobotDataset | MultiLeRobotDataset,
|
||||
LeRobotDataset | None,
|
||||
]:
|
||||
"""Create train and optional eval datasets by splitting episodes based on eval_split.
|
||||
|
||||
The last ceil(n_episodes * eval_split) episodes per task are held out for evaluation.
|
||||
@@ -181,6 +198,31 @@ def make_train_eval_datasets(
|
||||
ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None
|
||||
)
|
||||
|
||||
if cfg.dataset.streaming:
|
||||
train_dataset = StreamingLeRobotDataset(
|
||||
cfg.dataset.repo_id,
|
||||
root=cfg.dataset.root,
|
||||
episodes=train_episodes,
|
||||
delta_timestamps=delta_timestamps,
|
||||
image_transforms=train_image_transforms,
|
||||
revision=cfg.dataset.revision,
|
||||
max_num_shards=max(1, cfg.num_workers),
|
||||
tolerance_s=cfg.tolerance_s,
|
||||
return_uint8=True,
|
||||
depth_output_unit=cfg.dataset.depth_output_unit,
|
||||
video_backend=cfg.dataset.video_backend,
|
||||
data_root=cfg.dataset.streaming_data_root,
|
||||
episode_pool_size=cfg.dataset.streaming_episode_pool_size,
|
||||
prefetch_episodes=cfg.dataset.streaming_prefetch_episodes,
|
||||
byte_budget_gb=cfg.dataset.streaming_byte_budget_gb,
|
||||
decode_threads=cfg.dataset.streaming_decode_threads,
|
||||
decoded_queue_size=cfg.dataset.streaming_decoded_queue_size,
|
||||
max_open_decoders=cfg.dataset.streaming_max_open_decoders,
|
||||
native_http_connections=cfg.dataset.streaming_native_http_connections,
|
||||
native_http_subranges=cfg.dataset.streaming_native_http_subranges,
|
||||
repeat=True,
|
||||
)
|
||||
else:
|
||||
train_dataset = LeRobotDataset(
|
||||
cfg.dataset.repo_id,
|
||||
root=cfg.dataset.root,
|
||||
@@ -190,6 +232,7 @@ def make_train_eval_datasets(
|
||||
revision=cfg.dataset.revision,
|
||||
video_backend=cfg.dataset.video_backend,
|
||||
return_uint8=True,
|
||||
depth_output_unit=cfg.dataset.depth_output_unit,
|
||||
tolerance_s=cfg.tolerance_s,
|
||||
)
|
||||
|
||||
@@ -202,6 +245,7 @@ def make_train_eval_datasets(
|
||||
revision=cfg.dataset.revision,
|
||||
video_backend=cfg.dataset.video_backend,
|
||||
return_uint8=True,
|
||||
depth_output_unit=cfg.dataset.depth_output_unit,
|
||||
tolerance_s=cfg.tolerance_s,
|
||||
)
|
||||
|
||||
|
||||
@@ -17,8 +17,8 @@ from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
from lerobot.configs import PipelineFeatureType
|
||||
from lerobot.lerobot_types import RobotAction, RobotObservation
|
||||
from lerobot.processor import DataProcessorPipeline
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
|
||||
from lerobot.utils.feature_utils import hw_to_dataset_features
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,149 @@
|
||||
# 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
|
||||
|
||||
"""LeRobot metadata adapter for automatic MP4 sidecar resolution."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import fsspec
|
||||
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
from lerobot.streaming.manifest import EpisodeVideoManifest
|
||||
from lerobot.streaming.sidecar import SidecarSpec, ensure_mp4_sidecar, sidecar_cache_path
|
||||
from lerobot.utils.constants import HF_LEROBOT_HOME
|
||||
|
||||
DEFAULT_SIDECAR_CACHE = HF_LEROBOT_HOME / "streaming-sidecars"
|
||||
|
||||
|
||||
def range_backend_for_root(data_root: str) -> str:
|
||||
"""Use direct HTTP only for HF roots; all local and other fsspec protocols stay generic."""
|
||||
return "native-http" if data_root.startswith("hf://") else "fsspec"
|
||||
|
||||
|
||||
def streaming_data_root(
|
||||
meta: LeRobotDatasetMetadata,
|
||||
*,
|
||||
requested_root: str | Path | None,
|
||||
configured_data_root: str | None,
|
||||
) -> str:
|
||||
if configured_data_root is not None:
|
||||
return configured_data_root.rstrip("/")
|
||||
if requested_root is not None:
|
||||
return str(Path(requested_root).expanduser())
|
||||
return f"hf://datasets/{meta.repo_id}@{meta.revision}"
|
||||
|
||||
|
||||
def make_sidecar_spec(meta: LeRobotDatasetMetadata, data_root: str) -> SidecarSpec:
|
||||
relative_paths = sorted(
|
||||
{
|
||||
str(meta.get_video_file_path(episode_index, video_key))
|
||||
for episode_index in range(int(meta.total_episodes))
|
||||
for video_key in meta.video_keys
|
||||
}
|
||||
)
|
||||
root = Path(data_root).expanduser()
|
||||
source_files: tuple[tuple[str, int | None], ...]
|
||||
if root.is_dir():
|
||||
source_files = tuple((path, (root / path).stat().st_size) for path in relative_paths)
|
||||
else:
|
||||
source_files = tuple((path, None) for path in relative_paths)
|
||||
return SidecarSpec(
|
||||
repo_id=meta.repo_id,
|
||||
revision=str(meta.revision),
|
||||
data_root=data_root.rstrip("/"),
|
||||
source_files=source_files,
|
||||
)
|
||||
|
||||
|
||||
def build_mp4_sidecar(
|
||||
destination: str | Path,
|
||||
spec: SidecarSpec,
|
||||
*,
|
||||
workers: int = 8,
|
||||
range_backend: str = "native-http",
|
||||
max_probe_bytes: int = 64 * 1024 * 1024,
|
||||
token: str | bool | None = None,
|
||||
) -> None:
|
||||
EpisodeVideoManifest.write_file_sidecar(
|
||||
destination,
|
||||
[path for path, _size in spec.source_files],
|
||||
spec.data_root,
|
||||
spec=spec,
|
||||
range_backend=range_backend,
|
||||
workers=workers,
|
||||
max_probe_bytes=max_probe_bytes,
|
||||
token=token,
|
||||
)
|
||||
|
||||
|
||||
def published_sidecar_url(spec: SidecarSpec, cache_root: str | Path = DEFAULT_SIDECAR_CACHE) -> str:
|
||||
name = sidecar_cache_path(cache_root, spec).name
|
||||
return f"{spec.data_root}/meta/mp4-sidecars/{name}"
|
||||
|
||||
|
||||
def download_published_sidecar(
|
||||
destination: Path,
|
||||
spec: SidecarSpec,
|
||||
*,
|
||||
cache_root: str | Path = DEFAULT_SIDECAR_CACHE,
|
||||
token: str | bool | None = None,
|
||||
) -> bool:
|
||||
if Path(spec.data_root).expanduser().is_dir():
|
||||
return False
|
||||
source_url = published_sidecar_url(spec, cache_root)
|
||||
storage_options = {"token": token} if token is not None and source_url.startswith("hf://") else {}
|
||||
filesystem, source = fsspec.core.url_to_fs(source_url, **storage_options)
|
||||
if not filesystem.exists(source):
|
||||
return False
|
||||
with filesystem.open(source, "rb") as remote, destination.open("wb") as local:
|
||||
shutil.copyfileobj(remote, local)
|
||||
return True
|
||||
|
||||
|
||||
def ensure_dataset_mp4_sidecar(
|
||||
meta: LeRobotDatasetMetadata,
|
||||
data_root: str,
|
||||
*,
|
||||
cache_root: str | Path = DEFAULT_SIDECAR_CACHE,
|
||||
workers: int = 8,
|
||||
range_backend: str = "native-http",
|
||||
lock_timeout_s: float = 30 * 60,
|
||||
token: str | bool | None = None,
|
||||
) -> Path | None:
|
||||
if not meta.video_keys:
|
||||
return None
|
||||
|
||||
spec = make_sidecar_spec(meta, data_root)
|
||||
logging.info(
|
||||
"Resolving training-time MP4 sidecar for %s@%s (%d files)",
|
||||
spec.repo_id,
|
||||
spec.revision,
|
||||
len(spec.source_files),
|
||||
)
|
||||
return ensure_mp4_sidecar(
|
||||
spec,
|
||||
cache_root,
|
||||
build=lambda path, target_spec: build_mp4_sidecar(
|
||||
path,
|
||||
target_spec,
|
||||
workers=workers,
|
||||
range_backend=range_backend,
|
||||
token=token,
|
||||
),
|
||||
download=lambda path, target_spec: download_published_sidecar(
|
||||
path,
|
||||
target_spec,
|
||||
cache_root=cache_root,
|
||||
token=token,
|
||||
),
|
||||
lock_timeout_s=lock_timeout_s,
|
||||
)
|
||||
@@ -28,7 +28,7 @@ from dataclasses import asdict, dataclass, field
|
||||
from fractions import Fraction
|
||||
from pathlib import Path
|
||||
from threading import Lock
|
||||
from typing import Any, ClassVar
|
||||
from typing import Any, BinaryIO, ClassVar
|
||||
|
||||
import av
|
||||
import fsspec
|
||||
@@ -105,7 +105,7 @@ def decode_video_frames(
|
||||
|
||||
|
||||
def decode_video_frames_pyav(
|
||||
video_path: Path | str,
|
||||
video_path: Path | str | BinaryIO,
|
||||
timestamps: list[float],
|
||||
tolerance_s: float,
|
||||
log_loaded_timestamps: bool = False,
|
||||
@@ -124,7 +124,7 @@ def decode_video_frames_pyav(
|
||||
video can be adjusted at encoding time to trade off decoding speed against file size.
|
||||
|
||||
Args:
|
||||
video_path: Path to the video file.
|
||||
video_path: Path to the video file or a seekable binary file object.
|
||||
timestamps: List of timestamps (in seconds) to extract frames for.
|
||||
tolerance_s: Allowed deviation in seconds between a queried timestamp and the closest
|
||||
decoded frame.
|
||||
@@ -137,7 +137,8 @@ def decode_video_frames_pyav(
|
||||
torch.Tensor of shape (len(timestamps), C, H, W).
|
||||
"""
|
||||
# TODO(rcadene): also load audio stream at the same time
|
||||
video_path = str(video_path)
|
||||
video_source = str(video_path) if isinstance(video_path, (Path, str)) else video_path
|
||||
video_label = str(video_path) if isinstance(video_path, (Path, str)) else "<in-memory video>"
|
||||
|
||||
# set the first and last requested timestamps
|
||||
# Note: previous timestamps are usually loaded, since we need to access the previous key frame
|
||||
@@ -151,7 +152,7 @@ def decode_video_frames_pyav(
|
||||
# av.time_base units (microseconds). `backward=True` lands us on the nearest keyframe at or
|
||||
# before `first_ts`, so we can then decode forward until we cover `last_ts`. See:
|
||||
# https://pyav.basswood-io.com/docs/stable/api/container.html#av.container.InputContainer.seek
|
||||
with av.open(video_path) as container:
|
||||
with av.open(video_source) as container:
|
||||
stream = container.streams.video[0]
|
||||
# Seek to the nearest keyframe at or before `first_ts` with a 1 frame margin
|
||||
container.seek(
|
||||
@@ -180,7 +181,7 @@ def decode_video_frames_pyav(
|
||||
|
||||
if not loaded_frames:
|
||||
raise FrameTimestampError(
|
||||
f"No frames could be decoded from {video_path} in the timestamp range [{first_ts}, {last_ts}]."
|
||||
f"No frames could be decoded from {video_label} in the timestamp range [{first_ts}, {last_ts}]."
|
||||
)
|
||||
|
||||
query_ts = torch.tensor(timestamps)
|
||||
@@ -199,7 +200,7 @@ def decode_video_frames_pyav(
|
||||
" To be safe, we advise to ignore this item during training."
|
||||
f"\nqueried timestamps: {query_ts}"
|
||||
f"\nloaded timestamps: {loaded_ts_t}"
|
||||
f"\nvideo: {video_path}"
|
||||
f"\nvideo: {video_label}"
|
||||
f"\nbackend: pyav"
|
||||
)
|
||||
|
||||
|
||||
@@ -507,7 +507,7 @@ class MetaworldEnv(EnvConfig):
|
||||
class RoboCasaEnv(EnvConfig):
|
||||
task: str = "CloseFridge"
|
||||
fps: int = 20
|
||||
episode_length: int | None = None
|
||||
episode_length: int = 1000
|
||||
obs_type: str = "pixels_agent_pos"
|
||||
render_mode: str = "rgb_array"
|
||||
camera_name: str = "robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right"
|
||||
|
||||
@@ -30,7 +30,7 @@ from gymnasium import spaces
|
||||
from libero.libero import benchmark, get_libero_path
|
||||
from libero.libero.envs import OffScreenRenderEnv
|
||||
|
||||
from lerobot.lerobot_types import RobotObservation
|
||||
from lerobot.types import RobotObservation
|
||||
|
||||
from .utils import _LazyAsyncVectorEnv, parse_camera_names
|
||||
|
||||
@@ -384,12 +384,7 @@ class LiberoEnv(gym.Env):
|
||||
|
||||
def close(self):
|
||||
if self._env is not None:
|
||||
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(
|
||||
|
||||
@@ -25,7 +25,7 @@ import metaworld.policies as policies
|
||||
import numpy as np
|
||||
from gymnasium import spaces
|
||||
|
||||
from lerobot.lerobot_types import RobotObservation
|
||||
from lerobot.types import RobotObservation
|
||||
|
||||
from .utils import _LazyAsyncVectorEnv
|
||||
|
||||
|
||||
@@ -25,7 +25,7 @@ import gymnasium as gym
|
||||
import numpy as np
|
||||
from gymnasium import spaces
|
||||
|
||||
from lerobot.lerobot_types import RobotObservation
|
||||
from lerobot.types import RobotObservation
|
||||
|
||||
from .utils import _LazyAsyncVectorEnv, parse_camera_names
|
||||
|
||||
@@ -98,19 +98,6 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
|
||||
return names, None
|
||||
|
||||
|
||||
def _get_task_horizon(task: str) -> int:
|
||||
"""Return the rollout horizon registered by RoboCasa for a task."""
|
||||
from robocasa.utils.dataset_registry_utils import get_task_horizon
|
||||
|
||||
try:
|
||||
return int(get_task_horizon(task))
|
||||
except ValueError as exc:
|
||||
raise ValueError(
|
||||
f"No RoboCasa horizon is registered for task '{task}'. "
|
||||
"Set `--env.episode_length=<steps>` explicitly."
|
||||
) from exc
|
||||
|
||||
|
||||
def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
|
||||
"""Split a flat (12,) action vector into a RoboCasa action dict.
|
||||
|
||||
@@ -167,7 +154,7 @@ class RoboCasaEnv(gym.Env):
|
||||
|
||||
self.camera_name = parse_camera_names(camera_name)
|
||||
|
||||
self._max_episode_steps = episode_length if episode_length is not None else _get_task_horizon(task)
|
||||
self._max_episode_steps = episode_length if episode_length is not None else 1000
|
||||
|
||||
# Deferred — created on first reset() inside the worker subprocess
|
||||
# to avoid inheriting stale GPU/EGL contexts across fork().
|
||||
|
||||
@@ -28,7 +28,7 @@ import numpy as np
|
||||
import torch
|
||||
from gymnasium import spaces
|
||||
|
||||
from lerobot.lerobot_types import RobotObservation
|
||||
from lerobot.types import RobotObservation
|
||||
from lerobot.utils.import_utils import _scipy_available
|
||||
|
||||
from .utils import _LazyAsyncVectorEnv
|
||||
@@ -384,9 +384,7 @@ class RoboTwinEnv(gym.Env):
|
||||
|
||||
self._env: Any | None = None # deferred — created on first reset() inside worker
|
||||
self._step_count: int = 0
|
||||
self._black_frame: np.ndarray = np.zeros(
|
||||
(self.observation_height, self.observation_width, 3), dtype=np.uint8
|
||||
)
|
||||
self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
|
||||
|
||||
image_spaces = {
|
||||
cam: spaces.Box(
|
||||
|
||||
@@ -37,7 +37,7 @@ import numpy as np
|
||||
from gymnasium import spaces
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
from lerobot.lerobot_types import RobotObservation
|
||||
from lerobot.types import RobotObservation
|
||||
|
||||
from .utils import _LazyAsyncVectorEnv
|
||||
|
||||
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
|
||||
|
||||
if action.shape[0] != 7:
|
||||
# Unknown layout — fall back to zero-pad so the sim doesn't crash.
|
||||
padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
|
||||
padded = np.zeros(ctrl_dim, dtype=np.float64)
|
||||
padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
|
||||
return padded
|
||||
|
||||
|
||||
@@ -122,9 +122,6 @@ MODEL_ENCODING_TABLE = {
|
||||
"xm430-w350": X_SERIES_ENCODINGS_TABLE,
|
||||
"xm540-w270": 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}
|
||||
@@ -137,9 +134,6 @@ MODEL_RESOLUTION = {
|
||||
"xm430-w350": 4096,
|
||||
"xm540-w270": 4096,
|
||||
"xc430-w150": 4096,
|
||||
"xh540-w150": 4096,
|
||||
"xc330-t288": 4096,
|
||||
"xc330-t181": 4096,
|
||||
}
|
||||
|
||||
# {model: model_number}
|
||||
@@ -151,9 +145,6 @@ MODEL_NUMBER_TABLE = {
|
||||
"xm430-w350": 1020,
|
||||
"xm540-w270": 1120,
|
||||
"xc430-w150": 1070,
|
||||
"xh540-w150": 1110,
|
||||
"xc330-t288": 1220,
|
||||
"xc330-t181": 1210,
|
||||
}
|
||||
|
||||
# {model: available_operating_modes}
|
||||
@@ -165,9 +156,6 @@ MODEL_OPERATING_MODES = {
|
||||
"xm430-w350": [0, 1, 3, 4, 5, 16],
|
||||
"xm540-w270": [0, 1, 3, 4, 5, 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 = {
|
||||
@@ -178,9 +166,6 @@ MODEL_CONTROL_TABLE = {
|
||||
"xm430-w350": X_SERIES_CONTROL_TABLE,
|
||||
"xm540-w270": 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 = {
|
||||
@@ -191,9 +176,6 @@ MODEL_BAUDRATE_TABLE = {
|
||||
"xm430-w350": X_SERIES_BAUDRATE_TABLE,
|
||||
"xm540-w270": 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 = [
|
||||
|
||||
@@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, Any
|
||||
import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import TransitionKey
|
||||
from lerobot.processor import (
|
||||
ComplementaryDataProcessorStep,
|
||||
PolicyAction,
|
||||
@@ -32,6 +31,7 @@ from lerobot.processor import (
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.types import TransitionKey
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
from lerobot.utils.import_utils import _transformers_available, require_package
|
||||
|
||||
|
||||
@@ -42,9 +42,6 @@ class Evo1Policy(PreTrainedPolicy):
|
||||
config_class = Evo1Config
|
||||
name = "evo1"
|
||||
|
||||
def supports_rtc(self) -> bool:
|
||||
return True
|
||||
|
||||
def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
|
||||
super().__init__(config)
|
||||
config.validate_features()
|
||||
|
||||
@@ -21,7 +21,6 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, TransitionKey
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
@@ -41,6 +40,7 @@ from lerobot.processor.converters import (
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
ACTION,
|
||||
DONE,
|
||||
@@ -302,33 +302,6 @@ def _pad_evo1_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(
|
||||
config: Evo1Config,
|
||||
preprocessor: PolicyProcessorPipeline,
|
||||
@@ -336,19 +309,16 @@ def reconcile_evo1_processors(
|
||||
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
|
||||
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config.
|
||||
|
||||
Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
|
||||
(converters are plain functions and are never serialized), eval-time CLI overrides of the
|
||||
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
|
||||
(un)normalizer stats/features when the generic override path injects raw, unpadded dataset
|
||||
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.
|
||||
Two 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
|
||||
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
|
||||
restores the converter 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
|
||||
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
|
||||
preprocessor.to_transition = evo1_batch_to_transition
|
||||
|
||||
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
|
||||
|
||||
action_step = Evo1ActionProcessorStep(
|
||||
action_dim=_evo1_action_dim(config),
|
||||
binarize_gripper=config.binarize_gripper,
|
||||
|
||||
@@ -28,7 +28,6 @@ if TYPE_CHECKING:
|
||||
|
||||
from lerobot.configs import FeatureType, PreTrainedConfig
|
||||
from lerobot.envs import EnvConfig, env_to_policy_features
|
||||
from lerobot.lerobot_types import PolicyAction
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
PolicyProcessorPipeline,
|
||||
@@ -38,25 +37,19 @@ from lerobot.processor import (
|
||||
transition_to_batch,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.types import PolicyAction
|
||||
from lerobot.utils.constants import (
|
||||
ACTION,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
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 .groot.configuration_groot import GrootConfig
|
||||
from .pretrained import PreTrainedPolicy
|
||||
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(
|
||||
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
|
||||
@@ -341,15 +334,12 @@ def make_policy(
|
||||
# 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
|
||||
# adapter config first, then load the correct policy and then apply PEFT.
|
||||
require_package("peft", extra="peft")
|
||||
from peft import PeftConfig, PeftModel
|
||||
|
||||
logging.info("Loading policy's PEFT adapter.")
|
||||
|
||||
peft_pretrained_path = str(cfg.pretrained_path)
|
||||
peft_config = PeftConfig.from_pretrained(
|
||||
peft_pretrained_path,
|
||||
revision=cfg.pretrained_revision,
|
||||
)
|
||||
peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
|
||||
|
||||
kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
|
||||
if not kwargs["pretrained_name_or_path"]:
|
||||
@@ -360,14 +350,9 @@ def make_policy(
|
||||
"the adapter was trained."
|
||||
)
|
||||
|
||||
kwargs["revision"] = peft_config.revision
|
||||
policy = policy_cls.from_pretrained(**kwargs)
|
||||
policy = PeftModel.from_pretrained(
|
||||
policy,
|
||||
peft_pretrained_path,
|
||||
config=peft_config,
|
||||
revision=cfg.pretrained_revision,
|
||||
is_trainable=True,
|
||||
policy, peft_pretrained_path, config=peft_config, is_trainable=True
|
||||
)
|
||||
|
||||
else:
|
||||
|
||||
@@ -37,19 +37,13 @@ def is_image_feature(key: str) -> bool:
|
||||
@dataclass
|
||||
class ConcurrencyConfig:
|
||||
"""Configuration for the concurrency of the actor and learner.
|
||||
|
||||
Possible values are:
|
||||
- "threads": Use threads 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"
|
||||
learner: str = "threads"
|
||||
multiprocessing_context: str | None = "spawn"
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -68,9 +68,6 @@ class GrootPolicy(PreTrainedPolicy):
|
||||
name = "groot"
|
||||
config_class = GrootConfig
|
||||
|
||||
def supports_rtc(self) -> bool:
|
||||
return True
|
||||
|
||||
def __init__(self, config: GrootConfig, **kwargs):
|
||||
"""Initialize Groot policy wrapper."""
|
||||
require_package("transformers", extra="groot")
|
||||
|
||||
@@ -50,7 +50,6 @@ if TYPE_CHECKING or _datasets_available:
|
||||
else:
|
||||
LeRobotDataset = None
|
||||
|
||||
from lerobot.lerobot_types import EnvTransition, TransitionKey
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
@@ -67,6 +66,7 @@ from lerobot.processor import (
|
||||
transition_to_batch,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
ACTION,
|
||||
OBS_IMAGE,
|
||||
|
||||
@@ -43,22 +43,11 @@ from torch.distributions import Beta
|
||||
|
||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||
from lerobot.utils.constants import ACTION
|
||||
from lerobot.utils.import_utils import (
|
||||
_peft_available,
|
||||
_scipy_available,
|
||||
_transformers_available,
|
||||
require_package,
|
||||
)
|
||||
from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
|
||||
|
||||
from ..rtc.modeling_rtc import RTCProcessor
|
||||
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__)
|
||||
|
||||
|
||||
@@ -520,9 +509,6 @@ class MolmoAct2Policy(PreTrainedPolicy):
|
||||
config_class = MolmoAct2Config
|
||||
name = "molmoact2"
|
||||
|
||||
def supports_rtc(self) -> bool:
|
||||
return self.config.inference_action_mode == "continuous"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: MolmoAct2Config,
|
||||
@@ -1745,11 +1731,13 @@ class MolmoAct2Policy(PreTrainedPolicy):
|
||||
|
||||
def _build_inner_lora_config(self):
|
||||
require_package("peft", extra="molmoact2")
|
||||
from peft import LoraConfig
|
||||
|
||||
return LoraConfig(**self._get_inner_peft_targets())
|
||||
|
||||
def _apply_lora_adapters(self) -> None:
|
||||
require_package("peft", extra="molmoact2")
|
||||
from peft import get_peft_model
|
||||
|
||||
peft_config = self._build_inner_lora_config()
|
||||
self._validate_peft_config(peft_config)
|
||||
|
||||
@@ -36,7 +36,6 @@ import torch
|
||||
from torch import Tensor
|
||||
|
||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, TransitionKey
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
@@ -50,6 +49,7 @@ from lerobot.processor import (
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
ACTION,
|
||||
OBS_IMAGES,
|
||||
|
||||
@@ -749,9 +749,6 @@ class PI0Policy(PreTrainedPolicy):
|
||||
config_class = PI0Config
|
||||
name = "pi0"
|
||||
|
||||
def supports_rtc(self) -> bool:
|
||||
return True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PI0Config,
|
||||
|
||||
@@ -714,9 +714,6 @@ class PI05Policy(PreTrainedPolicy):
|
||||
config_class = PI05Config
|
||||
name = "pi05"
|
||||
|
||||
def supports_rtc(self) -> bool:
|
||||
return True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PI05Config,
|
||||
|
||||
@@ -22,7 +22,6 @@ import numpy as np
|
||||
import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, TransitionKey
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
PolicyAction,
|
||||
@@ -34,6 +33,7 @@ from lerobot.processor import (
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
|
||||
from .configuration_pi05 import PI05Config
|
||||
|
||||
@@ -22,7 +22,6 @@ import numpy as np
|
||||
import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, TransitionKey
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
ActionTokenizerProcessorStep,
|
||||
@@ -35,6 +34,7 @@ from lerobot.processor import (
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
|
||||
from .configuration_pi0_fast import PI0FastConfig
|
||||
|
||||
@@ -34,22 +34,14 @@ from lerobot.configs import PreTrainedConfig
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
from lerobot.utils.device_utils import resolve_safetensors_device
|
||||
from lerobot.utils.hub import HubMixin
|
||||
from lerobot.utils.import_utils import _peft_available, require_package
|
||||
|
||||
from .utils import log_model_loading_keys
|
||||
|
||||
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
|
||||
T = TypeVar("T", bound="PreTrainedPolicy")
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
|
||||
T = TypeVar("T", bound="PreTrainedPolicy")
|
||||
|
||||
|
||||
def _build_card_context(
|
||||
cfg: TrainPipelineConfig | None,
|
||||
@@ -249,10 +241,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def supports_rtc(self) -> bool:
|
||||
"""Whether this policy implements Real-Time Chunking inference semantics."""
|
||||
return False
|
||||
|
||||
# TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
|
||||
@abc.abstractmethod
|
||||
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
|
||||
@@ -396,7 +384,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
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.
|
||||
"""
|
||||
require_package("peft", extra="peft")
|
||||
from peft import get_peft_model
|
||||
|
||||
# If user provided a complete config, use it directly (with overrides)
|
||||
if peft_config is not None:
|
||||
@@ -467,7 +455,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
Returns:
|
||||
Preprocessed dict with renamed keys and init_type mapped to method-specific key.
|
||||
"""
|
||||
require_package("peft", extra="peft")
|
||||
from peft import PeftType
|
||||
|
||||
cli_overrides = cli_overrides.copy()
|
||||
|
||||
@@ -492,7 +480,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
|
||||
def _build_peft_config(self, cli_overrides: dict):
|
||||
"""Build a PEFT config from policy defaults and CLI overrides."""
|
||||
require_package("peft", extra="peft")
|
||||
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
|
||||
|
||||
# Determine PEFT method type (default to LORA)
|
||||
method_type_str = cli_overrides.get("method_type") or "lora"
|
||||
@@ -519,7 +507,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
|
||||
def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
|
||||
"""Apply CLI overrides to an existing PEFT config."""
|
||||
require_package("peft", extra="peft")
|
||||
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
|
||||
|
||||
# Get method type from existing config or CLI override
|
||||
method_type_str = cli_overrides.get("method_type")
|
||||
|
||||
@@ -145,9 +145,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
|
||||
config_class = SmolVLAConfig
|
||||
name = "smolvla"
|
||||
|
||||
def supports_rtc(self) -> bool:
|
||||
return True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: SmolVLAConfig,
|
||||
|
||||
@@ -168,23 +168,14 @@ class SmolVLMWithExpertModel(nn.Module):
|
||||
last_layers.append(self.num_vlm_layers - 2)
|
||||
frozen_layers = [
|
||||
"lm_head",
|
||||
"text_model.norm.weight",
|
||||
"text_model.model.norm.weight",
|
||||
]
|
||||
for layer in last_layers:
|
||||
frozen_layers.append(f"text_model.layers.{layer}.")
|
||||
frozen_layers.append(f"text_model.model.layers.{layer}.")
|
||||
|
||||
unmatched_patterns = set(frozen_layers)
|
||||
for name, params in self.vlm.named_parameters():
|
||||
matched_patterns = [k for k in frozen_layers if k in name]
|
||||
if matched_patterns:
|
||||
if any(k in name for k in frozen_layers):
|
||||
params.requires_grad = False
|
||||
unmatched_patterns.difference_update(matched_patterns)
|
||||
if unmatched_patterns:
|
||||
raise RuntimeError(
|
||||
"Some frozen layer patterns matched no VLM parameters, so the corresponding layers "
|
||||
"would silently remain trainable (parameter naming may have changed in transformers): "
|
||||
f"{sorted(unmatched_patterns)}"
|
||||
)
|
||||
# To avoid unused params issue with distributed training
|
||||
for name, params in self.lm_expert.named_parameters():
|
||||
if "lm_head" in name:
|
||||
|
||||
@@ -22,7 +22,7 @@ import torch
|
||||
from torch import nn
|
||||
|
||||
from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig
|
||||
from lerobot.lerobot_types import PolicyAction, RobotAction, RobotObservation
|
||||
from lerobot.types import PolicyAction, RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import ACTION, OBS_STR
|
||||
from lerobot.utils.feature_utils import build_dataset_frame
|
||||
|
||||
|
||||
@@ -150,7 +150,7 @@ class XVLAModel(nn.Module):
|
||||
# Freeze or unfreeze policy transformer
|
||||
if not self.config.train_policy_transformer:
|
||||
for name, param in self.transformer.named_parameters():
|
||||
if "soft_prompt" not in name:
|
||||
if "soft_prompts" not in name:
|
||||
param.requires_grad = False
|
||||
|
||||
# Freeze or unfreeze soft prompts
|
||||
|
||||
@@ -21,7 +21,6 @@ import numpy as np
|
||||
import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, TransitionKey
|
||||
from lerobot.processor import (
|
||||
ObservationProcessorStep,
|
||||
PolicyAction,
|
||||
@@ -32,6 +31,7 @@ from lerobot.processor import (
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
IMAGENET_STATS,
|
||||
OBS_IMAGES,
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from lerobot.lerobot_types import (
|
||||
from lerobot.types import (
|
||||
EnvAction,
|
||||
EnvTransition,
|
||||
PolicyAction,
|
||||
|
||||
@@ -25,7 +25,7 @@ from dataclasses import dataclass, field
|
||||
from torch import Tensor
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, PolicyAction
|
||||
from lerobot.types import EnvTransition, PolicyAction
|
||||
from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
|
||||
|
||||
from .pipeline import (
|
||||
|
||||
@@ -23,7 +23,7 @@ from typing import Any
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from lerobot.lerobot_types import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
|
||||
from lerobot.types import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
|
||||
from lerobot.utils.constants import ACTION, DONE, INFO, OBS_PREFIX, REWARD, TRUNCATED
|
||||
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import PolicyAction, RobotAction
|
||||
from lerobot.types import PolicyAction, RobotAction
|
||||
|
||||
from .pipeline import ActionProcessorStep, ProcessorStepRegistry, RobotActionProcessorStep
|
||||
|
||||
@@ -132,20 +132,10 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
for axis in ["x", "y", "z"]:
|
||||
for axis in ["x", "y", "z", "gripper"]:
|
||||
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",
|
||||
"gripper_vel",
|
||||
]:
|
||||
for feat in ["enabled", "target_x", "target_y", "target_z", "target_wx", "target_wy", "target_wz"]:
|
||||
features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
|
||||
type=FeatureType.ACTION, shape=(1,)
|
||||
)
|
||||
|
||||
@@ -25,7 +25,7 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey
|
||||
from lerobot.types import EnvTransition, PolicyAction, TransitionKey
|
||||
from lerobot.utils.device_utils import get_safe_torch_device
|
||||
|
||||
from .pipeline import ProcessorStep, ProcessorStepRegistry
|
||||
|
||||
@@ -20,7 +20,7 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.configs.policies import PreTrainedConfig
|
||||
from lerobot.lerobot_types import PolicyAction, RobotAction, RobotObservation
|
||||
from lerobot.types import PolicyAction, RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .batch_processor import AddBatchDimensionProcessorStep
|
||||
|
||||
@@ -17,7 +17,7 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvAction, EnvTransition, PolicyAction, TransitionKey
|
||||
from lerobot.types import EnvAction, EnvTransition, PolicyAction, TransitionKey
|
||||
|
||||
from .converters import to_tensor
|
||||
from .hil_processor import TELEOP_ACTION_KEY
|
||||
|
||||
@@ -29,7 +29,7 @@ from lerobot.teleoperators.utils import TeleopEvents
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.teleoperators.teleoperator import Teleoperator
|
||||
|
||||
from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey
|
||||
from lerobot.types import EnvTransition, PolicyAction, TransitionKey
|
||||
|
||||
from .pipeline import (
|
||||
ComplementaryDataProcessorStep,
|
||||
|
||||
@@ -25,7 +25,7 @@ import torch
|
||||
from torch import Tensor
|
||||
|
||||
from lerobot.configs import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey
|
||||
from lerobot.types import EnvTransition, PolicyAction, TransitionKey
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.datasets import LeRobotDataset
|
||||
|
||||
@@ -45,14 +45,7 @@ from huggingface_hub import hf_hub_download
|
||||
from safetensors.torch import load_file, save_file
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import (
|
||||
EnvAction,
|
||||
EnvTransition,
|
||||
PolicyAction,
|
||||
RobotAction,
|
||||
RobotObservation,
|
||||
TransitionKey,
|
||||
)
|
||||
from lerobot.types import EnvAction, EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
|
||||
from lerobot.utils.constants import HF_LEROBOT_HOME
|
||||
from lerobot.utils.hub import HubMixin
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import PolicyAction, RobotAction
|
||||
from lerobot.types import PolicyAction, RobotAction
|
||||
from lerobot.utils.constants import ACTION
|
||||
|
||||
from .pipeline import ActionProcessorStep, ProcessorStepRegistry
|
||||
|
||||
@@ -20,7 +20,7 @@ import torch
|
||||
from torch import Tensor
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, TransitionKey
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
|
||||
from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
|
||||
|
||||
@@ -23,7 +23,7 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.configs.recipe import TrainingRecipe
|
||||
from lerobot.datasets.language import LANGUAGE_EVENTS, LANGUAGE_PERSISTENT
|
||||
from lerobot.datasets.language_render import render_sample
|
||||
from lerobot.lerobot_types import EnvTransition, TransitionKey
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.utils import unwrap_scalar
|
||||
|
||||
from .pipeline import ProcessorStep, ProcessorStepRegistry
|
||||
|
||||
@@ -30,7 +30,7 @@ from typing import TYPE_CHECKING, Any
|
||||
import torch
|
||||
|
||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, RobotObservation, TransitionKey
|
||||
from lerobot.types import EnvTransition, RobotObservation, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
ACTION_TOKEN_MASK,
|
||||
ACTION_TOKENS,
|
||||
|
||||
@@ -57,10 +57,10 @@ import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from lerobot.datasets import LeRobotDataset
|
||||
from lerobot.lerobot_types import TransitionKey
|
||||
from lerobot.rewards.robometer.configuration_robometer import RobometerConfig
|
||||
from lerobot.rewards.robometer.modeling_robometer import RobometerRewardModel
|
||||
from lerobot.rewards.robometer.processor_robometer import RobometerEncoderProcessorStep
|
||||
from lerobot.types import TransitionKey
|
||||
|
||||
DEFAULT_OUTPUT_FILENAME = "robometer_progress.parquet"
|
||||
|
||||
|
||||
@@ -25,7 +25,6 @@ from PIL import Image
|
||||
from torch import Tensor
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, TransitionKey
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
@@ -40,6 +39,7 @@ from lerobot.rewards.robometer.configuration_robometer import (
|
||||
RobometerConfig,
|
||||
)
|
||||
from lerobot.rewards.robometer.modeling_robometer import ROBOMETER_FEATURE_PREFIX
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
OBS_IMAGES,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
|
||||
@@ -47,7 +47,6 @@ else:
|
||||
Faker = None # type: ignore[assignment, misc]
|
||||
|
||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
@@ -59,6 +58,7 @@ from lerobot.processor import (
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.types import EnvTransition, PolicyAction, TransitionKey
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_sarm import SARMConfig
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user