admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-06 09:37:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(lingbot_va): RoboTwin eef-pose eval, single-file model, Hub checkpoints

Browse Source 
Make the LingBot-VA port runnable on both LIBERO and RoboTwin and clean up the
package to LeRobot conventions.

- Consolidate all vendored Wan2.2 model code (transformer, attention, VAE helpers,
  flow-matching scheduler, grid utils, flex-attention) into a single
  modeling_lingbot_va.py; remove the separate wan_*/schedulers modules.
- Move the fixed action (un)normalization quantiles out of the config and into the
  post-processor (LIBERO 7-DoF + RoboTwin 16-d eef); remove the conversion script in
  favour of ready-to-use LeRobot-format checkpoints on the Hub.
- Fixes found via on-sim validation: undo LIBERO's 180-degree image flip
  (image_hflip), encode obs as a multi-frame streaming-VAE clip, reset the streaming
  VAE cache between episodes, run the transformer in config.dtype, lazy-load frozen
  VAE/UMT5 by subfolder with the text encoder on CPU.
- RoboTwin: add an end-effector-pose action mode to RoboTwinEnv (16-d per-arm
  xyz+quat+gripper deltas composed onto the initial eef pose, executed via CuRobo IK)
  and the robotwin_tshape latent layout (full-res head + half-res wrists via a second
  streaming VAE) with the upstream RoboTwin action quantiles + camera mapping.
- Predicted-video saving works for both benchmarks; docs + tests updated.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
pepijn223
2026-06-06 15:20:51 +02:00
committed by Maxime Ellerbach
parent d600a52943
commit b81909fc28
20 changed files with 2372 additions and 2834 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/lingbot_va.mdx
+34 -17
View File
@@ -28,7 +28,7 @@ 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.
- Checkpoint conversion from the released HuggingFace checkpoints.
- 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.
@@ -48,40 +48,57 @@ pip install -e ".[lingbot_va]"
pip install -e ".[lingbot_va,libero]"
```
## Checkpoint Conversion
## Checkpoints
The released checkpoints are diffusers-style directories
(`robbyant/lingbot-va-base`, `robbyant/lingbot-va-posttrain-robotwin`,
`robbyant/lingbot-va-posttrain-libero-long`). Convert one to LeRobot format with:
The released upstream checkpoints have been converted to LeRobot format and pushed to the Hub:
```bash
python -m lerobot.policies.lingbot_va.convert_lingbot_va_checkpoints \
--checkpoint robbyant/lingbot-va-posttrain-libero-long \
--variant libero \
--output_dir outputs/lingbot_va_libero_long
```
| 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 repo). Pass
`--bundle-frozen` to copy those sub-folders next to the converted checkpoint instead.
Run conversion on a Linux machine with a CUDA GPU and enough RAM/VRAM to materialize the
transformer.
`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=outputs/lingbot_va_libero_long \
--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