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
126 changed files with 6161 additions and 2766 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`.
On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
exposure or gain also affects the depth stream.
</hfoption>
</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 -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.
+5 -53
View File
@@ -11,10 +11,9 @@ LeRobot provides several utilities for manipulating datasets:
3. **Merge Datasets** - Combine multiple datasets into one. The datasets must have identical features, and episodes are concatenated in the order specified in `repo_ids`
4. **Add Features** - Add new features to a dataset
5. **Remove Features** - Remove features from a dataset
6. **Modify Tasks** - Change the natural-language task descriptions associated with episodes
7. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders)
8. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings
9. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc.
6. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders)
7. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings
8. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc.
The core implementation is in `lerobot.datasets.dataset_tools`.
An example script detailing how to use the tools API is available in `examples/dataset/use_dataset_tools.py`.
@@ -51,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 \
@@ -90,53 +89,6 @@ lerobot-edit-dataset \
--operation.feature_names "['observation.images.top']"
```
#### Modify Tasks
Change the natural-language task descriptions attached to episodes. This is useful for fixing typos, standardizing wording, or re-labeling episodes.
> [!WARNING]
> `modify_tasks` modifies the dataset **in-place** (updating `meta/tasks.parquet`, the `task_index` column in the data files, the `tasks` column in the episode metadata, and `total_tasks` in `meta/info.json`). The `--new_repo_id` and `--new_root` parameters are ignored for this operation.
```bash
# Set a single task for all episodes
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.new_task "Pick up the cube and place it"
# Set different tasks for specific episodes
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.episode_tasks '{"0": "Task A", "1": "Task B", "2": "Task A"}'
# Replace existing task strings wherever they appear
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
# Combine modes in a single run
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.new_task "Default task" \
--operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}' \
--operation.episode_tasks '{"5": "Special task for episode 5"}'
```
**Parameters:**
- `new_task`: A single task string used as the default for episodes not otherwise covered.
- `episode_tasks`: Mapping from episode index to task string.
- `task_replacements`: Mapping from existing task strings to their replacements, applied to episodes whose current task matches a key. Every key must be an existing task in the dataset.
The modes can be combined in a single run. Per episode, the task is resolved with the following precedence:
`episode_tasks` > `task_replacements` > `new_task` > original task
At least one of `new_task`, `episode_tasks`, or `task_replacements` must be specified. An episode that ends up with no task raises an error.
#### Convert to Video
Convert an image-based dataset to video format, creating a new LeRobotDataset where images are stored as videos. This is useful for reducing storage requirements and improving data loading performance. The new dataset will have the exact same structure as the original, but with images encoded as MP4 videos in the proper LeRobot format.
+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()
+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 -171
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,111 +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 sensor that controls the color stream.
Most RealSense cameras expose "RGB Camera" for color. The D405 has no
separate RGB module — its color stream comes from "Stereo Module".
We try RGB Camera first, then fall back to Stereo Module.
"""
if self.rs_profile is None:
raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
device = self.rs_profile.get_device()
sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
for name in ("RGB Camera", "Stereo Module"):
if name in sensors:
return sensors[name]
available = list(sensors.keys())
raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
"""Sets a sensor option, re-raising range errors with actionable diagnostics."""
try:
sensor.set_option(option, value)
except Exception as e:
range_info = ""
try:
option_range = sensor.get_option_range(option)
range_info = (
f" (supported range: min={option_range.min}, max={option_range.max}, "
f"step={option_range.step}, default={option_range.default})"
)
except Exception:
range_info = " (option range unavailable)"
raise ValueError(
f"{self}: failed to set {label} to {value}{range_info}. Original error: {e}"
) from e
def _configure_sensor_options(self) -> None:
"""Applies manual sensor options (exposure, gain, white balance) to the color sensor.
When exposure or gain is set, auto-exposure is disabled first. When white_balance
is set, auto white balance is disabled first. An omitted option is left unchanged,
and configuration is skipped entirely if all options are omitted.
Raises:
ValueError: If the sensor does not support a requested option or a requested
value is invalid. Invalid-value errors include the option name, requested
value, and supported range when available.
"""
if self.exposure is None and self.gain is None and self.white_balance is None:
return
color_sensor = self._get_color_sensor()
requested_options = (
(rs.option.exposure, self.exposure, "exposure"),
(rs.option.gain, self.gain, "gain"),
(rs.option.white_balance, self.white_balance, "white balance"),
)
unsupported_options = [
label
for option, value, label in requested_options
if value is not None and not color_sensor.supports(option)
]
if unsupported_options:
raise ValueError(
f"{self}: color sensor does not support requested manual options: {unsupported_options}."
)
manual_exposure_requested = self.exposure is not None or self.gain is not None
if manual_exposure_requested:
if color_sensor.supports(rs.option.enable_auto_exposure):
self._set_sensor_option(color_sensor, rs.option.enable_auto_exposure, 0, "auto-exposure")
logger.info(f"{self} auto-exposure disabled.")
else:
logger.warning(
f"{self} sensor does not support disabling auto-exposure; "
"applying manual exposure/gain directly."
)
if self.exposure is not None:
self._set_sensor_option(color_sensor, rs.option.exposure, self.exposure, "exposure")
logger.info(f"{self} exposure set to {self.exposure}.")
if self.gain is not None:
self._set_sensor_option(color_sensor, rs.option.gain, self.gain, "gain")
logger.info(f"{self} gain set to {self.gain}.")
if self.white_balance is not None:
if color_sensor.supports(rs.option.enable_auto_white_balance):
self._set_sensor_option(
color_sensor, rs.option.enable_auto_white_balance, 0, "auto white balance"
)
logger.info(f"{self} auto white balance disabled.")
else:
logger.warning(
f"{self} sensor does not support disabling auto white balance; "
"applying manual white balance directly."
)
self._set_sensor_option(
color_sensor, rs.option.white_balance, self.white_balance, "white balance"
)
logger.info(f"{self} white balance set to {self.white_balance}.")
@check_if_not_connected
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
"""
@@ -671,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():
@@ -835,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(
+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:
+15 -37
View File
@@ -1435,18 +1435,15 @@ def modify_tasks(
dataset: LeRobotDataset,
new_task: str | None = None,
episode_tasks: dict[int, str] | None = None,
task_replacements: dict[str, str] | None = None,
) -> LeRobotDataset:
"""Modify tasks in a LeRobotDataset.
This function allows you to either:
1. Set a single task for the entire dataset (using `new_task`)
2. Set specific tasks for specific episodes (using `episode_tasks`)
3. Replace existing task strings wherever they appear (using `task_replacements`)
Per episode, the task is resolved with precedence:
`episode_tasks` > `task_replacements` > `new_task` > original task. An episode that ends
up with no task (none of the above apply and it had no original task) raises an error.
You can combine both: `new_task` sets the default, and `episode_tasks` overrides
specific episodes.
The dataset is modified in-place, updating only the task-related files:
- meta/tasks.parquet
@@ -1456,14 +1453,11 @@ def modify_tasks(
Args:
dataset: The source LeRobotDataset to modify.
new_task: Default task applied to any episode not covered by `episode_tasks` or a
matching `task_replacements` entry.
episode_tasks: Optional dict mapping episode indices to task strings. Takes precedence
over both `task_replacements` and `new_task`.
task_replacements: Optional dict mapping existing task strings to new ones. Applied to
episodes whose current task matches a key. Every key must be an existing task.
new_task: A single task string to apply to all episodes. If None and episode_tasks
is also None, raises an error.
episode_tasks: Optional dict mapping episode indices to their task strings.
Overrides `new_task` for specific episodes.
At least one of `new_task`, `episode_tasks`, or `task_replacements` must be provided.
Examples:
Set a single task for all episodes:
@@ -1481,17 +1475,11 @@ def modify_tasks(
new_task="Default task",
episode_tasks={5: "Special task for episode 5"}
)
Replace existing task strings in-place:
dataset = modify_tasks(
dataset,
task_replacements={"Pick up the cube": "Lift the cube"}
)
"""
if not new_task and not episode_tasks and not task_replacements:
raise ValueError("Must specify at least one of new_task, episode_tasks, or task_replacements")
if new_task is None and episode_tasks is None:
raise ValueError("Must specify at least one of new_task or episode_tasks")
if episode_tasks:
if episode_tasks is not None:
valid_indices = set(range(dataset.meta.total_episodes))
invalid = set(episode_tasks.keys()) - valid_indices
if invalid:
@@ -1501,29 +1489,19 @@ def modify_tasks(
if dataset.meta.episodes is None:
dataset.meta.episodes = load_episodes(dataset.root)
if task_replacements:
current_tasks = set(dataset.meta.tasks.index)
invalid_tasks = set(task_replacements) - current_tasks
if invalid_tasks:
raise ValueError(f"Task replacements reference unknown tasks: {sorted(invalid_tasks)}")
# Build the mapping from episode index to task string
episode_to_task: dict[int, str] = {}
for ep_idx in range(dataset.meta.total_episodes):
original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
original_task = original_tasks[0] if original_tasks else None
if episode_tasks and ep_idx in episode_tasks:
episode_to_task[ep_idx] = episode_tasks[ep_idx]
elif task_replacements and original_task in task_replacements:
episode_to_task[ep_idx] = task_replacements[original_task]
elif new_task:
elif new_task is not None:
episode_to_task[ep_idx] = new_task
elif original_task:
# Keep original task if not overridden and no default provided
episode_to_task[ep_idx] = original_task
else:
raise ValueError(f"Episode {ep_idx} has no task; provide new_task or episode_tasks")
# Keep original task if not overridden and no default provided
original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
if not original_tasks:
raise ValueError(f"Episode {ep_idx} has no tasks and no default task was provided")
episode_to_task[ep_idx] = original_tasks[0]
# Collect all unique tasks and create new task mapping
unique_tasks = sorted(set(episode_to_task.values()))
+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,
)
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 -6
View File
@@ -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 -3
View File
@@ -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(
+1 -1
View File
@@ -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 = [
@@ -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()
+5 -35
View File
@@ -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,
+3 -18
View File
@@ -44,19 +44,12 @@ from lerobot.utils.constants import (
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")
@@ -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)
-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,
+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:
@@ -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,)
)
+4 -2
View File
@@ -91,7 +91,7 @@ from lerobot.robots import so_follower # noqa: F401
from lerobot.teleoperators import gamepad, so_leader # noqa: F401
from lerobot.teleoperators.utils import TeleopEvents
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.random_utils import set_seed
from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.transition import (
@@ -124,7 +124,9 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
cfg.validate()
display_pid = False
if not use_threads(cfg):
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
import torch.multiprocessing as mp
mp.set_start_method("spawn")
display_pid = True
# Create logs directory to ensure it exists
+2 -5
View File
@@ -16,7 +16,6 @@ from __future__ import annotations
import abc
import builtins
import json
import logging
import os
from dataclasses import dataclass, field
@@ -79,10 +78,8 @@ class RLAlgorithmConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
def _save_pretrained(self, save_directory: Path) -> None:
"""Serialize this config as ``config.json`` inside ``save_directory``."""
# 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, RLAlgorithmConfig), 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 -2
View File
@@ -18,7 +18,7 @@ import functools
import threading
from collections.abc import Callable, Sequence
from contextlib import suppress
from typing import NotRequired, TypedDict
from typing import TypedDict
import torch
import torch.nn.functional as F # noqa: N812
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
next_state: dict[str, torch.Tensor]
done: torch.Tensor
truncated: torch.Tensor
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
complementary_info: dict[str, torch.Tensor | float | int] | None = None
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
+4 -2
View File
@@ -102,7 +102,7 @@ from lerobot.utils.constants import (
)
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.io_utils import load_json, write_json
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
format_big_number,
@@ -123,7 +123,9 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
require_package("grpcio", extra="hilserl", import_name="grpc")
if not use_threads(cfg):
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
import torch.multiprocessing as mp
mp.set_start_method("spawn")
# Use the job_name from the config
train(
@@ -46,12 +46,6 @@ class SOFollowerConfig:
position_i_coefficient: int = 0
position_d_coefficient: int = 32
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
# at once, which otherwise aborts the control loop. Retries are immediate (no sleep) and only happen on
# failure, so the steady-state read cost is unchanged.
num_read_retries: int = 2
@RobotConfig.register_subclass("so101_follower")
@RobotConfig.register_subclass("so100_follower")
@@ -510,10 +510,10 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
# We only use the ee pose in the dataset, so we don't need the joint positions
for n in self.motor_names:
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
# Store end-effector features as actions in the dataset schema
# We specify the dataset features of this step that we want to be stored in the dataset
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
type=FeatureType.STATE, shape=(1,)
)
return features
@@ -180,7 +180,7 @@ class SOFollower(Robot):
def get_observation(self) -> RobotObservation:
# Read arm position
start = time.perf_counter()
obs_dict = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
obs_dict = self.bus.sync_read("Present_Position")
obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read state: {dt_ms:.1f}ms")
@@ -221,7 +221,7 @@ class SOFollower(Robot):
# Cap goal position when too far away from present position.
# /!\ Slower fps expected due to reading from the follower.
if self.config.max_relative_target is not None:
present_pos = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
present_pos = self.bus.sync_read("Present_Position")
goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
+3 -28
View File
@@ -24,7 +24,6 @@ from __future__ import annotations
import logging
from dataclasses import dataclass, field
from threading import Event
from typing import TYPE_CHECKING
import torch
@@ -48,7 +47,6 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor
from lerobot.robots import make_robot_from_config
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
from lerobot.utils.import_utils import _peft_available, require_package
from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
from .inference import (
@@ -57,15 +55,8 @@ from .inference import (
SyncInferenceConfig,
create_inference_engine,
)
from .inference.rtc import supports_rtc_inference
from .robot_wrapper import ThreadSafeRobot
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
logger = logging.getLogger(__name__)
@@ -180,7 +171,7 @@ def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy
revision=pretrained_revision,
)
require_package("peft", extra="peft")
from peft import PeftConfig, PeftModel
peft_path = policy_config.pretrained_path
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
@@ -227,12 +218,6 @@ def build_rollout_context(
policy = _load_pretrained_policy(policy_config)
if is_rtc:
if not supports_rtc_inference(policy):
raise ValueError(
f"RTC inference is not supported by policy type '{policy_config.type}': "
"the policy must implement RTC semantics and predict_action_chunk must accept "
"inference_delay and prev_chunk_left_over. Use '--inference.type=sync' instead."
)
policy.config.rtc_config = cfg.inference.rtc
if hasattr(policy, "init_rtc_processor"):
policy.init_rtc_processor()
@@ -309,22 +294,12 @@ def build_rollout_context(
# ``observation_features`` values are either a tuple (camera shape) or the
# ``float`` type itself used as a sentinel for scalar motor features —
# see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
# Keep cameras (tuple) plus both joint-position (.pos) and base-velocity (.vel)
# scalar state features. LeKiwi's observation.state is 9-dim (6 arm .pos +
# x/y/theta.vel) and the policy was trained/normalized on all 9; the old .pos-only
# filter fed a 6-dim state into a 9-dim normalizer → RuntimeError (size 6 vs 9).
# Pure-arm robots have no .vel state keys, so this is a no-op for them.
observation_features_hw = {
k: v
for k, v in all_obs_features.items()
if isinstance(v, tuple) or (v is float and k.endswith((".pos", ".vel")))
if isinstance(v, tuple) or (v is float and k.endswith(".pos"))
}
# Keep both joint-position (.pos) and base-velocity (.vel) action features so
# mobile manipulators command the base too (e.g. LeKiwi: 6 arm .pos +
# x/y/theta.vel = 9-dim action). Pure-arm robots have no .vel keys, so this is
# a no-op for them. Without the .vel keys the base velocities are silently
# dropped from dataset_features[ACTION]/ordered_action_keys and the base never moves.
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith((".pos", ".vel"))}
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith(".pos")}
# The action side is always needed: sync inference reads action names from
# ``dataset_features[ACTION]`` to map policy tensors back to robot actions.
-18
View File
@@ -22,7 +22,6 @@ way via ``notify_observation``.
from __future__ import annotations
import inspect
import logging
import math
import time
@@ -63,23 +62,6 @@ _RTC_JOIN_TIMEOUT_S: float = 3.0
# ---------------------------------------------------------------------------
def supports_rtc_inference(policy: PreTrainedPolicy) -> bool:
"""Whether a policy declares RTC support and accepts the RTC call shape."""
supports_rtc = getattr(policy, "supports_rtc", None)
if not callable(supports_rtc) or not supports_rtc():
return False
try:
inspect.signature(policy.predict_action_chunk).bind(
object(),
inference_delay=0,
prev_chunk_left_over=None,
)
except (TypeError, ValueError):
return False
return True
def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int) -> torch.Tensor:
"""Pad or truncate RTC prefix actions to a fixed length for stable compiled inference."""
if prev_actions.ndim != 2:
@@ -36,7 +36,6 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
import argparse
import concurrent.futures
import logging
import os
from pathlib import Path
import numpy as np
@@ -53,7 +52,6 @@ from lerobot.datasets import (
get_feature_stats,
write_stats,
)
from lerobot.datasets.compute_stats import sample_indices
from lerobot.utils.utils import init_logging
@@ -79,14 +77,12 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
return False
def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
"""Process a single episode and return its statistics.
Args:
dataset: The LeRobot dataset
episode_idx: Index of the episode to process
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
Returns:
Dictionary containing episode statistics
@@ -96,31 +92,16 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli
start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
episode_len = end_idx - start_idx
# Images/video are the memory hog, so sub-sample those frames per episode;
# numeric columns are cheap, so read them in full (exact).
image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")]
numeric_keys = [
k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
]
collected_data: dict[str, list] = {}
for idx in range(start_idx, end_idx):
item = dataset[idx]
for key, value in item.items():
if key not in dataset.features:
continue
# Numeric features: every frame, read directly from the underlying table.
if numeric_keys:
numeric_cols = dataset.hf_dataset.select_columns(numeric_keys)[start_idx:end_idx]
for key in numeric_keys:
collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]]
# Image/video features: decode only a sampled subset of frames.
if image_keys:
sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len))
for offset in sampled_offsets:
item = dataset[start_idx + offset]
for key in image_keys:
if key in item:
collected_data.setdefault(key, []).append(item[key])
if key not in collected_data:
collected_data[key] = []
collected_data[key].append(value)
ep_stats = {}
for key, data_list in collected_data.items():
@@ -150,13 +131,11 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli
return ep_stats
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]:
"""Compute quantile statistics for all episodes in the dataset.
Args:
dataset: The LeRobot dataset to compute statistics for
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
Returns:
Dictionary containing aggregated statistics with quantiles
@@ -174,15 +153,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bo
if has_videos:
logging.info("Dataset contains video keys - using sequential processing for thread safety")
for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
ep_stats = process_single_episode(dataset, episode_idx)
episode_stats_list.append(ep_stats)
else:
logging.info("Dataset has no video keys - using parallel processing for better performance")
max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))
max_workers = min(dataset.num_episodes, 16)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
future_to_episode = {
executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
executor.submit(process_single_episode, dataset, episode_idx): episode_idx
for episode_idx in range(dataset.num_episodes)
}
@@ -209,7 +188,6 @@ def augment_dataset_with_quantile_stats(
repo_id: str,
root: str | Path | None = None,
overwrite: bool = False,
use_sampling: bool = True,
) -> None:
"""Augment a dataset with quantile statistics if they are missing.
@@ -217,8 +195,6 @@ def augment_dataset_with_quantile_stats(
repo_id: Repository ID of the dataset
root: Local root directory for the dataset
overwrite: Overwrite existing quantile statistics if they already exist
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
"""
logging.info(f"Loading dataset: {repo_id}")
dataset = LeRobotDataset(
@@ -232,7 +208,7 @@ def augment_dataset_with_quantile_stats(
logging.info("Dataset does not contain quantile statistics. Computing them now...")
new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
new_stats = compute_quantile_stats_for_dataset(dataset)
logging.info("Updating dataset metadata with new quantile statistics")
dataset.meta.stats = new_stats
@@ -272,14 +248,6 @@ def main():
action="store_true",
help="Overwrite existing quantile statistics if they already exist",
)
parser.add_argument(
"--no-sampling",
action="store_true",
help=(
"Compute stats over every frame (exact, higher memory). By default, "
"image/video frames are sub-sampled per episode to bound memory."
),
)
args = parser.parse_args()
root = Path(args.root) if args.root else None
@@ -290,7 +258,6 @@ def main():
repo_id=args.repo_id,
root=root,
overwrite=args.overwrite,
use_sampling=not args.no_sampling,
)
@@ -94,8 +94,6 @@ from lerobot.datasets.video_utils import concatenate_video_files, get_video_dura
from lerobot.utils.constants import HF_LEROBOT_HOME
from lerobot.utils.utils import flatten_dict, init_logging
logger = logging.getLogger(__name__)
V21 = "v2.1"
V30 = "v3.0"
@@ -478,11 +476,11 @@ def convert_dataset(
# First check if the dataset already has a v3.0 version
if root is None and not force_conversion:
try:
logger.info("Trying to download v3.0 version of the dataset from the hub...")
print("Trying to download v3.0 version of the dataset from the hub...")
snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
return
except Exception:
logger.info("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
# Set root based on whether local dataset path is provided
use_local_dataset = False
@@ -490,7 +488,7 @@ def convert_dataset(
if root.exists():
validate_local_dataset_version(root)
use_local_dataset = True
logger.info(f"Using local dataset at {root}")
print(f"Using local dataset at {root}")
old_root = root.parent / f"{root.name}_old"
new_root = root.parent / f"{root.name}_v30"
@@ -525,7 +523,7 @@ def convert_dataset(
try:
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
except (HTTPError, RevisionNotFoundError) as e:
logger.warning(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
pass
hub_api.delete_files(
delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
+7 -6
View File
@@ -154,14 +154,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
repo_id = cfg.new_repo_id or cfg.repo_id
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
api = HfApi()
logger.info(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...")
print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True)
api.create_repo(
repo_id=repo_id,
repo_type="dataset",
private=cfg.push_private,
exist_ok=True,
)
logger.info(f"[lerobot-annotate] uploading {root} -> {repo_id}...")
print(f"[lerobot-annotate] uploading {root} -> {repo_id}...", flush=True)
commit_info = api.upload_folder(
folder_path=str(root),
repo_id=repo_id,
@@ -172,7 +172,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
# at the source dataset; a fresh card is generated below instead.
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
)
logger.info(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}")
print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
dataset_info = load_info(root)
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
@@ -200,13 +200,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
with suppress(RevisionNotFoundError):
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
api.create_tag(**tag_kwargs)
logger.info(f"[lerobot-annotate] tagged {repo_id} as {version_tag}")
print(f"[lerobot-annotate] tagged {repo_id} as {version_tag}", flush=True)
except Exception as exc: # noqa: BLE001
logger.warning(
print(
f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
"Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
"Run: from huggingface_hub import HfApi; "
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)"
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)",
flush=True,
)
+1 -3
View File
@@ -89,8 +89,6 @@ from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
DEFAULT_FOXGLOVE_PORT = 8765
DEFAULT_RERUN_PORT = 9090
@@ -301,7 +299,7 @@ def visualize_dataset(
while True:
time.sleep(1)
except KeyboardInterrupt:
logger.info("Ctrl-C received. Exiting.")
print("Ctrl-C received. Exiting.")
def main():
+2 -15
View File
@@ -127,12 +127,6 @@ Modify tasks - set default task with overrides for specific episodes (WARNING: m
--operation.new_task "Default task" \
--operation.episode_tasks '{"5": "Special task for episode 5"}'
Modify tasks - replace existing task strings in-place (WARNING: modifies in-place):
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
Convert image dataset to video format and save locally:
lerobot-edit-dataset \
--repo_id lerobot/pusht_image \
@@ -309,7 +303,6 @@ class RemoveFeatureConfig(OperationConfig):
class ModifyTasksConfig(OperationConfig):
new_task: str | None = None
episode_tasks: dict[str, str] | None = None
task_replacements: dict[str, str] | None = None
@OperationConfig.register_subclass("convert_image_to_video")
@@ -558,12 +551,9 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
new_task = cfg.operation.new_task
episode_tasks_raw = cfg.operation.episode_tasks
task_replacements = cfg.operation.task_replacements
if new_task is None and episode_tasks_raw is None and task_replacements is None:
raise ValueError(
"Must specify at least one of new_task, episode_tasks, or task_replacements for modify_tasks operation"
)
if new_task is None and episode_tasks_raw is None:
raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
if cfg.new_repo_id is not None or cfg.new_root is not None:
logging.warning(
@@ -583,14 +573,11 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
logging.info(f" Default task: '{new_task}'")
if episode_tasks:
logging.info(f" Episode-specific tasks: {episode_tasks}")
if task_replacements:
logging.info(f" Task replacements: {task_replacements}")
modified_dataset = modify_tasks(
dataset,
new_task=new_task,
episode_tasks=episode_tasks,
task_replacements=task_replacements,
)
logging.info(f"Dataset modified at {dataset.root}")
+14 -20
View File
@@ -62,7 +62,7 @@ from dataclasses import asdict
from functools import partial
from pathlib import Path
from pprint import pformat
from typing import TYPE_CHECKING, Any, TypedDict
from typing import Any, TypedDict
import einops
import gymnasium as gym
@@ -87,7 +87,7 @@ from lerobot.processor import PolicyProcessorPipeline
from lerobot.types import PolicyAction
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.io_utils import write_video
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
@@ -95,14 +95,6 @@ from lerobot.utils.utils import (
inside_slurm,
)
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
logger = logging.getLogger(__name__)
def _env_features_to_dataset_features(env_features: dict) -> dict:
"""Convert EnvConfig.features to the dict format expected by LeRobotDataset.create()."""
@@ -452,11 +444,13 @@ def eval_policy(
exc = ValueError(
f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
)
if not _peft_available:
raise exc
require_package("peft", extra="peft")
if not isinstance(policy, PeftModel):
raise exc
try:
from peft import PeftModel
if not isinstance(policy, PeftModel):
raise exc
except ImportError:
raise exc from None
start = time.time()
# Preserve the mode for direct callers. eval_policy_all scopes the mode
@@ -564,7 +558,7 @@ def eval_policy(
if seeds:
all_seeds.extend(seeds)
else:
all_seeds.extend([None] * env.num_envs)
all_seeds.append(None)
# FIXME: episode_data is either None or it doesn't exist
if return_episode_data:
@@ -802,13 +796,13 @@ def eval_main(cfg: EvalPipelineConfig):
recording_repo_id=cfg.eval.recording_repo_id,
recording_private=cfg.eval.recording_private,
)
logger.info("Overall Aggregated Metrics:")
logger.info(info["overall"])
print("Overall Aggregated Metrics:")
print(info["overall"])
# Print per-suite stats
for task_group, task_group_info in info.items():
logger.info(f"\nAggregated Metrics for {task_group}:")
logger.info(task_group_info)
print(f"\nAggregated Metrics for {task_group}:")
print(task_group_info)
# Close all vec envs
close_envs(envs)
+49 -41
View File
@@ -28,6 +28,7 @@ lerobot-find-cameras
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
import argparse
import concurrent.futures
import logging
import time
from pathlib import Path
@@ -39,7 +40,6 @@ from PIL import Image
from lerobot.cameras import ColorMode
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
@@ -132,7 +132,7 @@ def save_image(
camera_identifier: str | int,
images_dir: Path,
camera_type: str,
) -> None:
):
"""
Saves a single image to disk using Pillow. Handles color conversion if necessary.
"""
@@ -151,7 +151,7 @@ def save_image(
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None:
def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
"""Create and connect to a camera instance based on metadata."""
cam_type = cam_meta.get("type")
cam_id = cam_meta.get("id")
@@ -164,14 +164,12 @@ def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> di
cv_config = OpenCVCameraConfig(
index_or_path=cam_id,
color_mode=ColorMode.RGB,
warmup_s=warmup_s,
)
instance = OpenCVCamera(cv_config)
elif cam_type == "RealSense":
rs_config = RealSenseCameraConfig(
serial_number_or_name=cam_id,
color_mode=ColorMode.RGB,
warmup_s=warmup_s,
)
instance = RealSenseCamera(rs_config)
else:
@@ -189,7 +187,9 @@ def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> di
return None
def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None:
def process_camera_image(
cam_dict: dict[str, Any], output_dir: Path, current_time: float
) -> concurrent.futures.Future | None:
"""Capture and process an image from a single camera."""
cam = cam_dict["instance"]
meta = cam_dict["meta"]
@@ -199,7 +199,7 @@ def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_tim
try:
image_data = cam.read()
save_image(
return save_image(
image_data,
cam_id_str,
output_dir,
@@ -214,21 +214,21 @@ def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_tim
return None
def cleanup_camera(cam_dict: dict[str, Any]) -> None:
def cleanup_cameras(cameras_to_use: list[dict[str, Any]]):
"""Disconnect all cameras."""
logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...")
try:
if cam_dict["instance"] and cam_dict["instance"].is_connected:
cam_dict["instance"].disconnect()
except Exception as e:
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
logger.info(f"Disconnecting {len(cameras_to_use)} cameras...")
for cam_dict in cameras_to_use:
try:
if cam_dict["instance"] and cam_dict["instance"].is_connected:
cam_dict["instance"].disconnect()
except Exception as e:
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
def save_images_from_all_cameras(
output_dir: Path,
record_time_s: float = 2.0,
camera_type: str | None = None,
warmup_s: int = 1,
):
"""
Connects to detected cameras (optionally filtered by type) and saves images from each.
@@ -239,7 +239,6 @@ def save_images_from_all_cameras(
record_time_s: Duration in seconds to record images.
camera_type: Optional string to filter cameras ("realsense" or "opencv").
If None, uses all detected cameras.
warmup_s: Duration in seconds to warmup camera before recording images.
"""
output_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Saving images to {output_dir}")
@@ -249,32 +248,47 @@ def save_images_from_all_cameras(
logger.warning("No cameras detected matching the criteria. Cannot save images.")
return
logger.info(
f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras."
)
cameras_to_use = []
for cam_meta in all_camera_metadata:
camera_instance = create_camera_instance(cam_meta)
if camera_instance:
cameras_to_use.append(camera_instance)
try:
for cam_meta in all_camera_metadata:
cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s)
if cam_dict is None:
continue
start_time = time.perf_counter()
if not cameras_to_use:
logger.warning("No cameras could be connected. Aborting image save.")
return
logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
start_time = time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
try:
while time.perf_counter() - start_time < record_time_s:
futures = []
current_capture_time = time.perf_counter()
process_camera_image(cam_dict, output_dir, current_capture_time)
cleanup_camera(cam_dict)
except KeyboardInterrupt:
logger.info("Capture interrupted by user.")
finally:
print(f"Image capture finished. Images saved to {output_dir}")
for cam_dict in cameras_to_use:
future = process_camera_image(cam_dict, output_dir, current_capture_time)
if future:
futures.append(future)
if futures:
concurrent.futures.wait(futures)
except KeyboardInterrupt:
logger.info("Capture interrupted by user.")
finally:
print("\nFinalizing image saving...")
executor.shutdown(wait=True)
cleanup_cameras(cameras_to_use)
print(f"Image capture finished. Images saved to {output_dir}")
def main():
init_logging()
parser = argparse.ArgumentParser(
description="Unified camera utility script for listing cameras and capturing images."
)
parser.add_argument(
"camera_type",
type=str,
@@ -292,14 +306,8 @@ def main():
parser.add_argument(
"--record-time-s",
type=float,
default=2.0,
help="Time duration to attempt capturing frames. Default: 2 seconds.",
)
parser.add_argument(
"--warmup-s",
type=int,
default=1,
help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
default=6.0,
help="Time duration to attempt capturing frames. Default: 6 seconds.",
)
args = parser.parse_args()
save_images_from_all_cameras(**vars(args))
-1
View File
@@ -165,7 +165,6 @@ from lerobot.robots import ( # noqa: F401
earthrover_mini_plus,
hope_jr,
koch_follower,
lekiwi,
omx_follower,
openarm_follower,
reachy2,
+92 -37
View File
@@ -22,8 +22,7 @@ import dataclasses
import logging
import sys
import time
from collections.abc import Iterator
from contextlib import contextmanager, nullcontext
from contextlib import nullcontext
from pprint import pformat
from typing import TYPE_CHECKING, Any
@@ -42,6 +41,7 @@ from lerobot.common.train_utils import (
load_fsdp_optimizer_state,
load_training_batch_size,
load_training_num_processes,
load_training_num_workers,
load_training_state,
push_checkpoint_to_hub,
save_checkpoint,
@@ -58,7 +58,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
from lerobot.rewards import make_reward_pre_post_processors
from lerobot.utils.collate import lerobot_collate_fn
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
@@ -69,31 +69,16 @@ from lerobot.utils.utils import (
inside_slurm,
)
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
from .lerobot_eval import eval_policy_all
@contextmanager
def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
"""Create evaluation environments for one run and always dispose of them."""
envs = make_env(
cfg.env,
n_envs=cfg.eval.batch_size,
use_async_envs=cfg.eval.use_async_envs,
)
try:
yield envs
finally:
close_envs(envs)
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
def _dataloader_worker_kwargs(
cfg: TrainPipelineConfig,
*,
num_workers: int | None = None,
) -> dict[str, Any]:
"""Return worker-only DataLoader options, disabling them for single-process loading."""
workers_enabled = cfg.num_workers > 0
workers_enabled = (cfg.num_workers if num_workers is None else num_workers) > 0
return {
"prefetch_factor": cfg.prefetch_factor if workers_enabled else None,
"persistent_workers": cfg.persistent_workers and workers_enabled,
@@ -227,9 +212,11 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.job.is_remote:
return submit_to_hf(cfg)
from lerobot.utils.import_utils import require_package
require_package("accelerate", extra="training")
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs, DistributedType
from accelerate.utils import DistributedDataParallelKwargs, DistributedType, send_to_device
cfg.validate()
@@ -295,6 +282,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if not is_main_process:
dataset, eval_dataset = make_train_eval_datasets(cfg)
# Create environment used for evaluating checkpoints during training on simulation data.
# On real-world data, no need to create an environment as evaluations are done outside train.py,
# using the eval.py instead, with gym_dora environment and dora-rs.
eval_env = None
if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process:
logging.info("Creating env")
eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
if cfg.is_reward_model_training:
if is_main_process:
logging.info("Creating reward model")
@@ -322,7 +317,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.peft is not None:
if cfg.is_reward_model_training:
raise ValueError("PEFT is only supported for policy training. ")
require_package("peft", extra="peft")
from peft import PeftModel
if isinstance(policy, PeftModel):
logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
@@ -348,6 +343,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
preprocessor_overrides = {
"device_processor": {"device": device.type},
"normalizer_processor": {
"stats": dataset.meta.stats,
"features": {**policy.config.input_features, **policy.config.output_features},
"norm_map": policy.config.normalization_mapping,
},
@@ -355,17 +351,11 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
}
postprocessor_overrides = {
"unnormalizer_processor": {
"stats": dataset.meta.stats,
"features": policy.config.output_features,
"norm_map": policy.config.normalization_mapping,
},
}
# On resume, the checkpoint's saved processor stats are authoritative: they may have
# been adapted by the policy (e.g. EVO1 pads state/action stats to max_state_dim),
# and force-feeding raw dataset stats over them crashes normalization (#4006).
# This mirrors the `dataset_stats` kwarg above, which is also skipped on resume.
if not cfg.resume:
preprocessor_overrides["normalizer_processor"]["stats"] = dataset.meta.stats
postprocessor_overrides["unnormalizer_processor"]["stats"] = dataset.meta.stats
if getattr(active_cfg, "use_relative_actions", False):
preprocessor_overrides["relative_actions_processor"] = {
"enabled": True,
@@ -484,6 +474,53 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
else:
shuffle = True
sampler = None
# One dedicated DataLoader process owns the rank-level planner/cache. Its bounded result
# queue provides decoded-batch prefetch; max_num_shards controls the internal Parquet/range
# fetch executors independently.
train_num_workers = min(cfg.num_workers, 1)
if cfg.num_workers > 1 and is_main_process:
logging.info(
"Using one streaming DataLoader worker per rank; %d configured workers remain "
"available as the dataset's internal fetch concurrency.",
cfg.num_workers,
)
rank_frames = dataset.num_frames_for_rank(
accelerator.process_index,
accelerator.num_processes,
train_num_workers,
)
if rank_frames == 0:
raise ValueError("This rank owns no streaming episodes. Reduce the distributed world size.")
if cfg.resume and step > 0:
saved_num_processes = load_training_num_processes(cfg.checkpoint_path)
saved_batch_size = load_training_batch_size(cfg.checkpoint_path)
saved_num_workers = load_training_num_workers(cfg.checkpoint_path)
if saved_num_processes not in (None, accelerator.num_processes):
raise ValueError(
"Sample-exact streaming resume requires the checkpoint world size "
f"({saved_num_processes}) to match the current world size "
f"({accelerator.num_processes})."
)
if saved_batch_size not in (None, cfg.batch_size):
raise ValueError(
"Sample-exact streaming resume requires the checkpoint batch size "
f"({saved_batch_size}) to match the current batch size ({cfg.batch_size})."
)
if saved_num_workers not in (None, train_num_workers):
raise ValueError(
"Sample-exact streaming resume requires the checkpoint DataLoader worker count "
f"({saved_num_workers}) to match the current count ({train_num_workers})."
)
stream_offset = step * (saved_batch_size or cfg.batch_size)
dataset.load_state_dict(
{
"epoch": 0,
"offset": stream_offset,
"batch_size": cfg.batch_size,
}
)
if is_main_process:
logging.info(f"Resuming streaming data order at local sample {stream_offset}")
# Only swap in the language-aware collate when the dataset actually
# declares language columns; otherwise stay on PyTorch's default
@@ -491,14 +528,17 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None
dataloader = torch.utils.data.DataLoader(
dataset,
num_workers=cfg.num_workers,
num_workers=train_num_workers if cfg.dataset.streaming else cfg.num_workers,
batch_size=cfg.batch_size,
shuffle=shuffle and not cfg.dataset.streaming,
sampler=sampler,
pin_memory=device.type == "cuda",
drop_last=False,
collate_fn=collate_fn,
**_dataloader_worker_kwargs(cfg),
**_dataloader_worker_kwargs(
cfg,
num_workers=train_num_workers if cfg.dataset.streaming else cfg.num_workers,
),
)
# Build eval dataloader if a held-out split exists
@@ -529,7 +569,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
# Prepare everything with accelerator
accelerator.wait_for_everyone()
if eval_dataloader is not None:
if cfg.dataset.streaming:
# StreamingLeRobotDataset already assigns complete episodes disjointly to ranks and workers.
# Preparing this loader would make Accelerate shard its batches a second time.
if eval_dataloader is not None:
policy, optimizer, lr_scheduler, eval_dataloader = accelerator.prepare(
policy, optimizer, lr_scheduler, eval_dataloader
)
else:
policy, optimizer, lr_scheduler = accelerator.prepare(policy, optimizer, lr_scheduler)
elif eval_dataloader is not None:
policy, optimizer, dataloader, lr_scheduler, eval_dataloader = accelerator.prepare(
policy, optimizer, dataloader, lr_scheduler, eval_dataloader
)
@@ -592,6 +641,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
for _ in range(step, cfg.steps):
start_time = time.perf_counter()
batch = next(dl_iter)
if cfg.dataset.streaming:
batch = send_to_device(batch, device, non_blocking=device.type == "cuda")
for cam_key in dataset.meta.camera_keys:
if cam_key in batch and batch[cam_key].dtype == torch.uint8:
batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0
@@ -688,6 +739,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
postprocessor=postprocessor,
num_processes=accelerator.num_processes,
batch_size=cfg.batch_size,
num_workers=train_num_workers if cfg.dataset.streaming else cfg.num_workers,
model_state_dict=model_state_dict,
optim_state_dict=optim_state_dict,
)
@@ -707,7 +759,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
step_id = get_step_identifier(step, cfg.steps)
logging.info(f"Eval policy at step {step}")
with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
with torch.no_grad(), accelerator.autocast():
eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env
policy=accelerator.unwrap_model(policy),
@@ -755,6 +807,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
progbar.close()
if eval_env:
close_envs(eval_env)
is_fsdp = accelerator.distributed_type == DistributedType.FSDP
model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
if is_main_process:
+56 -58
View File
@@ -45,7 +45,6 @@ lerobot-train-tokenizer \
"""
import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING
@@ -64,9 +63,6 @@ else:
from lerobot.configs import NormalizationMode, parser
from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
@dataclass
@@ -278,8 +274,11 @@ def process_episode(args):
return action_chunks
except Exception:
logger.exception("Error processing episode %s", ep_idx)
except Exception as e:
print(f"Error processing episode {ep_idx}: {e}")
import traceback
traceback.print_exc()
return None
@@ -301,10 +300,10 @@ def train_fast_tokenizer(
Returns:
Trained FAST tokenizer
"""
logger.info(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
logger.info(f"Action chunk shape: {action_chunks.shape}")
logger.info(f"Vocab size: {vocab_size}")
logger.info(f"DCT scale: {scale}")
print(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
print(f"Action chunk shape: {action_chunks.shape}")
print(f"Vocab size: {vocab_size}")
print(f"DCT scale: {scale}")
# download the tokenizer source code (not pretrained weights)
# we'll train a new tokenizer on our own data
@@ -315,7 +314,7 @@ def train_fast_tokenizer(
# train the new tokenizer on our action data using .fit()
# this trains the BPE tokenizer on DCT coefficients
logger.info("Training new tokenizer (this may take a few minutes)...")
print("Training new tokenizer (this may take a few minutes)...")
tokenizer = base_tokenizer.fit(
action_data_list,
scale=scale,
@@ -323,21 +322,21 @@ def train_fast_tokenizer(
time_horizon=action_chunks.shape[1], # action_horizon
action_dim=action_chunks.shape[2], # encoded dimensions
)
logger.info("✓ Tokenizer training complete!")
print("✓ Tokenizer training complete!")
# validate it works
sample_chunk = action_chunks[0]
encoded = tokenizer(sample_chunk[None])[0]
if isinstance(encoded, list):
encoded = np.array(encoded)
logger.info(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
print(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
return tokenizer
def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
"""Compute compression statistics."""
logger.info("\nComputing compression statistics...")
print("\nComputing compression statistics...")
# sample for stats (use max 1000 chunks for speed)
sample_size = min(1000, len(action_chunks))
@@ -367,12 +366,12 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
"max_token_length": float(np.max(token_lengths)),
}
logger.info("Compression Statistics:")
logger.info(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
logger.info(f" Mean token length: {stats['mean_token_length']:.1f}")
logger.info(f" P99 token length: {stats['p99_token_length']:.0f}")
logger.info(f" Min token length: {stats['min_token_length']:.0f}")
logger.info(f" Max token length: {stats['max_token_length']:.0f}")
print("Compression Statistics:")
print(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
print(f" Mean token length: {stats['mean_token_length']:.1f}")
print(f" P99 token length: {stats['p99_token_length']:.0f}")
print(f" Min token length: {stats['min_token_length']:.0f}")
print(f" Max token length: {stats['max_token_length']:.0f}")
return stats
@@ -386,9 +385,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
cfg: TokenizerTrainingConfig dataclass with all configuration parameters
"""
# load dataset
logger.info(f"Loading dataset: {cfg.repo_id}")
print(f"Loading dataset: {cfg.repo_id}")
dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
logger.info(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
# parse normalization mode
try:
@@ -398,7 +397,7 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
f"Invalid normalization_mode: {cfg.normalization_mode}. "
f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
) from err
logger.info(f"Normalization mode: {norm_mode.value}")
print(f"Normalization mode: {norm_mode.value}")
# parse encoded dimensions
encoded_dim_ranges = []
@@ -407,38 +406,38 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
encoded_dim_ranges.append((start, end))
total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
logger.info(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
print(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
# parse relative dimensions
relative_dim_list = None
if cfg.relative_dims is not None and cfg.relative_dims.strip():
relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
logger.info(f"Relative dimensions: {relative_dim_list}")
print(f"Relative dimensions: {relative_dim_list}")
else:
logger.info("No relative dimensions specified")
print("No relative dimensions specified")
logger.info(f"Use relative transform: {cfg.use_relative_transform}")
print(f"Use relative transform: {cfg.use_relative_transform}")
if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
logger.warning(
print(
"Warning: use_relative_transform=True but no relative_dims specified. "
"No relative transform will be applied."
)
logger.info(f"Action horizon: {cfg.action_horizon}")
logger.info(f"State key: {cfg.state_key}")
print(f"Action horizon: {cfg.action_horizon}")
print(f"State key: {cfg.state_key}")
# determine episodes to process
num_episodes = dataset.num_episodes
if cfg.max_episodes is not None:
num_episodes = min(cfg.max_episodes, num_episodes)
logger.info(f"Processing {num_episodes} episodes...")
print(f"Processing {num_episodes} episodes...")
# process episodes sequentially (to avoid pickling issues with dataset)
all_chunks = []
for ep_idx in range(num_episodes):
if ep_idx % 10 == 0:
logger.info(f" Processing episode {ep_idx}/{num_episodes}...")
print(f" Processing episode {ep_idx}/{num_episodes}...")
chunks = process_episode(
(
@@ -456,19 +455,19 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
# concatenate all chunks
all_chunks = np.concatenate(all_chunks, axis=0)
logger.info(f"Collected {len(all_chunks)} action chunks")
print(f"Collected {len(all_chunks)} action chunks")
# extract only encoded dimensions FIRST (before normalization)
encoded_chunks = []
for start, end in encoded_dim_ranges:
encoded_chunks.append(all_chunks[:, :, start:end])
encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded]
logger.info(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
print(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
# apply normalization to encoded dimensions
logger.info("\nBefore normalization - overall stats:")
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
print("\nBefore normalization - overall stats:")
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
# get normalization stats from dataset
norm_stats = dataset.meta.stats
@@ -490,9 +489,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
encoded_stats[stat_name] = stat_array[encoded_dim_indices]
if encoded_stats:
logger.info(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
print(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
for stat_name, stat_values in encoded_stats.items():
logger.info(
print(
f" {stat_name}: shape={stat_values.shape}, "
f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
)
@@ -500,27 +499,27 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
# apply normalization based on mode
try:
encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
logger.info(f"\nApplied {norm_mode.value} normalization")
print(f"\nApplied {norm_mode.value} normalization")
except ValueError as e:
logger.warning(f"Warning: {e}. Using raw actions without normalization.")
print(f"Warning: {e}. Using raw actions without normalization.")
logger.info("\nAfter normalization - overall stats:")
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
print("\nAfter normalization - overall stats:")
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
logger.info("\nPer-dimension stats (after normalization):")
print("\nPer-dimension stats (after normalization):")
for d in range(encoded_chunks.shape[-1]):
dim_data = encoded_chunks[:, :, d]
logger.info(
print(
f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
)
else:
logger.warning("Warning: Could not extract stats for encoded dimensions, using raw actions")
print("Warning: Could not extract stats for encoded dimensions, using raw actions")
else:
logger.warning("Warning: No normalization stats found in dataset, using raw actions")
print("Warning: No normalization stats found in dataset, using raw actions")
logger.info(f"Encoded chunks shape: {encoded_chunks.shape}")
print(f"Encoded chunks shape: {encoded_chunks.shape}")
# train FAST tokenizer
tokenizer = train_fast_tokenizer(
@@ -562,8 +561,8 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
with open(output_path / "metadata.json", "w") as f:
json.dump(metadata, f, indent=2)
logger.info(f"\nSaved FAST tokenizer to {output_path}")
logger.info(f"Metadata: {json.dumps(metadata, indent=2)}")
print(f"\nSaved FAST tokenizer to {output_path}")
print(f"Metadata: {json.dumps(metadata, indent=2)}")
# push to Hugging Face Hub if requested
if cfg.push_to_hub:
@@ -571,10 +570,10 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
hub_repo_id = cfg.hub_repo_id
if hub_repo_id is None:
hub_repo_id = output_path.name
logger.info(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
print(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
logger.info(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
logger.info(f" Private: {cfg.hub_private}")
print(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
print(f" Private: {cfg.hub_private}")
try:
# use the tokenizer's push_to_hub method
@@ -594,15 +593,14 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
commit_message="Upload tokenizer metadata",
)
logger.info(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
print(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
except Exception as e:
logger.error(f"Error pushing to hub: {e}")
logger.error(" Make sure you're logged in with `huggingface-cli login`")
print(f"Error pushing to hub: {e}")
print(" Make sure you're logged in with `huggingface-cli login`")
def main():
"""CLI entry point that parses arguments and runs the tokenizer training."""
init_logging()
train_tokenizer()
+9
View File
@@ -0,0 +1,9 @@
# 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
"""Low-level primitives for training-time dataset streaming."""
+519
View File
@@ -0,0 +1,519 @@
# 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 io
import logging
import threading
import time
from collections import OrderedDict
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
from pathlib import Path
from types import SimpleNamespace
from typing import Any
from lerobot.streaming.manifest import EpisodeVideoManifest
from lerobot.streaming.mp4 import Mp4SampleSlice, synthesize_mp4
from lerobot.streaming.range_fetch import make_range_fetcher
logger = logging.getLogger(__name__)
class EpisodeByteCache:
def __init__(
self,
manifest: EpisodeVideoManifest,
data_root: str | Path,
*,
byte_budget: int = 80 * 1024**3,
workers: int = 8,
range_backend: str = "fsspec",
native_http_connections: int | None = None,
native_http_timeout: float = 60.0,
native_http_retries: int = 4,
native_http_subranges: int = 1,
open_decoders: bool = True,
max_open_decoders: int = 64,
video_backend: str = "torchcodec",
tolerance_s: float = 1e-4,
token: str | bool | None = None,
):
if byte_budget <= 0:
raise ValueError("byte_budget must be positive")
if max_open_decoders <= 0:
raise ValueError("max_open_decoders must be positive")
if video_backend == "video_reader":
video_backend = "pyav"
if video_backend not in {"torchcodec", "pyav"}:
raise ValueError(f"Unsupported video backend: {video_backend}")
if tolerance_s <= 0:
raise ValueError("tolerance_s must be positive")
self.manifest = manifest
self.fetcher = make_range_fetcher(
data_root,
range_backend=range_backend,
workers=workers,
native_http_connections=native_http_connections,
native_http_timeout=native_http_timeout,
native_http_retries=native_http_retries,
native_http_subranges=native_http_subranges,
token=token,
)
self.byte_budget = byte_budget
self.open_decoders = open_decoders
self.max_open_decoders = max_open_decoders
self.video_backend = video_backend
self.tolerance_s = tolerance_s
self._pool = ThreadPoolExecutor(max_workers=workers)
self._cache: OrderedDict[tuple[int, str], dict[str, Any]] = OrderedDict()
self._decoders: OrderedDict[tuple[int, str], Any] = OrderedDict()
self._decoder_locks: dict[tuple[int, str], threading.Lock] = {}
self._futures: dict[tuple[int, str], Future[dict[str, Any]]] = {}
self._retained_episodes: dict[int, int] = {}
self._decoder_fallback_count = 0
self._fallback_decoders: set[tuple[int, str]] = set()
self._fallback_warning_emitted = False
self._bytes = 0
self._lock = threading.Lock()
self._timing_totals = {
"lookup_s": 0.0,
"fetch_s": 0.0,
"synthesize_s": 0.0,
"store_s": 0.0,
"jobs": 0.0,
}
def close(self) -> None:
self._pool.shutdown(wait=True, cancel_futures=True)
with self._lock:
decoders = list(self._decoders.values())
self._cache.clear()
self._decoders.clear()
self._decoder_locks.clear()
self._futures.clear()
self._retained_episodes.clear()
self._fallback_decoders.clear()
self._bytes = 0
for decoder in decoders:
_close_decoder(decoder)
self.fetcher.close()
def __enter__(self) -> EpisodeByteCache:
return self
def __exit__(self, *_exc) -> None:
self.close()
def submit_prefetch(self, episode_index: int) -> None:
for camera_key in self.manifest.video_keys:
self._submit(episode_index, camera_key)
def retain_episode(self, episode_index: int) -> None:
with self._lock:
self._retained_episodes[episode_index] = self._retained_episodes.get(episode_index, 0) + 1
def release_episode(self, episode_index: int) -> None:
with self._lock:
count = self._retained_episodes.get(episode_index, 0)
if count <= 1:
self._retained_episodes.pop(episode_index, None)
else:
self._retained_episodes[episode_index] = count - 1
self._evict_locked()
@property
def resident_bytes(self) -> int:
with self._lock:
return self._bytes
@property
def open_decoder_count(self) -> int:
with self._lock:
return len(self._decoders)
@property
def decoder_fallback_count(self) -> int:
with self._lock:
return self._decoder_fallback_count
def ensure_ready(self, episode_index: int) -> None:
for camera_key in self.manifest.video_keys:
self.get_bytes(episode_index, camera_key)
if self.open_decoders:
self.get_decoder(episode_index, camera_key)
def is_ready(self, episode_index: int) -> bool:
"""Non-blocking: True when every camera of the episode is fetched (cached or future done).
Lets a consumer swap in replacements only when they are already resident, instead of
blocking the training hot path on a remote fetch (head-of-line stall).
"""
for camera_key in self.manifest.video_keys:
key = (episode_index, camera_key)
with self._lock:
if key in self._cache:
continue
future = self._futures.get(key)
if future is None or not future.done():
return False
return True
def get_bytes(self, episode_index: int, camera_key: str) -> bytes:
return self._get_entry(episode_index, camera_key)["bytes"]
def get_decoder(self, episode_index: int, camera_key: str) -> Any:
key = (episode_index, camera_key)
entry = self._get_entry(episode_index, camera_key)
with self._lock:
decoder = self._decoders.get(key)
if decoder is not None:
self._decoders.move_to_end(key)
return decoder
decoder = self._open_decoder(key, entry["bytes"])
with self._lock:
existing = self._decoders.get(key)
if existing is not None:
self._decoders.move_to_end(key)
_close_decoder(decoder)
return existing
self._decoders[key] = decoder
while len(self._decoders) > self.max_open_decoders:
evicted_key, evicted_decoder = self._decoders.popitem(last=False)
self._decoder_locks.pop(evicted_key, None)
self._fallback_decoders.discard(evicted_key)
_close_decoder(evicted_decoder)
return decoder
def _open_decoder(self, key: tuple[int, str], data: bytes) -> Any:
try:
if self.video_backend == "torchcodec":
return open_video_decoder(io.BytesIO(data))
return open_video_decoder(io.BytesIO(data), backend=self.video_backend)
except Exception as primary_error:
if self.video_backend != "torchcodec":
raise
try:
decoder = open_video_decoder(io.BytesIO(data), backend="pyav")
except Exception as fallback_error:
raise RuntimeError(
"Both TorchCodec and PyAV rejected synthesized episode video "
f"{key}: TorchCodec error: {primary_error}"
) from fallback_error
with self._lock:
self._decoder_fallback_count += 1
self._fallback_decoders.add(key)
should_warn = not self._fallback_warning_emitted
self._fallback_warning_emitted = True
if should_warn:
logger.warning(
"TorchCodec rejected a synthesized episode MP4; using the bounded PyAV "
"decoder fallback for affected videos. First error: %s",
primary_error,
)
else:
logger.debug("Using PyAV decoder fallback for synthesized episode video %s", key)
return decoder
def get_frames(self, episode_index: int, camera_key: str, timestamps: list[float]):
key = (episode_index, camera_key)
span = self.manifest.lookup(episode_index, camera_key)
local_ts = [ts - span.source_start_pts for ts in timestamps]
decoder, release = self._decoder_for_frames(episode_index, camera_key)
with self._lock:
uses_pyav_timestamps = self.video_backend == "pyav" and key not in self._fallback_decoders
decode_lock = self._decoder_locks.setdefault(key, threading.Lock())
try:
with decode_lock:
if uses_pyav_timestamps:
return decoder.get_frames_played_at(local_ts, tolerance_s=self.tolerance_s).data
metadata = decoder.metadata
fps = getattr(metadata, "average_fps", None)
if fps is None:
duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9)
fps = metadata.num_frames / duration
num_frames = getattr(metadata, "num_frames", None)
if num_frames is None:
duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9)
num_frames = round(duration * fps)
last_index = max(0, int(num_frames) - 1)
indices = [min(max(round(ts * fps), 0), last_index) for ts in local_ts]
return decoder.get_frames_at(indices=indices).data
finally:
if release is not None:
release()
def _decoder_for_frames(
self, episode_index: int, camera_key: str
) -> tuple[Any, Callable[[], None] | None]:
key = (episode_index, camera_key)
while True:
decoder = self.get_decoder(episode_index, camera_key)
acquire = getattr(decoder, "acquire", None)
if acquire is None:
return decoder, None
try:
acquire()
except RuntimeError:
# The decoder was evicted between lookup and lease acquisition. Remove a stale
# cached reference if it raced with close, then retry with a fresh decoder.
with self._lock:
if self._decoders.get(key) is decoder:
self._decoders.pop(key)
self._decoder_locks.pop(key, None)
self._fallback_decoders.discard(key)
continue
return decoder, decoder.release
def timing_summary(self) -> dict[str, float]:
with self._lock:
summary = dict(self._timing_totals)
summary["decoder_fallbacks"] = float(self._decoder_fallback_count)
fetcher_summary = getattr(self.fetcher, "timing_summary", None)
if fetcher_summary is not None:
summary.update(fetcher_summary())
return summary
def _submit(self, episode_index: int, camera_key: str) -> Future[dict[str, Any]]:
key = (episode_index, camera_key)
with self._lock:
if key in self._cache:
future: Future[dict[str, Any]] = Future()
future.set_result(self._cache[key])
return future
future = self._futures.get(key)
if future is None:
future = self._pool.submit(self._fetch_and_synthesize, episode_index, camera_key)
self._futures[key] = future
return future
def _get_entry(self, episode_index: int, camera_key: str) -> dict[str, Any]:
key = (episode_index, camera_key)
with self._lock:
entry = self._cache.get(key)
if entry is not None:
self._cache.move_to_end(key)
return entry
future = self._submit(episode_index, camera_key)
entry = future.result()
store_start = time.perf_counter()
with self._lock:
self._futures.pop(key, None)
existing = self._cache.get(key)
if existing is not None:
self._cache.move_to_end(key)
return existing
self._cache[key] = entry
self._bytes += len(entry["bytes"])
try:
self._evict_locked()
except MemoryError:
failed_entry = self._cache.pop(key, None)
if failed_entry is not None:
self._bytes -= len(failed_entry["bytes"])
raise
timings = entry.pop("_timings", None)
if timings is not None:
self._timing_totals["lookup_s"] += timings["lookup_s"]
self._timing_totals["fetch_s"] += timings["fetch_s"]
self._timing_totals["synthesize_s"] += timings["synthesize_s"]
self._timing_totals["store_s"] += time.perf_counter() - store_start
self._timing_totals["jobs"] += 1
return entry
def _evict_locked(self) -> None:
while self._bytes > self.byte_budget:
key = next(
(candidate for candidate in self._cache if candidate[0] not in self._retained_episodes),
None,
)
if key is None:
raise MemoryError(
f"Retained episode bytes exceed byte budget ({self._bytes} > {self.byte_budget})"
)
entry = self._cache.pop(key)
self._bytes -= len(entry["bytes"])
decoder = self._decoders.pop(key, None)
self._decoder_locks.pop(key, None)
self._fallback_decoders.discard(key)
if decoder is not None:
_close_decoder(decoder)
def _fetch_and_synthesize(self, episode_index: int, camera_key: str) -> dict[str, Any]:
lookup_start = time.perf_counter()
span = self.manifest.lookup(episode_index, camera_key)
file_record = self.manifest.file_lookup(span.file_id)
sample_slice = Mp4SampleSlice(
sample_lo=span.sample_lo,
sample_hi=span.sample_hi,
byte_offset=span.mdat_offset,
byte_length=span.mdat_length,
source_start_pts=span.source_start_pts,
)
lookup_s = time.perf_counter() - lookup_start
fetch_start = time.perf_counter()
payload = self.fetcher.read_range(file_record.file_path, span.mdat_offset, span.mdat_length)
fetch_s = time.perf_counter() - fetch_start
if len(payload) != span.mdat_length:
raise OSError(
f"Short read for {file_record.file_path}: expected {span.mdat_length}, got {len(payload)}"
)
synthesize_start = time.perf_counter()
mp4_bytes = synthesize_mp4(file_record.mp4, sample_slice, payload)
synthesize_s = time.perf_counter() - synthesize_start
entry: dict[str, Any] = {
"bytes": mp4_bytes,
"_timings": {
"lookup_s": lookup_s,
"fetch_s": fetch_s,
"synthesize_s": synthesize_s,
},
}
return entry
class _PyAVVideoDecoder:
"""Small seekable PyAV adapter matching the TorchCodec calls used by the byte cache."""
def __init__(self, file_like_or_bytesio: Any):
import av
self._source = file_like_or_bytesio
self._container = av.open(file_like_or_bytesio)
self._stream = self._container.streams.video[0]
average_rate = self._stream.average_rate
if average_rate is None:
raise ValueError("PyAV video stream does not expose an average frame rate")
self._fps = float(average_rate)
duration = (
float(self._stream.duration * self._stream.time_base)
if self._stream.duration is not None
else 0.0
)
self.metadata = SimpleNamespace(
average_fps=self._fps,
num_frames=int(self._stream.frames or round(duration * self._fps)),
begin_stream_seconds=0.0,
end_stream_seconds=duration,
)
self._decode_lock = threading.Lock()
self._state_lock = threading.Lock()
self._users = 0
self._close_requested = False
self._closed = False
def acquire(self) -> None:
with self._state_lock:
if self._close_requested or self._closed:
raise RuntimeError("PyAV decoder is closing")
self._users += 1
def release(self) -> None:
with self._state_lock:
self._users -= 1
if self._users < 0:
raise RuntimeError("Unbalanced PyAV decoder release")
if self._users == 0 and self._close_requested:
self._close_resources()
def get_frames_at(self, *, indices: list[int]) -> SimpleNamespace:
if not indices:
import torch
return SimpleNamespace(data=torch.empty((0, 3, 0, 0), dtype=torch.uint8))
timestamps = [index / self._fps for index in indices]
return self._get_frames_played_at(timestamps, tolerance_s=0.5 / self._fps + 1e-6)
def get_frames_played_at(
self,
timestamps: list[float],
*,
tolerance_s: float,
) -> SimpleNamespace:
return self._get_frames_played_at(timestamps, tolerance_s=tolerance_s)
def _get_frames_played_at(
self,
timestamps: list[float],
*,
tolerance_s: float,
) -> SimpleNamespace:
import torch
first_ts = min(timestamps)
last_ts = max(timestamps)
loaded_frames: list[torch.Tensor] = []
loaded_ts: list[float] = []
with self._decode_lock:
self._container.seek(
round(first_ts / self._stream.time_base) - 1,
backward=True,
any_frame=False,
stream=self._stream,
)
for frame in self._container.decode(self._stream):
if frame.pts is None:
continue
current_ts = float(frame.pts * self._stream.time_base)
array = frame.to_ndarray(format="rgb24")
loaded_frames.append(torch.from_numpy(array).permute(2, 0, 1).contiguous())
loaded_ts.append(current_ts)
if current_ts >= last_ts:
break
if not loaded_frames:
raise ValueError(f"PyAV decoded no frames for timestamps {timestamps}")
query_ts = torch.tensor(timestamps)
loaded_ts_tensor = torch.tensor(loaded_ts)
distances = torch.cdist(query_ts[:, None], loaded_ts_tensor[:, None], p=1)
minimum, closest = distances.min(1)
if not (minimum <= tolerance_s).all():
raise ValueError(
f"PyAV frame timestamps exceed tolerance {tolerance_s}: "
f"queries={query_ts}, decoded={loaded_ts_tensor}"
)
return SimpleNamespace(data=torch.stack([loaded_frames[index] for index in closest]))
def close(self) -> None:
with self._state_lock:
self._close_requested = True
if self._users == 0:
self._close_resources()
def _close_resources(self) -> None:
if self._closed:
return
self._container.close()
close = getattr(self._source, "close", None)
if close is not None:
close()
self._closed = True
def _close_decoder(decoder: Any) -> None:
close = getattr(decoder, "close", None)
if close is not None:
try:
close()
except Exception:
logger.debug("Failed to close video decoder", exc_info=True)
def open_video_decoder(file_like_or_bytesio, frame_mappings=None, *, backend: str = "torchcodec"):
if frame_mappings is not None:
raise ValueError("Synthesized episode videos use a local timeline; pass frame_mappings=None.")
if backend == "pyav":
return _PyAVVideoDecoder(file_like_or_bytesio)
if backend != "torchcodec":
raise ValueError(f"Unsupported video backend: {backend}")
from torchcodec.decoders import VideoDecoder
return VideoDecoder(file_like_or_bytesio, seek_mode="approximate")
+140
View File
@@ -0,0 +1,140 @@
# 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
"""Pure episode-scoped Parquet reads for training-time dataset streaming."""
from __future__ import annotations
import posixpath
import threading
import time
from collections.abc import Sequence
from pathlib import Path
from typing import Any
import fsspec
import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.parquet as pq
class EpisodeParquetReader:
"""Read complete episodes with column projection from local or fsspec roots."""
def __init__(
self,
data_root: str | Path,
*,
columns: Sequence[str],
token: str | bool | None = None,
max_retries: int = 4,
retry_backoff_s: float = 0.05,
):
if not columns:
raise ValueError("EpisodeParquetReader requires at least one projected column")
self.columns = tuple(dict.fromkeys(columns))
self._read_columns = (
self.columns if "episode_index" in self.columns else (*self.columns, "episode_index")
)
data_root_str = str(data_root)
if max_retries < 0:
raise ValueError("max_retries must be non-negative")
if retry_backoff_s < 0:
raise ValueError("retry_backoff_s must be non-negative")
self._max_retries = max_retries
self._retry_backoff_s = retry_backoff_s
self._open_lock = threading.Lock() if data_root_str.startswith("hf://") else None
storage_options = {"token": token} if token is not None and data_root_str.startswith("hf://") else {}
self._filesystem, self._root_path = fsspec.core.url_to_fs(data_root_str, **storage_options)
def read_episode(
self,
relative_path: str | Path,
*,
episode_index: int,
expected_rows: int,
) -> pa.Table:
if expected_rows <= 0:
raise ValueError(f"Episode {episode_index} must contain at least one row")
path = posixpath.join(self._root_path.rstrip("/"), str(relative_path).lstrip("/"))
with self._open_with_retry(path) as source:
parquet = pq.ParquetFile(source)
available = set(parquet.schema_arrow.names)
missing = sorted(set(self._read_columns) - available)
if missing:
raise ValueError(f"Parquet file {relative_path} is missing projected columns: {missing}")
row_group = self._matching_row_group(parquet, episode_index)
table = (
parquet.read_row_group(row_group, columns=list(self._read_columns))
if row_group is not None
else parquet.read(columns=list(self._read_columns))
)
if row_group is None:
table = table.filter(pc.equal(table.column("episode_index"), episode_index))
self._validate_complete_episode(table, episode_index, expected_rows, relative_path)
if "episode_index" not in self.columns:
table = table.drop_columns(["episode_index"])
return table
def _open_with_retry(self, path: str) -> Any:
for attempt in range(self._max_retries + 1):
try:
if self._open_lock is None:
return self._filesystem.open(path, "rb")
with self._open_lock:
return self._filesystem.open(path, "rb")
except FileNotFoundError:
if attempt == self._max_retries:
raise
self._filesystem.invalidate_cache(posixpath.dirname(path))
time.sleep(self._retry_backoff_s * 2**attempt)
raise RuntimeError("unreachable")
@staticmethod
def _matching_row_group(parquet: pq.ParquetFile, episode_index: int) -> int | None:
episode_column = next(
index
for index in range(parquet.metadata.num_columns)
if parquet.metadata.schema.column(index).path == "episode_index"
)
matches = []
for row_group in range(parquet.metadata.num_row_groups):
statistics = parquet.metadata.row_group(row_group).column(episode_column).statistics
if (
statistics is not None
and statistics.has_min_max
and int(statistics.min) == episode_index
and int(statistics.max) == episode_index
):
matches.append(row_group)
return matches[0] if len(matches) == 1 else None
@staticmethod
def _validate_complete_episode(
table: pa.Table,
episode_index: int,
expected_rows: int,
relative_path: str | Path,
) -> None:
actual_rows = len(table)
if actual_rows != expected_rows:
raise ValueError(
f"Parquet episode {episode_index} in {relative_path}: "
f"expected {expected_rows} rows, found {actual_rows}"
)
episodes = table.column("episode_index").to_pylist()
if any(int(value) != episode_index for value in episodes):
raise ValueError(f"Parquet file {relative_path} returned rows outside episode {episode_index}")
if "frame_index" in table.column_names:
frame_indices = [int(value) for value in table.column("frame_index").to_pylist()]
if frame_indices != list(range(expected_rows)):
raise ValueError(
f"Parquet episode {episode_index} in {relative_path} has non-contiguous frame indices"
)
+168
View File
@@ -0,0 +1,168 @@
# 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
from collections.abc import Mapping, Sequence
import numpy as np
class ExactCoveragePool:
"""Deterministic, exactly-once frame coverage over a byte-cache episode pool.
A with-replacement pool never guarantees a full
epoch: frames are drawn randomly and episodes rotate on a fixed cadence. This planner instead
enumerates *every frame of every episode exactly once per epoch* while keeping at most
``pool_size`` episodes resident, so batch mixing stays high but coverage is complete and
reproducible.
Mechanics (this is the "evict only when all frames sampled" model):
- Episodes are admitted in a seeded global permutation until either ``pool_size`` or the
optional indexed-byte budget is reached.
- Each resident episode carries a seeded shuffle of its own frame indices.
- Each draw picks a resident episode with probability proportional to its *remaining* frames
(i.e. a uniform draw over all remaining frames in the pool, the map-style ideal) and pops
one frame.
- An episode is evicted only when its last frame is emitted; a new episode is then admitted.
- The epoch ends when the admission order is exhausted and every resident episode is drained.
Newly admitted episodes are surfaced via :attr:`newly_admitted` (drain it to drive prefetch)
and evictions via :attr:`evicted` (drain to release cache bytes). The planner does no I/O and
is fully unit-testable. It yields ``(episode_index, frame_index)``; map to a decode timestamp
with ``frame_index / max(frame_count - 1, 1)``.
Determinism: the order is a pure function of ``(seed, epoch)``, the episode frame counts, and
optional byte sizes/budget. Resume is a deterministic fast-forward: re-instantiate with the
same inputs and skip ``n`` samples (tabular only, no decode).
"""
def __init__(
self,
episode_frame_counts: Sequence[tuple[int, int]],
pool_size: int,
*,
seed: int,
epoch: int = 0,
episode_byte_sizes: Mapping[int, int] | None = None,
byte_budget: int | None = None,
):
self._counts = {int(ep): int(n) for ep, n in episode_frame_counts if int(n) > 0}
self._rng = np.random.default_rng([seed, epoch])
order = np.array(sorted(self._counts), dtype=np.int64)
self._rng.shuffle(order)
self.pool_size = max(1, pool_size)
self._byte_budget = byte_budget
if byte_budget is not None and byte_budget <= 0:
raise ValueError("byte_budget must be positive")
if byte_budget is not None and episode_byte_sizes is None:
raise ValueError("episode_byte_sizes are required when byte_budget is set")
self._byte_sizes = {
episode: int(episode_byte_sizes[episode]) if episode_byte_sizes is not None else 0
for episode in self._counts
}
if any(size < 0 for size in self._byte_sizes.values()):
raise ValueError("episode byte sizes must be non-negative")
if byte_budget is not None:
oversized = next(
((episode, size) for episode, size in self._byte_sizes.items() if size > byte_budget),
None,
)
if oversized is not None:
episode, size = oversized
raise ValueError(
f"Episode {episode} requires {size} bytes, exceeding the byte budget {byte_budget}"
)
# Preserve the full seeded order for benchmark/tooling compatibility. Byte-aware admission
# may temporarily skip an entry, but every episode remains in this deterministic frontier.
self.admission_order: list[int] = order.tolist()
self._pending: list[int] = list(self.admission_order)
self._admitted_count = 0
self._remaining: dict[int, tuple[np.ndarray, int]] = {}
self._remaining_total = 0
self._resident_bytes = 0
self.newly_admitted: list[int] = []
self.evicted: list[int] = []
self._admit_available()
def _admit_available(self) -> None:
while len(self._remaining) < self.pool_size and self._pending:
available_bytes = None if self._byte_budget is None else self._byte_budget - self._resident_bytes
pending_index = next(
(
index
for index, episode in enumerate(self._pending)
if available_bytes is None or self._byte_sizes[episode] <= available_bytes
),
None,
)
if pending_index is None:
return
episode = self._pending.pop(pending_index)
frame_count = self._counts[episode]
frames = np.arange(frame_count, dtype=np.int64)
self._rng.shuffle(frames)
self._remaining[episode] = (frames, frame_count)
self._remaining_total += frame_count
self._resident_bytes += self._byte_sizes[episode]
self._admitted_count += 1
self.newly_admitted.append(episode)
@property
def remaining_total(self) -> int:
return self._remaining_total
@property
def admitted_count(self) -> int:
"""Number of episodes pulled from the admission order so far (pool fills + rotations)."""
return self._admitted_count
@property
def resident(self) -> list[int]:
return list(self._remaining)
@property
def resident_bytes(self) -> int:
return self._resident_bytes
def prefetch_candidates(self, count: int) -> list[int]:
"""Return the next deterministic pending frontier without admitting it."""
if count <= 0:
return []
return self._pending[:count]
def __iter__(self) -> ExactCoveragePool:
return self
def __next__(self) -> tuple[int, int]:
if self._remaining_total == 0:
raise StopIteration
# Uniform draw over all remaining frames in the pool: walk the residents by cumulative
# remaining count. O(pool_size) per draw (~1024) -> negligible next to decode.
target = int(self._rng.integers(self._remaining_total))
chosen = None
for ep, (_frames, remaining) in self._remaining.items():
if target < remaining:
chosen = ep
break
target -= remaining
frames, remaining = self._remaining[chosen]
remaining -= 1
frame_index = int(frames[remaining])
self._remaining_total -= 1
if remaining == 0:
del self._remaining[chosen]
self.evicted.append(chosen)
self._resident_bytes -= self._byte_sizes[chosen]
self._admit_available()
else:
self._remaining[chosen] = (frames, remaining)
return chosen, frame_index
+351
View File
@@ -0,0 +1,351 @@
# 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 json
import logging
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any
import numpy as np
from lerobot.streaming.mp4 import (
Mp4Index,
Mp4SampleSlice,
fetch_mp4_index,
synthesized_mp4_size,
)
from lerobot.streaming.range_fetch import make_range_fetcher
if TYPE_CHECKING:
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.streaming.sidecar import SidecarSpec
@dataclass(frozen=True)
class EpisodeVideoSpan:
file_id: int
mdat_offset: int
mdat_length: int
first_pts: float
last_pts: float
frame_count: int
sample_lo: int
sample_hi: int
source_start_pts: float
@dataclass(frozen=True)
class VideoFileRecord:
file_path: str
file_size: int
mp4: Mp4Index
class EpisodeVideoManifest:
_FILE_SIDECAR_CACHE: dict[str, tuple[tuple[int, int], dict[str, VideoFileRecord]]] = {}
_FILE_SIDECAR_CACHE_LOCK = threading.Lock()
def __init__(
self,
*,
video_keys: list[str],
files: list[VideoFileRecord],
spans: dict[str, np.ndarray],
):
self.video_keys = list(video_keys)
self._camera_to_id = {key: idx for idx, key in enumerate(self.video_keys)}
self.files = files
self.spans = spans
@classmethod
def build(
cls,
meta: LeRobotDatasetMetadata,
data_root: str | Path,
*,
episode_indices: list[int] | range | None = None,
range_backend: str = "fsspec",
workers: int = 8,
header_probe_bytes: int = 4 * 1024 * 1024,
max_probe_bytes: int = 64 * 1024 * 1024,
keyframe_pad_s: float = 0.1,
keyframe_pad_fraction: float = 0.05,
sidecar_path: str | Path | None = None,
token: str | bool | None = None,
) -> EpisodeVideoManifest:
meta.ensure_readable()
video_keys = list(meta.video_keys)
if episode_indices is None:
episode_indices = range(int(meta.total_episodes))
rel_paths = sorted(
{str(meta.get_video_file_path(ep_idx, key)) for ep_idx in episode_indices for key in video_keys}
)
path_to_id = {path: idx for idx, path in enumerate(rel_paths)}
if sidecar_path is None:
files = cls._build_file_records(
rel_paths,
data_root,
range_backend=range_backend,
workers=workers,
header_probe_bytes=header_probe_bytes,
max_probe_bytes=max_probe_bytes,
token=token,
)
else:
records = cls.load_file_sidecar(sidecar_path)
missing = [path for path in rel_paths if path not in records]
if missing:
raise ValueError(
f"Sidecar {sidecar_path} is missing {len(missing)} files, first: {missing[0]}"
)
files = [records[path] for path in rel_paths]
total = int(meta.total_episodes)
num_cameras = len(video_keys)
spans: dict[str, np.ndarray] = {
"file_id": np.zeros((total, num_cameras), dtype=np.int32),
"mdat_offset": np.zeros((total, num_cameras), dtype=np.int64),
"mdat_length": np.zeros((total, num_cameras), dtype=np.int64),
"first_pts": np.zeros((total, num_cameras), dtype=np.float64),
"last_pts": np.zeros((total, num_cameras), dtype=np.float64),
"frame_count": np.zeros((total, num_cameras), dtype=np.int32),
"sample_lo": np.zeros((total, num_cameras), dtype=np.int32),
"sample_hi": np.zeros((total, num_cameras), dtype=np.int32),
"source_start_pts": np.zeros((total, num_cameras), dtype=np.float64),
}
for ep_idx in episode_indices:
ep = meta.episodes[ep_idx]
for cam_idx, key in enumerate(video_keys):
rel_path = str(meta.get_video_file_path(ep_idx, key))
file_id = path_to_id[rel_path]
mp4 = files[file_id].mp4
from_ts = float(ep[f"videos/{key}/from_timestamp"])
to_ts = float(ep[f"videos/{key}/to_timestamp"])
sample_slice = mp4.sample_slice(
from_ts,
to_ts,
keyframe_pad_s=keyframe_pad_s,
keyframe_pad_fraction=keyframe_pad_fraction,
file_size=files[file_id].file_size,
)
spans["file_id"][ep_idx, cam_idx] = file_id
spans["mdat_offset"][ep_idx, cam_idx] = sample_slice.byte_offset
spans["mdat_length"][ep_idx, cam_idx] = sample_slice.byte_length
spans["first_pts"][ep_idx, cam_idx] = from_ts
spans["last_pts"][ep_idx, cam_idx] = to_ts
spans["frame_count"][ep_idx, cam_idx] = sample_slice.sample_hi - sample_slice.sample_lo + 1
spans["sample_lo"][ep_idx, cam_idx] = sample_slice.sample_lo
spans["sample_hi"][ep_idx, cam_idx] = sample_slice.sample_hi
spans["source_start_pts"][ep_idx, cam_idx] = sample_slice.source_start_pts
return cls(video_keys=video_keys, files=files, spans=spans)
@staticmethod
def _build_file_records(
rel_paths: list[str],
data_root: str | Path,
*,
range_backend: str,
workers: int,
header_probe_bytes: int,
max_probe_bytes: int,
token: str | bool | None,
) -> list[VideoFileRecord]:
fetcher = make_range_fetcher(
data_root,
range_backend=range_backend,
workers=workers,
token=token,
)
def build_file(path: str) -> VideoFileRecord:
file_size = fetcher.info_size(path)
mp4 = fetch_mp4_index(
path,
fetcher.read_range,
file_size=file_size,
header_probe_bytes=header_probe_bytes,
max_probe_bytes=max_probe_bytes,
)
return VideoFileRecord(path, file_size, mp4)
try:
with ThreadPoolExecutor(max_workers=workers) as pool:
futures = {pool.submit(build_file, path): path for path in rel_paths}
records = []
progress_interval = max(1, len(futures) // 20)
for completed, future in enumerate(as_completed(futures), start=1):
records.append(future.result())
if completed == len(futures) or completed % progress_interval == 0:
logging.info("Indexed %d/%d MP4 files for streaming sidecar", completed, len(futures))
return sorted(records, key=lambda record: record.file_path)
finally:
fetcher.close()
@classmethod
def write_file_sidecar(
cls,
sidecar_path: str | Path,
rel_paths: list[str],
data_root: str | Path,
*,
spec: SidecarSpec,
range_backend: str = "native-http",
workers: int = 8,
header_probe_bytes: int = 4 * 1024 * 1024,
max_probe_bytes: int = 64 * 1024 * 1024,
token: str | bool | None = None,
) -> None:
records = cls._build_file_records(
sorted(set(rel_paths)),
data_root,
range_backend=range_backend,
workers=workers,
header_probe_bytes=header_probe_bytes,
max_probe_bytes=max_probe_bytes,
token=token,
)
cls.save_file_sidecar(sidecar_path, records, spec=spec)
@staticmethod
def save_file_sidecar(
sidecar_path: str | Path,
records: list[VideoFileRecord],
*,
spec: SidecarSpec,
) -> None:
sidecar_path = Path(sidecar_path)
sidecar_path.parent.mkdir(parents=True, exist_ok=True)
payload = {
"version": 2,
"sidecar": spec.with_source_files(
tuple((record.file_path, record.file_size) for record in records)
).to_dict(),
"files": [
{"file_path": record.file_path, "file_size": record.file_size, "mp4": record.mp4.to_dict()}
for record in records
],
}
arrays = {}
for file_idx, record in enumerate(records):
arrays[f"{file_idx}/sample_pts"] = record.mp4.sample_pts
arrays[f"{file_idx}/sample_durations"] = record.mp4.sample_durations
arrays[f"{file_idx}/sample_sizes"] = record.mp4.sample_sizes
arrays[f"{file_idx}/sample_offsets"] = record.mp4.sample_offsets
arrays[f"{file_idx}/sync_samples"] = record.mp4.sync_samples
np.savez_compressed(sidecar_path, manifest_json=json.dumps(payload).encode("utf-8"), **arrays)
cache_key = str(sidecar_path.expanduser())
with EpisodeVideoManifest._FILE_SIDECAR_CACHE_LOCK:
EpisodeVideoManifest._FILE_SIDECAR_CACHE.pop(cache_key, None)
@staticmethod
def load_file_sidecar_metadata(sidecar_path: str | Path) -> dict[str, Any]:
with np.load(sidecar_path, allow_pickle=False) as data:
payload = json.loads(bytes(data["manifest_json"]).decode("utf-8"))
if payload.get("version") != 2 or not isinstance(payload.get("sidecar"), dict):
raise ValueError(f"Unsupported MP4 sidecar schema in {sidecar_path}")
return payload["sidecar"]
@staticmethod
def validate_file_sidecar(sidecar_path: str | Path, spec: SidecarSpec) -> bool:
try:
from lerobot.streaming.sidecar import SidecarSpec
candidate = SidecarSpec.from_dict(EpisodeVideoManifest.load_file_sidecar_metadata(sidecar_path))
if not spec.matches(candidate):
return False
records = EpisodeVideoManifest.load_file_sidecar(sidecar_path)
except (OSError, ValueError, KeyError, TypeError, json.JSONDecodeError):
return False
expected = dict(candidate.source_files)
actual = {path: record.file_size for path, record in records.items()}
return actual == expected
@staticmethod
def load_file_sidecar(sidecar_path: str | Path) -> dict[str, VideoFileRecord]:
path = Path(sidecar_path).expanduser()
cache_key = str(path)
stat = path.stat()
signature = (stat.st_mtime_ns, stat.st_size)
with EpisodeVideoManifest._FILE_SIDECAR_CACHE_LOCK:
cached = EpisodeVideoManifest._FILE_SIDECAR_CACHE.get(cache_key)
if cached is not None and cached[0] == signature:
return cached[1]
with np.load(path, allow_pickle=False) as data:
payload = json.loads(bytes(data["manifest_json"]).decode("utf-8"))
if payload.get("version") != 2:
raise ValueError(f"Unsupported MP4 sidecar schema in {path}")
records = {}
for file_idx, item in enumerate(payload["files"]):
arrays = {
name: data[f"{file_idx}/{name}"]
for name in [
"sample_pts",
"sample_durations",
"sample_sizes",
"sample_offsets",
"sync_samples",
]
}
mp4 = Mp4Index.from_dict(item["mp4"], arrays)
records[item["file_path"]] = VideoFileRecord(item["file_path"], int(item["file_size"]), mp4)
with EpisodeVideoManifest._FILE_SIDECAR_CACHE_LOCK:
EpisodeVideoManifest._FILE_SIDECAR_CACHE[cache_key] = (signature, records)
return records
def camera_id(self, camera_key: str) -> int:
return self._camera_to_id[camera_key]
def lookup(self, episode_index: int, camera_key: str) -> EpisodeVideoSpan:
cam = self.camera_id(camera_key)
return EpisodeVideoSpan(
file_id=int(self.spans["file_id"][episode_index, cam]),
mdat_offset=int(self.spans["mdat_offset"][episode_index, cam]),
mdat_length=int(self.spans["mdat_length"][episode_index, cam]),
first_pts=float(self.spans["first_pts"][episode_index, cam]),
last_pts=float(self.spans["last_pts"][episode_index, cam]),
frame_count=int(self.spans["frame_count"][episode_index, cam]),
sample_lo=int(self.spans["sample_lo"][episode_index, cam]),
sample_hi=int(self.spans["sample_hi"][episode_index, cam]),
source_start_pts=float(self.spans["source_start_pts"][episode_index, cam]),
)
def file_lookup(self, file_id: int) -> VideoFileRecord:
return self.files[file_id]
def mp4_index(self, episode_index: int, camera_key: str) -> Mp4Index:
return self.files[self.lookup(episode_index, camera_key).file_id].mp4
def sample_slice(self, episode_index: int, camera_key: str) -> Mp4SampleSlice:
span = self.lookup(episode_index, camera_key)
return Mp4SampleSlice(
sample_lo=span.sample_lo,
sample_hi=span.sample_hi,
byte_offset=span.mdat_offset,
byte_length=span.mdat_length,
source_start_pts=span.source_start_pts,
)
def episode_byte_size(self, episode_index: int) -> int:
"""Exact synthesized video bytes retained while an episode is active."""
return sum(
synthesized_mp4_size(
self.mp4_index(episode_index, camera_key),
self.sample_slice(episode_index, camera_key),
)
for camera_key in self.video_keys
)
+707
View File
@@ -0,0 +1,707 @@
# 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
"""MP4 indexing and in-memory episode synthesis primitives."""
from __future__ import annotations
import struct
from collections.abc import Callable, Iterable
from dataclasses import dataclass
import numpy as np
@dataclass(frozen=True)
class Box:
type: bytes
start: int
header_size: int
end: int
@property
def payload_start(self) -> int:
return self.start + self.header_size
@property
def size(self) -> int:
return self.end - self.start
@dataclass(frozen=True)
class Mp4SampleSlice:
sample_lo: int
sample_hi: int
byte_offset: int
byte_length: int
source_start_pts: float
@dataclass(frozen=True)
class Mp4Index:
file_path: str
file_size: int
ftyp: bytes
moov_offset: int
mdat_offset: int
mdat_payload_offset: int
mdat_payload_size: int
faststart: bool
codec: str
timescale: int
duration: int
track_id: int
width: int
height: int
stsd_body: bytes
sample_pts: np.ndarray
sample_durations: np.ndarray
sample_sizes: np.ndarray
sample_offsets: np.ndarray
sync_samples: np.ndarray
def sample_slice(
self,
from_ts: float,
to_ts: float,
*,
keyframe_pad_s: float = 0.1,
keyframe_pad_fraction: float = 0.05,
file_size: int | None = None,
) -> Mp4SampleSlice:
if to_ts < from_ts:
raise ValueError(f"Invalid timestamp span: {from_ts=} {to_ts=}")
if len(self.sample_pts) == 0:
raise ValueError(f"{self.file_path} contains no indexed samples")
pad = max(keyframe_pad_s, (to_ts - from_ts) * keyframe_pad_fraction)
lo_ts = max(0.0, from_ts - pad)
hi_ts = to_ts + pad
lo = int(np.searchsorted(self.sample_pts, lo_ts, side="left"))
hi = int(np.searchsorted(self.sample_pts, hi_ts, side="right")) - 1
lo = min(max(lo, 0), len(self.sample_pts) - 1)
hi = min(max(hi, lo), len(self.sample_pts) - 1)
if len(self.sync_samples):
prev_sync = self.sync_samples[self.sync_samples <= lo]
if len(prev_sync):
lo = int(prev_sync[-1])
else:
lo = int(self.sync_samples[0])
if lo > hi:
hi = lo
offsets = self.sample_offsets[lo : hi + 1]
sizes = self.sample_sizes[lo : hi + 1]
slice_lo = int(offsets.min())
slice_hi = int((offsets + sizes).max())
if file_size is not None:
slice_hi = min(slice_hi, int(file_size))
return Mp4SampleSlice(
sample_lo=lo,
sample_hi=hi,
byte_offset=slice_lo,
byte_length=slice_hi - slice_lo,
source_start_pts=float(self.sample_pts[lo]),
)
def to_dict(self) -> dict:
return {
"file_path": self.file_path,
"file_size": self.file_size,
"ftyp": self.ftyp.hex(),
"moov_offset": self.moov_offset,
"mdat_offset": self.mdat_offset,
"mdat_payload_offset": self.mdat_payload_offset,
"mdat_payload_size": self.mdat_payload_size,
"faststart": self.faststart,
"codec": self.codec,
"timescale": self.timescale,
"duration": self.duration,
"track_id": self.track_id,
"width": self.width,
"height": self.height,
"stsd_body": self.stsd_body.hex(),
}
@classmethod
def from_dict(cls, data: dict, arrays: dict[str, np.ndarray]) -> Mp4Index:
return cls(
file_path=data["file_path"],
file_size=int(data["file_size"]),
ftyp=bytes.fromhex(data["ftyp"]),
moov_offset=int(data["moov_offset"]),
mdat_offset=int(data["mdat_offset"]),
mdat_payload_offset=int(data["mdat_payload_offset"]),
mdat_payload_size=int(data["mdat_payload_size"]),
faststart=bool(data["faststart"]),
codec=data["codec"],
timescale=int(data["timescale"]),
duration=int(data["duration"]),
track_id=int(data["track_id"]),
width=int(data["width"]),
height=int(data["height"]),
stsd_body=bytes.fromhex(data["stsd_body"]),
sample_pts=arrays["sample_pts"],
sample_durations=arrays["sample_durations"],
sample_sizes=arrays["sample_sizes"],
sample_offsets=arrays["sample_offsets"],
sync_samples=arrays["sync_samples"],
)
def fetch_mp4_index(
path: str,
read_range: Callable[[str, int, int], bytes],
*,
file_size: int,
header_probe_bytes: int = 4 * 1024 * 1024,
max_probe_bytes: int = 64 * 1024 * 1024,
) -> Mp4Index:
probe_size = min(header_probe_bytes, file_size)
while True:
data = read_range(path, 0, probe_size)
top = list(iter_boxes(data, 0, len(data), absolute_base=0, allow_truncated=True))
has_mdat = any(box.type == b"mdat" for box in top)
has_moov = any(box.type == b"moov" and box.end <= len(data) for box in top)
if has_mdat and has_moov:
return parse_mp4_index(path, data, file_size=file_size)
if probe_size >= min(max_probe_bytes, file_size):
if has_mdat and not has_moov:
tail_index = _fetch_tail_moov_index(path, read_range, data, top, file_size, max_probe_bytes)
if tail_index is not None:
return tail_index
missing = []
if not has_mdat:
missing.append("mdat")
if not has_moov:
missing.append("moov")
raise ValueError(
f"Could not find complete {'/'.join(missing)} in first {probe_size} bytes of {path}"
)
probe_size = min(probe_size * 2, max_probe_bytes, file_size)
def _fetch_tail_moov_index(
path: str,
read_range: Callable[[str, int, int], bytes],
prefix: bytes,
top_boxes: list[Box],
file_size: int,
max_probe_bytes: int,
) -> Mp4Index | None:
mdat_box = _one(top_boxes, b"mdat")
if mdat_box is None or mdat_box.end >= file_size:
return None
tail_offset = mdat_box.end
tail_length = min(max_probe_bytes, file_size - tail_offset)
tail = read_range(path, tail_offset, tail_length)
tail_boxes = list(iter_boxes(tail, 0, len(tail), absolute_base=tail_offset, allow_truncated=True))
moov_box = next(
(box for box in tail_boxes if box.type == b"moov" and box.end <= tail_offset + len(tail)), None
)
if moov_box is None:
return None
ftyp_box = _one(top_boxes, b"ftyp", required=False)
ftyp = (
prefix[ftyp_box.start : ftyp_box.end]
if ftyp_box is not None
else _box(b"ftyp", b"isom\0\0\2\0isomiso2mp41")
)
moov_start = moov_box.payload_start - tail_offset
moov_end = moov_box.end - tail_offset
return _parse_mp4_index_from_layout(
path,
file_size=file_size,
ftyp=ftyp,
moov_offset=moov_box.start,
moov=tail[moov_start:moov_end],
mdat_box=mdat_box,
)
def parse_mp4_index(path: str, data: bytes, *, file_size: int | None = None) -> Mp4Index:
if file_size is None:
file_size = len(data)
top = list(iter_boxes(data, 0, len(data), absolute_base=0, allow_truncated=True))
ftyp_box = _one(top, b"ftyp", required=False)
moov_box = _one(top, b"moov")
mdat_box = _one(top, b"mdat")
if moov_box.end > len(data):
raise ValueError(f"{path}: moov box is truncated")
moov = data[moov_box.payload_start : moov_box.end]
ftyp = (
data[ftyp_box.start : ftyp_box.end]
if ftyp_box is not None
else _box(b"ftyp", b"isom\0\0\2\0isomiso2mp41")
)
return _parse_mp4_index_from_layout(
path,
file_size=file_size,
ftyp=ftyp,
moov_offset=moov_box.start,
moov=moov,
mdat_box=mdat_box,
)
def _parse_mp4_index_from_layout(
path: str,
*,
file_size: int,
ftyp: bytes,
moov_offset: int,
moov: bytes,
mdat_box: Box,
) -> Mp4Index:
mvhd_timescale, mvhd_duration = _parse_mvhd(_find_descendant(moov, [b"mvhd"]))
trak_box, trak_payload = _find_video_trak(moov)
_ = trak_box
tkhd = _parse_tkhd(_find_descendant(trak_payload, [b"tkhd"]))
mdhd_timescale, mdhd_duration = _parse_mdhd(_find_descendant(trak_payload, [b"mdia", b"mdhd"]))
stbl = _find_descendant(trak_payload, [b"mdia", b"minf", b"stbl"])
stsd = _find_child(stbl, b"stsd")
stsd_body = stbl[stsd.payload_start : stsd.end]
codec = _parse_stsd_codec(stsd_body)
stts = _parse_stts(_payload(stbl, b"stts"))
sample_sizes = _parse_stsz(_payload(stbl, b"stsz"))
stsc = _parse_stsc(_payload(stbl, b"stsc"))
chunk_offsets = _parse_chunk_offsets(stbl)
sync_samples = _parse_stss(stbl, len(sample_sizes))
sample_durations = _expand_stts(stts, len(sample_sizes))
sample_pts_units = np.empty(len(sample_durations), dtype=np.int64)
if len(sample_durations):
sample_pts_units[0] = 0
if len(sample_durations) > 1:
sample_pts_units[1:] = np.cumsum(sample_durations[:-1], dtype=np.int64)
sample_pts = sample_pts_units.astype(np.float64) / float(mdhd_timescale)
sample_offsets = _sample_offsets(stsc, chunk_offsets, sample_sizes)
return Mp4Index(
file_path=path,
file_size=file_size,
ftyp=ftyp,
moov_offset=moov_offset,
mdat_offset=mdat_box.start,
mdat_payload_offset=mdat_box.payload_start,
mdat_payload_size=mdat_box.end - mdat_box.payload_start
if mdat_box.end <= file_size
else file_size - mdat_box.payload_start,
faststart=moov_offset < mdat_box.start,
codec=codec,
timescale=mdhd_timescale,
duration=mdhd_duration or mvhd_duration,
track_id=tkhd["track_id"],
width=tkhd["width"],
height=tkhd["height"],
stsd_body=stsd_body,
sample_pts=sample_pts,
sample_durations=sample_durations,
sample_sizes=sample_sizes,
sample_offsets=sample_offsets,
sync_samples=sync_samples,
)
def synthesize_mp4(index: Mp4Index, sample_slice: Mp4SampleSlice, mdat_payload: bytes) -> bytes:
lo = sample_slice.sample_lo
hi = sample_slice.sample_hi + 1
if lo < 0 or hi > len(index.sample_sizes) or lo >= hi:
raise ValueError(f"Invalid sample range [{lo}, {hi}) for {index.file_path}")
offsets = index.sample_offsets[lo:hi]
sizes = index.sample_sizes[lo:hi]
rel_offsets = offsets - sample_slice.byte_offset
if int(rel_offsets.min()) != 0:
raise ValueError("Sample slice must start at the minimum referenced sample offset")
if int((rel_offsets + sizes).max()) > len(mdat_payload):
raise ValueError("Sample slice does not cover all referenced samples")
durations = index.sample_durations[lo:hi]
sync = index.sync_samples[(index.sync_samples >= lo) & (index.sync_samples < hi)] - lo + 1
moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=0)
header_size = len(index.ftyp) + len(moov)
mdat_header_size = 8 if len(mdat_payload) + 8 <= 0xFFFFFFFF else 16
moov = _make_moov(
index,
durations,
sizes,
rel_offsets,
sync,
mdat_data_offset=header_size + mdat_header_size,
)
return index.ftyp + moov + _box(b"mdat", mdat_payload)
def synthesized_mp4_size(index: Mp4Index, sample_slice: Mp4SampleSlice) -> int:
"""Return the exact synthesized mini-MP4 size without fetching its media payload."""
lo = sample_slice.sample_lo
hi = sample_slice.sample_hi + 1
if lo < 0 or hi > len(index.sample_sizes) or lo >= hi:
raise ValueError(f"Invalid sample range [{lo}, {hi}) for {index.file_path}")
offsets = index.sample_offsets[lo:hi]
sizes = index.sample_sizes[lo:hi]
rel_offsets = offsets - sample_slice.byte_offset
if int(rel_offsets.min()) != 0:
raise ValueError("Sample slice must start at the minimum referenced sample offset")
if int((rel_offsets + sizes).max()) > sample_slice.byte_length:
raise ValueError("Sample slice does not cover all referenced samples")
durations = index.sample_durations[lo:hi]
sync = index.sync_samples[(index.sync_samples >= lo) & (index.sync_samples < hi)] - lo + 1
moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=0)
header_size = len(index.ftyp) + len(moov)
mdat_header_size = 8 if sample_slice.byte_length + 8 <= 0xFFFFFFFF else 16
moov = _make_moov(
index,
durations,
sizes,
rel_offsets,
sync,
mdat_data_offset=header_size + mdat_header_size,
)
return len(index.ftyp) + len(moov) + mdat_header_size + sample_slice.byte_length
def iter_boxes(
data: bytes,
start: int,
end: int,
*,
absolute_base: int = 0,
allow_truncated: bool = False,
) -> Iterable[Box]:
pos = start
while pos + 8 <= end:
size = struct.unpack_from(">I", data, pos)[0]
typ = data[pos + 4 : pos + 8]
header_size = 8
if size == 1:
if pos + 16 > end:
break
size = struct.unpack_from(">Q", data, pos + 8)[0]
header_size = 16
elif size == 0:
size = end - pos
if size < header_size:
break
box_end = pos + size
if box_end > end and not allow_truncated:
break
yield Box(typ, absolute_base + pos, header_size, absolute_base + box_end)
pos = box_end
def _find_video_trak(moov: bytes) -> tuple[Box, bytes]:
for trak in _children(moov, 0, len(moov)):
if trak.type != b"trak":
continue
payload = moov[trak.payload_start : trak.end]
hdlr = _find_descendant(payload, [b"mdia", b"hdlr"])
if hdlr[8:12] == b"vide":
return trak, payload
raise ValueError("No video track found")
def _find_descendant(data: bytes, path: list[bytes]) -> bytes:
current = data
for typ in path:
box = _find_child(current, typ)
current = current[box.payload_start : box.end]
return current
def _find_child(data: bytes, typ: bytes) -> Box:
for box in _children(data, 0, len(data)):
if box.type == typ:
return box
raise ValueError(f"Missing MP4 box {typ.decode('latin1')}")
def _children(data: bytes, start: int, end: int) -> Iterable[Box]:
return iter_boxes(data, start, end, absolute_base=0)
def _one(boxes: list[Box], typ: bytes, *, required: bool = True) -> Box | None:
matches = [box for box in boxes if box.type == typ]
if not matches and required:
raise ValueError(f"Missing MP4 box {typ.decode('latin1')}")
return matches[0] if matches else None
def _payload(parent: bytes, typ: bytes) -> bytes:
box = _find_child(parent, typ)
return parent[box.payload_start : box.end]
def _parse_mvhd(payload: bytes) -> tuple[int, int]:
version = payload[0]
if version == 1:
return struct.unpack_from(">IQ", payload, 20)
return struct.unpack_from(">II", payload, 12)
def _parse_mdhd(payload: bytes) -> tuple[int, int]:
version = payload[0]
if version == 1:
return struct.unpack_from(">IQ", payload, 20)
return struct.unpack_from(">II", payload, 12)
def _parse_tkhd(payload: bytes) -> dict[str, int]:
version = payload[0]
if version == 1:
track_id = struct.unpack_from(">I", payload, 20)[0]
duration = struct.unpack_from(">Q", payload, 28)[0]
width, height = struct.unpack_from(">II", payload, 88)
else:
track_id = struct.unpack_from(">I", payload, 12)[0]
duration = struct.unpack_from(">I", payload, 20)[0]
width, height = struct.unpack_from(">II", payload, 76)
return {"track_id": track_id, "duration": duration, "width": width >> 16, "height": height >> 16}
def _parse_stsd_codec(stsd_body: bytes) -> str:
if len(stsd_body) < 16:
return "unknown"
return stsd_body[12:16].decode("latin1")
def _parse_stts(payload: bytes) -> list[tuple[int, int]]:
count = struct.unpack_from(">I", payload, 4)[0]
out = []
offset = 8
for _ in range(count):
out.append(struct.unpack_from(">II", payload, offset))
offset += 8
return out
def _expand_stts(entries: list[tuple[int, int]], sample_count: int) -> np.ndarray:
values = np.empty(sample_count, dtype=np.int64)
pos = 0
for count, delta in entries:
values[pos : pos + count] = delta
pos += count
if pos != sample_count:
raise ValueError(f"stts describes {pos} samples, stsz describes {sample_count}")
return values
def _parse_stsz(payload: bytes) -> np.ndarray:
sample_size, sample_count = struct.unpack_from(">II", payload, 4)
if sample_size:
return np.full(sample_count, sample_size, dtype=np.int64)
offset = 12
values = np.empty(sample_count, dtype=np.int64)
for idx in range(sample_count):
values[idx] = struct.unpack_from(">I", payload, offset)[0]
offset += 4
return values
def _parse_stsc(payload: bytes) -> list[tuple[int, int, int]]:
count = struct.unpack_from(">I", payload, 4)[0]
out = []
offset = 8
for _ in range(count):
out.append(struct.unpack_from(">III", payload, offset))
offset += 12
return out
def _parse_chunk_offsets(stbl: bytes) -> np.ndarray:
with_stco = None
with_co64 = None
for box in _children(stbl, 0, len(stbl)):
if box.type == b"stco":
with_stco = stbl[box.payload_start : box.end]
elif box.type == b"co64":
with_co64 = stbl[box.payload_start : box.end]
if with_co64 is not None:
count = struct.unpack_from(">I", with_co64, 4)[0]
return np.array(
[struct.unpack_from(">Q", with_co64, 8 + idx * 8)[0] for idx in range(count)], dtype=np.int64
)
if with_stco is None:
raise ValueError("Missing stco/co64 chunk offsets")
count = struct.unpack_from(">I", with_stco, 4)[0]
return np.array(
[struct.unpack_from(">I", with_stco, 8 + idx * 4)[0] for idx in range(count)], dtype=np.int64
)
def _parse_stss(stbl: bytes, sample_count: int) -> np.ndarray:
for box in _children(stbl, 0, len(stbl)):
if box.type == b"stss":
payload = stbl[box.payload_start : box.end]
count = struct.unpack_from(">I", payload, 4)[0]
return np.array(
[struct.unpack_from(">I", payload, 8 + idx * 4)[0] - 1 for idx in range(count)],
dtype=np.int64,
)
return np.arange(sample_count, dtype=np.int64)
def _sample_offsets(
stsc: list[tuple[int, int, int]], chunk_offsets: np.ndarray, sample_sizes: np.ndarray
) -> np.ndarray:
if not stsc:
raise ValueError("stsc is empty")
offsets = np.empty(len(sample_sizes), dtype=np.int64)
sample_idx = 0
for entry_idx, (first_chunk, samples_per_chunk, _desc_idx) in enumerate(stsc):
next_first = stsc[entry_idx + 1][0] if entry_idx + 1 < len(stsc) else len(chunk_offsets) + 1
for chunk_number in range(first_chunk, next_first):
if chunk_number < 1 or chunk_number > len(chunk_offsets):
raise ValueError("stsc references a chunk outside stco/co64")
chunk_pos = int(chunk_offsets[chunk_number - 1])
for _ in range(samples_per_chunk):
if sample_idx >= len(sample_sizes):
return offsets
offsets[sample_idx] = chunk_pos
chunk_pos += int(sample_sizes[sample_idx])
sample_idx += 1
if sample_idx != len(sample_sizes):
raise ValueError(f"stsc describes {sample_idx} samples, stsz describes {len(sample_sizes)}")
return offsets
def _make_moov(
index: Mp4Index,
durations: np.ndarray,
sizes: np.ndarray,
rel_offsets: np.ndarray,
sync_samples: np.ndarray,
*,
mdat_data_offset: int,
) -> bytes:
duration = int(durations.sum())
stco_values = [int(mdat_data_offset + value) for value in rel_offsets]
if any(value > 0xFFFFFFFF for value in stco_values):
offset_box = _co64(stco_values)
else:
offset_box = _stco(stco_values)
stbl = _box(
b"stbl",
_box(b"stsd", index.stsd_body)
+ _stts(durations)
+ _stsc_one_sample_per_chunk(len(sizes))
+ _stsz(sizes)
+ offset_box
+ (_stss(sync_samples) if len(sync_samples) else b""),
)
minf = _box(b"minf", _vmhd() + _dinf() + stbl)
mdia = _box(b"mdia", _mdhd(index.timescale, duration) + _hdlr() + minf)
trak = _box(b"trak", _tkhd(index.track_id, duration, index.width, index.height) + mdia)
return _box(b"moov", _mvhd(index.timescale, duration, index.track_id + 1) + trak)
def _full_box(typ: bytes, version: int, flags: int, payload: bytes = b"") -> bytes:
return _box(typ, bytes([version]) + flags.to_bytes(3, "big") + payload)
def _box(typ: bytes, payload: bytes) -> bytes:
size = len(payload) + 8
if size <= 0xFFFFFFFF:
return struct.pack(">I4s", size, typ) + payload
return struct.pack(">I4sQ", 1, typ, size + 8) + payload
def _mvhd(timescale: int, duration: int, next_track_id: int) -> bytes:
matrix = struct.pack(">9I", 0x00010000, 0, 0, 0, 0x00010000, 0, 0, 0, 0x40000000)
payload = (
struct.pack(">IIII", 0, 0, timescale, duration)
+ struct.pack(">IHH", 0x00010000, 0x0100, 0)
+ b"\0" * 8
+ matrix
+ b"\0" * 24
+ struct.pack(">I", next_track_id)
)
return _full_box(b"mvhd", 0, 0, payload)
def _tkhd(track_id: int, duration: int, width: int, height: int) -> bytes:
matrix = struct.pack(">9I", 0x00010000, 0, 0, 0, 0x00010000, 0, 0, 0, 0x40000000)
payload = (
struct.pack(">IIIII", 0, 0, track_id, 0, duration)
+ b"\0" * 8
+ struct.pack(">hhhh", 0, 0, 0, 0)
+ matrix
+ struct.pack(">II", width << 16, height << 16)
)
return _full_box(b"tkhd", 0, 7, payload)
def _mdhd(timescale: int, duration: int) -> bytes:
return _full_box(b"mdhd", 0, 0, struct.pack(">IIIIH", 0, 0, timescale, duration, 0x55C4) + b"\0\0")
def _hdlr() -> bytes:
return _full_box(b"hdlr", 0, 0, b"\0" * 4 + b"vide" + b"\0" * 12 + b"VideoHandler\0")
def _vmhd() -> bytes:
return _full_box(b"vmhd", 0, 1, struct.pack(">HHHH", 0, 0, 0, 0))
def _dinf() -> bytes:
url = _full_box(b"url ", 0, 1)
dref = _full_box(b"dref", 0, 0, struct.pack(">I", 1) + url)
return _box(b"dinf", dref)
def _stts(durations: np.ndarray) -> bytes:
runs = []
for duration in durations.tolist():
if runs and runs[-1][1] == int(duration):
runs[-1][0] += 1
else:
runs.append([1, int(duration)])
payload = struct.pack(">I", len(runs)) + b"".join(
struct.pack(">II", count, delta) for count, delta in runs
)
return _full_box(b"stts", 0, 0, payload)
def _stsc_one_sample_per_chunk(sample_count: int) -> bytes:
return _full_box(b"stsc", 0, 0, struct.pack(">IIII", 1, 1, 1, 1))
def _stsz(sizes: np.ndarray) -> bytes:
return _full_box(
b"stsz",
0,
0,
struct.pack(">II", 0, len(sizes)) + b"".join(struct.pack(">I", int(size)) for size in sizes.tolist()),
)
def _stco(values: list[int]) -> bytes:
return _full_box(
b"stco", 0, 0, struct.pack(">I", len(values)) + b"".join(struct.pack(">I", v) for v in values)
)
def _co64(values: list[int]) -> bytes:
return _full_box(
b"co64", 0, 0, struct.pack(">I", len(values)) + b"".join(struct.pack(">Q", v) for v in values)
)
def _stss(values: np.ndarray) -> bytes:
return _full_box(
b"stss",
0,
0,
struct.pack(">I", len(values)) + b"".join(struct.pack(">I", int(value)) for value in values.tolist()),
)
+743
View File
@@ -0,0 +1,743 @@
# 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 contextlib
import json
import os
import posixpath
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import UTC, datetime
from pathlib import Path
from types import MethodType
from typing import Any
from urllib.parse import quote, urljoin, urlparse
from uuid import uuid4
import fsspec
import httpx
from huggingface_hub import HfApi, HfFileSystem, constants
from huggingface_hub.utils import get_session, hf_raise_for_status
_HTTP_FAILURE_LOG_LOCK = threading.Lock()
def _get_header(headers: Any, name: str) -> str | None:
if hasattr(headers, "get"):
return headers.get(name)
lower_name = name.lower()
for key, value in headers.items():
if key.lower() == lower_name:
return value
return None
def _ensure_request_id(headers: dict[str, str]) -> str:
request_id = _get_header(headers, "X-Amzn-Trace-Id") or _get_header(headers, "X-Request-Id")
if request_id is None:
request_id = str(uuid4())
headers["X-Amzn-Trace-Id"] = request_id
return request_id
def _log_http_failure(
*,
backend: str,
method: str,
url: str,
headers: dict[str, str],
elapsed_s: float,
status_code: int | None = None,
exception: Exception | None = None,
attempt: int | None = None,
response_headers: Any | None = None,
) -> None:
log_path = os.environ.get("LEROBOT_HTTP_FAILURE_LOG")
if not log_path:
return
parsed = urlparse(url)
record = {
"ts": datetime.now(UTC).isoformat(),
"backend": backend,
"method": method,
"host": parsed.netloc,
"path": parsed.path,
"range": _get_header(headers, "Range") or _get_header(headers, "range"),
"request_id": _get_header(headers, "X-Amzn-Trace-Id") or _get_header(headers, "X-Request-Id"),
"elapsed_s": round(elapsed_s, 6),
}
if attempt is not None:
record["attempt"] = attempt
if status_code is not None:
record["status_code"] = status_code
if exception is not None:
record["exception_type"] = type(exception).__name__
record["exception"] = str(exception)
if response_headers is not None:
record["response_request_id"] = (
_get_header(response_headers, "x-request-id")
or _get_header(response_headers, "x-amz-cf-id")
or _get_header(response_headers, "x-amz-request-id")
)
record["cache_status"] = (
_get_header(response_headers, "x-cache")
or _get_header(response_headers, "cf-cache-status")
or _get_header(response_headers, "x-hf-cache")
)
record["content_range"] = _get_header(response_headers, "content-range")
record["content_length"] = _get_header(response_headers, "content-length")
path = Path(log_path).expanduser()
with _HTTP_FAILURE_LOG_LOCK:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("a") as out:
out.write(json.dumps(record, sort_keys=True) + "\n")
class ThreadLocalRangeFetcher:
"""Range reader that gives each worker thread independent file handles."""
def __init__(
self,
data_root: str | Path,
*,
block_size: int = 2**20,
cache_type: str = "none",
token: str | bool | None = None,
):
self.data_root = str(data_root).rstrip("/")
storage_options = {"token": token} if token is not None and self.data_root.startswith("hf://") else {}
self.fs, self._root_path = fsspec.core.url_to_fs(self.data_root, **storage_options)
self._is_local = self.fs.protocol in ("file", "local") or (
isinstance(self.fs.protocol, tuple) and "file" in self.fs.protocol
)
self.block_size = block_size
self.cache_type = cache_type
self._local = threading.local()
self._handles_lock = threading.Lock()
self._all_handles: dict[int, Any] = {}
self._timing_lock = threading.Lock()
self._timing_totals = {
"range_jobs": 0.0,
"range_bytes": 0.0,
"range_open_s": 0.0,
"range_seek_s": 0.0,
"range_read_s": 0.0,
}
def _url(self, relative_path: str) -> str:
if self._is_local:
return str(Path(self._root_path) / relative_path)
return posixpath.join(self._root_path.rstrip("/"), relative_path.lstrip("/"))
def _handle(self, relative_path: str):
handles = getattr(self._local, "handles", None)
if handles is None:
handles = {}
self._local.handles = handles
handle = handles.get(relative_path)
if handle is None or getattr(handle, "closed", False):
handle = self.fs.open(
self._url(relative_path), "rb", block_size=self.block_size, cache_type=self.cache_type
)
self._instrument_hf_handle(handle)
handles[relative_path] = handle
with self._handles_lock:
self._all_handles[id(handle)] = handle
return handle
def info_size(self, relative_path: str) -> int:
return int(self.fs.info(self._url(relative_path))["size"])
def read_range(self, relative_path: str, offset: int, length: int) -> bytes:
open_start = time.perf_counter()
handle = self._handle(relative_path)
open_s = time.perf_counter() - open_start
seek_start = time.perf_counter()
handle.seek(offset)
seek_s = time.perf_counter() - seek_start
read_start = time.perf_counter()
data = handle.read(length)
read_s = time.perf_counter() - read_start
self._record_timing(
range_jobs=1.0,
range_bytes=float(len(data)),
range_open_s=open_s,
range_seek_s=seek_s,
range_read_s=read_s,
)
return data
def _record_timing(self, **kwargs: float) -> None:
with self._timing_lock:
for key, value in kwargs.items():
self._timing_totals[key] = self._timing_totals.get(key, 0.0) + value
def _instrument_hf_handle(self, handle: Any) -> None:
if getattr(handle, "_lerobot_range_timing", False):
return
if not hasattr(handle, "_request_with_retry"):
return
def request_with_retry(
handle_self,
method: str,
url: str,
*,
headers: dict[str, str],
follow_redirects: bool | None = None,
max_retries: int = 5,
) -> httpx.Response:
from huggingface_hub.hf_file_system import _RANGE_RETRY_EXCEPTIONS, _RANGE_RETRY_STATUS_CODES
method_key = method.lower()
sleep_time = 1.0
retry_attempts = 0.0
retry_sleep_s = 0.0
failed_attempt_s = 0.0
exception_attempts = 0.0
extra_counts: dict[str, float] = {}
call_start = time.perf_counter()
request_kwargs: dict[str, Any] = {
"headers": headers,
"timeout": constants.HF_HUB_DOWNLOAD_TIMEOUT,
}
_ensure_request_id(headers)
if follow_redirects is not None:
request_kwargs["follow_redirects"] = follow_redirects
for attempt in range(max_retries + 1):
attempt_start = time.perf_counter()
try:
response = get_session().request(method, url, **request_kwargs)
except _RANGE_RETRY_EXCEPTIONS as exc:
attempt_s = time.perf_counter() - attempt_start
failed_attempt_s += attempt_s
exception_attempts += 1.0
key = f"range_hffs_{method_key}_exception_{type(exc).__name__}"
extra_counts[key] = extra_counts.get(key, 0.0) + 1.0
_log_http_failure(
backend="hffs",
method=method,
url=url,
headers=headers,
elapsed_s=attempt_s,
exception=exc,
attempt=attempt,
)
if attempt == max_retries:
self._record_hffs_request_timing(
method_key,
time.perf_counter() - call_start,
retry_attempts,
retry_sleep_s,
failed_attempt_s,
exception_attempts,
None,
0,
extra_counts,
)
raise
else:
elapsed = time.perf_counter() - attempt_start
if response.status_code not in _RANGE_RETRY_STATUS_CODES or attempt == max_retries:
self._record_hffs_request_timing(
method_key,
time.perf_counter() - call_start,
retry_attempts,
retry_sleep_s,
failed_attempt_s,
exception_attempts,
response.status_code,
len(response.content),
extra_counts,
)
return response
failed_attempt_s += elapsed
key = f"range_hffs_{method_key}_failed_status_{response.status_code}"
extra_counts[key] = extra_counts.get(key, 0.0) + 1.0
response.close()
_log_http_failure(
backend="hffs",
method=method,
url=url,
headers=headers,
elapsed_s=elapsed,
status_code=response.status_code,
attempt=attempt,
response_headers=response.headers,
)
time.sleep(sleep_time)
retry_attempts += 1.0
retry_sleep_s += sleep_time
sleep_time = min(8.0, sleep_time * 2)
raise RuntimeError("unreachable")
handle._request_with_retry = MethodType(request_with_retry, handle)
handle._lerobot_range_timing = True
def _record_hffs_request_timing(
self,
method_key: str,
total_s: float,
retry_attempts: float,
retry_sleep_s: float,
failed_attempt_s: float,
exception_attempts: float,
status_code: int | None,
byte_count: int,
extra_counts: dict[str, float],
) -> None:
timings = {
f"range_hffs_{method_key}_requests": 1.0,
f"range_hffs_{method_key}_s": total_s,
f"range_hffs_{method_key}_retries": retry_attempts,
f"range_hffs_{method_key}_retry_sleep_s": retry_sleep_s,
f"range_hffs_{method_key}_failed_attempt_s": failed_attempt_s,
f"range_hffs_{method_key}_exception_attempts": exception_attempts,
f"range_hffs_{method_key}_bytes": float(byte_count),
}
if status_code is not None:
timings[f"range_hffs_{method_key}_status_{status_code}"] = 1.0
timings.update(extra_counts)
self._record_timing(**timings)
def timing_summary(self) -> dict[str, float]:
with self._timing_lock:
return dict(self._timing_totals)
def close(self) -> None:
with self._handles_lock:
handles = list(self._all_handles.values())
self._all_handles.clear()
for handle in handles:
with contextlib.suppress(Exception):
handle.close()
local_handles = getattr(self._local, "handles", None)
if local_handles is not None:
local_handles.clear()
class NativeHTTPRangeFetcher:
"""Direct pooled HTTP range reader for hf:// paths."""
_GLOBAL_SOURCE_URLS: dict[tuple[str, str], str] = {}
_GLOBAL_RESOLVED_URLS: dict[tuple[str, str], str] = {}
_GLOBAL_SIZES: dict[tuple[str, str], int] = {}
_GLOBAL_LOCK = threading.Lock()
_RETRYABLE_EXCEPTIONS = (
httpx.ConnectError,
httpx.ConnectTimeout,
httpx.ReadError,
httpx.ReadTimeout,
httpx.RemoteProtocolError,
httpx.PoolTimeout,
)
_RETRYABLE_STATUS_CODES = {408, 425, 429, 500, 502, 503, 504}
def __init__(
self,
data_root: str | Path,
*,
max_connections: int = 32,
timeout: float = 60.0,
max_retries: int = 4,
subrange_parts: int = 1,
subrange_min_bytes: int = 8 * 1024 * 1024,
token: str | bool | None = None,
):
self.data_root = str(data_root).rstrip("/")
if not self.data_root.startswith("hf://"):
raise ValueError("NativeHTTPRangeFetcher only supports hf:// roots")
self.max_retries = max_retries
# Sub-range parallelism: split one large GET into `subrange_parts` concurrent GETs.
# Under a per-host throughput ceiling this adds no aggregate bandwidth, but divides
# per-request latency by ~parts - keep (in-flight jobs x parts) near the ceiling's
# connection sweet spot (~64 on the observed HF bucket path) rather than raising both.
self.subrange_parts = max(1, subrange_parts)
self.subrange_min_bytes = max(1, subrange_min_bytes)
self._subrange_pool = (
ThreadPoolExecutor(max_workers=max_connections, thread_name_prefix="subrange")
if self.subrange_parts > 1
else None
)
self.api = HfApi(token=token)
self.fs: HfFileSystem | None = None
self._bucket_id: str | None = None
self._bucket_prefix = ""
if self.data_root.startswith("hf://buckets/"):
bucket_root = self.data_root.removeprefix("hf://buckets/")
parts = bucket_root.split("/", 2)
if len(parts) < 2:
raise ValueError(f"Invalid bucket root: {self.data_root}")
self._bucket_id = f"{parts[0]}/{parts[1]}"
self._bucket_prefix = parts[2].strip("/") if len(parts) == 3 else ""
else:
self.fs = HfFileSystem(token=token)
self.client = httpx.Client(
timeout=timeout,
limits=httpx.Limits(max_connections=max_connections, max_keepalive_connections=max_connections),
follow_redirects=False,
)
self._resolved_urls: dict[str, str] = {}
self._source_urls: dict[str, str] = {}
self._sizes: dict[str, int] = {}
self._lock = threading.Lock()
self._timing_lock = threading.Lock()
self._timing_totals = {
"range_jobs": 0.0,
"range_bytes": 0.0,
"range_resolve_s": 0.0,
"range_header_s": 0.0,
"range_first_byte_s": 0.0,
"range_body_s": 0.0,
"range_retry_attempts": 0.0,
"range_retry_sleep_s": 0.0,
"range_failed_requests": 0.0,
}
def _request(self, method: str, url: str, **kwargs) -> httpx.Response:
last_exc: Exception | None = None
for attempt in range(self.max_retries + 1):
try:
return self.client.request(method, url, **kwargs)
except self._RETRYABLE_EXCEPTIONS as exc:
last_exc = exc
if attempt >= self.max_retries:
break
time.sleep(min(0.5 * 2**attempt, 5.0))
if last_exc is None:
raise RuntimeError("HTTP request failed without an exception")
raise last_exc
def _cache_key(self, relative_path: str) -> tuple[str, str]:
return self.data_root, relative_path
def _path(self, relative_path: str) -> str:
return f"{self.data_root}/{relative_path}"
def _bucket_path(self, relative_path: str) -> str:
if self._bucket_prefix:
return f"{self._bucket_prefix}/{relative_path}"
return relative_path
def _headers_for(self, request_url: str, source_url: str) -> dict[str, str]:
headers = self.api._build_hf_headers()
if urlparse(request_url).netloc != urlparse(source_url).netloc:
headers.pop("authorization", None)
headers.pop("Authorization", None)
return headers
def _source_url(self, relative_path: str) -> str:
with self._lock:
source = self._source_urls.get(relative_path)
if source is not None:
return source
key = self._cache_key(relative_path)
with self._GLOBAL_LOCK:
source = self._GLOBAL_SOURCE_URLS.get(key)
if source is None:
if self._bucket_id is not None:
source = (
f"{constants.ENDPOINT}/buckets/{self._bucket_id}/resolve/"
f"{quote(self._bucket_path(relative_path))}"
)
else:
if self.fs is None:
raise RuntimeError("HfFileSystem fallback was not initialized")
source = self.fs.url(self._path(relative_path))
with self._GLOBAL_LOCK:
self._GLOBAL_SOURCE_URLS[key] = source
with self._lock:
self._source_urls[relative_path] = source
return source
def _resolve_url(self, relative_path: str, *, refresh: bool = False) -> str:
with self._lock:
if not refresh and relative_path in self._resolved_urls:
return self._resolved_urls[relative_path]
key = self._cache_key(relative_path)
if not refresh:
with self._GLOBAL_LOCK:
resolved = self._GLOBAL_RESOLVED_URLS.get(key)
size = self._GLOBAL_SIZES.get(key)
if resolved is not None:
with self._lock:
self._resolved_urls[relative_path] = resolved
if size is not None:
self._sizes[relative_path] = size
return resolved
source = self._source_url(relative_path)
response = self._request("HEAD", source, headers=self.api._build_hf_headers(), follow_redirects=False)
try:
hf_raise_for_status(response)
location = response.headers.get("Location")
resolved = urljoin(source, location) if location else source
with self._lock:
self._resolved_urls[relative_path] = resolved
if "Content-Length" in response.headers:
self._sizes[relative_path] = int(response.headers["Content-Length"])
with self._GLOBAL_LOCK:
self._GLOBAL_RESOLVED_URLS[key] = resolved
if "Content-Length" in response.headers:
self._GLOBAL_SIZES[key] = int(response.headers["Content-Length"])
return resolved
finally:
response.close()
def info_size(self, relative_path: str) -> int:
with self._lock:
size = self._sizes.get(relative_path)
if size is not None:
return size
key = self._cache_key(relative_path)
with self._GLOBAL_LOCK:
size = self._GLOBAL_SIZES.get(key)
if size is not None:
with self._lock:
self._sizes[relative_path] = size
return size
resolved = self._resolve_url(relative_path)
source = self._source_url(relative_path)
response = self._request(
"HEAD", resolved, headers=self._headers_for(resolved, source), follow_redirects=True
)
try:
hf_raise_for_status(response)
size = int(response.headers["Content-Length"])
with self._lock:
self._sizes[relative_path] = size
with self._GLOBAL_LOCK:
self._GLOBAL_SIZES[key] = size
return size
finally:
response.close()
def read_range(self, relative_path: str, offset: int, length: int) -> bytes:
parts = self.subrange_parts
if self._subrange_pool is None or parts <= 1 or length < 2 * self.subrange_min_bytes:
return self._read_range_single(relative_path, offset, length)
parts = min(parts, max(1, length // self.subrange_min_bytes))
if parts <= 1:
return self._read_range_single(relative_path, offset, length)
step = (length + parts - 1) // parts
spans = [(offset + i * step, min(step, length - i * step)) for i in range(parts)]
futures = [
self._subrange_pool.submit(self._read_range_single, relative_path, span_off, span_len)
for span_off, span_len in spans
]
return b"".join(future.result() for future in futures)
def _read_range_single(self, relative_path: str, offset: int, length: int) -> bytes:
resolve_start = time.perf_counter()
resolved = self._resolve_url(relative_path)
source = self._source_url(relative_path)
resolve_s = time.perf_counter() - resolve_start
headers = self._headers_for(resolved, source)
headers["Range"] = f"bytes={offset}-{offset + length - 1}"
payload, status_code, timings = self._read_range_response(resolved, headers)
if status_code == 403:
refresh_start = time.perf_counter()
resolved = self._resolve_url(relative_path, refresh=True)
resolve_s += time.perf_counter() - refresh_start
headers = self._headers_for(resolved, source)
headers["Range"] = f"bytes={offset}-{offset + length - 1}"
payload, status_code, retry_timings = self._read_range_response(resolved, headers)
for key, value in retry_timings.items():
timings[key] += value
if status_code == 403:
raise PermissionError(f"HTTP range request returned 403 after URL refresh: {relative_path}")
if status_code != 206:
raise RuntimeError(f"HTTP range request returned {status_code} after retries: {relative_path}")
self._record_timing(
range_jobs=1.0,
range_bytes=float(len(payload)),
range_resolve_s=resolve_s,
**{f"range_status_{status_code}": 1.0},
**timings,
)
return payload
def _read_range_response(self, url: str, headers: dict[str, str]) -> tuple[bytes, int, dict[str, float]]:
last_exc: Exception | None = None
retry_attempts = 0.0
retry_sleep_s = 0.0
failed_attempt_s = 0.0
exception_attempts = 0.0
exception_counts: dict[str, float] = {}
_ensure_request_id(headers)
for attempt in range(self.max_retries + 1):
attempt_start = time.perf_counter()
try:
payload, status_code, timings = self._read_range_response_once(url, headers)
if status_code in self._RETRYABLE_STATUS_CODES:
attempt_s = time.perf_counter() - attempt_start
failed_attempt_s += attempt_s
exception_attempts += 1.0
status_key = f"range_failed_status_{status_code}"
exception_counts[status_key] = exception_counts.get(status_key, 0.0) + 1.0
_log_http_failure(
backend="native-http",
method="GET",
url=url,
headers=headers,
elapsed_s=attempt_s,
status_code=status_code,
attempt=attempt,
)
if attempt >= self.max_retries:
timings["range_retry_attempts"] = retry_attempts
timings["range_retry_sleep_s"] = retry_sleep_s
timings["range_failed_attempt_s"] = failed_attempt_s
timings["range_exception_attempts"] = exception_attempts
timings.update(exception_counts)
return payload, status_code, timings
retry_attempts += 1.0
sleep_s = min(0.5 * 2**attempt, 5.0)
retry_sleep_s += sleep_s
time.sleep(sleep_s)
continue
timings["range_retry_attempts"] = retry_attempts
timings["range_retry_sleep_s"] = retry_sleep_s
timings["range_failed_attempt_s"] = failed_attempt_s
timings["range_exception_attempts"] = exception_attempts
timings.update(exception_counts)
return payload, status_code, timings
except self._RETRYABLE_EXCEPTIONS as exc:
last_exc = exc
attempt_s = time.perf_counter() - attempt_start
failed_attempt_s += attempt_s
exception_attempts += 1.0
exception_key = f"range_exception_{type(exc).__name__}"
exception_counts[exception_key] = exception_counts.get(exception_key, 0.0) + 1.0
_log_http_failure(
backend="native-http",
method="GET",
url=url,
headers=headers,
elapsed_s=attempt_s,
exception=exc,
attempt=attempt,
)
if attempt >= self.max_retries:
break
retry_attempts += 1.0
sleep_s = min(0.5 * 2**attempt, 5.0)
retry_sleep_s += sleep_s
time.sleep(sleep_s)
self._record_timing(
range_failed_requests=1.0,
range_retry_attempts=retry_attempts,
range_retry_sleep_s=retry_sleep_s,
range_failed_attempt_s=failed_attempt_s,
range_exception_attempts=exception_attempts,
**exception_counts,
)
if last_exc is None:
raise RuntimeError("HTTP range request failed without an exception")
raise last_exc
def _read_range_response_once(
self, url: str, headers: dict[str, str]
) -> tuple[bytes, int, dict[str, float]]:
header_start = time.perf_counter()
with self.client.stream("GET", url, headers=headers) as response:
header_s = time.perf_counter() - header_start
if response.status_code == 403 or response.status_code in self._RETRYABLE_STATUS_CODES:
return (
b"",
response.status_code,
{
"range_header_s": header_s,
"range_first_byte_s": 0.0,
"range_body_s": 0.0,
},
)
hf_raise_for_status(response)
chunks = []
first_byte_s = 0.0
first_chunk = True
chunk_gap_s = 0.0
chunk_count = 0.0
previous_chunk_at = body_start = time.perf_counter()
for chunk in response.iter_bytes():
now = time.perf_counter()
if first_chunk:
first_byte_s = now - body_start
first_chunk = False
chunk_gap_s += now - previous_chunk_at
previous_chunk_at = now
chunk_count += 1.0
chunks.append(chunk)
body_s = time.perf_counter() - body_start
join_start = time.perf_counter()
payload = b"".join(chunks)
join_s = time.perf_counter() - join_start
return (
payload,
response.status_code,
{
"range_header_s": header_s,
"range_first_byte_s": first_byte_s,
"range_body_s": body_s,
"range_join_s": join_s,
"range_chunks": chunk_count,
"range_chunk_gap_s": chunk_gap_s,
},
)
def _record_timing(self, **kwargs: float) -> None:
with self._timing_lock:
for key, value in kwargs.items():
self._timing_totals[key] = self._timing_totals.get(key, 0.0) + value
def timing_summary(self) -> dict[str, float]:
with self._timing_lock:
return dict(self._timing_totals)
def close(self) -> None:
if self._subrange_pool is not None:
self._subrange_pool.shutdown(wait=False, cancel_futures=True)
self.client.close()
def make_range_fetcher(
data_root: str | Path,
*,
range_backend: str,
workers: int,
native_http_connections: int | None = None,
native_http_timeout: float = 60.0,
native_http_retries: int = 4,
native_http_subranges: int = 1,
token: str | bool | None = None,
):
if range_backend == "fsspec":
return ThreadLocalRangeFetcher(data_root, token=token)
if range_backend == "native-http":
max_connections = native_http_connections or max(8, workers)
return NativeHTTPRangeFetcher(
data_root,
max_connections=max_connections,
timeout=native_http_timeout,
max_retries=native_http_retries,
subrange_parts=native_http_subranges,
token=token,
)
raise ValueError(f"Unknown range backend: {range_backend}")
+176
View File
@@ -0,0 +1,176 @@
# 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
"""Revision-safe lifecycle for locally cached MP4 byte-index sidecars."""
from __future__ import annotations
import hashlib
import json
import logging
import os
import re
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from uuid import uuid4
from filelock import FileLock, Timeout
from lerobot.streaming.manifest import EpisodeVideoManifest
SIDECAR_SCHEMA_VERSION = 2
class SidecarLockTimeoutError(TimeoutError):
"""Raised when another process does not finish a sidecar build in time."""
@dataclass(frozen=True)
class SidecarSpec:
repo_id: str
revision: str
data_root: str
source_files: tuple[tuple[str, int | None], ...]
schema_version: int = SIDECAR_SCHEMA_VERSION
def __post_init__(self) -> None:
if not self.repo_id:
raise ValueError("repo_id must not be empty")
if not self.revision:
raise ValueError("revision must not be empty")
normalized = tuple(
sorted((str(path), None if size is None else int(size)) for path, size in self.source_files)
)
if any(not path or size is not None and size < 0 for path, size in normalized):
raise ValueError("source file paths must be non-empty and sizes must be non-negative")
object.__setattr__(self, "source_files", normalized)
def to_dict(self) -> dict[str, object]:
return {
"schema_version": self.schema_version,
"repo_id": self.repo_id,
"revision": self.revision,
"data_root": self.data_root,
"source_files": [{"path": path, "size": size} for path, size in self.source_files],
}
@classmethod
def from_dict(cls, data: dict[str, object]) -> SidecarSpec:
source_files = data.get("source_files")
if not isinstance(source_files, list):
raise ValueError("MP4 sidecar source_files must be a list")
parsed_files: list[tuple[str, int | None]] = []
for item in source_files:
if not isinstance(item, dict) or not isinstance(item.get("path"), str):
raise ValueError("Invalid MP4 sidecar source file entry")
size = item.get("size")
parsed_files.append((item["path"], None if size is None else int(size)))
return cls(
schema_version=int(data["schema_version"]),
repo_id=str(data["repo_id"]),
revision=str(data["revision"]),
data_root=str(data["data_root"]),
source_files=tuple(parsed_files),
)
def with_source_files(self, source_files: tuple[tuple[str, int], ...]) -> SidecarSpec:
return SidecarSpec(
repo_id=self.repo_id,
revision=self.revision,
data_root=self.data_root,
source_files=source_files,
schema_version=self.schema_version,
)
def matches(self, candidate: SidecarSpec) -> bool:
if (
self.schema_version != candidate.schema_version
or self.repo_id != candidate.repo_id
or self.revision != candidate.revision
or self.data_root != candidate.data_root
):
return False
expected = dict(self.source_files)
actual = dict(candidate.source_files)
if expected.keys() != actual.keys():
return False
return all(size is None or actual[path] == size for path, size in expected.items())
SidecarBuilder = Callable[[Path, SidecarSpec], None]
SidecarDownloader = Callable[[Path, SidecarSpec], bool]
def sidecar_cache_path(cache_root: str | Path, spec: SidecarSpec) -> Path:
identity = json.dumps(
{
"schema_version": spec.schema_version,
"repo_id": spec.repo_id,
"revision": spec.revision,
"data_root": spec.data_root,
},
sort_keys=True,
separators=(",", ":"),
)
digest = hashlib.sha256(identity.encode()).hexdigest()[:16]
repo_slug = re.sub(r"[^A-Za-z0-9_.-]+", "--", spec.repo_id).strip("-") or "dataset"
revision_slug = re.sub(r"[^A-Za-z0-9_.-]+", "-", spec.revision).strip("-")[:32] or "revision"
return (
Path(cache_root).expanduser() / repo_slug / f"mp4-v{spec.schema_version}-{revision_slug}-{digest}.npz"
)
def ensure_mp4_sidecar(
spec: SidecarSpec,
cache_root: str | Path,
*,
build: SidecarBuilder,
download: SidecarDownloader | None = None,
lock_timeout_s: float = 30 * 60,
) -> Path:
"""Return a valid local sidecar, downloading or building it exactly once.
This function never uploads. ``download`` and ``build`` must write only to the temporary path
provided to them; a validated file becomes visible at the cache path through ``os.replace``.
"""
destination = sidecar_cache_path(cache_root, spec)
if EpisodeVideoManifest.validate_file_sidecar(destination, spec):
return destination
destination.parent.mkdir(parents=True, exist_ok=True)
lock_path = destination.with_suffix(f"{destination.suffix}.lock")
try:
with FileLock(lock_path, timeout=lock_timeout_s):
if EpisodeVideoManifest.validate_file_sidecar(destination, spec):
return destination
temporary = destination.parent / f".{destination.name}.{uuid4().hex}.tmp.npz"
try:
if download is not None:
logging.info("Looking for published MP4 sidecar for %s@%s", spec.repo_id, spec.revision)
if download(temporary, spec) and EpisodeVideoManifest.validate_file_sidecar(
temporary, spec
):
os.replace(temporary, destination)
return destination
temporary.unlink(missing_ok=True)
logging.info("Building MP4 sidecar for %s@%s", spec.repo_id, spec.revision)
build(temporary, spec)
if not EpisodeVideoManifest.validate_file_sidecar(temporary, spec):
raise ValueError("Built MP4 sidecar failed revision and source validation")
os.replace(temporary, destination)
return destination
finally:
temporary.unlink(missing_ok=True)
except Timeout as exc:
raise SidecarLockTimeoutError(
f"Timed out waiting {lock_timeout_s:g}s for MP4 sidecar lock {lock_path}"
) from exc
@@ -67,15 +67,7 @@ class BiSOLeader(BimanualMixin, Teleoperator):
@cached_property
def feedback_features(self) -> dict[str, type]:
# Bimanual teleop has feedback (can be actuated for handover).
# Return the same structure as action_features for consistency with left/right arms.
left_arm_features = self.left_arm.feedback_features
right_arm_features = self.right_arm.feedback_features
return {
**{f"left_{k}": v for k, v in left_arm_features.items()},
**{f"right_{k}": v for k, v in right_arm_features.items()},
}
return {}
def setup_motors(self) -> None:
self.left_arm.setup_motors()
@@ -95,43 +87,6 @@ class BiSOLeader(BimanualMixin, Teleoperator):
return action_dict
def enable_torque(self) -> None:
"""Enable torque on both leader arms for smooth handover."""
self.left_arm.enable_torque()
self.right_arm.enable_torque()
def disable_torque(self) -> None:
"""Disable torque on both leader arms to allow human control."""
self.left_arm.disable_torque()
self.right_arm.disable_torque()
@check_if_not_connected
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Route bimanual feedback to left and right arms with proper prefix stripping.
Receives feedback dict with keys like: left_shoulder_pan.pos, right_shoulder_pan.pos, ...
Splits and routes to each arm by removing the prefix.
This enables DAgger smooth handover: when transitioning from policy control to human
intervention, both leader arms are commanded to the follower's current pose to avoid
discontinuities.
"""
# Split feedback by arm prefix
left_feedback = {}
right_feedback = {}
for key, value in feedback.items():
if key.startswith("left_"):
# Strip "left_" prefix and pass to left arm
stripped_key = key[5:] # len("left_") == 5
left_feedback[stripped_key] = value
elif key.startswith("right_"):
# Strip "right_" prefix and pass to right arm
stripped_key = key[6:] # len("right_") == 6
right_feedback[stripped_key] = value
# Send to each arm
if left_feedback:
self.left_arm.send_feedback(left_feedback)
if right_feedback:
self.right_arm.send_feedback(right_feedback)
# TODO: Implement force feedback
raise NotImplementedError
@@ -171,13 +171,7 @@ class IOSPhone(BasePhone, Teleoperator):
# HEBI provides orientation in w, x, y, z format.
# Scipy's Rotation expects x, y, z, w.
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
# ARKit can emit zero/NaN quaternions before tracking is ready or on a
# dropped packet. Rotation.from_quat now rejects those; degrade the same
# way as a missing pose so teleop stays alive mid-session.
try:
rot = Rotation.from_quat(quat_xyzw)
except ValueError:
return False, None, None, None
rot = Rotation.from_quat(quat_xyzw)
pos = ar_pos - rot.apply(self.config.camera_offset)
return True, pos, rot, pose
@@ -29,12 +29,6 @@ class SOLeaderConfig:
# Whether to use degrees for angles
use_degrees: bool = True
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
# at once, which otherwise aborts the teleoperation loop. Retries are immediate (no sleep) and only
# happen on failure, so the steady-state read cost is unchanged.
num_read_retries: int = 2
@TeleoperatorConfig.register_subclass("so101_leader")
@TeleoperatorConfig.register_subclass("so100_leader")
@@ -145,7 +145,7 @@ class SOLeader(Teleoperator):
@check_if_not_connected
def get_action(self) -> dict[str, float]:
start = time.perf_counter()
action = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
action = self.bus.sync_read("Present_Position")
action = {f"{motor}.pos": val for motor, val in action.items()}
dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read action: {dt_ms:.1f}ms")
+7 -7
View File
@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):
def __init__(
self,
transforms: Sequence[Callable[..., Any]],
transforms: Sequence[Callable],
p: list[float] | None = None,
n_subset: int | None = None,
random_order: bool = False,
@@ -50,7 +50,7 @@ class RandomSubsetApply(Transform):
if not isinstance(transforms, Sequence):
raise TypeError("Argument transforms should be a sequence of callables")
if p is None:
p = [1.0] * len(transforms)
p = [1] * len(transforms)
elif len(p) != len(transforms):
raise ValueError(
f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
@@ -69,7 +69,7 @@ class RandomSubsetApply(Transform):
self.n_subset = n_subset
self.random_order = random_order
self.selected_transforms: list[Callable[..., Any]] = []
self.selected_transforms = None
def forward(self, *inputs: Any) -> Any:
needs_unpacking = len(inputs) > 1
@@ -119,7 +119,7 @@ class SharpnessJitter(Transform):
super().__init__()
self.sharpness = self._check_input(sharpness)
def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]:
def _check_input(self, sharpness):
if isinstance(sharpness, (int | float)):
if sharpness < 0:
raise ValueError("If sharpness is a single number, it must be non negative.")
@@ -215,7 +215,7 @@ class ImageTransformsConfig:
)
def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
def make_transform_from_config(cfg: ImageTransformConfig):
if cfg.type == "SharpnessJitter":
return SharpnessJitter(**cfg.kwargs)
@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
super().__init__()
self._cfg = cfg
self.weights: list[float] = []
self.transforms: dict[str, Transform] = {}
self.weights = []
self.transforms = {}
for tf_name, tf_cfg in cfg.tfs.items():
if tf_cfg.weight <= 0.0:
continue
+4 -13
View File
@@ -37,25 +37,16 @@ def auto_select_torch_device() -> torch.device:
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
"""Given a string, return a torch.device with checks on whether the device is available.
Raises:
ValueError: If the requested device family is known but not available on
this machine (``AssertionError`` was previously used and is easy to
mistake for a programmer bug under ``python -O`` where asserts vanish).
"""
"""Given a string, return a torch.device with checks on whether the device is available."""
try_device = str(try_device)
if try_device.startswith("cuda"):
if not torch.cuda.is_available():
raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
assert torch.cuda.is_available()
device = torch.device(try_device)
elif try_device == "mps":
if not torch.backends.mps.is_available():
raise ValueError("Requested device 'mps' but MPS is not available.")
assert torch.backends.mps.is_available()
device = torch.device("mps")
elif try_device == "xpu":
if not torch.xpu.is_available():
raise ValueError("Requested device 'xpu' but XPU is not available.")
assert torch.xpu.is_available()
device = torch.device("xpu")
elif try_device == "cpu":
device = torch.device("cpu")
+5 -5
View File
@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
Returns:
Any: The data loaded from the JSON file.
"""
with open(fpath, encoding="utf-8") as f:
with open(fpath) as f:
return json.load(f)
def write_json(data: JsonLike, fpath: Path) -> None:
"""Write JSON-serializable data to a file.
def write_json(data: dict, fpath: Path) -> None:
"""Write data to a JSON file.
Creates parent directories if they don't exist.
Args:
data: JSON-serializable data to write.
data (dict): The dictionary to write.
fpath (Path): The path to the output JSON file.
"""
fpath.parent.mkdir(exist_ok=True, parents=True)
with open(fpath, "w", encoding="utf-8") as f:
with open(fpath, "w") as f:
json.dump(data, f, indent=4, ensure_ascii=False)
-28
View File
@@ -16,39 +16,11 @@
# limitations under the License.
import logging
import multiprocessing
import os
import signal
import sys
def ensure_multiprocessing_start_method(start_method: str | None) -> None:
"""Set a multiprocessing start method once, or verify the existing method matches.
Passing ``None`` leaves Python's process-wide default untouched. This is useful
when LeRobot is embedded in an application that owns multiprocessing setup.
"""
if start_method is None:
return
available_methods = multiprocessing.get_all_start_methods()
if start_method not in available_methods:
raise ValueError(
f"Multiprocessing start method must be one of {available_methods} on this platform, "
f"got {start_method!r}."
)
current_method = multiprocessing.get_start_method(allow_none=True)
if current_method is None:
multiprocessing.set_start_method(start_method)
elif current_method != start_method:
raise RuntimeError(
f"Multiprocessing start method is already {current_method!r}; cannot change it to "
f"{start_method!r}. Set the configured multiprocessing context to null to keep the "
"application's existing method, or launch LeRobot in a fresh process."
)
class ProcessSignalHandler:
"""Utility class to attach graceful shutdown signal handlers.
-4
View File
@@ -85,8 +85,6 @@ def serialize_torch_rng_state() -> dict[str, torch.Tensor]:
torch_rng_state_dict = {"torch_rng_state": torch.get_rng_state()}
if torch.cuda.is_available():
torch_rng_state_dict["torch_cuda_rng_state"] = torch.cuda.get_rng_state()
if torch.backends.mps.is_available():
torch_rng_state_dict["torch_mps_rng_state"] = torch.mps.get_rng_state()
return torch_rng_state_dict
@@ -97,8 +95,6 @@ def deserialize_torch_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None
torch.set_rng_state(rng_state_dict["torch_rng_state"])
if torch.cuda.is_available() and "torch_cuda_rng_state" in rng_state_dict:
torch.cuda.set_rng_state(rng_state_dict["torch_cuda_rng_state"])
if torch.backends.mps.is_available() and "torch_mps_rng_state" in rng_state_dict:
torch.mps.set_rng_state(rng_state_dict["torch_mps_rng_state"])
def serialize_rng_state() -> dict[str, torch.Tensor]:
-4
View File
@@ -30,10 +30,6 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
"""
if seconds <= 0:
return
if spin_threshold < 0:
raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
if sleep_margin < 0:
raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")
system = platform.system()
# On macOS and Windows the scheduler / sleep granularity can make
+3 -6
View File
@@ -29,13 +29,10 @@ class Rotation:
def __init__(self, quat: np.ndarray) -> None:
"""Initialize rotation from quaternion [x, y, z, w]."""
self._quat = np.asarray(quat, dtype=float)
if self._quat.shape != (4,):
raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
# Normalize quaternion. Reject the zero vector — it has no orientation.
# Normalize quaternion
norm = np.linalg.norm(self._quat)
if norm <= 0.0 or not np.isfinite(norm):
raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})")
self._quat = self._quat / norm
if norm > 0:
self._quat = self._quat / norm
@classmethod
def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":
+2 -2
View File
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import NotRequired, TypedDict
from typing import TypedDict
import torch
@@ -28,7 +28,7 @@ class Transition(TypedDict):
next_state: dict[str, torch.Tensor]
done: bool
truncated: bool
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
complementary_info: dict[str, torch.Tensor | float | int] | None = None
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
+12 -16
View File
@@ -24,6 +24,7 @@ import sys
import time
from collections.abc import Iterator
from copy import copy, deepcopy
from datetime import datetime
from pathlib import Path
from statistics import mean
from typing import TYPE_CHECKING, Any
@@ -60,16 +61,14 @@ def init_logging(
accelerator: Optional Accelerator instance (for multi-GPU detection)
"""
class LeRobotFormatter(logging.Formatter):
def format(self, record: logging.LogRecord) -> str:
record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:]
record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else ""
return super().format(record)
def custom_format(record: logging.LogRecord) -> str:
dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
fnameline = f"{record.pathname}:{record.lineno}"
pid_str = f"[PID: {os.getpid()}] " if display_pid else ""
return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}"
formatter = LeRobotFormatter(
"%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
formatter = logging.Formatter()
formatter.format = custom_format
logger = logging.getLogger()
logger.setLevel(logging.NOTSET)
@@ -134,13 +133,10 @@ def say(text: str, blocking: bool = False):
else:
raise RuntimeError("Unsupported operating system for text-to-speech.")
try:
if blocking:
subprocess.run(cmd, check=True, timeout=5)
else:
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
except (FileNotFoundError, subprocess.TimeoutExpired) as e:
logging.warning("Text-to-speech command failed: %s | Error: %s", cmd, e)
if blocking:
subprocess.run(cmd, check=True)
else:
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
+1 -68
View File
@@ -20,7 +20,7 @@
# ```
from pathlib import Path
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import cv2
import numpy as np
@@ -123,73 +123,6 @@ def test_invalid_width_connect():
camera.connect(warmup=False)
def test_connect_cleans_up_after_settings_failure_and_allows_retry():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
camera = OpenCVCamera(config)
opened_captures = []
def fail_settings():
opened_captures.append(camera.videocapture)
raise RuntimeError("settings failed")
with (
patch.object(camera, "_configure_capture_settings", side_effect=fail_settings),
pytest.raises(RuntimeError, match="settings failed"),
):
camera.connect(warmup=False)
assert camera.videocapture is None
assert camera.thread is None
assert not camera.is_connected
assert opened_captures[0] is not None
assert not opened_captures[0].isOpened()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=1)
camera = OpenCVCamera(config)
read_threads = []
def fail_warmup(*_args, **_kwargs):
read_threads.append(camera.thread)
raise TimeoutError("no frame")
with (
patch.object(camera, "async_read", side_effect=fail_warmup),
pytest.raises(TimeoutError, match="no frame"),
):
camera.connect()
assert camera.videocapture is None
assert camera.thread is None
assert not camera.is_connected
assert read_threads[0] is not None
assert not read_threads[0].is_alive()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_find_cameras_releases_unopened_handles():
module_path = OpenCVCamera.__module__
unopened_capture = MagicMock()
unopened_capture.isOpened.return_value = False
with (
patch(f"{module_path}.platform.system", return_value="Darwin"),
patch(f"{module_path}.MAX_OPENCV_INDEX", 1),
patch(f"{module_path}.cv2.VideoCapture", return_value=unopened_capture),
):
assert OpenCVCamera.find_cameras() == []
unopened_capture.release.assert_called_once_with()
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
def test_read(index_or_path):
config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
+1 -259
View File
@@ -20,7 +20,7 @@
# ```
from pathlib import Path
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import numpy as np
import pytest
@@ -30,8 +30,6 @@ from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnected
pytest.importorskip("pyrealsense2")
import pyrealsense2 as rs
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
@@ -63,17 +61,6 @@ def test_abc_implementation():
_ = RealSenseCamera(config)
@pytest.mark.parametrize("option", ["exposure", "gain", "white_balance"])
def test_manual_color_option_requires_rgb(option):
with pytest.raises(ValueError, match="use_rgb=True"):
RealSenseCameraConfig(
serial_number_or_name="042",
use_rgb=False,
use_depth=True,
**{option: 100},
)
def test_connect():
config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
@@ -96,27 +83,6 @@ def test_connect_invalid_camera_path(patch_realsense):
camera.connect(warmup=False)
def test_connect_cleans_up_when_sensor_configuration_fails():
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
camera = RealSenseCamera(config)
pipeline = MagicMock()
pipeline.start.return_value = MagicMock()
with (
patch("lerobot.cameras.realsense.camera_realsense.rs.pipeline", return_value=pipeline),
patch.object(camera, "_configure_rs_pipeline_config"),
patch.object(camera, "_configure_capture_settings"),
patch.object(camera, "_configure_sensor_options", side_effect=ValueError("invalid exposure")),
pytest.raises(ValueError, match="invalid exposure"),
):
camera.connect(warmup=False)
pipeline.stop.assert_called_once_with()
assert camera.rs_pipeline is None
assert camera.rs_profile is None
assert not camera.is_connected
def test_invalid_width_connect():
config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
camera = RealSenseCamera(config)
@@ -125,33 +91,6 @@ def test_invalid_width_connect():
camera.connect(warmup=False)
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
camera = RealSenseCamera(config)
read_threads = []
def fail_warmup(*_args, **_kwargs):
read_threads.append(camera.thread)
raise TimeoutError("no frame")
with (
patch.object(camera, "async_read", side_effect=fail_warmup),
pytest.raises(TimeoutError, match="no frame"),
):
camera.connect()
assert camera.rs_pipeline is None
assert camera.rs_profile is None
assert camera.thread is None
assert not camera.is_connected
assert read_threads[0] is not None
assert not read_threads[0].is_alive()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_read():
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
with RealSenseCamera(config) as camera:
@@ -289,203 +228,6 @@ def test_read_latest_too_old():
_ = camera.read_latest(max_age_ms=0) # immediately too old
def _make_mock_sensor(name: str, supported_options: set | None = None) -> MagicMock:
"""Build a fake rs.sensor that reports a name and a configurable supported-options set."""
supported = supported_options if supported_options is not None else set()
sensor = MagicMock()
sensor.get_info.return_value = name
sensor.supports.side_effect = lambda opt: opt in supported
return sensor
def _attach_mock_color_sensor(camera: RealSenseCamera, sensor: MagicMock) -> None:
"""Wire camera.rs_profile so _get_color_sensor finds the given sensor."""
profile = MagicMock()
device = MagicMock()
device.query_sensors.return_value = [sensor]
profile.get_device.return_value = device
camera.rs_profile = profile
def test_get_color_sensor_prefers_rgb_camera():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
rgb = _make_mock_sensor("RGB Camera")
stereo = _make_mock_sensor("Stereo Module")
profile = MagicMock()
device = MagicMock()
device.query_sensors.return_value = [stereo, rgb]
profile.get_device.return_value = device
camera.rs_profile = profile
assert camera._get_color_sensor() is rgb
def test_get_color_sensor_falls_back_to_stereo_module():
"""D405 has no separate RGB module; color comes from Stereo Module."""
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
stereo = _make_mock_sensor("Stereo Module")
_attach_mock_color_sensor(camera, stereo)
assert camera._get_color_sensor() is stereo
def test_get_color_sensor_raises_with_available_sensors():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
other = _make_mock_sensor("Motion Module")
_attach_mock_color_sensor(camera, other)
with pytest.raises(RuntimeError, match="Motion Module"):
camera._get_color_sensor()
def test_configure_sensor_options_skipped_when_none():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
with patch.object(RealSenseCamera, "_get_color_sensor") as mock_get:
camera._configure_sensor_options()
mock_get.assert_not_called()
def test_configure_sensor_options_applies_all_values():
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120, gain=64, white_balance=4600)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={
rs.option.enable_auto_exposure,
rs.option.exposure,
rs.option.gain,
rs.option.enable_auto_white_balance,
rs.option.white_balance,
},
)
_attach_mock_color_sensor(camera, sensor)
camera._configure_sensor_options()
sensor.set_option.assert_any_call(rs.option.enable_auto_exposure, 0)
sensor.set_option.assert_any_call(rs.option.exposure, 120)
sensor.set_option.assert_any_call(rs.option.gain, 64)
sensor.set_option.assert_any_call(rs.option.enable_auto_white_balance, 0)
sensor.set_option.assert_any_call(rs.option.white_balance, 4600)
@pytest.mark.parametrize(
("config_field", "option", "label"),
[
("exposure", rs.option.exposure, "exposure"),
("gain", rs.option.gain, "gain"),
("white_balance", rs.option.white_balance, "white balance"),
],
)
def test_configure_sensor_options_raises_when_requested_option_is_unsupported(config_field, option, label):
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: 100})
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options=set())
_attach_mock_color_sensor(camera, sensor)
with pytest.raises(ValueError, match=label):
camera._configure_sensor_options()
sensor.supports.assert_any_call(option)
sensor.set_option.assert_not_called()
@pytest.mark.parametrize(
("config_field", "option", "value"),
[
("exposure", rs.option.exposure, 120),
("gain", rs.option.gain, 64),
],
)
def test_configure_sensor_options_exposure_or_gain_disables_auto_exposure(config_field, option, value):
"""white_balance=None should not touch auto white balance."""
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: value})
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={rs.option.enable_auto_exposure, option},
)
_attach_mock_color_sensor(camera, sensor)
camera._configure_sensor_options()
calls = [call.args for call in sensor.set_option.call_args_list]
assert (rs.option.enable_auto_exposure, 0) in calls
assert (option, value) in calls
for opt, _ in calls:
assert opt != rs.option.enable_auto_white_balance
assert opt != rs.option.white_balance
def test_configure_sensor_options_warns_when_auto_exposure_control_is_unsupported(caplog):
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.exposure})
_attach_mock_color_sensor(camera, sensor)
with caplog.at_level("WARNING"):
camera._configure_sensor_options()
sensor.set_option.assert_called_once_with(rs.option.exposure, 120)
assert "does not support disabling auto-exposure" in caplog.text
def test_configure_sensor_options_warns_when_auto_white_balance_control_is_unsupported(caplog):
config = RealSenseCameraConfig(serial_number_or_name="042", white_balance=4600)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.white_balance})
_attach_mock_color_sensor(camera, sensor)
with caplog.at_level("WARNING"):
camera._configure_sensor_options()
sensor.set_option.assert_called_once_with(rs.option.white_balance, 4600)
assert "does not support disabling auto white balance" in caplog.text
def test_configure_sensor_options_out_of_range_raises_value_error():
"""set_option errors should be re-raised as ValueError with range diagnostics."""
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=999999)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={rs.option.enable_auto_exposure, rs.option.exposure},
)
def fake_set_option(option, value):
if option == rs.option.exposure:
raise RuntimeError("value out of range")
sensor.set_option.side_effect = fake_set_option
option_range = MagicMock(min=1, max=10000, step=1, default=156)
sensor.get_option_range.return_value = option_range
_attach_mock_color_sensor(camera, sensor)
with pytest.raises(ValueError, match="exposure") as exc_info:
camera._configure_sensor_options()
msg = str(exc_info.value)
assert "999999" in msg
assert "min=1" in msg
assert "max=10000" in msg
@pytest.mark.parametrize(
"rotation",
[
+22
View File
@@ -36,3 +36,25 @@ def test_dataset_config_none_episodes_ok():
def test_dataset_config_empty_episodes_ok():
DatasetConfig(repo_id="user/repo", episodes=[])
def test_dataset_config_derives_streaming_decoder_limit_by_default():
assert DatasetConfig(repo_id="user/repo").streaming_max_open_decoders is None
@pytest.mark.parametrize(
("field", "value", "message"),
[
("streaming_episode_pool_size", 0, "episode_pool_size"),
("streaming_prefetch_episodes", -1, "prefetch_episodes"),
("streaming_byte_budget_gb", 0, "byte_budget_gb"),
("streaming_decode_threads", 0, "decode_threads"),
("streaming_decoded_queue_size", 0, "decoded_queue_size"),
("streaming_max_open_decoders", 0, "max_open_decoders"),
("streaming_native_http_connections", 0, "native_http_connections"),
("streaming_native_http_subranges", 0, "native_http_subranges"),
],
)
def test_dataset_config_rejects_invalid_streaming_resource_limits(field, value, message):
with pytest.raises(ValueError, match=message):
DatasetConfig(repo_id="user/repo", **{field: value})
-24
View File
@@ -66,27 +66,3 @@ def test_from_pretrained_raises_when_no_root_config_and_no_checkpoints(monkeypat
with pytest.raises(FileNotFoundError, match="train_config.json not found"):
TrainPipelineConfig.from_pretrained("user/empty-repo")
@pytest.mark.parametrize("pass_dir", [False, True])
def test_resolve_resume_checkpoint_accepts_file_or_pretrained_model_dir(tmp_path, monkeypatch, pass_dir):
"""`--config_path` may point at the checkpoint's train_config.json or at its
pretrained_model/ directory; both must resolve `policy.pretrained_path` to the
pretrained_model/ directory (regression test for the directory case, which
previously resolved one level too high and failed on model.safetensors)."""
pretrained_dir = tmp_path / "checkpoints" / "000002" / "pretrained_model"
pretrained_dir.mkdir(parents=True)
(pretrained_dir / "train_config.json").touch()
target = pretrained_dir if pass_dir else pretrained_dir / "train_config.json"
from lerobot.policies.act.configuration_act import ACTConfig
cfg = tc.draccus.parse(TrainPipelineConfig, args=["--dataset.repo_id", "u/d"])
cfg.policy = ACTConfig()
cfg.resume = True
monkeypatch.setattr(tc.parser, "parse_arg", lambda name: str(target) if name == "config_path" else None)
cfg._resolve_resume_checkpoint()
assert cfg.policy.pretrained_path == pretrained_dir
assert cfg.checkpoint_path == pretrained_dir.parent
@@ -1,104 +0,0 @@
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import pytest
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.scripts.augment_dataset_quantile_stats import (
compute_quantile_stats_for_dataset,
has_quantile_stats,
)
def _numeric_keys(dataset):
return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
def _image_keys(dataset):
return [k for k, v in dataset.features.items() if v["dtype"] in ("image", "video")]
def test_numeric_stats_are_unaffected_by_sampling(tmp_path, lerobot_dataset_factory):
"""Sampling only touches image/video frames; numeric features are read in
full either way, so their stats must be identical with and without sampling."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
numeric_keys = _numeric_keys(dataset)
assert numeric_keys, "fixture should expose numeric features"
for key in numeric_keys:
if key not in exact:
continue
for stat in ("mean", "std", "q01", "q50", "q99"):
if stat in exact[key]:
np.testing.assert_allclose(
sampled[key][stat],
exact[key][stat],
rtol=1e-6,
atol=1e-6,
err_msg=f"numeric feature '{key}' stat '{stat}' changed under sampling",
)
def test_image_sampling_reduces_data_but_keeps_stats_close(tmp_path, lerobot_dataset_factory):
"""For images, sampling should reduce the number of samples considered while
keeping the resulting statistics close to the exact ones."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
image_keys = _image_keys(dataset)
assert image_keys, "fixture should expose at least one image feature"
for key in image_keys:
# sampling actually looked at fewer pixels
assert sampled[key]["count"][0] < exact[key]["count"][0]
# but per-channel mean stays close
np.testing.assert_allclose(
sampled[key]["mean"],
exact[key]["mean"],
rtol=0.15,
err_msg=f"image feature '{key}' mean drifted too far under sampling",
)
def test_short_episodes_use_all_frames(tmp_path, lerobot_dataset_factory):
"""With episodes shorter than the sampling floor, sampling is a no-op and
must produce exactly the same stats as the exact path."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=40, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
for key in _image_keys(dataset):
assert sampled[key]["count"][0] == exact[key]["count"][0]
def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory):
"""The computed stats should contain quantile keys for the dataset."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=200, use_videos=False
)
stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
assert has_quantile_stats(stats)

Some files were not shown because too many files have changed in this diff Show More