mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 21:19:40 +00:00
Compare commits
44 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 21321d7a27 | |||
| fdf6014214 | |||
| 4db03ed535 | |||
| 79cfb52a71 | |||
| c0e9b0bbff | |||
| b5a13e43ce | |||
| 1e7e0b6de5 | |||
| a3edab661b | |||
| 564ba6395a | |||
| 85086fed7a | |||
| ee06d1005f | |||
| 1c47809bf6 | |||
| 3d70b21aac | |||
| b85620657f | |||
| fbfc861cf2 | |||
| 06aa6a0425 | |||
| be64ded80f | |||
| 88843ed675 | |||
| f2b5c4a47b | |||
| 9202fcea96 | |||
| ef47c35178 | |||
| 6d6c82eb8c | |||
| 9201be92cb | |||
| 0064a06205 | |||
| 710171ccac | |||
| 0f8257443c | |||
| 3a09d0c48a | |||
| 03fc5e3ea9 | |||
| 28c3e095bf | |||
| 5bfb749a9b | |||
| 51c023a7a1 | |||
| 51ea18cb7a | |||
| 04ab43b8d2 | |||
| cdfe192491 | |||
| 3451e53452 | |||
| 30849ce74f | |||
| 7d6907c444 | |||
| d99e1fe89d | |||
| 7fcde61b69 | |||
| bdfe8f8ce9 | |||
| 34d0495d03 | |||
| 834c282631 | |||
| f132885cbc | |||
| d0686be2f5 |
+7
-11
@@ -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
|
||||
|
||||
@@ -68,16 +68,17 @@ ENV HOME=/home/user_lerobot \
|
||||
# issues with MuJoCo and OpenGL drivers.
|
||||
RUN uv venv --python python${PYTHON_VERSION}
|
||||
|
||||
# Install third-party dependencies separately for layer caching
|
||||
# Install Python dependencies for caching
|
||||
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
|
||||
RUN uv sync --locked --extra all --no-install-project --no-cache
|
||||
COPY --chown=user_lerobot:user_lerobot src/ src/
|
||||
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
|
||||
RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
|
||||
|
||||
# Copy the application source code and install the local project
|
||||
# Copy the rest of the application source code
|
||||
# Make sure to have the git-LFS files for testing
|
||||
COPY --chown=user_lerobot:user_lerobot . .
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
|
||||
# Set the default command
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
@@ -60,14 +60,15 @@ ENV HOME=/home/user_lerobot \
|
||||
# run other Python projects in the same container without dependency conflicts.
|
||||
RUN uv venv
|
||||
|
||||
# Install third-party dependencies separately for layer caching
|
||||
# Install Python dependencies for caching
|
||||
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
|
||||
RUN uv sync --locked --extra all --no-install-project --no-cache
|
||||
COPY --chown=user_lerobot:user_lerobot src/ src/
|
||||
|
||||
# Copy the application code and install the local project
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
|
||||
# Copy the rest of the application code
|
||||
# Make sure to have the git-LFS files for testing
|
||||
COPY --chown=user_lerobot:user_lerobot . .
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
|
||||
# Set the default command
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
@@ -37,6 +37,8 @@
|
||||
- sections:
|
||||
- local: lerobot-dataset-v3
|
||||
title: Using LeRobotDataset
|
||||
- local: training_dataset_streaming
|
||||
title: Training Dataset Streaming
|
||||
- local: porting_datasets_v3
|
||||
title: Porting Large Datasets
|
||||
- local: using_dataset_tools
|
||||
|
||||
@@ -58,7 +58,7 @@ final_action = postprocessor(action)
|
||||
|
||||
## Hardware API redesign
|
||||
|
||||
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is an overview of what changed and how you can continue to work with datasets created before this pull request.
|
||||
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request.
|
||||
|
||||
### What changed?
|
||||
|
||||
@@ -129,8 +129,8 @@ python examples/backward_compatibility/replay.py \
|
||||
|
||||
Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.
|
||||
|
||||
To find these transformations, we recommend first replaying an episode of the dataset your policy was trained on using the section above.
|
||||
Then, add these same transformations to your inference script (shown here in the `record.py` script):
|
||||
To find these transformations, we recommend to first try and and replay an episode of the dataset your policy was trained on using the section above.
|
||||
Then, add these same transformations on your inference script (shown here in the `record.py` script):
|
||||
|
||||
```diff
|
||||
action_values = predict_action(
|
||||
|
||||
@@ -136,10 +136,6 @@ config = RealSenseCameraConfig(
|
||||
height=480,
|
||||
color_mode=ColorMode.RGB,
|
||||
use_depth=True,
|
||||
# Optional fixed color controls. Omit them to leave the current sensor settings unchanged.
|
||||
exposure=120,
|
||||
gain=64,
|
||||
white_balance=4600,
|
||||
rotation=Cv2Rotation.NO_ROTATION
|
||||
)
|
||||
|
||||
@@ -158,15 +154,6 @@ finally:
|
||||
```
|
||||
<!-- prettier-ignore-end -->
|
||||
|
||||
Manual color controls disable the corresponding automatic exposure or white-balance mode. Their
|
||||
supported ranges vary by camera model; an invalid value raises an error at connection time that
|
||||
includes the range reported by the sensor. Requesting an unsupported control also raises an error.
|
||||
Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
|
||||
require `use_rgb=True`.
|
||||
|
||||
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>
|
||||
|
||||
|
||||
@@ -88,6 +88,20 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
|
||||
|
||||
The same policy can work with different environment processors, and the same environment processor can work with different policies:
|
||||
|
||||
````python
|
||||
# Use SmolVLA policy with LIBERO environment
|
||||
# Use SmolVLA policy with LIBERO environment
|
||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||
env_cfg=libero_cfg,
|
||||
policy_cfg=smolvla_cfg,
|
||||
)
|
||||
smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
|
||||
# Or use ACT policy with the same LIBERO environment
|
||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||
env_cfg=libero_cfg,
|
||||
policy_cfg=act_cfg,
|
||||
)
|
||||
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
||||
```python
|
||||
# Use SmolVLA policy with LIBERO environment
|
||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||
@@ -102,7 +116,6 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||
policy_cfg=act_cfg,
|
||||
)
|
||||
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
||||
```
|
||||
|
||||
### 3. **Easier Experimentation**
|
||||
|
||||
@@ -132,7 +145,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
|
||||
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
|
||||
gripper_pos, gripper_vel], dim=-1) # 14D
|
||||
return state
|
||||
```
|
||||
````
|
||||
|
||||
### 4. **Cleaner Environment Code**
|
||||
|
||||
|
||||
@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
|
||||
For each motor you want to update:
|
||||
|
||||
1. **Select the motor** from the list by clicking on it
|
||||
2. **Click the Upgrade tab**:
|
||||
3. **Click the Online button**:
|
||||
- If a potential firmware update is found, it will be displayed in the box
|
||||
4. **Click the Upgrade button**:
|
||||
2. **Click on Upgrade tab**:
|
||||
3. **Click on Online button**:
|
||||
- If an potential firmware update is found, it will be displayed in the box
|
||||
4. **Click on Upgrade button**:
|
||||
- The update progress will be displayed
|
||||
|
||||
## Step 6: Verify Update
|
||||
|
||||
@@ -59,7 +59,6 @@ The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with
|
||||
|
||||
- `bi_openarm_mini` - Bimanual OpenArm Mini
|
||||
- `so_leader` - SO100 / SO101 leader arm
|
||||
- `bi_so_leader` - Bimanual SO100 / SO101 leader arms
|
||||
|
||||
> [!IMPORTANT]
|
||||
> The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`.
|
||||
|
||||
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
|
||||
|
||||
### Record
|
||||
|
||||
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
|
||||
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
|
||||
|
||||
```bash
|
||||
lerobot-record \
|
||||
|
||||
@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
|
||||
- [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
|
||||
|
||||
Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
|
||||
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
|
||||
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
|
||||
|
||||
Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
|
||||
|
||||
|
||||
@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
|
||||
pip install -e ".[lekiwi]"
|
||||
```
|
||||
|
||||
Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
|
||||
Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
|
||||
Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
|
||||
|
||||
# Step-by-Step Assembly Instructions
|
||||
|
||||
@@ -131,15 +131,54 @@ for batch in data_loader:
|
||||
# model.forward(batch)
|
||||
```
|
||||
|
||||
## Stream a dataset (no downloads)
|
||||
## Stream a dataset during training
|
||||
|
||||
Use `StreamingLeRobotDataset` to iterate directly from the Hub without local copies. This allows to stream large datasets without the need to downloading them onto disk or loading them onto memory, and is a key feature of the new dataset format.
|
||||
Enable training-time streaming with the public `--dataset.streaming=true` flag:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=yaak-ai/L2D-v3 \
|
||||
--dataset.streaming=true \
|
||||
--policy.type=act \
|
||||
--output_dir=outputs/train/act_streaming
|
||||
```
|
||||
|
||||
This is separate from `--dataset.streaming_encoding=true`, which controls video encoding while
|
||||
recording. Training-time streaming leaves the map-style `LeRobotDataset` path unchanged.
|
||||
|
||||
`StreamingLeRobotDataset` assigns complete episodes disjointly across distributed ranks and
|
||||
DataLoader workers, reads only the selected episode rows from Parquet, and keeps a bounded pool of
|
||||
episode video bytes. Every selected frame is visited exactly once per streaming epoch.
|
||||
|
||||
On first use, LeRobot looks for a revision-matched MP4 index sidecar. If none is published with the
|
||||
dataset, it builds one in the revision-keyed local LeRobot cache under a process lock and installs it
|
||||
atomically. Training never uploads this sidecar. Dataset maintainers can build and publish one
|
||||
explicitly with `scripts/build_mp4_sidecar.py --push`.
|
||||
|
||||
The default memory cap is 8 GiB per DataLoader worker. It can be adjusted along with episode mixing
|
||||
and prefetch:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=yaak-ai/L2D-v3 \
|
||||
--dataset.streaming=true \
|
||||
--dataset.streaming_byte_budget_gb=4 \
|
||||
--dataset.streaming_episode_pool_size=16 \
|
||||
--dataset.streaming_prefetch_episodes=4 \
|
||||
--policy.type=act \
|
||||
--output_dir=outputs/train/act_streaming
|
||||
```
|
||||
|
||||
Use `--dataset.streaming_data_root=hf://buckets/OWNER/BUCKET/PREFIX` when metadata lives in a
|
||||
dataset repository but Parquet and MP4 data are mirrored in an HF Bucket.
|
||||
|
||||
The Python API uses the same implementation:
|
||||
|
||||
```python
|
||||
from lerobot.datasets import StreamingLeRobotDataset
|
||||
|
||||
repo_id = "yaak-ai/L2D-v3"
|
||||
dataset = StreamingLeRobotDataset(repo_id) # streams directly from the Hub
|
||||
dataset = StreamingLeRobotDataset(repo_id)
|
||||
```
|
||||
|
||||
<div style="display:flex; justify-content:center; gap:12px; flex-wrap:wrap;">
|
||||
|
||||
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
|
||||
|
||||
## Reproducing π₀Fast results
|
||||
|
||||
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
|
||||
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
|
||||
|
||||
The finetuned model can be found here:
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
|
||||
## Three pipelines
|
||||
|
||||
We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
|
||||
Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
|
||||
Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
|
||||
|
||||
1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
|
||||
2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
|
||||
@@ -74,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
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
|
||||
policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
|
||||
|
||||
# Now use predict_action_chunk with RTC parameters
|
||||
inference_delay = 4 # How many steps of inference latency, this value should be calculated based on the inference latency of the policy
|
||||
inference_delay = 4 # How many steps of inference latency, this values should be calculated based on the inference latency of the policy
|
||||
|
||||
# Initialize the action queue
|
||||
action_queue = ActionQueue(policy_cfg.rtc_config)
|
||||
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
|
||||
RTCConfig(execution_horizon=10)
|
||||
```
|
||||
|
||||
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is an optimal value.
|
||||
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value.
|
||||
|
||||
**`prefix_attention_schedule`**: How to weight consistency across the overlap region.
|
||||
|
||||
|
||||
@@ -93,7 +93,7 @@ lerobot-train --help
|
||||
|
||||
## Evaluate the finetuned model and run it in real-time
|
||||
|
||||
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
|
||||
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
|
||||
Once you are logged in, you can run inference in your setup by doing:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -338,7 +338,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
|
||||
<hfoption id="Leader">
|
||||
|
||||
- Mount the leader holder onto the wrist and secure it with 4 M3x6mm screws.
|
||||
- Attach the handle to the leader holder using 1 M2x6mm screw.
|
||||
- Attach the handle to motor 5 using 1 M2x6mm screw.
|
||||
- Insert the gripper motor, secure it with 2 M2x6mm screws on each side, attach a motor horn using a M3x6mm horn screw.
|
||||
- Attach the follower trigger with 4 M3x6mm screws.
|
||||
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
# Training Dataset Streaming
|
||||
|
||||
Training-time dataset streaming lets `lerobot-train` consume a LeRobotDataset v3 without first
|
||||
downloading its complete Parquet and video payload. Enable it through the existing public switch:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=OWNER/DATASET \
|
||||
--dataset.streaming=true \
|
||||
--policy.type=act \
|
||||
--output_dir=outputs/train/act_streaming
|
||||
```
|
||||
|
||||
This feature is independent of `--dataset.streaming_encoding=true`. `streaming_encoding` controls
|
||||
how videos are written while recording; `dataset.streaming` controls how an existing dataset is read
|
||||
during training. Recording and rollout encoders are not used by this training path.
|
||||
|
||||
## How an epoch is read
|
||||
|
||||
Each distributed rank owns a deterministic, frame-balanced set of complete episodes. One logical
|
||||
exact-coverage pool per rank mixes those episodes while visiting every selected frame once per
|
||||
rank-local coverage epoch. Parquet columns and compressed MP4 byte ranges are prefetched from the
|
||||
same byte-aware admission frontier. Temporal history and future windows are resolved inside the
|
||||
complete episode, including the same boundary padding masks as map-style loading.
|
||||
|
||||
When `--num_workers` is nonzero, training uses one dedicated DataLoader process per rank. Its
|
||||
bounded result queue holds decoded batches while the policy trains. The configured worker count is
|
||||
instead used as internal Parquet and byte-range fetch concurrency, so increasing it does not create
|
||||
independent samplers or multiply the cache budget.
|
||||
|
||||
The default map-style `LeRobotDataset` behavior is unchanged when `--dataset.streaming=false`.
|
||||
|
||||
## MP4 sidecars and the first run
|
||||
|
||||
Video streaming uses a small MP4 index sidecar. Dataset initialization first checks the
|
||||
revision-keyed local cache, then looks for a valid published sidecar. If neither is available,
|
||||
LeRobot builds the sidecar locally while holding a process lock and installs it atomically. A
|
||||
failed or interrupted build does not replace the previous valid file.
|
||||
|
||||
Training is read-only: it never uploads a sidecar or modifies the dataset repository. On a cluster
|
||||
with node-local caches, the first job may build once per node. A shared LeRobot cache avoids that
|
||||
duplication.
|
||||
|
||||
Dataset maintainers can build a sidecar ahead of time:
|
||||
|
||||
```bash
|
||||
uv run python scripts/build_mp4_sidecar.py \
|
||||
--repo-id=OWNER/DATASET \
|
||||
--revision=COMMIT_SHA \
|
||||
--data-root=hf://datasets/OWNER/DATASET@COMMIT_SHA \
|
||||
--output=/tmp/dataset-mp4-sidecar.npz
|
||||
```
|
||||
|
||||
Publication is always explicit. Add `--push` only after validating the complete-dataset sidecar.
|
||||
Subset sidecars cannot be published.
|
||||
|
||||
## Configuration
|
||||
|
||||
The production defaults are:
|
||||
|
||||
| Option | Default | Meaning |
|
||||
| ----------------------------------- | -------------: | --------------------------------------------------- |
|
||||
| `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank |
|
||||
| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool |
|
||||
| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank |
|
||||
| `streaming_decode_threads` | 2 | Parallel sample assembly and video decode workers |
|
||||
| `streaming_decoded_queue_size` | 8 | Decoded samples buffered ahead, in planner order |
|
||||
| `streaming_max_open_decoders` | pool × cameras | Independent open-decoder LRU cap per rank |
|
||||
| `streaming_native_http_connections` | unset | Native HTTP connection cap per rank |
|
||||
| `streaming_native_http_subranges` | 1 | Concurrent subranges per native HTTP range read |
|
||||
| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root |
|
||||
|
||||
The active episode set is capped by both episode count and the exact synthesized mini-MP4 sizes
|
||||
computed from the sidecar. An episode larger than the complete rank budget fails before training
|
||||
fetches its payload. The cache, decoder LRU, and decoded-batch queue remain independently bounded.
|
||||
Start with a smaller pool or budget on memory-constrained hosts:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=OWNER/DATASET \
|
||||
--dataset.streaming=true \
|
||||
--dataset.streaming_episode_pool_size=16 \
|
||||
--dataset.streaming_prefetch_episodes=4 \
|
||||
--dataset.streaming_byte_budget_gb=4 \
|
||||
--dataset.streaming_decode_threads=2 \
|
||||
--dataset.streaming_decoded_queue_size=8 \
|
||||
--num_workers=4 \
|
||||
--policy.type=act \
|
||||
--output_dir=outputs/train/act_streaming
|
||||
```
|
||||
|
||||
If metadata remains in a dataset repository while payload files are mirrored elsewhere, set
|
||||
`--dataset.streaming_data_root`. Supported values include local paths, revision-qualified
|
||||
`hf://datasets/...` roots, `hf://buckets/...` roots, and fsspec URLs.
|
||||
|
||||
## Resume and shuffle migration
|
||||
|
||||
The earlier streaming reader used a bounded row shuffle buffer. The episode reader instead has
|
||||
deterministic exact-coverage ordering derived from the seed and epoch. Checkpoint resume restores the
|
||||
per-rank sample offset using the checkpoint batch size. Changing distributed world size or
|
||||
batch size changes ownership or batch boundaries. For sample-exact comparisons, resume with the
|
||||
same world size and batch size. Keep the same internal fetch concurrency when comparing performance.
|
||||
|
||||
Streaming checkpoints created by the earlier multi-worker sampler record a different ownership
|
||||
topology and are rejected for sample-exact resume. Start a new run from the saved policy weights
|
||||
rather than claiming that its dataset stream resumes exactly.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- **The first batch takes a long time:** check logs for a sidecar build. Reuse a shared cache or
|
||||
publish a validated complete sidecar explicitly.
|
||||
- **A sidecar lock times out:** another process may still be indexing the same revision. Confirm it
|
||||
is healthy before removing a stale lock.
|
||||
- **A rank owns no data:** reduce the number of ranks so every rank owns at least one selected
|
||||
episode.
|
||||
- **The byte budget is exceeded:** lower the episode pool, increase the per-rank byte budget, or
|
||||
use an explicitly prepared payload layout with smaller episode ranges.
|
||||
- **Remote reads cannot authenticate:** verify the normal Hugging Face token or fsspec credentials
|
||||
are available in every worker environment. Credentials are never embedded in the sidecar.
|
||||
- **Refill stalls are high:** compare p95/p99 batch wait, reduce network contention, raise prefetch
|
||||
gradually, and verify that the dataset sidecar matches the exact revision.
|
||||
@@ -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.
|
||||
|
||||
@@ -12,8 +12,12 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""This script demonstrates how to train a Diffusion Policy on the PushT environment,
|
||||
using a dataset processed in streaming mode."""
|
||||
"""Train directly from episode-scoped Parquet and MP4 streams.
|
||||
|
||||
For normal training, prefer ``lerobot-train --dataset.streaming=true`` so distributed sharding,
|
||||
checkpoint resume, and device placement are configured by the training pipeline. This lower-level
|
||||
example shows the underlying Python API.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
@@ -64,18 +68,28 @@ def main():
|
||||
ACTION: [t / dataset_metadata.fps for t in range(cfg.n_action_steps)],
|
||||
}
|
||||
|
||||
# Instantiating the training dataset in streaming mode allows to not consume up memory as the data is fetched
|
||||
# iteratively rather than being load into memory all at once. Retrieved frames are shuffled across epochs
|
||||
dataset = StreamingLeRobotDataset(dataset_id, delta_timestamps=delta_timestamps, tolerance_s=1e-3)
|
||||
# The first run resolves or locally builds a revision-safe MP4 index sidecar. It is never
|
||||
# uploaded implicitly. Episode rows and video byte ranges are then prefetched together.
|
||||
dataset = StreamingLeRobotDataset(
|
||||
dataset_id,
|
||||
delta_timestamps=delta_timestamps,
|
||||
tolerance_s=1e-3,
|
||||
episode_pool_size=16,
|
||||
prefetch_episodes=4,
|
||||
byte_budget_gb=4,
|
||||
repeat=True,
|
||||
)
|
||||
|
||||
optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4)
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
dataset,
|
||||
num_workers=4,
|
||||
# One worker owns the rank-level pool. Internal fetch concurrency is configured through
|
||||
# StreamingLeRobotDataset.max_num_shards (the lerobot-train CLI derives it from num_workers).
|
||||
num_workers=1,
|
||||
batch_size=16,
|
||||
pin_memory=device.type != "cpu",
|
||||
drop_last=True,
|
||||
prefetch_factor=2, # loads batches with multiprocessing while policy trains
|
||||
prefetch_factor=2, # bounded decoded-batch queue while the policy trains
|
||||
persistent_workers=True,
|
||||
)
|
||||
|
||||
# Run training loop.
|
||||
|
||||
+5
-14
@@ -67,8 +67,11 @@ dependencies = [
|
||||
"einops>=0.8.0,<0.9.0",
|
||||
|
||||
# Config & Hub
|
||||
"draccus>=0.11.6,<0.12.0",
|
||||
"draccus==0.10.0", # TODO: Relax version constraint
|
||||
"huggingface-hub>=1.0.0,<2.0.0",
|
||||
"filelock>=3.12.0,<4.0.0",
|
||||
"fsspec>=2023.5.0,<2027.0.0",
|
||||
"httpx>=0.27.0,<1.0.0",
|
||||
"requests>=2.32.0,<3.0.0",
|
||||
|
||||
# Environments
|
||||
@@ -436,6 +439,7 @@ exclude_dirs = [
|
||||
skips = ["B101", "B311", "B404", "B603", "B615"]
|
||||
|
||||
[tool.typos]
|
||||
default.extend-words = { trak = "trak" }
|
||||
default.extend-ignore-re = [
|
||||
"(?Rm)^.*(#|//)\\s*spellchecker:disable-line$", # spellchecker:disable-line
|
||||
"(?s)(#|//)\\s*spellchecker:off.*?\\n\\s*(#|//)\\s*spellchecker:on", # spellchecker:<on|off>
|
||||
@@ -494,19 +498,6 @@ ignore_errors = true
|
||||
module = "lerobot.envs.*"
|
||||
ignore_errors = false
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = "lerobot.annotations.*"
|
||||
ignore_errors = false
|
||||
disallow_untyped_defs = true
|
||||
disallow_incomplete_defs = true
|
||||
check_untyped_defs = true
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = "lerobot.transforms.*"
|
||||
ignore_errors = false
|
||||
disallow_untyped_defs = true
|
||||
disallow_incomplete_defs = true
|
||||
check_untyped_defs = true
|
||||
|
||||
# [[tool.mypy.overrides]]
|
||||
# module = "lerobot.utils.*"
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import fsspec
|
||||
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
from lerobot.datasets.streaming_sidecar import (
|
||||
build_mp4_sidecar,
|
||||
make_sidecar_spec,
|
||||
published_sidecar_url,
|
||||
range_backend_for_root,
|
||||
)
|
||||
from lerobot.streaming.sidecar import SidecarSpec
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Build a reusable MP4 byte-index sidecar for streaming.")
|
||||
parser.add_argument("--repo-id", required=True)
|
||||
parser.add_argument("--revision", default=None)
|
||||
parser.add_argument("--data-root", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--episodes", type=int, default=None)
|
||||
parser.add_argument("--workers", type=int, default=8)
|
||||
parser.add_argument("--range-backend", choices=("fsspec", "native-http"), default=None)
|
||||
parser.add_argument("--max-probe-mb", type=int, default=64)
|
||||
parser.add_argument("--push", action="store_true", help="Explicitly publish the sidecar to data_root.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def push_sidecar(local_path: str, spec: SidecarSpec) -> list[str]:
|
||||
if not spec.data_root.startswith("hf://"):
|
||||
raise ValueError("--push currently supports only hf:// data roots")
|
||||
|
||||
fs = fsspec.filesystem("hf")
|
||||
remote = published_sidecar_url(spec)
|
||||
fs.put(str(Path(local_path)), remote)
|
||||
return [remote]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
meta = LeRobotDatasetMetadata(args.repo_id, revision=args.revision)
|
||||
meta.ensure_readable()
|
||||
total = (
|
||||
int(meta.total_episodes) if args.episodes is None else min(args.episodes, int(meta.total_episodes))
|
||||
)
|
||||
spec = make_sidecar_spec(meta, args.data_root)
|
||||
if total != int(meta.total_episodes):
|
||||
selected_paths = {
|
||||
str(meta.get_video_file_path(ep_idx, key)) for ep_idx in range(total) for key in meta.video_keys
|
||||
}
|
||||
spec = SidecarSpec(
|
||||
repo_id=spec.repo_id,
|
||||
revision=spec.revision,
|
||||
data_root=spec.data_root,
|
||||
source_files=tuple(item for item in spec.source_files if item[0] in selected_paths),
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
build_mp4_sidecar(
|
||||
args.output,
|
||||
spec,
|
||||
range_backend=args.range_backend or range_backend_for_root(args.data_root),
|
||||
workers=args.workers,
|
||||
max_probe_bytes=args.max_probe_mb * 1024 * 1024,
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
print(f"wrote {args.output}")
|
||||
print(f"episodes={total} files={len(spec.source_files)} elapsed_s={elapsed:.2f}")
|
||||
if args.push:
|
||||
if total != int(meta.total_episodes):
|
||||
raise ValueError("Only a complete dataset sidecar can be published")
|
||||
pushed = push_sidecar(args.output, spec)
|
||||
for remote in pushed:
|
||||
print(f"pushed {remote}")
|
||||
else:
|
||||
print("push_skipped: pass --push for explicit publication")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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.")
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -53,7 +53,11 @@ def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Pa
|
||||
|
||||
|
||||
def save_training_step(
|
||||
step: int, save_dir: Path, num_processes: int | None = None, batch_size: int | None = None
|
||||
step: int,
|
||||
save_dir: Path,
|
||||
num_processes: int | None = None,
|
||||
batch_size: int | None = None,
|
||||
num_workers: int | None = None,
|
||||
) -> None:
|
||||
state: dict = {"step": step}
|
||||
# num_processes and batch_size are recorded so a resumed run can detect a changed world size or
|
||||
@@ -64,6 +68,8 @@ def save_training_step(
|
||||
state["num_processes"] = num_processes
|
||||
if batch_size is not None:
|
||||
state["batch_size"] = batch_size
|
||||
if num_workers is not None:
|
||||
state["num_workers"] = num_workers
|
||||
write_json(state, save_dir / TRAINING_STEP)
|
||||
|
||||
|
||||
@@ -82,6 +88,11 @@ def load_training_batch_size(checkpoint_dir: Path) -> int | None:
|
||||
return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("batch_size")
|
||||
|
||||
|
||||
def load_training_num_workers(checkpoint_dir: Path) -> int | None:
|
||||
"""DataLoader worker count recorded at checkpoint time, or None for older checkpoints."""
|
||||
return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("num_workers")
|
||||
|
||||
|
||||
def update_last_checkpoint(checkpoint_dir: Path) -> Path:
|
||||
last_checkpoint_dir = checkpoint_dir.parent / LAST_CHECKPOINT_LINK
|
||||
if last_checkpoint_dir.is_symlink():
|
||||
@@ -101,6 +112,7 @@ def save_checkpoint(
|
||||
postprocessor: PolicyProcessorPipeline | None = None,
|
||||
num_processes: int | None = None,
|
||||
batch_size: int | None = None,
|
||||
num_workers: int | None = None,
|
||||
model_state_dict: dict | None = None,
|
||||
optim_state_dict: dict | None = None,
|
||||
) -> None:
|
||||
@@ -132,6 +144,8 @@ def save_checkpoint(
|
||||
resume. Defaults to None (not recorded).
|
||||
batch_size (int | None, optional): Per-process batch size to record for sample-exact
|
||||
resume. Defaults to None (not recorded).
|
||||
num_workers (int | None, optional): Per-process DataLoader worker count to record for
|
||||
sample-exact streaming resume. Defaults to None (not recorded).
|
||||
model_state_dict: Pre-gathered full (unsharded) model state dict. Required under FSDP,
|
||||
where `policy.state_dict()` would return sharded tensors; the caller gathers it via a
|
||||
cross-rank collective and passes it here so rank 0 can write it directly. It holds
|
||||
@@ -160,6 +174,7 @@ def save_checkpoint(
|
||||
scheduler,
|
||||
num_processes=num_processes,
|
||||
batch_size=batch_size,
|
||||
num_workers=num_workers,
|
||||
optim_state_dict=optim_state_dict,
|
||||
)
|
||||
|
||||
@@ -171,6 +186,7 @@ def save_training_state(
|
||||
scheduler: LRScheduler | None = None,
|
||||
num_processes: int | None = None,
|
||||
batch_size: int | None = None,
|
||||
num_workers: int | None = None,
|
||||
optim_state_dict: dict | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
@@ -185,12 +201,20 @@ def save_training_state(
|
||||
Defaults to None.
|
||||
num_processes (int | None, optional): Distributed world size to record. Defaults to None.
|
||||
batch_size (int | None, optional): Per-process batch size to record. Defaults to None.
|
||||
num_workers (int | None, optional): Per-process DataLoader worker count to record.
|
||||
Defaults to None.
|
||||
optim_state_dict: Pre-gathered full optimizer state dict (for FSDP). Saved instead of
|
||||
`optimizer.state_dict()` when provided. Defaults to None.
|
||||
"""
|
||||
save_dir = checkpoint_dir / TRAINING_STATE_DIR
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
save_training_step(train_step, save_dir, num_processes=num_processes, batch_size=batch_size)
|
||||
save_training_step(
|
||||
train_step,
|
||||
save_dir,
|
||||
num_processes=num_processes,
|
||||
batch_size=batch_size,
|
||||
num_workers=num_workers,
|
||||
)
|
||||
save_rng_state(save_dir)
|
||||
if optimizer is not None:
|
||||
save_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict)
|
||||
|
||||
@@ -71,19 +71,13 @@ class DatasetRecordConfig:
|
||||
# Number of threads per encoder instance. None = auto (codec default).
|
||||
# Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
|
||||
encoder_threads: int | None = None
|
||||
# Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
|
||||
# (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
|
||||
no_stamp: bool = False
|
||||
|
||||
def stamp_repo_id(self) -> None:
|
||||
"""Append a date-time tag to ``repo_id`` so each recording session gets a unique name.
|
||||
|
||||
Must be called explicitly at dataset *creation* time — not on resume,
|
||||
where the existing ``repo_id`` (already stamped) must be preserved.
|
||||
No-op when ``no_stamp`` is set, preserving a user-managed ``repo_id``.
|
||||
"""
|
||||
if self.no_stamp:
|
||||
return
|
||||
if self.repo_id:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
self.repo_id = f"{self.repo_id}_{timestamp}"
|
||||
|
||||
@@ -44,6 +44,22 @@ class DatasetConfig:
|
||||
# Has no effect on datasets without depth cameras.
|
||||
depth_output_unit: str = DEFAULT_DEPTH_UNIT
|
||||
streaming: bool = False
|
||||
# Optional data-plane root used by training-time streaming. Metadata still resolves from repo_id/root.
|
||||
streaming_data_root: str | None = None
|
||||
# Number of complete episodes mixed by the rank-level exact-coverage sampler.
|
||||
streaming_episode_pool_size: int = 32
|
||||
# Complete episodes fetched ahead of the current admission frontier.
|
||||
streaming_prefetch_episodes: int = 8
|
||||
# Hard per-rank cap for synthesized episode-video bytes.
|
||||
streaming_byte_budget_gb: float = 8.0
|
||||
# Parallel sample assembly/decode workers and their bounded in-order result queue.
|
||||
streaming_decode_threads: int = 2
|
||||
streaming_decoded_queue_size: int = 8
|
||||
# Independent decoder-state cap. None covers every camera in the configured episode pool.
|
||||
streaming_max_open_decoders: int | None = None
|
||||
# Per-rank native HTTP limits. None preserves the fetcher's worker-derived default.
|
||||
streaming_native_http_connections: int | None = None
|
||||
streaming_native_http_subranges: int = 1
|
||||
# Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
|
||||
eval_split: float = 0.0
|
||||
|
||||
@@ -54,6 +70,22 @@ class DatasetConfig:
|
||||
)
|
||||
if not (0.0 <= self.eval_split < 1.0):
|
||||
raise ValueError(f"eval_split must be in [0.0, 1.0), got {self.eval_split}")
|
||||
if self.streaming_episode_pool_size <= 0:
|
||||
raise ValueError("streaming_episode_pool_size must be positive")
|
||||
if self.streaming_prefetch_episodes < 0:
|
||||
raise ValueError("streaming_prefetch_episodes must be non-negative")
|
||||
if self.streaming_byte_budget_gb <= 0:
|
||||
raise ValueError("streaming_byte_budget_gb must be positive")
|
||||
if self.streaming_decode_threads <= 0:
|
||||
raise ValueError("streaming_decode_threads must be positive")
|
||||
if self.streaming_decoded_queue_size <= 0:
|
||||
raise ValueError("streaming_decoded_queue_size must be positive")
|
||||
if self.streaming_max_open_decoders is not None and self.streaming_max_open_decoders <= 0:
|
||||
raise ValueError("streaming_max_open_decoders must be positive")
|
||||
if self.streaming_native_http_connections is not None and self.streaming_native_http_connections <= 0:
|
||||
raise ValueError("streaming_native_http_connections must be positive")
|
||||
if self.streaming_native_http_subranges <= 0:
|
||||
raise ValueError("streaming_native_http_subranges must be positive")
|
||||
if self.episodes is not None:
|
||||
if any(ep < 0 for ep in self.episodes):
|
||||
raise ValueError(
|
||||
|
||||
@@ -163,10 +163,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
|
||||
return None
|
||||
|
||||
def _save_pretrained(self, save_directory: Path) -> None:
|
||||
# Encode against the base class so draccus includes the choice "type" key,
|
||||
# which `from_pretrained` needs to resolve the concrete subclass.
|
||||
with open(save_directory / CONFIG_NAME, "w") as f:
|
||||
json.dump(draccus.encode(self, PreTrainedConfig), f, indent=4)
|
||||
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
|
||||
draccus.dump(self, f, indent=4)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
|
||||
@@ -103,10 +103,8 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
|
||||
pass
|
||||
|
||||
def _save_pretrained(self, save_directory: Path) -> None:
|
||||
# Encode against the base class so draccus includes the choice "type" key,
|
||||
# which `from_pretrained` needs to resolve the concrete subclass.
|
||||
with open(save_directory / CONFIG_NAME, "w") as f:
|
||||
json.dump(draccus.encode(self, RewardModelConfig), f, indent=4)
|
||||
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
|
||||
draccus.dump(self, f, indent=4)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
|
||||
@@ -194,11 +194,7 @@ class TrainPipelineConfig(HubMixin):
|
||||
)
|
||||
|
||||
if Path(config_path).resolve().exists():
|
||||
# `config_path` may point at the checkpoint's train_config.json or at its
|
||||
# pretrained_model/ directory (both documented above) — resolve either to
|
||||
# the pretrained_model/ directory.
|
||||
config_path_obj = Path(config_path)
|
||||
policy_dir = config_path_obj.parent if config_path_obj.is_file() else config_path_obj
|
||||
policy_dir = Path(config_path).parent
|
||||
self.checkpoint_path = policy_dir.parent
|
||||
elif self.job.is_remote:
|
||||
return
|
||||
|
||||
@@ -19,7 +19,6 @@ import copy
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Any, NotRequired, TypedDict
|
||||
|
||||
import datasets
|
||||
import pandas as pd
|
||||
@@ -50,32 +49,8 @@ from .utils import (
|
||||
)
|
||||
from .video_utils import concatenate_video_files, get_video_duration_in_s
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
type FeatureDict = dict[str, dict[str, Any]]
|
||||
type ChunkFile = tuple[int, int]
|
||||
|
||||
|
||||
class IndexState(TypedDict):
|
||||
chunk: int
|
||||
file: int
|
||||
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
|
||||
|
||||
|
||||
class VideoIndex(TypedDict):
|
||||
chunk: int
|
||||
file: int
|
||||
latest_duration: float
|
||||
episode_duration: float
|
||||
src_to_offset: NotRequired[dict[ChunkFile, float]]
|
||||
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
|
||||
dst_file_durations: NotRequired[dict[ChunkFile, float]]
|
||||
|
||||
|
||||
type VideoIndexState = dict[str, VideoIndex]
|
||||
|
||||
|
||||
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> FeatureDict:
|
||||
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]:
|
||||
"""Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.
|
||||
|
||||
Args:
|
||||
@@ -84,14 +59,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
||||
Returns:
|
||||
dict: A dictionary of merged video feature info.
|
||||
"""
|
||||
merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
|
||||
merged_info = copy.deepcopy(all_metadata[0].features)
|
||||
video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
|
||||
|
||||
for vk in video_keys:
|
||||
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
|
||||
base_video_info = video_infos[0]
|
||||
|
||||
merged_encoder_info: dict[str, Any] = {}
|
||||
merged_encoder_info: dict = {}
|
||||
fallback_keys: list[str] = []
|
||||
for info_key in VIDEO_ENCODER_INFO_KEYS:
|
||||
values = [info.get(info_key, None) for info in video_infos]
|
||||
@@ -105,7 +80,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
||||
merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None
|
||||
|
||||
if fallback_keys:
|
||||
logger.warning(
|
||||
logging.warning(
|
||||
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
|
||||
f"Setting these keys to null: {fallback_keys}.",
|
||||
)
|
||||
@@ -117,7 +92,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
||||
return merged_info
|
||||
|
||||
|
||||
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
|
||||
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
|
||||
"""Validates that all dataset metadata have consistent properties.
|
||||
|
||||
Ensures all datasets have the same fps, robot_type, and features to guarantee
|
||||
@@ -154,9 +129,7 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[i
|
||||
return fps, robot_type, features
|
||||
|
||||
|
||||
def update_data_df(
|
||||
df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
|
||||
) -> pd.DataFrame:
|
||||
def update_data_df(df, src_meta, dst_meta):
|
||||
"""Updates a data DataFrame with new indices and task mappings for aggregation.
|
||||
|
||||
Adjusts episode indices, frame indices, and task indices to account for
|
||||
@@ -181,12 +154,12 @@ def update_data_df(
|
||||
|
||||
|
||||
def update_meta_data(
|
||||
df: pd.DataFrame,
|
||||
dst_meta: LeRobotDatasetMetadata,
|
||||
meta_idx: IndexState,
|
||||
data_idx: IndexState,
|
||||
videos_idx: VideoIndexState,
|
||||
) -> pd.DataFrame:
|
||||
df,
|
||||
dst_meta,
|
||||
meta_idx,
|
||||
data_idx,
|
||||
videos_idx,
|
||||
):
|
||||
"""Updates metadata DataFrame with new chunk, file, and timestamp indices.
|
||||
|
||||
Adjusts all indices and timestamps to account for previously aggregated
|
||||
@@ -316,7 +289,7 @@ def aggregate_datasets(
|
||||
chunk_size: int | None = None,
|
||||
concatenate_videos: bool = True,
|
||||
concatenate_data: bool = True,
|
||||
) -> None:
|
||||
):
|
||||
"""Aggregates multiple LeRobot datasets into a single unified dataset.
|
||||
|
||||
This is the main function that orchestrates the aggregation process by:
|
||||
@@ -336,7 +309,7 @@ def aggregate_datasets(
|
||||
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
|
||||
concatenate_data: When False, keep one parquet per source file instead of packing into shards.
|
||||
"""
|
||||
logger.info("Start aggregate_datasets")
|
||||
logging.info("Start aggregate_datasets")
|
||||
|
||||
if data_files_size_in_mb is None:
|
||||
data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
|
||||
@@ -368,15 +341,15 @@ def aggregate_datasets(
|
||||
video_files_size_in_mb=video_files_size_in_mb,
|
||||
)
|
||||
|
||||
logger.info("Find all tasks")
|
||||
logging.info("Find all tasks")
|
||||
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
|
||||
dst_meta.tasks = pd.DataFrame(
|
||||
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
|
||||
)
|
||||
|
||||
meta_idx: IndexState = {"chunk": 0, "file": 0}
|
||||
data_idx: IndexState = {"chunk": 0, "file": 0}
|
||||
videos_idx: VideoIndexState = {
|
||||
meta_idx = {"chunk": 0, "file": 0}
|
||||
data_idx = {"chunk": 0, "file": 0}
|
||||
videos_idx = {
|
||||
key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
|
||||
}
|
||||
|
||||
@@ -400,17 +373,12 @@ def aggregate_datasets(
|
||||
dst_meta.info.total_frames += src_meta.total_frames
|
||||
|
||||
finalize_aggregation(dst_meta, all_metadata)
|
||||
logger.info("Aggregation complete.")
|
||||
logging.info("Aggregation complete.")
|
||||
|
||||
|
||||
def aggregate_videos(
|
||||
src_meta: LeRobotDatasetMetadata,
|
||||
dst_meta: LeRobotDatasetMetadata,
|
||||
videos_idx: VideoIndexState,
|
||||
video_files_size_in_mb: float,
|
||||
chunk_size: int,
|
||||
concatenate_videos: bool = True,
|
||||
) -> VideoIndexState:
|
||||
src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
|
||||
):
|
||||
"""Aggregates video chunks from a source dataset into the destination dataset.
|
||||
|
||||
Handles video file concatenation and rotation based on file size limits.
|
||||
@@ -438,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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()))
|
||||
|
||||
@@ -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
@@ -0,0 +1,149 @@
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
"""LeRobot metadata adapter for automatic MP4 sidecar resolution."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import fsspec
|
||||
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
from lerobot.streaming.manifest import EpisodeVideoManifest
|
||||
from lerobot.streaming.sidecar import SidecarSpec, ensure_mp4_sidecar, sidecar_cache_path
|
||||
from lerobot.utils.constants import HF_LEROBOT_HOME
|
||||
|
||||
DEFAULT_SIDECAR_CACHE = HF_LEROBOT_HOME / "streaming-sidecars"
|
||||
|
||||
|
||||
def range_backend_for_root(data_root: str) -> str:
|
||||
"""Use direct HTTP only for HF roots; all local and other fsspec protocols stay generic."""
|
||||
return "native-http" if data_root.startswith("hf://") else "fsspec"
|
||||
|
||||
|
||||
def streaming_data_root(
|
||||
meta: LeRobotDatasetMetadata,
|
||||
*,
|
||||
requested_root: str | Path | None,
|
||||
configured_data_root: str | None,
|
||||
) -> str:
|
||||
if configured_data_root is not None:
|
||||
return configured_data_root.rstrip("/")
|
||||
if requested_root is not None:
|
||||
return str(Path(requested_root).expanduser())
|
||||
return f"hf://datasets/{meta.repo_id}@{meta.revision}"
|
||||
|
||||
|
||||
def make_sidecar_spec(meta: LeRobotDatasetMetadata, data_root: str) -> SidecarSpec:
|
||||
relative_paths = sorted(
|
||||
{
|
||||
str(meta.get_video_file_path(episode_index, video_key))
|
||||
for episode_index in range(int(meta.total_episodes))
|
||||
for video_key in meta.video_keys
|
||||
}
|
||||
)
|
||||
root = Path(data_root).expanduser()
|
||||
source_files: tuple[tuple[str, int | None], ...]
|
||||
if root.is_dir():
|
||||
source_files = tuple((path, (root / path).stat().st_size) for path in relative_paths)
|
||||
else:
|
||||
source_files = tuple((path, None) for path in relative_paths)
|
||||
return SidecarSpec(
|
||||
repo_id=meta.repo_id,
|
||||
revision=str(meta.revision),
|
||||
data_root=data_root.rstrip("/"),
|
||||
source_files=source_files,
|
||||
)
|
||||
|
||||
|
||||
def build_mp4_sidecar(
|
||||
destination: str | Path,
|
||||
spec: SidecarSpec,
|
||||
*,
|
||||
workers: int = 8,
|
||||
range_backend: str = "native-http",
|
||||
max_probe_bytes: int = 64 * 1024 * 1024,
|
||||
token: str | bool | None = None,
|
||||
) -> None:
|
||||
EpisodeVideoManifest.write_file_sidecar(
|
||||
destination,
|
||||
[path for path, _size in spec.source_files],
|
||||
spec.data_root,
|
||||
spec=spec,
|
||||
range_backend=range_backend,
|
||||
workers=workers,
|
||||
max_probe_bytes=max_probe_bytes,
|
||||
token=token,
|
||||
)
|
||||
|
||||
|
||||
def published_sidecar_url(spec: SidecarSpec, cache_root: str | Path = DEFAULT_SIDECAR_CACHE) -> str:
|
||||
name = sidecar_cache_path(cache_root, spec).name
|
||||
return f"{spec.data_root}/meta/mp4-sidecars/{name}"
|
||||
|
||||
|
||||
def download_published_sidecar(
|
||||
destination: Path,
|
||||
spec: SidecarSpec,
|
||||
*,
|
||||
cache_root: str | Path = DEFAULT_SIDECAR_CACHE,
|
||||
token: str | bool | None = None,
|
||||
) -> bool:
|
||||
if Path(spec.data_root).expanduser().is_dir():
|
||||
return False
|
||||
source_url = published_sidecar_url(spec, cache_root)
|
||||
storage_options = {"token": token} if token is not None and source_url.startswith("hf://") else {}
|
||||
filesystem, source = fsspec.core.url_to_fs(source_url, **storage_options)
|
||||
if not filesystem.exists(source):
|
||||
return False
|
||||
with filesystem.open(source, "rb") as remote, destination.open("wb") as local:
|
||||
shutil.copyfileobj(remote, local)
|
||||
return True
|
||||
|
||||
|
||||
def ensure_dataset_mp4_sidecar(
|
||||
meta: LeRobotDatasetMetadata,
|
||||
data_root: str,
|
||||
*,
|
||||
cache_root: str | Path = DEFAULT_SIDECAR_CACHE,
|
||||
workers: int = 8,
|
||||
range_backend: str = "native-http",
|
||||
lock_timeout_s: float = 30 * 60,
|
||||
token: str | bool | None = None,
|
||||
) -> Path | None:
|
||||
if not meta.video_keys:
|
||||
return None
|
||||
|
||||
spec = make_sidecar_spec(meta, data_root)
|
||||
logging.info(
|
||||
"Resolving training-time MP4 sidecar for %s@%s (%d files)",
|
||||
spec.repo_id,
|
||||
spec.revision,
|
||||
len(spec.source_files),
|
||||
)
|
||||
return ensure_mp4_sidecar(
|
||||
spec,
|
||||
cache_root,
|
||||
build=lambda path, target_spec: build_mp4_sidecar(
|
||||
path,
|
||||
target_spec,
|
||||
workers=workers,
|
||||
range_backend=range_backend,
|
||||
token=token,
|
||||
),
|
||||
download=lambda path, target_spec: download_published_sidecar(
|
||||
path,
|
||||
target_spec,
|
||||
cache_root=cache_root,
|
||||
token=token,
|
||||
),
|
||||
lock_timeout_s=lock_timeout_s,
|
||||
)
|
||||
@@ -28,7 +28,7 @@ from dataclasses import asdict, dataclass, field
|
||||
from fractions import Fraction
|
||||
from pathlib import Path
|
||||
from threading import Lock
|
||||
from typing import Any, ClassVar
|
||||
from typing import Any, BinaryIO, ClassVar
|
||||
|
||||
import av
|
||||
import fsspec
|
||||
@@ -105,7 +105,7 @@ def decode_video_frames(
|
||||
|
||||
|
||||
def decode_video_frames_pyav(
|
||||
video_path: Path | str,
|
||||
video_path: Path | str | BinaryIO,
|
||||
timestamps: list[float],
|
||||
tolerance_s: float,
|
||||
log_loaded_timestamps: bool = False,
|
||||
@@ -124,7 +124,7 @@ def decode_video_frames_pyav(
|
||||
video can be adjusted at encoding time to trade off decoding speed against file size.
|
||||
|
||||
Args:
|
||||
video_path: Path to the video file.
|
||||
video_path: Path to the video file or a seekable binary file object.
|
||||
timestamps: List of timestamps (in seconds) to extract frames for.
|
||||
tolerance_s: Allowed deviation in seconds between a queried timestamp and the closest
|
||||
decoded frame.
|
||||
@@ -137,7 +137,8 @@ def decode_video_frames_pyav(
|
||||
torch.Tensor of shape (len(timestamps), C, H, W).
|
||||
"""
|
||||
# TODO(rcadene): also load audio stream at the same time
|
||||
video_path = str(video_path)
|
||||
video_source = str(video_path) if isinstance(video_path, (Path, str)) else video_path
|
||||
video_label = str(video_path) if isinstance(video_path, (Path, str)) else "<in-memory video>"
|
||||
|
||||
# set the first and last requested timestamps
|
||||
# Note: previous timestamps are usually loaded, since we need to access the previous key frame
|
||||
@@ -151,7 +152,7 @@ def decode_video_frames_pyav(
|
||||
# av.time_base units (microseconds). `backward=True` lands us on the nearest keyframe at or
|
||||
# before `first_ts`, so we can then decode forward until we cover `last_ts`. See:
|
||||
# https://pyav.basswood-io.com/docs/stable/api/container.html#av.container.InputContainer.seek
|
||||
with av.open(video_path) as container:
|
||||
with av.open(video_source) as container:
|
||||
stream = container.streams.video[0]
|
||||
# Seek to the nearest keyframe at or before `first_ts` with a 1 frame margin
|
||||
container.seek(
|
||||
@@ -180,7 +181,7 @@ def decode_video_frames_pyav(
|
||||
|
||||
if not loaded_frames:
|
||||
raise FrameTimestampError(
|
||||
f"No frames could be decoded from {video_path} in the timestamp range [{first_ts}, {last_ts}]."
|
||||
f"No frames could be decoded from {video_label} in the timestamp range [{first_ts}, {last_ts}]."
|
||||
)
|
||||
|
||||
query_ts = torch.tensor(timestamps)
|
||||
@@ -199,7 +200,7 @@ def decode_video_frames_pyav(
|
||||
" To be safe, we advise to ignore this item during training."
|
||||
f"\nqueried timestamps: {query_ts}"
|
||||
f"\nloaded timestamps: {loaded_ts_t}"
|
||||
f"\nvideo: {video_path}"
|
||||
f"\nvideo: {video_label}"
|
||||
f"\nbackend: pyav"
|
||||
)
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
|
||||
|
||||
if action.shape[0] != 7:
|
||||
# Unknown layout — fall back to zero-pad so the sim doesn't crash.
|
||||
padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
|
||||
padded = np.zeros(ctrl_dim, dtype=np.float64)
|
||||
padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
|
||||
return padded
|
||||
|
||||
|
||||
@@ -122,9 +122,6 @@ MODEL_ENCODING_TABLE = {
|
||||
"xm430-w350": X_SERIES_ENCODINGS_TABLE,
|
||||
"xm540-w270": X_SERIES_ENCODINGS_TABLE,
|
||||
"xc430-w150": X_SERIES_ENCODINGS_TABLE,
|
||||
"xh540-w150": X_SERIES_ENCODINGS_TABLE,
|
||||
"xc330-t288": X_SERIES_ENCODINGS_TABLE,
|
||||
"xc330-t181": X_SERIES_ENCODINGS_TABLE,
|
||||
}
|
||||
|
||||
# {model: model_resolution}
|
||||
@@ -137,9 +134,6 @@ MODEL_RESOLUTION = {
|
||||
"xm430-w350": 4096,
|
||||
"xm540-w270": 4096,
|
||||
"xc430-w150": 4096,
|
||||
"xh540-w150": 4096,
|
||||
"xc330-t288": 4096,
|
||||
"xc330-t181": 4096,
|
||||
}
|
||||
|
||||
# {model: model_number}
|
||||
@@ -151,9 +145,6 @@ MODEL_NUMBER_TABLE = {
|
||||
"xm430-w350": 1020,
|
||||
"xm540-w270": 1120,
|
||||
"xc430-w150": 1070,
|
||||
"xh540-w150": 1110,
|
||||
"xc330-t288": 1220,
|
||||
"xc330-t181": 1210,
|
||||
}
|
||||
|
||||
# {model: available_operating_modes}
|
||||
@@ -165,9 +156,6 @@ MODEL_OPERATING_MODES = {
|
||||
"xm430-w350": [0, 1, 3, 4, 5, 16],
|
||||
"xm540-w270": [0, 1, 3, 4, 5, 16],
|
||||
"xc430-w150": [1, 3, 4, 16],
|
||||
"xh540-w150": [0, 1, 3, 4, 5, 16],
|
||||
"xc330-t288": [0, 1, 3, 4, 5, 16],
|
||||
"xc330-t181": [0, 1, 3, 4, 5, 16],
|
||||
}
|
||||
|
||||
MODEL_CONTROL_TABLE = {
|
||||
@@ -178,9 +166,6 @@ MODEL_CONTROL_TABLE = {
|
||||
"xm430-w350": X_SERIES_CONTROL_TABLE,
|
||||
"xm540-w270": X_SERIES_CONTROL_TABLE,
|
||||
"xc430-w150": X_SERIES_CONTROL_TABLE,
|
||||
"xh540-w150": X_SERIES_CONTROL_TABLE,
|
||||
"xc330-t288": X_SERIES_CONTROL_TABLE,
|
||||
"xc330-t181": X_SERIES_CONTROL_TABLE,
|
||||
}
|
||||
|
||||
MODEL_BAUDRATE_TABLE = {
|
||||
@@ -191,9 +176,6 @@ MODEL_BAUDRATE_TABLE = {
|
||||
"xm430-w350": X_SERIES_BAUDRATE_TABLE,
|
||||
"xm540-w270": X_SERIES_BAUDRATE_TABLE,
|
||||
"xc430-w150": X_SERIES_BAUDRATE_TABLE,
|
||||
"xh540-w150": X_SERIES_BAUDRATE_TABLE,
|
||||
"xc330-t288": X_SERIES_BAUDRATE_TABLE,
|
||||
"xc330-t181": X_SERIES_BAUDRATE_TABLE,
|
||||
}
|
||||
|
||||
AVAILABLE_BAUDRATES = [
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,)
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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*"],
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -165,7 +165,6 @@ from lerobot.robots import ( # noqa: F401
|
||||
earthrover_mini_plus,
|
||||
hope_jr,
|
||||
koch_follower,
|
||||
lekiwi,
|
||||
omx_follower,
|
||||
openarm_follower,
|
||||
reachy2,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
|
||||
@@ -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."""
|
||||
@@ -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")
|
||||
@@ -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"
|
||||
)
|
||||
@@ -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
|
||||
@@ -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
|
||||
)
|
||||
@@ -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()),
|
||||
)
|
||||
@@ -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}")
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
[
|
||||
|
||||
@@ -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})
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user