Compare commits

..

44 Commits

Author SHA1 Message Date
Pepijn 21321d7a27 Merge remote-tracking branch 'origin/main' into codex/episode-video-streaming-byte-cache
# Conflicts:
#	src/lerobot/scripts/lerobot_train.py
2026-07-28 11:18:45 +02:00
Pepijn fdf6014214 Remove streaming benchmark prototypes 2026-07-28 11:14:10 +02:00
Pepijn 4db03ed535 Derive streaming decoder cap from episode pool 2026-07-27 20:49:19 +02:00
Pepijn 79cfb52a71 Harden production dataset streaming pipeline 2026-07-27 16:07:48 +02:00
Pepijn c0e9b0bbff Preserve private Hub auth in streaming readers 2026-07-27 09:41:57 +02:00
Pepijn b5a13e43ce Merge origin/main into streaming byte-cache branch 2026-07-27 09:37:33 +02:00
Pepijn 1e7e0b6de5 Fix streaming Parquet imports and video fallback 2026-07-24 16:41:54 +02:00
Pepijn a3edab661b chore(streaming): remove obsolete benchmark flags 2026-07-24 13:36:17 +02:00
Pepijn 564ba6395a Merge remote-tracking branch 'origin/main' into codex/episode-video-streaming-byte-cache 2026-07-24 13:26:42 +02:00
Pepijn 85086fed7a refactor(streaming): split episode data plane 2026-07-24 13:26:36 +02:00
Pepijn ee06d1005f refactor(dataset): make streaming pool rank-owned 2026-07-23 17:12:04 +02:00
Pepijn 1c47809bf6 chore: Merge origin/main into streaming branch 2026-07-23 16:10:46 +02:00
Pepijn 3d70b21aac feat(dataset): integrate episode streaming into training 2026-07-23 16:10:39 +02:00
Pepijn b85620657f chore: Merge origin/main into streaming branch 2026-07-23 14:27:30 +02:00
Pepijn fbfc861cf2 refactor(streaming): exact coverage is the only pool mode
Drop the with-replacement sampled path: delete run_pool_stream_simulation
and the --coverage flag; the streaming keep-up sim always uses
run_exact_coverage_stream (ExactCoveragePool), so every frame of every
episode is decoded exactly once per epoch. --pool-samples-per-episode is
kept as a deprecated no-op so existing commands still parse (exact mode
evicts an episode only when all its frames are emitted, so a turnover
cadence no longer applies).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-03 15:07:11 +02:00
Pepijn 06aa6a0425 feat(streaming): exact-once epoch coverage for the byte-cache episode pool
The pool path sampled frames with replacement and never guaranteed a full
epoch (episodes rotated on a fixed cadence; frames drawn randomly, none
tracked). Add ExactCoveragePool: a deterministic planner that enumerates
every frame of every episode exactly once per epoch while keeping at most
pool_size episodes resident, so batch mixing stays high (uniform draw over
all remaining frames in the pool) but coverage is complete and reproducible.

Mechanics (the "evict only when all frames sampled" model): episodes are
admitted in a seeded global permutation; each resident episode carries a
seeded frame-index shuffle; each draw picks a resident episode with
probability proportional to its remaining frames and pops one; an episode
is evicted only when its last frame is emitted, then a new one is admitted;
the epoch ends when admission is exhausted and every resident episode drains.
Order is a pure function of (seed, epoch) -> resumable by deterministic
fast-forward. The planner does no I/O and exposes admission_order so callers
can prefetch episodes ahead of the sampling frontier.

Wired into the benchmark as --coverage {sampled,exact}: run_exact_coverage_stream
prefetches stream_prefetch_episodes beyond the frontier so a freshly admitted
episode's bytes are resident before it is drawn, then decodes each frame once,
paced to target.

Tests: 7 planner unit tests (exact-once coverage incl. a 45k-frame epoch,
pool-size bound, per-(seed,epoch) determinism with coverage preserved,
admission/eviction events, coupon-collector mixing, zero-length episodes) and
a mocked-cache structural test of run_exact_coverage_stream asserting
every-frame-once-per-camera plus the prefetch-before-decode invariant.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-03 15:01:35 +02:00
Pepijn be64ded80f perf(streaming): sub-range parallel fetch + non-blocking pool replacement
The 64-vs-128-worker benchmark pair proved a per-host throughput ceiling
(~270 MiB/s) on the HF bucket path: doubling connections exactly halved
per-connection speed (4.8 -> 2.2 MiB/s) and left the aggregate flat,
while per-episode latency doubled (5.7s -> 12s) and keep-up worsened.
Steady-state demand (148 MiB/s) is well below the ceiling; the keep-up
misses come entirely from consumer stalls (refill_wait 14-19s of ~84s):
the sim blocks the training hot path on ensure_ready() for the FIFO-head
replacement while episodes take 5.7-12s to arrive.

Two fixes:

- Non-blocking replacements: EpisodeByteCache.is_ready() (all cameras
  cached or futures done, no blocking) and the stream sim now swaps a
  replacement only when it is already resident, deferring otherwise;
  fetch capacity (~2x demand) repays the debt on later batches. A
  deferred_swaps metric is reported.
- Sub-range parallel fetch (native-http): --range-subranges N splits one
  camera GET into N concurrent sub-range GETs. Under a per-host ceiling
  this adds no bandwidth but divides per-episode latency by ~N. Keep
  workers x subranges near the ~64-connection saturation point (e.g.
  --workers 16 --range-subranges 4).

Verified: sub-range span math + order-preserving concat and is_ready
semantics (unit-level, network stubbed).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-03 14:05:46 +02:00
Pepijn 88843ed675 Fix episode pool stream benchmark pacing 2026-06-22 20:40:33 +02:00
Pepijn f2b5c4a47b Add distributed episode pool benchmark summaries 2026-06-22 17:08:02 +02:00
Pepijn 9202fcea96 Fix pool sampling camera timestamps 2026-06-22 16:44:37 +02:00
Pepijn ef47c35178 Benchmark random sampling from episode pool 2026-06-22 16:26:26 +02:00
Pepijn 6d6c82eb8c Add GOP window range benchmark 2026-06-22 15:10:21 +02:00
Pepijn 9201be92cb Log failed HTTP range attempts 2026-06-22 12:42:30 +02:00
Pepijn 0064a06205 Instrument HfFileSystem range requests 2026-06-22 12:07:59 +02:00
Pepijn 710171ccac Clarify native HTTP exception timing 2026-06-22 12:02:12 +02:00
Pepijn 0f8257443c Retry native HTTP timeout statuses 2026-06-22 11:35:28 +02:00
Pepijn 3a09d0c48a Track native HTTP failed attempt timing 2026-06-22 11:27:08 +02:00
Pepijn 03fc5e3ea9 Allow dynamic range timing counters 2026-06-19 09:43:57 +02:00
Pepijn 28c3e095bf Report native HTTP chunk timing 2026-06-19 09:34:22 +02:00
Pepijn 5bfb749a9b Log episode cache fill progress 2026-06-18 18:43:43 +02:00
Pepijn 51c023a7a1 Tune native HTTP range diagnostics 2026-06-17 21:50:05 +02:00
Pepijn 51ea18cb7a Allow native HTTP sidecar range diagnostics 2026-06-17 21:36:57 +02:00
Pepijn 04ab43b8d2 Report range read timing breakdown 2026-06-17 21:20:08 +02:00
Pepijn cdfe192491 Remove random frame benchmark path 2026-06-17 21:14:42 +02:00
Pepijn 3451e53452 Use HfFileSystem for sidecar episode benchmark 2026-06-17 21:01:43 +02:00
Pepijn 30849ce74f Report memory usage in cache benchmarks 2026-06-17 20:54:12 +02:00
Pepijn 7d6907c444 Add random frame range fetch benchmark 2026-06-17 20:48:46 +02:00
Pepijn d99e1fe89d Report episode cache fill stage timings 2026-06-17 20:29:57 +02:00
Pepijn 7fcde61b69 Report full dataset estimate in episode cache benchmark 2026-06-17 20:25:21 +02:00
Pepijn bdfe8f8ce9 Use full MP4 sidecar for episode cache benchmark 2026-06-17 20:22:04 +02:00
Pepijn 34d0495d03 Retry transient native HTTP range failures 2026-06-17 20:19:54 +02:00
Pepijn 834c282631 Make episode cache benchmark fetch-only by default 2026-06-17 20:16:30 +02:00
Pepijn f132885cbc Pin Hub range cache and datasets main sources 2026-06-17 19:46:41 +02:00
Pepijn d0686be2f5 Add episode video streaming byte cache 2026-06-17 19:31:02 +02:00
227 changed files with 6297 additions and 2902 deletions
+7 -11
View File
@@ -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
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
+5 -4
View File
@@ -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"]
+5 -4
View File
@@ -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"]
+2
View File
@@ -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
+3 -3
View File
@@ -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(
-13
View File
@@ -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>
+15 -2
View File
@@ -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**
+4 -4
View File
@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
For each motor you want to update:
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
-1
View File
@@ -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`.
+1 -1
View File
@@ -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 \
+1 -1
View File
@@ -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):
+1 -1
View File
@@ -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
+42 -3
View File
@@ -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;">
+1 -1
View File
@@ -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:
+4 -4
View File
@@ -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,15 +74,15 @@ In the phone to SO-100 follower examples we use the following adapters:
- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
- `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.
Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
## Dataset feature contracts
Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
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(
-2
View File
@@ -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
View File
@@ -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.
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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.
+121
View File
@@ -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.
+2 -2
View File
@@ -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 \
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+22 -8
View File
@@ -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
View File
@@ -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.*"
+95
View File
@@ -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()
+1 -1
View File
@@ -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
+48 -78
View File
@@ -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()
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})"
def _reset_connection_settings(self) -> None:
"""Restore settings that may have been auto-detected during a failed connection."""
self.fps = self.config.fps
self.width = self.config.width
self.height = self.config.height
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
@property
def is_connected(self) -> bool:
"""Checks if the camera is currently connected and opened."""
@@ -172,25 +164,17 @@ 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()
self._configure_capture_settings()
self._start_read_thread()
if warmup and self.warmup_s > 0:
start_time = time.time()
while time.time() - start_time < self.warmup_s:
self.async_read(timeout_ms=self.warmup_s * 1000)
time.sleep(0.1)
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
if warmup and self.warmup_s > 0:
start_time = time.time()
while time.time() - start_time < self.warmup_s:
self.async_read(timeout_ms=self.warmup_s * 1000)
time.sleep(0.1)
with self.frame_lock:
if self.latest_frame is None:
raise ConnectionError(f"{self} failed to capture frames during warmup.")
logger.info(f"{self} connected.")
@@ -328,36 +312,32 @@ 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))
default_fps = camera.get(cv2.CAP_PROP_FPS)
default_format = camera.get(cv2.CAP_PROP_FORMAT)
if camera.isOpened():
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
default_fps = camera.get(cv2.CAP_PROP_FPS)
default_format = camera.get(cv2.CAP_PROP_FORMAT)
# 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)]
)
# 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)])
camera_info = {
"name": f"OpenCV Camera @ {target}",
"type": "OpenCV",
"id": target,
"backend_api": camera.getBackendName(),
"default_stream_profile": {
"format": default_format,
"fourcc": default_fourcc,
"width": default_width,
"height": default_height,
"fps": default_fps,
},
}
camera_info = {
"name": f"OpenCV Camera @ {target}",
"type": "OpenCV",
"id": target,
"backend_api": camera.getBackendName(),
"default_stream_profile": {
"format": default_format,
"fourcc": default_fourcc,
"width": default_width,
"height": default_height,
"fps": default_fps,
},
}
found_cameras_info.append(camera_info)
finally:
found_cameras_info.append(camera_info)
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.")
+33 -174
View File
@@ -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()
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})"
def _reset_connection_settings(self) -> None:
"""Restore settings that may have been auto-detected during a failed connection."""
self.fps = self.config.fps
self.width = self.config.width
self.height = self.config.height
self.warmup_s = self.config.warmup_s
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
@property
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,31 +190,22 @@ 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()
self._configure_capture_settings()
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.
self.warmup_s = max(self.warmup_s, 1)
# NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
self.warmup_s = max(self.warmup_s, 1)
warmup_read = self.async_read if self.use_rgb else self.async_read_depth
start_time = time.time()
while time.time() - start_time < self.warmup_s:
warmup_read(timeout_ms=self.warmup_s * 1000)
time.sleep(0.1)
with self.frame_lock:
if (self.use_rgb and self.latest_color_frame is None) or (
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
warmup_read = self.async_read if self.use_rgb else self.async_read_depth
start_time = time.time()
while time.time() - start_time < self.warmup_s:
warmup_read(timeout_ms=self.warmup_s * 1000)
time.sleep(0.1)
with self.frame_lock:
if (self.use_rgb and self.latest_color_frame is None) or (
self.use_depth and self.latest_depth_frame is None
):
raise ConnectionError(f"{self} failed to capture frames during warmup.")
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(
+1 -1
View File
@@ -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(
+26 -2
View File
@@ -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)
-6
View File
@@ -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}"
+32
View File
@@ -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(
+2 -4
View File
@@ -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(
+2 -4
View File
@@ -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(
+1 -5
View File
@@ -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
+58 -114
View File
@@ -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,16 +406,15 @@ def aggregate_videos(
videos_idx[key]["dst_file_durations"] = {}
for key, video_idx in videos_idx.items():
unique_chunk_file_pairs: list[ChunkFile] = sorted(
{
(chunk, file)
for chunk, file in zip(
src_meta.episodes[f"videos/{key}/chunk_index"],
src_meta.episodes[f"videos/{key}/file_index"],
strict=False,
)
}
)
unique_chunk_file_pairs = {
(chunk, file)
for chunk, file in zip(
src_meta.episodes[f"videos/{key}/chunk_index"],
src_meta.episodes[f"videos/{key}/file_index"],
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(
{
(c, f)
for c, f in zip(
src_meta.episodes["data/chunk_index"],
src_meta.episodes["data/file_index"],
strict=False,
)
}
)
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
)
}
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,16 +580,16 @@ def aggregate_metadata(
Returns:
dict: Updated meta_idx with current chunk and file indices.
"""
chunk_file_ids: list[ChunkFile] = sorted(
{
(c, f)
for c, f in zip(
src_meta.episodes["meta/episodes/chunk_index"],
src_meta.episodes["meta/episodes/file_index"],
strict=False,
)
}
)
chunk_file_ids = {
(c, f)
for c, f in zip(
src_meta.episodes["meta/episodes/chunk_index"],
src_meta.episodes["meta/episodes/file_index"],
strict=False,
)
}
chunk_file_ids = sorted(chunk_file_ids)
for chunk_idx, file_idx in chunk_file_ids:
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)
+2 -2
View File
@@ -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:
+2 -4
View File
@@ -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
+58 -14
View File
@@ -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,17 +198,43 @@ def make_train_eval_datasets(
ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None
)
train_dataset = LeRobotDataset(
cfg.dataset.repo_id,
root=cfg.dataset.root,
episodes=train_episodes,
delta_timestamps=delta_timestamps,
image_transforms=train_image_transforms,
revision=cfg.dataset.revision,
video_backend=cfg.dataset.video_backend,
return_uint8=True,
tolerance_s=cfg.tolerance_s,
)
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,
episodes=train_episodes,
delta_timestamps=delta_timestamps,
image_transforms=train_image_transforms,
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,
)
eval_dataset = LeRobotDataset(
cfg.dataset.repo_id,
@@ -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,
)
+1 -1
View File
@@ -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
+149
View File
@@ -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,
)
+8 -7
View File
@@ -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"
)
+1 -1
View File
@@ -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"
+2 -7
View File
@@ -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
self._env.close()
def _make_env_fns(
+1 -1
View File
@@ -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
+2 -15
View File
@@ -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().
+2 -4
View File
@@ -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(
+2 -2
View File
@@ -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
-18
View File
@@ -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 = [
+1 -1
View File
@@ -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()
+6 -36
View File
@@ -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,
+4 -19
View File
@@ -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,
-3
View File
@@ -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,
+1 -1
View File
@@ -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
+5 -17
View File
@@ -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:
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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,
+1 -1
View File
@@ -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,
+1 -1
View File
@@ -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 (
+1 -1
View File
@@ -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,)
)
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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,
+1 -1
View File
@@ -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
+1 -8
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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,
+1 -1
View File
@@ -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