mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-06 09:37:06 +00:00
Compare commits
10 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Martino Russi
|
f7afeb9cb2 | ||
|
Martino Russi
|
3b6de2fdf8 | ||
|
Martino Russi
|
744f3667c0 | ||
|
Martino Russi
|
fdde436776 | ||
|
Martino Russi
|
5c683c65c6 | ||
|
Martino Russi
|
dfbc25c58f | ||
|
Martino Russi
|
804c76bcc2 | ||
|
Martino Russi
|
e6afa69be9 | ||
|
Martino Russi
|
31d1439e29 | ||
|
Martino Russi
|
1c118c6359 |
+1
-1
@@ -138,7 +138,7 @@ lerobot-replay --robot.type=so101_follower --robot.port=<FOLLOWER_PORT> --robot.
|
||||
--dataset.repo_id=${HF_USER}/my_task --dataset.episode=0
|
||||
```
|
||||
|
||||
**4.9 Train** (default: ACT — fastest, lowest memory). Apple silicon: `--policy.device=mps`. No local GPU? Add `--job.target=<flavor>` (e.g. `a10g-small`, list them with `hf jobs hardware`) to run on Hugging Face Jobs instead. See §6/§7 for policy and duration.
|
||||
**4.9 Train** (default: ACT — fastest, lowest memory). Apple silicon: `--policy.device=mps`. See §6/§7 for policy and duration.
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
|
||||
@@ -69,8 +69,6 @@
|
||||
title: VLA-JEPA
|
||||
- local: eo1
|
||||
title: EO-1
|
||||
- local: fastwam
|
||||
title: FastWAM
|
||||
- local: groot
|
||||
title: NVIDIA GR00T N1.5
|
||||
- local: xvla
|
||||
|
||||
@@ -157,14 +157,6 @@ finally:
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
### Working with depth
|
||||
|
||||
The Intel RealSense and Reachy 2 cameras can capture both color and depth in lockstep. Calling `read()` returns the **color** frame as `(H, W, 3)` `uint8`. Calling `read_depth()` returns the **depth map** as `(H, W, 1)` `uint16`, where each pixel value is the distance from the sensor expressed in **millimetres**. A pixel value of `0` typically means "no measurement available" (out-of-range, occluded, or low-confidence).
|
||||
|
||||
During recording, the control loop peeks the freshest buffered frames non-blockingly via `read_latest()` (color) and `read_latest_depth()` (depth), adding the depth map as a sibling feature (e.g. `front_depth` next to `front`).
|
||||
|
||||
For how depth streams are stored and encoded when recording a dataset, see the [Depth streams](./video_encoding_parameters#depth-streams) section of the video encoding guide.
|
||||
|
||||
## Use your phone's camera
|
||||
|
||||
<hfoptions id="use phone">
|
||||
|
||||
@@ -89,36 +89,6 @@ Control the data recording flow using keyboard shortcuts:
|
||||
- Press **Left Arrow (`←`)**: Delete current episode and retry.
|
||||
- Press **Escape (`ESC`)**: Stop, encode videos, and upload.
|
||||
|
||||
### Recording depth
|
||||
|
||||
Intel RealSense cameras (`type: intelrealsense`) record a depth stream when you set `use_depth: true`. Depth is quantized to 12-bit codes and stored as its own video.
|
||||
|
||||
```bash
|
||||
lerobot-record \
|
||||
... \
|
||||
--robot.cameras="{ head: {type: intelrealsense, serial_number_or_name: \"0123456789\", width: 640, height: 480, fps: 30, use_depth: true} }" \
|
||||
--dataset.repo_id=${HF_USER}/so101_depth_test \
|
||||
--dataset.single_task="put the red brick in a bowl" \
|
||||
--dataset.depth_encoder.depth_min=0.01 \
|
||||
--dataset.depth_encoder.depth_max=10.0 \
|
||||
--dataset.depth_encoder.shift=0.0 \
|
||||
--dataset.depth_encoder.use_log=true
|
||||
```
|
||||
|
||||
### Video encoding parameters
|
||||
|
||||
RGB and depth streams are encoded independently via the `--dataset.rgb_encoder.*` and `--dataset.depth_encoder.*` keys.
|
||||
|
||||
```bash
|
||||
lerobot-record \
|
||||
... \
|
||||
--dataset.rgb_encoder.vcodec=h264 \
|
||||
--dataset.rgb_encoder.pix_fmt=yuv420p \
|
||||
--dataset.rgb_encoder.crf=23 \
|
||||
--dataset.depth_encoder.vcodec=hevc \
|
||||
--dataset.depth_encoder.extra_options='{"x265-params": "lossless=1"}'
|
||||
```
|
||||
|
||||
### Training
|
||||
|
||||
Depending on your hardware training the policy might take a few hours. That's how you train simple `ACT` policy:
|
||||
@@ -150,14 +120,6 @@ lerobot-train \
|
||||
--steps=20000
|
||||
```
|
||||
|
||||
No local GPU? Add `--job.target=<flavor>` (e.g. `a10g-small`) to either command and `lerobot-train` runs it on [Hugging Face Jobs](https://huggingface.co/docs/hub/jobs) instead — it uploads a local-only dataset for you and pushes the trained model. List flavors with `hf jobs hardware`.
|
||||
|
||||
To resume, point `--config_path` at a checkpoint and add `--resume=true`. It accepts a local path or a Hub repo id (the latest checkpoint is fetched), and works locally or on a job by adding `--job.target=<flavor>`:
|
||||
|
||||
```bash
|
||||
lerobot-train --config_path=${HF_USER}/policy_test --resume=true --job.target=a10g-small
|
||||
```
|
||||
|
||||
### Inference
|
||||
|
||||
Inference means running the trained policy/model on a robot. For that we use `lerobot-rollout`. You will need to provide a path to your policy. It can be a local path or a path to Hugging Face for example "lerobot/folding_latest". Your cameras configuration needs to match what was used when collecting the dataset. Duration is in seconds if unspecified, it will run forever.
|
||||
|
||||
@@ -194,7 +194,7 @@ lerobot-record \
|
||||
--dataset.single_task="Navigate around obstacles" \
|
||||
--dataset.streaming_encoding=true \
|
||||
--dataset.encoder_threads=2 \
|
||||
# --dataset.rgb_encoder.vcodec=auto \
|
||||
# --dataset.camera_encoder.vcodec=auto \
|
||||
--display_data=true
|
||||
```
|
||||
|
||||
|
||||
@@ -1,167 +0,0 @@
|
||||
# FastWAM
|
||||
|
||||
FastWAM is a World Action Model policy for robot control. The LeRobot integration exposes FastWAM through the standard policy API so it can be configured with `policy.type=fastwam`, trained with `lerobot-train`, and loaded through the LeRobot pretrained policy interface.
|
||||
|
||||
## Model Overview
|
||||
|
||||
FastWAM keeps video modeling during training, but uses direct action prediction at inference time instead of iteratively generating future observations. This LeRobot policy wraps the FastWAM action model, adapts LeRobot batches to FastWAM training samples, and provides the standard processor pipeline for normalization and action postprocessing.
|
||||
|
||||
The implementation initializes the visual world-model components from `Wan-AI/Wan2.2-TI2V-5B` by default and predicts action chunks with shape `[batch, action_horizon, action_dim]`.
|
||||
|
||||
### What the LeRobot Integration Covers
|
||||
|
||||
- Standard `policy.type=fastwam` configuration through LeRobot
|
||||
- Image, state, action, and language-task batch adaptation
|
||||
- Action chunk inference through `select_action` and `predict_action_chunk`
|
||||
- Checkpoint save/load through the LeRobot policy APIs
|
||||
- Configurable LIBERO gripper action postprocessing
|
||||
|
||||
## Installation Requirements
|
||||
|
||||
Install LeRobot from source, then install FastWAM dependencies:
|
||||
|
||||
```bash
|
||||
pip install -e ".[fastwam]"
|
||||
```
|
||||
|
||||
This installs the FastWAM policy extra from `pyproject.toml`: `transformers`,
|
||||
`diffusers`, `ftfy`, and `regex`, plus LeRobot's base dependencies.
|
||||
|
||||
For LIBERO evaluation, install the benchmark dependencies too:
|
||||
|
||||
```bash
|
||||
pip install -e ".[fastwam,libero]"
|
||||
```
|
||||
|
||||
This installs both extras. In addition to the FastWAM dependencies above, the
|
||||
`libero` extra installs LeRobot dataset dependencies, `hf-libero` on Linux, and
|
||||
`scipy`.
|
||||
|
||||
FastWAM uses the Wan2.2 TI2V backbone. The default model id is:
|
||||
|
||||
```python
|
||||
policy.model_id=Wan-AI/Wan2.2-TI2V-5B
|
||||
```
|
||||
|
||||
## Data Requirements
|
||||
|
||||
FastWAM expects a LeRobot dataset with:
|
||||
|
||||
- one or more visual observations whose widths concatenate to `policy.image_size[1]`
|
||||
- `observation.state` when `policy.proprio_dim` is not `None`
|
||||
- `action`
|
||||
- a language task instruction through the dataset task field, or precomputed `context` and `context_mask` tensors
|
||||
|
||||
The default visual setup is one image feature named `observation.images.image` with shape `(3, 224, 448)`. If the dataset uses two cameras, configure `policy.input_features` so their heights match `224` and their widths sum to `448`.
|
||||
|
||||
## Usage
|
||||
|
||||
Create a new FastWAM policy with:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=your-org/your-dataset \
|
||||
--policy.type=fastwam \
|
||||
--policy.action_dim=7 \
|
||||
--policy.proprio_dim=8 \
|
||||
--policy.action_horizon=32 \
|
||||
--policy.n_action_steps=10 \
|
||||
--policy.image_size='[224,448]' \
|
||||
--output_dir=./outputs/fastwam_training \
|
||||
--job_name=fastwam_training \
|
||||
--steps=300000 \
|
||||
--batch_size=8 \
|
||||
--policy.device=cuda
|
||||
```
|
||||
|
||||
Evaluate an existing LeRobot-format checkpoint on LIBERO-10 with:
|
||||
|
||||
```bash
|
||||
lerobot-eval \
|
||||
--policy.path=ZibinDong/fastwam_libero_uncond_2cam224 \
|
||||
--policy.device=cuda \
|
||||
--policy.torch_dtype=float32 \
|
||||
--policy.n_action_steps=10 \
|
||||
--env.type=libero \
|
||||
--env.task=libero_10 \
|
||||
--env.observation_height=224 \
|
||||
--env.observation_width=224 \
|
||||
--eval.batch_size=1 \
|
||||
--eval.n_episodes=50 \
|
||||
--seed=0 \
|
||||
--env.episode_length=600
|
||||
```
|
||||
|
||||
For `libero_goal`, `libero_spatial`, and `libero_object`, use
|
||||
`--env.episode_length=300`.
|
||||
|
||||
For real-robot rollout, use the same checkpoint path:
|
||||
|
||||
```bash
|
||||
lerobot-rollout \
|
||||
--robot.type=so101_follower \
|
||||
--robot.port=/dev/ttyACM0 \
|
||||
--policy.path=your-org/fastwam-real-robot
|
||||
```
|
||||
|
||||
## Configuration Notes
|
||||
|
||||
### Image Features
|
||||
|
||||
`policy.image_size` is the size of the concatenated FastWAM image tensor as `(height, width)`. Each configured image feature must have shape `(3, height, camera_width)`, and all camera widths must sum to the configured width.
|
||||
|
||||
### Action Chunking
|
||||
|
||||
`policy.action_horizon` controls the number of future actions supervised during training and predicted during inference. `policy.n_action_steps` controls how many actions are consumed before the policy predicts a fresh chunk. `policy.n_action_steps` must be less than or equal to `policy.action_horizon`.
|
||||
|
||||
### Wan Components
|
||||
|
||||
FastWAM loads the Wan VAE, video DiT, text encoder, and tokenizer from the configured Wan model directory or Hugging Face Hub model id. LeRobot-format FastWAM checkpoints saved by `save_pretrained` also copy the local Wan component files needed by `from_pretrained`.
|
||||
|
||||
### Attention Backend
|
||||
|
||||
FastWAM's DiT uses PyTorch's `scaled_dot_product_attention` (SDPA) for all attention. It does **not** use FlashAttention: its Mixture-of-Transformers (MoT) routing needs arbitrary boolean `[query, key]` attention masks, which the FlashAttention varlen API cannot express. Installing the `flash-attn` package therefore has no effect on the FastWAM path. (Note that SDPA itself may still select PyTorch's own flash / memory-efficient / math kernel internally — this is unrelated to the `flash-attn` package.)
|
||||
|
||||
### LIBERO Action Toggle
|
||||
|
||||
FastWAM LIBERO checkpoints use `policy.toggle_action_dimensions=[-1]` by
|
||||
default to match the gripper action convention used by the original FastWAM
|
||||
evaluation pipeline:
|
||||
|
||||
```bash
|
||||
--policy.toggle_action_dimensions='[-1]'
|
||||
```
|
||||
|
||||
## Results
|
||||
|
||||
Evaluated on LIBERO with [`ZibinDong/fastwam_libero_uncond_2cam224`](https://huggingface.co/ZibinDong/fastwam_libero_uncond_2cam224):
|
||||
|
||||
| Suite | Success rate | n_episodes |
|
||||
| -------------- | -----------: | ---------: |
|
||||
| libero_spatial | 97.6% | 500 |
|
||||
| libero_object | 99.0% | 500 |
|
||||
| libero_goal | 95.0% | 500 |
|
||||
| libero_10 | 94.0% | 500 |
|
||||
| **average** | **96.4%** | 2000 |
|
||||
|
||||
Reproduce: `lerobot-eval --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 --policy.device=cuda --policy.torch_dtype=float32 --policy.n_action_steps=10 --env.type=libero --env.task=libero_spatial --env.observation_height=256 --env.observation_width=256 --eval.batch_size=1 --eval.n_episodes=50 --seed=0 --env.episode_length=300` (1x H20 140 GB).
|
||||
|
||||
## References
|
||||
|
||||
- [Fast-WAM paper](https://arxiv.org/abs/2603.16666)
|
||||
- [Fast-WAM project page](https://yuantianyuan01.github.io/FastWAM/)
|
||||
- [Fast-WAM code](https://github.com/yuantianyuan01/FastWAM)
|
||||
- [Released upstream checkpoints](https://huggingface.co/yuanty/fastwam)
|
||||
- [Wan2.2 TI2V 5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B)
|
||||
|
||||
## Citation
|
||||
|
||||
```bibtex
|
||||
@article{yuan2026fastwam,
|
||||
title = {Fast-WAM: Do World Action Models Need Test-time Future Imagination?},
|
||||
author = {Tianyuan Yuan and Zibin Dong and Yicheng Liu and Hang Zhao},
|
||||
journal = {arXiv preprint arXiv:2603.16666},
|
||||
year = {2026},
|
||||
url = {https://arxiv.org/abs/2603.16666}
|
||||
}
|
||||
```
|
||||
@@ -124,7 +124,7 @@ lerobot-rollout\
|
||||
--dataset.single_task="Grab and handover the red cube to the other arm" \
|
||||
--dataset.streaming_encoding=true \
|
||||
--dataset.encoder_threads=2 \
|
||||
# --dataset.rgb_encoder.vcodec=auto \
|
||||
# --dataset.camera_encoder.vcodec=auto \
|
||||
--policy.path=<user>/groot-bimanual \ # your trained model
|
||||
--duration=600
|
||||
```
|
||||
|
||||
@@ -96,4 +96,3 @@ Notes:
|
||||
- The leading `nvidia-smi` is a quick sanity check that CUDA is visible inside the container — useful to fail fast if the flavor or driver mismatched.
|
||||
- The default Job timeout is 30 minutes; pass `--timeout 4h` (or longer) for real training.
|
||||
- `--flavor` maps onto the table above: `t4-small`/`t4-medium` (T4, ACT only), `l4x1`/`l4x4` (L4 24 GB), `a10g-small/large/largex2/largex4` (A10G 24 GB scaled out), `a100-large` (A100). For the current full catalogue + pricing see [https://huggingface.co/docs/hub/jobs](https://huggingface.co/docs/hub/jobs).
|
||||
- Prefer not to write the `hf jobs run` wrapper yourself? `lerobot-train` can submit the job for you: just add `--job.target=<flavor>` to a normal training command and it handles dataset upload, log streaming, and the final model push. See the [imitation-learning training guide](./il_robots).
|
||||
|
||||
@@ -232,7 +232,7 @@ lerobot-record \
|
||||
--dataset.private=true \
|
||||
--dataset.streaming_encoding=true \
|
||||
--dataset.encoder_threads=2 \
|
||||
# --dataset.rgb_encoder.vcodec=auto \
|
||||
# --dataset.camera_encoder.vcodec=auto \
|
||||
--display_data=true
|
||||
```
|
||||
|
||||
@@ -278,6 +278,6 @@ lerobot-record \
|
||||
--dataset.num_episodes=10 \
|
||||
--dataset.streaming_encoding=true \
|
||||
--dataset.encoder_threads=2 \
|
||||
# --dataset.rgb_encoder.vcodec=auto \
|
||||
# --dataset.camera_encoder.vcodec=auto \
|
||||
--policy.path=outputs/train/hopejr_hand/checkpoints/last/pretrained_model
|
||||
```
|
||||
|
||||
@@ -207,7 +207,7 @@ lerobot-record \
|
||||
--dataset.num_episodes=5 \
|
||||
--dataset.single_task="Grab the black cube" \
|
||||
--dataset.streaming_encoding=true \
|
||||
# --dataset.rgb_encoder.vcodec=auto \
|
||||
# --dataset.camera_encoder.vcodec=auto \
|
||||
--dataset.encoder_threads=2
|
||||
```
|
||||
</hfoption>
|
||||
@@ -514,12 +514,6 @@ lerobot-train \
|
||||
--resume=true
|
||||
```
|
||||
|
||||
`--config_path` also accepts a **Hub repo id**: if a run pushed its checkpoints to the Hub (with `--save_checkpoint_to_hub=true`), you can resume straight from the repo — its latest checkpoint is downloaded and training continues, restoring the optimizer, scheduler, step counter and data order:
|
||||
|
||||
```bash
|
||||
lerobot-train --config_path=${HF_USER}/my_policy --resume=true
|
||||
```
|
||||
|
||||
If you do not want to push your model to the hub after training use `--policy.push_to_hub=false`.
|
||||
|
||||
Additionally you can provide extra `tags` or specify a `license` for your model or make the model repo `private` by adding this: `--policy.private=true --policy.tags=\[ppo,rl\] --policy.license=mit`
|
||||
@@ -532,9 +526,7 @@ If your local computer doesn't have a powerful GPU you could utilize Google Cola
|
||||
|
||||
Hugging Face jobs let's you easily select hardware and run the training in the cloud. So if you don't have a powerful GPU or you need more VRAM or just want to train a model much faster use HF Jobs! It's pay as you go and you simply pay for each second of use, you can see the pricing and additional information [here](https://huggingface.co/docs/hub/jobs).
|
||||
|
||||
> **Tip:** if you just want to launch a standard training run, you can skip building the command below and use the integrated **Train on HF Jobs via `--job.target`** flow described further down — `lerobot-train` then submits the job, uploads a local-only dataset for you, and streams the logs.
|
||||
|
||||
To run the training manually use this command:
|
||||
To run the training use this command:
|
||||
|
||||
<hfoptions id="train_with_hf_jobs">
|
||||
<hfoption id="Command">
|
||||
@@ -607,51 +599,6 @@ Once the training is started you can go to [Jobs](https://huggingface.co/setting
|
||||
|
||||
After training the model will be pushed to hub and you can use it as any other model with LeRobot.
|
||||
|
||||
#### Train on HF Jobs via `--job.target` (integrated CLI)
|
||||
|
||||
`lerobot-train` runs locally by default. To run on a HuggingFace GPU without constructing the Docker command yourself, pass `--job.target` with a hardware flavor name:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=${HF_USER}/so101_test \
|
||||
--policy.type=act \
|
||||
--policy.repo_id=${HF_USER}/my_policy \
|
||||
--job.target=a10g-small
|
||||
```
|
||||
|
||||
List available flavors and pricing with `hf jobs hardware`. The run streams its logs to your terminal; press Ctrl-C to detach (the job keeps running in the cloud). Re-attach or cancel with:
|
||||
|
||||
```bash
|
||||
hf jobs logs <job-id>
|
||||
hf jobs cancel <job-id>
|
||||
```
|
||||
|
||||
If your dataset exists only locally (not yet on the Hub), it is automatically pushed to a **private** Hub repo so the job can download it by `repo_id` (nothing is made public). The trained model is pushed to the model repo at the end of the run. To also push every intermediate checkpoint to the Hub as it is saved (so you can monitor progress mid-run), add `--save_checkpoint_to_hub=true` — this requires a runtime image that includes this feature.
|
||||
|
||||
Every job (and any dataset pushed by the run) is tagged `lerobot` so it's easy to find on the Hub. Add your own with `--job.tags '["my-tag"]'`.
|
||||
|
||||
By default the job is capped at `2d` (48h) of wall-clock. Override it with an HF Jobs duration string, e.g. `--job.timeout=4h` to fail faster or `--job.timeout=7d` for a longer run.
|
||||
|
||||
> **Note:** the model repo is created up front (it holds the staged training config the job runs from). If a run fails before the model is pushed, that repo is left on the Hub so you can inspect it — it is not deleted automatically, so repeated failures can leave empty repos behind. Remove one with `hf repo delete <repo-id>`.
|
||||
|
||||
**Prerequisites:** run `hf auth login` before submitting. For Weights & Biases integration, run `wandb login` or set `WANDB_API_KEY` on your machine — the key is forwarded to the job automatically.
|
||||
|
||||
**Resuming on a job.** Adding `--job.target` to a resume command runs the resume in the cloud — the same command works locally or remotely. The checkpoint repo is the source of truth, and new checkpoints continue the lineage in the same repo:
|
||||
|
||||
```bash
|
||||
# resume a Hub run on a job (its checkpoints are already on the Hub)
|
||||
lerobot-train --config_path=${HF_USER}/my_policy --resume=true --job.target=a10g-small
|
||||
|
||||
# resume a LOCAL run on a job — the checkpoint is uploaded to a private Hub repo first,
|
||||
# then the job resumes from it (a local-only dataset is uploaded the same way)
|
||||
lerobot-train \
|
||||
--config_path=outputs/train/act_so101_test/checkpoints/last/pretrained_model/train_config.json \
|
||||
--resume=true \
|
||||
--job.target=a10g-small
|
||||
```
|
||||
|
||||
Job settings come from the current command, so override `--job.target`, `--job.timeout`, etc. as needed; for the resumed run to itself be resumable later, keep `--save_checkpoint_to_hub=true`.
|
||||
|
||||
#### Upload policy checkpoints
|
||||
|
||||
Once training is done, upload the latest checkpoint with:
|
||||
@@ -673,8 +620,6 @@ hf upload ${HF_USER}/act_so101_test${CKPT} \
|
||||
|
||||
Use `lerobot-rollout` to deploy a trained policy on your robot. You can choose different strategies depending on your needs:
|
||||
|
||||
The examples below load the model from `--policy.path`. To pin a specific pushed version — useful once `--save_checkpoint_to_hub=true` has committed several checkpoints — add `--policy.pretrained_revision` with a commit hash, branch, or tag. Each pushed checkpoint is tagged with its step (e.g. `--policy.pretrained_revision=010000`), so you can recover a checkpoint by step without looking up its commit sha.
|
||||
|
||||
<hfoptions id="eval">
|
||||
<hfoption id="Base mode (no recording)">
|
||||
```bash
|
||||
|
||||
@@ -44,7 +44,7 @@ lerobot-record \
|
||||
--dataset.num_episodes=5 \
|
||||
--dataset.single_task="Grab the black cube" \
|
||||
--dataset.streaming_encoding=true \
|
||||
# --dataset.rgb_encoder.vcodec=auto \
|
||||
# --dataset.camera_encoder.vcodec=auto \
|
||||
--dataset.encoder_threads=2
|
||||
```
|
||||
|
||||
|
||||
@@ -386,68 +386,6 @@ These results demonstrate MolmoAct2's strong performance across diverse robotic
|
||||
manipulation tasks. To reproduce them, follow the instructions in the LIBERO
|
||||
evaluation section.
|
||||
|
||||
## Hardware Deployment (lerobot-rollout)
|
||||
|
||||
LeRobot-format checkpoints are available on the Hub for direct use with
|
||||
`lerobot-rollout`. Each checkpoint uses specific camera names that must
|
||||
match your robot's camera configuration.
|
||||
|
||||
### Camera naming convention
|
||||
|
||||
Each checkpoint expects specific `observation.images.*` keys.
|
||||
If your robot cameras have different names, use `--rename_map` to map them:
|
||||
|
||||
| Checkpoint | Camera keys | Description |
|
||||
| ----------------------------- | ---------------------- | ------------------------ |
|
||||
| MolmoAct2-LIBERO-LeRobot | `image`, `wrist_image` | LIBERO sim cameras |
|
||||
| MolmoAct2-BimanualYAM-LeRobot | `top`, `left`, `right` | YAM 3-camera setup |
|
||||
| MolmoAct2-DROID-LeRobot | `cam0`, `cam1` | External + wrist |
|
||||
| MolmoAct2-SO100_101-LeRobot | `cam0`, `cam1` | Primary + secondary view |
|
||||
|
||||
Example with an SO-100 robot using top and side cameras:
|
||||
|
||||
```bash
|
||||
lerobot-rollout \
|
||||
--policy.path=lerobot/MolmoAct2-SO100_101-LeRobot \
|
||||
--rename_map='{"observation.images.top": "observation.images.cam0", "observation.images.side": "observation.images.cam1"}' \
|
||||
--robot.type=so100_follower \
|
||||
--robot.port=/dev/ttyACM0 \
|
||||
--robot.cameras='{
|
||||
top: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30},
|
||||
side: {type: opencv, index_or_path: 2, width: 640, height: 480, fps: 30}
|
||||
}' \
|
||||
--task="pick up the red cube" --duration=30
|
||||
```
|
||||
|
||||
To use a wrist camera instead, just change the rename mapping:
|
||||
|
||||
```bash
|
||||
--rename_map='{"observation.images.top": "observation.images.cam0", "observation.images.wrist": "observation.images.cam1"}'
|
||||
```
|
||||
|
||||
### Joint frame transform (SO-100/101 zero-shot)
|
||||
|
||||
<Tip warning={true}>
|
||||
The MolmoAct2-SO100_101 checkpoint was trained on data that uses a different
|
||||
joint calibration convention than LeRobot >= 0.5.0. Without a frame
|
||||
correction, the arm may move in the wrong direction.
|
||||
|
||||
This affects both **zero-shot deployment** and **fine-tuning** from the
|
||||
original checkpoint. The pretrained weights expect the old convention, so
|
||||
all joint data (observations and actions) must be transformed to match.
|
||||
|
||||
The converted LeRobot checkpoint (`lerobot/MolmoAct2-SO100_101-LeRobot`)
|
||||
already includes this correction in its processor pipeline. If you convert
|
||||
or fine-tune the checkpoint yourself, set the following in the policy config (`configuration_molmoact2.py`):
|
||||
|
||||
- `joint_signs`: `[1, -1, 1, 1, 1, 1]` (flips shoulder_lift direction)
|
||||
- `joint_offsets`: `[0, 90, 90, 0, 0, 0]` (shifts shoulder_lift and elbow_flex by 90°)
|
||||
|
||||
See the [backward compatibility guide](./backwardcomp) for details on the
|
||||
calibration change.
|
||||
|
||||
</Tip>
|
||||
|
||||
## Differences From the Original Implementation
|
||||
|
||||
This LeRobot port is intended to match MolmoAct2 behavior while using LeRobot's
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
## Research Paper
|
||||
|
||||
Paper: https://arxiv.org/abs/2603.16666
|
||||
|
||||
## Repository
|
||||
|
||||
Code: https://github.com/yuantianyuan01/FastWAM
|
||||
|
||||
Project page: https://yuantianyuan01.github.io/FastWAM/
|
||||
|
||||
## Citation
|
||||
|
||||
```bibtex
|
||||
@article{yuan2026fastwam,
|
||||
title = {Fast-WAM: Do World Action Models Need Test-time Future Imagination?},
|
||||
author = {Tianyuan Yuan and Zibin Dong and Yicheng Liu and Hang Zhao},
|
||||
journal = {arXiv preprint arXiv:2603.16666},
|
||||
year = {2026},
|
||||
url = {https://arxiv.org/abs/2603.16666}
|
||||
}
|
||||
```
|
||||
|
||||
## Additional Resources
|
||||
|
||||
Base video model: https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B
|
||||
|
||||
Released upstream checkpoints: https://huggingface.co/yuanty/fastwam
|
||||
|
||||
## Results
|
||||
|
||||
Evaluated on LIBERO with [`ZibinDong/fastwam_libero_uncond_2cam224`](https://huggingface.co/ZibinDong/fastwam_libero_uncond_2cam224):
|
||||
|
||||
| Suite | Success rate | n_episodes |
|
||||
| -------------- | -----------: | ---------: |
|
||||
| libero_spatial | 97.6% | 500 |
|
||||
| libero_object | 99.0% | 500 |
|
||||
| libero_goal | 95.0% | 500 |
|
||||
| libero_10 | 94.0% | 500 |
|
||||
| **average** | **96.4%** | 2000 |
|
||||
|
||||
Reproduce: `lerobot-eval --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 --policy.device=cuda --policy.torch_dtype=float32 --policy.n_action_steps=10 --env.type=libero --env.task=libero_spatial --env.observation_height=256 --env.observation_width=256 --eval.batch_size=1 --eval.n_episodes=50 --seed=0 --env.episode_length=300`.
|
||||
|
||||
For LIBERO-10, use `--env.task=libero_10 --env.episode_length=600`:
|
||||
|
||||
```bash
|
||||
lerobot-eval \
|
||||
--policy.path=ZibinDong/fastwam_libero_uncond_2cam224 \
|
||||
--policy.device=cuda \
|
||||
--policy.torch_dtype=float32 \
|
||||
--policy.n_action_steps=10 \
|
||||
--env.type=libero \
|
||||
--env.task=libero_10 --env.observation_height=256 --env.observation_width=256 \
|
||||
--eval.batch_size=1 \
|
||||
--eval.n_episodes=50 \
|
||||
--seed=0 --env.episode_length=600
|
||||
```
|
||||
@@ -161,7 +161,7 @@ lerobot-record \
|
||||
--dataset.private=true \
|
||||
--dataset.streaming_encoding=true \
|
||||
--dataset.encoder_threads=2 \
|
||||
# --dataset.rgb_encoder.vcodec=auto \
|
||||
# --dataset.camera_encoder.vcodec=auto \
|
||||
--display_data=true
|
||||
```
|
||||
|
||||
@@ -203,7 +203,7 @@ lerobot-record \
|
||||
--dataset.private=true \
|
||||
--dataset.streaming_encoding=true \
|
||||
--dataset.encoder_threads=2 \
|
||||
# --dataset.rgb_encoder.vcodec=auto \
|
||||
# --dataset.camera_encoder.vcodec=auto \
|
||||
--display_data=true
|
||||
```
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ This makes `save_episode()` near-instant (the video is already encoded by the ti
|
||||
| Parameter | CLI Flag | Type | Default | Description |
|
||||
| ----------------------- | --------------------------------- | ------------- | ------------- | ----------------------------------------------------------------- |
|
||||
| `streaming_encoding` | `--dataset.streaming_encoding` | `bool` | `True` | Enable real-time encoding during capture |
|
||||
| `vcodec` | `--dataset.rgb_encoder.vcodec` | `str` | `"libsvtav1"` | Video codec. `"auto"` detects best HW encoder |
|
||||
| `vcodec` | `--dataset.camera_encoder.vcodec` | `str` | `"libsvtav1"` | Video codec. `"auto"` detects best HW encoder |
|
||||
| `encoder_threads` | `--dataset.encoder_threads` | `int \| None` | `None` (auto) | Threads per encoder instance. `None` will leave the vcoded decide |
|
||||
| `encoder_queue_maxsize` | `--dataset.encoder_queue_maxsize` | `int` | `30` | Max buffered frames per camera (~1s at 30fps). Consumes RAM |
|
||||
|
||||
@@ -82,15 +82,15 @@ Use HW encoding when:
|
||||
|
||||
### Available HW Encoders
|
||||
|
||||
| Encoder | Platform | Hardware | CLI Value |
|
||||
| ------------------- | ------------- | ------------------------------------------------------------------------------------------------ | ------------------------------------------------ |
|
||||
| `h264_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.rgb_encoder.vcodec=h264_videotoolbox` |
|
||||
| `hevc_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.rgb_encoder.vcodec=hevc_videotoolbox` |
|
||||
| `h264_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.rgb_encoder.vcodec=h264_nvenc` |
|
||||
| `hevc_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.rgb_encoder.vcodec=hevc_nvenc` |
|
||||
| `h264_vaapi` | Linux | Intel/AMD GPU | `--dataset.rgb_encoder.vcodec=h264_vaapi` |
|
||||
| `h264_qsv` | Linux/Windows | Intel Quick Sync | `--dataset.rgb_encoder.vcodec=h264_qsv` |
|
||||
| `auto` | Any | Probes the system for available HW encoders. Falls back to `libsvtav1` if no HW encoder is found | `--dataset.rgb_encoder.vcodec=auto` |
|
||||
| Encoder | Platform | Hardware | CLI Value |
|
||||
| ------------------- | ------------- | ------------------------------------------------------------------------------------------------ | --------------------------------------------------- |
|
||||
| `h264_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.camera_encoder.vcodec=h264_videotoolbox` |
|
||||
| `hevc_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.camera_encoder.vcodec=hevc_videotoolbox` |
|
||||
| `h264_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.camera_encoder.vcodec=h264_nvenc` |
|
||||
| `hevc_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.camera_encoder.vcodec=hevc_nvenc` |
|
||||
| `h264_vaapi` | Linux | Intel/AMD GPU | `--dataset.camera_encoder.vcodec=h264_vaapi` |
|
||||
| `h264_qsv` | Linux/Windows | Intel Quick Sync | `--dataset.camera_encoder.vcodec=h264_qsv` |
|
||||
| `auto` | Any | Probes the system for available HW encoders. Falls back to `libsvtav1` if no HW encoder is found | `--dataset.camera_encoder.vcodec=auto` |
|
||||
|
||||
> [!NOTE]
|
||||
> In order to use the HW accelerated encoders you might need to upgrade your GPU drivers.
|
||||
@@ -100,15 +100,15 @@ Use HW encoding when:
|
||||
|
||||
## 5. Troubleshooting
|
||||
|
||||
| Symptom | Likely Cause | Fix |
|
||||
| ------------------------------------------------------------------ | -------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
|
||||
| System freezes or choppy robot movement or Rerun visualization lag | CPU starved (100% load usage) | Close other apps, reduce encoding throughput, lower `encoder_threads`, use `h264`, use `display_data=False`. If the CPU continues to be at 100% then it might be insufficient for your setup, consider `--dataset.streaming_encoding=false` or HW encoding (`--dataset.rgb_encoder.vcodec=auto`) |
|
||||
| "Encoder queue full" warnings or dropped frames in dataset | Encoder can't keep up (Queue overflow) | If CPU is not at 100%: Increase `encoder_threads`, increase `encoder_queue_maxsize` or use HW encoding (`--dataset.rgb_encoder.vcodec=auto`). |
|
||||
| High RAM usage | Queue filling faster than encoding | `encoder_threads` too low or CPU insufficient. Reduce `encoder_queue_maxsize` or use HW encoding |
|
||||
| Large video files | Using HW encoder or H.264 | Expected trade-off. Switch to `libsvtav1` if CPU allows |
|
||||
| `save_episode()` still slow | `streaming_encoding` is `False` | Set `--dataset.streaming_encoding=true` |
|
||||
| Encoder thread crash | Codec not available or invalid settings | Check `vcodec` is installed, try `--dataset.rgb_encoder.vcodec=auto` |
|
||||
| Recorded dataset is missing frames | CPU/GPU starvation or occasional load spikes | If ~5% of frames are missing, your system is likely overloaded — follow the recommendations above. If fewer frames are missing (~2%), they are probably due to occasional transient load spikes (often at startup) and can be considered expected. |
|
||||
| Symptom | Likely Cause | Fix |
|
||||
| ------------------------------------------------------------------ | -------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| System freezes or choppy robot movement or Rerun visualization lag | CPU starved (100% load usage) | Close other apps, reduce encoding throughput, lower `encoder_threads`, use `h264`, use `display_data=False`. If the CPU continues to be at 100% then it might be insufficient for your setup, consider `--dataset.streaming_encoding=false` or HW encoding (`--dataset.camera_encoder.vcodec=auto`) |
|
||||
| "Encoder queue full" warnings or dropped frames in dataset | Encoder can't keep up (Queue overflow) | If CPU is not at 100%: Increase `encoder_threads`, increase `encoder_queue_maxsize` or use HW encoding (`--dataset.camera_encoder.vcodec=auto`). |
|
||||
| High RAM usage | Queue filling faster than encoding | `encoder_threads` too low or CPU insufficient. Reduce `encoder_queue_maxsize` or use HW encoding |
|
||||
| Large video files | Using HW encoder or H.264 | Expected trade-off. Switch to `libsvtav1` if CPU allows |
|
||||
| `save_episode()` still slow | `streaming_encoding` is `False` | Set `--dataset.streaming_encoding=true` |
|
||||
| Encoder thread crash | Codec not available or invalid settings | Check `vcodec` is installed, try `--dataset.camera_encoder.vcodec=auto` |
|
||||
| Recorded dataset is missing frames | CPU/GPU starvation or occasional load spikes | If ~5% of frames are missing, your system is likely overloaded — follow the recommendations above. If fewer frames are missing (~2%), they are probably due to occasional transient load spikes (often at startup) and can be considered expected. |
|
||||
|
||||
## 6. Recommended Configurations
|
||||
|
||||
@@ -146,7 +146,7 @@ On very constrained systems, streaming encoding may compete too heavily with the
|
||||
# 2camsx 640x480x3 @30fps: Requires some tuning.
|
||||
|
||||
# Use H.264, disable streaming, consider batching encoding
|
||||
lerobot-record --dataset.rgb_encoder.vcodec=h264 --dataset.streaming_encoding=false ...
|
||||
lerobot-record --dataset.camera_encoder.vcodec=h264 --dataset.streaming_encoding=false ...
|
||||
```
|
||||
|
||||
## 7. Closing note
|
||||
|
||||
@@ -11,9 +11,8 @@ 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. **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.
|
||||
6. **Convert to Video** - Convert image-based datasets to video format for efficient storage
|
||||
7. **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`.
|
||||
@@ -118,19 +117,10 @@ lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht_image \
|
||||
--operation.type convert_image_to_video \
|
||||
--operation.output_dir outputs/pusht_video \
|
||||
--operation.rgb_encoder.vcodec libsvtav1 \
|
||||
--operation.rgb_encoder.pix_fmt yuv420p \
|
||||
--operation.rgb_encoder.g 2 \
|
||||
--operation.rgb_encoder.crf 30
|
||||
|
||||
# Convert a dataset that includes depth maps, customizing the depth encoder
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht_image \
|
||||
--operation.type convert_image_to_video \
|
||||
--operation.output_dir outputs/pusht_video \
|
||||
--operation.depth_encoder.depth_min 0.01 \
|
||||
--operation.depth_encoder.depth_max 10.0 \
|
||||
--operation.depth_encoder.use_log true
|
||||
--operation.camera_encoder.vcodec libsvtav1 \
|
||||
--operation.camera_encoder.pix_fmt yuv420p \
|
||||
--operation.camera_encoder.g 2 \
|
||||
--operation.camera_encoder.crf 30
|
||||
|
||||
# Convert only specific episodes
|
||||
lerobot-edit-dataset \
|
||||
@@ -157,42 +147,11 @@ lerobot-edit-dataset \
|
||||
**Parameters:**
|
||||
|
||||
- `output_dir`: Custom output directory (optional - by default uses `new_repo_id` or `{repo_id}_video`)
|
||||
- `rgb_encoder`: Video encoder settings applied to RGB cameras — all sub-fields accessible via `--operation.rgb_encoder.<field>`. See [Video Encoding Parameters](./video_encoding_parameters) for more details.
|
||||
- `depth_encoder`: Video encoder settings applied to depth-map cameras (e.g. from an Intel RealSense). In addition to the standard encoder fields it exposes the depth quantization knobs (`depth_min`, `depth_max`, `shift`, `use_log`), accessible via `--operation.depth_encoder.<field>`. These quantization settings are persisted to the dataset metadata so depth can be dequantized back to physical units on load. See the [Depth streams](./video_encoding_parameters#depth-streams) section for details.
|
||||
- `camera_encoder`: Video encoder settings — all sub-fields accessible via `--operation.camera_encoder.<field>. See [Video Encoding Parameters](./video_encoding_parameters) for more details.
|
||||
- `episode_indices`: List of specific episodes to convert (default: all episodes)
|
||||
- `num_workers`: Number of parallel workers for processing (default: 4)
|
||||
|
||||
**Note:** The resulting dataset will be a proper LeRobotDataset with all cameras encoded as videos in the `videos/` directory, with parquet files containing only metadata (no raw image data). Depth-map cameras are detected automatically and routed to the `depth_encoder`, while RGB cameras use the `rgb_encoder`. All episodes, stats, and tasks are preserved.
|
||||
|
||||
#### Re-encode Videos
|
||||
|
||||
Re-encode the videos of an existing video dataset with different encoder settings, without going back to raw frames. RGB videos use the `rgb_encoder` and depth videos use the `depth_encoder`. Provide only the encoder(s) you want to re-encode; the other stream type is left untouched.
|
||||
|
||||
```bash
|
||||
# Re-encode all RGB videos with new settings (saves to lerobot/pusht_reencoded by default)
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht \
|
||||
--operation.type reencode_videos \
|
||||
--operation.rgb_encoder.vcodec h264 \
|
||||
--operation.rgb_encoder.pix_fmt yuv420p \
|
||||
--operation.rgb_encoder.crf 23
|
||||
|
||||
# Re-encode both RGB and depth videos in a dataset with depth maps
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht_depth \
|
||||
--operation.type reencode_videos \
|
||||
--operation.rgb_encoder.vcodec h264 \
|
||||
--operation.depth_encoder.crf 50
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
|
||||
- `rgb_encoder`: Encoder settings applied to every RGB video. Omit to skip re-encoding RGB videos.
|
||||
- `depth_encoder`: Encoder settings applied to every depth video. Omit to skip re-encoding depth videos.
|
||||
- `num_workers`: Number of parallel workers for processing.
|
||||
|
||||
> [!NOTE]
|
||||
> When re-encoding depth videos, the existing depth quantization parameters (`depth_min`, `depth_max`, `shift`, `use_log`) and the `is_depth_map` flag are **preserved** — re-encoding only changes the codec/quality of the stored stream, not how depth is dequantized on load.
|
||||
**Note:** The resulting dataset will be a proper LeRobotDataset with all cameras encoded as videos in the `videos/` directory, with parquet files containing only metadata (no raw image data). All episodes, stats, and tasks are preserved.
|
||||
|
||||
### Show the information of datasets
|
||||
|
||||
|
||||
@@ -2,15 +2,15 @@
|
||||
|
||||
When video storage is enabled, LeRobot stores each camera stream as an **MP4** file instead of saving one image file per timestep. Video encoding compresses across time, which usually cuts dataset size and I/O compared to a pile of PNG, while keeping MP4 — a format every player and loader understands.
|
||||
|
||||
Encoding frames into an MP4 is a full FFmpeg pipeline: choice of encoder, pixel format, GOP/keyframes, quality vs. speed, and optional extra encoder flags. Most of these knobs are user-tunable through `rgb_encoder`, a nested `RGBEncoderConfig` (`lerobot.configs.video.RGBEncoderConfig`) passed through PyAV.
|
||||
Encoding frames into an MP4 is a full FFmpeg pipeline: choice of encoder, pixel format, GOP/keyframes, quality vs. speed, and optional extra encoder flags. Most of these knobs are user-tunable through `camera_encoder`, a nested `VideoEncoderConfig` (`lerobot.configs.video.VideoEncoderConfig`) passed through PyAV.
|
||||
|
||||
You can set these parameters from the CLI with `--dataset.rgb_encoder.<field>` (e.g. with `lerobot-record` or `lerobot-rollout`). The same block applies to every camera video stream in that run.
|
||||
You can set these parameters from the CLI with `--dataset.camera_encoder.<field>` (e.g. with `lerobot-record` or `lerobot-rollout`). The same block applies to every camera video stream in that run.
|
||||
|
||||
<Tip>
|
||||
Video storage must be on for `rgb_encoder` to have any effect —
|
||||
Video storage must be on for `camera_encoder` to have any effect —
|
||||
`use_videos=True` in Python APIs, or `--dataset.video=true` on the CLI (the
|
||||
recording default). With video off, inputs stay as images and `rgb_encoder` is
|
||||
ignored.
|
||||
recording default). With video off, inputs stay as images and `camera_encoder`
|
||||
is ignored.
|
||||
</Tip>
|
||||
|
||||
For details on **when** frames are written vs. encoded (streaming vs. post-episode), queues, and other top-level `--dataset.*` switches, see [Streaming Video Encoding](./streaming_video_encoding). For an encoding-parameter comparison and experiments, see the [video-benchmark Space](https://huggingface.co/spaces/lerobot/video-benchmark).
|
||||
@@ -33,9 +33,9 @@ lerobot-record \
|
||||
--dataset.single_task="Grab the cube" \
|
||||
--dataset.streaming_encoding=true \
|
||||
--dataset.encoder_threads=2 \
|
||||
--dataset.rgb_encoder.vcodec=h264 \
|
||||
--dataset.rgb_encoder.preset=fast \
|
||||
--dataset.rgb_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \
|
||||
--dataset.camera_encoder.vcodec=h264 \
|
||||
--dataset.camera_encoder.preset=fast \
|
||||
--dataset.camera_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \
|
||||
--display_data=true
|
||||
```
|
||||
|
||||
@@ -50,7 +50,7 @@ Only override these parameters if you have a specific reason to, and measure the
|
||||
|
||||
</Tip>
|
||||
|
||||
All flags below are prefixed with `--dataset.rgb_encoder.` on the CLI.
|
||||
All flags below are prefixed with `--dataset.camera_encoder.` on the CLI.
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| --------------- | ---------------- | ------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
@@ -65,80 +65,6 @@ All flags below are prefixed with `--dataset.rgb_encoder.` on the CLI.
|
||||
|
||||
---
|
||||
|
||||
## Depth streams
|
||||
|
||||
Depth maps (Intel RealSense, Reachy 2) are stored as their **own video streams** alongside the RGB streams. Raw depth (`uint16` millimetres or `float32` metres) can't survive an 8-bit codec, so LeRobot **quantizes** each map to a 12-bit code (`[0, 4095]`) — logarithmically by default, to match the `1/depth` error profile of depth sensors — then packs it into a high-bit-depth pixel format (`gray12le`) and encodes it with a 12-bit codec.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A["Raw depth (uint16 mm / float32 m)"] --> B["Clip to depth_min, depth_max"]
|
||||
B --> C["Quantize to 12-bit code 0–4095 (log or linear)"]
|
||||
C --> D["Pack into gray12le"]
|
||||
D --> E["Encode video (hevc Main 12)"]
|
||||
E --> F[("MP4 + metadata: depth_min/max, shift, use_log")]
|
||||
F -. "load time (depth_output_unit)" .-> G["Dequantize to mm or m"]
|
||||
|
||||
classDef input fill:#e3f2fd,stroke:#1565c0,color:#0d47a1;
|
||||
classDef encode fill:#ede7f6,stroke:#5e35b1,color:#311b92;
|
||||
classDef store fill:#fff8e1,stroke:#f9a825,color:#e65100;
|
||||
classDef load fill:#e8f5e9,stroke:#2e7d32,color:#1b5e20;
|
||||
|
||||
class A input;
|
||||
class B,C,D,E encode;
|
||||
class F store;
|
||||
class G load;
|
||||
```
|
||||
|
||||
Configure the depth pipeline through a parallel **`depth_encoder`** block (`DepthEncoderConfig`). It shares every `RGBEncoderConfig` field (`vcodec`, `pix_fmt`, `crf`, …) and adds four quantizer knobs, set via `--dataset.depth_encoder.<field>`:
|
||||
|
||||
```bash
|
||||
lerobot-record \
|
||||
... \
|
||||
--dataset.depth_encoder.vcodec=hevc \
|
||||
--dataset.depth_encoder.depth_min=0.05 \
|
||||
--dataset.depth_encoder.depth_max=5.0 \
|
||||
--dataset.depth_encoder.use_log=true
|
||||
```
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| --------------- | ------- | ------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `vcodec` | `str` | `"hevc"` | HEVC Main 12 (a 12-bit-capable codec, MP4-compatible). |
|
||||
| `extra_options` | `dict` | `{"x265-params": "lossless=1"}` | **Depth defaults to lossless** (exact round-trip); `crf` is ignored. Pass `extra_options={}` and set `crf` for a smaller lossy stream. |
|
||||
| `pix_fmt` | `str` | `"gray12le"` | Single-channel 12-bit pixel format used to carry the quantized codes. |
|
||||
| `depth_min` | `float` | `0.01` | Depth in metres mapped to quantum `0`. Values below are clipped on decode. |
|
||||
| `depth_max` | `float` | `10.0` | Depth in metres mapped to quantum `4095`. Values above are clipped on decode. |
|
||||
| `shift` | `float` | `3.5` | Pre-log offset (metres) used in logarithmic quantization for numerical stability near zero. Must satisfy `depth_min + shift > 0`. |
|
||||
| `use_log` | `bool` | `True` | If `true`, quantize in log-space (recommended for typical depth sensors). Set to `false` for uniform/linear quantization. |
|
||||
|
||||
> [!TIP]
|
||||
> `depth_min`, `depth_max`, and `shift` are always interpreted in **metres**, regardless of the input depth's unit. Inputs are auto-detected: integer arrays (e.g. `uint16` millimetres straight from a RealSense) are treated as millimetres, floating arrays as metres.
|
||||
> Pick `depth_min` / `depth_max` to bracket the actual working range of your sensor — quanta outside that range saturate, which can crush detail at the boundaries.
|
||||
|
||||
Depth features are flagged with `"is_depth_map": true` in `meta/info.json`, and their quantizer settings (`video.depth_min`, `video.depth_max`, `video.shift`, `video.use_log`) are persisted — which is what lets depth be **dequantized back to physical units** on load.
|
||||
|
||||
### Output unit at load time
|
||||
|
||||
`depth_encoder` is a **record-time** concern. The unit that depth maps are dequantized to on _load_ (e.g. during training) is set separately by the read-time flag `--dataset.depth_output_unit`:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=<my_username>/<my_dataset_name> \
|
||||
--dataset.depth_output_unit=m \
|
||||
--policy.type=act
|
||||
```
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| ------------------- | ----- | ------- | -------------------------------------------------------------------------------------------- |
|
||||
| `depth_output_unit` | `str` | `"mm"` | Physical unit depth maps are dequantized to on load: `"mm"` (millimetres) or `"m"` (metres). |
|
||||
|
||||
> [!TIP]
|
||||
> This is purely a decode-time presentation choice — it does **not** alter the stored video or its metadata, so the same dataset can be read as `mm` or `m` without re-encoding. It has no effect on datasets without depth cameras.
|
||||
|
||||
> [!IMPORTANT]
|
||||
> Depth statistics in `meta/stats.json` are always computed in **millimetres**, regardless of the raw frame dtype.
|
||||
|
||||
---
|
||||
|
||||
## Persistence in dataset metadata
|
||||
|
||||
After the first episode of a video stream is encoded, the encoder configuration is **persisted into the dataset metadata** (`meta/info.json`) under each video feature, alongside the values probed from the file itself. For a video feature `observation.images.<camera>`, the layout in `info.json` is:
|
||||
@@ -156,7 +82,7 @@ After the first episode of a video stream is encoded, the encoder configuration
|
||||
"video.pix_fmt": "yuv420p",
|
||||
"video.fps": 30,
|
||||
"video.channels": 3,
|
||||
"is_depth_map": false,
|
||||
"video.is_depth_map": false,
|
||||
"video.g": 2,
|
||||
"video.crf": 30,
|
||||
"video.preset": "fast",
|
||||
@@ -171,12 +97,12 @@ After the first episode of a video stream is encoded, the encoder configuration
|
||||
|
||||
Two sources contribute to the `info` block:
|
||||
|
||||
- **Stream-derived** (read back from the encoded MP4 with PyAV): `video.height`, `video.width`, `video.codec`, `video.pix_fmt`, `video.fps`, `video.channels`, `is_depth_map`, plus `audio.*` if an audio stream is present.
|
||||
- **Encoder-derived** (taken from `RGBEncoderConfig` or `DepthEncoderConfig`): `video.g`, `video.crf`, `video.preset`, `video.fast_decode`, `video.video_backend`, `video.extra_options`.
|
||||
- **Stream-derived** (read back from the encoded MP4 with PyAV): `video.height`, `video.width`, `video.codec`, `video.pix_fmt`, `video.fps`, `video.channels`, `video.is_depth_map`, plus `audio.*` if an audio stream is present.
|
||||
- **Encoder-derived** (taken from `VideoEncoderConfig`): `video.g`, `video.crf`, `video.preset`, `video.fast_decode`, `video.video_backend`, `video.extra_options`.
|
||||
|
||||
<Tip>
|
||||
This block is populated **once**, from the **first** episode. It assumes every
|
||||
episode in the dataset was encoded with the same `rgb_encoder`. Changing
|
||||
episode in the dataset was encoded with the same `camera_encoder`. Changing
|
||||
encoder settings partway through a recording is not supported — the
|
||||
`info.json` will only reflect the parameters used for the first episode.
|
||||
</Tip>
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Load an SMPL motion clip and expose it in SONIC's encoder format.
|
||||
|
||||
SONIC's whole-body tracking mode (``encode_mode == 2``) consumes a flat
|
||||
720-vector ``smpl_joints_10frame_step1`` = 10 consecutive frames x 24 SMPL
|
||||
joints x 3 (xyz) at 50 Hz.
|
||||
|
||||
IMPORTANT - frame convention: the encoder expects each frame's joints with the
|
||||
body's *root orientation removed* (per-frame canonical), exactly like the live
|
||||
deploy stream's ``smpl_joints_local`` (see ``process_smpl_joints`` in the GEAR
|
||||
PICO teleop and ``smpl_joints_multi_future_local`` in training). The reference
|
||||
``smpl_filtered`` clips instead store **world-frame** joints (heading retained),
|
||||
so feeding them raw makes the robot move but track poorly / never face-forward.
|
||||
This loader therefore canonicalizes on load using the clip's per-frame root
|
||||
orientation (``pose_aa[:, :3]``):
|
||||
|
||||
A = Rx(+90deg) * rotvec(pose_aa[:, :3]) # y-up -> z-up root quat
|
||||
local = base120 * A^-1 * joints # remove root orient
|
||||
|
||||
with ``base120 = quat(0.5,0.5,0.5,0.5)`` (SMPL base rotation). This reproduces
|
||||
the deployed transform (verified: per-frame hip-heading std -> 0).
|
||||
|
||||
Clip is read from a numpy ``.npz``. Expected keys:
|
||||
smpl_joints : (T, 24, 3) float32 -- world-frame joint positions, 50 fps
|
||||
pose_aa : (T, 72) float32 -- SMPL axis-angle (root = [:, :3])
|
||||
transl : (T, 3) float32 -- global root translation (optional)
|
||||
fps : scalar
|
||||
|
||||
Example:
|
||||
python examples/unitree_g1/motion_loader.py \
|
||||
--motion examples/unitree_g1/motions/walk_forward.npz
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
||||
import numpy as np
|
||||
|
||||
WINDOW = 10 # frames per encoder window (smpl_joints_10frame_step1)
|
||||
N_JOINTS = 24
|
||||
JOINT_DIM = 3
|
||||
SMPL_OBS_DIM = WINDOW * N_JOINTS * JOINT_DIM # 720
|
||||
|
||||
|
||||
def canonicalize_smpl_joints(smpl_joints: np.ndarray, root_aa: np.ndarray) -> np.ndarray:
|
||||
"""Remove per-frame root orientation -> SONIC ``smpl_joints_local`` format.
|
||||
|
||||
Args:
|
||||
smpl_joints: (T, 24, 3) world-frame (z-up) SMPL joint positions.
|
||||
root_aa: (T, 3) SMPL global-orient axis-angle (y-up convention).
|
||||
|
||||
Returns:
|
||||
(T, 24, 3) per-frame root-orientation-removed joints.
|
||||
"""
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
rx90 = Rotation.from_euler("x", 90, degrees=True) # smpl_root_ytoz_up
|
||||
base120 = Rotation.from_quat([0.5, 0.5, 0.5, 0.5]) # remove_smpl_base_rot
|
||||
a = rx90 * Rotation.from_rotvec(root_aa) # z-up root quat (left-mult)
|
||||
b_inv = base120 * a.inv() # inv(remove_smpl_base_rot(a))
|
||||
return np.einsum("tij,tkj->tki", b_inv.as_matrix(), smpl_joints).astype(np.float32)
|
||||
|
||||
|
||||
class SmplMotion:
|
||||
"""A single SMPL clip with SONIC-format windowing."""
|
||||
|
||||
def __init__(self, path: str, loop: bool = True, canonicalize: bool = True):
|
||||
data = np.load(path)
|
||||
smpl_joints = data["smpl_joints"].astype(np.float32) # (T, 24, 3)
|
||||
self.pose_aa = data["pose_aa"].astype(np.float32) if "pose_aa" in data.files else None
|
||||
self.transl = data["transl"].astype(np.float32) if "transl" in data.files else None
|
||||
self.fps = float(data["fps"]) if "fps" in data.files else 50.0
|
||||
self.loop = loop
|
||||
|
||||
if smpl_joints.ndim != 3 or smpl_joints.shape[1:] != (N_JOINTS, JOINT_DIM):
|
||||
raise ValueError(f"Expected smpl_joints (T, {N_JOINTS}, {JOINT_DIM}), got {smpl_joints.shape}")
|
||||
|
||||
# Reference clips store world-frame joints; the encoder wants per-frame
|
||||
# root-orientation-removed joints. Canonicalize when we have the root pose.
|
||||
self.canonicalized = False
|
||||
if canonicalize and self.pose_aa is not None:
|
||||
smpl_joints = canonicalize_smpl_joints(smpl_joints, self.pose_aa[:, :3])
|
||||
self.canonicalized = True
|
||||
self.smpl_joints = smpl_joints
|
||||
|
||||
self.num_frames = self.smpl_joints.shape[0]
|
||||
self._cursor = 0
|
||||
|
||||
def window(self, start: int) -> np.ndarray:
|
||||
"""Return the 720-vector for the 10-frame window beginning at ``start``.
|
||||
|
||||
Frames are laid out oldest->newest, joint-major within a frame:
|
||||
[f0_j0_xyz, f0_j1_xyz, ..., f9_j23_xyz].
|
||||
"""
|
||||
idx = np.arange(start, start + WINDOW)
|
||||
idx = np.mod(idx, self.num_frames) if self.loop else np.clip(idx, 0, self.num_frames - 1)
|
||||
return self.smpl_joints[idx].reshape(-1).astype(np.float32)
|
||||
|
||||
def reset(self):
|
||||
self._cursor = 0
|
||||
|
||||
def step(self) -> np.ndarray:
|
||||
"""Advance one frame and return the current 720-vector window."""
|
||||
w = self.window(self._cursor)
|
||||
self._cursor += 1
|
||||
if self.loop:
|
||||
self._cursor %= self.num_frames
|
||||
return w
|
||||
|
||||
@property
|
||||
def done(self) -> bool:
|
||||
return (not self.loop) and (self._cursor + WINDOW >= self.num_frames)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--motion", required=True, help="Path to motion .npz")
|
||||
parser.add_argument("--no-loop", action="store_true")
|
||||
parser.add_argument(
|
||||
"--no-canon", action="store_true", help="Skip canonicalization (feed raw stored joints)"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
m = SmplMotion(args.motion, loop=not args.no_loop, canonicalize=not args.no_canon)
|
||||
duration = m.num_frames / m.fps
|
||||
print(f"Loaded '{args.motion}'")
|
||||
print(f" frames={m.num_frames} fps={m.fps:.1f} duration={duration:.1f}s")
|
||||
print(
|
||||
f" smpl_joints={m.smpl_joints.shape} canonicalized={m.canonicalized} "
|
||||
f"pose_aa={None if m.pose_aa is None else m.pose_aa.shape} "
|
||||
f"transl={None if m.transl is None else m.transl.shape}"
|
||||
)
|
||||
|
||||
# Sanity: after canonicalization the per-frame body heading should be fixed.
|
||||
j = m.smpl_joints
|
||||
v = j[:, 2, :2] - j[:, 1, :2] # R_hip - L_hip, horizontal
|
||||
a = np.arctan2(v[:, 1], v[:, 0])
|
||||
rlen = np.clip(np.hypot(np.cos(a).mean(), np.sin(a).mean()), 1e-9, 1.0)
|
||||
circ_std = np.degrees(np.sqrt(-2 * np.log(rlen)))
|
||||
print(f" hip-heading circ-std={circ_std:.1f} deg (~0 => orientation removed; large => world-frame)")
|
||||
|
||||
w0 = m.window(0)
|
||||
print(f" window(0): shape={w0.shape} (expected {SMPL_OBS_DIM}) min={w0.min():.3f} max={w0.max():.3f}")
|
||||
assert w0.shape == (SMPL_OBS_DIM,), "window must be 720-dim for obs[922:1642]"
|
||||
|
||||
# Simulate a few control ticks.
|
||||
print(" stepping 5 ticks:")
|
||||
for t in range(5):
|
||||
w = m.step()
|
||||
print(f" t={t} cursor={m._cursor} window_norm={np.linalg.norm(w):.2f}")
|
||||
|
||||
print("OK: motion loads and yields SONIC-format 720-vec windows.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,88 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Convert a GEAR-SONIC / BONES-SEED ``smpl_filtered`` clip (.pkl) to .npz.
|
||||
|
||||
The reference clips are zlib-compressed joblib pickles holding a dict with
|
||||
``pose_aa`` (T, 72), ``transl`` (T, 3), ``smpl_joints`` (T, 24, 3), ``fps``.
|
||||
``motion_loader.SmplMotion`` consumes the .npz form so the runtime needs no
|
||||
joblib dependency. Canonicalization (root-orientation removal) happens at load
|
||||
time in ``motion_loader``, so this converter just repackages the raw arrays.
|
||||
|
||||
Run this in an environment that has ``joblib`` (e.g. the sonic teleop venv):
|
||||
|
||||
python examples/unitree_g1/pkl_to_npz.py \
|
||||
--pkl sample_data/smpl_filtered/walk_forward_amateur_001__A001.pkl \
|
||||
--out examples/unitree_g1/motions/walk_forward.npz
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def load_pkl(path: str) -> dict:
|
||||
try:
|
||||
import joblib
|
||||
|
||||
return joblib.load(path)
|
||||
except Exception:
|
||||
# joblib clips are zlib-compressed pickles; fall back to manual inflate.
|
||||
import contextlib
|
||||
import pickle # nosec B403 - loads trusted local SMPL clips authored by the user
|
||||
import zlib
|
||||
|
||||
with open(path, "rb") as f:
|
||||
raw = f.read()
|
||||
with contextlib.suppress(zlib.error):
|
||||
raw = zlib.decompress(raw)
|
||||
return pickle.loads(raw) # nosec B301 - local, user-provided motion files only
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--pkl", required=True, help="Input smpl_filtered .pkl")
|
||||
parser.add_argument("--out", required=True, help="Output .npz path")
|
||||
args = parser.parse_args()
|
||||
|
||||
d = load_pkl(args.pkl)
|
||||
if not isinstance(d, dict) or "smpl_joints" not in d:
|
||||
raise ValueError(f"Unexpected pkl structure; keys={list(d) if isinstance(d, dict) else type(d)}")
|
||||
|
||||
smpl_joints = np.asarray(d["smpl_joints"], np.float32)
|
||||
if smpl_joints.ndim != 3 or smpl_joints.shape[1:] != (24, 3):
|
||||
raise ValueError(f"smpl_joints must be (T,24,3), got {smpl_joints.shape}")
|
||||
|
||||
out = {"smpl_joints": smpl_joints, "fps": np.float32(d.get("fps", 50.0))}
|
||||
if "pose_aa" in d:
|
||||
out["pose_aa"] = np.asarray(d["pose_aa"], np.float32)
|
||||
else:
|
||||
print("[warn] no pose_aa -> loader cannot canonicalize (will feed raw)")
|
||||
if "transl" in d:
|
||||
out["transl"] = np.asarray(d["transl"], np.float32)
|
||||
|
||||
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
|
||||
np.savez_compressed(args.out, **out)
|
||||
dur = smpl_joints.shape[0] / float(out["fps"])
|
||||
print(f"Wrote {args.out}")
|
||||
print(
|
||||
f" frames={smpl_joints.shape[0]} fps={float(out['fps']):.1f} duration={dur:.1f}s keys={sorted(out)}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,294 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
SONIC planner with full mode control.
|
||||
|
||||
Keyboard controls:
|
||||
N / P - next / previous motion set
|
||||
1-8 - select mode within current set
|
||||
WASD - movement direction
|
||||
Q / E - rotate facing left / right
|
||||
9 / 0 - decrease / increase speed
|
||||
- / = - decrease / increase height
|
||||
R - force replan
|
||||
M - toggle SMPL motion playback <-> locomotion (needs --motion-file)
|
||||
Space - emergency stop -> IDLE
|
||||
Esc - quit
|
||||
|
||||
Gamepad controls (Unitree wireless controller):
|
||||
Left stick Y - speed (forward = fast, back = stop)
|
||||
Left stick X - movement direction (offset from facing)
|
||||
Right stick X - facing direction (incremental rotation)
|
||||
Right stick Y - height (up = tall 0.8m, down = low 0.1m)
|
||||
Buttons - unused (mode selection is keyboard-only)
|
||||
|
||||
For teleop integration use --robot.controller=SonicWholeBodyController instead.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import faulthandler
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
from motion_loader import SmplMotion
|
||||
|
||||
from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
|
||||
from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
|
||||
CONTROL_DT,
|
||||
DEFAULT_ANGLES,
|
||||
LM,
|
||||
MOTION_SETS,
|
||||
RawKeyboard,
|
||||
compute_kp_kd,
|
||||
drain_keyboard,
|
||||
)
|
||||
from lerobot.robots.unitree_g1.controllers.sonic_whole_body import SonicRuntime
|
||||
from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
|
||||
from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="SONIC planner with keyboard + gamepad control")
|
||||
parser.add_argument(
|
||||
"--ip",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Robot IP for real hardware (e.g. 192.168.123.164). Omit for simulation.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-csv",
|
||||
action="store_true",
|
||||
help="Write /tmp/sonic_pose_log.csv (disabled by default for teleop perf)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cpu",
|
||||
action="store_true",
|
||||
help="Force CPU ONNX Runtime (skip CUDA even if onnxruntime-gpu is installed)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--headless", action="store_true", help="Ignored for sim (stock UnitreeG1 uses hub MuJoCo defaults)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gamepad",
|
||||
action="store_true",
|
||||
help="Read Unitree wireless gamepad in sim (default: keyboard-only in sim)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--keyboard-only", action="store_true", help="Ignore wireless gamepad (terminal keyboard only)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--motion-file",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Play an SMPL motion clip (.npz) via SONIC whole-body mode "
|
||||
"(encode_mode=2) instead of locomotion planning.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-loop", action="store_true", help="With --motion-file, play once instead of looping"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Surface native crashes (onnxruntime / mujoco) with a real traceback, and
|
||||
# avoid losing buffered diagnostics if the process dies mid-loop.
|
||||
faulthandler.enable()
|
||||
with contextlib.suppress(Exception):
|
||||
sys.stdout.reconfigure(line_buffering=True)
|
||||
|
||||
print("=" * 60)
|
||||
print("SONIC planner - full mode control")
|
||||
print(" N/P cycle sets | 1-8 select mode | WASD move")
|
||||
print(" Q/E rotate | 9/0 speed | -/= height")
|
||||
print(" R replan | Space IDLE | Esc quit")
|
||||
if args.ip:
|
||||
print(f" Robot IP: {args.ip}")
|
||||
else:
|
||||
print(" Mode: simulation")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
cfg = UnitreeG1Config(controller=None) # full-body SONIC; standalone loop owns publish
|
||||
if args.ip:
|
||||
cfg.is_simulation = False
|
||||
cfg.robot_ip = args.ip
|
||||
else:
|
||||
cfg.is_simulation = True
|
||||
if args.headless:
|
||||
print("[Note] --headless ignored: sim uses stock UnitreeG1 + hub env")
|
||||
robot = UnitreeG1(cfg)
|
||||
robot.connect()
|
||||
kp, kd = compute_kp_kd()
|
||||
robot.kp = kp.copy()
|
||||
robot.kd = kd.copy()
|
||||
|
||||
runtime = SonicRuntime(force_cpu=args.cpu)
|
||||
controller = runtime.controller
|
||||
ms = runtime.ms
|
||||
|
||||
motion = None
|
||||
if args.motion_file:
|
||||
motion = SmplMotion(args.motion_file, loop=not args.no_loop)
|
||||
controller.smpl_motion = motion # lets 'M' key toggle playback
|
||||
controller.encode_mode = 2 # start in SONIC whole-body SMPL imitation
|
||||
dur = motion.num_frames / motion.fps
|
||||
print(f"\n[Motion] SMPL whole-body playback: {args.motion_file}")
|
||||
print(
|
||||
f" frames={motion.num_frames} fps={motion.fps:.1f} "
|
||||
f"duration={dur:.1f}s loop={not args.no_loop} encode_mode=2"
|
||||
)
|
||||
print(" Press 'M' to toggle SMPL playback <-> locomotion at runtime.")
|
||||
|
||||
runtime.controller.print_input_diagnostics()
|
||||
|
||||
print(f"\nStarting: {MOTION_SETS[0][0]} (default mode: {LM(ms.mode).name})")
|
||||
[print(f" {i + 1}: {m.name}") for i, m in enumerate(MOTION_SETS[0][1])]
|
||||
print(
|
||||
"\n[Ready] Click THIS terminal, then W/A/S/D to move. 1-6 change mode, 9/0 speed, Esc quit.\n",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# Sim hub publishes wireless_remote bytes that can fight terminal WASD.
|
||||
base_joystick = not args.keyboard_only and (args.gamepad or args.ip is not None)
|
||||
|
||||
with RawKeyboard() as kb:
|
||||
try:
|
||||
gc.disable()
|
||||
gc_timer = 0.0
|
||||
robot.reset(CONTROL_DT, DEFAULT_ANGLES)
|
||||
time.sleep(1.0)
|
||||
|
||||
last_status = time.time() - 2.1
|
||||
loop_t, enc_t, dec_t, obs_t, act_t = [], [], [], [], []
|
||||
slow_n = blend_n = 0
|
||||
stall_src = ""
|
||||
did_blend = False
|
||||
t_start = time.time()
|
||||
|
||||
log_path = os.path.join(tempfile.gettempdir(), "sonic_pose_log.csv")
|
||||
jnames = [m.name for m in G1_29_JointIndex]
|
||||
log_ctx = open(log_path, "w") if args.log_csv else None # noqa: SIM115
|
||||
if log_ctx:
|
||||
log_ctx.write(
|
||||
"t,step,cursor,ts,blend,mode,"
|
||||
+ ",".join(f"q{i}" for i in range(29))
|
||||
+ ","
|
||||
+ ",".join(f"ref{i}" for i in range(29))
|
||||
+ ","
|
||||
+ ",".join(f"act{i}" for i in range(29))
|
||||
+ ",delta_max,action_norm,token_norm\n"
|
||||
)
|
||||
|
||||
try:
|
||||
while not robot._shutdown_event.is_set():
|
||||
t0 = time.time()
|
||||
if drain_keyboard(kb, ms, controller):
|
||||
break
|
||||
|
||||
obs = robot.get_observation()
|
||||
t_obs = time.time()
|
||||
obs_t.append(1000 * (t_obs - t0))
|
||||
if not obs:
|
||||
runtime.tick({}, use_joystick=False)
|
||||
time.sleep(max(0.0, CONTROL_DT - (time.time() - t0)))
|
||||
continue
|
||||
|
||||
# SMPL playback only while in whole-body mode; 'M' toggles it.
|
||||
motion_active = motion is not None and controller.encode_mode == 2
|
||||
if motion_active:
|
||||
controller.smpl_joints_10frame_step1 = motion.step()
|
||||
if motion.done:
|
||||
print("\n[Motion] clip finished")
|
||||
break
|
||||
|
||||
step_before = runtime.step
|
||||
t_step = time.time()
|
||||
action = runtime.tick(obs, use_joystick=base_joystick and not motion_active)
|
||||
step_ms = 1000 * (time.time() - t_step)
|
||||
do_enc = step_before % 5 == 0
|
||||
(enc_t if do_enc else dec_t).append(step_ms)
|
||||
|
||||
t_act = time.time()
|
||||
robot.send_action(action)
|
||||
act_t.append(1000 * (time.time() - t_act))
|
||||
|
||||
if log_ctx and runtime.step % 5 == 0:
|
||||
t_rel = time.time() - t_start
|
||||
q_r = np.array([obs.get(f"{n}.q", 0) for n in jnames])
|
||||
a_v = np.array([action.get(f"{n}.q", 0) for n in jnames])
|
||||
cur, ts = controller.ref_cursor, controller.motion_timesteps
|
||||
q_ref = (
|
||||
controller.motion_joint_positions[min(cur, ts - 1)] if ts > 0 else np.zeros(29)
|
||||
)
|
||||
log_ctx.write(
|
||||
f"{t_rel:.4f},{runtime.step},{cur},{ts},{int(did_blend)},{ms.mode},"
|
||||
+ ",".join(f"{v:.6f}" for v in q_r)
|
||||
+ ","
|
||||
+ ",".join(f"{v:.6f}" for v in q_ref)
|
||||
+ ","
|
||||
+ ",".join(f"{v:.6f}" for v in a_v)
|
||||
+ ","
|
||||
+ f"{np.max(np.abs(a_v - q_r)):.6f},"
|
||||
f"{np.linalg.norm(a_v):.6f},"
|
||||
f"{np.linalg.norm(controller.token):.6f}\n"
|
||||
)
|
||||
did_blend = False
|
||||
|
||||
now = time.time()
|
||||
loop_ms = 1000 * (now - t0)
|
||||
if loop_ms > 50:
|
||||
stall_src = (
|
||||
f"[STALL] {loop_ms:.0f}ms: "
|
||||
f"obs={obs_t[-1]:.0f} step={step_ms:.0f} act={act_t[-1]:.0f}"
|
||||
)
|
||||
if loop_ms > CONTROL_DT * 1500:
|
||||
slow_n += 1
|
||||
|
||||
if now - last_status > 2.0:
|
||||
|
||||
def _avg(lst):
|
||||
return sum(lst) / len(lst) if lst else 0
|
||||
|
||||
hz = 1000 / _avg(loop_t) if _avg(loop_t) else 0
|
||||
print(
|
||||
f"\r {ms.status_line()} step={runtime.step} "
|
||||
f"ref={controller.ref_cursor}/{controller.motion_timesteps} "
|
||||
f"loop={_avg(loop_t):.1f}ms(max={max(loop_t, default=0):.1f}) hz={hz:.0f} "
|
||||
f"enc={_avg(enc_t):.1f} dec={_avg(dec_t):.1f} obs={_avg(obs_t):.1f} "
|
||||
f"slow={slow_n} blends={blend_n}",
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
if stall_src:
|
||||
print(f"\n {stall_src}")
|
||||
stall_src = ""
|
||||
last_status = now
|
||||
loop_t, enc_t, dec_t, obs_t, act_t = [], [], [], [], []
|
||||
slow_n = blend_n = 0
|
||||
|
||||
gc_timer += CONTROL_DT
|
||||
if gc_timer >= 10.0:
|
||||
gc.collect()
|
||||
gc_timer = 0.0
|
||||
loop_t.append(loop_ms)
|
||||
time.sleep(max(0.0, CONTROL_DT - (time.time() - t0)))
|
||||
finally:
|
||||
if log_ctx:
|
||||
log_ctx.close()
|
||||
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
finally:
|
||||
gc.enable()
|
||||
if args.log_csv:
|
||||
print(f"\n[Log] Saved to {log_path}")
|
||||
runtime.shutdown()
|
||||
print("\nStopping...")
|
||||
if robot.is_connected:
|
||||
robot.disconnect()
|
||||
print("Done.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+2
-8
@@ -124,7 +124,7 @@ hardware = [
|
||||
"lerobot[deepdiff-dep]",
|
||||
]
|
||||
viz = [
|
||||
"rerun-sdk>=0.24.0,<0.34.0",
|
||||
"rerun-sdk>=0.24.0,<0.27.0",
|
||||
]
|
||||
# ── User-facing composite extras (map to CLI scripts) ─────
|
||||
# lerobot-record, lerobot-replay, lerobot-calibrate, lerobot-teleoperate, etc.
|
||||
@@ -229,10 +229,6 @@ robometer = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]", "lerobot
|
||||
topreward = ["lerobot[transformers-dep]"]
|
||||
xvla = ["lerobot[transformers-dep]"]
|
||||
eo1 = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]"]
|
||||
fastwam = [
|
||||
"lerobot[transformers-dep]",
|
||||
"lerobot[diffusers-dep]",
|
||||
]
|
||||
hilserl = ["lerobot[transformers-dep]", "lerobot[dataset]", "gym-hil>=0.1.14,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
|
||||
vla_jepa = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[qwen-vl-utils-dep]"]
|
||||
|
||||
@@ -312,7 +308,6 @@ all = [
|
||||
"lerobot[pi]",
|
||||
"lerobot[molmoact2]",
|
||||
"lerobot[smolvla]",
|
||||
"lerobot[fastwam]",
|
||||
# "lerobot[groot]", TODO(Steven): Gr00t requires specific installation instructions for flash-attn
|
||||
"lerobot[xvla]",
|
||||
"lerobot[hilserl]",
|
||||
@@ -449,8 +444,7 @@ default.extend-ignore-identifiers-re = [
|
||||
"is_compileable",
|
||||
"ROBOTIS",
|
||||
"OT_VALUE",
|
||||
"VanderBilt",
|
||||
"seperated_timestep",
|
||||
"VanderBilt"
|
||||
]
|
||||
|
||||
# TODO: Uncomment when ready to use
|
||||
|
||||
@@ -36,7 +36,7 @@ from typing import Any, Protocol
|
||||
import PIL.Image
|
||||
import torch
|
||||
|
||||
from lerobot.configs import RGBEncoderConfig
|
||||
from lerobot.configs.video import VideoEncoderConfig
|
||||
from lerobot.datasets.video_utils import decode_video_frames, reencode_video
|
||||
|
||||
from .reader import EpisodeRecord, snap_to_frame
|
||||
@@ -164,9 +164,7 @@ class VideoFrameProvider:
|
||||
# only for video-stored cameras. Image-stored cameras (also in
|
||||
# ``camera_keys``) would KeyError, so restrict the list — and the
|
||||
# default — to video keys.
|
||||
# Depth cameras are excluded from the annotation pipeline for now.
|
||||
depth_keys = set(self._meta.depth_keys)
|
||||
keys = [key for key in self._meta.video_keys if key not in depth_keys]
|
||||
keys = list(self._meta.video_keys)
|
||||
# Last-resort fallback: if metadata didn't surface any video keys but
|
||||
# the caller explicitly named a camera (``--vlm.camera_key=...``),
|
||||
# trust them — the key is by definition known to exist on the dataset.
|
||||
@@ -278,12 +276,12 @@ class VideoFrameProvider:
|
||||
from_timestamp = float(ep[f"videos/{self.camera_key}/from_timestamp"])
|
||||
to_timestamp = float(ep[f"videos/{self.camera_key}/to_timestamp"])
|
||||
src = self.root / self._meta.get_video_file_path(record.episode_index, self.camera_key)
|
||||
encoder = RGBEncoderConfig(vcodec="h264", pix_fmt="yuv420p", g=None, crf=23, preset="ultrafast")
|
||||
encoder = VideoEncoderConfig(vcodec="h264", pix_fmt="yuv420p", g=None, crf=23, preset="ultrafast")
|
||||
try:
|
||||
reencode_video(
|
||||
src,
|
||||
out_path,
|
||||
video_encoder=encoder,
|
||||
camera_encoder=encoder,
|
||||
overwrite=True,
|
||||
start_time_s=from_timestamp,
|
||||
end_time_s=to_timestamp,
|
||||
|
||||
@@ -105,9 +105,8 @@ def raw_observation_to_observation(
|
||||
|
||||
|
||||
def prepare_image(image: torch.Tensor) -> torch.Tensor:
|
||||
"""Minimal preprocessing to turn RGB uint8 images to float32 in [0, 1], and create a memory-contiguous tensor"""
|
||||
if image.dtype == torch.uint8:
|
||||
image = image.type(torch.float32) / 255
|
||||
"""Minimal preprocessing to turn int8 images to float32 in [0, 1], and create a memory-contiguous tensor"""
|
||||
image = image.type(torch.float32) / 255
|
||||
image = image.contiguous()
|
||||
|
||||
return image
|
||||
|
||||
@@ -436,7 +436,7 @@ class OpenCVCamera(Camera):
|
||||
Internal loop run by the background thread for asynchronous reading.
|
||||
|
||||
On each iteration:
|
||||
1. Reads a color frame (blocking call)
|
||||
1. Reads a color frame
|
||||
2. Stores result in latest_frame and updates timestamp (thread-safe)
|
||||
3. Sets new_frame_event to notify listeners
|
||||
|
||||
@@ -485,8 +485,6 @@ class OpenCVCamera(Camera):
|
||||
|
||||
if self.thread is not None and self.thread.is_alive():
|
||||
self.thread.join(timeout=2.0)
|
||||
if self.thread.is_alive():
|
||||
logger.warning(f"{self} read thread did not terminate within timeout.")
|
||||
|
||||
self.thread = None
|
||||
self.stop_event = None
|
||||
|
||||
@@ -128,7 +128,6 @@ class RealSenseCamera(Camera):
|
||||
|
||||
self.fps = config.fps
|
||||
self.color_mode = config.color_mode
|
||||
self.use_rgb = config.use_rgb
|
||||
self.use_depth = config.use_depth
|
||||
self.warmup_s = config.warmup_s
|
||||
|
||||
@@ -196,15 +195,12 @@ class RealSenseCamera(Camera):
|
||||
# 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)
|
||||
self.async_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
|
||||
):
|
||||
if 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.")
|
||||
@@ -272,13 +268,13 @@ class RealSenseCamera(Camera):
|
||||
)
|
||||
|
||||
if len(found_devices) > 1:
|
||||
serial_numbers = [dev["id"] for dev in found_devices]
|
||||
serial_numbers = [dev["serial_number"] for dev in found_devices]
|
||||
raise ValueError(
|
||||
f"Multiple RealSense cameras found with name '{name}'. "
|
||||
f"Please use a unique serial number instead. Found SNs: {serial_numbers}"
|
||||
)
|
||||
|
||||
serial_number = str(found_devices[0]["id"])
|
||||
serial_number = str(found_devices[0]["serial_number"])
|
||||
return serial_number
|
||||
|
||||
def _configure_rs_pipeline_config(self, rs_config: Any) -> None:
|
||||
@@ -286,17 +282,15 @@ class RealSenseCamera(Camera):
|
||||
rs.config.enable_device(rs_config, self.serial_number)
|
||||
|
||||
if self.width and self.height and self.fps:
|
||||
if self.use_rgb:
|
||||
rs_config.enable_stream(
|
||||
rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps
|
||||
)
|
||||
rs_config.enable_stream(
|
||||
rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps
|
||||
)
|
||||
if self.use_depth:
|
||||
rs_config.enable_stream(
|
||||
rs.stream.depth, self.capture_width, self.capture_height, rs.format.z16, self.fps
|
||||
)
|
||||
else:
|
||||
if self.use_rgb:
|
||||
rs_config.enable_stream(rs.stream.color)
|
||||
rs_config.enable_stream(rs.stream.color)
|
||||
if self.use_depth:
|
||||
rs_config.enable_stream(rs.stream.depth)
|
||||
|
||||
@@ -304,9 +298,8 @@ class RealSenseCamera(Camera):
|
||||
def _configure_capture_settings(self) -> None:
|
||||
"""Sets fps, width, and height from device stream if not already configured.
|
||||
|
||||
Uses the color stream profile (or the depth stream profile when the color
|
||||
stream is disabled) to update unset attributes. Handles rotation by swapping
|
||||
width/height when needed. Original capture dimensions are always stored.
|
||||
Uses the color stream profile to update unset attributes. Handles rotation by
|
||||
swapping width/height when needed. Original capture dimensions are always stored.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If device is not connected.
|
||||
@@ -315,8 +308,7 @@ class RealSenseCamera(Camera):
|
||||
if self.rs_profile is None:
|
||||
raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
|
||||
|
||||
rs_stream = rs.stream.color if self.use_rgb else rs.stream.depth
|
||||
stream = self.rs_profile.get_stream(rs_stream).as_video_stream_profile()
|
||||
stream = self.rs_profile.get_stream(rs.stream.color).as_video_stream_profile()
|
||||
|
||||
if self.fps is None:
|
||||
self.fps = stream.fps()
|
||||
@@ -331,14 +323,6 @@ class RealSenseCamera(Camera):
|
||||
self.width, self.height = actual_width, actual_height
|
||||
self.capture_width, self.capture_height = actual_width, actual_height
|
||||
|
||||
def _read(self, read_depth: bool = False) -> NDArray[Any]:
|
||||
"""Shared helper for :meth:`read`/:meth:`read_depth`: wait for a fresh color or depth frame."""
|
||||
if self.thread is None or not self.thread.is_alive():
|
||||
raise RuntimeError(f"{self} read thread is not running.")
|
||||
|
||||
self.new_frame_event.clear()
|
||||
return self._async_read(timeout_ms=10000, read_depth=read_depth)
|
||||
|
||||
@check_if_not_connected
|
||||
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
|
||||
"""
|
||||
@@ -348,8 +332,8 @@ class RealSenseCamera(Camera):
|
||||
from the camera hardware via the RealSense pipeline.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The depth map as a NumPy array (height, width, 1)
|
||||
of type `np.uint16` (raw depth values in millimeters).
|
||||
np.ndarray: The depth map as a NumPy array (height, width)
|
||||
of type `np.uint16` (raw depth values in millimeters) and rotation.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the camera is not connected.
|
||||
@@ -365,7 +349,20 @@ class RealSenseCamera(Camera):
|
||||
f"Failed to capture depth frame '.read_depth()'. Depth stream is not enabled for {self}."
|
||||
)
|
||||
|
||||
return self._read(read_depth=True)
|
||||
if self.thread is None or not self.thread.is_alive():
|
||||
raise RuntimeError(f"{self} read thread is not running.")
|
||||
|
||||
self.new_frame_event.clear()
|
||||
|
||||
_ = self.async_read(timeout_ms=10000)
|
||||
|
||||
with self.frame_lock:
|
||||
depth_map = self.latest_depth_frame
|
||||
|
||||
if depth_map is None:
|
||||
raise RuntimeError("No depth frame available. Ensure camera is streaming.")
|
||||
|
||||
return depth_map
|
||||
|
||||
def _read_from_hardware(self):
|
||||
if self.rs_pipeline is None:
|
||||
@@ -408,10 +405,12 @@ class RealSenseCamera(Camera):
|
||||
f"{self} read() timeout_ms parameter is deprecated and will be removed in future versions."
|
||||
)
|
||||
|
||||
if not self.use_rgb:
|
||||
raise RuntimeError(f"{self}: cannot read color — camera was configured with use_rgb=False.")
|
||||
if self.thread is None or not self.thread.is_alive():
|
||||
raise RuntimeError(f"{self} read thread is not running.")
|
||||
|
||||
frame = self._read()
|
||||
self.new_frame_event.clear()
|
||||
|
||||
frame = self.async_read(timeout_ms=10000)
|
||||
|
||||
read_duration_ms = (time.perf_counter() - start_time) * 1e3
|
||||
logger.debug(f"{self} read took: {read_duration_ms:.1f}ms")
|
||||
@@ -466,8 +465,8 @@ class RealSenseCamera(Camera):
|
||||
Internal loop run by the background thread for asynchronous reading.
|
||||
|
||||
On each iteration:
|
||||
1. Reads a color/depth frame (blocking call with 10s timeout)
|
||||
2. Stores result in latest_color_frame/latest_depth_frame and updates timestamp (thread-safe)
|
||||
1. Reads a color frame with 500ms timeout
|
||||
2. Stores result in latest_frame and updates timestamp (thread-safe)
|
||||
3. Sets new_frame_event to notify listeners
|
||||
|
||||
Stops on DeviceNotConnectedError, logs other errors and continues.
|
||||
@@ -480,24 +479,19 @@ class RealSenseCamera(Camera):
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
frame = self._read_from_hardware()
|
||||
|
||||
if self.use_rgb:
|
||||
color_frame_raw = frame.get_color_frame()
|
||||
color_frame = np.asanyarray(color_frame_raw.get_data())
|
||||
processed_color_frame = self._postprocess_image(color_frame)
|
||||
color_frame_raw = frame.get_color_frame()
|
||||
color_frame = np.asanyarray(color_frame_raw.get_data())
|
||||
processed_color_frame = self._postprocess_image(color_frame)
|
||||
|
||||
if self.use_depth:
|
||||
depth_frame_raw = frame.get_depth_frame()
|
||||
depth_frame = np.asanyarray(depth_frame_raw.get_data())
|
||||
processed_depth_frame = self._postprocess_image(depth_frame, depth_frame=True)
|
||||
if processed_depth_frame.ndim == 2: # (H, W) -> (H, W, 1)
|
||||
processed_depth_frame = processed_depth_frame[..., np.newaxis]
|
||||
|
||||
capture_time = time.perf_counter()
|
||||
|
||||
with self.frame_lock:
|
||||
if self.use_rgb:
|
||||
self.latest_color_frame = processed_color_frame
|
||||
self.latest_color_frame = processed_color_frame
|
||||
if self.use_depth:
|
||||
self.latest_depth_frame = processed_depth_frame
|
||||
self.latest_timestamp = capture_time
|
||||
@@ -529,8 +523,6 @@ class RealSenseCamera(Camera):
|
||||
|
||||
if self.thread is not None and self.thread.is_alive():
|
||||
self.thread.join(timeout=2.0)
|
||||
if self.thread.is_alive(): # pragma: no cover
|
||||
logger.warning(f"{self} read thread did not terminate within timeout.")
|
||||
|
||||
self.thread = None
|
||||
self.stop_event = None
|
||||
@@ -541,26 +533,7 @@ class RealSenseCamera(Camera):
|
||||
self.latest_timestamp = None
|
||||
self.new_frame_event.clear()
|
||||
|
||||
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():
|
||||
raise RuntimeError(f"{self} read thread is not running.")
|
||||
|
||||
if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0):
|
||||
raise TimeoutError(
|
||||
f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. "
|
||||
f"Read thread alive: {self.thread.is_alive()}."
|
||||
)
|
||||
|
||||
with self.frame_lock:
|
||||
frame = self.latest_depth_frame if read_depth else self.latest_color_frame
|
||||
self.new_frame_event.clear()
|
||||
|
||||
if frame is None:
|
||||
raise RuntimeError(f"Internal error: Event set but no frame available for {self}.")
|
||||
|
||||
return frame
|
||||
|
||||
# NOTE(Steven): Missing implementation for depth for now
|
||||
@check_if_not_connected
|
||||
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
|
||||
"""
|
||||
@@ -585,31 +558,25 @@ class RealSenseCamera(Camera):
|
||||
RuntimeError: If the background thread died unexpectedly or another error occurs.
|
||||
"""
|
||||
|
||||
if not self.use_rgb:
|
||||
raise RuntimeError(f"{self}: cannot read color — camera was configured with use_rgb=False.")
|
||||
|
||||
return self._async_read(timeout_ms=timeout_ms)
|
||||
|
||||
def _read_latest(self, max_age_ms: int, read_depth: bool = False) -> NDArray[Any]:
|
||||
"""Shared helper for :meth:`read_latest`/:meth:`read_latest_depth`: peek the latest buffered frame."""
|
||||
if self.thread is None or not self.thread.is_alive():
|
||||
raise RuntimeError(f"{self} read thread is not running.")
|
||||
|
||||
with self.frame_lock:
|
||||
frame = self.latest_depth_frame if read_depth else self.latest_color_frame
|
||||
timestamp = self.latest_timestamp
|
||||
|
||||
if frame is None or timestamp is None:
|
||||
raise RuntimeError(f"{self} has not captured any frames yet.")
|
||||
|
||||
age_ms = (time.perf_counter() - timestamp) * 1e3
|
||||
if age_ms > max_age_ms:
|
||||
if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0):
|
||||
raise TimeoutError(
|
||||
f"{self} latest frame is too old: {age_ms:.1f} ms (max allowed: {max_age_ms} ms)."
|
||||
f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. "
|
||||
f"Read thread alive: {self.thread.is_alive()}."
|
||||
)
|
||||
|
||||
with self.frame_lock:
|
||||
frame = self.latest_color_frame
|
||||
self.new_frame_event.clear()
|
||||
|
||||
if frame is None:
|
||||
raise RuntimeError(f"Internal error: Event set but no frame available for {self}.")
|
||||
|
||||
return frame
|
||||
|
||||
# NOTE(Steven): Missing implementation for depth for now
|
||||
@check_if_not_connected
|
||||
def read_latest(self, max_age_ms: int = 500) -> NDArray[Any]:
|
||||
"""Return the most recent (color) frame captured immediately (Peeking).
|
||||
@@ -626,48 +593,24 @@ class RealSenseCamera(Camera):
|
||||
DeviceNotConnectedError: If the camera is not connected.
|
||||
RuntimeError: If the camera is connected but has not captured any frames yet.
|
||||
"""
|
||||
if not self.use_rgb:
|
||||
raise RuntimeError(f"{self}: cannot read color — camera was configured with use_rgb=False.")
|
||||
|
||||
return self._read_latest(max_age_ms=max_age_ms)
|
||||
if self.thread is None or not self.thread.is_alive():
|
||||
raise RuntimeError(f"{self} read thread is not running.")
|
||||
|
||||
@check_if_not_connected
|
||||
def async_read_depth(self, timeout_ms: float = 200) -> NDArray[np.uint16]:
|
||||
"""Read the latest depth frame asynchronously, in millimeters.
|
||||
with self.frame_lock:
|
||||
frame = self.latest_color_frame
|
||||
timestamp = self.latest_timestamp
|
||||
|
||||
Mirrors :meth:`async_read` but returns the depth stream rather than the
|
||||
color stream. Output is ``np.uint16`` of shape ``(H, W, 1)``, where each
|
||||
pixel is the distance from the sensor in millimeters.
|
||||
if frame is None or timestamp is None:
|
||||
raise RuntimeError(f"{self} has not captured any frames yet.")
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the camera is not connected.
|
||||
RuntimeError: If ``use_depth`` is ``False`` for this camera, or if
|
||||
the background read thread is not running.
|
||||
TimeoutError: If no frame becomes available within ``timeout_ms``.
|
||||
"""
|
||||
if not self.use_depth:
|
||||
raise RuntimeError(f"{self}: cannot read depth — camera was configured with use_depth=False.")
|
||||
age_ms = (time.perf_counter() - timestamp) * 1e3
|
||||
if age_ms > max_age_ms:
|
||||
raise TimeoutError(
|
||||
f"{self} latest frame is too old: {age_ms:.1f} ms (max allowed: {max_age_ms} ms)."
|
||||
)
|
||||
|
||||
return self._async_read(timeout_ms=timeout_ms, read_depth=True)
|
||||
|
||||
@check_if_not_connected
|
||||
def read_latest_depth(self, max_age_ms: int = 500) -> NDArray[Any]:
|
||||
"""Return the most recent depth frame in millimeters (peeking).
|
||||
|
||||
Non-blocking counterpart of :meth:`read_latest` for the depth stream.
|
||||
Output is ``np.uint16`` of shape ``(H, W, 1)``, where each pixel is the
|
||||
distance from the sensor in millimeters.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the camera is not connected.
|
||||
RuntimeError: If ``use_depth`` is ``False`` for this camera, or if
|
||||
no depth frame has been captured yet.
|
||||
TimeoutError: If the latest depth frame is older than ``max_age_ms``.
|
||||
"""
|
||||
if not self.use_depth:
|
||||
raise RuntimeError(f"{self}: cannot read depth — camera was configured with use_depth=False.")
|
||||
|
||||
return self._read_latest(max_age_ms=max_age_ms, read_depth=True)
|
||||
return frame
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""
|
||||
|
||||
@@ -42,14 +42,12 @@ class RealSenseCameraConfig(CameraConfig):
|
||||
height: Requested frame height in pixels for the color stream.
|
||||
serial_number_or_name: Unique serial number or human-readable name to identify the camera.
|
||||
color_mode: Color mode for image output (RGB or BGR). Defaults to RGB.
|
||||
use_rgb: Whether to enable the color stream. Defaults to True.
|
||||
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)
|
||||
|
||||
Note:
|
||||
- Either name or serial_number must be specified.
|
||||
- At least one of `use_rgb` or `use_depth` must be enabled.
|
||||
- Depth stream configuration (if enabled) will use the same FPS as the color stream.
|
||||
- The actual resolution and FPS may be adjusted by the camera to the nearest supported mode.
|
||||
- For `fps`, `width` and `height`, either all of them need to be set, or none of them.
|
||||
@@ -57,7 +55,6 @@ class RealSenseCameraConfig(CameraConfig):
|
||||
|
||||
serial_number_or_name: str
|
||||
color_mode: ColorMode = ColorMode.RGB
|
||||
use_rgb: bool = True
|
||||
use_depth: bool = False
|
||||
rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
|
||||
warmup_s: int = 1
|
||||
@@ -66,9 +63,6 @@ class RealSenseCameraConfig(CameraConfig):
|
||||
self.color_mode = ColorMode(self.color_mode)
|
||||
self.rotation = Cv2Rotation(self.rotation)
|
||||
|
||||
if not self.use_rgb and not self.use_depth:
|
||||
raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")
|
||||
|
||||
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(
|
||||
|
||||
@@ -293,8 +293,6 @@ class ZMQCamera(Camera):
|
||||
|
||||
if self.thread is not None and self.thread.is_alive():
|
||||
self.thread.join(timeout=2.0)
|
||||
if self.thread.is_alive():
|
||||
logger.warning(f"{self} read thread did not terminate within timeout.")
|
||||
|
||||
self.thread = None
|
||||
self.stop_event = None
|
||||
|
||||
@@ -15,7 +15,6 @@
|
||||
# limitations under the License.
|
||||
from pathlib import Path
|
||||
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
from torch.optim import Optimizer
|
||||
from torch.optim.lr_scheduler import LRScheduler
|
||||
|
||||
@@ -36,7 +35,6 @@ from lerobot.utils.constants import (
|
||||
TRAINING_STATE_DIR,
|
||||
TRAINING_STEP,
|
||||
)
|
||||
from lerobot.utils.hub import find_latest_hub_checkpoint
|
||||
from lerobot.utils.io_utils import load_json, write_json
|
||||
from lerobot.utils.random_utils import load_rng_state, save_rng_state
|
||||
|
||||
@@ -285,61 +283,3 @@ def load_fsdp_optimizer_state(model, optimizer, checkpoint_dir: Path) -> None:
|
||||
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg):
|
||||
sharded_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=full_osd)
|
||||
optimizer.load_state_dict(sharded_osd)
|
||||
|
||||
|
||||
def push_checkpoint_to_hub(
|
||||
checkpoint_dir: Path,
|
||||
repo_id: str,
|
||||
*,
|
||||
private: bool | None = None,
|
||||
) -> None:
|
||||
"""Upload a saved checkpoint directory to the Hub under checkpoints/<name>/.
|
||||
|
||||
Called once per save step when save_checkpoint_to_hub is enabled, so a
|
||||
timed-out or crashed run still leaves recoverable checkpoints on the Hub.
|
||||
The model repo is created idempotently, and the commit is tagged with the
|
||||
checkpoint step so a checkpoint can be recovered with
|
||||
--policy.pretrained_revision=<step> instead of a commit sha.
|
||||
"""
|
||||
api = HfApi()
|
||||
api.create_repo(repo_id=repo_id, repo_type="model", private=private, exist_ok=True)
|
||||
commit = api.upload_folder(
|
||||
folder_path=str(checkpoint_dir),
|
||||
repo_id=repo_id,
|
||||
repo_type="model",
|
||||
path_in_repo=f"checkpoints/{checkpoint_dir.name}",
|
||||
commit_message=f"checkpoint {checkpoint_dir.name}",
|
||||
)
|
||||
api.create_tag(
|
||||
repo_id=repo_id,
|
||||
tag=checkpoint_dir.name,
|
||||
revision=commit.oid,
|
||||
repo_type="model",
|
||||
exist_ok=True,
|
||||
)
|
||||
|
||||
|
||||
def resolve_resume_checkpoint(repo_id: str, output_dir: Path) -> Path:
|
||||
"""Download the latest checkpoint of a Hub training repo into a local run dir.
|
||||
|
||||
The symmetric counterpart to `push_checkpoint_to_hub`: given a model repo holding
|
||||
`checkpoints/<step>/{pretrained_model,training_state}` subtrees, download the highest-numbered step
|
||||
into `output_dir/checkpoints/<step>/`, recreate the local `last` symlink, and return that local
|
||||
checkpoint dir. Used to resume training from the Hub on a machine (or HF Jobs pod) that does not
|
||||
have the original local run dir.
|
||||
"""
|
||||
latest = find_latest_hub_checkpoint(repo_id)
|
||||
if latest is None:
|
||||
raise FileNotFoundError(
|
||||
f"No checkpoint found in '{repo_id}' under '{CHECKPOINTS_DIR}/'. "
|
||||
"Was the run trained with --save_checkpoint_to_hub?"
|
||||
)
|
||||
snapshot_download(
|
||||
repo_id=repo_id,
|
||||
repo_type="model",
|
||||
allow_patterns=f"{latest}/*",
|
||||
local_dir=str(output_dir),
|
||||
)
|
||||
checkpoint_dir = output_dir / latest
|
||||
update_last_checkpoint(checkpoint_dir)
|
||||
return checkpoint_dir
|
||||
|
||||
@@ -22,7 +22,7 @@ Import them directly: ``from lerobot.configs.train import TrainPipelineConfig``
|
||||
"""
|
||||
|
||||
from .dataset import DatasetRecordConfig
|
||||
from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
|
||||
from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
|
||||
from .policies import PreTrainedConfig
|
||||
from .recipe import MessageTurn, TrainingRecipe, load_recipe
|
||||
from .types import (
|
||||
@@ -33,15 +33,10 @@ from .types import (
|
||||
RTCAttentionSchedule,
|
||||
)
|
||||
from .video import (
|
||||
DEFAULT_DEPTH_UNIT,
|
||||
VALID_VIDEO_CODECS,
|
||||
VIDEO_ENCODER_INFO_KEYS,
|
||||
DepthEncoderConfig,
|
||||
RGBEncoderConfig,
|
||||
VideoEncoderConfig,
|
||||
depth_encoder_defaults,
|
||||
encoder_config_from_video_info,
|
||||
rgb_encoder_defaults,
|
||||
camera_encoder_defaults,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
@@ -55,7 +50,6 @@ __all__ = [
|
||||
"DatasetRecordConfig",
|
||||
"DatasetConfig",
|
||||
"EvalConfig",
|
||||
"JobConfig",
|
||||
"MessageTurn",
|
||||
"PeftConfig",
|
||||
"PreTrainedConfig",
|
||||
@@ -63,15 +57,9 @@ __all__ = [
|
||||
"WandBConfig",
|
||||
"load_recipe",
|
||||
"VideoEncoderConfig",
|
||||
"RGBEncoderConfig",
|
||||
"DepthEncoderConfig",
|
||||
# Defaults
|
||||
"rgb_encoder_defaults",
|
||||
"depth_encoder_defaults",
|
||||
# Factories
|
||||
"encoder_config_from_video_info",
|
||||
"camera_encoder_defaults",
|
||||
# Constants
|
||||
"DEFAULT_DEPTH_UNIT",
|
||||
"VALID_VIDEO_CODECS",
|
||||
"VIDEO_ENCODER_INFO_KEYS",
|
||||
]
|
||||
|
||||
@@ -18,7 +18,7 @@ from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from .video import DepthEncoderConfig, RGBEncoderConfig, depth_encoder_defaults, rgb_encoder_defaults
|
||||
from .video import VideoEncoderConfig, camera_encoder_defaults
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -58,10 +58,8 @@ class DatasetRecordConfig:
|
||||
# Set to 1 for immediate encoding (default behavior), or higher for batched encoding
|
||||
video_encoding_batch_size: int = 1
|
||||
# Video encoder settings for camera MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys,
|
||||
# e.g. ``--dataset.rgb_encoder.vcodec=h264`` (see ``RGBEncoderConfig``).
|
||||
rgb_encoder: RGBEncoderConfig = field(default_factory=rgb_encoder_defaults)
|
||||
# Video encoder settings for depth-map MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys.
|
||||
depth_encoder: DepthEncoderConfig = field(default_factory=depth_encoder_defaults)
|
||||
# e.g. ``--dataset.camera_encoder.vcodec=h264`` (see ``VideoEncoderConfig``).
|
||||
camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults)
|
||||
# Enable streaming video encoding: encode frames in real-time during capture instead
|
||||
# of writing PNG images first. Makes save_episode() near-instant. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding
|
||||
streaming_encoding: bool = False
|
||||
|
||||
@@ -19,8 +19,6 @@ from dataclasses import dataclass, field
|
||||
from lerobot.transforms import ImageTransformsConfig
|
||||
from lerobot.utils.import_utils import get_safe_default_video_backend
|
||||
|
||||
from .video import DEFAULT_DEPTH_UNIT, DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT
|
||||
|
||||
|
||||
@dataclass
|
||||
class DatasetConfig:
|
||||
@@ -37,21 +35,14 @@ class DatasetConfig:
|
||||
revision: str | None = None
|
||||
use_imagenet_stats: bool = True
|
||||
video_backend: str = field(default_factory=get_safe_default_video_backend)
|
||||
# When True, RGB video frames are returned as uint8 tensors (0-255) instead of float32 (0.0-1.0).
|
||||
# When True, video frames are returned as uint8 tensors (0-255) instead of float32 (0.0-1.0).
|
||||
# This reduces memory and speeds up DataLoader IPC. The training pipeline handles the conversion.
|
||||
return_uint8: bool = False
|
||||
# Physical unit depth maps are dequantized to at load time: "mm" (millimeters) or "m" (metres).
|
||||
# Has no effect on datasets without depth cameras.
|
||||
depth_output_unit: str = DEFAULT_DEPTH_UNIT
|
||||
streaming: bool = False
|
||||
# Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
|
||||
eval_split: float = 0.0
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.depth_output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT):
|
||||
raise ValueError(
|
||||
f"depth_output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {self.depth_output_unit!r}"
|
||||
)
|
||||
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.episodes is not None:
|
||||
@@ -145,35 +136,3 @@ class PeftConfig:
|
||||
# If None, the PEFT library defaults to alpha=8, which may dampen high-rank adapters.
|
||||
# Common values are r (alpha == rank) or 2*r.
|
||||
lora_alpha: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobConfig:
|
||||
# Where training runs. None (omitted) or "local" runs on this machine.
|
||||
# Any other value is an HF Jobs flavor and submits the run to HF Jobs.
|
||||
# List available flavors + pricing with `hf jobs hardware` command.
|
||||
target: str | None = None
|
||||
# Runtime image for the remote job (ignored for local runs).
|
||||
image: str = "huggingface/lerobot-gpu:latest"
|
||||
# Max wall-clock for the remote job as an HF Jobs duration string (e.g. "2h").
|
||||
# Defaults to "2d": We pass an explicit, generous cap instead. Set a smaller
|
||||
# value to fail fast, or a larger one for long runs.
|
||||
timeout: str | None = "2d"
|
||||
# Submit and exit instead of streaming the job logs in the foreground.
|
||||
detach: bool = False
|
||||
# Extra tags attached to the HF job and to any dataset this run pushes to the
|
||||
# Hub. A "lerobot" tag is always added; e.g. --job.tags '["lelab"]' adds more.
|
||||
tags: list[str] = field(default_factory=list)
|
||||
|
||||
# Two entry points to the same predicate: the staticmethod tests a raw target string
|
||||
# straight from argv (before any JobConfig exists, to decide dispatch early), while the
|
||||
# property is the ergonomic accessor for code that already holds a config instance.
|
||||
@staticmethod
|
||||
def is_remote_target(target: str | None) -> bool:
|
||||
"""True when `target` names an HF Jobs flavor rather than a local run."""
|
||||
return target not in (None, "local")
|
||||
|
||||
@property
|
||||
def is_remote(self) -> bool:
|
||||
"""True when training should run on HF Jobs rather than this machine."""
|
||||
return self.is_remote_target(self.target)
|
||||
|
||||
@@ -26,12 +26,11 @@ from huggingface_hub.errors import HfHubHTTPError
|
||||
|
||||
from lerobot import envs
|
||||
from lerobot.optim import LRSchedulerConfig, OptimizerConfig
|
||||
from lerobot.utils.constants import PRETRAINED_MODEL_DIR
|
||||
from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint
|
||||
from lerobot.utils.hub import HubMixin
|
||||
from lerobot.utils.sample_weighting import SampleWeightingConfig
|
||||
|
||||
from . import parser
|
||||
from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
|
||||
from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
|
||||
from .policies import PreTrainedConfig
|
||||
from .rewards import RewardModelConfig
|
||||
|
||||
@@ -84,11 +83,10 @@ class TrainPipelineConfig(HubMixin):
|
||||
# with the same value for `dir` its contents will be overwritten unless you set `resume` to true.
|
||||
output_dir: Path | None = None
|
||||
job_name: str | None = None
|
||||
# Set `resume` to true to resume a previous run. Pass `--config_path` pointing at either a local
|
||||
# checkpoint's train_config.json or a Hub repo id holding `checkpoints/<step>/` subtrees (the
|
||||
# latest checkpoint is downloaded and resumed from). Note that when resuming, the default behavior
|
||||
# is to use the configuration from the checkpoint, regardless of what's provided with the training
|
||||
# command at the time of resumption (CLI `--*` flags still override).
|
||||
# Set `resume` to true to resume a previous run. In order for this to work, you will need to make sure
|
||||
# `dir` is the directory of an existing run with at least one checkpoint in it.
|
||||
# Note that when resuming a run, the default behavior is to use the configuration from the checkpoint,
|
||||
# regardless of what's provided with the training command at the time of resumption.
|
||||
resume: bool = False
|
||||
# `seed` is used for training (eg: model initialization, dataset shuffling)
|
||||
# AND for the evaluation environments.
|
||||
@@ -120,13 +118,6 @@ class TrainPipelineConfig(HubMixin):
|
||||
wandb: WandBConfig = field(default_factory=WandBConfig)
|
||||
peft: PeftConfig | None = None
|
||||
|
||||
# Where to run training (local default, or an HF Jobs flavor). See JobConfig.
|
||||
job: JobConfig = field(default_factory=JobConfig)
|
||||
# Push each saved checkpoint to the Hub (policy.repo_id) as it is written, not
|
||||
# just the final model (useful to monitor progress mid-run). Optional; the
|
||||
# final model is pushed regardless. Works the same locally and remotely.
|
||||
save_checkpoint_to_hub: bool = False
|
||||
|
||||
# Sample weighting configuration (e.g., for RA-BC training)
|
||||
sample_weighting: SampleWeightingConfig | None = None
|
||||
|
||||
@@ -146,17 +137,10 @@ class TrainPipelineConfig(HubMixin):
|
||||
return self.reward_model # type: ignore[return-value]
|
||||
return self.policy # type: ignore[return-value]
|
||||
|
||||
def _resolve_pretrained_from_cli(self) -> None:
|
||||
"""Resolve the pretrained source passed on the CLI into a loaded config.
|
||||
|
||||
The pretrained paths (`--policy.path`, `--reward_model.path`) and
|
||||
`--config_path` are only recoverable by re-reading the CLI args: draccus
|
||||
has already consumed them by the time `validate()` runs, so they are not
|
||||
reflected on `self`. Exactly one source applies, in priority order:
|
||||
reward-model path, policy path, then resume.
|
||||
"""
|
||||
reward_model_path = parser.get_path_arg("reward_model")
|
||||
def validate(self) -> None:
|
||||
# HACK: We parse again the cli args here to get the pretrained paths if there was some.
|
||||
policy_path = parser.get_path_arg("policy")
|
||||
reward_model_path = parser.get_path_arg("reward_model")
|
||||
|
||||
if reward_model_path:
|
||||
cli_overrides = parser.get_cli_overrides("reward_model")
|
||||
@@ -165,54 +149,31 @@ class TrainPipelineConfig(HubMixin):
|
||||
)
|
||||
self.reward_model.pretrained_path = str(Path(reward_model_path))
|
||||
elif policy_path:
|
||||
overrides = parser.get_yaml_overrides("policy") + (parser.get_cli_overrides("policy") or [])
|
||||
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=overrides)
|
||||
yaml_overrides = parser.get_yaml_overrides("policy")
|
||||
cli_overrides = parser.get_cli_overrides("policy") or []
|
||||
self.policy = PreTrainedConfig.from_pretrained(
|
||||
policy_path, cli_overrides=yaml_overrides + cli_overrides
|
||||
)
|
||||
self.policy.pretrained_path = Path(policy_path)
|
||||
elif self.resume:
|
||||
self._resolve_resume_checkpoint()
|
||||
config_path = parser.parse_arg("config_path")
|
||||
if not config_path:
|
||||
raise ValueError(
|
||||
f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}"
|
||||
)
|
||||
|
||||
def _resolve_resume_checkpoint(self) -> None:
|
||||
"""Point the trainable config at the checkpoint named by `--config_path`.
|
||||
if not Path(config_path).resolve().exists():
|
||||
raise NotADirectoryError(
|
||||
f"{config_path=} is expected to be a local path. "
|
||||
"Resuming from the hub is not supported for now."
|
||||
)
|
||||
|
||||
`config_path` is either a local path (to a checkpoint's train_config.json or its
|
||||
pretrained_model/ dir) or a Hub repo id. For a Hub repo, the latest checkpoint is downloaded
|
||||
into a fresh local run dir and resumed from there. The download is skipped when dispatching to
|
||||
an HF Job (`job.is_remote`): the pod performs it when it runs the resume locally, and
|
||||
`submit_to_hf` resolves the source repo for the remote command.
|
||||
"""
|
||||
config_path = parser.parse_arg("config_path")
|
||||
if not config_path:
|
||||
raise ValueError(
|
||||
f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}"
|
||||
)
|
||||
|
||||
if Path(config_path).resolve().exists():
|
||||
policy_dir = Path(config_path).parent
|
||||
if self.policy is not None:
|
||||
self.policy.pretrained_path = policy_dir
|
||||
if self.reward_model is not None:
|
||||
self.reward_model.pretrained_path = str(policy_dir)
|
||||
self.checkpoint_path = policy_dir.parent
|
||||
elif self.job.is_remote:
|
||||
return
|
||||
else:
|
||||
from lerobot.common.train_utils import resolve_resume_checkpoint
|
||||
|
||||
# `self.output_dir` was loaded from the checkpoint's config and points at the original
|
||||
# run's (now-absent) local dir. Resume into a fresh local dir instead, unless the user
|
||||
# passed --output_dir explicitly.
|
||||
cli_output_dir = parser.parse_arg("output_dir")
|
||||
if cli_output_dir:
|
||||
self.output_dir = Path(cli_output_dir)
|
||||
else:
|
||||
now = dt.datetime.now()
|
||||
self.output_dir = Path("outputs/train") / f"{now:%Y-%m-%d}/{now:%H-%M-%S}_resume"
|
||||
self.checkpoint_path = resolve_resume_checkpoint(config_path, self.output_dir)
|
||||
policy_dir = self.checkpoint_path / PRETRAINED_MODEL_DIR
|
||||
|
||||
if self.policy is not None:
|
||||
self.policy.pretrained_path = policy_dir
|
||||
if self.reward_model is not None:
|
||||
self.reward_model.pretrained_path = str(policy_dir)
|
||||
|
||||
def validate(self) -> None:
|
||||
self._resolve_pretrained_from_cli()
|
||||
|
||||
if self.policy is None and self.reward_model is None:
|
||||
raise ValueError(
|
||||
@@ -255,19 +216,9 @@ class TrainPipelineConfig(HubMixin):
|
||||
if self.eval_steps > 0 and self.dataset.eval_split == 0.0:
|
||||
raise ValueError("eval_steps > 0 requires dataset.eval_split > 0.0 to hold out eval data.")
|
||||
|
||||
# Remote runs auto-generate the repo_id in submit_to_hf (the policy may only be
|
||||
# resolved here, from --policy.path), so don't demand it up front for them.
|
||||
if (
|
||||
hasattr(active_cfg, "push_to_hub")
|
||||
and active_cfg.push_to_hub
|
||||
and not active_cfg.repo_id
|
||||
and not self.job.is_remote
|
||||
):
|
||||
if hasattr(active_cfg, "push_to_hub") and active_cfg.push_to_hub and not active_cfg.repo_id:
|
||||
raise ValueError("'repo_id' argument missing. Please specify it to push the model to the hub.")
|
||||
|
||||
if self.save_checkpoint_to_hub and not (self.policy is not None and self.policy.repo_id):
|
||||
raise ValueError("save_checkpoint_to_hub requires --policy.repo_id.")
|
||||
|
||||
@classmethod
|
||||
def __get_path_fields__(cls) -> list[str]:
|
||||
"""Keys for draccus pretrained-path loading."""
|
||||
@@ -304,30 +255,22 @@ class TrainPipelineConfig(HubMixin):
|
||||
elif Path(model_id).is_file():
|
||||
config_file = model_id
|
||||
else:
|
||||
dl_kwargs = {
|
||||
"repo_id": model_id,
|
||||
"revision": revision,
|
||||
"cache_dir": cache_dir,
|
||||
"force_download": force_download,
|
||||
"proxies": proxies,
|
||||
"resume_download": resume_download,
|
||||
"token": token,
|
||||
"local_files_only": local_files_only,
|
||||
}
|
||||
try:
|
||||
config_file = hf_hub_download(filename=TRAIN_CONFIG_NAME, **dl_kwargs)
|
||||
except HfHubHTTPError as e:
|
||||
# No root train_config.json: this is a repo of periodic checkpoints from an
|
||||
# interrupted run. Fall back to the latest checkpoint's config so the run can be
|
||||
# resumed straight from the repo with `--config_path=<repo>`.
|
||||
latest = find_latest_hub_checkpoint(model_id, token=token, revision=revision)
|
||||
if latest is None:
|
||||
raise FileNotFoundError(
|
||||
f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
|
||||
) from e
|
||||
config_file = hf_hub_download(
|
||||
filename=f"{latest}/{PRETRAINED_MODEL_DIR}/{TRAIN_CONFIG_NAME}", **dl_kwargs
|
||||
repo_id=model_id,
|
||||
filename=TRAIN_CONFIG_NAME,
|
||||
revision=revision,
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
proxies=proxies,
|
||||
resume_download=resume_download,
|
||||
token=token,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
except HfHubHTTPError as e:
|
||||
raise FileNotFoundError(
|
||||
f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
|
||||
) from e
|
||||
|
||||
cli_args = kwargs.pop("cli_args", [])
|
||||
# Legacy RA-BC migration only applies to framework-saved checkpoints (always JSON).
|
||||
|
||||
+36
-123
@@ -20,7 +20,7 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, ClassVar, Self
|
||||
from typing import Any
|
||||
|
||||
from lerobot.utils.import_utils import require_package
|
||||
|
||||
@@ -40,6 +40,7 @@ VALID_VIDEO_CODECS: frozenset[str] = frozenset({"h264", "hevc", "libsvtav1", "au
|
||||
# Aliases for legacy video codec names.
|
||||
VIDEO_CODECS_ALIASES: dict[str, str] = {"av1": "libsvtav1"}
|
||||
|
||||
|
||||
LIBSVTAV1_DEFAULT_PRESET: int = 12
|
||||
|
||||
# Keys persisted under ``features[*]["info"]`` as ``video.<name>`` (from :class:`VideoEncoderConfig`).
|
||||
@@ -51,45 +52,40 @@ VIDEO_ENCODER_INFO_KEYS: frozenset[str] = frozenset(
|
||||
f"video.{name}" for name in VIDEO_ENCODER_INFO_FIELD_NAMES
|
||||
)
|
||||
|
||||
# Default depth quantization and encoding parameters.
|
||||
DEPTH_QUANT_BITS: int = 12
|
||||
DEPTH_QMAX: int = (1 << DEPTH_QUANT_BITS) - 1 # 4095
|
||||
|
||||
DEFAULT_DEPTH_MIN: float = 0.01
|
||||
DEFAULT_DEPTH_MAX: float = 10.0
|
||||
DEFAULT_DEPTH_SHIFT: float = 3.5
|
||||
DEFAULT_DEPTH_USE_LOG: bool = True
|
||||
DEFAULT_DEPTH_PIX_FMT: str = "gray12le"
|
||||
|
||||
DEPTH_METER_UNIT: str = "m"
|
||||
DEPTH_MILLIMETER_UNIT: str = "mm"
|
||||
DEFAULT_DEPTH_UNIT: str = DEPTH_MILLIMETER_UNIT
|
||||
|
||||
# Depth-specific tuning fields persisted under ``features[*]["info"]`` as ``video.<name>``.
|
||||
DEPTH_ENCODER_INFO_FIELD_NAMES: frozenset[str] = frozenset({"depth_min", "depth_max", "shift", "use_log"})
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoEncoderConfig:
|
||||
"""Video encoder configuration."""
|
||||
"""Video encoder configuration.
|
||||
|
||||
vcodec: str = "libsvtav1" # Video codec name. "auto" picks a hardware codec if available, else libsvtav1.
|
||||
pix_fmt: str = "yuv420p" # Pixel format (e.g. yuv420p).
|
||||
g: int | None = 2 # GOP size (keyframe interval).
|
||||
crf: int | float | None = 30 # Quality level. Lower means better quality and larger files.
|
||||
preset: int | str | None = None # Speed/quality preset. Accepted values are codec-specific.
|
||||
fast_decode: int = 0 # Fast-decode tuning. Accepted values are codec-specific, 0 disables it.
|
||||
Attributes:
|
||||
vcodec: Video encoder name. ``"auto"`` is resolved during
|
||||
construction (HW encoder if available, else ``libsvtav1``).
|
||||
pix_fmt: Pixel format (e.g. ``"yuv420p"``).
|
||||
g: GOP size (keyframe interval).
|
||||
crf: Quality level — mapped to the native quality parameter of the
|
||||
codec (``crf`` for software, ``qp`` for NVENC/VAAPI,
|
||||
``q:v`` for VideoToolbox, ``global_quality`` for QSV).
|
||||
preset: Speed/quality preset. Accepted type is per-codec.
|
||||
fast_decode: Fast-decode tuning. For ``libsvtav1`` this is a level (0-2)
|
||||
embedded in ``svtav1-params``. For ``h264`` and ``hevc`` non-zero values
|
||||
set ``tune=fastdecode``. Ignored for other codecs.
|
||||
video_backend: Python to be used for encoding. Only ``"pyav"``
|
||||
is currently supported.
|
||||
extra_options: Free-form dictionary of additional video encoder options
|
||||
(e.g. ``{"tune": "film", "profile:v": "high", "bf": 2}``).
|
||||
"""
|
||||
|
||||
vcodec: str = "libsvtav1" # TODO(CarolinePascal): rename to codec ?
|
||||
pix_fmt: str = "yuv420p"
|
||||
g: int | None = 2
|
||||
crf: int | float | None = 30
|
||||
preset: int | str | None = None
|
||||
fast_decode: int = 0
|
||||
# TODO(CarolinePascal): add torchcodec support + find a way to unify the
|
||||
# two backends (encoding and decoding).
|
||||
video_backend: str = "pyav" # Encoding backend. Only "pyav" is currently supported.
|
||||
# Extra codec options merged last, e.g. {"tune": "film"}.
|
||||
video_backend: str = "pyav"
|
||||
extra_options: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# Source-data channel count this encoder is expected to handle. ``None``
|
||||
# disables the pix_fmt channel-count check; concrete subclasses set it
|
||||
# (3 for RGB, 1 for depth, etc.).
|
||||
_DEFAULT_CHANNELS: ClassVar[int | None] = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.resolve_vcodec()
|
||||
# Empty-constructor ergonomics: ``VideoEncoderConfig()`` must "just work".
|
||||
@@ -98,9 +94,9 @@ class VideoEncoderConfig:
|
||||
self.validate()
|
||||
|
||||
@classmethod
|
||||
def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]:
|
||||
"""Parse the ``video.*`` keys of a feature ``info`` block into
|
||||
constructor kwargs.
|
||||
def from_video_info(cls, video_info: dict | None) -> VideoEncoderConfig:
|
||||
"""Reconstruct a :class:`VideoEncoderConfig` from a video feature's ``info`` block.
|
||||
Missing or ``None`` values fall back to the class defaults.
|
||||
"""
|
||||
video_info = video_info or {}
|
||||
kwargs: dict[str, Any] = {}
|
||||
@@ -119,15 +115,7 @@ class VideoEncoderConfig:
|
||||
continue
|
||||
kwargs[field_name] = value
|
||||
|
||||
return kwargs
|
||||
|
||||
@classmethod
|
||||
def from_video_info(cls, video_info: dict | None) -> Self:
|
||||
"""Reconstruct an encoder config from a video feature's ``info`` block.
|
||||
|
||||
Missing or ``None`` values fall back to the class defaults.
|
||||
"""
|
||||
return cls(**cls._kwargs_from_video_info(video_info))
|
||||
return cls(**kwargs)
|
||||
|
||||
def detect_available_encoders(self, encoders: list[str] | str) -> list[str]:
|
||||
"""Return the subset of available encoders based on the specified video backend.
|
||||
@@ -150,9 +138,7 @@ class VideoEncoderConfig:
|
||||
require_package("av", extra="dataset")
|
||||
from lerobot.datasets import check_video_encoder_parameters_pyav
|
||||
|
||||
check_video_encoder_parameters_pyav(
|
||||
self.vcodec, self.pix_fmt, self.get_codec_options(), channels=self._DEFAULT_CHANNELS
|
||||
)
|
||||
check_video_encoder_parameters_pyav(self.vcodec, self.pix_fmt, self.get_codec_options())
|
||||
|
||||
def resolve_vcodec(self) -> None:
|
||||
"""Check ``vcodec`` and, when it is ``"auto"``, pick a concrete encoder.
|
||||
@@ -244,79 +230,6 @@ class VideoEncoderConfig:
|
||||
return opts
|
||||
|
||||
|
||||
@dataclass
|
||||
class RGBEncoderConfig(VideoEncoderConfig):
|
||||
"""Encoder configuration for RGB camera streams.
|
||||
|
||||
Identical to :class:`VideoEncoderConfig` but declares the 3-channel
|
||||
source-data layout so ``pix_fmt`` is validated against RGB inputs.
|
||||
"""
|
||||
|
||||
_DEFAULT_CHANNELS: ClassVar[int] = 3
|
||||
|
||||
|
||||
def rgb_encoder_defaults() -> RGBEncoderConfig:
|
||||
"""Return a :class:`RGBEncoderConfig` with RGB-camera defaults."""
|
||||
return RGBEncoderConfig()
|
||||
|
||||
|
||||
@dataclass
|
||||
class DepthEncoderConfig(VideoEncoderConfig):
|
||||
"""Encoder configuration for depth-map streams.
|
||||
|
||||
Inherits the full :class:`VideoEncoderConfig` surface (codec, GOP, CRF,
|
||||
preset, ``extra_options``…) and adds the parameters of the depth quantizer.
|
||||
Defaults flip ``vcodec`` to ``"hevc"`` (Main 12 profile) and ``pix_fmt`` to
|
||||
``"gray12le"``.
|
||||
"""
|
||||
|
||||
vcodec: str = "hevc" # Video codec name. Defaults to HEVC Main 12 (a 12-bit-capable codec).
|
||||
pix_fmt: str = "gray12le" # Pixel format. Defaults to 12-bit grayscale.
|
||||
extra_options: dict[str, Any] = field(default_factory=lambda: {"x265-params": "lossless=1"})
|
||||
|
||||
depth_min: float = DEFAULT_DEPTH_MIN # Minimum depth in meters, mapped to the lowest quantum.
|
||||
depth_max: float = DEFAULT_DEPTH_MAX # Maximum depth in meters, mapped to the highest quantum.
|
||||
shift: float = DEFAULT_DEPTH_SHIFT # Pre-log offset in meters for numerical stability near zero.
|
||||
use_log: bool = DEFAULT_DEPTH_USE_LOG # Use logarithmic quantization (True) or linear (False).
|
||||
|
||||
_DEFAULT_CHANNELS: ClassVar[int] = 1
|
||||
|
||||
@classmethod
|
||||
def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]:
|
||||
"""Layer the depth-specific tuning (``depth_min`` / ``depth_max`` /
|
||||
``shift`` / ``use_log``) on top of the base parser. Missing keys
|
||||
fall back to the class defaults.
|
||||
"""
|
||||
kwargs = super()._kwargs_from_video_info(video_info)
|
||||
video_info = video_info or {}
|
||||
for name in DEPTH_ENCODER_INFO_FIELD_NAMES:
|
||||
value = video_info.get(f"video.{name}")
|
||||
if value is not None:
|
||||
kwargs[name] = value
|
||||
return kwargs
|
||||
|
||||
|
||||
def depth_encoder_defaults() -> DepthEncoderConfig:
|
||||
"""Return a :class:`DepthEncoderConfig` with depth-camera defaults."""
|
||||
return DepthEncoderConfig()
|
||||
|
||||
|
||||
def encoder_config_from_video_info(video_info: dict | None) -> VideoEncoderConfig:
|
||||
"""Build the appropriate encoder config from a feature's ``info`` block.
|
||||
|
||||
Dispatches to :class:`DepthEncoderConfig` when the dict marks the feature
|
||||
as a depth map and to :class:`RGBEncoderConfig`
|
||||
otherwise.
|
||||
|
||||
Args:
|
||||
video_info: A feature's ``info`` dict as persisted in ``info.json``,
|
||||
or ``None`` (treated as an empty dict).
|
||||
|
||||
Returns:
|
||||
A :class:`DepthEncoderConfig` for depth features, otherwise a
|
||||
:class:`RGBEncoderConfig`.
|
||||
"""
|
||||
video_info = video_info or {}
|
||||
is_depth = bool(video_info.get("is_depth_map") or video_info.get("video.is_depth_map"))
|
||||
cls: type[VideoEncoderConfig] = DepthEncoderConfig if is_depth else RGBEncoderConfig
|
||||
return cls.from_video_info(video_info)
|
||||
def camera_encoder_defaults() -> VideoEncoderConfig:
|
||||
"""Return a :class:`VideoEncoderConfig` with RGB-camera defaults."""
|
||||
return VideoEncoderConfig()
|
||||
|
||||
@@ -22,7 +22,6 @@ import numpy as np
|
||||
from lerobot.processor import RelativeActionsProcessorStep
|
||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||
|
||||
from .depth_utils import MM_PER_METRE
|
||||
from .io_utils import load_image_as_numpy
|
||||
|
||||
DEFAULT_QUANTILES = [0.01, 0.10, 0.50, 0.90, 0.99]
|
||||
@@ -243,12 +242,12 @@ def sample_images(image_paths: list[str]) -> np.ndarray:
|
||||
images = None
|
||||
for i, idx in enumerate(sampled_indices):
|
||||
path = image_paths[idx]
|
||||
# we load RGB images as uint8 to reduce memory usage; depth keeps its native dtype
|
||||
# we load as uint8 to reduce memory usage
|
||||
img = load_image_as_numpy(path, dtype=np.uint8, channel_first=True)
|
||||
img = auto_downsample_height_width(img)
|
||||
|
||||
if images is None:
|
||||
images = np.empty((len(sampled_indices), *img.shape), dtype=img.dtype)
|
||||
images = np.empty((len(sampled_indices), *img.shape), dtype=np.uint8)
|
||||
|
||||
images[i] = img
|
||||
|
||||
@@ -507,10 +506,8 @@ def compute_episode_stats(
|
||||
Each statistics dictionary contains min, max, mean, std, count, and quantiles.
|
||||
|
||||
Note:
|
||||
For 'image'/'video' features, stats are computed per channel and kept with a
|
||||
leading channel axis (e.g. shape (3, 1, 1) for RGB). RGB stats are divided by
|
||||
255 to land in [0, 1]; depth maps (features flagged with ``is_depth_map``) are
|
||||
instead canonicalized to millimetres regardless of the raw frame unit.
|
||||
Image statistics are normalized to [0,1] range and have shape (3,1,1) for
|
||||
per-channel values when dtype is 'image' or 'video'.
|
||||
"""
|
||||
if quantile_list is None:
|
||||
quantile_list = DEFAULT_QUANTILES
|
||||
@@ -534,17 +531,8 @@ def compute_episode_stats(
|
||||
)
|
||||
|
||||
if features[key]["dtype"] in ["image", "video"]:
|
||||
if (features[key].get("info") or {}).get("is_depth_map", False):
|
||||
# Depth stats are canonically stored in millimetres; metre (float) depth is
|
||||
# scaled up, integer (millimetre) depth is left as-is.
|
||||
normalization_factor = (
|
||||
1.0 / MM_PER_METRE if np.issubdtype(ep_ft_array.dtype, np.floating) else 1.0
|
||||
)
|
||||
else:
|
||||
normalization_factor = 255.0
|
||||
ep_stats[key] = {
|
||||
k: v if k == "count" else np.squeeze(v / normalization_factor, axis=0)
|
||||
for k, v in ep_stats[key].items()
|
||||
k: v if k == "count" else np.squeeze(v / 255.0, axis=0) for k, v in ep_stats[key].items()
|
||||
}
|
||||
|
||||
return ep_stats
|
||||
@@ -564,10 +552,8 @@ def _validate_stat_value(value: np.ndarray, key: str, feature_key: str) -> None:
|
||||
if key == "count" and value.shape != (1,):
|
||||
raise ValueError(f"Shape of 'count' must be (1), but is {value.shape} instead.")
|
||||
|
||||
if "image" in feature_key and key != "count" and value.shape not in ((3, 1, 1), (1, 1, 1)):
|
||||
raise ValueError(
|
||||
f"Shape of quantile '{key}' must be (3,1,1) or (1,1,1) but is {value.shape} instead."
|
||||
)
|
||||
if "image" in feature_key and key != "count" and value.shape != (3, 1, 1):
|
||||
raise ValueError(f"Shape of quantile '{key}' must be (3,1,1), but is {value.shape} instead.")
|
||||
|
||||
|
||||
def _assert_type_and_shape(stats_list: list[dict[str, dict]]):
|
||||
|
||||
@@ -14,8 +14,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import contextlib
|
||||
import logging
|
||||
from collections.abc import Callable, Iterable
|
||||
from collections.abc import Callable
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
|
||||
@@ -339,25 +338,6 @@ class LeRobotDatasetMetadata:
|
||||
"""Keys to access visual modalities stored as videos."""
|
||||
return [key for key, ft in self.features.items() if ft["dtype"] == "video"]
|
||||
|
||||
@property
|
||||
def depth_keys(self) -> list[str]:
|
||||
"""Keys to access depth-map modalities stored as videos or images.
|
||||
|
||||
A depth key is a feature whose ``info`` dict carries ``"is_depth_map": True``
|
||||
(or the legacy ``"video.is_depth_map"`` inside ``info`` or ``video_info``).
|
||||
"""
|
||||
|
||||
def _is_depth(ft: dict) -> bool:
|
||||
info = ft.get("info") or {}
|
||||
video_info = ft.get("video_info") or {}
|
||||
return (
|
||||
info.get("is_depth_map", False)
|
||||
or info.get("video.is_depth_map", False)
|
||||
or video_info.get("video.is_depth_map", False)
|
||||
)
|
||||
|
||||
return [key for key, ft in self.features.items() if _is_depth(ft)]
|
||||
|
||||
@property
|
||||
def camera_keys(self) -> list[str]:
|
||||
"""Keys to access visual modalities (regardless of their storage method)."""
|
||||
@@ -601,48 +581,29 @@ class LeRobotDatasetMetadata:
|
||||
def update_video_info(
|
||||
self,
|
||||
video_key: str | None = None,
|
||||
video_encoder: VideoEncoderConfig | None = None,
|
||||
preserve_keys: Iterable[str] | None = None,
|
||||
camera_encoder: VideoEncoderConfig | None = None,
|
||||
) -> None:
|
||||
"""Populate or refresh per-feature video info in ``info.json``.
|
||||
"""Populate per-feature video info in ``info.json``.
|
||||
|
||||
Warning: this function writes info from first episode videos, implicitly assuming that all videos have
|
||||
been encoded the same way. Also, this means it assumes the first episode exists.
|
||||
|
||||
Always re-probes the videos and overwrites existing info for every recomputed
|
||||
key. ``preserve_keys`` lists keys whose existing values must be kept (e.g.
|
||||
data-intrinsic entries like ``is_depth_map`` and depth quantization params)
|
||||
instead of being recomputed.
|
||||
|
||||
Args:
|
||||
video_key: If provided, only update this video key. Otherwise update
|
||||
all video keys in the dataset.
|
||||
video_encoder: Encoder configuration used to produce the
|
||||
camera_encoder: Encoder configuration used to produce the
|
||||
videos. When provided, its fields are recorded as
|
||||
``video.<field>`` entries alongside the stream-derived
|
||||
``video.*`` entries (see :func:`get_video_info`).
|
||||
preserve_keys: Keys whose existing values are kept instead of being
|
||||
recomputed. ``None`` (default) recomputes every key.
|
||||
"""
|
||||
if video_key is not None and video_key not in self.video_keys:
|
||||
raise ValueError(f"Video key {video_key} not found in dataset")
|
||||
|
||||
video_keys = [video_key] if video_key is not None else self.video_keys
|
||||
preserve_set = set(preserve_keys or ())
|
||||
for key in video_keys:
|
||||
existing = self.features[key].get("info") or {}
|
||||
video_path = self.root / self.video_path.format(video_key=key, chunk_index=0, file_index=0)
|
||||
new_info = get_video_info(video_path, video_encoder=video_encoder)
|
||||
# Drop preserved keys so the existing values win on merge.
|
||||
new_info = {k: v for k, v in new_info.items() if k not in preserve_set}
|
||||
merged = {**existing, **new_info}
|
||||
# Migrate the legacy depth marker to the canonical key.
|
||||
if "video.is_depth_map" in merged:
|
||||
logging.warning(
|
||||
f"Migrating legacy 'video.is_depth_map' to 'is_depth_map' for feature {key!r}."
|
||||
)
|
||||
merged.setdefault("is_depth_map", merged.pop("video.is_depth_map"))
|
||||
self.info.features[key]["info"] = merged
|
||||
if not self.features[key].get("info", None):
|
||||
video_path = self.root / self.video_path.format(video_key=key, chunk_index=0, file_index=0)
|
||||
self.info.features[key]["info"] = get_video_info(video_path, camera_encoder=camera_encoder)
|
||||
|
||||
def update_chunk_settings(
|
||||
self,
|
||||
|
||||
@@ -22,10 +22,7 @@ from pathlib import Path
|
||||
import datasets
|
||||
import torch
|
||||
|
||||
from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig
|
||||
|
||||
from .dataset_metadata import LeRobotDatasetMetadata
|
||||
from .depth_utils import dequantize_depth
|
||||
from .feature_utils import (
|
||||
check_delta_timestamps,
|
||||
get_delta_indices,
|
||||
@@ -54,7 +51,6 @@ class DatasetReader:
|
||||
delta_timestamps: dict[str, list[float]] | None,
|
||||
image_transforms: Callable | None,
|
||||
return_uint8: bool = False,
|
||||
depth_output_unit: str = DEFAULT_DEPTH_UNIT,
|
||||
):
|
||||
"""Initialize the reader with metadata, filtering, and transform config.
|
||||
|
||||
@@ -72,10 +68,6 @@ class DatasetReader:
|
||||
relative timestamp offsets for temporal context windows.
|
||||
image_transforms: Optional torchvision v2 transform applied to
|
||||
visual features.
|
||||
return_uint8: If True, return RGB video frames as raw uint8 tensors
|
||||
instead of normalized float32.
|
||||
depth_output_unit: Physical unit depth maps are dequantized to
|
||||
(``"m"`` or ``"mm"``). Defaults to ``"mm"``.
|
||||
"""
|
||||
self._meta = meta
|
||||
self.root = root
|
||||
@@ -86,7 +78,6 @@ class DatasetReader:
|
||||
raise TypeError("image_transforms must be callable or None.")
|
||||
self._image_transforms = image_transforms
|
||||
self._return_uint8 = return_uint8
|
||||
self._depth_output_unit = depth_output_unit
|
||||
|
||||
self.hf_dataset: datasets.Dataset | None = None
|
||||
self._absolute_to_relative_idx: dict[int, int] | None = None
|
||||
@@ -97,11 +88,6 @@ class DatasetReader:
|
||||
check_delta_timestamps(delta_timestamps, meta.fps, tolerance_s)
|
||||
self.delta_indices = get_delta_indices(delta_timestamps, meta.fps)
|
||||
|
||||
self._depth_encoder_configs: dict[str, DepthEncoderConfig] = {
|
||||
vid_key: DepthEncoderConfig.from_video_info(self._meta.features[vid_key].get("info"))
|
||||
for vid_key in self._meta.depth_keys
|
||||
}
|
||||
|
||||
def set_image_transforms(self, image_transforms: Callable | None) -> None:
|
||||
"""Replace the transform applied to visual observations."""
|
||||
if image_transforms is not None and not callable(image_transforms):
|
||||
@@ -273,18 +259,7 @@ class DatasetReader:
|
||||
self._tolerance_s,
|
||||
self._video_backend,
|
||||
return_uint8=self._return_uint8,
|
||||
is_depth=vid_key in self._meta.depth_keys,
|
||||
)
|
||||
if vid_key in self._meta.depth_keys:
|
||||
depth_encoder = self._depth_encoder_configs[vid_key]
|
||||
frames = dequantize_depth(
|
||||
frames,
|
||||
depth_min=depth_encoder.depth_min,
|
||||
depth_max=depth_encoder.depth_max,
|
||||
shift=depth_encoder.shift,
|
||||
use_log=depth_encoder.use_log,
|
||||
output_unit=self._depth_output_unit,
|
||||
)
|
||||
return vid_key, frames.squeeze(0)
|
||||
|
||||
items = list(query_timestamps.items())
|
||||
@@ -324,9 +299,8 @@ class DatasetReader:
|
||||
item = {**video_frames, **item}
|
||||
|
||||
if self._image_transforms is not None:
|
||||
for cam in self._meta.camera_keys:
|
||||
if cam in self._meta.depth_keys:
|
||||
continue
|
||||
image_keys = self._meta.camera_keys
|
||||
for cam in image_keys:
|
||||
item[cam] = self._image_transforms(item[cam])
|
||||
|
||||
# Add task as a string
|
||||
|
||||
@@ -37,15 +37,7 @@ import pyarrow.parquet as pq
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from lerobot.configs import (
|
||||
DepthEncoderConfig,
|
||||
RGBEncoderConfig,
|
||||
VideoEncoderConfig,
|
||||
depth_encoder_defaults,
|
||||
encoder_config_from_video_info,
|
||||
rgb_encoder_defaults,
|
||||
)
|
||||
from lerobot.configs.video import DEPTH_ENCODER_INFO_FIELD_NAMES
|
||||
from lerobot.configs import VideoEncoderConfig, camera_encoder_defaults
|
||||
from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE
|
||||
from lerobot.utils.utils import flatten_dict
|
||||
|
||||
@@ -56,7 +48,6 @@ from .compute_stats import (
|
||||
compute_relative_action_stats,
|
||||
)
|
||||
from .dataset_metadata import LeRobotDatasetMetadata
|
||||
from .image_writer import write_image
|
||||
from .io_utils import (
|
||||
get_parquet_file_size_in_mb,
|
||||
load_episodes,
|
||||
@@ -71,13 +62,12 @@ from .utils import (
|
||||
DEFAULT_DATA_FILE_SIZE_IN_MB,
|
||||
DEFAULT_DATA_PATH,
|
||||
DEFAULT_EPISODES_PATH,
|
||||
DEPTH_FILE_PATTERN,
|
||||
IMAGE_FILE_PATTERN,
|
||||
VIDEO_DIR,
|
||||
update_chunk_file_indices,
|
||||
)
|
||||
from .video_utils import (
|
||||
encode_video_frames,
|
||||
get_video_info,
|
||||
reencode_video,
|
||||
)
|
||||
|
||||
@@ -611,7 +601,7 @@ def _keep_episodes_from_video_with_av(
|
||||
output_path: Path,
|
||||
episodes_to_keep: list[tuple[int, int]],
|
||||
fps: float,
|
||||
video_encoder: VideoEncoderConfig,
|
||||
camera_encoder: VideoEncoderConfig,
|
||||
) -> None:
|
||||
"""Keep only specified episodes from a video file using PyAV.
|
||||
|
||||
@@ -625,7 +615,7 @@ def _keep_episodes_from_video_with_av(
|
||||
Ranges are half-open intervals: [start_frame, end_frame), where start_frame
|
||||
is inclusive and end_frame is exclusive.
|
||||
fps: Frame rate of the video.
|
||||
video_encoder: Video encoder settings used to re-encode the kept frames.
|
||||
camera_encoder: Video encoder settings used to re-encode the kept frames.
|
||||
"""
|
||||
from fractions import Fraction
|
||||
|
||||
@@ -650,13 +640,13 @@ def _keep_episodes_from_video_with_av(
|
||||
|
||||
# Convert fps to Fraction for PyAV compatibility.
|
||||
fps_fraction = Fraction(fps).limit_denominator(1000)
|
||||
codec_options = video_encoder.get_codec_options(as_strings=True)
|
||||
v_out = out.add_stream(video_encoder.vcodec, rate=fps_fraction, options=codec_options)
|
||||
codec_options = camera_encoder.get_codec_options(as_strings=True)
|
||||
v_out = out.add_stream(camera_encoder.vcodec, rate=fps_fraction, options=codec_options)
|
||||
|
||||
# PyAV type stubs don't distinguish video streams from audio/subtitle streams.
|
||||
v_out.width = v_in.codec_context.width
|
||||
v_out.height = v_in.codec_context.height
|
||||
v_out.pix_fmt = video_encoder.pix_fmt
|
||||
v_out.pix_fmt = camera_encoder.pix_fmt
|
||||
|
||||
# Set time_base to match the frame rate for proper timestamp handling.
|
||||
v_out.time_base = Fraction(1, int(fps))
|
||||
@@ -743,7 +733,7 @@ def _copy_and_reindex_videos(
|
||||
|
||||
for video_key in src_dataset.meta.video_keys:
|
||||
logging.info(f"Processing videos for {video_key}")
|
||||
video_encoder = encoder_config_from_video_info(
|
||||
camera_encoder = VideoEncoderConfig.from_video_info(
|
||||
src_dataset.meta.info.features.get(video_key, {}).get("info")
|
||||
)
|
||||
|
||||
@@ -827,7 +817,7 @@ def _copy_and_reindex_videos(
|
||||
dst_video_path,
|
||||
episodes_to_keep_ranges,
|
||||
src_dataset.meta.fps,
|
||||
video_encoder,
|
||||
camera_encoder,
|
||||
)
|
||||
|
||||
cumulative_ts = 0.0
|
||||
@@ -884,11 +874,11 @@ def _copy_and_reindex_episodes_metadata(
|
||||
episode_meta.update(video_metadata[new_idx])
|
||||
|
||||
# Extract episode statistics from parquet metadata.
|
||||
# When pandas/pyarrow serializes numpy arrays with shape (C, 1, 1) to parquet,
|
||||
# Note (maractingi): When pandas/pyarrow serializes numpy arrays with shape (3, 1, 1) to parquet,
|
||||
# they are being deserialized as nested object arrays like:
|
||||
# array([array([array([0.])]), array([array([0.])]), array([array([0.])])])
|
||||
# This happens particularly with image/video statistics. We need to detect and flatten
|
||||
# these nested structures back to proper (C, 1, 1) arrays so aggregate_stats can process them.
|
||||
# these nested structures back to proper (3, 1, 1) arrays so aggregate_stats can process them.
|
||||
episode_stats = {}
|
||||
for key in src_episode_full:
|
||||
if key.startswith("stats/"):
|
||||
@@ -904,16 +894,15 @@ def _copy_and_reindex_episodes_metadata(
|
||||
if feature_name in src_dataset.meta.features:
|
||||
feature_dtype = src_dataset.meta.features[feature_name]["dtype"]
|
||||
if feature_dtype in ["image", "video"] and stat_name != "count":
|
||||
# Stats are channel-first (C, 1, 1)
|
||||
if isinstance(value, np.ndarray) and value.dtype == object:
|
||||
flat_values = []
|
||||
for item in value:
|
||||
while isinstance(item, np.ndarray):
|
||||
item = item.flatten()[0]
|
||||
flat_values.append(item)
|
||||
value = np.array(flat_values, dtype=np.float64).reshape(-1, 1, 1)
|
||||
elif isinstance(value, np.ndarray) and value.ndim == 1:
|
||||
value = value.reshape(-1, 1, 1)
|
||||
value = np.array(flat_values, dtype=np.float64).reshape(3, 1, 1)
|
||||
elif isinstance(value, np.ndarray) and value.shape == (3,):
|
||||
value = value.reshape(3, 1, 1)
|
||||
|
||||
episode_stats[feature_name][stat_name] = value
|
||||
|
||||
@@ -1164,15 +1153,15 @@ def _save_episode_images_for_video(
|
||||
# Get all items for this episode
|
||||
episode_dataset = imgs_dataset.select(range(from_idx, to_idx))
|
||||
|
||||
is_depth = img_key in dataset.meta.depth_keys
|
||||
frame_pattern = DEPTH_FILE_PATTERN if is_depth else IMAGE_FILE_PATTERN
|
||||
|
||||
# Define function to save a single image
|
||||
def save_single_image(i_item_tuple):
|
||||
i, item = i_item_tuple
|
||||
write_image(item[img_key], imgs_dir / frame_pattern.format(frame_index=i))
|
||||
img = item[img_key]
|
||||
# Use frame-XXXXXX.png format to match encode_video_frames expectations
|
||||
img.save(str(imgs_dir / f"frame-{i:06d}.png"), quality=100)
|
||||
return i
|
||||
|
||||
# Save images with proper naming convention for encode_video_frames (frame-XXXXXX.png)
|
||||
items = list(enumerate(episode_dataset))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=num_workers) as executor:
|
||||
@@ -1204,14 +1193,13 @@ def _save_batch_episodes_images(
|
||||
hf_dataset = dataset.hf_dataset.with_format(None)
|
||||
imgs_dataset = hf_dataset.select_columns(img_key)
|
||||
|
||||
is_depth = img_key in dataset.meta.depth_keys
|
||||
frame_pattern = DEPTH_FILE_PATTERN if is_depth else IMAGE_FILE_PATTERN
|
||||
|
||||
# Define function to save a single image with global frame index
|
||||
# Defined once outside the loop to avoid repeated closure creation
|
||||
def save_single_image(i_item_tuple, base_frame_idx, img_key_param):
|
||||
i, item = i_item_tuple
|
||||
write_image(item[img_key_param], imgs_dir / frame_pattern.format(frame_index=base_frame_idx + i))
|
||||
img = item[img_key_param]
|
||||
# Use global frame index for naming
|
||||
img.save(str(imgs_dir / f"frame-{base_frame_idx + i:06d}.png"), quality=100)
|
||||
return i
|
||||
|
||||
episode_durations = []
|
||||
@@ -1302,7 +1290,7 @@ def _estimate_frame_size_via_calibration(
|
||||
episode_indices: list[int],
|
||||
temp_dir: Path,
|
||||
fps: int,
|
||||
video_encoder: VideoEncoderConfig,
|
||||
camera_encoder: VideoEncoderConfig,
|
||||
num_calibration_frames: int = 30,
|
||||
) -> float:
|
||||
"""Estimate MB per frame by encoding a small calibration sample.
|
||||
@@ -1316,7 +1304,7 @@ def _estimate_frame_size_via_calibration(
|
||||
episode_indices: List of episode indices being processed.
|
||||
temp_dir: Temporary directory for calibration files.
|
||||
fps: Frames per second for video encoding.
|
||||
video_encoder: Video encoder settings used for calibration encoding.
|
||||
camera_encoder: Video encoder settings used for calibration encoding.
|
||||
num_calibration_frames: Number of frames to use for calibration (default: 30).
|
||||
|
||||
Returns:
|
||||
@@ -1341,11 +1329,10 @@ def _estimate_frame_size_via_calibration(
|
||||
hf_dataset = dataset.hf_dataset.with_format(None)
|
||||
sample_indices = range(from_idx, from_idx + num_frames)
|
||||
|
||||
# Save calibration frames using the suffix/format the encoder expects.
|
||||
is_depth = img_key in dataset.meta.depth_keys
|
||||
frame_pattern = DEPTH_FILE_PATTERN if is_depth else IMAGE_FILE_PATTERN
|
||||
# Save calibration frames
|
||||
for i, idx in enumerate(sample_indices):
|
||||
write_image(hf_dataset[idx][img_key], calibration_dir / frame_pattern.format(frame_index=i))
|
||||
img = hf_dataset[idx][img_key]
|
||||
img.save(str(calibration_dir / f"frame-{i:06d}.png"), quality=100)
|
||||
|
||||
# Encode calibration video
|
||||
calibration_video_path = calibration_dir / "calibration.mp4"
|
||||
@@ -1353,7 +1340,7 @@ def _estimate_frame_size_via_calibration(
|
||||
imgs_dir=calibration_dir,
|
||||
video_path=calibration_video_path,
|
||||
fps=fps,
|
||||
video_encoder=video_encoder,
|
||||
camera_encoder=camera_encoder,
|
||||
overwrite=True,
|
||||
)
|
||||
|
||||
@@ -1626,7 +1613,6 @@ def recompute_stats(
|
||||
raise ValueError(f"No parquet files found in {data_dir}")
|
||||
|
||||
all_episode_stats = []
|
||||
# TODO: enable image and video stats re-computation
|
||||
numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]]
|
||||
|
||||
for parquet_path in tqdm(parquet_files, desc="Computing stats from data files"):
|
||||
@@ -1672,8 +1658,7 @@ def convert_image_to_video_dataset(
|
||||
dataset: LeRobotDataset,
|
||||
output_dir: Path | None = None,
|
||||
repo_id: str | None = None,
|
||||
rgb_encoder: RGBEncoderConfig | None = None,
|
||||
depth_encoder: DepthEncoderConfig | None = None,
|
||||
camera_encoder: VideoEncoderConfig | None = None,
|
||||
episode_indices: list[int] | None = None,
|
||||
num_workers: int = 4,
|
||||
max_episodes_per_batch: int | None = None,
|
||||
@@ -1685,32 +1670,21 @@ def convert_image_to_video_dataset(
|
||||
LeRobot dataset structure with videos stored in chunked MP4 files.
|
||||
|
||||
Args:
|
||||
dataset: The source LeRobot dataset with images.
|
||||
output_dir: Root directory where the converted dataset will be stored. When
|
||||
``None``, defaults to ``$HF_LEROBOT_HOME/repo_id``. Equivalent to
|
||||
``new_root`` in ``EditDatasetConfig``.
|
||||
repo_id: Converted dataset identifier. Equivalent to ``new_repo_id`` in
|
||||
``EditDatasetConfig``.
|
||||
rgb_encoder: Video encoder settings applied to RGB cameras. When ``None``,
|
||||
:func:`~lerobot.configs.video.rgb_encoder_defaults` is used.
|
||||
depth_encoder: Video encoder settings applied to depth-map cameras, including
|
||||
the quantization parameters persisted to the dataset metadata. When
|
||||
``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used.
|
||||
episode_indices: Episode indices to convert. When ``None``, all episodes are
|
||||
converted.
|
||||
num_workers: Number of threads for parallel processing.
|
||||
max_episodes_per_batch: Maximum episodes per video batch, to bound memory use.
|
||||
``None`` means no limit.
|
||||
max_frames_per_batch: Maximum frames per video batch, to bound memory use.
|
||||
``None`` means no limit.
|
||||
dataset: The source LeRobot dataset with images
|
||||
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||
camera_encoder: Video encoder settings
|
||||
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
||||
episode_indices: List of episode indices to convert (None = all episodes)
|
||||
num_workers: Number of threads for parallel processing (default: 4)
|
||||
max_episodes_per_batch: Maximum episodes per video batch to avoid memory issues (None = no limit)
|
||||
max_frames_per_batch: Maximum frames per video batch to avoid memory issues (None = no limit)
|
||||
|
||||
Returns:
|
||||
A new :class:`LeRobotDataset` with images encoded as videos.
|
||||
New LeRobotDataset with images encoded as videos
|
||||
"""
|
||||
if rgb_encoder is None:
|
||||
rgb_encoder = rgb_encoder_defaults()
|
||||
if depth_encoder is None:
|
||||
depth_encoder = depth_encoder_defaults()
|
||||
if camera_encoder is None:
|
||||
camera_encoder = camera_encoder_defaults()
|
||||
|
||||
# Check that it's an image dataset
|
||||
if len(dataset.meta.video_keys) > 0:
|
||||
@@ -1735,7 +1709,10 @@ def convert_image_to_video_dataset(
|
||||
logging.info(
|
||||
f"Converting {len(episode_indices)} episodes with {len(img_keys)} cameras from {dataset.repo_id}"
|
||||
)
|
||||
logging.info(f"RGB video encoder: {rgb_encoder}, depth video encoder: {depth_encoder}")
|
||||
logging.info(
|
||||
f"Video codec: {camera_encoder.vcodec}, pixel format: {camera_encoder.pix_fmt}, "
|
||||
f"GOP: {camera_encoder.g}, CRF: {camera_encoder.crf}"
|
||||
)
|
||||
|
||||
# Create new features dict, converting image features to video features
|
||||
new_features = {}
|
||||
@@ -1797,8 +1774,6 @@ def convert_image_to_video_dataset(
|
||||
episode_lengths = {ep_idx: dataset.meta.episodes["length"][ep_idx] for ep_idx in episode_indices}
|
||||
|
||||
for img_key in tqdm(img_keys, desc="Processing cameras"):
|
||||
target_encoder = depth_encoder if img_key in dataset.meta.depth_keys else rgb_encoder
|
||||
|
||||
# Estimate size per frame by encoding a small calibration sample
|
||||
# This provides accurate compression ratio for the specific codec parameters
|
||||
size_per_frame_mb = _estimate_frame_size_via_calibration(
|
||||
@@ -1807,7 +1782,7 @@ def convert_image_to_video_dataset(
|
||||
episode_indices=episode_indices,
|
||||
temp_dir=temp_dir,
|
||||
fps=fps,
|
||||
video_encoder=target_encoder,
|
||||
camera_encoder=camera_encoder,
|
||||
)
|
||||
|
||||
logging.info(f"Processing camera: {img_key}")
|
||||
@@ -1849,7 +1824,7 @@ def convert_image_to_video_dataset(
|
||||
imgs_dir=imgs_dir,
|
||||
video_path=video_path,
|
||||
fps=fps,
|
||||
video_encoder=target_encoder,
|
||||
camera_encoder=camera_encoder,
|
||||
overwrite=True,
|
||||
)
|
||||
|
||||
@@ -1888,11 +1863,16 @@ def convert_image_to_video_dataset(
|
||||
new_meta.info.total_tasks = dataset.meta.total_tasks
|
||||
new_meta.info.splits = {"train": f"0:{len(episode_indices)}"}
|
||||
|
||||
# Update video info for all image keys (now videos). They are registered as
|
||||
# video features above, so update_video_info populates their (still-empty) info.
|
||||
# Update video info for all image keys (now videos)
|
||||
# We need to manually set video info since update_video_info() checks video_keys first
|
||||
for img_key in img_keys:
|
||||
target_encoder = depth_encoder if img_key in dataset.meta.depth_keys else rgb_encoder
|
||||
new_meta.update_video_info(video_key=img_key, video_encoder=target_encoder)
|
||||
if not new_meta.features[img_key].get("info", None):
|
||||
video_path = new_meta.root / new_meta.video_path.format(
|
||||
video_key=img_key, chunk_index=0, file_index=0
|
||||
)
|
||||
new_meta.info.features[img_key]["info"] = get_video_info(
|
||||
video_path, camera_encoder=camera_encoder
|
||||
)
|
||||
|
||||
write_info(new_meta.info, new_meta.root)
|
||||
|
||||
@@ -1919,11 +1899,11 @@ def convert_image_to_video_dataset(
|
||||
|
||||
def _reencode_video_worker(args: tuple) -> Path:
|
||||
"""Picklable worker for :func:`reencode_dataset`'s process pool."""
|
||||
video_path, video_encoder, encoder_threads = args
|
||||
video_path, camera_encoder, encoder_threads = args
|
||||
reencode_video(
|
||||
input_video_path=video_path,
|
||||
output_video_path=video_path,
|
||||
video_encoder=video_encoder,
|
||||
camera_encoder=camera_encoder,
|
||||
encoder_threads=encoder_threads,
|
||||
overwrite=True,
|
||||
)
|
||||
@@ -1932,8 +1912,7 @@ def _reencode_video_worker(args: tuple) -> Path:
|
||||
|
||||
def reencode_dataset(
|
||||
dataset: LeRobotDataset,
|
||||
rgb_encoder: RGBEncoderConfig | None = None,
|
||||
depth_encoder: DepthEncoderConfig | None = None,
|
||||
camera_encoder: VideoEncoderConfig,
|
||||
encoder_threads: int | None = None,
|
||||
num_workers: int | None = None,
|
||||
) -> LeRobotDataset:
|
||||
@@ -1944,11 +1923,8 @@ def reencode_dataset(
|
||||
Args:
|
||||
dataset: An existing :class:`LeRobotDataset` whose videos will be
|
||||
re-encoded.
|
||||
rgb_encoder: Target encoder configuration applied to every RGB video
|
||||
file. If ``None``, re-encoding is skipped for RGB videos.
|
||||
depth_encoder: Target encoder configuration applied to every depth video
|
||||
file. If ``None``, re-encoding is skipped for depth videos.
|
||||
Quantization parameters will not override the ones in the current dataset.
|
||||
camera_encoder: Target encoder configuration applied to every video
|
||||
file.
|
||||
encoder_threads: Per-encoder thread count forwarded to
|
||||
:func:`reencode_video`. ``None`` lets the codec decide.
|
||||
num_workers: Number of parallel processes. ``None`` or ``0`` means
|
||||
@@ -1960,35 +1936,23 @@ def reencode_dataset(
|
||||
on disk.
|
||||
"""
|
||||
meta = dataset.meta
|
||||
video_keys_encoders_dict = {}
|
||||
video_keys_paths_dict = {}
|
||||
|
||||
if rgb_encoder is None and depth_encoder is None:
|
||||
raise ValueError("Either rgb_encoder or depth_encoder must be provided")
|
||||
video_paths_list = []
|
||||
|
||||
# Only re-encode if the videos are not already encoded with the given video encoding parameters
|
||||
for video_key in meta.video_keys:
|
||||
current_info = meta.info.features[video_key].get("info", {})
|
||||
current_encoder = encoder_config_from_video_info(current_info)
|
||||
target_encoder = depth_encoder if video_key in meta.depth_keys else rgb_encoder
|
||||
if target_encoder is None:
|
||||
logging.info(f"No encoder provided for {video_key} video. Skipping re-encoding.")
|
||||
elif current_encoder != target_encoder:
|
||||
video_keys_paths_dict[video_key] = list((meta.root / VIDEO_DIR / video_key).rglob("*.mp4"))
|
||||
video_keys_encoders_dict[video_key] = target_encoder
|
||||
current_encoder = VideoEncoderConfig.from_video_info(current_info)
|
||||
if current_encoder != camera_encoder:
|
||||
video_paths_list.extend((meta.root / VIDEO_DIR / video_key).rglob("*.mp4"))
|
||||
else:
|
||||
logging.info(f"{video_key} videos are already encoded with {target_encoder}. Nothing to do.")
|
||||
logging.info(f"{video_key} videos are already encoded with {camera_encoder}. Nothing to do.")
|
||||
|
||||
if len(video_keys_paths_dict) == 0:
|
||||
if len(video_paths_list) == 0:
|
||||
logging.warning("Dataset has no videos to re-encode.")
|
||||
return dataset
|
||||
logging.info(f"Re-encoding {sum(len(paths) for paths in video_keys_paths_dict.values())} video file(s).")
|
||||
logging.info(f"Re-encoding {len(video_paths_list)} video file(s) with {camera_encoder}")
|
||||
|
||||
worker_args = [
|
||||
(path, encoder, encoder_threads)
|
||||
for video_key, encoder in video_keys_encoders_dict.items()
|
||||
for path in video_keys_paths_dict[video_key]
|
||||
]
|
||||
worker_args = [(vp, camera_encoder, encoder_threads) for vp in video_paths_list]
|
||||
if num_workers and num_workers > 1:
|
||||
with ProcessPoolExecutor(max_workers=num_workers) as pool:
|
||||
futures = [pool.submit(_reencode_video_worker, args) for args in worker_args]
|
||||
@@ -2002,15 +1966,10 @@ def reencode_dataset(
|
||||
for args in tqdm(worker_args, desc="Re-encoding videos"):
|
||||
_reencode_video_worker(args)
|
||||
|
||||
# Refresh video info in metadata for every re-encoded key. Re-encoding only
|
||||
# changes codec/container params, so for depth videos we preserve ``is_depth_map``
|
||||
# and the depth quantization params (``video.depth_min`` / ``video.depth_max`` /
|
||||
# ...), which describe the data rather than the codec and must survive a transcode.
|
||||
# RGB videos pass an empty set: still a refresh, but nothing to preserve.
|
||||
depth_preserve_keys = {"is_depth_map", *(f"video.{n}" for n in DEPTH_ENCODER_INFO_FIELD_NAMES)}
|
||||
for video_key, encoder in video_keys_encoders_dict.items():
|
||||
preserve_keys = depth_preserve_keys if video_key in meta.depth_keys else set()
|
||||
meta.update_video_info(video_key=video_key, video_encoder=encoder, preserve_keys=preserve_keys)
|
||||
# Refresh video info in metadata for every video key.
|
||||
for vid_key in meta.video_keys:
|
||||
video_path = meta.root / meta.get_video_file_path(0, vid_key)
|
||||
meta.info.features[vid_key]["info"] = get_video_info(video_path, camera_encoder=camera_encoder)
|
||||
|
||||
write_info(meta.info, meta.root)
|
||||
logging.info("Dataset metadata updated.")
|
||||
|
||||
@@ -31,13 +31,7 @@ import PIL.Image
|
||||
import pyarrow.parquet as pq
|
||||
import torch
|
||||
|
||||
from lerobot.configs import (
|
||||
DepthEncoderConfig,
|
||||
RGBEncoderConfig,
|
||||
VideoEncoderConfig,
|
||||
depth_encoder_defaults,
|
||||
rgb_encoder_defaults,
|
||||
)
|
||||
from lerobot.configs import VideoEncoderConfig, camera_encoder_defaults
|
||||
|
||||
from .compute_stats import compute_episode_stats
|
||||
from .dataset_metadata import LeRobotDatasetMetadata
|
||||
@@ -54,7 +48,6 @@ from .io_utils import (
|
||||
write_info,
|
||||
)
|
||||
from .utils import (
|
||||
DEFAULT_DEPTH_PATH,
|
||||
DEFAULT_EPISODES_PATH,
|
||||
DEFAULT_IMAGE_PATH,
|
||||
update_chunk_file_indices,
|
||||
@@ -74,22 +67,17 @@ def _encode_video_worker(
|
||||
episode_index: int,
|
||||
root: Path,
|
||||
fps: int,
|
||||
video_encoder: VideoEncoderConfig | None = None,
|
||||
camera_encoder: VideoEncoderConfig | None = None,
|
||||
encoder_threads: int | None = None,
|
||||
) -> Path:
|
||||
temp_path = Path(tempfile.mkdtemp(dir=root)) / f"{video_key}_{episode_index:03d}.mp4"
|
||||
path_template = (
|
||||
DEFAULT_DEPTH_PATH
|
||||
if video_encoder is not None and isinstance(video_encoder, DepthEncoderConfig)
|
||||
else DEFAULT_IMAGE_PATH
|
||||
)
|
||||
fpath = path_template.format(image_key=video_key, episode_index=episode_index, frame_index=0)
|
||||
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=episode_index, frame_index=0)
|
||||
img_dir = (root / fpath).parent
|
||||
encode_video_frames(
|
||||
img_dir,
|
||||
temp_path,
|
||||
fps,
|
||||
video_encoder=video_encoder,
|
||||
camera_encoder=camera_encoder,
|
||||
encoder_threads=encoder_threads,
|
||||
overwrite=True,
|
||||
)
|
||||
@@ -108,8 +96,7 @@ class DatasetWriter:
|
||||
self,
|
||||
meta: LeRobotDatasetMetadata,
|
||||
root: Path,
|
||||
rgb_encoder: RGBEncoderConfig | None,
|
||||
depth_encoder: DepthEncoderConfig | None,
|
||||
camera_encoder: VideoEncoderConfig | None,
|
||||
encoder_threads: int | None,
|
||||
batch_encoding_size: int,
|
||||
streaming_encoder: StreamingVideoEncoder | None = None,
|
||||
@@ -121,11 +108,8 @@ class DatasetWriter:
|
||||
meta: Dataset metadata instance (used for feature schema, chunk
|
||||
settings, and episode persistence).
|
||||
root: Local dataset root directory.
|
||||
rgb_encoder: Video encoder settings applied to RGB cameras. When
|
||||
``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` is used.
|
||||
depth_encoder: Video encoder settings applied to depth cameras, including
|
||||
the quantization parameters. When ``None``,
|
||||
:func:`~lerobot.configs.video.depth_encoder_defaults` is used.
|
||||
camera_encoder: Video encoder settings applied to all cameras.
|
||||
``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`.
|
||||
encoder_threads: Number of encoder threads (global). ``None``
|
||||
lets the codec decide.
|
||||
batch_encoding_size: Number of episodes to accumulate before
|
||||
@@ -136,8 +120,7 @@ class DatasetWriter:
|
||||
"""
|
||||
self._meta = meta
|
||||
self._root = root
|
||||
self._rgb_encoder = rgb_encoder or rgb_encoder_defaults()
|
||||
self._depth_encoder = depth_encoder or depth_encoder_defaults()
|
||||
self._camera_encoder = camera_encoder or camera_encoder_defaults()
|
||||
self._encoder_threads = encoder_threads
|
||||
self._batch_encoding_size = batch_encoding_size
|
||||
self._streaming_encoder = streaming_encoder
|
||||
@@ -162,8 +145,7 @@ class DatasetWriter:
|
||||
return ep_buffer
|
||||
|
||||
def _get_image_file_path(self, episode_index: int, image_key: str, frame_index: int) -> Path:
|
||||
path_template = DEFAULT_DEPTH_PATH if image_key in self._meta.depth_keys else DEFAULT_IMAGE_PATH
|
||||
fpath = path_template.format(
|
||||
fpath = DEFAULT_IMAGE_PATH.format(
|
||||
image_key=image_key, episode_index=episode_index, frame_index=frame_index
|
||||
)
|
||||
return self._root / fpath
|
||||
@@ -213,7 +195,6 @@ class DatasetWriter:
|
||||
if frame_index == 0 and self._streaming_encoder is not None:
|
||||
self._streaming_encoder.start_episode(
|
||||
video_keys=list(self._meta.video_keys),
|
||||
depth_video_keys=list(self._meta.depth_keys),
|
||||
temp_dir=self._root,
|
||||
)
|
||||
|
||||
@@ -301,13 +282,10 @@ class DatasetWriter:
|
||||
if use_streaming:
|
||||
streaming_results = self._streaming_encoder.finish_episode()
|
||||
for video_key in self._meta.video_keys:
|
||||
normalization_factor = 255.0 if video_key not in self._meta.depth_keys else 1.0
|
||||
temp_path, video_stats = streaming_results[video_key]
|
||||
if video_stats is not None:
|
||||
ep_stats[video_key] = {
|
||||
k: v
|
||||
if k == "count"
|
||||
else np.squeeze(v.reshape(1, -1, 1, 1) / normalization_factor, axis=0)
|
||||
k: v if k == "count" else np.squeeze(v.reshape(1, -1, 1, 1) / 255.0, axis=0)
|
||||
for k, v in video_stats.items()
|
||||
}
|
||||
ep_metadata.update(self._save_episode_video(video_key, episode_index, temp_path=temp_path))
|
||||
@@ -322,7 +300,7 @@ class DatasetWriter:
|
||||
episode_index,
|
||||
self._root,
|
||||
self._meta.fps,
|
||||
self._depth_encoder if video_key in self._meta.depth_keys else self._rgb_encoder,
|
||||
self._camera_encoder,
|
||||
self._encoder_threads,
|
||||
): video_key
|
||||
for video_key in self._meta.video_keys
|
||||
@@ -533,12 +511,7 @@ class DatasetWriter:
|
||||
|
||||
# Update video info (only needed when first episode is encoded)
|
||||
if episode_index == 0:
|
||||
self._meta.update_video_info(
|
||||
video_key,
|
||||
video_encoder=self._depth_encoder
|
||||
if video_key in self._meta.depth_keys
|
||||
else self._rgb_encoder,
|
||||
)
|
||||
self._meta.update_video_info(video_key, camera_encoder=self._camera_encoder)
|
||||
write_info(self._meta.info, self._meta.root)
|
||||
|
||||
metadata = {
|
||||
@@ -605,14 +578,13 @@ class DatasetWriter:
|
||||
self.image_writer.wait_until_done()
|
||||
|
||||
def _encode_temporary_episode_video(self, video_key: str, episode_index: int) -> Path:
|
||||
"""Use ffmpeg to convert frames stored as png/tiff into mp4 videos."""
|
||||
is_depth = video_key in self._meta.depth_keys
|
||||
"""Use ffmpeg to convert frames stored as png into mp4 videos."""
|
||||
return _encode_video_worker(
|
||||
video_key,
|
||||
episode_index,
|
||||
self._root,
|
||||
self._meta.fps,
|
||||
self._depth_encoder if is_depth else self._rgb_encoder,
|
||||
self._camera_encoder,
|
||||
self._encoder_threads,
|
||||
)
|
||||
|
||||
|
||||
@@ -1,268 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
Depth encoding/decoding helpers for :class:`DepthEncoderConfig`.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import Literal
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from lerobot.configs.video import (
|
||||
DEFAULT_DEPTH_MAX,
|
||||
DEFAULT_DEPTH_MIN,
|
||||
DEFAULT_DEPTH_PIX_FMT,
|
||||
DEFAULT_DEPTH_SHIFT,
|
||||
DEFAULT_DEPTH_USE_LOG,
|
||||
DEPTH_METER_UNIT,
|
||||
DEPTH_MILLIMETER_UNIT,
|
||||
DEPTH_QMAX,
|
||||
)
|
||||
|
||||
from .image_writer import squeeze_single_channel
|
||||
from .pyav_utils import write_u16_plane
|
||||
|
||||
MM_PER_METRE = 1000.0
|
||||
_UINT16_MAX = 65535
|
||||
|
||||
|
||||
def _validate_log_quant_params(depth_min: float, shift: float) -> None:
|
||||
"""Ensure ``log(depth_min + shift)`` is finite."""
|
||||
if depth_min + shift <= 0:
|
||||
raise ValueError(
|
||||
f"depth_min + shift must be positive for logarithmic quantization, "
|
||||
f"got depth_min={depth_min} + shift={shift} = {depth_min + shift}"
|
||||
)
|
||||
|
||||
|
||||
def _depth_input_to_float32_and_unit(
|
||||
depth: NDArray[np.integer] | NDArray[np.floating],
|
||||
input_unit: Literal["auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT],
|
||||
) -> tuple[NDArray[np.float32], Literal[DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT]]:
|
||||
"""Convert depth to float32 in the chosen unit, and return the resolved unit."""
|
||||
resolved_unit = (
|
||||
(DEPTH_METER_UNIT if np.issubdtype(depth.dtype, np.floating) else DEPTH_MILLIMETER_UNIT)
|
||||
if input_unit == "auto"
|
||||
else input_unit
|
||||
)
|
||||
return depth.astype(np.float32, order="K"), resolved_unit
|
||||
|
||||
|
||||
def quantize_depth(
|
||||
depth: NDArray[np.uint16] | NDArray[np.float32] | torch.Tensor,
|
||||
depth_min: float = DEFAULT_DEPTH_MIN,
|
||||
depth_max: float = DEFAULT_DEPTH_MAX,
|
||||
shift: float = DEFAULT_DEPTH_SHIFT,
|
||||
use_log: bool = DEFAULT_DEPTH_USE_LOG,
|
||||
pix_fmt: str = DEFAULT_DEPTH_PIX_FMT,
|
||||
video_backend: str | None = "pyav",
|
||||
input_unit: Literal["auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT] = "auto",
|
||||
) -> NDArray[np.uint16] | av.VideoFrame:
|
||||
"""Quantize depth to 12-bit codes (``uint16``, values ``0…DEPTH_QMAX``).
|
||||
|
||||
Depth maps are packed into 12-bit integer frames so they fit in standard
|
||||
high-bit-depth pixel formats (e.g. ``yuv420p12le`` / ``gray12le``)
|
||||
and can be encoded by widely supported video codecs (e.g. HEVC Main 12).
|
||||
Logarithmic quantization is the default because it allocates more quanta
|
||||
to near-range depth, which matches the (1/depth) error profile of typical
|
||||
depth sensors. Math is ported from BEHAVIOR-1K's ``obs_utils.py``.
|
||||
|
||||
**Input units**:
|
||||
|
||||
- ``input_unit="auto"`` (default): infer from dtype (floating = m, non-floating = mm).
|
||||
- ``input_unit="mm"``: interpret input values as millimetres.
|
||||
- ``input_unit="m"``: interpret input values as metres.
|
||||
|
||||
Quantization math runs in the **resolved input unit**.
|
||||
|
||||
``depth_min``, ``depth_max``, and ``shift`` are always in **metres**.
|
||||
|
||||
Args:
|
||||
depth: Depth map; ``torch.Tensor`` is moved to CPU for conversion.
|
||||
depth_min: Depth (metres) at quantum ``0``.
|
||||
depth_max: Depth (metres) at quantum :data:`DEPTH_QMAX`.
|
||||
shift: Depth shift (metres); used in log mode. Must satisfy ``depth_min + shift > 0``.
|
||||
use_log: If ``True`` (default), quantize in log space.
|
||||
video_backend: Video backend to use for encoding. Defaults to "pyav".
|
||||
input_unit: Input unit policy (``"auto"``, ``"mm"``, ``"m"``).
|
||||
|
||||
Returns:
|
||||
``numpy.ndarray``, ``dtype=uint16``, same shape as ``depth``, values in
|
||||
``[0, DEPTH_QMAX]``.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``input_unit`` is not ``"auto"``, ``"mm"``, or ``"m"``.
|
||||
ValueError: If ``use_log=True`` and ``depth_min + shift <= 0``.
|
||||
"""
|
||||
if input_unit not in ("auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT):
|
||||
raise ValueError(
|
||||
f"input_unit must be 'auto', '{DEPTH_METER_UNIT}', or '{DEPTH_MILLIMETER_UNIT}', got {input_unit!r}"
|
||||
)
|
||||
|
||||
if isinstance(depth, torch.Tensor):
|
||||
depth = depth.detach().cpu().numpy()
|
||||
|
||||
# Squeeze single-channel dim: (H, W, 1) or (1, H, W) → (H, W)
|
||||
depth = squeeze_single_channel(depth)
|
||||
|
||||
depth_f, resolved_unit = _depth_input_to_float32_and_unit(depth, input_unit=input_unit)
|
||||
|
||||
# Convert depth_min, depth_max, and shift to the resolved input unit.
|
||||
depth_min_u = (
|
||||
np.float32(depth_min) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_min * MM_PER_METRE)
|
||||
)
|
||||
depth_max_u = (
|
||||
np.float32(depth_max) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_max * MM_PER_METRE)
|
||||
)
|
||||
shift_u = np.float32(shift) if resolved_unit == DEPTH_METER_UNIT else np.float32(shift * MM_PER_METRE)
|
||||
|
||||
# Normalization and quantization is performed in the resolved input unit.
|
||||
if use_log:
|
||||
_validate_log_quant_params(depth_min, shift)
|
||||
log_min = math.log(float(depth_min_u + shift_u))
|
||||
log_max = math.log(float(depth_max_u + shift_u))
|
||||
norm = (np.log(depth_f + shift_u) - log_min) / (log_max - log_min)
|
||||
else:
|
||||
norm = (depth_f - depth_min_u) / (depth_max_u - depth_min_u)
|
||||
|
||||
quantized = np.rint(norm * DEPTH_QMAX).clip(0, DEPTH_QMAX).astype(np.uint16, copy=False)
|
||||
|
||||
if video_backend == "pyav":
|
||||
frame = av.VideoFrame.from_ndarray(quantized, format=pix_fmt)
|
||||
write_u16_plane(frame.planes[0], quantized)
|
||||
return frame
|
||||
else:
|
||||
return quantized
|
||||
|
||||
|
||||
def dequantize_depth(
|
||||
quantized: NDArray[np.uint16] | av.VideoFrame | torch.Tensor,
|
||||
depth_min: float = DEFAULT_DEPTH_MIN,
|
||||
depth_max: float = DEFAULT_DEPTH_MAX,
|
||||
shift: float = DEFAULT_DEPTH_SHIFT,
|
||||
use_log: bool = DEFAULT_DEPTH_USE_LOG,
|
||||
pix_fmt: str = DEFAULT_DEPTH_PIX_FMT,
|
||||
output_unit: Literal[DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT] = DEPTH_MILLIMETER_UNIT,
|
||||
output_tensor: bool = True,
|
||||
output_channel_last: bool = False,
|
||||
) -> NDArray[np.uint16] | NDArray[np.float32] | torch.Tensor:
|
||||
"""Inverse of :func:`quantize_depth`.
|
||||
|
||||
Decoding inverts the same normalized code mapping as :func:`quantize_depth`
|
||||
using ``depth_min`` / ``depth_max`` / ``shift`` (in metres), then returns
|
||||
the requested output unit. Tuning arguments **must match** :func:`quantize_depth`.
|
||||
|
||||
Accepted input layouts :
|
||||
|
||||
- ``(H, W, 1)`` or ``(H, W)`` — single frame with channel-last.
|
||||
- ``(..., 1, H, W)`` — batched frames with channel-first.
|
||||
- ``(..., H, W, 1)`` — batched frames with channel-last.
|
||||
Output layout is determined by ``output_channel_last``.
|
||||
|
||||
Args:
|
||||
quantized: 12-bit codes in ``[0, DEPTH_QMAX]``. ``np.ndarray``,
|
||||
``av.VideoFrame``, or ``torch.Tensor`` (any integer or float dtype).
|
||||
depth_min, depth_max, shift, use_log: Same as :func:`quantize_depth` (metres).
|
||||
pix_fmt: Pixel format used to extract the plane from an ``av.VideoFrame``.
|
||||
output_unit: ``"mm"`` returns ``uint16`` millimetres (rint, clip
|
||||
``[0, 65535]``) when returning a numpy array, or ``float32`` mm when
|
||||
``output_tensor=True``. ``"m"`` returns ``float32`` metres in
|
||||
``[depth_min, depth_max]``.
|
||||
output_tensor: If True, return a ``torch.Tensor`` instead of a numpy array.
|
||||
|
||||
Returns:
|
||||
Depth map in the requested unit and dtype.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``output_unit`` is not ``"m"`` or ``"mm"``.
|
||||
ValueError: If ``use_log=True`` and ``depth_min + shift <= 0``.
|
||||
"""
|
||||
if output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT):
|
||||
raise ValueError(
|
||||
f"output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {output_unit!r}"
|
||||
)
|
||||
if use_log:
|
||||
_validate_log_quant_params(depth_min, shift)
|
||||
|
||||
if isinstance(quantized, av.VideoFrame):
|
||||
quantized = quantized.to_ndarray(format=pix_fmt)
|
||||
|
||||
# Compute the scale and offset first.
|
||||
depth_min_m = float(depth_min)
|
||||
depth_max_m = float(depth_max)
|
||||
shift_m = float(shift)
|
||||
if use_log:
|
||||
log_min = math.log(depth_min_m + shift_m)
|
||||
log_max = math.log(depth_max_m + shift_m)
|
||||
scale = (log_max - log_min) / DEPTH_QMAX
|
||||
offset = log_min
|
||||
else:
|
||||
scale = (depth_max_m - depth_min_m) / DEPTH_QMAX
|
||||
offset = depth_min_m
|
||||
|
||||
# ── Torch path: stay on the input device, single fp32 allocation. ────────
|
||||
if isinstance(quantized, torch.Tensor):
|
||||
if quantized.ndim >= 3:
|
||||
# Drop the single-channel dimension so the math runs on (..., H, W).
|
||||
quantized = quantized.squeeze(-3) if quantized.shape[-3] == 1 else quantized.squeeze(-1)
|
||||
|
||||
# Single allocation we own; everything else is in-place.
|
||||
buf = quantized.to(dtype=torch.float32, copy=True)
|
||||
buf.mul_(scale).add_(offset)
|
||||
if use_log:
|
||||
buf.exp_().sub_(shift_m)
|
||||
buf.clamp_(depth_min_m, depth_max_m)
|
||||
buf.unsqueeze_(-1) if output_channel_last else buf.unsqueeze_(-3)
|
||||
|
||||
if output_unit == DEPTH_METER_UNIT:
|
||||
return buf if output_tensor else buf.cpu().numpy()
|
||||
|
||||
# mm path: round + clamp in float32, skipping the uint16 round-trip
|
||||
# when returning a tensor (torch.uint16 is poorly supported).
|
||||
buf.mul_(MM_PER_METRE).round_().clamp_(0.0, _UINT16_MAX)
|
||||
if output_tensor:
|
||||
return buf
|
||||
return buf.cpu().numpy().astype(np.uint16, copy=False)
|
||||
|
||||
# ── NumPy path: single fp32 allocation, ``out=`` for in-place math. ─────
|
||||
arr = np.asarray(quantized)
|
||||
if arr.ndim >= 3:
|
||||
# Drop the single-channel dimension so the math runs on (..., H, W).
|
||||
arr = np.squeeze(arr, axis=-3) if arr.shape[-3] == 1 else np.squeeze(arr, axis=-1)
|
||||
|
||||
buf = np.empty(arr.shape, dtype=np.float32)
|
||||
np.multiply(arr, scale, out=buf)
|
||||
np.add(buf, offset, out=buf)
|
||||
if use_log:
|
||||
np.exp(buf, out=buf)
|
||||
np.subtract(buf, shift_m, out=buf)
|
||||
np.clip(buf, depth_min_m, depth_max_m, out=buf)
|
||||
buf = np.expand_dims(buf, axis=-1) if output_channel_last else np.expand_dims(buf, axis=-3)
|
||||
|
||||
if output_unit == DEPTH_METER_UNIT:
|
||||
return torch.from_numpy(buf) if output_tensor else buf
|
||||
|
||||
np.multiply(buf, MM_PER_METRE, out=buf)
|
||||
np.rint(buf, out=buf)
|
||||
np.clip(buf, 0.0, _UINT16_MAX, out=buf)
|
||||
if output_tensor:
|
||||
# torch.uint16 support is very limited; return float32 millimetres.
|
||||
return torch.from_numpy(buf)
|
||||
return buf.astype(np.uint16, copy=False)
|
||||
@@ -97,7 +97,6 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
|
||||
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,
|
||||
)
|
||||
else:
|
||||
@@ -128,8 +127,6 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
|
||||
|
||||
if cfg.dataset.use_imagenet_stats:
|
||||
for key in dataset.meta.camera_keys:
|
||||
if key in dataset.meta.depth_keys:
|
||||
continue # Exclude depth keys from ImageNet stats
|
||||
for stats_type, stats in IMAGENET_STATS.items():
|
||||
dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)
|
||||
|
||||
|
||||
@@ -336,7 +336,7 @@ def validate_feature_image_or_video(
|
||||
|
||||
Args:
|
||||
name (str): The name of the feature.
|
||||
expected_shape (list[str]): The expected shape, e.g. (C, H, W) or (H, W, C).
|
||||
expected_shape (list[str]): The expected shape (C, H, W).
|
||||
value: The image data to validate.
|
||||
|
||||
Returns:
|
||||
|
||||
@@ -41,51 +41,11 @@ def safe_stop_image_writer(func):
|
||||
return wrapper
|
||||
|
||||
|
||||
def squeeze_single_channel(array: np.ndarray) -> np.ndarray:
|
||||
"""Drop a leading or trailing singleton channel dim: ``(1, H, W)`` / ``(H, W, 1)`` -> ``(H, W)``.
|
||||
|
||||
Unlike ``array.squeeze()``, this only removes the channel axis, never an ``H`` or ``W`` of size 1.
|
||||
"""
|
||||
if array.ndim == 3:
|
||||
if array.shape[0] == 1:
|
||||
return array[0]
|
||||
if array.shape[-1] == 1:
|
||||
return array[..., 0]
|
||||
return array
|
||||
|
||||
|
||||
def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True) -> PIL.Image.Image:
|
||||
"""Convert a NumPy array to a PIL Image, preserving precision for grayscale.
|
||||
# TODO(aliberts): handle 1 channel and 4 for depth images
|
||||
if image_array.ndim != 3:
|
||||
raise ValueError(f"The array has {image_array.ndim} dimensions, but 3 is expected for an image.")
|
||||
|
||||
Behaviour by shape:
|
||||
|
||||
- ``(H, W)`` or ``(1, H, W)`` / ``(H, W, 1)``: single-channel grayscale.
|
||||
The native dtype is preserved using the matching PIL mode
|
||||
(``I;16`` / ``F``). This is the path used for raw depth maps (no rescaling, clamping, or downcasting)
|
||||
- ``(3, H, W)`` / ``(H, W, 3)``: RGB. Channels-first inputs are transposed
|
||||
to channels-last. Float inputs in ``[0, 1]`` are scaled to ``uint8``
|
||||
(existing behaviour, gated by ``range_check``).
|
||||
|
||||
Other shapes / channel counts raise ``NotImplementedError`` or
|
||||
``ValueError``.
|
||||
"""
|
||||
# TODO(CarolinePascal): 4 dimensions RGB-D images
|
||||
if image_array.ndim not in (2, 3):
|
||||
raise ValueError(f"The array has {image_array.ndim} dimensions, but 2 or 3 is expected for an image.")
|
||||
|
||||
# Squeeze 3D single-channel inputs to 2D so depth maps work whether the
|
||||
# caller emits (H, W), (1, H, W), or (H, W, 1).
|
||||
image_array = squeeze_single_channel(image_array)
|
||||
|
||||
if image_array.ndim == 2:
|
||||
if image_array.dtype not in [np.uint16, np.float32]:
|
||||
raise ValueError(
|
||||
f"Unsupported single-channel image dtype: {image_array.dtype}. "
|
||||
f"Supported dtypes: {sorted(str(d) for d in [np.uint16, np.float32])}."
|
||||
)
|
||||
return PIL.Image.fromarray(np.ascontiguousarray(image_array))
|
||||
|
||||
# 3D path: must be RGB (3 channels), channels-first or channels-last.
|
||||
if image_array.shape[0] == 3:
|
||||
# Transpose from pytorch convention (C, H, W) to (H, W, C)
|
||||
image_array = image_array.transpose(1, 2, 0)
|
||||
@@ -111,29 +71,13 @@ def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True)
|
||||
return PIL.Image.fromarray(image_array)
|
||||
|
||||
|
||||
def save_kwargs_for_path(fpath: Path, compress_level: int) -> dict:
|
||||
"""Pick the right format-specific kwargs for :meth:`PIL.Image.Image.save`.
|
||||
|
||||
PNG uses ``compress_level`` (0-9, zlib). TIFF uses ``compression`` (raw) for lossless raw depth maps.
|
||||
"""
|
||||
suffix = Path(fpath).suffix.lower()
|
||||
if suffix == ".png":
|
||||
return {"compress_level": compress_level}
|
||||
if suffix in (".tif", ".tiff"):
|
||||
return {"compression": "raw"}
|
||||
else:
|
||||
raise ValueError(f"Unsupported image file extension: {suffix}")
|
||||
|
||||
|
||||
def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1):
|
||||
"""
|
||||
Saves a NumPy array or PIL Image to a file.
|
||||
|
||||
This function handles both NumPy arrays and PIL Image objects, converting
|
||||
the former to a PIL Image before saving. It includes error handling for
|
||||
the save operation. The output format is inferred from the *fpath*
|
||||
extension: ``.png`` → PNG with ``compress_level``, ``.tiff`` / ``.tif``
|
||||
→ lossless raw depth maps (TIFF).
|
||||
the save operation.
|
||||
|
||||
Args:
|
||||
image (np.ndarray | PIL.Image.Image): The image data to save.
|
||||
@@ -157,7 +101,7 @@ def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level
|
||||
img = image
|
||||
else:
|
||||
raise TypeError(f"Unsupported image type: {type(image)}")
|
||||
img.save(fpath, **save_kwargs_for_path(fpath, compress_level))
|
||||
img.save(fpath, compress_level=compress_level)
|
||||
except Exception as e:
|
||||
logger.error("Error writing image %s: %s", fpath, e)
|
||||
|
||||
|
||||
@@ -226,50 +226,28 @@ def load_image_as_numpy(
|
||||
Args:
|
||||
fpath (str | Path): Path to the image file.
|
||||
dtype (np.dtype): The desired data type of the output array. If floating,
|
||||
pixels are scaled to [0, 1]. Only used for RGB images.
|
||||
pixels are scaled to [0, 1].
|
||||
channel_first (bool): If True, converts the image to (C, H, W) format.
|
||||
Otherwise, it remains in (H, W, C) format.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image as a numpy array.
|
||||
"""
|
||||
is_depth = fpath.endswith(".tiff") or fpath.endswith(".tif")
|
||||
if is_depth:
|
||||
# Preserve the native depth dtype (uint16 -> "I;16", float32 -> "F").
|
||||
img = PILImage.open(fpath)
|
||||
img_array = np.array(img)
|
||||
else:
|
||||
img = PILImage.open(fpath).convert("RGB")
|
||||
img_array = np.array(img, dtype=dtype)
|
||||
if np.issubdtype(dtype, np.floating):
|
||||
img_array /= 255.0
|
||||
img = PILImage.open(fpath).convert("RGB")
|
||||
img_array = np.array(img, dtype=dtype)
|
||||
if channel_first: # (H, W, C) -> (C, H, W)
|
||||
img_array = img_array[np.newaxis, ...] if img_array.ndim == 2 else np.transpose(img_array, (2, 0, 1))
|
||||
img_array = np.transpose(img_array, (2, 0, 1))
|
||||
if np.issubdtype(dtype, np.floating):
|
||||
img_array /= 255.0
|
||||
return img_array
|
||||
|
||||
|
||||
# PIL modes for 16-bit unsigned depth maps.
|
||||
UINT16_PIL_MODES = {"I;16", "I;16B", "I;16L"}
|
||||
|
||||
|
||||
def pil_to_chw_tensor(img: PILImage.Image) -> torch.Tensor:
|
||||
"""Convert a PIL image to a channel-first tensor.
|
||||
|
||||
``uint16`` depth maps become ``float32 (1, H, W)`` in native units (``ToTensor``
|
||||
would overflow them to ``int16``); all other modes use the standard ``ToTensor`` path.
|
||||
"""
|
||||
if img.mode in UINT16_PIL_MODES:
|
||||
return torch.from_numpy(np.array(img, dtype=np.float32))[None, ...]
|
||||
return transforms.ToTensor()(img)
|
||||
|
||||
|
||||
def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[torch.Tensor | str]]:
|
||||
"""Convert a batch from a Hugging Face dataset to torch tensors.
|
||||
|
||||
This transform function converts items from Hugging Face dataset format (pyarrow)
|
||||
to torch tensors. RGB images are converted from PIL objects (H, W, C, uint8)
|
||||
to a torch image representation (C, H, W, float32) in the range [0, 1]. Depth
|
||||
maps are returned as float32 (1, H, W) in their native units. Other
|
||||
to torch tensors. Importantly, images are converted from PIL objects (H, W, C, uint8)
|
||||
to a torch image representation (C, H, W, float32) in the range [0, 1]. Other
|
||||
types are converted to torch.tensor.
|
||||
|
||||
Args:
|
||||
@@ -284,7 +262,8 @@ def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[to
|
||||
continue
|
||||
first_item = items_dict[key][0]
|
||||
if isinstance(first_item, PILImage.Image):
|
||||
items_dict[key] = [pil_to_chw_tensor(img) for img in items_dict[key]]
|
||||
to_tensor = transforms.ToTensor()
|
||||
items_dict[key] = [to_tensor(img) for img in items_dict[key]]
|
||||
elif first_item is None or isinstance(first_item, dict):
|
||||
pass
|
||||
else:
|
||||
@@ -350,11 +329,7 @@ def item_to_torch(item: dict) -> dict:
|
||||
"""
|
||||
skip_keys = {"task", *LANGUAGE_COLUMNS}
|
||||
for key, val in item.items():
|
||||
if key in skip_keys:
|
||||
continue
|
||||
if isinstance(val, PILImage.Image):
|
||||
item[key] = pil_to_chw_tensor(val)
|
||||
elif isinstance(val, (np.ndarray | list)):
|
||||
if isinstance(val, (np.ndarray | list)) and key not in skip_keys:
|
||||
# Convert numpy arrays and lists to torch tensors
|
||||
item[key] = torch.tensor(val)
|
||||
return item
|
||||
|
||||
@@ -24,7 +24,7 @@ import torch.utils
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
from huggingface_hub.errors import RevisionNotFoundError
|
||||
|
||||
from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig, RGBEncoderConfig
|
||||
from lerobot.configs import VideoEncoderConfig
|
||||
from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE
|
||||
|
||||
from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
|
||||
@@ -58,10 +58,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
download_videos: bool = True,
|
||||
video_backend: str | None = None,
|
||||
return_uint8: bool = False,
|
||||
depth_output_unit: str = DEFAULT_DEPTH_UNIT,
|
||||
batch_encoding_size: int = 1,
|
||||
rgb_encoder: RGBEncoderConfig | None = None,
|
||||
depth_encoder: DepthEncoderConfig | None = None,
|
||||
camera_encoder: VideoEncoderConfig | None = None,
|
||||
encoder_threads: int | None = None,
|
||||
streaming_encoding: bool = False,
|
||||
encoder_queue_maxsize: int = 30,
|
||||
@@ -185,11 +183,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision.
|
||||
batch_encoding_size (int, optional): Number of episodes to accumulate before batch encoding videos.
|
||||
Set to 1 for immediate encoding (default), or higher for batched encoding. Defaults to 1.
|
||||
rgb_encoder (RGBEncoderConfig | None, optional): Video encoder settings for cameras
|
||||
(codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults`
|
||||
is used by the writer.
|
||||
depth_encoder (DepthEncoderConfig | None, optional): Video encoder settings for depth cameras
|
||||
(codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults`
|
||||
camera_encoder (VideoEncoderConfig | None, optional): Video encoder settings for cameras
|
||||
(codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults`
|
||||
is used by the writer.
|
||||
encoder_threads (int | None, optional): Number of encoder threads (global). ``None`` lets the
|
||||
codec decide.
|
||||
@@ -211,7 +206,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
self.revision = revision if revision else CODEBASE_VERSION
|
||||
self._video_backend = video_backend if video_backend else get_safe_default_video_backend()
|
||||
self._return_uint8 = return_uint8
|
||||
self._depth_output_unit = depth_output_unit
|
||||
self._batch_encoding_size = batch_encoding_size
|
||||
self._encoder_threads = encoder_threads
|
||||
|
||||
@@ -252,7 +246,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
delta_timestamps=delta_timestamps,
|
||||
image_transforms=image_transforms,
|
||||
return_uint8=self._return_uint8,
|
||||
depth_output_unit=self._depth_output_unit,
|
||||
)
|
||||
self.image_transforms = image_transforms
|
||||
|
||||
@@ -278,16 +271,14 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
if streaming_encoding and len(self.meta.video_keys) > 0:
|
||||
streaming_enc = self._build_streaming_encoder(
|
||||
self.meta.fps,
|
||||
rgb_encoder,
|
||||
depth_encoder,
|
||||
camera_encoder,
|
||||
encoder_queue_maxsize,
|
||||
encoder_threads,
|
||||
)
|
||||
self.writer = DatasetWriter(
|
||||
meta=self.meta,
|
||||
root=self.root,
|
||||
rgb_encoder=rgb_encoder,
|
||||
depth_encoder=depth_encoder,
|
||||
camera_encoder=camera_encoder,
|
||||
encoder_threads=encoder_threads,
|
||||
batch_encoding_size=batch_encoding_size,
|
||||
streaming_encoder=streaming_enc,
|
||||
@@ -323,22 +314,19 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
delta_timestamps=self.delta_timestamps,
|
||||
image_transforms=self.image_transforms,
|
||||
return_uint8=self._return_uint8,
|
||||
depth_output_unit=self._depth_output_unit,
|
||||
)
|
||||
return self.reader
|
||||
|
||||
@staticmethod
|
||||
def _build_streaming_encoder(
|
||||
fps: int,
|
||||
rgb_encoder: RGBEncoderConfig | None,
|
||||
depth_encoder: DepthEncoderConfig | None,
|
||||
camera_encoder: VideoEncoderConfig | None,
|
||||
encoder_queue_maxsize: int,
|
||||
encoder_threads: int | None,
|
||||
) -> StreamingVideoEncoder:
|
||||
return StreamingVideoEncoder(
|
||||
fps=fps,
|
||||
rgb_encoder=rgb_encoder,
|
||||
depth_encoder=depth_encoder,
|
||||
camera_encoder=camera_encoder,
|
||||
queue_maxsize=encoder_queue_maxsize,
|
||||
encoder_threads=encoder_threads,
|
||||
)
|
||||
@@ -667,8 +655,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
image_writer_threads: int = 0,
|
||||
video_backend: str | None = None,
|
||||
batch_encoding_size: int = 1,
|
||||
rgb_encoder: RGBEncoderConfig | None = None,
|
||||
depth_encoder: DepthEncoderConfig | None = None,
|
||||
camera_encoder: VideoEncoderConfig | None = None,
|
||||
metadata_buffer_size: int = 10,
|
||||
streaming_encoding: bool = False,
|
||||
encoder_queue_maxsize: int = 30,
|
||||
@@ -699,10 +686,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
video_backend: Video decoding backend (used when reading back).
|
||||
batch_encoding_size: Number of episodes to accumulate before
|
||||
batch-encoding videos. ``1`` means encode immediately.
|
||||
rgb_encoder: Video encoder settings for cameras (codec, quality, etc.).
|
||||
When ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` is used.
|
||||
depth_encoder: Video encoder settings for depth cameras (codec, quality, etc.).
|
||||
When ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used.
|
||||
camera_encoder: Video encoder settings for cameras (codec, quality, etc.).
|
||||
When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults` is used.
|
||||
encoder_threads: Number of encoder threads (global). ``None``
|
||||
lets the codec decide.
|
||||
metadata_buffer_size: Number of episode metadata records to buffer
|
||||
@@ -737,7 +722,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
obj.episodes = None
|
||||
obj._video_backend = video_backend if video_backend is not None else get_safe_default_video_backend()
|
||||
obj._return_uint8 = False
|
||||
obj._depth_output_unit = DEFAULT_DEPTH_UNIT
|
||||
obj._batch_encoding_size = batch_encoding_size
|
||||
obj._encoder_threads = encoder_threads
|
||||
|
||||
@@ -747,13 +731,12 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
streaming_enc = None
|
||||
if streaming_encoding and len(obj.meta.video_keys) > 0:
|
||||
streaming_enc = cls._build_streaming_encoder(
|
||||
fps, rgb_encoder, depth_encoder, encoder_queue_maxsize, encoder_threads
|
||||
fps, camera_encoder, encoder_queue_maxsize, encoder_threads
|
||||
)
|
||||
obj.writer = DatasetWriter(
|
||||
meta=obj.meta,
|
||||
root=obj.root,
|
||||
rgb_encoder=rgb_encoder,
|
||||
depth_encoder=depth_encoder,
|
||||
camera_encoder=camera_encoder,
|
||||
encoder_threads=encoder_threads,
|
||||
batch_encoding_size=batch_encoding_size,
|
||||
streaming_encoder=streaming_enc,
|
||||
@@ -776,8 +759,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
force_cache_sync: bool = False,
|
||||
video_backend: str | None = None,
|
||||
batch_encoding_size: int = 1,
|
||||
rgb_encoder: RGBEncoderConfig | None = None,
|
||||
depth_encoder: DepthEncoderConfig | None = None,
|
||||
camera_encoder: VideoEncoderConfig | None = None,
|
||||
encoder_threads: int | None = None,
|
||||
image_writer_processes: int = 0,
|
||||
image_writer_threads: int = 0,
|
||||
@@ -805,10 +787,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
video_backend: Video decoding backend for reading back data.
|
||||
batch_encoding_size: Number of episodes to accumulate before
|
||||
batch-encoding videos.
|
||||
rgb_encoder: Video encoder settings for cameras (codec, quality, etc.).
|
||||
When ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` is used.
|
||||
depth_encoder: Video encoder settings for depth cameras (codec, quality, etc.).
|
||||
When ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used.
|
||||
camera_encoder: Video encoder settings for cameras (codec, quality, etc.).
|
||||
When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults` is used.
|
||||
encoder_threads: Number of encoder threads (global). ``None``
|
||||
lets the codec decide.
|
||||
image_writer_processes: Subprocesses for async image writing.
|
||||
@@ -836,7 +816,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
obj.episodes = None
|
||||
obj._video_backend = video_backend if video_backend else get_safe_default_video_backend()
|
||||
obj._return_uint8 = False
|
||||
obj._depth_output_unit = DEFAULT_DEPTH_UNIT
|
||||
obj._batch_encoding_size = batch_encoding_size
|
||||
|
||||
if obj._requested_root is not None:
|
||||
@@ -856,13 +835,12 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
streaming_enc = None
|
||||
if streaming_encoding and len(obj.meta.video_keys) > 0:
|
||||
streaming_enc = cls._build_streaming_encoder(
|
||||
obj.meta.fps, rgb_encoder, depth_encoder, encoder_queue_maxsize, encoder_threads
|
||||
obj.meta.fps, camera_encoder, encoder_queue_maxsize, encoder_threads
|
||||
)
|
||||
obj.writer = DatasetWriter(
|
||||
meta=obj.meta,
|
||||
root=obj.root,
|
||||
rgb_encoder=rgb_encoder,
|
||||
depth_encoder=depth_encoder,
|
||||
camera_encoder=camera_encoder,
|
||||
encoder_threads=encoder_threads,
|
||||
batch_encoding_size=batch_encoding_size,
|
||||
streaming_encoder=streaming_enc,
|
||||
|
||||
@@ -24,7 +24,6 @@ import logging
|
||||
from typing import Any
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -32,34 +31,6 @@ FFMPEG_NUMERIC_OPTION_TYPES = ("INT", "INT64", "UINT64", "FLOAT", "DOUBLE")
|
||||
FFMPEG_INTEGER_OPTION_TYPES = ("INT", "INT64", "UINT64")
|
||||
|
||||
|
||||
def write_u16_plane(plane: av.video.plane.VideoPlane, src: np.ndarray, fill_value: int | None = None) -> None:
|
||||
"""Copy a 2D ``uint16`` image into the plane's memory buffer, row by row.
|
||||
|
||||
For speed, each row is padded to a wider size than ``width``, so the true row width in
|
||||
memory is ``plane.line_size`` (bytes), not ``width``. Copying as one straight stream
|
||||
would skew the image, so we write only the first ``width`` columns of each row and
|
||||
leave the padding untouched.
|
||||
|
||||
Args:
|
||||
plane: Destination 16-bit plane.
|
||||
src: Source image, shape ``(height, width)``, dtype ``uint16``.
|
||||
fill_value: If given, every pixel (padding included) is set to this first, so the
|
||||
padding holds clean data instead of garbage.
|
||||
"""
|
||||
height, width = src.shape
|
||||
stride_u16 = plane.line_size // np.dtype(np.uint16).itemsize
|
||||
dst = np.frombuffer(plane, dtype=np.uint16).reshape(height, stride_u16)
|
||||
if fill_value is not None:
|
||||
dst.fill(fill_value)
|
||||
dst[:, :width] = src
|
||||
|
||||
|
||||
@functools.cache
|
||||
def get_pix_fmt_channels(pix_fmt: str) -> int:
|
||||
"""Return the number of components (channels) for *pix_fmt*."""
|
||||
return len(av.VideoFormat(pix_fmt).components)
|
||||
|
||||
|
||||
@functools.cache
|
||||
def get_codec(vcodec: str) -> av.codec.Codec | None:
|
||||
"""PyAV write-mode ``Codec`` for *vcodec*, or ``None`` if unavailable."""
|
||||
@@ -121,7 +92,7 @@ def _check_option_value(vcodec: str, label: str, value: Any, opt: av.option.Opti
|
||||
f"{label}={value!r} is not numeric; codec {vcodec!r} expects a number for this option."
|
||||
) from e
|
||||
elif isinstance(value, (float, int)):
|
||||
num_val = float(value)
|
||||
num_val = value
|
||||
else:
|
||||
raise ValueError(
|
||||
f"{label}={value!r} is not numeric; codec {vcodec!r} expects a number for this option."
|
||||
@@ -171,16 +142,6 @@ def _check_pixel_format(vcodec: str, pix_fmt: str) -> None:
|
||||
)
|
||||
|
||||
|
||||
def _check_pix_fmt_channels(pix_fmt: str, channels: int) -> None:
|
||||
"""Ensure *pix_fmt* can carry at least *channels* components."""
|
||||
pix_fmt_channels = get_pix_fmt_channels(pix_fmt)
|
||||
if pix_fmt_channels < channels:
|
||||
raise ValueError(
|
||||
f"pix_fmt={pix_fmt!r} carries only {pix_fmt_channels} component(s) "
|
||||
f"but the source data has {channels} channel(s)."
|
||||
)
|
||||
|
||||
|
||||
def _check_codec_options(vcodec: str, codec_options: dict[str, Any]) -> None:
|
||||
"""Validate merged encoder options (typed) against the codec's published AVOptions."""
|
||||
supported_options = _get_codec_options_by_name(vcodec)
|
||||
@@ -195,18 +156,12 @@ def _check_codec_options(vcodec: str, codec_options: dict[str, Any]) -> None:
|
||||
_check_option_value(vcodec, key, value, supported_options[key])
|
||||
|
||||
|
||||
def check_video_encoder_parameters_pyav(
|
||||
vcodec: str,
|
||||
pix_fmt: str,
|
||||
codec_options: dict[str, Any],
|
||||
channels: int | None = None,
|
||||
) -> None:
|
||||
def check_video_encoder_parameters_pyav(vcodec: str, pix_fmt: str, codec_options: dict[str, Any]) -> None:
|
||||
"""Verify *config* is compatible with the bundled FFmpeg build.
|
||||
|
||||
Checks pixel format, abstract tuning-field compatibility, and each merged
|
||||
encoder option from :meth:`~lerobot.configs.video.VideoEncoderConfig.get_codec_options`
|
||||
against PyAV (including numeric ``extra_options`` present in that dict).
|
||||
When given, additionally verify that *pix_fmt* carries as many components as the source data channels.
|
||||
No-op when ``config.vcodec`` isn't in the local FFmpeg build.
|
||||
|
||||
Raises:
|
||||
@@ -216,6 +171,4 @@ def check_video_encoder_parameters_pyav(
|
||||
if not options:
|
||||
raise ValueError(f"Codec {vcodec!r} is not available in the bundled FFmpeg build")
|
||||
_check_pixel_format(vcodec, pix_fmt)
|
||||
if channels is not None:
|
||||
_check_pix_fmt_channels(pix_fmt, channels)
|
||||
_check_codec_options(vcodec, codec_options)
|
||||
|
||||
@@ -22,11 +22,9 @@ import numpy as np
|
||||
import torch
|
||||
from datasets import load_dataset
|
||||
|
||||
from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig
|
||||
from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
|
||||
|
||||
from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
|
||||
from .depth_utils import dequantize_depth
|
||||
from .feature_utils import get_delta_indices
|
||||
from .io_utils import item_to_torch
|
||||
from .utils import (
|
||||
@@ -37,7 +35,6 @@ from .utils import (
|
||||
)
|
||||
from .video_utils import (
|
||||
VideoDecoderCache,
|
||||
decode_video_frames,
|
||||
decode_video_frames_torchcodec,
|
||||
)
|
||||
|
||||
@@ -255,7 +252,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
rng: np.random.Generator | None = None,
|
||||
shuffle: bool = True,
|
||||
return_uint8: bool = False,
|
||||
depth_output_unit: str = DEFAULT_DEPTH_UNIT,
|
||||
):
|
||||
"""Initialize a StreamingLeRobotDataset.
|
||||
|
||||
@@ -276,8 +272,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
seed (int, optional): Reproducibility random seed.
|
||||
rng (np.random.Generator | None, optional): Random number generator.
|
||||
shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
|
||||
depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm").
|
||||
Defaults to "mm".
|
||||
"""
|
||||
super().__init__()
|
||||
self.repo_id = repo_id
|
||||
@@ -296,7 +290,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
self.streaming = streaming
|
||||
self.buffer_size = buffer_size
|
||||
self._return_uint8 = return_uint8
|
||||
self._depth_output_unit = depth_output_unit
|
||||
|
||||
# We cache the video decoders to avoid re-initializing them at each frame (avoiding a ~10x slowdown)
|
||||
self.video_decoder_cache = None
|
||||
@@ -313,11 +306,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
# Check version
|
||||
check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION)
|
||||
|
||||
self._depth_encoder_configs: dict[str, DepthEncoderConfig] = {
|
||||
vid_key: DepthEncoderConfig.from_video_info(self.meta.features[vid_key].get("info"))
|
||||
for vid_key in self.meta.depth_keys
|
||||
}
|
||||
|
||||
self.delta_timestamps = None
|
||||
self.delta_indices = None
|
||||
|
||||
@@ -566,34 +554,13 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
for video_key, query_ts in query_timestamps.items():
|
||||
root = self.meta.url_root if self.streaming and not self.streaming_from_local else self.root
|
||||
video_path = f"{root}/{self.meta.get_video_file_path(ep_idx, video_key)}"
|
||||
if video_key in self.meta.depth_keys:
|
||||
# Depth maps are 12-bit quantized and only decodable via pyav; dequantize back
|
||||
# to physical units to match the non-streaming reader.
|
||||
frames = decode_video_frames(
|
||||
video_path,
|
||||
query_ts,
|
||||
self.tolerance_s,
|
||||
backend="pyav",
|
||||
return_uint8=False,
|
||||
is_depth=True,
|
||||
)
|
||||
depth_encoder = self._depth_encoder_configs[video_key]
|
||||
frames = dequantize_depth(
|
||||
frames,
|
||||
depth_min=depth_encoder.depth_min,
|
||||
depth_max=depth_encoder.depth_max,
|
||||
shift=depth_encoder.shift,
|
||||
use_log=depth_encoder.use_log,
|
||||
output_unit=self._depth_output_unit,
|
||||
)
|
||||
else:
|
||||
frames = decode_video_frames_torchcodec(
|
||||
video_path,
|
||||
query_ts,
|
||||
self.tolerance_s,
|
||||
decoder_cache=self.video_decoder_cache,
|
||||
return_uint8=self._return_uint8,
|
||||
)
|
||||
frames = decode_video_frames_torchcodec(
|
||||
video_path,
|
||||
query_ts,
|
||||
self.tolerance_s,
|
||||
decoder_cache=self.video_decoder_cache,
|
||||
return_uint8=self._return_uint8,
|
||||
)
|
||||
|
||||
item[video_key] = frames.squeeze(0) if len(query_ts) == 1 else frames
|
||||
|
||||
|
||||
@@ -87,14 +87,11 @@ DATA_DIR = "data"
|
||||
VIDEO_DIR = "videos"
|
||||
|
||||
CHUNK_FILE_PATTERN = "chunk-{chunk_index:03d}/file-{file_index:03d}"
|
||||
IMAGE_FILE_PATTERN = "frame-{frame_index:06d}.png"
|
||||
DEPTH_FILE_PATTERN = "frame-{frame_index:06d}.tiff"
|
||||
DEFAULT_TASKS_PATH = "meta/tasks.parquet"
|
||||
DEFAULT_EPISODES_PATH = EPISODES_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet"
|
||||
DEFAULT_DATA_PATH = DATA_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet"
|
||||
DEFAULT_VIDEO_PATH = VIDEO_DIR + "/{video_key}/" + CHUNK_FILE_PATTERN + ".mp4"
|
||||
DEFAULT_IMAGE_PATH = "images/{image_key}/episode-{episode_index:06d}/" + IMAGE_FILE_PATTERN
|
||||
DEFAULT_DEPTH_PATH = "images/{image_key}/episode-{episode_index:06d}/" + DEPTH_FILE_PATTERN
|
||||
DEFAULT_IMAGE_PATH = "images/{image_key}/episode-{episode_index:06d}/frame-{frame_index:06d}.png"
|
||||
|
||||
LEGACY_EPISODES_PATH = "meta/episodes.jsonl"
|
||||
LEGACY_EPISODES_STATS_PATH = "meta/episodes_stats.jsonl"
|
||||
|
||||
@@ -39,17 +39,11 @@ from datasets.features.features import register_feature
|
||||
from PIL import Image
|
||||
|
||||
from lerobot.configs import (
|
||||
DepthEncoderConfig,
|
||||
RGBEncoderConfig,
|
||||
VideoEncoderConfig,
|
||||
depth_encoder_defaults,
|
||||
rgb_encoder_defaults,
|
||||
camera_encoder_defaults,
|
||||
)
|
||||
from lerobot.utils.import_utils import get_safe_default_video_backend
|
||||
|
||||
from .depth_utils import MM_PER_METRE, quantize_depth
|
||||
from .pyav_utils import get_pix_fmt_channels
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -59,7 +53,6 @@ def decode_video_frames(
|
||||
tolerance_s: float,
|
||||
backend: str | None = None,
|
||||
return_uint8: bool = False,
|
||||
is_depth: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Decodes video frames using the specified backend.
|
||||
@@ -71,35 +64,23 @@ def decode_video_frames(
|
||||
backend (str, optional): Backend to use for decoding. Defaults to "torchcodec" when available
|
||||
in the platform; otherwise, defaults to "pyav". The legacy value "video_reader" is
|
||||
accepted for one release as an alias for "pyav" and will be removed in a future version.
|
||||
return_uint8 (bool): For RGB videos, if True return raw uint8 frames without float32 normalization.
|
||||
return_uint8 (bool): If True, return raw uint8 frames without float32 normalization.
|
||||
This reduces memory for DataLoader IPC; normalization can be done on GPU afterward.
|
||||
is_depth (bool): Set to True if the video is a depth map (1 channel, uint12).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Decoded frames (RGB: float32 in [0,1] by default, or uint8 if return_uint8=True, Depth: uint12).
|
||||
torch.Tensor: Decoded frames (float32 in [0,1] by default, or uint8 if return_uint8=True).
|
||||
|
||||
Currently supports torchcodec on cpu and pyav.
|
||||
"""
|
||||
if backend != "pyav" and is_depth:
|
||||
logger.debug("Decoding depth maps is only supported with the 'pyav' backend, falling back to pyav.")
|
||||
# We do not actually return uint8 here, but we avoid the 255 normalization step.
|
||||
return decode_video_frames_pyav(
|
||||
video_path, timestamps, tolerance_s, return_uint8=False, is_depth=True
|
||||
)
|
||||
|
||||
if backend is None:
|
||||
backend = get_safe_default_video_backend()
|
||||
if backend == "torchcodec":
|
||||
return decode_video_frames_torchcodec(video_path, timestamps, tolerance_s, return_uint8=return_uint8)
|
||||
elif backend == "pyav":
|
||||
return decode_video_frames_pyav(
|
||||
video_path, timestamps, tolerance_s, return_uint8=return_uint8, is_depth=is_depth
|
||||
)
|
||||
return decode_video_frames_pyav(video_path, timestamps, tolerance_s, return_uint8=return_uint8)
|
||||
elif backend == "video_reader":
|
||||
logger.warning("backend='video_reader' is deprecated and now aliases to 'pyav'.")
|
||||
return decode_video_frames_pyav(
|
||||
video_path, timestamps, tolerance_s, return_uint8=return_uint8, is_depth=is_depth
|
||||
)
|
||||
return decode_video_frames_pyav(video_path, timestamps, tolerance_s, return_uint8=return_uint8)
|
||||
else:
|
||||
raise ValueError(f"Unsupported video backend: {backend}")
|
||||
|
||||
@@ -110,7 +91,6 @@ def decode_video_frames_pyav(
|
||||
tolerance_s: float,
|
||||
log_loaded_timestamps: bool = False,
|
||||
return_uint8: bool = False,
|
||||
is_depth: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Loads frames associated to the requested timestamps of a video using PyAV.
|
||||
|
||||
@@ -129,9 +109,8 @@ def decode_video_frames_pyav(
|
||||
tolerance_s: Allowed deviation in seconds between a queried timestamp and the closest
|
||||
decoded frame.
|
||||
log_loaded_timestamps: When True, log every decoded frame's timestamp at INFO level.
|
||||
return_uint8: For RGB videos, if True return raw uint8 frames (C, H, W).
|
||||
Otherwise, return float32 in [0, 1] range.
|
||||
is_depth: Set to True if the video is a depth map (1 channel, uint12).
|
||||
return_uint8: When True, return raw uint8 frames (C, H, W). Otherwise, return float32 in
|
||||
[0, 1] range.
|
||||
|
||||
Returns:
|
||||
torch.Tensor of shape (len(timestamps), C, H, W).
|
||||
@@ -153,13 +132,7 @@ def decode_video_frames_pyav(
|
||||
# https://pyav.basswood-io.com/docs/stable/api/container.html#av.container.InputContainer.seek
|
||||
with av.open(video_path) as container:
|
||||
stream = container.streams.video[0]
|
||||
# Seek to the nearest keyframe at or before `first_ts` with a 1 frame margin
|
||||
container.seek(
|
||||
round(first_ts / stream.time_base) - 1,
|
||||
backward=True,
|
||||
any_frame=False,
|
||||
stream=stream,
|
||||
)
|
||||
container.seek(int(first_ts * av.time_base), backward=True)
|
||||
|
||||
for frame in container.decode(stream):
|
||||
if frame.pts is None:
|
||||
@@ -167,13 +140,9 @@ def decode_video_frames_pyav(
|
||||
current_ts = float(frame.pts * stream.time_base)
|
||||
if log_loaded_timestamps:
|
||||
logger.info(f"frame loaded at timestamp={current_ts:.4f}")
|
||||
if is_depth:
|
||||
arr = frame.to_ndarray(format="gray12le") # (H, W) uint12
|
||||
loaded_frames.append(torch.from_numpy(arr).unsqueeze(0).contiguous())
|
||||
else:
|
||||
arr = frame.to_ndarray(format="rgb24") # (H, W, 3)
|
||||
# Convert to CHW uint8 to match torchcodec's output layout.
|
||||
loaded_frames.append(torch.from_numpy(arr).permute(2, 0, 1).contiguous())
|
||||
# Convert to CHW uint8 to match torchcodec's output layout.
|
||||
arr = frame.to_ndarray(format="rgb24") # H, W, 3
|
||||
loaded_frames.append(torch.from_numpy(arr).permute(2, 0, 1).contiguous())
|
||||
loaded_ts.append(current_ts)
|
||||
if current_ts >= last_ts:
|
||||
break
|
||||
@@ -216,7 +185,7 @@ def decode_video_frames_pyav(
|
||||
f"number of queried timestamps ({len(timestamps)})"
|
||||
)
|
||||
|
||||
if return_uint8 or is_depth:
|
||||
if return_uint8:
|
||||
return closest_frames
|
||||
|
||||
# convert to the pytorch format which is float32 in [0,1] range (and channel first)
|
||||
@@ -437,38 +406,17 @@ def encode_video_frames(
|
||||
imgs_dir: Path | str,
|
||||
video_path: Path | str,
|
||||
fps: int,
|
||||
video_encoder: VideoEncoderConfig | None = None,
|
||||
camera_encoder: VideoEncoderConfig | None = None,
|
||||
encoder_threads: int | None = None,
|
||||
*,
|
||||
log_level: int | None = av.logging.WARNING,
|
||||
overwrite: bool = False,
|
||||
) -> None:
|
||||
"""Encode a directory of image frames into an MP4 video.
|
||||
|
||||
When ``video_encoder`` is a :class:`~lerobot.configs.video.DepthEncoderConfig`,
|
||||
frames are read from ``.tiff`` files and quantized to 12-bit depth codes using the
|
||||
encoder's ``depth_min`` / ``depth_max`` / ``shift`` / ``use_log``; otherwise ``.png``
|
||||
RGB frames are encoded directly.
|
||||
|
||||
Args:
|
||||
imgs_dir: Directory containing the frames to encode, named ``frame-000000``
|
||||
onwards (``.png`` for RGB, ``.tiff`` for depth).
|
||||
video_path: Output path for the encoded ``.mp4`` file.
|
||||
fps: Frame rate of the output video.
|
||||
video_encoder: Encoder settings (codec, pixel format, quality, ...). When
|
||||
``None``, :func:`rgb_encoder_defaults` is used. Pass a
|
||||
:class:`~lerobot.configs.video.DepthEncoderConfig` to encode depth frames.
|
||||
encoder_threads: Per-encoder thread count forwarded to the codec. ``None``
|
||||
lets the codec decide.
|
||||
log_level: libav log level to set while encoding, or ``None`` to leave the
|
||||
current logging configuration unchanged.
|
||||
overwrite: When ``False`` and ``video_path`` already exists, skip encoding and
|
||||
log a warning. When ``True``, re-encode and replace the existing file.
|
||||
"""
|
||||
if video_encoder is None:
|
||||
video_encoder = rgb_encoder_defaults()
|
||||
vcodec = video_encoder.vcodec
|
||||
pix_fmt = video_encoder.pix_fmt
|
||||
"""More info on ffmpeg arguments tuning on `benchmark/video/README.md`"""
|
||||
if camera_encoder is None:
|
||||
camera_encoder = camera_encoder_defaults()
|
||||
vcodec = camera_encoder.vcodec
|
||||
pix_fmt = camera_encoder.pix_fmt
|
||||
|
||||
video_path = Path(video_path)
|
||||
imgs_dir = Path(imgs_dir)
|
||||
@@ -480,19 +428,17 @@ def encode_video_frames(
|
||||
video_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Get input frames
|
||||
is_depth = isinstance(video_encoder, DepthEncoderConfig)
|
||||
suffix = ".png" if not is_depth else ".tiff"
|
||||
template = "frame-" + ("[0-9]" * 6) + suffix
|
||||
template = "frame-" + ("[0-9]" * 6) + ".png"
|
||||
input_list = sorted(
|
||||
glob.glob(str(imgs_dir / template)), key=lambda x: int(x.split("-")[-1].split(".")[0])
|
||||
)
|
||||
|
||||
if len(input_list) == 0:
|
||||
raise FileNotFoundError(f"No images with suffix {suffix} found in {imgs_dir}.")
|
||||
raise FileNotFoundError(f"No images found in {imgs_dir}.")
|
||||
with Image.open(input_list[0]) as dummy_image:
|
||||
width, height = dummy_image.size
|
||||
|
||||
video_options = video_encoder.get_codec_options(encoder_threads, as_strings=True)
|
||||
video_options = camera_encoder.get_codec_options(encoder_threads, as_strings=True)
|
||||
|
||||
# Set logging level
|
||||
if log_level is not None:
|
||||
@@ -509,19 +455,8 @@ def encode_video_frames(
|
||||
# Loop through input frames and encode them
|
||||
for input_data in input_list:
|
||||
with Image.open(input_data) as input_image:
|
||||
if is_depth:
|
||||
input_frame = quantize_depth(
|
||||
np.array(input_image),
|
||||
depth_min=video_encoder.depth_min,
|
||||
depth_max=video_encoder.depth_max,
|
||||
shift=video_encoder.shift,
|
||||
use_log=video_encoder.use_log,
|
||||
pix_fmt=video_encoder.pix_fmt,
|
||||
video_backend="pyav",
|
||||
)
|
||||
else:
|
||||
input_image = input_image.convert("RGB")
|
||||
input_frame = av.VideoFrame.from_image(input_image)
|
||||
input_image = input_image.convert("RGB")
|
||||
input_frame = av.VideoFrame.from_image(input_image)
|
||||
packet = output_stream.encode(input_frame)
|
||||
if packet:
|
||||
output.mux(packet)
|
||||
@@ -542,7 +477,7 @@ def encode_video_frames(
|
||||
def reencode_video(
|
||||
input_video_path: Path | str,
|
||||
output_video_path: Path | str,
|
||||
video_encoder: VideoEncoderConfig | None = None,
|
||||
camera_encoder: VideoEncoderConfig | None = None,
|
||||
encoder_threads: int | None = None,
|
||||
log_level: int | None = av.logging.WARNING,
|
||||
overwrite: bool = False,
|
||||
@@ -554,7 +489,7 @@ def reencode_video(
|
||||
Args:
|
||||
input_video_path: Existing video file to read.
|
||||
output_video_path: Path for the re-encoded file.
|
||||
video_encoder: Encoder configuration. Defaults to :func:`rgb_encoder_defaults`.
|
||||
camera_encoder: Encoder configuration. Defaults to :func:`camera_encoder_defaults`.
|
||||
encoder_threads: Optional thread count forwarded to :meth:`VideoEncoderConfig.get_codec_options`.
|
||||
log_level: libav log level while encoding, or ``None`` to leave logging unchanged. Defaults to WARNING.
|
||||
overwrite: When ``False`` and ``output_video_path`` already exists, skip and log a warning.
|
||||
@@ -562,7 +497,7 @@ def reencode_video(
|
||||
end_time_s: When set, trim the output to end at this timestamp (seconds, exclusive).
|
||||
"""
|
||||
|
||||
video_encoder = video_encoder or rgb_encoder_defaults()
|
||||
camera_encoder = camera_encoder or camera_encoder_defaults()
|
||||
|
||||
if (start_time_s is not None and start_time_s < 0) or (end_time_s is not None and end_time_s < 0):
|
||||
raise ValueError(f"Trim times must be non-negative, got start={start_time_s}, end={end_time_s}.")
|
||||
@@ -577,9 +512,9 @@ def reencode_video(
|
||||
|
||||
output_video_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
video_options = video_encoder.get_codec_options(encoder_threads, as_strings=True)
|
||||
vcodec = video_encoder.vcodec
|
||||
pix_fmt = video_encoder.pix_fmt
|
||||
video_options = camera_encoder.get_codec_options(encoder_threads, as_strings=True)
|
||||
vcodec = camera_encoder.vcodec
|
||||
pix_fmt = camera_encoder.pix_fmt
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_named_file:
|
||||
tmp_output_video_path = tmp_named_file.name
|
||||
@@ -761,21 +696,22 @@ class _CameraEncoderThread(threading.Thread):
|
||||
self,
|
||||
video_path: Path,
|
||||
fps: int,
|
||||
video_encoder: VideoEncoderConfig,
|
||||
vcodec: str,
|
||||
pix_fmt: str,
|
||||
codec_options: dict[str, str],
|
||||
frame_queue: queue.Queue,
|
||||
result_queue: queue.Queue,
|
||||
stop_event: threading.Event,
|
||||
encoder_threads: int | None = None,
|
||||
):
|
||||
super().__init__(daemon=True)
|
||||
self.video_path = video_path
|
||||
self.fps = fps
|
||||
self.video_encoder = video_encoder
|
||||
self.is_depth = isinstance(video_encoder, DepthEncoderConfig)
|
||||
self.vcodec = vcodec
|
||||
self.pix_fmt = pix_fmt
|
||||
self.codec_options = codec_options
|
||||
self.frame_queue = frame_queue
|
||||
self.result_queue = result_queue
|
||||
self.stop_event = stop_event
|
||||
self.encoder_threads = encoder_threads
|
||||
|
||||
def run(self) -> None:
|
||||
from .compute_stats import RunningQuantileStats, auto_downsample_height_width
|
||||
@@ -800,12 +736,12 @@ class _CameraEncoderThread(threading.Thread):
|
||||
# Sentinel: flush and close
|
||||
break
|
||||
|
||||
# Ensure HWC (RGB or depth) uint8 (RGB only) numpy array
|
||||
# Ensure HWC uint8 numpy array
|
||||
if isinstance(frame_data, np.ndarray):
|
||||
if frame_data.ndim == 3 and frame_data.shape[0] in (1, 3):
|
||||
if frame_data.ndim == 3 and frame_data.shape[0] == 3:
|
||||
# CHW -> HWC
|
||||
frame_data = frame_data.transpose(1, 2, 0)
|
||||
if not self.is_depth and frame_data.dtype != np.uint8:
|
||||
if frame_data.dtype != np.uint8:
|
||||
frame_data = (frame_data * 255).astype(np.uint8)
|
||||
|
||||
# Open container on first frame (to get width/height)
|
||||
@@ -813,29 +749,15 @@ class _CameraEncoderThread(threading.Thread):
|
||||
height, width = frame_data.shape[:2]
|
||||
Path(self.video_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
container = av.open(str(self.video_path), "w")
|
||||
output_stream = container.add_stream(
|
||||
self.video_encoder.vcodec,
|
||||
self.fps,
|
||||
options=self.video_encoder.get_codec_options(self.encoder_threads, as_strings=True),
|
||||
)
|
||||
output_stream.pix_fmt = self.video_encoder.pix_fmt
|
||||
output_stream = container.add_stream(self.vcodec, self.fps, options=self.codec_options)
|
||||
output_stream.pix_fmt = self.pix_fmt
|
||||
output_stream.width = width
|
||||
output_stream.height = height
|
||||
output_stream.time_base = Fraction(1, self.fps)
|
||||
|
||||
# Encode frame with explicit timestamps
|
||||
if not self.is_depth:
|
||||
pil_img = Image.fromarray(frame_data)
|
||||
video_frame = av.VideoFrame.from_image(pil_img)
|
||||
else:
|
||||
video_frame = quantize_depth(
|
||||
frame_data,
|
||||
depth_min=self.video_encoder.depth_min,
|
||||
depth_max=self.video_encoder.depth_max,
|
||||
shift=self.video_encoder.shift,
|
||||
use_log=self.video_encoder.use_log,
|
||||
video_backend=self.video_encoder.video_backend,
|
||||
)
|
||||
pil_img = Image.fromarray(frame_data)
|
||||
video_frame = av.VideoFrame.from_image(pil_img)
|
||||
video_frame.pts = frame_count
|
||||
video_frame.time_base = Fraction(1, self.fps)
|
||||
packet = output_stream.encode(video_frame)
|
||||
@@ -848,9 +770,6 @@ class _CameraEncoderThread(threading.Thread):
|
||||
# Reshape CHW to (H*W, C) for per-channel stats
|
||||
channels = img_downsampled.shape[0]
|
||||
img_for_stats = img_downsampled.transpose(1, 2, 0).reshape(-1, channels)
|
||||
# Depth stats are canonically stored in millimetres; metre (float) depth is scaled up.
|
||||
if self.is_depth and np.issubdtype(frame_data.dtype, np.floating):
|
||||
img_for_stats = img_for_stats * MM_PER_METRE
|
||||
stats_tracker.update(img_for_stats)
|
||||
|
||||
frame_count += 1
|
||||
@@ -896,27 +815,22 @@ class StreamingVideoEncoder:
|
||||
def __init__(
|
||||
self,
|
||||
fps: int,
|
||||
rgb_encoder: RGBEncoderConfig | None = None,
|
||||
depth_encoder: DepthEncoderConfig | None = None,
|
||||
camera_encoder: VideoEncoderConfig | None = None,
|
||||
queue_maxsize: int = 30,
|
||||
encoder_threads: int | None = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
fps: Frames per second for the output videos.
|
||||
rgb_encoder: Video encoder settings applied to all RGB cameras.
|
||||
When ``None``, :func:`rgb_encoder_defaults` is used.
|
||||
depth_encoder: Video encoder settings applied to all depth cameras,
|
||||
including the depth quantization parameters. When ``None``,
|
||||
:func:`depth_encoder_defaults` is used.
|
||||
queue_maxsize: Max frames to buffer per camera before
|
||||
back-pressure drops frames.
|
||||
camera_encoder: Video encoder settings applied to all cameras.
|
||||
When ``None``, :func:`camera_encoder_defaults` is used.
|
||||
encoder_threads: Number of encoder threads (global setting).
|
||||
``None`` lets the codec decide.
|
||||
queue_maxsize: Max frames to buffer per camera before
|
||||
back-pressure drops frames.
|
||||
"""
|
||||
self.fps = fps
|
||||
self._rgb_encoder = rgb_encoder or rgb_encoder_defaults()
|
||||
self._depth_encoder = depth_encoder or depth_encoder_defaults()
|
||||
self._camera_encoder = camera_encoder or camera_encoder_defaults()
|
||||
self._encoder_threads = encoder_threads
|
||||
self.queue_maxsize = queue_maxsize
|
||||
|
||||
@@ -929,25 +843,18 @@ class StreamingVideoEncoder:
|
||||
self._episode_active = False
|
||||
self._closed = False
|
||||
|
||||
def start_episode(
|
||||
self, video_keys: list[str], temp_dir: Path, depth_video_keys: list[str] | None = None
|
||||
) -> None:
|
||||
def start_episode(self, video_keys: list[str], temp_dir: Path) -> None:
|
||||
"""Start encoder threads for a new episode.
|
||||
|
||||
Args:
|
||||
video_keys: List of video feature keys (e.g. ["observation.images.laptop"])
|
||||
temp_dir: Base directory for temporary MP4 files
|
||||
depth_video_keys: List of video or image feature keys that carry depth maps (e.g.
|
||||
["observation.images.laptop_depth"]). Defaults to ``[]`` (no depth keys).
|
||||
"""
|
||||
if self._episode_active:
|
||||
self.cancel_episode()
|
||||
|
||||
self._dropped_frames.clear()
|
||||
|
||||
if depth_video_keys is None:
|
||||
depth_video_keys = []
|
||||
|
||||
for video_key in video_keys:
|
||||
frame_queue: queue.Queue = queue.Queue(maxsize=self.queue_maxsize)
|
||||
result_queue: queue.Queue = queue.Queue(maxsize=1)
|
||||
@@ -956,15 +863,17 @@ class StreamingVideoEncoder:
|
||||
temp_video_dir = Path(tempfile.mkdtemp(dir=temp_dir))
|
||||
video_path = temp_video_dir / f"{video_key.replace('/', '_')}_streaming.mp4"
|
||||
|
||||
encoder = self._depth_encoder if video_key in depth_video_keys else self._rgb_encoder
|
||||
vcodec = self._camera_encoder.vcodec
|
||||
codec_options = self._camera_encoder.get_codec_options(self._encoder_threads, as_strings=True)
|
||||
encoder_thread = _CameraEncoderThread(
|
||||
video_path=video_path,
|
||||
fps=self.fps,
|
||||
video_encoder=encoder,
|
||||
vcodec=vcodec,
|
||||
pix_fmt=self._camera_encoder.pix_fmt,
|
||||
codec_options=codec_options,
|
||||
frame_queue=frame_queue,
|
||||
result_queue=result_queue,
|
||||
stop_event=stop_event,
|
||||
encoder_threads=self._encoder_threads,
|
||||
)
|
||||
encoder_thread.start()
|
||||
|
||||
@@ -1171,23 +1080,15 @@ def get_audio_info(video_path: Path | str) -> dict:
|
||||
|
||||
def get_video_info(
|
||||
video_path: Path | str,
|
||||
video_encoder: VideoEncoderConfig | None = None,
|
||||
camera_encoder: VideoEncoderConfig | None = None,
|
||||
) -> dict:
|
||||
"""Build the ``video.*`` / ``audio.*`` info dict persisted in ``info.json``.
|
||||
|
||||
Args:
|
||||
video_path: Path to the encoded video file to probe.
|
||||
video_encoder: If provided, record the exact encoder settings used to encode this
|
||||
camera_encoder: If provided, record the exact encoder settings used to encode this
|
||||
video. Stream-derived values take precedence — encoder fields are only written for keys
|
||||
not already populated from the video file itself. When a
|
||||
:class:`~lerobot.configs.video.DepthEncoderConfig` is passed, the depth
|
||||
quantization parameters (``depth_min`` / ``depth_max`` / ``shift`` /
|
||||
``use_log``) are recorded so frames can be dequantized on read.
|
||||
|
||||
Returns:
|
||||
The ``video.*`` / ``audio.*`` info dict, including ``is_depth_map`` which is
|
||||
``True`` only when ``video_encoder`` is a
|
||||
:class:`~lerobot.configs.video.DepthEncoderConfig`.
|
||||
not already populated from the video file itself.
|
||||
"""
|
||||
logging.getLogger("libav").setLevel(av.logging.WARNING)
|
||||
|
||||
@@ -1205,10 +1106,13 @@ def get_video_info(
|
||||
video_info["video.width"] = video_stream.width
|
||||
video_info["video.codec"] = video_stream.codec.canonical_name
|
||||
video_info["video.pix_fmt"] = video_stream.pix_fmt
|
||||
video_info["video.is_depth_map"] = False
|
||||
|
||||
# Calculate fps from r_frame_rate
|
||||
video_info["video.fps"] = int(video_stream.base_rate)
|
||||
video_info["video.channels"] = get_pix_fmt_channels(video_stream.pix_fmt)
|
||||
|
||||
pixel_channels = get_video_pixel_channels(video_stream.pix_fmt)
|
||||
video_info["video.channels"] = pixel_channels
|
||||
|
||||
# Reset logging level
|
||||
av.logging.restore_default_callback()
|
||||
@@ -1217,18 +1121,27 @@ def get_video_info(
|
||||
video_info.update(**get_audio_info(video_path))
|
||||
|
||||
# Add additional encoder configuration if provided
|
||||
if video_encoder is not None:
|
||||
for field_name, field_value in asdict(video_encoder).items():
|
||||
if camera_encoder is not None:
|
||||
for field_name, field_value in asdict(camera_encoder).items():
|
||||
# vcodec is already populated from the video stream
|
||||
if field_name == "vcodec":
|
||||
continue
|
||||
video_info.setdefault(f"video.{field_name}", field_value)
|
||||
|
||||
video_info["is_depth_map"] = isinstance(video_encoder, DepthEncoderConfig)
|
||||
|
||||
return video_info
|
||||
|
||||
|
||||
def get_video_pixel_channels(pix_fmt: str) -> int:
|
||||
if "gray" in pix_fmt or "depth" in pix_fmt or "monochrome" in pix_fmt:
|
||||
return 1
|
||||
elif "rgba" in pix_fmt or "yuva" in pix_fmt:
|
||||
return 4
|
||||
elif "rgb" in pix_fmt or "yuv" in pix_fmt:
|
||||
return 3
|
||||
else:
|
||||
raise ValueError("Unknown format")
|
||||
|
||||
|
||||
def get_video_duration_in_s(video_path: Path | str) -> float:
|
||||
"""
|
||||
Get the duration of a video file in seconds using PyAV.
|
||||
@@ -1289,13 +1202,10 @@ class VideoEncodingManager:
|
||||
img_dir = self.dataset.root / "images"
|
||||
if img_dir.exists():
|
||||
png_files = list(img_dir.rglob("*.png"))
|
||||
tiff_files = list(img_dir.rglob("*.tiff"))
|
||||
if len(png_files) == 0 and len(tiff_files) == 0:
|
||||
if len(png_files) == 0:
|
||||
shutil.rmtree(img_dir)
|
||||
logger.debug("Cleaned up empty images directory")
|
||||
else:
|
||||
logger.debug(
|
||||
f"Images directory is not empty, containing {len(png_files)} PNG and {len(tiff_files)} TIFF files"
|
||||
)
|
||||
logger.debug(f"Images directory is not empty, containing {len(png_files)} PNG files")
|
||||
|
||||
return False # Don't suppress the original exception
|
||||
|
||||
@@ -1,23 +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.
|
||||
|
||||
from lerobot.utils.import_utils import require_package
|
||||
|
||||
# LeRobotDataset (imported at module top in dataset.py) pulls in heavy dataset deps;
|
||||
# guard the optional dependency here so importing this package fails loudly if it's missing.
|
||||
require_package("datasets", extra="dataset")
|
||||
|
||||
from .hf import submit_to_hf
|
||||
|
||||
__all__ = ["submit_to_hf"]
|
||||
@@ -1,53 +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.
|
||||
"""Make a training dataset reachable from an HF Job pod.
|
||||
|
||||
The pod can't see the host's ~/.cache/huggingface/lerobot, so the dataset has to
|
||||
live on the Hub: the pod downloads it by repo_id at train time (the forwarded
|
||||
HF_TOKEN covers private datasets). A dataset already on the Hub is used as-is; a
|
||||
local-only dataset is pushed to a PRIVATE repo first (never public).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from lerobot.datasets import LeRobotDataset
|
||||
from lerobot.utils.constants import HF_LEROBOT_HOME
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
|
||||
def ensure_dataset_available(repo_id: str, *, api: HfApi, tags: list[str] | None = None) -> None:
|
||||
"""Ensure repo_id resolves on the Hub, pushing a local-only dataset privately first.
|
||||
|
||||
`tags` are attached to the dataset only when we push it (an already-on-Hub
|
||||
dataset is left untouched). Raises RuntimeError if the dataset is neither on
|
||||
the Hub nor in the local cache.
|
||||
"""
|
||||
if api.repo_exists(repo_id, repo_type="dataset"):
|
||||
return
|
||||
|
||||
local_present = (HF_LEROBOT_HOME / repo_id / "meta" / "info.json").is_file()
|
||||
if not local_present:
|
||||
raise RuntimeError(
|
||||
f"Dataset '{repo_id}' is not in the local cache ({HF_LEROBOT_HOME}) and could not be "
|
||||
f"reached on the Hub — it may not exist, or be private and inaccessible with your "
|
||||
f"token. Record or download it first, or run `hf auth login`."
|
||||
)
|
||||
|
||||
print(f"[dataset] '{repo_id}' is local-only; pushing to a PRIVATE Hub repo...")
|
||||
LeRobotDataset(repo_id).push_to_hub(private=True, tags=tags)
|
||||
print(f"[dataset] '{repo_id}' uploaded (private). The job will download it by repo_id.")
|
||||
@@ -1,425 +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.
|
||||
"""Run a lerobot training on HF Jobs (HuggingFace GPUs).
|
||||
|
||||
Ported and simplified from lelab's runners/hf_cloud.py: no UI log queue, no
|
||||
registry — just submit and stream to stdout.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import datetime as dt
|
||||
import json
|
||||
import netrc
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
from huggingface_hub import (
|
||||
HfApi,
|
||||
create_repo,
|
||||
fetch_job_logs,
|
||||
get_token,
|
||||
inspect_job,
|
||||
run_job,
|
||||
upload_file,
|
||||
)
|
||||
|
||||
from lerobot.common.train_utils import push_checkpoint_to_hub
|
||||
from lerobot.configs import parser
|
||||
|
||||
from .dataset import ensure_dataset_available
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
|
||||
_SLUG_RE = re.compile(r"[^a-zA-Z0-9._-]+")
|
||||
|
||||
_TERMINAL_STAGES = {"COMPLETED", "CANCELED", "ERROR", "DELETED"}
|
||||
|
||||
# huggingface_hub 1.x runs on httpx: transient HTTP/transport failures surface as
|
||||
# httpx.HTTPError and socket-level errors as OSError. Catching only these keeps real
|
||||
# bugs (TypeError, AttributeError, ...) from being silently retried or counted as
|
||||
# job failures.
|
||||
_TRANSIENT_NET_ERRORS = (OSError, httpx.HTTPError)
|
||||
|
||||
# Always attached to remote jobs and pushed datasets so LeRobot-originated work
|
||||
# is identifiable on the Hub; callers (e.g. LeLab) add their own via --job.tags.
|
||||
LEROBOT_TAG = "lerobot"
|
||||
|
||||
|
||||
def resolve_job_tags(extra: list[str] | None) -> list[str]:
|
||||
"""Return the tag list for a run: the lerobot tag plus any extras, deduped, order-stable."""
|
||||
tags = [LEROBOT_TAG, *(extra or [])]
|
||||
seen: set[str] = set()
|
||||
return [t for t in tags if not (t in seen or seen.add(t))]
|
||||
|
||||
|
||||
def resolve_wandb_api_key() -> str | None:
|
||||
"""Host's wandb key for forwarding to the job: $WANDB_API_KEY, else ~/.netrc."""
|
||||
key = os.environ.get("WANDB_API_KEY")
|
||||
if key:
|
||||
return key
|
||||
try:
|
||||
rc = netrc.netrc()
|
||||
except (FileNotFoundError, netrc.NetrcParseError, OSError):
|
||||
return None
|
||||
auth = rc.authenticators("api.wandb.ai")
|
||||
if auth is None:
|
||||
return None
|
||||
_login, _account, password = auth
|
||||
return password or None
|
||||
|
||||
|
||||
def build_repo_id(username: str, job_name: str, now: dt.datetime) -> str:
|
||||
"""Generate the model repo id for a remote run: <user>/<job_name>_<timestamp>."""
|
||||
slug = _SLUG_RE.sub("-", job_name).strip("-") or "train"
|
||||
stamp = now.strftime("%Y-%m-%d_%H-%M-%S")
|
||||
return f"{username}/{slug}_{stamp}"
|
||||
|
||||
|
||||
def build_remote_config_file(cfg, repo_id: str, dest: Path, tags: list[str] | None = None) -> Path:
|
||||
"""Write a train_config.json for the pod, with remote overrides applied.
|
||||
|
||||
The pod runs `lerobot-train --config_path=<dest>` and downloads the dataset
|
||||
by repo_id into its own cache. Client-only fields are stripped so the config
|
||||
is accepted by the trainer image: `job` (pure client orchestration) is always
|
||||
removed, and `save_checkpoint_to_hub` is removed unless explicitly enabled —
|
||||
older lerobot images reject unknown keys, so the default keeps the config
|
||||
compatible with the released `lerobot-gpu` image. `tags` are merged into
|
||||
policy.tags so the trained model the pod pushes carries them too.
|
||||
"""
|
||||
remote = copy.deepcopy(cfg)
|
||||
remote.policy.push_to_hub = True
|
||||
remote.policy.repo_id = repo_id
|
||||
# Don't pin the client's resolved device (e.g. "mps"); let the pod auto-detect its GPU.
|
||||
remote.policy.device = None
|
||||
# Drop any host-local dataset root; the pod resolves the dataset by repo_id.
|
||||
remote.dataset.root = None
|
||||
if tags:
|
||||
existing = list(remote.policy.tags or [])
|
||||
remote.policy.tags = existing + [t for t in tags if t not in existing]
|
||||
|
||||
# Encode to the canonical, pod-parseable dict, then drop the keys the released
|
||||
# trainer image doesn't know about.
|
||||
data = remote.to_dict()
|
||||
data.pop("job", None)
|
||||
if not remote.save_checkpoint_to_hub:
|
||||
data.pop("save_checkpoint_to_hub", None)
|
||||
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
dest.write_text(json.dumps(data, indent=4))
|
||||
return dest
|
||||
|
||||
|
||||
def _stage_config_on_hub(cfg, repo_id: str, token: str, tags: list[str] | None = None) -> str:
|
||||
"""Upload train_config.json to the model repo and return the repo_id for --config_path."""
|
||||
create_repo(repo_id, repo_type="model", private=True, exist_ok=True, token=token)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
config_path = build_remote_config_file(cfg, repo_id, Path(tmp) / "train_config.json", tags=tags)
|
||||
upload_file(
|
||||
path_or_fileobj=config_path,
|
||||
path_in_repo="train_config.json",
|
||||
repo_id=repo_id,
|
||||
repo_type="model",
|
||||
token=token,
|
||||
)
|
||||
return repo_id
|
||||
|
||||
|
||||
def _tail_logs(
|
||||
job_id: str,
|
||||
done: threading.Event,
|
||||
success_marker: str | None = None,
|
||||
success_event: threading.Event | None = None,
|
||||
) -> None:
|
||||
"""Stream job logs to stdout, reconnecting on dropped streams until done is set.
|
||||
|
||||
Each reconnect re-fetches the full buffered log, so we track how many lines
|
||||
were already printed and skip them — otherwise a fast-failing job's traceback
|
||||
gets reprinted on every reconnect.
|
||||
|
||||
When `success_marker` appears in a line, set `success_event` and `done` so the
|
||||
caller can finish as soon as the trained model lands on the Hub, rather than
|
||||
waiting out the platform's post-run finalization (which can add ~30s).
|
||||
"""
|
||||
printed = 0
|
||||
while not done.is_set():
|
||||
try:
|
||||
seen = 0
|
||||
for line in fetch_job_logs(job_id=job_id, follow=True):
|
||||
seen += 1
|
||||
if seen <= printed:
|
||||
continue # already shown on a previous connection
|
||||
printed = seen
|
||||
# fetch_job_logs yields SSE data without trailing newlines, so add one
|
||||
# per entry — otherwise all log lines concatenate onto a single line.
|
||||
print(line.rstrip("\n"), flush=True)
|
||||
if success_marker and success_event is not None and success_marker in line:
|
||||
success_event.set()
|
||||
done.set()
|
||||
return
|
||||
if done.is_set():
|
||||
return
|
||||
# Stream closed cleanly. Wait a moment so the status poller can mark
|
||||
# the job terminal before we reconnect (avoids re-tailing the buffer).
|
||||
if done.wait(3):
|
||||
return
|
||||
except _TRANSIENT_NET_ERRORS:
|
||||
if done.wait(2):
|
||||
return
|
||||
|
||||
|
||||
def _poll_until_done(
|
||||
job_id: str,
|
||||
done: threading.Event,
|
||||
poll_interval: float = 5.0,
|
||||
status_holder: dict | None = None,
|
||||
max_failures: int = 6,
|
||||
) -> str | None:
|
||||
"""Poll inspect_job until a terminal stage or until `done` is set.
|
||||
|
||||
Returns the terminal stage string, or None if `done` was set first (detach)
|
||||
or after `max_failures` consecutive inspect_job errors. When a terminal stage
|
||||
is reached and `status_holder` is given, records `status_holder["message"]`
|
||||
(the platform's status message, e.g. "Job timeout").
|
||||
"""
|
||||
failures = 0
|
||||
while not done.is_set():
|
||||
try:
|
||||
info = inspect_job(job_id=job_id)
|
||||
failures = 0
|
||||
# `stage` is an enum in some huggingface_hub versions and a plain str in others.
|
||||
stage = getattr(info.status.stage, "value", info.status.stage)
|
||||
if stage in _TERMINAL_STAGES:
|
||||
if status_holder is not None:
|
||||
status_holder["message"] = getattr(info.status, "message", None)
|
||||
done.set()
|
||||
return stage
|
||||
except _TRANSIENT_NET_ERRORS:
|
||||
failures += 1
|
||||
if failures >= max_failures:
|
||||
done.set()
|
||||
return None
|
||||
done.wait(poll_interval)
|
||||
return None
|
||||
|
||||
|
||||
def _pod_forwarded_args(
|
||||
argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = ()
|
||||
) -> list[str]:
|
||||
"""User CLI overrides to replay on the pod, minus flags the submitter sets itself.
|
||||
|
||||
Handles both `--name=value` and `--name value` forms. Forwarding the user's overrides (e.g.
|
||||
`--steps`, `--save_checkpoint_to_hub`) makes a remote resume behave like the same local command.
|
||||
"""
|
||||
out: list[str] = []
|
||||
skip_next = False
|
||||
for i, tok in enumerate(argv):
|
||||
if skip_next:
|
||||
skip_next = False
|
||||
continue
|
||||
name = tok.split("=", 1)[0]
|
||||
if name in drop_names or any(name.startswith(p) for p in drop_prefixes):
|
||||
if "=" not in tok and i + 1 < len(argv) and not argv[i + 1].startswith("--"):
|
||||
skip_next = True # also drop the space-separated value
|
||||
continue
|
||||
out.append(tok)
|
||||
return out
|
||||
|
||||
|
||||
def _build_resume_job(cfg: TrainPipelineConfig, username: str) -> tuple[str, list[str]]:
|
||||
"""Resolve the model repo and pod command to resume a run on a job.
|
||||
|
||||
A Hub `config_path` is resumed from directly: its checkpoint config already targets that repo,
|
||||
so new checkpoints continue the lineage there. A local `config_path` has its checkpoint uploaded
|
||||
to a new PRIVATE repo first, and the resumed run is forced to push back to it. The pod command
|
||||
always carries `--job.target=local` so the checkpoint's saved `job.target` can't make the pod
|
||||
re-dispatch itself.
|
||||
"""
|
||||
config_path = parser.parse_arg("config_path")
|
||||
forwarded = _pod_forwarded_args(
|
||||
sys.argv[1:],
|
||||
drop_names=("--config_path", "--policy.repo_id", "--policy.push_to_hub", "--dataset.root"),
|
||||
drop_prefixes=("--job.",),
|
||||
)
|
||||
|
||||
if Path(config_path).exists():
|
||||
# Local checkpoint: stage it on the Hub so the pod can resume from it, and push back there.
|
||||
# Resolve so a `last` symlink uploads under its real step name (digit), which the pod's
|
||||
# latest-checkpoint lookup keys on.
|
||||
checkpoint_dir = Path(cfg.checkpoint_path).resolve()
|
||||
source_repo = build_repo_id(username, cfg.job_name or "train", dt.datetime.now(dt.UTC))
|
||||
push_checkpoint_to_hub(checkpoint_dir, source_repo, private=True)
|
||||
extra = [f"--policy.repo_id={source_repo}", "--policy.push_to_hub=true"]
|
||||
else:
|
||||
source_repo = config_path
|
||||
extra = []
|
||||
|
||||
command = [
|
||||
"lerobot-train",
|
||||
*forwarded,
|
||||
f"--config_path={source_repo}",
|
||||
"--job.target=local",
|
||||
*extra,
|
||||
]
|
||||
return source_repo, command
|
||||
|
||||
|
||||
def submit_to_hf(cfg: TrainPipelineConfig) -> None:
|
||||
"""Submit a training job to HF Jobs infrastructure.
|
||||
|
||||
Validates cfg, resolves credentials, ensures the dataset is on the Hub, then either stages a
|
||||
sanitized config (fresh run) or resumes from a checkpoint repo, submits the job, and tails logs
|
||||
until completion or detaches immediately. Ctrl-C detaches without cancelling the remote job.
|
||||
"""
|
||||
token = get_token()
|
||||
if not token:
|
||||
raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.")
|
||||
|
||||
api = HfApi(token=token)
|
||||
user_info = api.whoami(token=token)
|
||||
username = user_info["name"]
|
||||
|
||||
now = dt.datetime.now(dt.UTC)
|
||||
fresh_repo_id: str | None = None
|
||||
if not cfg.resume:
|
||||
# Resolve the model repo and mark it for push BEFORE validate(): validate() requires repo_id
|
||||
# to be set whenever push_to_hub is True. (A resume reuses the checkpoint's repo instead.)
|
||||
if cfg.policy is not None:
|
||||
base_name = cfg.job_name or cfg.policy.type
|
||||
fresh_repo_id = cfg.policy.repo_id or build_repo_id(username, base_name, now)
|
||||
cfg.policy.repo_id = fresh_repo_id
|
||||
cfg.policy.push_to_hub = True
|
||||
else:
|
||||
# Path-based policy is resolved inside validate(); fall back to a generic slug.
|
||||
fresh_repo_id = build_repo_id(username, cfg.job_name or "train", now)
|
||||
|
||||
cfg.validate()
|
||||
|
||||
if cfg.is_reward_model_training:
|
||||
raise ValueError(
|
||||
"Remote training via --job.target only supports policy training, not reward models. "
|
||||
"Run reward-model training locally."
|
||||
)
|
||||
|
||||
secrets: dict[str, str] = {"HF_TOKEN": token}
|
||||
if cfg.wandb.enable:
|
||||
wandb_key = resolve_wandb_api_key()
|
||||
if wandb_key is None:
|
||||
raise ValueError(
|
||||
"wandb is enabled but no WANDB_API_KEY found. "
|
||||
"Set it via `export WANDB_API_KEY=...` or add it to ~/.netrc."
|
||||
)
|
||||
secrets["WANDB_API_KEY"] = wandb_key
|
||||
|
||||
tags = resolve_job_tags(cfg.job.tags)
|
||||
# The dataset must be reachable from the pod for both fresh and resumed runs; a local-only
|
||||
# dataset is pushed PRIVATE here. Hoisted before the resume/fresh branch since it applies to both.
|
||||
ensure_dataset_available(cfg.dataset.repo_id, api=api, tags=tags)
|
||||
|
||||
if cfg.resume:
|
||||
repo_id, command = _build_resume_job(cfg, username)
|
||||
else:
|
||||
config_repo_id = _stage_config_on_hub(cfg, fresh_repo_id, token, tags=tags)
|
||||
repo_id = fresh_repo_id
|
||||
command = ["lerobot-train", f"--config_path={config_repo_id}"]
|
||||
|
||||
print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...")
|
||||
job_info = run_job(
|
||||
image=cfg.job.image,
|
||||
command=command,
|
||||
flavor=cfg.job.target,
|
||||
secrets=secrets,
|
||||
timeout=cfg.job.timeout,
|
||||
# HF Jobs labels are key/value; expose each tag as a queryable label.
|
||||
labels=dict.fromkeys(tags, "true"),
|
||||
)
|
||||
job_id = job_info.id
|
||||
job_url = getattr(job_info, "url", None)
|
||||
print(f"Job submitted: {job_id}")
|
||||
if job_url:
|
||||
print(f" Job page: {job_url}")
|
||||
print(f" Model repo: https://huggingface.co/{repo_id}")
|
||||
print(f" Monitor: hf jobs logs {job_id}")
|
||||
print(f" Cancel: hf jobs cancel {job_id}")
|
||||
|
||||
if cfg.job.detach:
|
||||
return
|
||||
|
||||
done = threading.Event()
|
||||
detached = threading.Event()
|
||||
pushed_ok = threading.Event()
|
||||
stage_holder: dict[str, str | None] = {}
|
||||
|
||||
def _poll() -> None:
|
||||
stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder)
|
||||
|
||||
poll_thread = threading.Thread(target=_poll, daemon=True)
|
||||
poll_thread.start()
|
||||
# Finish as soon as the model is pushed, rather than waiting out the platform's
|
||||
# post-run finalization before the job stage flips to COMPLETED. This matches the
|
||||
# exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay
|
||||
# in sync. If it ever stops matching we just fall back to stage-based completion
|
||||
# (~30s slower), so the contract is an optimization, not a correctness requirement.
|
||||
success_marker = f"Model pushed to https://huggingface.co/{repo_id}"
|
||||
log_thread = threading.Thread(
|
||||
target=_tail_logs, args=(job_id, done, success_marker, pushed_ok), daemon=True
|
||||
)
|
||||
log_thread.start()
|
||||
|
||||
def _detach(sig, frame):
|
||||
detached.set()
|
||||
done.set()
|
||||
print("\nDetached. Job is still running.")
|
||||
print(f" Monitor: hf jobs logs {job_id}")
|
||||
print(f" Cancel: hf jobs cancel {job_id}")
|
||||
|
||||
# signal.signal only works on the main thread; when called from a worker thread
|
||||
# (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels
|
||||
# handler rather than crashing with ValueError.
|
||||
install_sigint = threading.current_thread() is threading.main_thread()
|
||||
original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None
|
||||
if install_sigint:
|
||||
signal.signal(signal.SIGINT, _detach)
|
||||
try:
|
||||
# Timeout-based join so SIGINT is delivered to the main thread promptly.
|
||||
while poll_thread.is_alive():
|
||||
poll_thread.join(timeout=0.5)
|
||||
log_thread.join(timeout=5)
|
||||
finally:
|
||||
if install_sigint:
|
||||
signal.signal(signal.SIGINT, original_sigint)
|
||||
|
||||
if detached.is_set():
|
||||
return
|
||||
|
||||
if pushed_ok.is_set():
|
||||
print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}")
|
||||
return
|
||||
|
||||
stage = stage_holder.get("stage")
|
||||
if stage != "COMPLETED":
|
||||
message = stage_holder.get("message")
|
||||
detail = f" ({message})" if message else ""
|
||||
raise RuntimeError(
|
||||
f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}"
|
||||
)
|
||||
@@ -18,7 +18,6 @@ from .act.configuration_act import ACTConfig as ACTConfig
|
||||
from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
|
||||
from .eo1.configuration_eo1 import EO1Config as EO1Config
|
||||
from .factory import get_policy_class, make_policy, make_policy_config, make_pre_post_processors
|
||||
from .fastwam.configuration_fastwam import FastWAMConfig as FastWAMConfig
|
||||
from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig as GaussianActorConfig
|
||||
from .groot.configuration_groot import GrootConfig as GrootConfig
|
||||
from .molmoact2.configuration_molmoact2 import MolmoAct2Config as MolmoAct2Config
|
||||
@@ -43,7 +42,6 @@ __all__ = [
|
||||
"ACTConfig",
|
||||
"DiffusionConfig",
|
||||
"EO1Config",
|
||||
"FastWAMConfig",
|
||||
"GaussianActorConfig",
|
||||
"GrootConfig",
|
||||
"MolmoAct2Config",
|
||||
|
||||
@@ -47,7 +47,6 @@ from lerobot.utils.feature_utils import dataset_to_policy_features
|
||||
from .act.configuration_act import ACTConfig
|
||||
from .diffusion.configuration_diffusion import DiffusionConfig
|
||||
from .eo1.configuration_eo1 import EO1Config
|
||||
from .fastwam.configuration_fastwam import FastWAMConfig
|
||||
from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
|
||||
from .groot.configuration_groot import GrootConfig
|
||||
from .molmoact2.configuration_molmoact2 import MolmoAct2Config
|
||||
@@ -163,10 +162,6 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
|
||||
from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
|
||||
|
||||
return VLAJEPAPolicy
|
||||
elif name == "fastwam":
|
||||
from .fastwam.modeling_fastwam import FastWAMPolicy
|
||||
|
||||
return FastWAMPolicy
|
||||
else:
|
||||
try:
|
||||
return _get_policy_cls_from_policy_name(name=name)
|
||||
@@ -223,8 +218,6 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
return MolmoAct2Config(**kwargs)
|
||||
elif policy_type == "vla_jepa":
|
||||
return VLAJEPAConfig(**kwargs)
|
||||
elif policy_type == "fastwam":
|
||||
return FastWAMConfig(**kwargs)
|
||||
else:
|
||||
try:
|
||||
config_cls = PreTrainedConfig.get_choice_class(policy_type)
|
||||
@@ -458,14 +451,6 @@ def make_pre_post_processors(
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, FastWAMConfig):
|
||||
from .fastwam.processor_fastwam import make_fastwam_pre_post_processors
|
||||
|
||||
processors = make_fastwam_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
else:
|
||||
try:
|
||||
processors = _make_processors_from_policy_config(
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
../../../../docs/source/policy_fastwam_README.md
|
||||
@@ -1,23 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .configuration_fastwam import FastWAMConfig
|
||||
from .modeling_fastwam import FastWAMPolicy
|
||||
from .processor_fastwam import make_fastwam_pre_post_processors
|
||||
|
||||
__all__ = [
|
||||
"FastWAMConfig",
|
||||
"FastWAMPolicy",
|
||||
"make_fastwam_pre_post_processors",
|
||||
]
|
||||
@@ -1,399 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from lerobot.configs import (
|
||||
FeatureType,
|
||||
NormalizationMode,
|
||||
PolicyFeature,
|
||||
PreTrainedConfig,
|
||||
)
|
||||
from lerobot.optim import AdamWConfig
|
||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||
|
||||
WAN22_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B"
|
||||
WAN22_DIFFUSERS_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
|
||||
FASTWAM_BASE_MODEL_ID = "lerobot/fastwam_base"
|
||||
WAN_T5_TOKENIZER_ID = "google/umt5-xxl"
|
||||
|
||||
|
||||
_FASTWAM_VIDEO_BASE_COMPAT_KEYS = (
|
||||
"patch_size",
|
||||
"in_dim",
|
||||
"hidden_dim",
|
||||
"ffn_dim",
|
||||
"freq_dim",
|
||||
"text_dim",
|
||||
"out_dim",
|
||||
"num_heads",
|
||||
"attn_head_dim",
|
||||
"num_layers",
|
||||
)
|
||||
|
||||
_FASTWAM_ACTION_BASE_COMPAT_KEYS = (
|
||||
"hidden_dim",
|
||||
"ffn_dim",
|
||||
"num_heads",
|
||||
"attn_head_dim",
|
||||
"num_layers",
|
||||
"text_dim",
|
||||
"freq_dim",
|
||||
)
|
||||
|
||||
|
||||
def default_video_dit_config(action_dim: int) -> dict[str, Any]:
|
||||
return {
|
||||
"patch_size": [1, 2, 2],
|
||||
"in_dim": 48,
|
||||
"hidden_dim": 3072,
|
||||
"ffn_dim": 14336,
|
||||
"freq_dim": 256,
|
||||
"text_dim": 4096,
|
||||
"out_dim": 48,
|
||||
"num_heads": 24,
|
||||
"attn_head_dim": 128,
|
||||
"num_layers": 30,
|
||||
"eps": 1.0e-6,
|
||||
"seperated_timestep": True,
|
||||
"use_gradient_checkpointing": False,
|
||||
"video_attention_mask_mode": "first_frame_causal",
|
||||
"action_conditioned": False,
|
||||
"action_dim": action_dim,
|
||||
"action_group_causal_mask_mode": "group_diagonal",
|
||||
"fp32_attention": True,
|
||||
}
|
||||
|
||||
|
||||
def default_action_dit_config(action_dim: int) -> dict[str, Any]:
|
||||
return {
|
||||
"action_dim": action_dim,
|
||||
"hidden_dim": 1024,
|
||||
"ffn_dim": 4096,
|
||||
"num_heads": 24,
|
||||
"attn_head_dim": 128,
|
||||
"num_layers": 30,
|
||||
"text_dim": 4096,
|
||||
"freq_dim": 256,
|
||||
"eps": 1.0e-6,
|
||||
"use_gradient_checkpointing": False,
|
||||
"fp32_attention": True,
|
||||
}
|
||||
|
||||
|
||||
def _coerce_enum(enum_cls: type, value: Any) -> Any:
|
||||
if isinstance(value, enum_cls):
|
||||
return value
|
||||
try:
|
||||
return enum_cls(value)
|
||||
except (TypeError, ValueError) as exc:
|
||||
member = getattr(enum_cls, str(value), None)
|
||||
if member is None:
|
||||
raise ValueError(f"Cannot coerce {value!r} into {enum_cls.__name__}.") from exc
|
||||
return member
|
||||
|
||||
|
||||
def _coerce_policy_features(features: dict[str, Any] | None) -> dict[str, PolicyFeature] | None:
|
||||
if features is None:
|
||||
return None
|
||||
coerced = {}
|
||||
for name, feature in features.items():
|
||||
if isinstance(feature, PolicyFeature):
|
||||
coerced[name] = feature
|
||||
continue
|
||||
coerced[name] = PolicyFeature(
|
||||
type=_coerce_enum(FeatureType, feature["type"]),
|
||||
shape=tuple(feature["shape"]),
|
||||
)
|
||||
return coerced
|
||||
|
||||
|
||||
def _is_local_model_id(value: str) -> bool:
|
||||
path = Path(value).expanduser()
|
||||
return path.is_absolute() or value.startswith(("./", "../", "~")) or path.exists()
|
||||
|
||||
|
||||
def _validate_wan_model_id(value: str, field_name: str) -> str:
|
||||
if value == WAN22_MODEL_ID or _is_local_model_id(value):
|
||||
return value
|
||||
raise ValueError(f"`{field_name}` must be `{WAN22_MODEL_ID}` or an explicit local path, got `{value}`.")
|
||||
|
||||
|
||||
def is_fastwam_base_compatible_config(config: FastWAMConfig) -> bool:
|
||||
"""Return whether `fastwam_base` partial weights can initialize this config."""
|
||||
|
||||
default_video_config = default_video_dit_config(config.action_dim)
|
||||
default_action_config = default_action_dit_config(config.action_dim)
|
||||
return all(
|
||||
config.video_dit_config.get(key) == default_video_config.get(key)
|
||||
for key in _FASTWAM_VIDEO_BASE_COMPAT_KEYS
|
||||
) and all(
|
||||
config.action_dit_config.get(key) == default_action_config.get(key)
|
||||
for key in _FASTWAM_ACTION_BASE_COMPAT_KEYS
|
||||
)
|
||||
|
||||
|
||||
@PreTrainedConfig.register_subclass("fastwam")
|
||||
@dataclass
|
||||
class FastWAMConfig(PreTrainedConfig):
|
||||
"""Configuration for the FastWAM LeRobot policy.
|
||||
|
||||
Args:
|
||||
action_dim (int): Number of scalar action channels per timestep.
|
||||
proprio_dim (int | None): Number of proprioception channels used as an
|
||||
extra text-context token. `None` disables proprio conditioning.
|
||||
action_horizon (int): Number of actions predicted by one policy call.
|
||||
num_video_frames (int): Raw video sampling window (in dataset frames). The
|
||||
model actually operates on `model_video_frames` frames after subsampling
|
||||
by `action_video_freq_ratio`.
|
||||
action_video_freq_ratio (int): Actions are sampled at this multiple of the
|
||||
video frame rate. Video frames are taken every `action_video_freq_ratio`-th
|
||||
raw frame, so the model sees `(num_video_frames - 1) // ratio + 1` frames
|
||||
spanning the same time window as `action_horizon` actions (ratio actions
|
||||
per video frame).
|
||||
image_size (tuple[int, int]): Concatenated image size as `(height, width)`.
|
||||
context_len (int): Maximum text embedding token length.
|
||||
video_dit_config (dict[str, Any] | None): Wan video expert config.
|
||||
action_dit_config (dict[str, Any] | None): Action expert config.
|
||||
use_gradient_checkpointing (bool): Enable activation checkpointing in both DiT
|
||||
experts (trades compute for memory; propagated into the DiT configs).
|
||||
freeze_video_expert (bool): Freeze the ~5B Wan video expert
|
||||
(`model.video_expert`) so only the action expert + proprio encoder train.
|
||||
Cuts the AdamW optimizer footprint substantially; the video expert keeps its
|
||||
pretrained weights. (If enabled, also set `loss.lambda_video=0` to skip the
|
||||
now-gradient-free video loss compute.)
|
||||
"""
|
||||
|
||||
n_obs_steps: int = 1
|
||||
action_dim: int = 7
|
||||
proprio_dim: int | None = 8
|
||||
action_horizon: int = 32
|
||||
n_action_steps: int = 32
|
||||
num_video_frames: int = 33
|
||||
action_video_freq_ratio: int = 4
|
||||
image_size: tuple[int, int] = (224, 448)
|
||||
context_len: int = 128
|
||||
model_id: str = WAN22_MODEL_ID
|
||||
tokenizer_model_id: str = WAN_T5_TOKENIZER_ID
|
||||
text_encoder_model_id: str = WAN22_DIFFUSERS_MODEL_ID
|
||||
base_model_id: str | None = FASTWAM_BASE_MODEL_ID
|
||||
tokenizer_max_len: int = 128
|
||||
load_text_encoder: bool = True
|
||||
mot_checkpoint_mixed_attn: bool = False
|
||||
torch_dtype: str = "bfloat16"
|
||||
prompt_template: str = (
|
||||
"A video recorded from a robot's point of view executing the following instruction: {task}"
|
||||
)
|
||||
num_inference_steps: int = 10
|
||||
inference_seed: int | None = 42
|
||||
rand_device: str = "cpu"
|
||||
text_cfg_scale: float = 1.0
|
||||
negative_prompt: str = ""
|
||||
sigma_shift: float | None = None
|
||||
tiled: bool = False
|
||||
fp32_attention: bool = True
|
||||
use_gradient_checkpointing: bool = False
|
||||
freeze_video_expert: bool = False
|
||||
toggle_action_dimensions: list[int] = field(default_factory=list)
|
||||
video_scheduler: dict[str, float | int] = field(
|
||||
default_factory=lambda: {"train_shift": 5.0, "infer_shift": 5.0, "num_train_timesteps": 1000}
|
||||
)
|
||||
action_scheduler: dict[str, float | int] = field(
|
||||
default_factory=lambda: {"train_shift": 5.0, "infer_shift": 5.0, "num_train_timesteps": 1000}
|
||||
)
|
||||
loss: dict[str, float] = field(default_factory=lambda: {"lambda_video": 1.0, "lambda_action": 1.0})
|
||||
video_dit_config: dict[str, Any] | None = None
|
||||
action_dit_config: dict[str, Any] | None = None
|
||||
normalization_mapping: dict[str, NormalizationMode] = field(
|
||||
default_factory=lambda: {
|
||||
"VISUAL": NormalizationMode.IDENTITY,
|
||||
"STATE": NormalizationMode.MEAN_STD,
|
||||
"ACTION": NormalizationMode.MEAN_STD,
|
||||
}
|
||||
)
|
||||
input_features: dict[str, PolicyFeature] | None = None
|
||||
output_features: dict[str, PolicyFeature] | None = None
|
||||
optimizer_lr: float = 1.0e-4
|
||||
optimizer_weight_decay: float = 1.0e-2
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
self.image_size = tuple(self.image_size)
|
||||
self.model_id = _validate_wan_model_id(self.model_id, "model_id")
|
||||
self.input_features = _coerce_policy_features(self.input_features)
|
||||
self.output_features = _coerce_policy_features(self.output_features)
|
||||
self.toggle_action_dimensions = [int(dim) for dim in self.toggle_action_dimensions]
|
||||
self.video_dit_config = self.video_dit_config or default_video_dit_config(self.action_dim)
|
||||
self.action_dit_config = self.action_dit_config or default_action_dit_config(self.action_dim)
|
||||
self.video_dit_config["fp32_attention"] = bool(self.fp32_attention)
|
||||
self.action_dit_config["fp32_attention"] = bool(self.fp32_attention)
|
||||
self.video_dit_config["use_gradient_checkpointing"] = bool(self.use_gradient_checkpointing)
|
||||
self.action_dit_config["use_gradient_checkpointing"] = bool(self.use_gradient_checkpointing)
|
||||
if self.input_features is None:
|
||||
height, width = self.image_size
|
||||
self.input_features = {
|
||||
"observation.images.image": PolicyFeature(
|
||||
type=FeatureType.VISUAL,
|
||||
shape=(3, height, width),
|
||||
)
|
||||
}
|
||||
if self.proprio_dim is not None:
|
||||
self.input_features[OBS_STATE] = PolicyFeature(
|
||||
type=FeatureType.STATE,
|
||||
shape=(self.proprio_dim,),
|
||||
)
|
||||
if self.output_features is None:
|
||||
self.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(self.action_dim,))}
|
||||
self.validate_features()
|
||||
if self.pretrained_path or self.use_peft or not self.base_model_id:
|
||||
return
|
||||
if not is_fastwam_base_compatible_config(self):
|
||||
return
|
||||
self.pretrained_path = Path(self.base_model_id)
|
||||
self._auto_pretrained_path = True
|
||||
|
||||
def _save_pretrained(self, save_directory: Path) -> None:
|
||||
if not getattr(self, "_auto_pretrained_path", False):
|
||||
super()._save_pretrained(save_directory)
|
||||
return
|
||||
|
||||
pretrained_path = self.pretrained_path
|
||||
self.pretrained_path = None
|
||||
try:
|
||||
super()._save_pretrained(save_directory)
|
||||
finally:
|
||||
self.pretrained_path = pretrained_path
|
||||
|
||||
def get_optimizer_preset(self) -> AdamWConfig:
|
||||
return AdamWConfig(lr=self.optimizer_lr, weight_decay=self.optimizer_weight_decay)
|
||||
|
||||
def get_scheduler_preset(self) -> None:
|
||||
return None
|
||||
|
||||
def set_dataset_feature_metadata(self, dataset_features: dict[str, Any]) -> None:
|
||||
"""Rebuild visual input features from the dataset's real camera keys.
|
||||
|
||||
FastWAM's `__post_init__` installs a synthetic single-image default
|
||||
(`observation.images.image` at full `image_size` width). For datasets
|
||||
with one or more separately-named cameras (e.g. `observation.images.top`,
|
||||
`observation.images.wrist`), this hook — invoked by `make_policy` once the
|
||||
dataset metadata is known — replaces that default with the actual camera
|
||||
keys, each declared at the policy's native per-camera resolution
|
||||
(`image_size[0]` x `image_size[1] // num_cameras`). The accompanying
|
||||
resize step in `make_fastwam_pre_post_processors` resizes raw frames to
|
||||
match, so heterogeneous source resolutions (e.g. 480x640) are supported.
|
||||
"""
|
||||
image_keys = sorted(
|
||||
key
|
||||
for key, feature in dataset_features.items()
|
||||
if key.startswith("observation.images.") and feature.get("dtype") in ("video", "image")
|
||||
)
|
||||
if not image_keys:
|
||||
return
|
||||
height, total_width = self.image_size
|
||||
per_cam_width = total_width // len(image_keys)
|
||||
new_inputs: dict[str, PolicyFeature] = {
|
||||
key: PolicyFeature(type=FeatureType.VISUAL, shape=(3, height, per_cam_width))
|
||||
for key in image_keys
|
||||
}
|
||||
if self.proprio_dim is not None and OBS_STATE in dataset_features:
|
||||
new_inputs[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(self.proprio_dim,))
|
||||
self.input_features = new_inputs
|
||||
self.validate_features()
|
||||
|
||||
def validate_features(self) -> None:
|
||||
if self.action_dim <= 0:
|
||||
raise ValueError(f"`action_dim` must be positive, got {self.action_dim}.")
|
||||
if self.action_horizon <= 0:
|
||||
raise ValueError(f"`action_horizon` must be positive, got {self.action_horizon}.")
|
||||
if self.n_action_steps > self.action_horizon:
|
||||
raise ValueError("`n_action_steps` cannot exceed `action_horizon`.")
|
||||
if self.action_video_freq_ratio <= 0:
|
||||
raise ValueError(
|
||||
f"`action_video_freq_ratio` must be positive, got {self.action_video_freq_ratio}."
|
||||
)
|
||||
# Video frames are subsampled by action_video_freq_ratio; the resulting model frame
|
||||
# count must satisfy T % 4 == 1 for the VAE temporal tokenization (mirrors the
|
||||
# original FastWAM dataset asserts).
|
||||
if (self.num_video_frames - 1) % self.action_video_freq_ratio != 0:
|
||||
raise ValueError(
|
||||
f"`num_video_frames - 1` ({self.num_video_frames - 1}) must be divisible by "
|
||||
f"`action_video_freq_ratio` ({self.action_video_freq_ratio})."
|
||||
)
|
||||
if ((self.num_video_frames - 1) // self.action_video_freq_ratio) % 4 != 0:
|
||||
raise ValueError(
|
||||
f"Subsampled video transitions ({(self.num_video_frames - 1) // self.action_video_freq_ratio}) "
|
||||
"must be divisible by 4 for VAE tokenization (i.e. model_video_frames % 4 == 1)."
|
||||
)
|
||||
if self.action_horizon % ((self.num_video_frames - 1) // self.action_video_freq_ratio) != 0:
|
||||
raise ValueError(
|
||||
f"`action_horizon` ({self.action_horizon}) must be divisible by the number of "
|
||||
f"video transitions ({(self.num_video_frames - 1) // self.action_video_freq_ratio})."
|
||||
)
|
||||
if not self.image_features:
|
||||
raise ValueError("FastWAM requires at least one image feature.")
|
||||
if self.action_feature is None:
|
||||
raise ValueError("FastWAM requires `action` in output_features.")
|
||||
action_shape = tuple(self.action_feature.shape)
|
||||
if action_shape != (self.action_dim,):
|
||||
raise ValueError(
|
||||
f"FastWAM action feature shape must be ({self.action_dim},), got {action_shape}."
|
||||
)
|
||||
if self.proprio_dim is not None:
|
||||
state_feature = self.robot_state_feature
|
||||
if state_feature is None:
|
||||
raise ValueError("FastWAM requires `observation.state` when `proprio_dim` is set.")
|
||||
state_shape = tuple(state_feature.shape)
|
||||
if state_shape != (self.proprio_dim,):
|
||||
raise ValueError(
|
||||
f"FastWAM state feature shape must be ({self.proprio_dim},), got {state_shape}."
|
||||
)
|
||||
height, width = self.image_size
|
||||
image_width_sum = 0
|
||||
for name, feature in self.image_features.items():
|
||||
shape = tuple(feature.shape)
|
||||
if len(shape) != 3 or shape[0] != 3:
|
||||
raise ValueError(f"FastWAM image feature `{name}` must have shape (3, H, W), got {shape}.")
|
||||
if shape[1] != height:
|
||||
raise ValueError(f"FastWAM image feature `{name}` height must be {height}, got {shape[1]}.")
|
||||
image_width_sum += shape[2]
|
||||
if image_width_sum != width:
|
||||
raise ValueError(f"FastWAM image feature widths must sum to {width}, got {image_width_sum}.")
|
||||
|
||||
@property
|
||||
def model_video_frames(self) -> int:
|
||||
"""Number of video frames the model actually operates on, after subsampling the
|
||||
raw `num_video_frames` window by `action_video_freq_ratio` (e.g. 33 -> 9)."""
|
||||
return (self.num_video_frames - 1) // self.action_video_freq_ratio + 1
|
||||
|
||||
@property
|
||||
def observation_delta_indices(self) -> list[int]:
|
||||
# Load the video frames the model is supervised on: the future window subsampled by
|
||||
# action_video_freq_ratio (e.g. [0, 4, 8, ..., 32] -> 9 frames). Each video frame is
|
||||
# thus `action_video_freq_ratio` actions apart, while actions load at the full rate
|
||||
# (`action_delta_indices` = range(action_horizon)). Returning None would load only the
|
||||
# current frame, making the video target a static repeat (degenerate supervision).
|
||||
return list(range(0, self.num_video_frames, self.action_video_freq_ratio))
|
||||
|
||||
@property
|
||||
def action_delta_indices(self) -> list[int]:
|
||||
return list(range(self.action_horizon))
|
||||
|
||||
@property
|
||||
def reward_delta_indices(self) -> None:
|
||||
return None
|
||||
@@ -1,440 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections import deque
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
from lerobot.utils.import_utils import require_package
|
||||
|
||||
from .configuration_fastwam import FastWAMConfig
|
||||
from .wan import (
|
||||
ActionDiT,
|
||||
FastWAM,
|
||||
MoT,
|
||||
WanVideoDiT,
|
||||
build_wan_tokenizer,
|
||||
load_pretrained_wan_text_encoder,
|
||||
load_pretrained_wan_vae,
|
||||
)
|
||||
|
||||
|
||||
class FastWAMPolicy(PreTrainedPolicy):
|
||||
"""LeRobot policy wrapper for FastWAM.
|
||||
|
||||
Attention backend: FastWAM's DiT uses ``torch.nn.functional.scaled_dot_product_attention``
|
||||
(SDPA) for all attention. It does not use FlashAttention, because MoT routing requires
|
||||
arbitrary boolean ``[query, key]`` masks that the FlashAttention varlen API cannot express;
|
||||
installing ``flash-attn`` has no effect on the FastWAM path. (SDPA may still dispatch to
|
||||
PyTorch's own flash/mem-efficient/math kernel internally, unrelated to the ``flash-attn`` package.)
|
||||
|
||||
Args:
|
||||
config (FastWAMConfig): FastWAM policy configuration.
|
||||
dataset_stats (dict[str, dict[str, Tensor]] | None): Optional LeRobot
|
||||
dataset statistics passed by the training/evaluation stack.
|
||||
"""
|
||||
|
||||
config_class = FastWAMConfig
|
||||
name = "fastwam"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: FastWAMConfig,
|
||||
dataset_stats: dict[str, dict[str, Tensor]] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
# FastWAM's Wan2.2 backbone needs transformers (UMT5 text encoder/tokenizer) and
|
||||
# diffusers (Wan VAE), both behind the `fastwam` extra. Fail fast with an actionable
|
||||
# message in base installs rather than deep in Wan component construction.
|
||||
require_package("transformers", extra="fastwam")
|
||||
require_package("diffusers", extra="fastwam")
|
||||
# `make_policy`/`from_pretrained` forward extra kwargs (e.g. `dataset_meta`); the
|
||||
# dataset feature metadata is already applied to `config` by make_policy upstream,
|
||||
# so we accept and ignore them, matching the other LeRobot policies.
|
||||
super().__init__(config, dataset_stats)
|
||||
config.validate_features()
|
||||
self.config = config
|
||||
self.dataset_stats = dataset_stats
|
||||
self.model = self._build_core_model(config)
|
||||
if config.freeze_video_expert and getattr(self.model, "video_expert", None) is not None:
|
||||
# Freeze the ~5B Wan video expert; get_optim_params filters on requires_grad,
|
||||
# so its params drop out of the optimizer (and DDP skips them).
|
||||
self.model.video_expert.requires_grad_(False)
|
||||
# The transformer blocks are re-parented onto the MoTLayers (single FSDP owner), so
|
||||
# `video_expert.requires_grad_` no longer reaches them — freeze them via the layers.
|
||||
mot = getattr(self.model, "mot", None)
|
||||
if mot is not None and getattr(mot, "layers", None) is not None:
|
||||
for layer in mot.layers:
|
||||
if "video" in layer.blocks:
|
||||
layer.blocks["video"].requires_grad_(False)
|
||||
self.reset()
|
||||
|
||||
@classmethod
|
||||
def _load_as_safetensor(cls, model, model_file: str, map_location: str, strict: bool):
|
||||
"""Shape-aware load that supports cross-embodiment fine-tuning.
|
||||
|
||||
`safetensors.load_model(strict=False)` ignores missing/unexpected keys but
|
||||
still raises on a shape mismatch for a shared key. When fine-tuning from a
|
||||
checkpoint trained on a different embodiment (e.g. the LIBERO 7-DoF / 8-dim
|
||||
checkpoint adapted to a 6-DoF / 6-dim arm), the action encoder/head and
|
||||
proprio encoder legitimately differ in shape. With `strict=False` we drop
|
||||
only those shape-mismatched tensors — leaving them at their freshly
|
||||
initialized values — and load every compatible tensor. With `strict=True`
|
||||
the standard exact-match loader is used.
|
||||
"""
|
||||
from safetensors import safe_open
|
||||
|
||||
model_state_dict = model.state_dict()
|
||||
mismatched = []
|
||||
with safe_open(model_file, framework="pt") as f:
|
||||
checkpoint_keys = list(f.keys())
|
||||
for key in checkpoint_keys:
|
||||
if key in model_state_dict and tuple(model_state_dict[key].shape) != tuple(
|
||||
f.get_slice(key).get_shape()
|
||||
):
|
||||
mismatched.append(key)
|
||||
|
||||
if not mismatched:
|
||||
return super()._load_as_safetensor(model, model_file, map_location, strict)
|
||||
if strict:
|
||||
raise RuntimeError(
|
||||
f"FastWAM: {len(mismatched)} checkpoint tensors have a shape mismatch under "
|
||||
f"strict=True: {mismatched}"
|
||||
)
|
||||
|
||||
from safetensors.torch import load_file
|
||||
|
||||
logging.warning(
|
||||
"FastWAM cross-embodiment load: reinitializing %d shape-mismatched tensor(s), keeping "
|
||||
"every compatible weight: %s",
|
||||
len(mismatched),
|
||||
mismatched,
|
||||
)
|
||||
state_dict = load_file(model_file, device="cpu")
|
||||
for key in mismatched:
|
||||
state_dict.pop(key, None)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
if map_location and map_location != "cpu":
|
||||
model.to(map_location)
|
||||
return model
|
||||
|
||||
def get_optim_params(self) -> list[Tensor]:
|
||||
# Return the trainable tensors directly (a single param group). The optimizer
|
||||
# builder wraps these in a param group; returning a bare {"params": [...]} dict
|
||||
# instead would make `list(...)` yield the key string "params".
|
||||
params = (
|
||||
list(self.model.dit.parameters()) if hasattr(self.model, "dit") else list(self.model.parameters())
|
||||
)
|
||||
proprio_encoder = getattr(self.model, "proprio_encoder", None)
|
||||
if proprio_encoder is not None:
|
||||
params.extend(list(proprio_encoder.parameters()))
|
||||
return [p for p in params if p.requires_grad]
|
||||
|
||||
def reset(self) -> None:
|
||||
self._action_queue: deque[Tensor] = deque([], maxlen=self.config.n_action_steps)
|
||||
|
||||
def _batch_to_training_sample(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
|
||||
"""Adapt a standard LeRobot batch to the FastWAM-native sample that
|
||||
`FastWAM.build_inputs` consumes (`video`, `action`, `context`/`context_mask`,
|
||||
per-frame `proprio`).
|
||||
|
||||
The LeRobot training loop passes raw `observation.images.*`, a single-step
|
||||
`observation.state` `[B, D]`, `action`, and a language `task` string. We do
|
||||
only the translation `build_inputs` can't: stack the camera frames into a
|
||||
video, encode the prompt with the (frozen) text encoder (mirroring inference,
|
||||
so language-conditioned datasets need no precomputed context), and give proprio
|
||||
the per-frame axis `build_inputs` indexes. All shape/presence validation is
|
||||
left to `build_inputs`, the single authority on the contract.
|
||||
"""
|
||||
sample = dict(batch)
|
||||
if "video" not in sample:
|
||||
sample["video"] = _stack_video_from_images(batch, self.config)
|
||||
if "context" not in sample or "context_mask" not in sample:
|
||||
prompt = _prompt_from_batch(batch=batch, config=self.config)
|
||||
if prompt is None:
|
||||
raise KeyError(
|
||||
"FastWAM training requires a `task`/`prompt` to encode text context, "
|
||||
"or precomputed `context`/`context_mask` in the batch."
|
||||
)
|
||||
sample["context"], sample["context_mask"] = self.model.encode_prompt(prompt)
|
||||
if self.config.proprio_dim is not None and "proprio" not in sample:
|
||||
state = sample.get(OBS_STATE)
|
||||
if state is not None:
|
||||
# LeRobot gives a single-step state [B, D]; build_inputs expects
|
||||
# per-frame [B, T, D] and uses frame 0, so add a T=1 axis.
|
||||
sample["proprio"] = state.unsqueeze(1) if state.ndim == 2 else state
|
||||
return sample
|
||||
|
||||
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]:
|
||||
"""Compute FastWAM training loss for a LeRobot batch.
|
||||
|
||||
Args:
|
||||
batch (dict[str, Tensor]): Batch containing FastWAM-ready keys
|
||||
(`video`, `action`, `context`, `context_mask`) or LeRobot keys
|
||||
that can be adapted (`observation.images.*`, `observation.state`,
|
||||
`action`, `action_is_pad`).
|
||||
|
||||
Returns:
|
||||
tuple[Tensor, dict[str, Any]]: The scalar loss to backprop, and a dict of
|
||||
logging metrics (e.g. `loss_video`, `loss_action`) — the `(loss, output_dict)`
|
||||
contract the LeRobot training loop expects.
|
||||
"""
|
||||
|
||||
sample = self._batch_to_training_sample(batch)
|
||||
loss, metrics = self.model.training_loss(sample)
|
||||
return loss, dict(metrics or {})
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_action_chunk(self, batch: dict[str, Tensor], **_: Any) -> Tensor:
|
||||
"""Predict a chunk of actions from the current FastWAM observation.
|
||||
|
||||
Args:
|
||||
batch (dict[str, Tensor]): Inference batch with `input_image` or
|
||||
image observation keys, plus `context/context_mask` or `prompt`.
|
||||
|
||||
Returns:
|
||||
Tensor: Action chunk with shape `[B, action_horizon, action_dim]`.
|
||||
"""
|
||||
|
||||
self.eval()
|
||||
infer_kwargs = _batch_to_infer_kwargs(batch=batch, config=self.config)
|
||||
batch_size = _infer_kwargs_batch_size(infer_kwargs)
|
||||
if batch_size == 1:
|
||||
action = _action_from_model_output(self.model.infer_action(**infer_kwargs))
|
||||
else:
|
||||
action = torch.cat(
|
||||
[
|
||||
_action_from_model_output(
|
||||
self.model.infer_action(
|
||||
**_slice_infer_kwargs(infer_kwargs, index=i, batch_size=batch_size)
|
||||
)
|
||||
)
|
||||
for i in range(batch_size)
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
return action.to(device=batch_device(batch), dtype=torch.float32)
|
||||
|
||||
@torch.no_grad()
|
||||
def select_action(self, batch: dict[str, Tensor], **kwargs: Any) -> Tensor:
|
||||
self.eval()
|
||||
if len(self._action_queue) == 0:
|
||||
actions = self.predict_action_chunk(batch, **kwargs)[:, : self.config.n_action_steps]
|
||||
self._action_queue.extend(actions.transpose(0, 1))
|
||||
return self._action_queue.popleft()
|
||||
|
||||
def _build_core_model(self, config: FastWAMConfig) -> FastWAM:
|
||||
"""Build the FastWAM core for training / inference.
|
||||
|
||||
Only the trainable parts (the MoT DiT and the proprio encoder) are
|
||||
materialized empty here and then filled from the policy's
|
||||
`model.safetensors` by the base `from_pretrained`. The *frozen* Wan2.2 VAE
|
||||
and UMT5 text encoder are loaded with their real weights from the
|
||||
`Wan-AI/Wan2.2-TI2V-5B-Diffusers` repo (cached in the HF cache, shared
|
||||
across checkpoints) and are intentionally excluded from `model.safetensors`
|
||||
— see `FastWAM.__init__`. The tokenizer comes from `google/umt5-xxl`.
|
||||
"""
|
||||
dtype = _dtype_from_name(config.torch_dtype)
|
||||
device = config.device
|
||||
video_expert = WanVideoDiT(**config.video_dit_config).to(device=device, dtype=dtype)
|
||||
action_expert = ActionDiT(**config.action_dit_config).to(device=device, dtype=dtype)
|
||||
mot = MoT(
|
||||
mixtures={"video": video_expert, "action": action_expert},
|
||||
mot_checkpoint_mixed_attn=config.mot_checkpoint_mixed_attn,
|
||||
)
|
||||
text_encoder = (
|
||||
load_pretrained_wan_text_encoder(
|
||||
model_id=config.text_encoder_model_id, torch_dtype=dtype, device=device
|
||||
)
|
||||
if config.load_text_encoder
|
||||
else None
|
||||
)
|
||||
return FastWAM(
|
||||
video_expert=video_expert,
|
||||
action_expert=action_expert,
|
||||
mot=mot,
|
||||
vae=load_pretrained_wan_vae(torch_dtype=dtype, device=device),
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=build_wan_tokenizer(
|
||||
model_id=config.tokenizer_model_id, tokenizer_max_len=config.tokenizer_max_len
|
||||
),
|
||||
text_dim=int(config.video_dit_config["text_dim"]),
|
||||
proprio_dim=config.proprio_dim,
|
||||
device=device,
|
||||
torch_dtype=dtype,
|
||||
video_train_shift=float(config.video_scheduler["train_shift"]),
|
||||
video_infer_shift=float(config.video_scheduler["infer_shift"]),
|
||||
video_num_train_timesteps=int(config.video_scheduler["num_train_timesteps"]),
|
||||
action_train_shift=float(config.action_scheduler["train_shift"]),
|
||||
action_infer_shift=float(config.action_scheduler["infer_shift"]),
|
||||
action_num_train_timesteps=int(config.action_scheduler["num_train_timesteps"]),
|
||||
loss_lambda_video=float(config.loss["lambda_video"]),
|
||||
loss_lambda_action=float(config.loss["lambda_action"]),
|
||||
)
|
||||
|
||||
|
||||
def _scalar(value: Any) -> Any:
|
||||
"""Unwrap a 0-/1-element tensor (e.g. from DataLoader collation) to a Python scalar."""
|
||||
return value.item() if isinstance(value, Tensor) else value
|
||||
|
||||
|
||||
def _batch_to_infer_kwargs(batch: dict[str, Tensor], config: FastWAMConfig) -> dict[str, Any]:
|
||||
return {
|
||||
"prompt": _prompt_from_batch(batch=batch, config=config),
|
||||
"input_image": _input_image_from_batch(batch, config),
|
||||
"action_horizon": config.action_horizon,
|
||||
"proprio": batch.get("proprio", batch.get(OBS_STATE)),
|
||||
"context": batch.get("context"),
|
||||
"context_mask": batch.get("context_mask"),
|
||||
"negative_prompt": batch.get("negative_prompt", config.negative_prompt),
|
||||
"text_cfg_scale": float(_scalar(batch.get("text_cfg_scale", config.text_cfg_scale))),
|
||||
"num_inference_steps": int(_scalar(batch.get("num_inference_steps", config.num_inference_steps))),
|
||||
"sigma_shift": batch.get("sigma_shift", config.sigma_shift),
|
||||
"seed": batch.get("seed", config.inference_seed),
|
||||
"rand_device": batch.get("rand_device", config.rand_device),
|
||||
"tiled": bool(batch.get("tiled", config.tiled)),
|
||||
}
|
||||
|
||||
|
||||
def _prompt_from_batch(batch: dict[str, Tensor], config: FastWAMConfig) -> Any:
|
||||
prompt = batch.get("prompt")
|
||||
if prompt is not None:
|
||||
return prompt
|
||||
|
||||
task = batch.get("task")
|
||||
if task is None:
|
||||
return None
|
||||
if isinstance(task, str):
|
||||
return config.prompt_template.format(task=task)
|
||||
if isinstance(task, (list, tuple)):
|
||||
return [config.prompt_template.format(task=str(item)) for item in task]
|
||||
return config.prompt_template.format(task=str(task))
|
||||
|
||||
|
||||
def _action_from_model_output(output: Any) -> Tensor:
|
||||
action = output["action"] if isinstance(output, dict) else output
|
||||
if action.ndim == 2:
|
||||
action = action.unsqueeze(0)
|
||||
return action
|
||||
|
||||
|
||||
def _infer_kwargs_batch_size(infer_kwargs: dict[str, Any]) -> int:
|
||||
image = infer_kwargs["input_image"]
|
||||
if not isinstance(image, Tensor):
|
||||
raise TypeError(f"`input_image` must be a tensor, got {type(image).__name__}.")
|
||||
if image.ndim == 3:
|
||||
return 1
|
||||
if image.ndim == 4:
|
||||
return int(image.shape[0])
|
||||
raise ValueError(f"`input_image` must be [B,C,H,W] or [C,H,W], got {tuple(image.shape)}.")
|
||||
|
||||
|
||||
def _slice_infer_kwargs(infer_kwargs: dict[str, Any], *, index: int, batch_size: int) -> dict[str, Any]:
|
||||
return {
|
||||
key: _slice_infer_value(value, index=index, batch_size=batch_size)
|
||||
for key, value in infer_kwargs.items()
|
||||
}
|
||||
|
||||
|
||||
def _slice_infer_value(value: Any, *, index: int, batch_size: int) -> Any:
|
||||
if isinstance(value, Tensor) and value.ndim > 0 and value.shape[0] == batch_size:
|
||||
return value[index : index + 1]
|
||||
if isinstance(value, (list, tuple)) and len(value) == batch_size:
|
||||
return value[index]
|
||||
return value
|
||||
|
||||
|
||||
def _dtype_from_name(name: str) -> torch.dtype:
|
||||
dtype_map = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}
|
||||
if name not in dtype_map:
|
||||
raise ValueError(f"Unsupported torch dtype `{name}`.")
|
||||
return dtype_map[name]
|
||||
|
||||
|
||||
def batch_device(batch: dict[str, Any]) -> torch.device:
|
||||
for value in batch.values():
|
||||
if isinstance(value, Tensor):
|
||||
return value.device
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
def _resize_frames(frames: Tensor, size: tuple[int, int]) -> Tensor:
|
||||
"""Resize a frame tensor to `size` (H, W), tolerating a leading temporal/batch stack.
|
||||
|
||||
`interpolate` only accepts a single leading batch dim (`[N, C, H, W]`), but FastWAM camera
|
||||
tensors arrive as `[B, C, H, W]` (live eval) or `[B, T, C, H, W]` (temporal stack), so flatten
|
||||
any leading dims into the batch, resize, then restore. A no-op when already at `size`.
|
||||
"""
|
||||
if tuple(frames.shape[-2:]) == size:
|
||||
return frames
|
||||
lead = frames.shape[:-3]
|
||||
flat = frames.reshape(-1, *frames.shape[-3:])
|
||||
flat = torch.nn.functional.interpolate(
|
||||
flat, size=size, mode="bilinear", align_corners=False, antialias=True
|
||||
)
|
||||
return flat.reshape(*lead, *flat.shape[-3:])
|
||||
|
||||
|
||||
def _stack_video_from_images(batch: dict[str, Tensor], config: FastWAMConfig) -> Tensor:
|
||||
# Exclude the `*_is_pad` companion tensors that delta-timestamp loading adds alongside
|
||||
# each camera (shape [B, T]); they share the `observation.images.` prefix but are not frames.
|
||||
image_keys = sorted(k for k in batch if k.startswith("observation.images.") and not k.endswith("_is_pad"))
|
||||
if not image_keys:
|
||||
raise KeyError("FastWAM batch must contain `video` or `observation.images.*` keys.")
|
||||
per_cam = (int(config.image_size[0]), int(config.image_size[1]) // len(image_keys))
|
||||
images = [_resize_frames(batch[key], per_cam) for key in image_keys]
|
||||
# Cameras concatenate along width (last dim) in both the single-frame and temporal case.
|
||||
image = torch.cat(images, dim=-1) if len(images) > 1 else images[0]
|
||||
if image.ndim == 4:
|
||||
# [B, C, H, W]: a single frame (e.g. the live eval observation) -> repeat across time.
|
||||
image = image.unsqueeze(2).repeat(1, 1, config.model_video_frames, 1, 1)
|
||||
elif image.ndim == 5:
|
||||
# [B, T, C, H, W]: temporal stack from delta-timestamp loading -> [B, C, T, H, W].
|
||||
image = image.permute(0, 2, 1, 3, 4)
|
||||
else:
|
||||
raise ValueError(f"Expected image batch [B,C,H,W] or temporal [B,T,C,H,W], got {tuple(image.shape)}.")
|
||||
return image
|
||||
|
||||
|
||||
def _input_image_from_batch(batch: dict[str, Tensor], config: FastWAMConfig) -> Tensor:
|
||||
if "input_image" in batch:
|
||||
return _prepare_infer_image(batch["input_image"], config)
|
||||
video = batch.get("video")
|
||||
if video is None:
|
||||
video = _stack_video_from_images(batch, config)
|
||||
if video.ndim == 5:
|
||||
return _prepare_infer_image(video[:, :, 0], config)
|
||||
if video.ndim == 4:
|
||||
return _prepare_infer_image(video, config)
|
||||
raise ValueError(f"Cannot build input image from tensor with shape {tuple(video.shape)}.")
|
||||
|
||||
|
||||
def _prepare_infer_image(image: Tensor, config: FastWAMConfig) -> Tensor:
|
||||
if image.ndim == 3:
|
||||
image = image.unsqueeze(0)
|
||||
if image.ndim != 4:
|
||||
raise ValueError(f"Expected image tensor [B,C,H,W] or [C,H,W], got {tuple(image.shape)}.")
|
||||
|
||||
# Resize to the full configured resolution (no-op when the video path already produced it, but
|
||||
# also covers a directly-supplied `input_image`). The model owns its input resolution — see
|
||||
# `_stack_video_from_images` — so we resize rather than assert on a mismatch.
|
||||
target_h, target_w = int(config.image_size[0]), int(config.image_size[1])
|
||||
return _resize_frames(image, (target_h, target_w))
|
||||
@@ -1,142 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
ActionProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import (
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
|
||||
from .configuration_fastwam import FastWAMConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
@ProcessorStepRegistry.register(name="fastwam_action_toggle_processor")
|
||||
class FastWAMActionToggleProcessorStep(ActionProcessorStep):
|
||||
"""Apply FastWAM LIBERO toggle semantics to configured action dimensions."""
|
||||
|
||||
toggle_dimensions: list[int]
|
||||
|
||||
def action(self, action: PolicyAction) -> PolicyAction:
|
||||
if not self.toggle_dimensions:
|
||||
return action
|
||||
processed_action = action.clone()
|
||||
action_dim = int(processed_action.shape[-1])
|
||||
for dim in self.toggle_dimensions:
|
||||
resolved_dim = dim if dim >= 0 else action_dim + dim
|
||||
if resolved_dim < 0 or resolved_dim >= action_dim:
|
||||
raise ValueError(
|
||||
f"FastWAM action toggle dimension {dim} is out of bounds for action dim {action_dim}."
|
||||
)
|
||||
value = processed_action[..., resolved_dim]
|
||||
value = value * 2.0 - 1.0
|
||||
processed_action[..., resolved_dim] = torch.sign(-value)
|
||||
return processed_action
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
return {"toggle_dimensions": self.toggle_dimensions}
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
return features
|
||||
|
||||
|
||||
def make_fastwam_pre_post_processors(
|
||||
config: FastWAMConfig,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
|
||||
"""Create LeRobot pre- and post-processing pipelines for FastWAM.
|
||||
|
||||
Args:
|
||||
config (FastWAMConfig): Policy configuration controlling device and
|
||||
normalization feature metadata.
|
||||
dataset_stats (dict[str, dict[str, torch.Tensor]] | None): Optional
|
||||
LeRobot dataset statistics used by normalization processors.
|
||||
|
||||
Returns:
|
||||
tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: Input and
|
||||
output processor pipelines discoverable by LeRobot.
|
||||
"""
|
||||
|
||||
# NOTE: no visual normalization here. VISUAL is IDENTITY (see configuration_fastwam.normalization_mapping)
|
||||
# — images pass through in [0, 1] and the model maps them to the Wan VAE's [-1, 1] at the encode
|
||||
# boundary. This is deliberate: `lerobot_train.py` overrides the normalizer stats with
|
||||
# `dataset.meta.stats` when fine-tuning, and a real dataset's per-channel image std is the tiny
|
||||
# frame-to-frame brightness variance, which would blow images far outside [-1,1] and saturate them.
|
||||
# STATE/ACTION still normalize with dataset stats below.
|
||||
normalization_stats: dict[str, dict[str, Any]] = dict(dataset_stats or {})
|
||||
|
||||
# NOTE: no resize step here. The model is the single authority on input resolution: it resizes
|
||||
# each camera to the per-camera target (image_size split across cameras) in
|
||||
# `_stack_video_from_images` / `_prepare_infer_image`, on every path (train forward, rollout and
|
||||
# eval select_action). A preprocessor resize step would be both redundant (the model re-resizes
|
||||
# anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base
|
||||
# checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide.
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=normalization_stats,
|
||||
device=config.device,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=normalization_stats,
|
||||
),
|
||||
]
|
||||
if config.toggle_action_dimensions:
|
||||
output_steps.append(
|
||||
FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions)
|
||||
)
|
||||
output_steps.append(DeviceProcessorStep(device="cpu"))
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
@@ -1,34 +0,0 @@
|
||||
# FastWAM `wan` package
|
||||
|
||||
This package holds FastWAM's model implementation. It mixes a small **vendored
|
||||
subset of the official Wan2.2 source tree** with FastWAM's own code, kept flat in
|
||||
a single directory.
|
||||
|
||||
## Vendored from Wan2.2
|
||||
|
||||
- Upstream repository: https://github.com/Wan-Video/Wan2.2
|
||||
- Upstream commit: `42bf4cfaa384bc21833865abc2f9e6c0e67233dc`
|
||||
- License: Apache-2.0, matching the license in `LICENSE.txt` from the upstream repository
|
||||
|
||||
Copied files:
|
||||
|
||||
- `model.py` (was `wan/modules/model.py`), trimmed: the flash-attention path
|
||||
(the vendored `attention.py` and the block/model `forward`s) was removed.
|
||||
FastWAM's DiT uses SDPA instead (see `video_dit.py`).
|
||||
- `get_sampling_sigmas` in `video_dit.py` (was `wan/utils/fm_solvers.py`), inlined
|
||||
next to its only caller.
|
||||
|
||||
This subset only backs FastWAM's **custom MoT video DiT**. The Wan2.2 VAE,
|
||||
UMT5 text encoder, and tokenizer are no longer vendored - they come from
|
||||
`diffusers.AutoencoderKLWan`, `transformers.UMT5EncoderModel`, and
|
||||
`transformers.AutoTokenizer` (see `components.py` and `adapters.py`).
|
||||
|
||||
## FastWAM's own code
|
||||
|
||||
- `video_dit.py` builds on `model` (`sinusoidal_embedding_1d`, `rope_params`,
|
||||
`rope_apply`, …) and computes attention with SDPA (`fastwam_masked_attention`). Its
|
||||
`WanContinuousFlowMatchScheduler` uses `get_sampling_sigmas` for Wan-compatible
|
||||
inference timesteps.
|
||||
- `components.py` / `adapters.py` load the VAE, text encoder, tokenizer, and the
|
||||
custom DiT weights.
|
||||
- `modular.py` defines the FastWAM model (`ActionDiT`, `MoT`, `FastWAM`, …).
|
||||
@@ -1,33 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .adapters import WanVideoVAE38
|
||||
from .components import (
|
||||
build_wan_tokenizer,
|
||||
load_pretrained_wan_text_encoder,
|
||||
load_pretrained_wan_vae,
|
||||
)
|
||||
from .modular import ActionDiT, FastWAM, MoT
|
||||
from .video_dit import WanVideoDiT
|
||||
|
||||
__all__ = [
|
||||
"ActionDiT",
|
||||
"FastWAM",
|
||||
"MoT",
|
||||
"WanVideoDiT",
|
||||
"WanVideoVAE38",
|
||||
"build_wan_tokenizer",
|
||||
"load_pretrained_wan_text_encoder",
|
||||
"load_pretrained_wan_vae",
|
||||
]
|
||||
@@ -1,108 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from diffusers import AutoencoderKLWan
|
||||
|
||||
|
||||
class WanVideoVAE38(torch.nn.Module):
|
||||
"""FastWAM VAE contract over `diffusers.AutoencoderKLWan` (Wan2.2-TI2V-5B).
|
||||
|
||||
16x spatial / 4x temporal compression, 48 latent channels. diffusers'
|
||||
`AutoencoderKLWan` returns *raw* latents (it does not apply `latents_mean`/
|
||||
`latents_std`), so `encode`/`decode` here apply the same standardization the
|
||||
Wan reference uses — `(latents - mean) / std` — done in fp32 for stability.
|
||||
`encode` uses the deterministic posterior mode, matching the original VAE
|
||||
which returned the latent mean `mu`.
|
||||
"""
|
||||
|
||||
upsampling_factor = 16
|
||||
temporal_downsample_factor = 4
|
||||
z_dim = 48
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
device: str | torch.device = "cuda",
|
||||
*,
|
||||
pretrained: AutoencoderKLWan,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
# The Wan2.2 VAE is a fixed pretrained model — it is never trained from scratch,
|
||||
# so a real `AutoencoderKLWan` (with weights) must always be supplied (loaded from
|
||||
# the diffusers repo by `load_pretrained_wan_vae`). No random/offline build path.
|
||||
self.vae = pretrained.to(device=device, dtype=dtype)
|
||||
|
||||
# Read the standardization stats from the VAE's own config (diffusers populates
|
||||
# these from vae/config.json) — single source of truth, no local copy. diffusers'
|
||||
# encode/decode return *raw* latents, so we apply (latent - mean) / std ourselves.
|
||||
# Non-persistent: kept out of state_dict.
|
||||
self.register_buffer(
|
||||
"latents_mean",
|
||||
torch.tensor(self.vae.config.latents_mean).view(1, self.z_dim, 1, 1, 1),
|
||||
persistent=False,
|
||||
)
|
||||
self.register_buffer(
|
||||
"latents_std",
|
||||
torch.tensor(self.vae.config.latents_std).view(1, self.z_dim, 1, 1, 1),
|
||||
persistent=False,
|
||||
)
|
||||
|
||||
def _device_dtype(self) -> tuple[torch.device, torch.dtype]:
|
||||
param = next(self.vae.parameters())
|
||||
return param.device, param.dtype
|
||||
|
||||
def encode(
|
||||
self,
|
||||
videos: list[torch.Tensor] | torch.Tensor,
|
||||
device: str | torch.device | None = None,
|
||||
tiled: bool = False,
|
||||
tile_size: tuple[int, int] = (34, 34),
|
||||
tile_stride: tuple[int, int] = (18, 16),
|
||||
) -> torch.Tensor:
|
||||
del device, tile_size, tile_stride
|
||||
if tiled:
|
||||
raise NotImplementedError("Tiled Wan2.2 VAE encoding is not supported by the FastWAM adapter.")
|
||||
if isinstance(videos, (list, tuple)):
|
||||
videos = torch.stack(list(videos))
|
||||
dev, dtype = self._device_dtype()
|
||||
mu = self.vae.encode(videos.to(device=dev, dtype=dtype)).latent_dist.mode().float()
|
||||
mean = self.latents_mean.float().to(mu.device)
|
||||
std = self.latents_std.float().to(mu.device)
|
||||
return (mu - mean) / std
|
||||
|
||||
def decode(
|
||||
self,
|
||||
hidden_states: list[torch.Tensor] | torch.Tensor,
|
||||
device: str | torch.device | None = None,
|
||||
tiled: bool = False,
|
||||
tile_size: tuple[int, int] = (34, 34),
|
||||
tile_stride: tuple[int, int] = (18, 16),
|
||||
) -> torch.Tensor:
|
||||
del device, tile_size, tile_stride
|
||||
if tiled:
|
||||
raise NotImplementedError("Tiled Wan2.2 VAE decoding is not supported by the FastWAM adapter.")
|
||||
if isinstance(hidden_states, (list, tuple)):
|
||||
hidden_states = torch.stack(list(hidden_states))
|
||||
dev, dtype = self._device_dtype()
|
||||
z = hidden_states.float()
|
||||
z = z * self.latents_std.float().to(z.device) + self.latents_mean.float().to(z.device)
|
||||
out = self.vae.decode(z.to(device=dev, dtype=dtype)).sample
|
||||
return out.float().clamp_(-1.0, 1.0)
|
||||
@@ -1,175 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from lerobot.utils.import_utils import _diffusers_available, _transformers_available, require_package
|
||||
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers import AutoTokenizer, UMT5EncoderModel
|
||||
else:
|
||||
AutoTokenizer = None
|
||||
UMT5EncoderModel = None
|
||||
|
||||
if TYPE_CHECKING or _diffusers_available:
|
||||
from diffusers import AutoencoderKLWan
|
||||
else:
|
||||
AutoencoderKLWan = None
|
||||
|
||||
from .adapters import WanVideoVAE38
|
||||
from .video_dit import WanVideoDiT
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# The custom MoT video DiT still ships in the original (non-diffusers) Wan2.2
|
||||
# repo as sharded `diffusion_pytorch_model*.safetensors`; the VAE and UMT5 text
|
||||
# encoder come from the diffusers conversion. Tokenizer is the stock UMT5 one.
|
||||
WAN_DIT_PATTERN = "diffusion_pytorch_model*.safetensors"
|
||||
WAN_T5_TOKENIZER = "google/umt5-xxl"
|
||||
WAN22_DIFFUSERS_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
|
||||
|
||||
|
||||
class WanTextEncoder(torch.nn.Module):
|
||||
"""FastWAM text-encoder contract over `transformers.UMT5EncoderModel`.
|
||||
|
||||
Exposes `.dim` (hidden size) and `forward(ids, mask) -> [B, L, dim]`, matching
|
||||
the call in `FastWAM.encode_prompt`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
device: str | torch.device = "cuda",
|
||||
*,
|
||||
pretrained: torch.nn.Module,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
# UMT5-XXL is a fixed pretrained encoder — never trained from scratch, so a real
|
||||
# `UMT5EncoderModel` (with weights) must always be supplied (loaded from the
|
||||
# diffusers repo by `load_pretrained_wan_text_encoder`). No random/offline build.
|
||||
self.model = pretrained.to(device=device, dtype=dtype)
|
||||
self.dim = int(self.model.config.d_model)
|
||||
|
||||
def forward(self, ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
|
||||
return self.model(input_ids=ids, attention_mask=mask.long()).last_hidden_state
|
||||
|
||||
|
||||
class WanTokenizer:
|
||||
"""UMT5 tokenizer wrapper returning `(input_ids, attention_mask)` like the
|
||||
FastWAM call site expects."""
|
||||
|
||||
def __init__(self, name: str = WAN_T5_TOKENIZER, seq_len: int = 512) -> None:
|
||||
require_package("transformers", extra="fastwam")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(name)
|
||||
self.seq_len = int(seq_len)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
sequence: str | Sequence[str],
|
||||
return_mask: bool = False,
|
||||
add_special_tokens: bool = True,
|
||||
**_: Any,
|
||||
):
|
||||
if isinstance(sequence, str):
|
||||
sequence = [sequence]
|
||||
out = self.tokenizer(
|
||||
list(sequence),
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
max_length=self.seq_len,
|
||||
add_special_tokens=add_special_tokens,
|
||||
return_tensors="pt",
|
||||
)
|
||||
if return_mask:
|
||||
return out.input_ids, out.attention_mask
|
||||
return out.input_ids
|
||||
|
||||
|
||||
def build_wan_tokenizer(*, model_id: str = WAN_T5_TOKENIZER, tokenizer_max_len: int) -> WanTokenizer:
|
||||
return WanTokenizer(name=model_id, seq_len=int(tokenizer_max_len))
|
||||
|
||||
|
||||
def load_pretrained_wan_vae(*, torch_dtype: torch.dtype, device: str) -> WanVideoVAE38:
|
||||
"""Load real Wan2.2 VAE weights from the diffusers repo (offline base creation)."""
|
||||
require_package("diffusers", extra="fastwam")
|
||||
vae = AutoencoderKLWan.from_pretrained(WAN22_DIFFUSERS_MODEL_ID, subfolder="vae", torch_dtype=torch_dtype)
|
||||
return WanVideoVAE38(dtype=torch_dtype, device=device, pretrained=vae)
|
||||
|
||||
|
||||
def load_pretrained_wan_text_encoder(
|
||||
*,
|
||||
model_id: str = WAN22_DIFFUSERS_MODEL_ID,
|
||||
subfolder: str | None = "text_encoder",
|
||||
torch_dtype: torch.dtype,
|
||||
device: str,
|
||||
) -> WanTextEncoder:
|
||||
"""Load UMT5-XXL encoder weights (defaults to the Wan2.2 diffusers repo).
|
||||
|
||||
Must stay compatible with the tokenizer (see `build_wan_tokenizer`): the encoder's
|
||||
embedding table is indexed by the tokenizer's vocabulary.
|
||||
"""
|
||||
require_package("transformers", extra="fastwam")
|
||||
encoder = UMT5EncoderModel.from_pretrained(model_id, subfolder=subfolder, torch_dtype=torch_dtype)
|
||||
return WanTextEncoder(dtype=torch_dtype, device=device, pretrained=encoder)
|
||||
|
||||
|
||||
def resolve_wan_dit_paths(
|
||||
model_id_or_path: str | Path,
|
||||
*,
|
||||
cache_dir: str | Path | None = None,
|
||||
local_files_only: bool = False,
|
||||
revision: str | None = None,
|
||||
) -> list[Path]:
|
||||
"""Resolve the custom MoT DiT shards from the original Wan2.2 repo or a local dir."""
|
||||
path = Path(model_id_or_path).expanduser()
|
||||
if path.is_dir():
|
||||
return sorted(path.glob(WAN_DIT_PATTERN))
|
||||
|
||||
snapshot_path = snapshot_download(
|
||||
repo_id=str(model_id_or_path),
|
||||
revision=revision,
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
allow_patterns=[WAN_DIT_PATTERN],
|
||||
)
|
||||
return sorted(Path(snapshot_path).glob(WAN_DIT_PATTERN))
|
||||
|
||||
|
||||
def load_wan_video_dit(
|
||||
paths: list[str | Path],
|
||||
*,
|
||||
dit_config: dict[str, Any],
|
||||
torch_dtype: torch.dtype,
|
||||
device: str,
|
||||
) -> WanVideoDiT:
|
||||
model = WanVideoDiT(**dit_config)
|
||||
state_dict = _read_wan_dit_safetensors(paths)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
return model.to(device=device, dtype=torch_dtype)
|
||||
|
||||
|
||||
def _read_wan_dit_safetensors(paths: list[str | Path]) -> dict[str, torch.Tensor]:
|
||||
state_dict = {}
|
||||
for path in paths:
|
||||
state_dict.update(load_file(str(path), device="cpu"))
|
||||
return state_dict
|
||||
@@ -1,341 +0,0 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def sinusoidal_embedding_1d(dim, position):
|
||||
# preprocess
|
||||
if dim % 2 != 0:
|
||||
raise ValueError(f"dim must be even, got {dim}.")
|
||||
half = dim // 2
|
||||
position = position.type(torch.float64)
|
||||
|
||||
# calculation
|
||||
sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
|
||||
x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
|
||||
return x
|
||||
|
||||
|
||||
@torch.amp.autocast("cuda", enabled=False)
|
||||
def rope_params(max_seq_len, dim, theta=10000):
|
||||
if dim % 2 != 0:
|
||||
raise ValueError(f"dim must be even, got {dim}.")
|
||||
freqs = torch.outer(
|
||||
torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim))
|
||||
)
|
||||
freqs = torch.polar(torch.ones_like(freqs), freqs)
|
||||
return freqs
|
||||
|
||||
|
||||
@torch.amp.autocast("cuda", enabled=False)
|
||||
def rope_apply(x, grid_sizes, freqs):
|
||||
n, c = x.size(2), x.size(3) // 2
|
||||
|
||||
# split freqs
|
||||
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
|
||||
|
||||
# loop over samples
|
||||
output = []
|
||||
for i, (f, h, w) in enumerate(grid_sizes.tolist()):
|
||||
seq_len = f * h * w
|
||||
|
||||
# precompute multipliers
|
||||
x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2))
|
||||
freqs_i = torch.cat(
|
||||
[
|
||||
freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
|
||||
freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
|
||||
freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
|
||||
],
|
||||
dim=-1,
|
||||
).reshape(seq_len, 1, -1)
|
||||
|
||||
# apply rotary embedding
|
||||
x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
|
||||
x_i = torch.cat([x_i, x[i, seq_len:]])
|
||||
|
||||
# append to collection
|
||||
output.append(x_i)
|
||||
return torch.stack(output).float()
|
||||
|
||||
|
||||
class WanRMSNorm(nn.Module):
|
||||
def __init__(self, dim, eps=1e-5):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x):
|
||||
r"""
|
||||
Args:
|
||||
x(Tensor): Shape [B, L, C]
|
||||
"""
|
||||
return self._norm(x.float()).type_as(x) * self.weight
|
||||
|
||||
def _norm(self, x):
|
||||
return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
|
||||
|
||||
|
||||
class WanLayerNorm(nn.LayerNorm):
|
||||
def __init__(self, dim, eps=1e-6, elementwise_affine=False):
|
||||
super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
|
||||
|
||||
def forward(self, x):
|
||||
r"""
|
||||
Args:
|
||||
x(Tensor): Shape [B, L, C]
|
||||
"""
|
||||
return super().forward(x.float()).type_as(x)
|
||||
|
||||
|
||||
class WanSelfAttention(nn.Module):
|
||||
def __init__(self, dim, num_heads, qk_norm=True, eps=1e-6):
|
||||
if dim % num_heads != 0:
|
||||
raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads}).")
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
|
||||
# layers
|
||||
self.q = nn.Linear(dim, dim)
|
||||
self.k = nn.Linear(dim, dim)
|
||||
self.v = nn.Linear(dim, dim)
|
||||
self.o = nn.Linear(dim, dim)
|
||||
self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
|
||||
self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
|
||||
|
||||
# NOTE: FastWAM never runs the upstream Wan attention forward. FastWAMAttentionBlock
|
||||
# reuses only the q/k/v/o/norm submodules defined above and computes attention via
|
||||
# `fastwam_masked_attention` (SDPA). The original flash-attention forward was removed,
|
||||
# which also collapsed the former WanCrossAttention subclass into this class (it only
|
||||
# differed by its forward): self- and cross-attention now share the same projection module.
|
||||
|
||||
|
||||
class WanAttentionBlock(nn.Module):
|
||||
def __init__(self, dim, ffn_dim, num_heads, qk_norm=True, cross_attn_norm=False, eps=1e-6):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.ffn_dim = ffn_dim
|
||||
self.num_heads = num_heads
|
||||
self.qk_norm = qk_norm
|
||||
self.cross_attn_norm = cross_attn_norm
|
||||
self.eps = eps
|
||||
|
||||
# layers
|
||||
self.norm1 = WanLayerNorm(dim, eps)
|
||||
self.self_attn = WanSelfAttention(dim, num_heads, qk_norm, eps)
|
||||
self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
|
||||
self.cross_attn = WanSelfAttention(dim, num_heads, qk_norm, eps)
|
||||
self.norm2 = WanLayerNorm(dim, eps)
|
||||
self.ffn = nn.Sequential(
|
||||
nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim)
|
||||
)
|
||||
|
||||
# modulation
|
||||
self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
|
||||
|
||||
# NOTE: The upstream Wan block forward (self-attention + cross-attention + FFN via
|
||||
# flash-attention) was removed. FastWAM subclasses this block as FastWAMAttentionBlock
|
||||
# and overrides forward to use SDPA with explicit boolean masks; only __init__ (the
|
||||
# norm/attention/ffn submodules) is reused here.
|
||||
|
||||
|
||||
class Head(nn.Module):
|
||||
def __init__(self, dim, out_dim, patch_size, eps=1e-6):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.out_dim = out_dim
|
||||
self.patch_size = patch_size
|
||||
self.eps = eps
|
||||
|
||||
# layers
|
||||
out_dim = math.prod(patch_size) * out_dim
|
||||
self.norm = WanLayerNorm(dim, eps)
|
||||
self.head = nn.Linear(dim, out_dim)
|
||||
|
||||
# modulation
|
||||
self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
|
||||
|
||||
def forward(self, x, e):
|
||||
r"""
|
||||
Args:
|
||||
x(Tensor): Shape [B, L1, C]
|
||||
e(Tensor): Shape [B, L1, C]
|
||||
"""
|
||||
with torch.amp.autocast("cuda", dtype=torch.float32):
|
||||
e = (self.modulation.unsqueeze(0) + e.unsqueeze(2)).chunk(2, dim=2)
|
||||
x = self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2))
|
||||
return x
|
||||
|
||||
|
||||
class WanModel(nn.Module):
|
||||
r"""
|
||||
Wan diffusion backbone supporting both text-to-video and image-to-video.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_type="t2v",
|
||||
patch_size=(1, 2, 2),
|
||||
text_len=512,
|
||||
in_dim=16,
|
||||
dim=2048,
|
||||
ffn_dim=8192,
|
||||
freq_dim=256,
|
||||
text_dim=4096,
|
||||
out_dim=16,
|
||||
num_heads=16,
|
||||
num_layers=32,
|
||||
qk_norm=True,
|
||||
cross_attn_norm=True,
|
||||
eps=1e-6,
|
||||
):
|
||||
r"""
|
||||
Initialize the diffusion model backbone.
|
||||
|
||||
Args:
|
||||
model_type (`str`, *optional*, defaults to 't2v'):
|
||||
Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
|
||||
patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
|
||||
3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
|
||||
text_len (`int`, *optional*, defaults to 512):
|
||||
Fixed length for text embeddings
|
||||
in_dim (`int`, *optional*, defaults to 16):
|
||||
Input video channels (C_in)
|
||||
dim (`int`, *optional*, defaults to 2048):
|
||||
Hidden dimension of the transformer
|
||||
ffn_dim (`int`, *optional*, defaults to 8192):
|
||||
Intermediate dimension in feed-forward network
|
||||
freq_dim (`int`, *optional*, defaults to 256):
|
||||
Dimension for sinusoidal time embeddings
|
||||
text_dim (`int`, *optional*, defaults to 4096):
|
||||
Input dimension for text embeddings
|
||||
out_dim (`int`, *optional*, defaults to 16):
|
||||
Output video channels (C_out)
|
||||
num_heads (`int`, *optional*, defaults to 16):
|
||||
Number of attention heads
|
||||
num_layers (`int`, *optional*, defaults to 32):
|
||||
Number of transformer blocks
|
||||
qk_norm (`bool`, *optional*, defaults to True):
|
||||
Enable query/key normalization
|
||||
cross_attn_norm (`bool`, *optional*, defaults to False):
|
||||
Enable cross-attention normalization
|
||||
eps (`float`, *optional*, defaults to 1e-6):
|
||||
Epsilon value for normalization layers
|
||||
"""
|
||||
|
||||
super().__init__()
|
||||
|
||||
if model_type not in ["t2v", "i2v", "ti2v", "s2v"]:
|
||||
raise ValueError(f"model_type must be one of ['t2v', 'i2v', 'ti2v', 's2v'], got {model_type!r}.")
|
||||
self.model_type = model_type
|
||||
|
||||
self.patch_size = patch_size
|
||||
self.text_len = text_len
|
||||
self.in_dim = in_dim
|
||||
self.dim = dim
|
||||
self.ffn_dim = ffn_dim
|
||||
self.freq_dim = freq_dim
|
||||
self.text_dim = text_dim
|
||||
self.out_dim = out_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_layers = num_layers
|
||||
self.qk_norm = qk_norm
|
||||
self.cross_attn_norm = cross_attn_norm
|
||||
self.eps = eps
|
||||
|
||||
# embeddings
|
||||
self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size)
|
||||
self.text_embedding = nn.Sequential(
|
||||
nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim)
|
||||
)
|
||||
|
||||
self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
|
||||
self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
|
||||
|
||||
# blocks
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
WanAttentionBlock(dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
# head
|
||||
self.head = Head(dim, out_dim, patch_size, eps)
|
||||
|
||||
# buffers (don't use register_buffer otherwise dtype will be changed in to())
|
||||
if (dim % num_heads) != 0 or (dim // num_heads) % 2 != 0:
|
||||
raise ValueError(
|
||||
f"dim ({dim}) must be divisible by num_heads ({num_heads}) with an even head dim."
|
||||
)
|
||||
d = dim // num_heads
|
||||
self.freqs = torch.cat(
|
||||
[
|
||||
rope_params(1024, d - 4 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6)),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
# initialize weights
|
||||
self.init_weights()
|
||||
|
||||
# NOTE: The upstream Wan diffusion forward (flash-attention based) was removed.
|
||||
# FastWAM's WanVideoDiT subclasses this model, rebuilds `self.blocks` with
|
||||
# FastWAMAttentionBlock, and provides its own SDPA-based forward. Only the
|
||||
# constructor (embeddings, blocks, head, rope buffers) and the helpers below
|
||||
# (unpatchify / init_weights) are reused. WanModel is never run directly.
|
||||
|
||||
def unpatchify(self, x, grid_sizes):
|
||||
r"""
|
||||
Reconstruct video tensors from patch embeddings.
|
||||
|
||||
Args:
|
||||
x (List[Tensor]):
|
||||
List of patchified features, each with shape [L, C_out * prod(patch_size)]
|
||||
grid_sizes (Tensor):
|
||||
Original spatial-temporal grid dimensions before patching,
|
||||
shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
|
||||
|
||||
Returns:
|
||||
List[Tensor]:
|
||||
Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
|
||||
"""
|
||||
|
||||
c = self.out_dim
|
||||
out = []
|
||||
for u, v in zip(x, grid_sizes.tolist(), strict=False):
|
||||
u = u[: math.prod(v)].view(*v, *self.patch_size, c)
|
||||
u = torch.einsum("fhwpqrc->cfphqwr", u)
|
||||
u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=False)])
|
||||
out.append(u)
|
||||
return out
|
||||
|
||||
def init_weights(self):
|
||||
r"""
|
||||
Initialize model parameters using Xavier initialization.
|
||||
"""
|
||||
|
||||
# basic init
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
# init embeddings
|
||||
nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
|
||||
for m in self.text_embedding.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, std=0.02)
|
||||
for m in self.time_embedding.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, std=0.02)
|
||||
|
||||
# init output layer
|
||||
nn.init.zeros_(self.head.head.weight)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,800 +0,0 @@
|
||||
# Copyright 2024 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 logging
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as functional
|
||||
from einops import rearrange
|
||||
|
||||
from .model import (
|
||||
WanAttentionBlock,
|
||||
WanLayerNorm,
|
||||
WanModel,
|
||||
WanRMSNorm,
|
||||
rope_apply,
|
||||
rope_params,
|
||||
sinusoidal_embedding_1d,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_sampling_sigmas(sampling_steps, shift):
|
||||
# Vendored from Wan2.2 (formerly wan/utils/fm_solvers.py); computes the
|
||||
# noise-level (sigma) schedule for Wan-compatible flow-matching inference.
|
||||
sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps]
|
||||
sigma = shift * sigma / (1 + (shift - 1) * sigma)
|
||||
return sigma
|
||||
|
||||
|
||||
def create_custom_forward(module):
|
||||
def custom_forward(*inputs, **kwargs):
|
||||
return module(*inputs, **kwargs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
|
||||
def gradient_checkpoint_forward(
|
||||
model,
|
||||
use_gradient_checkpointing,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if use_gradient_checkpointing:
|
||||
model_output = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(model),
|
||||
*args,
|
||||
**kwargs,
|
||||
use_reentrant=False,
|
||||
)
|
||||
else:
|
||||
model_output = model(*args, **kwargs)
|
||||
return model_output
|
||||
|
||||
|
||||
def fastwam_masked_attention(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
num_heads: int,
|
||||
ctx_mask: torch.Tensor | None = None,
|
||||
fp32_attention: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""FastWAM masked attention wrapper for MoT masks and CPU test coverage.
|
||||
|
||||
The official Wan attention implementation is still used as the source of
|
||||
the projection/norm modules. This wrapper only replaces the final attention
|
||||
kernel because FastWAM needs explicit boolean masks for video/action MoT
|
||||
routing, while the upstream FlashAttention path accepts sequence lengths
|
||||
but not arbitrary [query, key] masks.
|
||||
"""
|
||||
|
||||
q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
|
||||
k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
|
||||
v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
|
||||
if fp32_attention:
|
||||
q = q.float()
|
||||
k = k.float()
|
||||
v = v.float()
|
||||
else:
|
||||
q = q.to(dtype=v.dtype)
|
||||
k = k.to(dtype=v.dtype)
|
||||
x = functional.scaled_dot_product_attention(q, k, v, attn_mask=ctx_mask)
|
||||
return rearrange(x, "b n s d -> b s (n d)", n=num_heads)
|
||||
|
||||
|
||||
def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
|
||||
return x * (1 + scale) + shift
|
||||
|
||||
|
||||
class WanContinuousFlowMatchScheduler:
|
||||
"""Continuous-time Flow-Matching scheduler with shift-based Wan sampling."""
|
||||
|
||||
def __init__(self, num_train_timesteps: int = 1000, shift: float = 5.0, eps: float = 1e-10):
|
||||
if num_train_timesteps <= 0:
|
||||
raise ValueError(f"`num_train_timesteps` must be positive, got {num_train_timesteps}")
|
||||
if shift <= 0:
|
||||
raise ValueError(f"`shift` must be positive, got {shift}")
|
||||
self.num_train_timesteps = int(num_train_timesteps)
|
||||
self.shift = float(shift)
|
||||
self.eps = float(eps)
|
||||
self._y_min, self._weight_norm_const = self._precompute_training_weight_stats()
|
||||
|
||||
@staticmethod
|
||||
def _phi(u: torch.Tensor, shift: float) -> torch.Tensor:
|
||||
return shift * u / (1.0 + (shift - 1.0) * u)
|
||||
|
||||
def _precompute_training_weight_stats(self) -> tuple[float, float]:
|
||||
steps = self.num_train_timesteps
|
||||
u_grid = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64)[:-1]
|
||||
t_grid = self._phi(u_grid, self.shift) * float(steps)
|
||||
y_grid = torch.exp(-2.0 * ((t_grid - (steps / 2.0)) / steps) ** 2)
|
||||
y_min = float(y_grid.min().item())
|
||||
y_shifted_grid = y_grid - y_min
|
||||
norm_const = float(y_shifted_grid.mean().item())
|
||||
return y_min, norm_const
|
||||
|
||||
def sample_training_t(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
||||
if batch_size <= 0:
|
||||
raise ValueError(f"`batch_size` must be positive, got {batch_size}")
|
||||
u = torch.rand((batch_size,), device=device, dtype=torch.float32)
|
||||
sigma = self._phi(u, self.shift)
|
||||
timestep = sigma * float(self.num_train_timesteps)
|
||||
return timestep.to(dtype=dtype)
|
||||
|
||||
def training_weight(self, timestep: torch.Tensor) -> torch.Tensor:
|
||||
t = timestep.to(dtype=torch.float32)
|
||||
steps = float(self.num_train_timesteps)
|
||||
y = torch.exp(-2.0 * ((t - (steps / 2.0)) / steps) ** 2)
|
||||
y_shifted = y - self._y_min
|
||||
weight = y_shifted / (self._weight_norm_const + self.eps)
|
||||
if weight.numel() == 1:
|
||||
return weight.reshape(())
|
||||
return weight
|
||||
|
||||
def add_noise(
|
||||
self, original_samples: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
sigma = (timestep / float(self.num_train_timesteps)).to(
|
||||
original_samples.device, dtype=original_samples.dtype
|
||||
)
|
||||
if sigma.ndim == 0:
|
||||
return (1 - sigma) * original_samples + sigma * noise
|
||||
sigma = sigma.view(-1, *([1] * (original_samples.ndim - 1)))
|
||||
return (1 - sigma) * original_samples + sigma * noise
|
||||
|
||||
@staticmethod
|
||||
def training_target(sample: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
|
||||
del timestep
|
||||
return noise - sample
|
||||
|
||||
def build_inference_schedule(
|
||||
self,
|
||||
num_inference_steps: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
shift_override: float | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if num_inference_steps <= 0:
|
||||
raise ValueError(f"`num_inference_steps` must be positive, got {num_inference_steps}")
|
||||
shift = self.shift if shift_override is None else float(shift_override)
|
||||
if shift <= 0:
|
||||
raise ValueError(f"`shift` must be positive, got {shift}")
|
||||
|
||||
sigma_steps = torch.as_tensor(
|
||||
get_sampling_sigmas(num_inference_steps, shift),
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
timesteps = sigma_steps * float(self.num_train_timesteps)
|
||||
sigma_next = torch.cat([sigma_steps[1:], sigma_steps.new_zeros(1)])
|
||||
deltas = sigma_next - sigma_steps
|
||||
return timesteps.to(dtype=dtype), deltas.to(dtype=dtype)
|
||||
|
||||
@staticmethod
|
||||
def step(model_output: torch.Tensor, delta: torch.Tensor, sample: torch.Tensor) -> torch.Tensor:
|
||||
delta = delta.to(sample.device, dtype=sample.dtype)
|
||||
if delta.ndim == 0:
|
||||
return sample + model_output * delta
|
||||
delta = delta.view(-1, *([1] * (sample.ndim - 1)))
|
||||
return sample + model_output * delta
|
||||
|
||||
|
||||
def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0):
|
||||
return rope_params(end, dim, theta)
|
||||
|
||||
|
||||
def apply_dense_rope(x: torch.Tensor, freqs: torch.Tensor, num_heads: int) -> torch.Tensor:
|
||||
x = rearrange(x, "b s (n d) -> b s n d", n=num_heads)
|
||||
x_out = torch.view_as_complex(x.to(torch.float32).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2))
|
||||
freqs = freqs.to(torch.complex64) if freqs.device.type == "npu" else freqs
|
||||
x_out = torch.view_as_real(x_out * freqs).flatten(2)
|
||||
return x_out.to(x.dtype)
|
||||
|
||||
|
||||
def _linear_input(linear: nn.Linear, x: torch.Tensor) -> torch.Tensor:
|
||||
return x.to(dtype=linear.weight.dtype)
|
||||
|
||||
|
||||
def _wan_layer_norm(norm: nn.Module, x: torch.Tensor) -> torch.Tensor:
|
||||
if isinstance(norm, WanLayerNorm) and norm.weight is not None:
|
||||
weight = norm.weight.float()
|
||||
bias = norm.bias.float() if norm.bias is not None else None
|
||||
return functional.layer_norm(x.float(), norm.normalized_shape, weight, bias, norm.eps).to(
|
||||
dtype=x.dtype
|
||||
)
|
||||
return norm(x)
|
||||
|
||||
|
||||
def create_group_causal_attn_mask(
|
||||
num_temporal_groups: int, num_query_per_group: int, num_key_per_group: int, mode: str = "causal"
|
||||
) -> torch.Tensor:
|
||||
if mode not in ["causal", "group_diagonal"]:
|
||||
raise ValueError(f"`mode` must be 'causal' or 'group_diagonal', got {mode}.")
|
||||
if num_temporal_groups <= 0:
|
||||
raise ValueError(f"`num_temporal_groups` must be positive, got {num_temporal_groups}.")
|
||||
if num_query_per_group <= 0:
|
||||
raise ValueError(f"`num_query_per_group` must be positive, got {num_query_per_group}.")
|
||||
if num_key_per_group <= 0:
|
||||
raise ValueError(f"`num_key_per_group` must be positive, got {num_key_per_group}.")
|
||||
|
||||
total_num_query_tokens = num_temporal_groups * num_query_per_group
|
||||
total_num_key_tokens = num_temporal_groups * num_key_per_group
|
||||
query_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_query_per_group).unsqueeze(1)
|
||||
key_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_key_per_group).unsqueeze(0)
|
||||
|
||||
if mode == "causal":
|
||||
attn_mask = query_time_indices >= key_time_indices
|
||||
else:
|
||||
attn_mask = query_time_indices == key_time_indices
|
||||
|
||||
if attn_mask.shape != (total_num_query_tokens, total_num_key_tokens):
|
||||
raise RuntimeError("Attention mask shape mismatch.")
|
||||
return attn_mask
|
||||
|
||||
|
||||
class FastWAMAttentionBlock(WanAttentionBlock):
|
||||
"""Wan attention block with FastWAM's arbitrary boolean mask support."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_dim: int,
|
||||
attn_head_dim: int,
|
||||
num_heads: int,
|
||||
ffn_dim: int,
|
||||
eps: float = 1e-6,
|
||||
fp32_attention: bool = True,
|
||||
):
|
||||
attention_dim = attn_head_dim * num_heads
|
||||
if hidden_dim == attention_dim:
|
||||
super().__init__(
|
||||
dim=hidden_dim,
|
||||
ffn_dim=ffn_dim,
|
||||
num_heads=num_heads,
|
||||
qk_norm=True,
|
||||
cross_attn_norm=True,
|
||||
eps=eps,
|
||||
)
|
||||
else:
|
||||
nn.Module.__init__(self)
|
||||
self.dim = hidden_dim
|
||||
self.ffn_dim = ffn_dim
|
||||
self.num_heads = num_heads
|
||||
self.qk_norm = True
|
||||
self.cross_attn_norm = True
|
||||
self.eps = eps
|
||||
self.norm1 = WanLayerNorm(hidden_dim, eps)
|
||||
self.self_attn = _FastWAMProjectedAttention(hidden_dim, attention_dim, num_heads, eps)
|
||||
self.norm3 = WanLayerNorm(hidden_dim, eps, elementwise_affine=True)
|
||||
self.cross_attn = _FastWAMProjectedAttention(hidden_dim, attention_dim, num_heads, eps)
|
||||
self.norm2 = WanLayerNorm(hidden_dim, eps)
|
||||
self.ffn = nn.Sequential(
|
||||
nn.Linear(hidden_dim, ffn_dim),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Linear(ffn_dim, hidden_dim),
|
||||
)
|
||||
self.modulation = nn.Parameter(torch.randn(1, 6, hidden_dim) / hidden_dim**0.5)
|
||||
self.attn_head_dim = attn_head_dim
|
||||
self.fp32_attention = bool(fp32_attention)
|
||||
|
||||
@staticmethod
|
||||
def split_modulation(block, t_mod: torch.Tensor):
|
||||
has_seq = len(t_mod.shape) == 4
|
||||
chunk_dim = 2 if has_seq else 1
|
||||
|
||||
base_mod = block.modulation.to(dtype=t_mod.dtype, device=t_mod.device)
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (base_mod + t_mod).chunk(
|
||||
6, dim=chunk_dim
|
||||
)
|
||||
if has_seq:
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
shift_msa.squeeze(2),
|
||||
scale_msa.squeeze(2),
|
||||
gate_msa.squeeze(2),
|
||||
shift_mlp.squeeze(2),
|
||||
scale_mlp.squeeze(2),
|
||||
gate_mlp.squeeze(2),
|
||||
)
|
||||
return shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp
|
||||
|
||||
def project_self_attention(
|
||||
self, x: torch.Tensor, freqs: torch.Tensor | dict[str, torch.Tensor]
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
q = self.self_attn.norm_q(self.self_attn.q(x))
|
||||
k = self.self_attn.norm_k(self.self_attn.k(x))
|
||||
v = self.self_attn.v(x)
|
||||
if isinstance(freqs, dict):
|
||||
b, s = x.shape[:2]
|
||||
q = rope_apply(
|
||||
q.view(b, s, self.num_heads, self.attn_head_dim),
|
||||
freqs["grid_sizes"],
|
||||
freqs["freqs"],
|
||||
).flatten(2)
|
||||
k = rope_apply(
|
||||
k.view(b, s, self.num_heads, self.attn_head_dim),
|
||||
freqs["grid_sizes"],
|
||||
freqs["freqs"],
|
||||
).flatten(2)
|
||||
else:
|
||||
q = apply_dense_rope(q, freqs, self.num_heads)
|
||||
k = apply_dense_rope(k, freqs, self.num_heads)
|
||||
return q, k, v
|
||||
|
||||
def apply_cross_attention(
|
||||
self, x: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor | None = None
|
||||
) -> torch.Tensor:
|
||||
if context_mask is not None and context_mask.dim() == 3:
|
||||
context_mask = context_mask.unsqueeze(1)
|
||||
attn = self.cross_attn
|
||||
b, n, d = x.size(0), attn.num_heads, attn.head_dim
|
||||
q = attn.norm_q(attn.q(x)).view(b, -1, n * d)
|
||||
k = attn.norm_k(attn.k(context)).view(b, -1, n * d)
|
||||
v = attn.v(context).view(b, -1, n * d)
|
||||
x = fastwam_masked_attention(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
num_heads=n,
|
||||
ctx_mask=context_mask,
|
||||
fp32_attention=self.fp32_attention,
|
||||
)
|
||||
return attn.o(_linear_input(attn.o, x))
|
||||
|
||||
def project_self_attention_output(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.self_attn.o(_linear_input(self.self_attn.o, x))
|
||||
|
||||
def apply_norm1(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return _wan_layer_norm(self.norm1, x)
|
||||
|
||||
def apply_norm2(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return _wan_layer_norm(self.norm2, x)
|
||||
|
||||
def apply_norm3(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return _wan_layer_norm(self.norm3, x)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
context: torch.Tensor,
|
||||
t_mod: torch.Tensor,
|
||||
freqs: torch.Tensor,
|
||||
context_mask: torch.Tensor | None = None,
|
||||
self_attn_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.split_modulation(self, t_mod)
|
||||
residual_x = x
|
||||
attn_input = modulate(self.apply_norm1(x), shift_msa, scale_msa)
|
||||
q, k, v = self.project_self_attention(attn_input, freqs)
|
||||
y = fastwam_masked_attention(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
num_heads=self.num_heads,
|
||||
ctx_mask=self_attn_mask,
|
||||
fp32_attention=self.fp32_attention,
|
||||
)
|
||||
x = residual_x + gate_msa * self.project_self_attention_output(y)
|
||||
x = x + self.apply_cross_attention(self.apply_norm3(x), context, context_mask=context_mask)
|
||||
mlp_input = modulate(self.apply_norm2(x), shift_mlp, scale_mlp)
|
||||
return x + gate_mlp * self.ffn(mlp_input)
|
||||
|
||||
|
||||
class _FastWAMProjectedAttention(nn.Module):
|
||||
def __init__(self, hidden_dim: int, attention_dim: int, num_heads: int, eps: float):
|
||||
super().__init__()
|
||||
self.dim = hidden_dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = attention_dim // num_heads
|
||||
self.q = nn.Linear(hidden_dim, attention_dim)
|
||||
self.k = nn.Linear(hidden_dim, attention_dim)
|
||||
self.v = nn.Linear(hidden_dim, attention_dim)
|
||||
self.o = nn.Linear(attention_dim, hidden_dim)
|
||||
self.norm_q = WanRMSNorm(attention_dim, eps=eps)
|
||||
self.norm_k = WanRMSNorm(attention_dim, eps=eps)
|
||||
|
||||
|
||||
class WanVideoDiT(WanModel):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_dim: int,
|
||||
in_dim: int,
|
||||
ffn_dim: int,
|
||||
out_dim: int,
|
||||
text_dim: int,
|
||||
freq_dim: int,
|
||||
eps: float,
|
||||
patch_size: tuple[int, int, int],
|
||||
num_heads: int,
|
||||
attn_head_dim: int,
|
||||
num_layers: int,
|
||||
has_image_input: bool = False,
|
||||
has_image_pos_emb: bool = False,
|
||||
has_ref_conv: bool = False,
|
||||
add_control_adapter: bool = False,
|
||||
in_dim_control_adapter: int = 24,
|
||||
seperated_timestep: bool = False,
|
||||
require_vae_embedding: bool = False,
|
||||
require_clip_embedding: bool = False,
|
||||
fuse_vae_embedding_in_latents: bool = True,
|
||||
action_conditioned: bool = False,
|
||||
action_dim: int = 7,
|
||||
action_group_causal_mask_mode="causal",
|
||||
video_attention_mask_mode: str = "bidirectional",
|
||||
use_gradient_checkpointing: bool = False,
|
||||
fp32_attention: bool = True,
|
||||
):
|
||||
del in_dim_control_adapter
|
||||
if has_image_input:
|
||||
raise ValueError("FastWAM currently expects Wan2.2 TI2V latents with fused image conditioning.")
|
||||
if has_image_pos_emb:
|
||||
raise ValueError("FastWAM does not support extra image positional embeddings in WanVideoDiT.")
|
||||
if has_ref_conv:
|
||||
raise ValueError("FastWAM does not support reference convolutions in WanVideoDiT.")
|
||||
if add_control_adapter:
|
||||
raise ValueError("FastWAM does not support control adapters in WanVideoDiT.")
|
||||
if require_clip_embedding:
|
||||
raise ValueError("FastWAM does not support CLIP embedding conditioning in WanVideoDiT.")
|
||||
if require_vae_embedding or not fuse_vae_embedding_in_latents:
|
||||
raise ValueError("FastWAM expects VAE conditioning to be fused in latents.")
|
||||
if attn_head_dim != hidden_dim // num_heads:
|
||||
raise ValueError(
|
||||
"`attn_head_dim` must match the upstream Wan head dimension `hidden_dim // num_heads`; "
|
||||
f"got {attn_head_dim} vs {hidden_dim // num_heads}."
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
model_type="ti2v",
|
||||
patch_size=patch_size,
|
||||
text_len=512,
|
||||
in_dim=in_dim,
|
||||
dim=hidden_dim,
|
||||
ffn_dim=ffn_dim,
|
||||
freq_dim=freq_dim,
|
||||
text_dim=text_dim,
|
||||
out_dim=out_dim,
|
||||
num_heads=num_heads,
|
||||
num_layers=num_layers,
|
||||
qk_norm=True,
|
||||
cross_attn_norm=True,
|
||||
eps=eps,
|
||||
)
|
||||
self.blocks = torch.nn.ModuleList(
|
||||
[
|
||||
FastWAMAttentionBlock(
|
||||
hidden_dim=hidden_dim,
|
||||
attn_head_dim=attn_head_dim,
|
||||
num_heads=num_heads,
|
||||
ffn_dim=ffn_dim,
|
||||
eps=eps,
|
||||
fp32_attention=fp32_attention,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
self.init_weights()
|
||||
|
||||
self.hidden_dim = hidden_dim
|
||||
self.attn_head_dim = attn_head_dim
|
||||
self.seperated_timestep = seperated_timestep
|
||||
self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents
|
||||
self.video_attention_mask_mode = str(video_attention_mask_mode)
|
||||
self.action_conditioned = action_conditioned
|
||||
self.action_dim = action_dim
|
||||
self.fp32_attention = bool(fp32_attention)
|
||||
|
||||
if self.action_conditioned:
|
||||
self.action_embedding = torch.nn.Linear(action_dim, hidden_dim)
|
||||
self.action_group_causal_mask_mode = action_group_causal_mask_mode
|
||||
|
||||
self.use_gradient_checkpointing = use_gradient_checkpointing
|
||||
if self.use_gradient_checkpointing:
|
||||
logger.info(
|
||||
"Using gradient checkpointing for DiT blocks. This will save memory but use more computation."
|
||||
)
|
||||
|
||||
def patchify(self, x: torch.Tensor):
|
||||
return self.patch_embedding(x)
|
||||
|
||||
def _validate_forward_inputs(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
context: torch.Tensor,
|
||||
context_mask: torch.Tensor | None,
|
||||
action: torch.Tensor | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
if x.ndim != 5:
|
||||
raise ValueError(f"`latents` must be 5D [B, C, T, H, W], got shape {tuple(x.shape)}")
|
||||
num_latent_frames = x.shape[2]
|
||||
if context.ndim != 3:
|
||||
raise ValueError(f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}")
|
||||
if timestep.ndim != 1:
|
||||
raise ValueError(f"`timestep` must be 1D [B] or [1], got shape {tuple(timestep.shape)}")
|
||||
if self.action_conditioned:
|
||||
allow_text_only_single_frame = num_latent_frames == 1 and action is None
|
||||
if not allow_text_only_single_frame:
|
||||
if action is None:
|
||||
raise ValueError("Action input is required for action-conditioned model.")
|
||||
if action.ndim != 3:
|
||||
raise ValueError(
|
||||
f"`action` must be 3D [B, action_horizon, action_dim], got shape {tuple(action.shape)}"
|
||||
)
|
||||
if action.shape[2] != self.action_dim:
|
||||
raise ValueError(
|
||||
f"`action` last dimension must be {self.action_dim}, got {action.shape[2]}"
|
||||
)
|
||||
if num_latent_frames <= 1:
|
||||
raise ValueError(
|
||||
f"video length must be > 1 for action-conditioned model, got {num_latent_frames}"
|
||||
)
|
||||
if action.shape[1] % (num_latent_frames - 1) != 0:
|
||||
raise ValueError(
|
||||
"action horizon must be divisible by (num_latent_frames - 1), "
|
||||
f"got action_horizon={action.shape[1]}"
|
||||
)
|
||||
if context_mask is None:
|
||||
context_mask = torch.ones(
|
||||
(context.shape[0], context.shape[1]), dtype=torch.bool, device=context.device
|
||||
)
|
||||
else:
|
||||
if context_mask.ndim != 2:
|
||||
raise ValueError(f"`context_mask` must be 2D [B, L], got shape {tuple(context_mask.shape)}")
|
||||
if context_mask.shape[0] != context.shape[0] or context_mask.shape[1] != context.shape[1]:
|
||||
raise ValueError(
|
||||
"`context_mask` shape must match `context` shape [B, L], "
|
||||
f"got {tuple(context_mask.shape)} vs {tuple(context.shape)}"
|
||||
)
|
||||
|
||||
batch_size = x.shape[0]
|
||||
if batch_size != context.shape[0]:
|
||||
if not self.training and batch_size == 1:
|
||||
x = x.expand(context.shape[0], -1, -1, -1, -1)
|
||||
batch_size = context.shape[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Batch mismatch between latents and context: {batch_size} vs {context.shape[0]}."
|
||||
)
|
||||
|
||||
if timestep.shape[0] not in (1, batch_size):
|
||||
raise ValueError(
|
||||
f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}"
|
||||
)
|
||||
if timestep.shape[0] == 1 and batch_size > 1:
|
||||
if self.training:
|
||||
raise ValueError("During training, timestep length must match batch_size.")
|
||||
timestep = timestep.expand(batch_size)
|
||||
return x, timestep, context_mask
|
||||
|
||||
def build_video_to_video_mask(
|
||||
self,
|
||||
video_seq_len: int,
|
||||
video_tokens_per_frame: int,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
if video_seq_len <= 0:
|
||||
raise ValueError(f"`video_seq_len` must be positive, got {video_seq_len}")
|
||||
if video_tokens_per_frame <= 0:
|
||||
raise ValueError(f"`video_tokens_per_frame` must be positive, got {video_tokens_per_frame}")
|
||||
|
||||
if self.video_attention_mask_mode == "bidirectional":
|
||||
return torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device)
|
||||
|
||||
if self.video_attention_mask_mode == "per_frame_causal":
|
||||
if video_seq_len % video_tokens_per_frame != 0:
|
||||
raise ValueError(
|
||||
"`video_seq_len` must be divisible by `video_tokens_per_frame` in `per_frame_causal` mode, "
|
||||
f"got {video_seq_len} and {video_tokens_per_frame}"
|
||||
)
|
||||
num_video_frames = video_seq_len // video_tokens_per_frame
|
||||
frame_causal = torch.tril(
|
||||
torch.ones((num_video_frames, num_video_frames), dtype=torch.bool, device=device)
|
||||
)
|
||||
return frame_causal.repeat_interleave(video_tokens_per_frame, dim=0).repeat_interleave(
|
||||
video_tokens_per_frame, dim=1
|
||||
)
|
||||
|
||||
if self.video_attention_mask_mode == "first_frame_causal":
|
||||
video_mask = torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device)
|
||||
first_frame_tokens = min(video_tokens_per_frame, video_seq_len)
|
||||
video_mask[:first_frame_tokens, first_frame_tokens:] = False
|
||||
return video_mask
|
||||
|
||||
raise ValueError(f"Unsupported video attention mask mode: {self.video_attention_mask_mode}")
|
||||
|
||||
def pre_dit(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
context: torch.Tensor,
|
||||
context_mask: torch.Tensor | None = None,
|
||||
action: torch.Tensor | None = None,
|
||||
fuse_vae_embedding_in_latents: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
x, timestep, context_mask = self._validate_forward_inputs(
|
||||
x=x,
|
||||
timestep=timestep,
|
||||
context=context,
|
||||
context_mask=context_mask,
|
||||
action=action,
|
||||
)
|
||||
model_dtype = self.patch_embedding.weight.dtype
|
||||
x = x.to(dtype=model_dtype)
|
||||
context = context.to(dtype=model_dtype)
|
||||
if action is not None:
|
||||
action = action.to(dtype=model_dtype)
|
||||
|
||||
batch_size = x.shape[0]
|
||||
patch_h = int(self.patch_size[1])
|
||||
patch_w = int(self.patch_size[2])
|
||||
if x.shape[3] % patch_h != 0 or x.shape[4] % patch_w != 0:
|
||||
raise ValueError(
|
||||
"Latent spatial shape must be divisible by DiT patch size, "
|
||||
f"got HxW=({x.shape[3]}, {x.shape[4]}), patch=({patch_h}, {patch_w})"
|
||||
)
|
||||
tokens_per_frame = (x.shape[3] // patch_h) * (x.shape[4] // patch_w)
|
||||
|
||||
if not (self.seperated_timestep and fuse_vae_embedding_in_latents):
|
||||
raise NotImplementedError(
|
||||
"FastWAM currently requires separated timesteps with fused VAE latents."
|
||||
)
|
||||
|
||||
token_timesteps = torch.ones(
|
||||
(batch_size, x.shape[2], tokens_per_frame),
|
||||
dtype=model_dtype,
|
||||
device=timestep.device,
|
||||
) * timestep.to(dtype=model_dtype).view(batch_size, 1, 1)
|
||||
token_timesteps[:, 0, :] = 0
|
||||
token_timesteps = token_timesteps.reshape(batch_size, -1)
|
||||
# Wan keeps the time embedding in fp32: the AdaLN modulation in the vendored
|
||||
# Head/Block asserts e.dtype == float32 (numerical stability of the scale/shift).
|
||||
# Upstream guarantees this via an fp32 autocast region, so it holds even when the
|
||||
# model runs in bf16. Mirror that here, then cast the per-block modulation back to
|
||||
# model_dtype so the bf16 attention blocks are not upcast to fp32.
|
||||
with torch.amp.autocast("cuda", dtype=torch.float32):
|
||||
token_t_emb = sinusoidal_embedding_1d(self.freq_dim, token_timesteps.reshape(-1)).float()
|
||||
t = self.time_embedding(token_t_emb).reshape(batch_size, -1, self.hidden_dim)
|
||||
t_mod = self.time_projection(t).unflatten(2, (6, self.hidden_dim))
|
||||
t_mod = t_mod.to(dtype=model_dtype)
|
||||
|
||||
x = self.patchify(x)
|
||||
f, h, w = x.shape[2:]
|
||||
|
||||
context = self.text_embedding(context)
|
||||
context_len = context.shape[1]
|
||||
if self.action_conditioned and action is not None:
|
||||
action_len = action.shape[1]
|
||||
action_emb = self.action_embedding(action)
|
||||
action_pos_embed = sinusoidal_embedding_1d(
|
||||
self.hidden_dim, torch.arange(action_len, device=action_emb.device)
|
||||
).to(dtype=action_emb.dtype)
|
||||
action_emb = action_emb + action_pos_embed.unsqueeze(0)
|
||||
context = torch.cat([context, action_emb], dim=1)
|
||||
|
||||
num_temporal_groups = f - 1
|
||||
if num_temporal_groups <= 0:
|
||||
raise ValueError(
|
||||
"Action-conditioned context mask requires at least 2 latent frames when `action` is provided."
|
||||
)
|
||||
if action_emb.shape[1] % num_temporal_groups != 0:
|
||||
raise ValueError(
|
||||
f"Action embedding length {action_emb.shape[1]} must be divisible by "
|
||||
f"number of temporal groups {num_temporal_groups}"
|
||||
)
|
||||
action_group_mask = create_group_causal_attn_mask(
|
||||
num_temporal_groups=num_temporal_groups,
|
||||
num_query_per_group=tokens_per_frame,
|
||||
num_key_per_group=action_len // num_temporal_groups,
|
||||
mode=self.action_group_causal_mask_mode,
|
||||
).to(context.device)
|
||||
|
||||
seq_len = f * h * w
|
||||
final_context_mask = torch.zeros(
|
||||
(batch_size, seq_len, context.shape[1]), dtype=torch.bool, device=context.device
|
||||
)
|
||||
final_context_mask[:, :, :context_len] = context_mask.unsqueeze(1).expand(-1, seq_len, -1)
|
||||
final_context_mask[:, tokens_per_frame:, context_len:] = action_group_mask.unsqueeze(0).expand(
|
||||
batch_size, -1, -1
|
||||
)
|
||||
context_mask = final_context_mask
|
||||
elif self.action_conditioned and action is None:
|
||||
if f != 1:
|
||||
raise ValueError(
|
||||
"Action-conditioned model requires `action` unless running single-frame text-only mode "
|
||||
"with num_latent_frames=1."
|
||||
)
|
||||
context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1)
|
||||
else:
|
||||
context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1)
|
||||
|
||||
x_tokens = rearrange(x, "b c f h w -> b (f h w) c").contiguous()
|
||||
grid_sizes = torch.tensor([[f, h, w]] * batch_size, dtype=torch.long, device=x_tokens.device)
|
||||
freqs = {"grid_sizes": grid_sizes, "freqs": self.freqs.to(x_tokens.device)}
|
||||
|
||||
return {
|
||||
"tokens": x_tokens,
|
||||
"freqs": freqs,
|
||||
"t": t,
|
||||
"t_mod": t_mod,
|
||||
"context": context,
|
||||
"context_mask": context_mask,
|
||||
"meta": {
|
||||
"grid_sizes": grid_sizes,
|
||||
"tokens_per_frame": tokens_per_frame,
|
||||
"batch_size": batch_size,
|
||||
},
|
||||
}
|
||||
|
||||
def post_dit(self, x_tokens: torch.Tensor, pre_state: dict[str, Any]) -> torch.Tensor:
|
||||
x = self.head(x_tokens, pre_state["t"])
|
||||
return torch.stack(super().unpatchify(x, pre_state["meta"]["grid_sizes"]))
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
context: torch.Tensor,
|
||||
context_mask: torch.Tensor | None = None,
|
||||
action: torch.Tensor | None = None,
|
||||
fuse_vae_embedding_in_latents: bool = False,
|
||||
):
|
||||
pre_state = self.pre_dit(
|
||||
x=x,
|
||||
timestep=timestep,
|
||||
context=context,
|
||||
context_mask=context_mask,
|
||||
action=action,
|
||||
fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents,
|
||||
)
|
||||
x_tokens = pre_state["tokens"]
|
||||
context_emb = pre_state["context"]
|
||||
t_mod = pre_state["t_mod"]
|
||||
freqs = pre_state["freqs"]
|
||||
context_attn_mask = pre_state["context_mask"]
|
||||
self_attn_mask = (
|
||||
self.build_video_to_video_mask(
|
||||
video_seq_len=x_tokens.shape[1],
|
||||
video_tokens_per_frame=int(pre_state["meta"]["tokens_per_frame"]),
|
||||
device=x_tokens.device,
|
||||
)
|
||||
if self.video_attention_mask_mode != "bidirectional"
|
||||
else None
|
||||
)
|
||||
|
||||
for block in self.blocks:
|
||||
if self.use_gradient_checkpointing:
|
||||
x_tokens = gradient_checkpoint_forward(
|
||||
block,
|
||||
self.use_gradient_checkpointing,
|
||||
x_tokens,
|
||||
context_emb,
|
||||
t_mod,
|
||||
freqs,
|
||||
context_mask=context_attn_mask,
|
||||
self_attn_mask=self_attn_mask,
|
||||
)
|
||||
else:
|
||||
x_tokens = block(
|
||||
x_tokens,
|
||||
context_emb,
|
||||
t_mod,
|
||||
freqs,
|
||||
context_mask=context_attn_mask,
|
||||
self_attn_mask=self_attn_mask,
|
||||
)
|
||||
|
||||
return self.post_dit(x_tokens, pre_state)
|
||||
@@ -79,15 +79,6 @@ class MolmoAct2Config(PreTrainedConfig):
|
||||
eval_seed: int | None = None
|
||||
rtc_config: RTCConfig | None = None
|
||||
|
||||
# Joint frame transform for cross-calibration compatibility.
|
||||
# Some MolmoAct2 checkpoints were trained on data using a different joint
|
||||
# convention than the current LeRobot calibration. Set both to apply a
|
||||
# sign/offset correction at runtime (state before model, action after).
|
||||
# See: https://huggingface.co/docs/lerobot/backwardcomp
|
||||
# Default is None (no transform). Both must be set together.
|
||||
joint_signs: list[float] | None = None
|
||||
joint_offsets: list[float] | None = None
|
||||
|
||||
# Default is full finetuning with gradients from the action expert flowing into the VLM.
|
||||
enable_lora_vlm: bool = False
|
||||
lora_rank: int = 64
|
||||
@@ -132,10 +123,6 @@ class MolmoAct2Config(PreTrainedConfig):
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if (self.joint_signs is None) != (self.joint_offsets is None):
|
||||
raise ValueError("joint_signs and joint_offsets must both be set or both be None.")
|
||||
if self.joint_signs is not None and len(self.joint_signs) != len(self.joint_offsets):
|
||||
raise ValueError("joint_signs and joint_offsets must have the same length.")
|
||||
if self.action_mode not in {"continuous", "discrete", "both"}:
|
||||
raise ValueError(
|
||||
f"Unsupported action_mode={self.action_mode!r}. "
|
||||
|
||||
@@ -1005,93 +1005,6 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep):
|
||||
return features
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="molmoact2_state_frame_transform")
|
||||
@dataclass
|
||||
class MolmoAct2StateFrameTransformStep(ProcessorStep):
|
||||
"""Convert robot state from arm frame to model frame before normalization.
|
||||
|
||||
Required for zero-shot deployment of MolmoAct2-SO100_101 on SO-100/101
|
||||
arms calibrated with LeRobot >= 0.5.0 (v3.0 convention). The checkpoint
|
||||
was trained on data using a different joint convention (sign flip on
|
||||
shoulder_lift, 90 deg offset on shoulder_lift and elbow_flex).
|
||||
|
||||
No-op when joint_signs and joint_offsets are None (default), so this
|
||||
step has no effect on fine-tuned models or other embodiments.
|
||||
|
||||
state_model = signs * arm_state + offsets
|
||||
|
||||
See: https://huggingface.co/docs/lerobot/backwardcomp
|
||||
"""
|
||||
|
||||
joint_signs: list[float] | None = None
|
||||
joint_offsets: list[float] | None = None
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
if self.joint_signs is None or self.joint_offsets is None:
|
||||
return transition
|
||||
observation = transition.get(TransitionKey.OBSERVATION)
|
||||
if not isinstance(observation, dict) or OBS_STATE not in observation:
|
||||
return transition
|
||||
transition = transition.copy()
|
||||
observation = observation.copy()
|
||||
state = torch.as_tensor(observation[OBS_STATE], dtype=torch.float32).clone()
|
||||
n = len(self.joint_signs)
|
||||
signs = torch.tensor(self.joint_signs, dtype=torch.float32, device=state.device)
|
||||
offsets = torch.tensor(self.joint_offsets, dtype=torch.float32, device=state.device)
|
||||
state[..., :n] = signs * state[..., :n] + offsets
|
||||
observation[OBS_STATE] = state
|
||||
transition[TransitionKey.OBSERVATION] = observation
|
||||
return transition
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
return features
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
return {"joint_signs": self.joint_signs, "joint_offsets": self.joint_offsets}
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="molmoact2_action_frame_transform")
|
||||
@dataclass
|
||||
class MolmoAct2ActionFrameTransformStep(ProcessorStep):
|
||||
"""Convert model action from model frame back to arm frame after unnormalization.
|
||||
|
||||
Inverse of MolmoAct2StateFrameTransformStep. Required for zero-shot
|
||||
MolmoAct2-SO100_101 on SO-100/101 arms. No-op when both fields are None.
|
||||
|
||||
action_arm = signs * (model_action - offsets)
|
||||
|
||||
See: https://huggingface.co/docs/lerobot/backwardcomp
|
||||
"""
|
||||
|
||||
joint_signs: list[float] | None = None
|
||||
joint_offsets: list[float] | None = None
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
if self.joint_signs is None or self.joint_offsets is None:
|
||||
return transition
|
||||
action = transition.get(TransitionKey.ACTION)
|
||||
if action is None:
|
||||
return transition
|
||||
transition = transition.copy()
|
||||
action = torch.as_tensor(action, dtype=torch.float32).clone()
|
||||
n = len(self.joint_signs)
|
||||
signs = torch.tensor(self.joint_signs, dtype=torch.float32, device=action.device)
|
||||
offsets = torch.tensor(self.joint_offsets, dtype=torch.float32, device=action.device)
|
||||
action[..., :n] = signs * (action[..., :n] - offsets)
|
||||
transition[TransitionKey.ACTION] = action
|
||||
return transition
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
return features
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
return {"joint_signs": self.joint_signs, "joint_offsets": self.joint_offsets}
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="molmoact2_clamp_action")
|
||||
@dataclass
|
||||
class MolmoAct2ClampActionProcessorStep(ProcessorStep):
|
||||
@@ -1154,10 +1067,6 @@ def make_molmoact2_pre_post_processors(
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
MolmoAct2StateFrameTransformStep(
|
||||
joint_signs=config.joint_signs,
|
||||
joint_offsets=config.joint_offsets,
|
||||
),
|
||||
MolmoAct2MaskedNormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
@@ -1193,10 +1102,6 @@ def make_molmoact2_pre_post_processors(
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=masked_dataset_stats,
|
||||
),
|
||||
MolmoAct2ActionFrameTransformStep(
|
||||
joint_signs=config.joint_signs,
|
||||
joint_offsets=config.joint_offsets,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
|
||||
@@ -11,8 +11,6 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import builtins
|
||||
import dataclasses
|
||||
@@ -21,7 +19,7 @@ import os
|
||||
from importlib.resources import files
|
||||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack
|
||||
from typing import TypedDict, TypeVar, Unpack
|
||||
|
||||
import packaging
|
||||
import safetensors
|
||||
@@ -40,13 +38,10 @@ from .utils import log_model_loading_keys
|
||||
|
||||
T = TypeVar("T", bound="PreTrainedPolicy")
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
|
||||
|
||||
def _build_card_context(
|
||||
cfg: TrainPipelineConfig | None,
|
||||
dataset_meta: LeRobotDatasetMetadata | None,
|
||||
dataset_repo_id: str | None,
|
||||
input_features: dict | None,
|
||||
output_features: dict | None,
|
||||
) -> dict:
|
||||
@@ -77,16 +72,30 @@ def _build_card_context(
|
||||
"lerobot_version": __version__,
|
||||
}
|
||||
|
||||
if dataset_meta is not None:
|
||||
context["dataset"] = {
|
||||
"repo_id": dataset_meta.repo_id,
|
||||
"episodes": dataset_meta.total_episodes,
|
||||
"frames": dataset_meta.total_frames,
|
||||
"fps": dataset_meta.fps,
|
||||
"tasks": [str(task) for task in dataset_meta.tasks.index],
|
||||
}
|
||||
context["robot_type"] = dataset_meta.robot_type
|
||||
context["cameras"] = [key.split(".")[-1] for key in dataset_meta.camera_keys]
|
||||
if dataset_repo_id:
|
||||
dataset_cfg = getattr(cfg, "dataset", None)
|
||||
try:
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
|
||||
meta = LeRobotDatasetMetadata(
|
||||
dataset_repo_id,
|
||||
root=getattr(dataset_cfg, "root", None),
|
||||
revision=getattr(dataset_cfg, "revision", None),
|
||||
)
|
||||
context["dataset"] = {
|
||||
"repo_id": dataset_repo_id,
|
||||
"episodes": meta.total_episodes,
|
||||
"frames": meta.total_frames,
|
||||
"fps": meta.fps,
|
||||
"tasks": [str(task) for task in meta.tasks.index],
|
||||
}
|
||||
context["robot_type"] = meta.robot_type
|
||||
context["cameras"] = [key.split(".")[-1] for key in meta.camera_keys]
|
||||
except Exception as e: # noqa: BLE001 — dataset details are optional, never fail the push
|
||||
logging.warning(
|
||||
f"Could not load dataset metadata for '{dataset_repo_id}'; those sections will be "
|
||||
f"omitted from the model card. ({e})"
|
||||
)
|
||||
|
||||
return context
|
||||
|
||||
@@ -295,7 +304,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
cfg: TrainPipelineConfig,
|
||||
peft_model=None,
|
||||
state_dict: dict[str, Tensor] | None = None,
|
||||
dataset_meta: LeRobotDatasetMetadata | None = None,
|
||||
):
|
||||
api = HfApi()
|
||||
repo_id = api.create_repo(
|
||||
@@ -317,12 +325,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
self.save_pretrained(saved_path, state_dict=state_dict)
|
||||
|
||||
card = self.generate_model_card(
|
||||
cfg.dataset.repo_id,
|
||||
self.config.type,
|
||||
self.config.license,
|
||||
self.config.tags,
|
||||
cfg=cfg,
|
||||
dataset_meta=dataset_meta,
|
||||
cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg
|
||||
)
|
||||
card.save(str(saved_path / "README.md"))
|
||||
|
||||
@@ -337,9 +340,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
ignore_patterns=["*.tmp", "*.log"],
|
||||
)
|
||||
|
||||
# Contract: lerobot.jobs.hf.submit_to_hf watches for this exact
|
||||
# "Model pushed to <url>" line to end a remote run early. Keep the wording
|
||||
# and URL format in sync (it falls back to status polling if they drift).
|
||||
logging.info(f"Model pushed to {commit_info.repo_url.url}")
|
||||
|
||||
def generate_model_card(
|
||||
@@ -349,7 +349,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
license: str | None,
|
||||
tags: list[str] | None,
|
||||
cfg: TrainPipelineConfig | None = None,
|
||||
dataset_meta: LeRobotDatasetMetadata | None = None,
|
||||
) -> ModelCard:
|
||||
base_model_mapping = {
|
||||
"smolvla": "lerobot/smolvla_base",
|
||||
@@ -370,7 +369,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
)
|
||||
|
||||
context = _build_card_context(
|
||||
cfg, dataset_meta, self.config.input_features, self.config.output_features
|
||||
cfg, dataset_repo_id, self.config.input_features, self.config.output_features
|
||||
)
|
||||
# Used by the template to pre-fill commands and the "Fine-tuned from" line.
|
||||
context["policy_repo_id"] = getattr(self.config, "repo_id", None)
|
||||
@@ -387,7 +386,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
self,
|
||||
peft_config=None,
|
||||
peft_cli_overrides: dict | None = None,
|
||||
) -> PreTrainedPolicy:
|
||||
) -> "PreTrainedPolicy":
|
||||
"""
|
||||
Wrap this policy with PEFT adapters for parameter-efficient fine-tuning.
|
||||
|
||||
|
||||
@@ -126,8 +126,7 @@ def prepare_observation_for_inference(
|
||||
for name in observation:
|
||||
observation[name] = torch.from_numpy(observation[name])
|
||||
if "image" in name:
|
||||
if observation[name].dtype == torch.uint8:
|
||||
observation[name] = observation[name].type(torch.float32) / 255
|
||||
observation[name] = observation[name].type(torch.float32) / 255
|
||||
observation[name] = observation[name].permute(2, 0, 1).contiguous()
|
||||
observation[name] = observation[name].unsqueeze(0)
|
||||
observation[name] = observation[name].to(device)
|
||||
|
||||
@@ -17,10 +17,12 @@ from __future__ import annotations
|
||||
import logging
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F # noqa: N812
|
||||
from PIL import Image
|
||||
from torch import Tensor, nn
|
||||
|
||||
from lerobot.policies.pretrained import PreTrainedPolicy, T
|
||||
@@ -53,13 +55,12 @@ class VLAJEPAModel(nn.Module):
|
||||
- DiT-B: flow-matching action head for future action prediction
|
||||
- V-JEPA: world model for video frame prediction
|
||||
|
||||
Inputs are batched tensors kept on the model device
|
||||
- images: List[List[Tensor [C, H, W]]] (float [0,1]) — per sample, per view (Qwen messages)
|
||||
- instructions: List[str]
|
||||
- videos: Tensor [B, V, T, C, H, W] (float [0,1], world model only)
|
||||
- actions: Tensor [B, T, action_dim] (optional, training only)
|
||||
- state: Tensor [B, 1, state_dim] (optional)
|
||||
- action_is_pad: Tensor [B, T] (optional)
|
||||
Input: List[dict] native format (same as original starVLA)
|
||||
- "image": List[PIL.Image] (multi-view images)
|
||||
- "video": np.ndarray [V, T, H, W, 3]
|
||||
- "lang": str (task instruction)
|
||||
- "action": np.ndarray [T, action_dim] (optional, training only)
|
||||
- "state": np.ndarray [1, state_dim] (optional)
|
||||
"""
|
||||
|
||||
def __init__(self, config: VLAJEPAConfig) -> None:
|
||||
@@ -74,11 +75,6 @@ class VLAJEPAModel(nn.Module):
|
||||
self.action_tokens, self.action_token_ids, self.embodied_action_token_id = (
|
||||
self.qwen.expand_tokenizer()
|
||||
)
|
||||
self.register_buffer(
|
||||
"_action_token_ids_t",
|
||||
torch.tensor(self.action_token_ids, dtype=torch.long),
|
||||
persistent=False,
|
||||
)
|
||||
|
||||
# Action head (flow-matching DiT)
|
||||
self.action_model = VLAJEPAActionHead(config, cross_attention_dim=self.qwen.model.config.hidden_size)
|
||||
@@ -165,123 +161,166 @@ class VLAJEPAModel(nn.Module):
|
||||
|
||||
# ---- Native VLA-JEPA forward (follows original VLA_JEPA.py) ----
|
||||
|
||||
def _encode_qwen(
|
||||
self, images: list[list[Tensor]], instructions: list[str], *, need_action_tokens: bool
|
||||
) -> tuple[Tensor, Tensor, Tensor | None]:
|
||||
"""Run Qwen and gather the embodied-action (and optionally action) token hidden states."""
|
||||
def forward(self, examples: list[dict]) -> dict[str, Tensor]:
|
||||
"""
|
||||
Native forward pass following original starVLA VLA_JEPA.forward.
|
||||
|
||||
Args:
|
||||
examples: List of per-sample dicts with keys:
|
||||
"image" : List[PIL.Image] — multi-view images
|
||||
"video" : np.ndarray [V, T, H, W, 3]
|
||||
"lang" : str — task instruction
|
||||
"action" : np.ndarray [T, action_dim] (optional)
|
||||
"state" : np.ndarray [1, state_dim] (optional)
|
||||
|
||||
Returns:
|
||||
dict with "action_loss" and "wm_loss" keys (scalar Tensors).
|
||||
"""
|
||||
# Unpack native format (same pattern as original VLA_JEPA.py)
|
||||
batch_images = [ex["image"] for ex in examples] # List[List[PIL.Image]]
|
||||
batch_videos = [ex["video"] for ex in examples] # List[np.ndarray]
|
||||
instructions = [ex["lang"] for ex in examples] # List[str]
|
||||
has_action = "action" in examples[0] and examples[0]["action"] is not None
|
||||
actions = [ex["action"] for ex in examples] if has_action else None
|
||||
has_state = "state" in examples[0] and examples[0]["state"] is not None
|
||||
state = [ex["state"] for ex in examples] if has_state else None
|
||||
action_is_pad = (
|
||||
[ex["action_is_pad"] for ex in examples]
|
||||
if has_action and "action_is_pad" in examples[0] and examples[0]["action_is_pad"] is not None
|
||||
else None
|
||||
)
|
||||
|
||||
# Stack videos: [B, V, T, H, W, 3] -> [B, V, T, 3, H, W]
|
||||
batch_videos = np.stack(batch_videos)
|
||||
batch_videos = batch_videos.transpose(0, 1, 2, 5, 3, 4) # [B, V, T, 3, H, W]
|
||||
|
||||
# Adjust number of views for the world model:
|
||||
# - fewer views than expected: duplicate the first view to fill up
|
||||
# - more views than expected: keep only the first num_views_world_model views
|
||||
num_views_world_model = self.config.jepa_tubelet_size
|
||||
if batch_videos.shape[1] < num_views_world_model:
|
||||
num_missing_views = num_views_world_model - batch_videos.shape[1]
|
||||
first_view = np.repeat(batch_videos[:, :1], num_missing_views, axis=1)
|
||||
batch_videos = np.concatenate([batch_videos, first_view], axis=1)
|
||||
elif batch_videos.shape[1] > num_views_world_model:
|
||||
batch_videos = batch_videos[:, :num_views_world_model]
|
||||
|
||||
# ---- Step 1: QwenVL encode (same as original) ----
|
||||
qwen_inputs = self.qwen.build_inputs(
|
||||
images=images,
|
||||
images=batch_images,
|
||||
instructions=instructions,
|
||||
action_prompt=self.replace_prompt,
|
||||
embodied_prompt=self.embodied_replace_prompt,
|
||||
)
|
||||
input_ids = qwen_inputs["input_ids"]
|
||||
embodied_idx = (input_ids == self.embodied_action_token_id).nonzero(as_tuple=True)
|
||||
action_idx = None
|
||||
if need_action_tokens:
|
||||
action_mask = torch.isin(input_ids, self._action_token_ids_t)
|
||||
action_idx = action_mask.nonzero(as_tuple=True)
|
||||
|
||||
# Locate embodied-action tokens (always needed for action head)
|
||||
embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id
|
||||
embodied_indices = embodied_mask.nonzero(as_tuple=True)
|
||||
|
||||
# Locate action tokens (only needed for world model predictor)
|
||||
if self.config.enable_world_model:
|
||||
action_mask = torch.isin(
|
||||
qwen_inputs["input_ids"],
|
||||
torch.tensor(self.action_token_ids, device=qwen_inputs["input_ids"].device),
|
||||
)
|
||||
action_indices = action_mask.nonzero(as_tuple=True)
|
||||
|
||||
device_type = next(self.parameters()).device.type
|
||||
|
||||
with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
|
||||
last_hidden = self._qwen_last_decoder_hidden(qwen_inputs) # [B, seq_len, H]
|
||||
b, _, h = last_hidden.shape
|
||||
embodied_action_tokens = last_hidden[embodied_idx[0], embodied_idx[1], :].view(b, -1, h)
|
||||
action_tokens = (
|
||||
last_hidden[action_idx[0], action_idx[1], :].view(b, -1, h)
|
||||
if action_idx is not None
|
||||
else None
|
||||
)
|
||||
return embodied_action_tokens, action_tokens
|
||||
|
||||
def _world_model_loss(self, videos: Tensor, action_tokens: Tensor) -> Tensor:
|
||||
"""JEPA encode + predictor L1 loss. `videos` is [B, V, T, C, H, W] float in [0, 1]."""
|
||||
# Match the world model's expected view count: pad with the first view, or trim extras.
|
||||
num_views = self.config.jepa_tubelet_size
|
||||
if videos.shape[1] < num_views:
|
||||
missing = num_views - videos.shape[1]
|
||||
videos = torch.cat([videos, videos[:, :1].repeat(1, missing, 1, 1, 1, 1)], dim=1)
|
||||
elif videos.shape[1] > num_views:
|
||||
videos = videos[:, :num_views]
|
||||
if self.config.enable_world_model:
|
||||
action_tokens = last_hidden[action_indices[0], action_indices[1], :].view(b, -1, h)
|
||||
|
||||
b, v, t_frames, c, h_img, w_img = videos.shape
|
||||
flat = videos.reshape(b * v, t_frames, c, h_img, w_img)
|
||||
# Fast (torchvision) video processor on-device, do_rescale=False (frames already in [0, 1]).
|
||||
video_pixels = self.video_processor(
|
||||
videos=list(flat),
|
||||
return_tensors="pt",
|
||||
device=self.video_encoder.device,
|
||||
do_rescale=False,
|
||||
)["pixel_values_videos"] # [B*V, T, C, H, W]
|
||||
embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h)
|
||||
|
||||
with torch.no_grad():
|
||||
video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels)
|
||||
# Merge views: [B*V, ...] -> [B, ..., V*embed_dim]
|
||||
video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2)
|
||||
|
||||
tubelet_size = self.video_encoder.config.tubelet_size
|
||||
# num_video_frames raw frames → t_enc_total temporal positions after tubelet compression
|
||||
t_enc_total = self.config.num_video_frames // tubelet_size
|
||||
if t_enc_total < 2:
|
||||
return torch.zeros((), device=video_embeddings.device)
|
||||
|
||||
# Shift-by-one JEPA split: input_states = positions 0..T-2, gt_states = positions 1..T-1
|
||||
t_enc_ctx = t_enc_total - 1
|
||||
tokens_per_frame = video_embeddings.shape[1] // t_enc_total
|
||||
input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :]
|
||||
gt_states = video_embeddings[:, tokens_per_frame:, :]
|
||||
|
||||
expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep
|
||||
if action_tokens.shape[1] < expected_actions:
|
||||
pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1)
|
||||
action_tokens = torch.cat([action_tokens, pad], dim=1)
|
||||
|
||||
predicted_states = self.video_predictor(
|
||||
input_states.float(), action_tokens[:, :expected_actions].float()
|
||||
)
|
||||
return F.l1_loss(predicted_states, gt_states.float(), reduction="mean")
|
||||
|
||||
def _action_loss(
|
||||
self,
|
||||
embodied_action_tokens: Tensor,
|
||||
actions: Tensor,
|
||||
state: Tensor | None,
|
||||
action_is_pad: Tensor | None,
|
||||
) -> Tensor:
|
||||
"""Flow-matching action-head loss, repeated over `repeated_diffusion_steps`."""
|
||||
device_type = next(self.parameters()).device.type
|
||||
with torch.autocast(device_type=device_type, dtype=torch.float32):
|
||||
r = self.config.repeated_diffusion_steps
|
||||
horizon = self.config.chunk_size
|
||||
actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1)
|
||||
embodied = embodied_action_tokens.repeat(r, 1, 1)
|
||||
state_rep = state.to(embodied_action_tokens.dtype).repeat(r, 1, 1) if state is not None else None
|
||||
pad_rep = action_is_pad[:, -horizon:].repeat(r, 1) if action_is_pad is not None else None
|
||||
return self.action_model(embodied, actions_target, state_rep, pad_rep)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
images: list[list[Tensor]],
|
||||
instructions: list[str],
|
||||
videos: Tensor | None = None,
|
||||
actions: Tensor | None = None,
|
||||
state: Tensor | None = None,
|
||||
action_is_pad: Tensor | None = None,
|
||||
) -> dict[str, Tensor]:
|
||||
"""Native forward: Qwen encode → optional world-model loss → optional action-head loss."""
|
||||
embodied_action_tokens, action_tokens = self._encode_qwen(
|
||||
images, instructions, need_action_tokens=self.config.enable_world_model
|
||||
)
|
||||
|
||||
if self.config.enable_world_model and videos is not None:
|
||||
wm_loss = self._world_model_loss(videos, action_tokens)
|
||||
# ---- Step 2+3: JEPA Encoder + Predictor ----
|
||||
device_wm = last_hidden.device
|
||||
if not self.config.enable_world_model:
|
||||
wm_loss = torch.tensor(0.0, device=device_wm)
|
||||
else:
|
||||
wm_loss = torch.zeros((), device=embodied_action_tokens.device)
|
||||
b, v, t_frames, c, h_img, w_img = batch_videos.shape
|
||||
batch_videos_flat = batch_videos.reshape(b * v, t_frames, c, h_img, w_img)
|
||||
|
||||
if actions is None:
|
||||
video_pixels = self.video_processor(videos=list(batch_videos_flat), return_tensors="pt")[
|
||||
"pixel_values_videos"
|
||||
].to(self.video_encoder.device) # [B*V, T, C, H, W]
|
||||
|
||||
with torch.no_grad():
|
||||
video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels)
|
||||
# Merge views: [B*V, ...] -> [B, ..., V*embed_dim]
|
||||
video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2)
|
||||
|
||||
tubelet_size = self.video_encoder.config.tubelet_size
|
||||
device_wm = video_embeddings.device
|
||||
# num_video_frames raw frames → t_enc_total temporal positions after tubelet compression
|
||||
t_enc_total = self.config.num_video_frames // tubelet_size
|
||||
|
||||
if t_enc_total < 2:
|
||||
wm_loss = torch.tensor(0.0, device=device_wm)
|
||||
else:
|
||||
# Shift-by-one JEPA split (matches original VLA_JEPA.py lines 231-232):
|
||||
# input_states: positions 0..T-2, gt_states: positions 1..T-1
|
||||
t_enc_ctx = t_enc_total - 1
|
||||
tokens_per_frame = video_embeddings.shape[1] // t_enc_total
|
||||
|
||||
input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :]
|
||||
gt_states = video_embeddings[:, tokens_per_frame:, :]
|
||||
|
||||
expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep
|
||||
if action_tokens.shape[1] < expected_actions:
|
||||
pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1)
|
||||
action_tokens = torch.cat([action_tokens, pad], dim=1)
|
||||
|
||||
predicted_states = self.video_predictor(
|
||||
input_states.float(),
|
||||
action_tokens[:, :expected_actions].float(),
|
||||
)
|
||||
|
||||
wm_loss = F.l1_loss(predicted_states, gt_states.float(), reduction="mean")
|
||||
|
||||
if not has_action:
|
||||
return {"wm_loss": wm_loss}
|
||||
|
||||
action_loss = self._action_loss(embodied_action_tokens, actions, state, action_is_pad)
|
||||
# ---- Step 4: Action Head ----
|
||||
with torch.autocast(device_type=device_type, dtype=torch.float32):
|
||||
actions_tensor = torch.tensor(
|
||||
np.array(actions), device=last_hidden.device, dtype=torch.float32
|
||||
) # [B, T_full, action_dim]
|
||||
action_horizon = self.config.chunk_size
|
||||
actions_target = actions_tensor[:, -action_horizon:, :]
|
||||
|
||||
state_tensor = None
|
||||
if state is not None:
|
||||
state_tensor = torch.tensor(
|
||||
np.array(state), device=last_hidden.device, dtype=last_hidden.dtype
|
||||
) # [B, 1, state_dim]
|
||||
|
||||
repeated_diffusion_steps = self.config.repeated_diffusion_steps
|
||||
actions_target = actions_target.repeat(repeated_diffusion_steps, 1, 1)
|
||||
embodied_action_tokens = embodied_action_tokens.repeat(repeated_diffusion_steps, 1, 1)
|
||||
if state_tensor is not None:
|
||||
state_tensor = state_tensor.repeat(repeated_diffusion_steps, 1, 1)
|
||||
|
||||
action_is_pad_rep = None
|
||||
if action_is_pad is not None:
|
||||
pad_tensor = torch.stack(
|
||||
[
|
||||
p.to(actions_target.device)
|
||||
if isinstance(p, Tensor)
|
||||
else torch.tensor(p, device=actions_target.device)
|
||||
for p in action_is_pad
|
||||
]
|
||||
) # [B, T_full]
|
||||
pad_tensor = pad_tensor[:, -action_horizon:] # [B, action_horizon]
|
||||
action_is_pad_rep = pad_tensor.repeat(repeated_diffusion_steps, 1) # [B*R, action_horizon]
|
||||
|
||||
action_loss = self.action_model(
|
||||
embodied_action_tokens, actions_target, state_tensor, action_is_pad_rep
|
||||
)
|
||||
|
||||
return {"action_loss": action_loss, "wm_loss": wm_loss * self.config.world_model_loss_weight}
|
||||
|
||||
# ---- Native predict_action (follows original VLA_JEPA.predict_action) ----
|
||||
@@ -289,23 +328,58 @@ class VLAJEPAModel(nn.Module):
|
||||
@torch.no_grad()
|
||||
def predict_action(
|
||||
self,
|
||||
images: list[list[Tensor]],
|
||||
batch_images: list[list[Image.Image]],
|
||||
instructions: list[str],
|
||||
state: Tensor | None = None,
|
||||
) -> Tensor:
|
||||
"""Predict an action chunk. `images` is per-sample, per-view float [0,1] [C, H, W] tensors."""
|
||||
state: np.ndarray | None = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Native action prediction following original VLA_JEPA.predict_action.
|
||||
|
||||
Args:
|
||||
batch_images: List of samples; each is List[PIL.Image] (multi-view).
|
||||
instructions: Task instructions, one per sample.
|
||||
state: Optional [B, state_dim] numpy array.
|
||||
|
||||
Returns:
|
||||
np.ndarray [B, action_horizon, action_dim] — predicted actions.
|
||||
"""
|
||||
if self.config.resize_images_to is not None:
|
||||
height, width = self.config.resize_images_to
|
||||
images = [
|
||||
[F.interpolate(img[None], size=(height, width), mode="area")[0] for img in views]
|
||||
for views in images
|
||||
resampling = getattr(Image, "Resampling", Image).BOX
|
||||
batch_images = [
|
||||
[image.resize((width, height), resample=resampling) for image in sample_images]
|
||||
for sample_images in batch_images
|
||||
]
|
||||
|
||||
embodied_action_tokens, _ = self._encode_qwen(images, instructions, need_action_tokens=False)
|
||||
return self.action_model.predict_action(
|
||||
embodied_action_tokens.float(), state.float() if state is not None else None
|
||||
qwen_inputs = self.qwen.build_inputs(
|
||||
images=batch_images,
|
||||
instructions=instructions,
|
||||
action_prompt=self.replace_prompt,
|
||||
embodied_prompt=self.embodied_replace_prompt,
|
||||
)
|
||||
|
||||
embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id
|
||||
embodied_indices = embodied_mask.nonzero(as_tuple=True)
|
||||
|
||||
device_type = next(self.parameters()).device.type
|
||||
|
||||
with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
|
||||
last_hidden = self._qwen_last_decoder_hidden(qwen_inputs) # [B, seq_len, H]
|
||||
b, _, h = last_hidden.shape
|
||||
embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h)
|
||||
|
||||
state_tensor = None
|
||||
if state is not None:
|
||||
state_tensor = torch.from_numpy(np.array(state)).to(
|
||||
device=last_hidden.device, dtype=last_hidden.dtype
|
||||
)
|
||||
|
||||
pred_actions = self.action_model.predict_action(
|
||||
embodied_action_tokens.float(), state_tensor.float() if state_tensor is not None else None
|
||||
) # [B, action_horizon, action_dim]
|
||||
|
||||
return pred_actions.detach().cpu().numpy()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# LeRobot Adapter Layer - converts between LeRobot batch format and native VLA-JEPA format
|
||||
@@ -316,9 +390,9 @@ class VLAJEPAPolicy(PreTrainedPolicy):
|
||||
"""
|
||||
LeRobot adapter for VLA-JEPA.
|
||||
|
||||
Converts LeRobot's standard batch format (dict[str, Tensor]) to the batched tensors
|
||||
the native model expects (keeping everything on-device), calls the native model, and
|
||||
converts outputs back to LeRobot format.
|
||||
Converts LeRobot's standard batch format (dict[str, Tensor]) to the native
|
||||
VLA-JEPA format (List[dict]), calls the native model, and converts outputs
|
||||
back to LeRobot format.
|
||||
"""
|
||||
|
||||
config_class = VLAJEPAConfig
|
||||
@@ -345,8 +419,9 @@ class VLAJEPAPolicy(PreTrainedPolicy):
|
||||
|
||||
# ---- Format Conversion: LeRobot → Native ----
|
||||
|
||||
def _prepare_model_inputs(self, batch: dict[str, Tensor], training=True) -> dict[str, Any]:
|
||||
"""Convert a LeRobot batch to the model's batched, on-device inputs.
|
||||
def _prepare_model_inputs(self, batch: dict[str, Tensor]) -> list[dict]:
|
||||
"""
|
||||
Convert LeRobot batch format to native VLA-JEPA examples format.
|
||||
|
||||
LeRobot format:
|
||||
batch = {
|
||||
@@ -356,25 +431,65 @@ class VLAJEPAPolicy(PreTrainedPolicy):
|
||||
"task": str | List[str], (optional instruction)
|
||||
}
|
||||
|
||||
Returns the kwargs for `VLAJEPAModel.forward` / `.predict_action` (everything stays
|
||||
on the batch device; no per-sample shredding): `images` (per-sample, per-view list for
|
||||
Qwen messages), `instructions`, and the batched `videos` / `actions` / `state` /
|
||||
`action_is_pad` when present.
|
||||
Native format (List[dict]):
|
||||
{
|
||||
"image": List[PIL.Image], # multi-view images per sample
|
||||
"video": np.ndarray [V, T, H, W, 3],
|
||||
"lang": str, # task instruction
|
||||
"action": np.ndarray [T, action_dim], # optional
|
||||
"state": np.ndarray [1, state_dim], # optional
|
||||
}
|
||||
"""
|
||||
# Determine batch size from the first image feature
|
||||
image_keys = list(self.config.image_features.keys())
|
||||
if not image_keys:
|
||||
raise ValueError("VLAJEPA requires at least one image feature.")
|
||||
batch_size = batch[image_keys[0]].shape[0]
|
||||
first_key = image_keys[0]
|
||||
first_tensor = batch[first_key]
|
||||
batch_size = first_tensor.shape[0]
|
||||
|
||||
# Current-frame image per view ([B, C, H, W]); regroup per sample for Qwen messages.
|
||||
frames = []
|
||||
# ---- Collect images per sample ----
|
||||
# images_per_sample[b][v] = PIL.Image for view v
|
||||
images_per_sample: list[list[Image.Image]] = [[] for _ in range(batch_size)]
|
||||
for key in image_keys:
|
||||
t = batch[key]
|
||||
if t.ndim == 5: # [B, T, C, H, W] -> current observation (delta=0)
|
||||
t = t[:, 0]
|
||||
frames.append(self.model.qwen.to_pixel_values(t))
|
||||
images = [[frame[b] for frame in frames] for b in range(batch_size)]
|
||||
tensor = batch[key] # [B, C, H, W] or [B, T, C, H, W]
|
||||
if tensor.ndim == 5:
|
||||
# observation_delta_indices = [0, 1, ..., num_video_frames-1]
|
||||
# index 0 is the current observation (delta=0)
|
||||
tensor = tensor[:, 0]
|
||||
for b in range(batch_size):
|
||||
images_per_sample[b].append(self.model.qwen.tensor_to_pil(tensor[b]))
|
||||
|
||||
# ---- Collect videos per sample ----
|
||||
# Build video arrays: for each sample, stack views as [V, T, H, W, 3]
|
||||
# Check whether any image feature has a time dimension
|
||||
video_source = None
|
||||
for k in image_keys:
|
||||
if k in batch:
|
||||
video_source = batch[k] # Use first available for shape inspection
|
||||
break
|
||||
|
||||
if video_source is None:
|
||||
raise ValueError("No image data found in batch for video construction.")
|
||||
|
||||
videos_per_sample = []
|
||||
for b in range(batch_size):
|
||||
sample_views = []
|
||||
for k in image_keys:
|
||||
t = batch[k][b] # [C, H, W] or [T, C, H, W]
|
||||
if t.ndim == 3:
|
||||
t = t.unsqueeze(0) # [1, C, H, W]
|
||||
# Convert to [T, H, W, 3] numpy
|
||||
t_np = t.permute(0, 2, 3, 1).detach().cpu().float().numpy()
|
||||
# Clamp to [0, 255]
|
||||
if t_np.max() <= 1.0:
|
||||
t_np = t_np * 255.0
|
||||
t_np = np.rint(t_np.clip(0, 255)).astype(np.uint8)
|
||||
sample_views.append(t_np)
|
||||
# Stack views: [V, T, H, W, 3]
|
||||
videos_per_sample.append(np.stack(sample_views, axis=0))
|
||||
|
||||
# ---- Collect instructions ----
|
||||
tasks = batch.get("task")
|
||||
if tasks is None:
|
||||
instructions = ["Execute the robot action."] * batch_size
|
||||
@@ -383,32 +498,52 @@ class VLAJEPAPolicy(PreTrainedPolicy):
|
||||
else:
|
||||
instructions = list(tasks)
|
||||
|
||||
inputs: dict[str, Any] = {"images": images, "instructions": instructions}
|
||||
# ---- Collect actions (training only) ----
|
||||
actions_list = None
|
||||
action_is_pad_list = None
|
||||
actions_tensor = batch.get(ACTION)
|
||||
if actions_tensor is not None:
|
||||
if actions_tensor.ndim == 2:
|
||||
actions_tensor = actions_tensor.unsqueeze(1)
|
||||
actions_list = [actions_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)]
|
||||
action_is_pad_tensor = batch.get("action_is_pad")
|
||||
if action_is_pad_tensor is not None:
|
||||
action_is_pad_list = [action_is_pad_tensor[b].detach().cpu() for b in range(batch_size)]
|
||||
|
||||
# Videos [B, V, T, C, H, W] - only assembled during training when the world model consumes them.
|
||||
if self.model.config.enable_world_model and training:
|
||||
views = [batch[k].unsqueeze(1) if batch[k].ndim == 4 else batch[k] for k in image_keys]
|
||||
inputs["videos"] = self.model.qwen.to_pixel_values(torch.stack(views, dim=1))
|
||||
# ---- Collect state ----
|
||||
state_list = None
|
||||
state_tensor = batch.get(OBS_STATE)
|
||||
if state_tensor is not None:
|
||||
if state_tensor.ndim > 2:
|
||||
state_tensor = state_tensor[:, -1, :]
|
||||
if state_tensor.ndim == 2:
|
||||
state_tensor = state_tensor.unsqueeze(1) # [B, 1, state_dim]
|
||||
state_list = [state_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)]
|
||||
|
||||
actions = batch.get(ACTION)
|
||||
if actions is not None:
|
||||
inputs["actions"] = (actions.unsqueeze(1) if actions.ndim == 2 else actions).float()
|
||||
if (pad := batch.get("action_is_pad")) is not None:
|
||||
inputs["action_is_pad"] = pad
|
||||
# ---- Assemble native examples ----
|
||||
examples = []
|
||||
for b in range(batch_size):
|
||||
example = {
|
||||
"image": images_per_sample[b],
|
||||
"video": videos_per_sample[b],
|
||||
"lang": instructions[b],
|
||||
}
|
||||
if actions_list is not None:
|
||||
example["action"] = actions_list[b]
|
||||
if action_is_pad_list is not None:
|
||||
example["action_is_pad"] = action_is_pad_list[b]
|
||||
if state_list is not None:
|
||||
example["state"] = state_list[b]
|
||||
examples.append(example)
|
||||
|
||||
state = batch.get(OBS_STATE)
|
||||
if state is not None:
|
||||
if state.ndim > 2:
|
||||
state = state[:, -1, :]
|
||||
inputs["state"] = (state.unsqueeze(1) if state.ndim == 2 else state).float() # [B, 1, dim]
|
||||
|
||||
return inputs
|
||||
return examples
|
||||
|
||||
# ---- LeRobot Policy Interface ----
|
||||
|
||||
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
|
||||
"""LeRobot train forward: convert → native forward → aggregate losses."""
|
||||
native_output = self.model.forward(**self._prepare_model_inputs(batch, training=True))
|
||||
examples = self._prepare_model_inputs(batch)
|
||||
native_output = self.model.forward(examples)
|
||||
|
||||
ref = next(iter(native_output.values()))
|
||||
zero = torch.zeros((), device=ref.device, dtype=ref.dtype)
|
||||
@@ -426,9 +561,16 @@ class VLAJEPAPolicy(PreTrainedPolicy):
|
||||
self.eval()
|
||||
self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
|
||||
|
||||
inputs = self._prepare_model_inputs(batch, training=False)
|
||||
actions = self.model.predict_action(inputs["images"], inputs["instructions"], inputs.get("state"))
|
||||
return actions.to(device=self.config.device, dtype=torch.float32)
|
||||
examples = self._prepare_model_inputs(batch)
|
||||
batch_images = [ex["image"] for ex in examples]
|
||||
instructions = [ex["lang"] for ex in examples]
|
||||
|
||||
state_np = None
|
||||
if "state" in examples[0] and examples[0]["state"] is not None:
|
||||
state_np = np.stack([ex["state"] for ex in examples])
|
||||
|
||||
actions_np = self.model.predict_action(batch_images, instructions, state_np)
|
||||
return torch.from_numpy(actions_np).to(device=self.config.device, dtype=torch.float32)
|
||||
|
||||
@torch.no_grad()
|
||||
def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
|
||||
|
||||
@@ -17,7 +17,9 @@ from __future__ import annotations
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from lerobot.utils.import_utils import _transformers_available
|
||||
|
||||
@@ -76,7 +78,7 @@ class Qwen3VLInterface(torch.nn.Module):
|
||||
|
||||
def build_inputs(
|
||||
self,
|
||||
images: Sequence[Sequence[torch.Tensor]],
|
||||
images: Sequence[Sequence[Image.Image]],
|
||||
instructions: Sequence[str],
|
||||
action_prompt: str,
|
||||
embodied_prompt: str,
|
||||
@@ -92,42 +94,24 @@ class Qwen3VLInterface(torch.nn.Module):
|
||||
content.append({"type": "text", "text": prompt})
|
||||
messages.append([{"role": "user", "content": content}])
|
||||
|
||||
# The Qwen image processor is a torchvision-backed fast processor: passing the
|
||||
# images as GPU tensors (with `device`) keeps the whole vision pipeline on-device
|
||||
# and avoids a GPU->CPU->GPU roundtrip. The image tensors are forwarded through
|
||||
# apply_chat_template untouched into Qwen3VLProcessor.__call__.
|
||||
# do_rescale=False: images already arrive as float in [0, 1] (the dataset decoder
|
||||
# yields float32/255 and VISUAL normalization is IDENTITY), so we skip the
|
||||
# processor's /255 rescale instead of round-tripping through uint8.
|
||||
batch_inputs = self.processor.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_dict=True,
|
||||
processor_kwargs={
|
||||
"padding": True,
|
||||
"return_tensors": "pt",
|
||||
"device": self.model.device,
|
||||
"do_rescale": False,
|
||||
},
|
||||
processor_kwargs={"padding": True, "return_tensors": "pt"},
|
||||
)
|
||||
return batch_inputs.to(self.model.device)
|
||||
|
||||
@staticmethod
|
||||
def to_pixel_values(image_tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""Prepare an image/video tensor for the fast processors (used with do_rescale=False).
|
||||
|
||||
The dataset decoder yields float32 in [0, 1] (channels-first) and VISUAL
|
||||
normalization is IDENTITY, so the tensor already arrives in [0, 1]; we pass it
|
||||
through as float and let the processors normalize (no rescale, no uint8
|
||||
quantization). A single channel is expanded to 3 to match the RGB processors.
|
||||
|
||||
Works for any channels-first layout (channel dim is -3): [C, H, W], [B, C, H, W],
|
||||
[T, C, H, W], [B, V, T, C, H, W], ...
|
||||
"""
|
||||
image = image_tensor.detach().float()
|
||||
if image.shape[-3] == 1:
|
||||
repeats = [1] * image.ndim
|
||||
repeats[-3] = 3
|
||||
image = image.repeat(*repeats)
|
||||
return image
|
||||
def tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image:
|
||||
image = image_tensor.detach().cpu()
|
||||
if image.ndim == 3 and image.shape[0] in (1, 3):
|
||||
image = image.permute(1, 2, 0)
|
||||
image = image.float()
|
||||
if image.max() <= 1.0:
|
||||
image = image * 255.0
|
||||
image = image.clamp(0, 255).round().to(torch.uint8).numpy()
|
||||
if image.shape[-1] == 1:
|
||||
image = np.repeat(image, 3, axis=-1)
|
||||
return Image.fromarray(image)
|
||||
|
||||
@@ -65,13 +65,7 @@ class BiRebotB601Follower(BimanualMixin, Robot):
|
||||
cameras=left_arm_cameras,
|
||||
motor_can_ids=config.left_arm_config.motor_can_ids,
|
||||
pos_vel_velocity=config.left_arm_config.pos_vel_velocity,
|
||||
control_mode=config.left_arm_config.control_mode,
|
||||
mit_kp=config.left_arm_config.mit_kp,
|
||||
mit_kd=config.left_arm_config.mit_kd,
|
||||
gripper_control_mode=config.left_arm_config.gripper_control_mode,
|
||||
gripper_torque_ratio=config.left_arm_config.gripper_torque_ratio,
|
||||
gripper_mit_kp=config.left_arm_config.gripper_mit_kp,
|
||||
gripper_mit_kd=config.left_arm_config.gripper_mit_kd,
|
||||
joint_limits=config.left_arm_config.joint_limits,
|
||||
)
|
||||
|
||||
@@ -86,13 +80,7 @@ class BiRebotB601Follower(BimanualMixin, Robot):
|
||||
cameras=config.right_arm_config.cameras,
|
||||
motor_can_ids=config.right_arm_config.motor_can_ids,
|
||||
pos_vel_velocity=config.right_arm_config.pos_vel_velocity,
|
||||
control_mode=config.right_arm_config.control_mode,
|
||||
mit_kp=config.right_arm_config.mit_kp,
|
||||
mit_kd=config.right_arm_config.mit_kd,
|
||||
gripper_control_mode=config.right_arm_config.gripper_control_mode,
|
||||
gripper_torque_ratio=config.right_arm_config.gripper_torque_ratio,
|
||||
gripper_mit_kp=config.right_arm_config.gripper_mit_kp,
|
||||
gripper_mit_kd=config.right_arm_config.gripper_mit_kd,
|
||||
joint_limits=config.right_arm_config.joint_limits,
|
||||
)
|
||||
|
||||
|
||||
@@ -66,14 +66,9 @@ class HopeJrArm(Robot):
|
||||
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
features: dict[str, tuple] = {}
|
||||
for cam in self.cameras:
|
||||
cfg = self.config.cameras[cam]
|
||||
if getattr(cfg, "use_rgb", True):
|
||||
features[cam] = (cfg.height, cfg.width, 3)
|
||||
if getattr(cfg, "use_depth", False):
|
||||
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
||||
return features
|
||||
return {
|
||||
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
|
||||
}
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
@@ -144,17 +139,10 @@ class HopeJrArm(Robot):
|
||||
|
||||
# Capture images from cameras
|
||||
for cam_key, cam in self.cameras.items():
|
||||
if getattr(cam, "use_rgb", True):
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
if getattr(cam, "use_depth", False):
|
||||
start = time.perf_counter()
|
||||
obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms")
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
return obs_dict
|
||||
|
||||
|
||||
@@ -102,14 +102,9 @@ class HopeJrHand(Robot):
|
||||
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
features: dict[str, tuple] = {}
|
||||
for cam in self.cameras:
|
||||
cfg = self.config.cameras[cam]
|
||||
if getattr(cfg, "use_rgb", True):
|
||||
features[cam] = (cfg.height, cfg.width, 3)
|
||||
if getattr(cfg, "use_depth", False):
|
||||
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
||||
return features
|
||||
return {
|
||||
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
|
||||
}
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
@@ -175,17 +170,10 @@ class HopeJrHand(Robot):
|
||||
|
||||
# Capture images from cameras
|
||||
for cam_key, cam in self.cameras.items():
|
||||
if getattr(cam, "use_rgb", True):
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
if getattr(cam, "use_depth", False):
|
||||
start = time.perf_counter()
|
||||
obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms")
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
return obs_dict
|
||||
|
||||
|
||||
@@ -68,14 +68,9 @@ class KochFollower(Robot):
|
||||
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
features: dict[str, tuple] = {}
|
||||
for cam in self.cameras:
|
||||
cfg = self.config.cameras[cam]
|
||||
if getattr(cfg, "use_rgb", True):
|
||||
features[cam] = (cfg.height, cfg.width, 3)
|
||||
if getattr(cfg, "use_depth", False):
|
||||
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
||||
return features
|
||||
return {
|
||||
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
|
||||
}
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
@@ -197,17 +192,10 @@ class KochFollower(Robot):
|
||||
|
||||
# Capture images from cameras
|
||||
for cam_key, cam in self.cameras.items():
|
||||
if getattr(cam, "use_rgb", True):
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
if getattr(cam, "use_depth", False):
|
||||
start = time.perf_counter()
|
||||
obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms")
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
return obs_dict
|
||||
|
||||
|
||||
@@ -72,12 +72,6 @@ class LeKiwi(Robot):
|
||||
)
|
||||
self.arm_motors = [motor for motor in self.bus.motors if motor.startswith("arm")]
|
||||
self.base_motors = [motor for motor in self.bus.motors if motor.startswith("base")]
|
||||
depth_cameras = [name for name, cfg in config.cameras.items() if getattr(cfg, "use_depth", False)]
|
||||
if depth_cameras:
|
||||
raise NotImplementedError(
|
||||
f"Depth cameras are not supported on LeKiwi (got depth-enabled cameras: {depth_cameras}). "
|
||||
"The host/client transport only carries color frames."
|
||||
)
|
||||
self.cameras = make_cameras_from_configs(config.cameras)
|
||||
|
||||
@property
|
||||
|
||||
@@ -44,13 +44,6 @@ class LeKiwiClient(Robot):
|
||||
self.id = config.id
|
||||
self.robot_type = config.type
|
||||
|
||||
depth_cameras = [name for name, cfg in config.cameras.items() if getattr(cfg, "use_depth", False)]
|
||||
if depth_cameras:
|
||||
raise NotImplementedError(
|
||||
f"Depth cameras are not supported on LeKiwi (got depth-enabled cameras: {depth_cameras}). "
|
||||
"The host/client transport only carries color frames."
|
||||
)
|
||||
|
||||
self.remote_ip = config.remote_ip
|
||||
self.port_zmq_cmd = config.port_zmq_cmd
|
||||
self.port_zmq_observations = config.port_zmq_observations
|
||||
|
||||
@@ -68,14 +68,9 @@ class OmxFollower(Robot):
|
||||
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
features: dict[str, tuple] = {}
|
||||
for cam in self.cameras:
|
||||
cfg = self.config.cameras[cam]
|
||||
if getattr(cfg, "use_rgb", True):
|
||||
features[cam] = (cfg.height, cfg.width, 3)
|
||||
if getattr(cfg, "use_depth", False):
|
||||
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
||||
return features
|
||||
return {
|
||||
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
|
||||
}
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
@@ -180,17 +175,10 @@ class OmxFollower(Robot):
|
||||
|
||||
# Capture images from cameras
|
||||
for cam_key, cam in self.cameras.items():
|
||||
if getattr(cam, "use_rgb", True):
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
if getattr(cam, "use_depth", False):
|
||||
start = time.perf_counter()
|
||||
obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms")
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
return obs_dict
|
||||
|
||||
|
||||
@@ -101,14 +101,9 @@ class OpenArmFollower(Robot):
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
"""Camera features for observation space."""
|
||||
features: dict[str, tuple] = {}
|
||||
for cam in self.cameras:
|
||||
cfg = self.config.cameras[cam]
|
||||
if getattr(cfg, "use_rgb", True):
|
||||
features[cam] = (cfg.height, cfg.width, 3)
|
||||
if getattr(cfg, "use_depth", False):
|
||||
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
||||
return features
|
||||
return {
|
||||
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
|
||||
}
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
@@ -247,17 +242,10 @@ class OpenArmFollower(Robot):
|
||||
|
||||
# Capture images from cameras
|
||||
for cam_key, cam in self.cameras.items():
|
||||
if getattr(cam, "use_rgb", True):
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
if getattr(cam, "use_depth", False):
|
||||
start = time.perf_counter()
|
||||
obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms")
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} get_observation took: {dt_ms:.1f}ms")
|
||||
|
||||
@@ -65,33 +65,18 @@ class RebotB601FollowerConfig:
|
||||
}
|
||||
)
|
||||
|
||||
# Max speed (deg/s) per joint for POS_VEL arms and FORCE_POS gripper (motor order).
|
||||
pos_vel_velocity: float | list[float] = field(
|
||||
default_factory=lambda: [150.0, 150.0, 150.0, 150.0, 150.0, 150.0, 900.0]
|
||||
)
|
||||
# Target velocity for joints running in POS_VEL mode, in degrees/s. A scalar is
|
||||
# applied to every joint; a list provides one value per joint (in motor order).
|
||||
pos_vel_velocity: float | list[float] = field(default_factory=lambda: [150.0] * 7)
|
||||
|
||||
# Arm control: "mit" or "pos_vel".
|
||||
control_mode: str = "mit"
|
||||
|
||||
# MIT kp/kd per arm joint (motor order). Unused when control_mode="pos_vel".
|
||||
mit_kp: float | list[float] = field(default_factory=lambda: [45.0, 45.0, 45.0, 8.0, 9.0, 8.0, 8.0])
|
||||
mit_kd: float | list[float] = field(default_factory=lambda: [12.0, 12.0, 12.0, 1.0, 1.0, 1.0, 1.0])
|
||||
|
||||
# Gripper control: "force_pos" or "mit".
|
||||
gripper_control_mode: str = "force_pos"
|
||||
|
||||
# FORCE_POS only: max grip force, in [0, 1].
|
||||
gripper_torque_ratio: float = 0.07
|
||||
|
||||
# MIT only.
|
||||
gripper_mit_kp: float = 8.0
|
||||
gripper_mit_kd: float = 0.3
|
||||
# Torque/current ratio for the gripper's FORCE_POS mode, in range [0, 1].
|
||||
gripper_torque_ratio: float = 0.1
|
||||
|
||||
# Soft joint limits (degrees). These are clipped against on every action.
|
||||
joint_limits: dict[str, tuple[float, float]] = field(
|
||||
default_factory=lambda: {
|
||||
"shoulder_pan": (-150.0, 150.0),
|
||||
"shoulder_lift": (-200.0, 1.0),
|
||||
"shoulder_pan": (-145.0, 145.0),
|
||||
"shoulder_lift": (-170.0, 1.0),
|
||||
"elbow_flex": (-200.0, 1.0),
|
||||
"wrist_flex": (-80.0, 90.0),
|
||||
"wrist_yaw": (-90.0, 90.0),
|
||||
|
||||
@@ -80,14 +80,9 @@ class RebotB601Follower(Robot):
|
||||
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
features: dict[str, tuple] = {}
|
||||
for cam in self.cameras:
|
||||
cfg = self.config.cameras[cam]
|
||||
if getattr(cfg, "use_rgb", True):
|
||||
features[cam] = (cfg.height, cfg.width, 3)
|
||||
if getattr(cfg, "use_depth", False):
|
||||
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
||||
return features
|
||||
return {
|
||||
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
|
||||
}
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
@@ -174,25 +169,11 @@ class RebotB601Follower(Robot):
|
||||
print(f"Calibration saved to {self.calibration_fpath}")
|
||||
|
||||
def configure(self) -> None:
|
||||
if self.config.control_mode not in ("pos_vel", "mit"):
|
||||
raise ValueError(
|
||||
f"Unsupported control_mode '{self.config.control_mode}'. Use 'pos_vel' or 'mit'."
|
||||
)
|
||||
if self.config.gripper_control_mode not in ("force_pos", "mit"):
|
||||
raise ValueError(
|
||||
f"Unsupported gripper_control_mode '{self.config.gripper_control_mode}'. "
|
||||
"Use 'force_pos' or 'mit'."
|
||||
)
|
||||
use_mit = self.config.control_mode == "mit"
|
||||
gripper_use_mit = self.config.gripper_control_mode == "mit"
|
||||
self.bus.enable_all()
|
||||
for motor_name, motor in self.motors.items():
|
||||
if motor_name == GRIPPER_MOTOR:
|
||||
target_mode = MotorBridgeMode.MIT if gripper_use_mit else MotorBridgeMode.FORCE_POS
|
||||
elif use_mit:
|
||||
target_mode = MotorBridgeMode.MIT
|
||||
else:
|
||||
target_mode = MotorBridgeMode.POS_VEL
|
||||
target_mode = (
|
||||
MotorBridgeMode.FORCE_POS if motor_name == GRIPPER_MOTOR else MotorBridgeMode.POS_VEL
|
||||
)
|
||||
for attempt in range(_ENSURE_MODE_RETRIES + 1):
|
||||
try:
|
||||
motor.ensure_mode(target_mode)
|
||||
@@ -232,17 +213,10 @@ class RebotB601Follower(Robot):
|
||||
logger.debug(f"{self} read state: {dt_ms:.1f}ms")
|
||||
|
||||
for cam_key, cam in self.cameras.items():
|
||||
if getattr(cam, "use_rgb", True):
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
if getattr(cam, "use_depth", False):
|
||||
start = time.perf_counter()
|
||||
obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms")
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
return obs_dict
|
||||
|
||||
@@ -278,34 +252,22 @@ class RebotB601Follower(Robot):
|
||||
goal_present_pos = {key: (g, present_pos.get(key, g)) for key, g in goal_pos.items()}
|
||||
goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
|
||||
|
||||
use_mit = self.config.control_mode == "mit"
|
||||
for motor_name, position_deg in goal_pos.items():
|
||||
motor = self.motors.get(motor_name)
|
||||
if motor is None:
|
||||
continue
|
||||
idx = self.motor_names.index(motor_name)
|
||||
vel_deg_s = (
|
||||
self.config.pos_vel_velocity[idx]
|
||||
if isinstance(self.config.pos_vel_velocity, list)
|
||||
else self.config.pos_vel_velocity
|
||||
)
|
||||
pos_rad = math.radians(position_deg)
|
||||
vel_rad = math.radians(vel_deg_s)
|
||||
if motor_name == GRIPPER_MOTOR:
|
||||
if self.config.gripper_control_mode == "mit":
|
||||
motor.send_mit(pos_rad, 0.0, self.config.gripper_mit_kp, self.config.gripper_mit_kd, 0.0)
|
||||
else:
|
||||
vel_deg_s = (
|
||||
self.config.pos_vel_velocity[idx]
|
||||
if isinstance(self.config.pos_vel_velocity, list)
|
||||
else self.config.pos_vel_velocity
|
||||
)
|
||||
motor.send_force_pos(pos_rad, math.radians(vel_deg_s), self.config.gripper_torque_ratio)
|
||||
elif use_mit:
|
||||
kp = self.config.mit_kp[idx] if isinstance(self.config.mit_kp, list) else self.config.mit_kp
|
||||
kd = self.config.mit_kd[idx] if isinstance(self.config.mit_kd, list) else self.config.mit_kd
|
||||
motor.send_mit(pos_rad, 0.0, kp, kd, 0.0)
|
||||
motor.send_force_pos(pos_rad, vel_rad, self.config.gripper_torque_ratio)
|
||||
else:
|
||||
vel_deg_s = (
|
||||
self.config.pos_vel_velocity[idx]
|
||||
if isinstance(self.config.pos_vel_velocity, list)
|
||||
else self.config.pos_vel_velocity
|
||||
)
|
||||
motor.send_pos_vel(pos_rad, math.radians(vel_deg_s))
|
||||
motor.send_pos_vel(pos_rad, vel_rad)
|
||||
|
||||
return {f"{motor}.pos": val for motor, val in goal_pos.items()}
|
||||
|
||||
|
||||
@@ -68,13 +68,9 @@ class SOFollower(Robot):
|
||||
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
features: dict[str, tuple] = {}
|
||||
for cam in self.cameras:
|
||||
if getattr(self.cameras[cam], "use_rgb", True):
|
||||
features[cam] = (self.cameras[cam].height, self.cameras[cam].width, 3)
|
||||
if getattr(self.cameras[cam], "use_depth", False):
|
||||
features[f"{cam}_depth"] = (self.cameras[cam].height, self.cameras[cam].width, 1)
|
||||
return features
|
||||
return {
|
||||
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
|
||||
}
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
@@ -189,17 +185,10 @@ class SOFollower(Robot):
|
||||
|
||||
# Capture images from cameras
|
||||
for cam_key, cam in self.cameras.items():
|
||||
if getattr(cam, "use_rgb", True):
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
if getattr(cam, "use_depth", False):
|
||||
start = time.perf_counter()
|
||||
obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms")
|
||||
start = time.perf_counter()
|
||||
obs_dict[cam_key] = cam.read_latest()
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||
|
||||
return obs_dict
|
||||
|
||||
|
||||
@@ -68,6 +68,6 @@ class UnitreeG1Config(RobotConfig):
|
||||
# Compensates for gravity on the unitree's arms using the arm ik solver
|
||||
gravity_compensation: bool = False
|
||||
|
||||
# Lower-body controller class name, e.g. "GrootLocomotionController" or
|
||||
# "HolosomaLocomotionController". None disables it.
|
||||
# Locomotion controller class name, e.g. "GrootLocomotionController",
|
||||
# "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
|
||||
controller: str | None = None
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
@@ -12,6 +14,12 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# Importing concrete policy configs registers their draccus `--policy.type`
|
||||
# choices (e.g. "act") so tests can parse them.
|
||||
from lerobot.policies.act.configuration_act import ACTConfig # noqa: F401
|
||||
"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC, mjlab motion imitation)."""
|
||||
|
||||
__all__ = [
|
||||
"GrootLocomotionController",
|
||||
"HolosomaLocomotionController",
|
||||
"SonicWholeBodyController",
|
||||
"SonicRuntime",
|
||||
"MjlabMotionImitationController",
|
||||
]
|
||||
+1
-1
@@ -21,7 +21,7 @@ import numpy as np
|
||||
import onnxruntime as ort
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
from .g1_utils import (
|
||||
from lerobot.robots.unitree_g1.g1_utils import (
|
||||
REMOTE_AXES,
|
||||
REMOTE_BUTTONS,
|
||||
G1_29_JointIndex,
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user