mirror of
https://github.com/huggingface/lerobot.git
synced 2026-06-18 08:47:05 +00:00
pepijn223
f9d12db9cf
- Add tests/policies/lingbot_va/__init__.py so the test files don't clash by basename with tests/policies/vla_jepa/* under pytest's default import mode (fast-test collection error). - Fix vendored typos flagged by the typos hook (pach_scale->patch_scale, total_tolen-> total_token_len, stablized->stabilized) and a mypy union-attr in RoboTwinEnv._read_eef_pose. - Apply Prettier formatting to docs/source/lingbot_va.mdx. Co-authored-by: Cursor <cursoragent@cursor.com>
163 lines
8.6 KiB
Plaintext
163 lines
8.6 KiB
Plaintext
# LingBot-VA
|
||
|
||
LingBot-VA is an **autoregressive video-action world-model policy** built on the **Wan2.2**
|
||
video-diffusion stack. It interleaves, in one autoregressive sequence, the prediction of
|
||
future **video latents** and **robot actions** ("VA" = Video-Action). The LeRobot
|
||
integration wires LingBot-VA into the standard training, evaluation and processor
|
||
interfaces.
|
||
|
||
## Model Overview
|
||
|
||
LingBot-VA is a **dual-stream "mixture-of-transformers"**: a video/latent stream
|
||
(`patch_embedding_mlp → blocks → proj_out`) and an action stream
|
||
(`action_embedder → blocks → action_proj_out`) share the same 30 transformer blocks and
|
||
text conditioning. Actions are produced by the dedicated `action_proj_out` head — they are
|
||
**not** decoded from predicted pixels, though video and action are co-trained.
|
||
|
||
| Component | Class | Role |
|
||
| ------------------------ | ----------------------- | -------------------------------------------------------------------------------------- |
|
||
| DiT backbone (trainable) | `WanTransformer3DModel` | ~5B-param dual-stream transformer (the only weights stored in the LeRobot checkpoint). |
|
||
| VAE (frozen) | `AutoencoderKLWan` | Wan2.2 VAE, `z_dim=48`. Lazy-pulled from the source repo. |
|
||
| Text encoder (frozen) | `UMT5EncoderModel` | UMT5-XXL, `d_model=4096`. Lazy-pulled from the source repo. |
|
||
|
||
At inference the policy runs an autoregressive loop per chunk: it denoises the video-latent
|
||
stream (CFG, ~20 steps) and the action stream (~50 steps) with two independent
|
||
flow-matching schedulers, maintaining a KV cache across chunks. Real observed keyframes are
|
||
fed back into the KV cache as the chunk is executed (closed-loop world modeling).
|
||
|
||
### What the LeRobot Integration Covers
|
||
|
||
- Standard `policy.type=lingbot_va` configuration through LeRobot.
|
||
- Ready-to-use LeRobot-format checkpoints on the Hub (converted from the released upstream ones).
|
||
- Autoregressive dual-stream inference behind the standard `select_action` interface
|
||
(single-environment eval, `--eval.batch_size=1`).
|
||
- Opt-in saving of the policy's **predicted (imagined) videos** during eval / training.
|
||
- Evaluation with `lerobot-eval` on LIBERO and RoboTwin.
|
||
- Training / fine-tuning via the dual-stream flow-matching loss (`policy.forward`), see below.
|
||
|
||
## Installation
|
||
|
||
1. Install LeRobot by following the [Installation Guide](./installation).
|
||
2. Install the LingBot-VA extra (brings in `diffusers>=0.36` for the Wan2.2 stack):
|
||
|
||
```bash
|
||
pip install -e ".[lingbot_va]"
|
||
# For LIBERO evaluation (Linux only):
|
||
pip install -e ".[lingbot_va,libero]"
|
||
```
|
||
|
||
## Checkpoints
|
||
|
||
The released upstream checkpoints have been converted to LeRobot format and pushed to the Hub:
|
||
|
||
| Variant | LeRobot checkpoint |
|
||
| ---------------------- | ---------------------------------- |
|
||
| LIBERO-Long post-train | `pepijn223/lingbot_va_libero_long` |
|
||
| RoboTwin post-train | `pepijn223/lingbot_va_robotwin` |
|
||
| Pretrained base | `pepijn223/lingbot_va_base` |
|
||
|
||
**Packaging:** only the trainable ~5B transformer is stored in the LeRobot
|
||
`model.safetensors`. The frozen VAE + UMT5 + tokenizer (~20 GB) are **lazily pulled** from
|
||
`config.wan_pretrained_path` at load time (defaults to the source `robbyant/*` repo). The
|
||
UMT5-XXL text encoder runs on CPU by default (`config.text_encoder_device`) so the 5B
|
||
transformer + VAE fit on a single 24–32 GB GPU.
|
||
|
||
## Evaluation (LIBERO)
|
||
|
||
```bash
|
||
lerobot-eval \
|
||
--policy.path=pepijn223/lingbot_va_libero_long \
|
||
--policy.device=cuda \
|
||
--env.type=libero --env.task=libero_10 \
|
||
--env.observation_height=128 --env.observation_width=128 \
|
||
--eval.n_episodes=50 --eval.batch_size=1 \
|
||
--output_dir=outputs/eval/lingbot_va_libero
|
||
```
|
||
|
||
Native LeRobot eval reproduces **96% success on `libero_10` (LIBERO-Long)** (48/50 episodes).
|
||
|
||
LingBot-VA's streaming inference (KV cache + observed-keyframe feedback) is implemented for
|
||
single-environment eval; use `--eval.batch_size=1`.
|
||
|
||
## Evaluation (RoboTwin)
|
||
|
||
RoboTwin 2.0 needs the SAPIEN + CuRobo simulator stack — use the benchmark Docker image
|
||
(`docker/Dockerfile.benchmark.robotwin`, which also needs `warp-lang==1.3.1` and CuRobo built
|
||
with the GPU's compute capability in `TORCH_CUDA_ARCH_LIST`). RoboTwin uses **end-effector-pose
|
||
control**, so run with `--env.action_mode=ee`: the policy predicts per-arm `xyz+quaternion+gripper`
|
||
deltas (`robotwin_tshape` latent layout) that are composed onto the episode's initial eef pose and
|
||
executed via CuRobo IK.
|
||
|
||
```bash
|
||
lerobot-eval \
|
||
--policy.path=pepijn223/lingbot_va_robotwin \
|
||
--policy.device=cuda \
|
||
--env.type=robotwin --env.task=beat_block_hammer --env.action_mode=ee \
|
||
--eval.n_episodes=10 --eval.batch_size=1 \
|
||
--output_dir=outputs/eval/lingbot_va_robotwin
|
||
```
|
||
|
||
### Saving predicted (imagined) videos
|
||
|
||
Set `--policy.save_predicted_video=true` to additionally VAE-decode the predicted video
|
||
latents and write `pred_episode_*.mp4` next to the env-rendered `eval_episode_*.mp4` videos.
|
||
The same flag works for the periodic eval during `lerobot-train`.
|
||
|
||
## Training / fine-tuning
|
||
|
||
`LingBotVAPolicy.forward(batch)` implements the dual-stream **flow-matching** loss
|
||
(`latent_loss + action_loss`, timestep-weighted, action-masked) from the paper: it VAE-encodes
|
||
the camera clips into video latents, UMT5-encodes the task, noises both streams, runs the
|
||
transformer's block-causal training pass and returns `(loss, metrics)`. Optimizer preset is AdamW
|
||
with a linear-warmup-then-constant schedule (matching upstream).
|
||
|
||
Requirements:
|
||
|
||
- The block-causal masks use PyTorch **flex-attention**, so build the policy with
|
||
`--policy.attn_mode=flex` for training (the default `torch` SDPA is inference-only).
|
||
- The full 5B DiT does not fit a single 24–32 GB GPU under AdamW; fine-tune with **LoRA**
|
||
(`--policy.use_peft=true`) and/or optimizer offload. `get_optim_params` returns only the
|
||
trainable (e.g. adapter) parameters; the VAE + UMT5 text encoder stay frozen.
|
||
|
||
```bash
|
||
lerobot-train \
|
||
--policy.path=pepijn223/lingbot_va_libero_long --policy.attn_mode=flex \
|
||
--policy.use_peft=true \
|
||
--dataset.repo_id=<your LeRobot-format dataset> \
|
||
--batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va
|
||
```
|
||
|
||
The dataset must provide camera clips (a temporal window per camera, VAE-encoded to
|
||
`frame_chunk_size` latent frames) and `frame_chunk_size * action_per_frame` action steps per item.
|
||
|
||
## Inference Hyperparameters (LIBERO)
|
||
|
||
| Key | Value |
|
||
| -------------------------------------- | --------------------------------------------------------------------------------- |
|
||
| height × width | 128 × 128 |
|
||
| cameras | `observation.images.image` (agentview), `observation.images.image2` (eye-in-hand) |
|
||
| action channels used | 0–6 (7-DoF arm + gripper) |
|
||
| action_per_frame / frame_chunk_size | 4 / 4 |
|
||
| attn_window | 30 |
|
||
| video / action denoising steps | 20 / 50 |
|
||
| guidance_scale / action_guidance_scale | 5 / 1 |
|
||
| snr_shift / action_snr_shift | 5.0 / 0.05 |
|
||
|
||
These are the defaults of `LingBotVAConfig`; override any of them via `--policy.<name>=...`.
|
||
|
||
## Notes & Limitations
|
||
|
||
- **Correctness gate:** matching the upstream LIBERO success rate requires validating the
|
||
converted checkpoint on a GPU and tensor-diffing intermediate activations against the
|
||
upstream implementation. The most sensitive parts are the action quantile normalization,
|
||
the camera ordering, the `action_per_frame`/`frame_chunk_size` alignment, and `attn_mode`.
|
||
- **Attention backend:** inference uses the `torch` SDPA backend (always available). The
|
||
`flashattn` and `flex` backends are optional; `flex` is only needed for training.
|
||
- **Model size:** the DiT is ~5B params and the frozen VAE+UMT5 add ~20 GB; inference needs
|
||
roughly 18–24 GB of VRAM.
|
||
|
||
## License
|
||
|
||
LingBot-VA is released under Apache-2.0. See the
|
||
[upstream repository](https://github.com/Robbyant/lingbot-va).
|