fix(advantage): update frame count calculation in constant mode

- Add recipe_path, advantage_prefix, cfg_beta to MolmoAct2Config
- Place advantage clause in the assistant section of _build_robot_text
- Add MolmoAct2NormalizeTaskStep for consistent task normalization
- Parse recipe-rendered advantage in MolmoAct2PackInputsProcessorStep
- Insert RenderMessagesStep pipeline when recipe_path is configured
- Add recap_advantage_molmoact2.yaml recipe

.dockerignore

@@ -22,6 +22,10 @@ outputs
 rl
 media
 
+# Local virtualenvs (the image provides its own)
+.venv
+venv
+
 
 # Logging
 logs

AGENT_GUIDE.md

@@ -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`. See §6/§7 for policy and duration.
+**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.
 
 ```bash
 lerobot-train \

docs/source/_toctree.yml

@@ -69,6 +69,10 @@
     title: VLA-JEPA
   - local: eo1
     title: EO-1
+  - local: lingbot_va
+    title: LingBot-VA
+  - local: fastwam
+    title: FastWAM
   - local: groot
     title: NVIDIA GR00T N1.5
   - local: xvla

docs/source/cheat-sheet.mdx

@@ -150,6 +150,14 @@ 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.

docs/source/fastwam.mdx

@@ -0,0 +1,167 @@
+# 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}
+}
+```

docs/source/hardware_guide.mdx

@@ -82,17 +82,18 @@ VRAM is the first filter. Within a tier, pick by budget and availability — the
 
 ### Hugging Face Jobs
 
-[Hugging Face Jobs](https://huggingface.co/docs/hub/jobs) lets you run training on managed HF infrastructure, billed by the second. The repo publishes a ready-to-use image: **`huggingface/lerobot-gpu:latest`**, rebuilt **every night at 02:00 UTC from `main`** ([`docker_publish.yml`](https://github.com/huggingface/lerobot/blob/main/.github/workflows/docker_publish.yml)) — so it tracks the current state of the repo, not a tagged release.
+[Hugging Face Jobs](https://huggingface.co/docs/hub/jobs) lets you run training on managed HF infrastructure, billed by the second, without owning a GPU. `lerobot-train` submits and streams the job for you — just add `--job.target=<flavor>` to a normal training command:
 
 ```bash
-hf jobs run --flavor a10g-large huggingface/lerobot-gpu:latest \
-  bash -c "nvidia-smi && lerobot-train \
-    --policy.type=act --dataset.repo_id=<USER>/<DATASET> \
-    --policy.repo_id=<USER>/act_<task> --batch_size=8 --steps=50000"
+lerobot-train \
+  --policy.type=act --dataset.repo_id=<USER>/<DATASET> \
+  --policy.repo_id=<USER>/act_<task> \
+  --job.target=a10g-large
 ```
 
 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).
+- Run `hf auth login` once before submitting, the job runs under your token.
+- `--job.target` 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). List the current catalogue with pricing via `hf jobs hardware`, or see [https://huggingface.co/docs/hub/jobs](https://huggingface.co/docs/hub/jobs).
+- The job defaults to a `2d` (48h) timeout. Override it with `--job.timeout=4h` (or any other valid duration string) to shorten or extend the timeout. The job automatically stops when the command completes.
+- For the full walkthrough — dataset upload, checkpoint streaming, resuming a run on a job — see the [imitation-learning training guide](./il_robots#train-using-hugging-face-jobs).

docs/source/il_robots.mdx

@@ -126,7 +126,7 @@ import time
 from lerobot.teleoperators.so_leader import SO101Leader, SO101LeaderConfig
 from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig
 from lerobot.cameras.opencv import OpenCVCameraConfig
-from lerobot.utils.visualization_utils import init_rerun, log_rerun_data, shutdown_rerun
+from lerobot.utils.visualization_utils import init_visualization, log_visualization_data, shutdown_visualization
 
 robot_config = SO101FollowerConfig(
     port="/dev/tty.usbmodem5AB90687491",
@@ -142,7 +142,7 @@ teleop_config = SO101LeaderConfig(
     id="my_leader_arm",
 )
 
-init_rerun(session_name="teleoperation")
+init_visualization("rerun", session_name="teleoperation")  # pass "foxglove" to stream to Foxglove instead
 
 robot = SO101Follower(robot_config)
 teleop_device = SO101Leader(teleop_config)
@@ -158,7 +158,7 @@ while True:
     observation = robot.get_observation()
     action = teleop_device.get_action()
     robot.send_action(action)
-    log_rerun_data(observation=observation, action=action)
+    log_visualization_data("rerun", observation=observation, action=action)
 
     elapsed_time = time.perf_counter() - start_time
     sleep_time = TIME_PER_FRAME - elapsed_time
@@ -223,7 +223,7 @@ from lerobot.teleoperators.so_leader.config_so_leader import SO101LeaderConfig
 from lerobot.teleoperators.so_leader.so_leader import SO101Leader
 from lerobot.common.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import init_visualization
 from lerobot.scripts.lerobot_record import record_loop
 from lerobot.processor import make_default_processors
 
@@ -270,7 +270,7 @@ def main():
 
     # Initialize the keyboard listener and rerun visualization
     _, events = init_keyboard_listener()
-    init_rerun(session_name="recording")
+    init_visualization("rerun", session_name="recording")
 
     # Connect the robot and teleoperator
     robot.connect()
@@ -514,6 +514,12 @@ 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`
@@ -526,78 +532,48 @@ 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).
 
-To run the training use this command:
+`lerobot-train` runs locally by default. To run on a HuggingFace GPU, pass `--job.target` with a hardware flavor name:
 
-<hfoptions id="train_with_hf_jobs">
-<hfoption id="Command">
 ```bash
-hf jobs run \
-  --flavor a10g-small \
-  --timeout 4h \
-  --secrets HF_TOKEN \
-  huggingface/lerobot-gpu:latest \
-  -- \
-  python -m lerobot.scripts.lerobot_train \
-    --dataset.repo_id=username/dataset \
-    --policy.type=act \
-    --steps=5000 \
-    --batch_size=16 \
-    --policy.device=cuda \
-    --policy.repo_id=username/your_policy \
-    --log_freq=100
+lerobot-train \
+  --dataset.repo_id=${HF_USER}/so101_test \
+  --policy.type=act \
+  --policy.repo_id=${HF_USER}/my_policy \
+  --job.target=a10g-small
 ```
-</hfoption>
-<hfoption id="API example">
 
-<!-- prettier-ignore-start -->
-```python
-from huggingface_hub import run_job, get_token
+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:
 
-run_name = "act_so101_hf_jobs"
-dataset_id = "username/dataset"
-user_hub_id = "username"
-
-command_args = [
-    "python", "-m", "lerobot.scripts.lerobot_train",
-    "--dataset.repo_id", dataset_id,
-    "--policy.type", "act",
-    "--steps", "5000",
-    "--batch_size", "16",
-    "--num_workers", "4",
-    "--policy.device", "cuda",
-    "--log_freq", "100",
-    "--save_freq", "1000",
-    "--save_checkpoint", "true",
-    "--wandb.enable", "false",
-    "--policy.repo_id", f"{user_hub_id}/{run_name}"
-]
-
-print(f"Submitting job '{run_name}' to Hugging Face Infrastructure...")
-
-job_info = run_job(
-    image="huggingface/lerobot-gpu:latest",
-    command=command_args,
-    flavor="a10g-small",
-    timeout="4h",
-    secrets={"HF_TOKEN": get_token()}
-)
-
-print("\n🚀 Job successfully launched!")
-print(f"🔹 Job ID: {job_info.id}")
-print(f"🔗 Live UI Dashboard & Logs: {job_info.url}")
+```bash
+hf jobs logs <job-id>
+hf jobs cancel <job-id>
 ```
-<!-- prettier-ignore-end -->
 
-</hfoption>
-</hfoptions>
+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.
 
-You can modify the `--flavor` to use different hardware, for example: `t4-small`, `a100-large`, `h200`. Use `hf jobs hardware` to see the full list with pricing.
-Depending on the model you want to train and the hardware you selected you can also modify the `--batch_size` and `--number_of_workers`.
-For longer training sessions increase the timeout.
+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"]'`.
 
-Once the training is started you can go to [Jobs](https://huggingface.co/settings/jobs) and see if your jobs is running as well as all the outputs. Sometimes it takes a few minutes to schedule your job so be patient.
+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.
 
-After training the model will be pushed to hub and you can use it as any other model with LeRobot.
+> **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
 
@@ -620,6 +596,8 @@ 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

docs/source/lingbot_va.mdx

@@ -0,0 +1,187 @@
+# 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.
+
+| Component                | Class                   | Role                                                        |
+| ------------------------ | ----------------------- | ----------------------------------------------------------- |
+| DiT backbone (trainable) | `WanTransformer3DModel` | ~5B-param dual-stream transformer.                          |
+| 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:
+
+```bash
+pip install -e ".[lingbot_va]"
+```
+
+## Checkpoints
+
+The released upstream checkpoints have been converted to LeRobot format and pushed to the Hub:
+
+| Variant                | LeRobot checkpoint               |
+| ---------------------- | -------------------------------- |
+| LIBERO-Long post-train | `lerobot/lingbot_va_libero_long` |
+| RoboTwin post-train    | `lerobot/lingbot_va_robotwin`    |
+| Pretrained base        | `lerobot/lingbot_va_base`        |
+
+Only the trainable ~5B transformer is stored in the LeRobot
+`model.safetensors`. The frozen VAE + UMT5 + tokenizer (~20 GB) are 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=lerobot/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
+```
+
+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. You can 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=lerobot/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=lerobot/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.
+
+## Data format (action channels & camera order)
+
+LingBot-VA is an **end-effector (Cartesian) pose** policy, it predicts EEF poses + gripper, not
+joint positions. Actions live in a fixed multi-embodiment **30-dim** layout; map your robot's
+action dimensions into these channels and pad the rest with `0` (`used_action_channel_ids` selects
+the channels a given checkpoint actually uses):
+
+| channels | meaning                                               |
+| -------- | ----------------------------------------------------- |
+| 0–6      | Left-arm end-effector pose                            |
+| 7–13     | Right-arm end-effector pose                           |
+| 14–20    | Left-arm joints (unused by the released checkpoints)  |
+| 21–27    | Right-arm joints (unused by the released checkpoints) |
+| 28       | Left gripper                                          |
+| 29       | Right gripper                                         |
+
+- **LIBERO** uses channels `0–6`: a 6-DoF EEF delta (xyz + rotation) + gripper (single arm).
+- **RoboTwin** uses channels `[0–6, 28, 7–13, 29]`: left EEF (xyz + quaternion) + left gripper +
+  right EEF + right gripper (16 dims). The env converts these poses to joint trajectories via
+  CuRobo IK — joints are never predicted.
+
+Joint-space datasets (or a different EEF convention) must be remapped into this schema before
+fine-tuning these checkpoints.
+
+**Camera order is fixed and order-sensitive**, per-camera latents are concatenated spatially in
+`obs_cam_keys` order, so the physical camera→slot mapping must match training:
+
+| benchmark | `obs_cam_keys` (in order)                                                                             | `camera_layout`                                                     |
+| --------- | ----------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
+| LIBERO    | `observation.images.image` (agentview / 3rd-person), `observation.images.image2` (eye-in-hand wrist)  | `width_concat` (latents concatenated on width)                      |
+| RoboTwin  | `observation.images.head_camera`, `observation.images.left_camera`, `observation.images.right_camera` | `robotwin_tshape` (full-res head below, two half-res wrists on top) |
+
+The first camera is the exterior/head view and the rest are wrist views.
+
+## 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
+
+- **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).

docs/source/molmoact2.mdx

@@ -386,6 +386,68 @@ 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

docs/source/policy_fastwam_README.md

@@ -0,0 +1,56 @@
+## 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
+```

docs/source/using_dataset_tools.mdx

@@ -265,6 +265,8 @@ lerobot-dataset-viz \
 
 Once executed, the tool opens `rerun.io` and displays the camera streams, robot states, and actions for the selected episode.
 
+To use [Foxglove](https://foxglove.dev) instead of Rerun, install the extra add `--display-mode foxglove`. This starts a WebSocket server (connect the Foxglove app to `ws://127.0.0.1:8765`) that serves the episode as a seekable timeline you can play/pause and scrub.
+
 For advanced usage—including visualizing datasets stored on a remote server—run:
 
 ```bash

pyproject.toml

@@ -124,7 +124,8 @@ hardware = [
     "lerobot[deepdiff-dep]",
 ]
 viz = [
-    "rerun-sdk>=0.24.0,<0.27.0",
+    "rerun-sdk>=0.24.0,<0.34.0",
+    "foxglove-sdk>=0.25.1,<0.26.0",
 ]
 # ── User-facing composite extras (map to CLI scripts) ─────
 # lerobot-record, lerobot-replay, lerobot-calibrate, lerobot-teleoperate, etc.
@@ -227,10 +228,16 @@ groot = [
 sarm = ["lerobot[transformers-dep]", "pydantic>=2.0.0,<3.0.0", "faker>=33.0.0,<35.0.0", "lerobot[matplotlib-dep]", "lerobot[qwen-vl-utils-dep]"]
 robometer = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]", "lerobot[peft-dep]"]
 topreward = ["lerobot[transformers-dep]"]
+recap = ["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]"]
+lingbot_va = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[accelerate-dep]"]
 
 # Features
 async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"]
@@ -308,10 +315,12 @@ 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]",
     "lerobot[vla_jepa]",
+    "lerobot[lingbot_va]",
     "lerobot[async]",
     "lerobot[dev]",
     "lerobot[test]",
@@ -324,6 +333,7 @@ all = [
     "lerobot[sarm]",
     "lerobot[robometer]",
     "lerobot[topreward]",
+    "lerobot[recap]",
     "lerobot[peft]",
     # "lerobot[unitree_g1]", TODO: Unitree requires specific installation instructions for unitree_sdk2
 ]
@@ -347,6 +357,7 @@ lerobot-edit-dataset="lerobot.scripts.lerobot_edit_dataset:main"
 lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main"
 lerobot-annotate="lerobot.scripts.lerobot_annotate:main"
 lerobot-rollout="lerobot.scripts.lerobot_rollout:main"
+lerobot-compute-returns="lerobot.scripts.lerobot_compute_returns:main"
 
 # ---------------- Tool Configurations ----------------
 
@@ -444,7 +455,8 @@ default.extend-ignore-identifiers-re = [
     "is_compileable",
     "ROBOTIS",
     "OT_VALUE",
-    "VanderBilt"
+    "VanderBilt",
+    "seperated_timestep",
 ]
 
 # TODO: Uncomment when ready to use

src/lerobot/annotations/steerable_pipeline/config.py

@@ -169,6 +169,51 @@ class ExecutorConfig:
     episode_parallelism: int = 16
 
 
+@dataclass
+class AdvantageConfig:
+    """``advantage`` module: RECAP advantage scoring via frozen value function."""
+
+    enabled: bool = True
+
+    # Constant advantage label for all frames (e.g. "positive" for SFT iteration 0).
+    # Skips VF inference, dropout still applies for CFG.
+    constant_value: str | None = None
+
+    # Trained value function checkpoint (local path or Hub repo ID).
+    # Ignored when constant_value is set.
+    value_function_path: str = ""
+
+    # Device to run the value function on.
+    device: str = "cuda"
+
+    # N-step lookahead for advantage estimation.
+    # None = MC (N=T): A_t = R_t - V(s_t), using mc_return from dataset.
+    # 50 = fine-tuning mode: A_t = Σ r_{t:t+N} + V(s_{t+N}) - V(s_t).
+    n_step: int | None = None
+
+    # Per-task percentile for binarization threshold ε_ℓ.
+    # Actions with advantage > ε_ℓ get I_t = True (positive).
+    threshold_percentile: float = 0.3
+
+    # Fraction of frames to randomly omit advantage labels (enables CFG).
+    dropout_rate: float = 0.3
+
+    # Force I_t = True for frames marked as human interventions.
+    force_positive_on_intervention: bool = True
+
+    # Column name in dataset for intervention flag.
+    intervention_key: str = "intervention"
+
+    # Column name for pre-computed MC returns (from lerobot-compute-returns).
+    mc_return_key: str = "mc_return"
+
+    # Batch size for value function inference.
+    batch_size: int = 32
+
+    # Random seed for dropout reproducibility.
+    seed: int = 1729
+
+
 @dataclass
 class AnnotationPipelineConfig:
     """Top-level config for ``lerobot-annotate`` (rewrites data shards in place)."""
@@ -190,6 +235,7 @@ class AnnotationPipelineConfig:
     plan: PlanConfig = field(default_factory=PlanConfig)
     interjections: InterjectionsConfig = field(default_factory=InterjectionsConfig)
     vqa: VqaConfig = field(default_factory=VqaConfig)
+    advantage: AdvantageConfig = field(default_factory=AdvantageConfig)
 
     vlm: VlmConfig = field(default_factory=VlmConfig)
     executor: ExecutorConfig = field(default_factory=ExecutorConfig)

src/lerobot/annotations/steerable_pipeline/executor.py

@@ -15,20 +15,24 @@
 # limitations under the License.
 """In-process executor that runs the annotation phases.
 
-The executor runs **six phases** in dependency order:
+The executor runs **seven phases** in dependency order:
 
     phase 1: ``plan`` module (plan + subtasks + memory)
     phase 2: ``interjections`` module (interjections + speech)
     phase 3: ``plan`` plan-update pass — re-runs plan emission at every
              interjection timestamp produced by phase 2
     phase 4: ``vqa`` module (VQA)
-    phase 5: validator
-    phase 6: writer
+    phase 5: ``advantage`` module (advantage scoring via frozen VF)
+    phase 6: validator
+    phase 7: writer
 
 Phase 3 is why the ``plan`` module must be re-entered after the
 ``interjections`` module — to refresh ``plan`` rows at interjection
 timestamps.
 
+Phase 5 (advantage) does not depend on the VLM modules, it uses a frozen
+distributional value function to compute per-frame advantage indicators.
+
 Distributed execution is provided by Hugging Face Jobs (see
 ``examples/annotations/run_hf_job.py``); the runner inside the job
 invokes ``lerobot-annotate`` which uses this in-process executor.
@@ -74,7 +78,7 @@ class PipelineRunSummary:
 
 @dataclass
 class Executor:
-    """Run all six phases over a dataset root in-process.
+    """Run all seven phases over a dataset root in-process.
 
     Episode-level concurrency comes from ``ExecutorConfig.episode_parallelism``
     (a thread pool); cluster-level concurrency comes from running this
@@ -86,6 +90,7 @@ class Executor:
     plan: Any  # PlanSubtasksMemoryModule
     interjections: Any  # InterjectionsAndSpeechModule
     vqa: Any  # GeneralVqaModule
+    advantage: Any  # AdvantageModule
     writer: LanguageColumnsWriter
     validator: StagingValidator
 
@@ -112,6 +117,8 @@ class Executor:
         phases.append(self._run_plan_update_phase(records, staging_dir))
         # Phase 4: ``vqa`` module (VQA)
         phases.append(self._run_module_phase("vqa", records, staging_dir, self.vqa))
+        # Phase 5: ``advantage`` module (advantage scoring via frozen VF)
+        phases.append(self._run_module_phase("advantage", records, staging_dir, self.advantage))
 
         print("[annotate] running validator...", flush=True)
         report = self.validator.validate(records, staging_dir)
@@ -179,7 +186,7 @@ class Executor:
         staging_dir: Path,
         module: Any,
     ) -> PhaseResult:
-        if not module.enabled:
+        if module is None or not module.enabled:
             print(f"[annotate] phase={name} skipped (module disabled)", flush=True)
             return PhaseResult(name=name, episodes_processed=0, episodes_skipped=len(records))
         n = len(records)
@@ -231,7 +238,7 @@ class Executor:
         ``plan`` module with the interjection timestamps so its existing
         prompt path is reused.
         """
-        if not self.plan.enabled or not self.interjections.enabled:
+        if not self.plan or not self.plan.enabled or not self.interjections or not self.interjections.enabled:
             return PhaseResult(name="plan_update", episodes_processed=0, episodes_skipped=len(records))
         processed = 0
         for record in records:

src/lerobot/annotations/steerable_pipeline/modules/__init__.py

@@ -14,11 +14,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .advantage import AdvantageModule
 from .general_vqa import GeneralVqaModule
 from .interjections_and_speech import InterjectionsAndSpeechModule
 from .plan_subtasks_memory import PlanSubtasksMemoryModule
 
 __all__ = [
+    "AdvantageModule",
     "GeneralVqaModule",
     "InterjectionsAndSpeechModule",
     "PlanSubtasksMemoryModule",

src/lerobot/annotations/steerable_pipeline/modules/advantage.py

@@ -0,0 +1,298 @@
+#!/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.
+
+"""Advantage scoring module for RECAP.
+
+Computes per-frame advantage values using a frozen distributional value function,
+binarizes them into improvement indicators (I_t), and emits ``style="advantage"``
+persistent rows for policy conditioning.
+
+Paper reference: pi*0.6, Section IV-B and Appendix F.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+import torch
+
+from ..config import AdvantageConfig
+from ..reader import EpisodeRecord
+from ..staging import EpisodeStaging
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AdvantageModule:
+    """Compute advantage indicators and emit persistent annotation rows.
+
+    The module loads a frozen distributional value function and scores each
+    frame in an episode. Advantages are binarized into ``positive``/``negative``
+    indicators using a per-task threshold, then written as ``style="advantage"``
+    persistent rows into the staging area.
+
+    Requires ``mc_return`` column in the dataset (from lerobot-compute-returns).
+    """
+
+    config: AdvantageConfig
+    _model: Any = field(default=None, init=False, repr=False)
+    _preprocessor: Any = field(default=None, init=False, repr=False)
+    _threshold: float | None = field(default=None, init=False, repr=False)
+
+    @property
+    def enabled(self) -> bool:
+        return self.config.enabled
+
+    def _ensure_model_loaded(self) -> None:
+        """Lazy-load the frozen value function on first use."""
+        if self._model is not None:
+            return
+
+        from lerobot.rewards import (
+            make_reward_model,
+            make_reward_model_config,
+            make_reward_pre_post_processors,
+        )
+
+        cfg = make_reward_model_config(
+            "distributional_value_function",
+            pretrained_path=self.config.value_function_path,
+            device=self.config.device,
+        )
+        self._model = make_reward_model(cfg)
+        self._model.eval()
+        for p in self._model.parameters():
+            p.requires_grad_(False)
+
+        self._preprocessor, _ = make_reward_pre_post_processors(cfg)
+        logger.info("Loaded frozen VF from %s on %s", self.config.value_function_path, self.config.device)
+
+    def compute_advantages_for_episode(self, record: EpisodeRecord) -> tuple[np.ndarray, np.ndarray]:
+        """Compute raw advantage values for all frames in an episode.
+
+        Returns:
+            (advantages, intervention_mask) both shape [num_frames].
+            advantages[t] = A_t, intervention_mask[t] = True if frame is intervention.
+        """
+        self._ensure_model_loaded()
+
+        df = record.frames_df()
+        num_frames = len(df)
+
+        mc_return_key = self.config.mc_return_key
+        if mc_return_key not in df.columns:
+            raise KeyError(
+                f"Column '{mc_return_key}' not found in episode {record.episode_index}. "
+                "Run lerobot-compute-returns first."
+            )
+
+        mc_returns = df[mc_return_key].values.astype(np.float32)
+
+        intervention_mask = np.zeros(num_frames, dtype=bool)
+        if self.config.intervention_key in df.columns:
+            intervention_mask = df[self.config.intervention_key].values.astype(bool)
+
+        # Skip VF inference on intervention frames — they're always "positive"
+        # regardless of advantage value, so V(s_t) is never used for them.
+        skip_mask = intervention_mask if self.config.force_positive_on_intervention else None
+        values = self._compute_values(record, skip_mask=skip_mask)
+
+        if self.config.n_step is None:
+            advantages = mc_returns - values
+        else:
+            advantages = self._compute_n_step_advantages(mc_returns, values, record, n=self.config.n_step)
+
+        return advantages, intervention_mask
+
+    def _compute_values(self, record: EpisodeRecord, skip_mask: np.ndarray | None = None) -> np.ndarray:
+        """Run frozen VF over all frames to get V(s_t) predictions.
+
+        Args:
+            record: Episode data.
+            skip_mask: Optional boolean mask [num_frames]. Frames where True are
+                skipped (left as 0.0) to avoid unnecessary inference.
+        """
+        df = record.frames_df()
+        num_frames = len(df)
+        values = np.zeros(num_frames, dtype=np.float32)
+
+        image_key = self._resolve_image_key(df)
+        if image_key is None:
+            logger.warning("No image key found for episode %d; returning zero values.", record.episode_index)
+            return values
+
+        # Determine which frame indices actually need inference
+        infer_indices = np.where(~skip_mask)[0] if skip_mask is not None else np.arange(num_frames)
+
+        if len(infer_indices) == 0:
+            return values
+
+        task_text = record.episode_task
+
+        for batch_start in range(0, len(infer_indices), self.config.batch_size):
+            batch_end = min(batch_start + self.config.batch_size, len(infer_indices))
+            batch_indices = infer_indices[batch_start:batch_end]
+            batch_images = []
+
+            for idx in batch_indices:
+                img_val = df.iloc[idx][image_key]
+                if isinstance(img_val, np.ndarray):
+                    img_tensor = torch.from_numpy(img_val).float()
+                elif isinstance(img_val, torch.Tensor):
+                    img_tensor = img_val.float()
+                else:
+                    img_tensor = torch.zeros(3, 224, 224)
+                batch_images.append(img_tensor)
+
+            batch_images_tensor = torch.stack(batch_images)
+            batch_size = batch_images_tensor.shape[0]
+
+            raw_batch = {
+                image_key: batch_images_tensor,
+                "task": [task_text] * batch_size,
+            }
+
+            processed = self._preprocessor(raw_batch)
+
+            with torch.no_grad():
+                v_values = self._model.compute_reward(processed)
+
+            values[batch_indices] = v_values.cpu().numpy()
+
+        return values
+
+    def _compute_n_step_advantages(
+        self, mc_returns: np.ndarray, values: np.ndarray, record: EpisodeRecord, n: int
+    ) -> np.ndarray:
+        """Compute N-step advantage: A_t = Σ r_{t:t+N-1} + V(s_{t+N}) - V(s_t).
+
+        When t+N exceeds episode length, truncates to MC (uses mc_return directly).
+        """
+        num_frames = len(values)
+        advantages = np.zeros(num_frames, dtype=np.float32)
+
+        for t in range(num_frames):
+            if t + n >= num_frames:
+                advantages[t] = mc_returns[t] - values[t]
+            else:
+                n_step_return = mc_returns[t] - mc_returns[t + n]
+                advantages[t] = n_step_return + values[t + n] - values[t]
+
+        return advantages
+
+    def _resolve_image_key(self, df) -> str | None:
+        """Find the first image observation key in the dataframe columns."""
+        for col in df.columns:
+            if col.startswith("observation.images."):
+                return col
+        return None
+
+    def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
+        """Score one episode and write advantage rows to staging."""
+        if self.config.constant_value:
+            self._run_constant_mode(record, staging)
+            return
+
+        if not self.config.value_function_path:
+            logger.warning("No value_function_path or constant_value configured; skipping advantage scoring.")
+            return
+
+        advantages, intervention_mask = self.compute_advantages_for_episode(record)
+        num_frames = len(advantages)
+
+        threshold = self._compute_threshold(advantages, intervention_mask)
+
+        rng = np.random.default_rng(seed=self.config.seed + record.episode_index)
+
+        rows: list[dict[str, Any]] = []
+        for t in range(num_frames):
+            if rng.random() < self.config.dropout_rate:
+                continue
+
+            if (
+                self.config.force_positive_on_intervention
+                and intervention_mask[t]
+                or advantages[t] > threshold
+            ):
+                indicator = "positive"
+            else:
+                indicator = "negative"
+
+            timestamp = float(record.frame_timestamps[t]) if t < len(record.frame_timestamps) else 0.0
+
+            rows.append(
+                {
+                    "role": "user",
+                    "content": indicator,
+                    "style": "advantage",
+                    "timestamp": timestamp,
+                    "camera": None,
+                    "tool_calls": None,
+                }
+            )
+
+        staging.write("advantage", rows)
+        logger.debug(
+            "Episode %d: %d/%d frames scored (threshold=%.4f, %d positive, %d negative)",
+            record.episode_index,
+            len(rows),
+            num_frames,
+            threshold,
+            sum(1 for r in rows if r["content"] == "positive"),
+            sum(1 for r in rows if r["content"] == "negative"),
+        )
+
+    def _run_constant_mode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
+        """Emit a fixed advantage value for every frame (with dropout for CFG)."""
+        num_frames = len(record.frame_timestamps)
+        rng = np.random.default_rng(seed=self.config.seed + record.episode_index)
+
+        rows: list[dict[str, Any]] = []
+        for t in range(num_frames):
+            if rng.random() < self.config.dropout_rate:
+                continue
+
+            rows.append(
+                {
+                    "role": "user",
+                    "content": self.config.constant_value,
+                    "style": "advantage",
+                    "timestamp": float(record.frame_timestamps[t]),
+                    "camera": None,
+                    "tool_calls": None,
+                }
+            )
+
+        staging.write("advantage", rows)
+        logger.debug(
+            "Episode %d: %d/%d frames labeled constant '%s' (dropout=%.2f)",
+            record.episode_index,
+            len(rows),
+            num_frames,
+            self.config.constant_value,
+            self.config.dropout_rate,
+        )
+
+    def _compute_threshold(self, advantages: np.ndarray, intervention_mask: np.ndarray) -> float:
+        """Compute the binarization threshold as the configured percentile of advantages."""
+        non_intervention = advantages[~intervention_mask] if intervention_mask.any() else advantages
+        if len(non_intervention) == 0:
+            return 0.0
+        return float(np.percentile(non_intervention, self.config.threshold_percentile * 100))

src/lerobot/annotations/steerable_pipeline/staging.py

@@ -39,6 +39,7 @@ _MODULES: tuple[ModuleName, ...] = (
     "plan",
     "interjections",
     "vqa",
+    "advantage",
 )
 
 

src/lerobot/common/train_utils.py

@@ -15,6 +15,7 @@
 # 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
 
@@ -35,6 +36,7 @@ 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
 
@@ -283,3 +285,61 @@ 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

src/lerobot/configs/__init__.py

@@ -22,7 +22,7 @@ Import them directly: ``from lerobot.configs.train import TrainPipelineConfig``
 """
 
 from .dataset import DatasetRecordConfig
-from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
+from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
 from .recipe import MessageTurn, TrainingRecipe, load_recipe
 from .types import (
@@ -34,6 +34,8 @@ from .types import (
 )
 from .video import (
     DEFAULT_DEPTH_UNIT,
+    DEPTH_METER_UNIT,
+    DEPTH_MILLIMETER_UNIT,
     VALID_VIDEO_CODECS,
     VIDEO_ENCODER_INFO_KEYS,
     DepthEncoderConfig,
@@ -41,6 +43,7 @@ from .video import (
     VideoEncoderConfig,
     depth_encoder_defaults,
     encoder_config_from_video_info,
+    infer_depth_unit,
     rgb_encoder_defaults,
 )
 
@@ -55,6 +58,7 @@ __all__ = [
     "DatasetRecordConfig",
     "DatasetConfig",
     "EvalConfig",
+    "JobConfig",
     "MessageTurn",
     "PeftConfig",
     "PreTrainedConfig",
@@ -69,8 +73,11 @@ __all__ = [
     "depth_encoder_defaults",
     # Factories
     "encoder_config_from_video_info",
+    "infer_depth_unit",
     # Constants
     "DEFAULT_DEPTH_UNIT",
+    "DEPTH_METER_UNIT",
+    "DEPTH_MILLIMETER_UNIT",
     "VALID_VIDEO_CODECS",
     "VIDEO_ENCODER_INFO_KEYS",
 ]

src/lerobot/configs/default.py

@@ -145,3 +145,35 @@ 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)

src/lerobot/configs/recipe.py

@@ -32,6 +32,7 @@ DEFAULT_BINDINGS = {
     "interjection": "emitted_at(t, style=interjection)",
     "vqa": "emitted_at(t, style=vqa, role=assistant)",
     "vqa_query": "emitted_at(t, style=vqa, role=user)",
+    "advantage": "active_at(t, style=advantage)",
 }
 
 PLACEHOLDER_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")

src/lerobot/configs/recipes/recap_advantage.yaml

@@ -0,0 +1,30 @@
+# RECAP advantage-conditioned recipe.
+#
+# Composes task + advantage indicator into the prompt for conditional SFT.
+# The advantage binding resolves to "positive" or "negative" from the
+# language_persistent column (written by lerobot-annotate --advantage).
+# When advantage is absent (30% dropout), the advantage turn is skipped
+# entirely via if_present, training the unconditional branch for CFG.
+#
+# This recipe is policy-agnostic: any VLA that consumes chat-style messages
+# can use it. Override bindings or add blend components for task-specific needs.
+#
+# Paper: pi*0.6, Section IV-B (conditional policy training with I_t).
+
+bindings:
+  advantage: "active_at(t, style=advantage)"
+
+messages:
+  - role: user
+    content: "${task}"
+    stream: high_level
+
+  - role: user
+    content: "Advantage: ${advantage}"
+    stream: high_level
+    if_present: advantage
+
+  - role: assistant
+    content: "${subtask}"
+    stream: low_level
+    target: true

src/lerobot/configs/recipes/recap_advantage_blend.yaml

@@ -0,0 +1,41 @@
+# RECAP full recipe with advantage conditioning and subtask blending.
+#
+# Blend of two training modes:
+# 1. advantage_conditioned (70%): Task + advantage indicator → action
+# 2. unconditional (30%): Task only → action (no advantage, trains CFG baseline)
+#
+# This achieves the same effect as per-frame dropout in the annotation module
+# but at the recipe level, giving explicit control over the conditioning ratio.
+# Use this instead of annotation-level dropout if you want a fixed split.
+#
+# Paper: pi*0.6, Appendix E (classifier-free guidance requires both branches).
+
+blend:
+  advantage_conditioned:
+    weight: 0.7
+    messages:
+      - role: user
+        content: "${task}\nAdvantage: ${advantage}"
+        stream: high_level
+        if_present: advantage
+
+      - role: user
+        content: "${task}"
+        stream: high_level
+
+      - role: assistant
+        content: "${subtask}"
+        stream: low_level
+        target: true
+
+  unconditional:
+    weight: 0.3
+    messages:
+      - role: user
+        content: "${task}"
+        stream: high_level
+
+      - role: assistant
+        content: "${subtask}"
+        stream: low_level
+        target: true

src/lerobot/configs/recipes/recap_advantage_molmoact2.yaml

@@ -0,0 +1,28 @@
+# RECAP advantage recipe for MolmoAct2.
+#
+# Renders task + advantage into the task field as "<task> Advantage: <value>".
+# MolmoAct2PackInputsProcessorStep parses this, extracts the advantage value,
+# and places it AFTER the full user prompt but BEFORE action tokens — matching
+# the RECAP paper (Section V-B): "The advantage indicator appears in the training
+# sequence after ˆℓ but before the actions, such that only the action
+# log-likelihoods are affected."
+#
+# Final prompt layout:
+#   <images><|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\nAdvantage: positive. <action_output>...
+#
+# When advantage is absent (CFG dropout), if_present guard skips this message
+# and RenderedMessagesToTaskStep leaves the task unchanged — no advantage clause.
+
+bindings:
+  advantage: "active_at(t, style=advantage)"
+
+messages:
+  - role: user
+    content: "${task} Advantage: ${advantage}"
+    stream: high_level
+    if_present: advantage
+
+  - role: assistant
+    content: ""
+    stream: low_level
+    target: true

src/lerobot/configs/train.py

@@ -26,11 +26,12 @@ from huggingface_hub.errors import HfHubHTTPError
 
 from lerobot import envs
 from lerobot.optim import LRSchedulerConfig, OptimizerConfig
-from lerobot.utils.hub import HubMixin
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR
+from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint
 from lerobot.utils.sample_weighting import SampleWeightingConfig
 
 from . import parser
-from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
+from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
 from .rewards import RewardModelConfig
 
@@ -83,10 +84,11 @@ 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. 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.
+    # 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).
     resume: bool = False
     # `seed` is used for training (eg: model initialization, dataset shuffling)
     # AND for the evaluation environments.
@@ -118,6 +120,13 @@ 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
 
@@ -137,10 +146,17 @@ class TrainPipelineConfig(HubMixin):
             return self.reward_model  # type: ignore[return-value]
         return self.policy  # type: ignore[return-value]
 
-    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")
+    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")
+        policy_path = parser.get_path_arg("policy")
 
         if reward_model_path:
             cli_overrides = parser.get_cli_overrides("reward_model")
@@ -149,31 +165,54 @@ class TrainPipelineConfig(HubMixin):
             )
             self.reward_model.pretrained_path = str(Path(reward_model_path))
         elif policy_path:
-            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
-            )
+            overrides = parser.get_yaml_overrides("policy") + (parser.get_cli_overrides("policy") or [])
+            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=overrides)
             self.policy.pretrained_path = Path(policy_path)
         elif self.resume:
-            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}"
-                )
+            self._resolve_resume_checkpoint()
 
-            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."
-                )
+    def _resolve_resume_checkpoint(self) -> None:
+        """Point the trainable config at the checkpoint named by `--config_path`.
 
+        `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(
@@ -216,9 +255,19 @@ 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.")
 
-        if hasattr(active_cfg, "push_to_hub") and active_cfg.push_to_hub and not active_cfg.repo_id:
+        # 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
+        ):
             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."""
@@ -255,22 +304,30 @@ 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(
-                    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,
-                )
+                config_file = hf_hub_download(filename=TRAIN_CONFIG_NAME, **dl_kwargs)
             except HfHubHTTPError as e:
-                raise FileNotFoundError(
-                    f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
-                ) from 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
+                )
 
         cli_args = kwargs.pop("cli_args", [])
         # Legacy RA-BC migration only applies to framework-saved checkpoints (always JSON).

src/lerobot/configs/video.py

@@ -22,6 +22,8 @@ import logging
 from dataclasses import dataclass, field
 from typing import Any, ClassVar, Self
 
+import numpy as np
+
 from lerobot.utils.import_utils import require_package
 
 logger = logging.getLogger(__name__)
@@ -36,7 +38,9 @@ HW_VIDEO_CODECS = [
     "h264_vaapi",  # Linux Intel/AMD
     "h264_qsv",  # Intel Quick Sync
 ]
-VALID_VIDEO_CODECS: frozenset[str] = frozenset({"h264", "hevc", "libsvtav1", "auto", *HW_VIDEO_CODECS})
+VALID_VIDEO_CODECS: frozenset[str] = frozenset(
+    {"h264", "hevc", "libsvtav1", "libaom-av1", "auto", *HW_VIDEO_CODECS}
+)
 # Aliases for legacy video codec names.
 VIDEO_CODECS_ALIASES: dict[str, str] = {"av1": "libsvtav1"}
 
@@ -65,6 +69,15 @@ DEPTH_METER_UNIT: str = "m"
 DEPTH_MILLIMETER_UNIT: str = "mm"
 DEFAULT_DEPTH_UNIT: str = DEPTH_MILLIMETER_UNIT
 
+
+def infer_depth_unit(dtype: np.dtype | type) -> str:
+    """Infer the physical unit of raw depth frames from their dtype.
+
+    Floating-point frames are assumed to be in metres, integer frames in millimetres.
+    """
+    return DEPTH_METER_UNIT if np.issubdtype(np.dtype(dtype), np.floating) else 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"})
 
@@ -213,18 +226,24 @@ class VideoEncoderConfig:
             if encoder_threads is not None:
                 svtav1_parts.append(f"lp={encoder_threads}")
             if svtav1_parts:
-                opts["svtav1-params"] = ":".join(svtav1_parts)
+                set_if("svtav1-params", ":".join(svtav1_parts))
         elif self.vcodec in ("h264", "hevc"):
             set_if("crf", self.crf)
             set_if("preset", self.preset)
             if self.fast_decode:
-                opts["tune"] = "fastdecode"
+                set_if("tune", "fastdecode")
             set_if("threads", encoder_threads)
+        elif self.vcodec == "libaom-av1":
+            set_if("crf", self.crf)
+            set_if("preset", self.preset)
+            if encoder_threads is not None:
+                set_if("threads", encoder_threads)
+                set_if("row-mt", 1)
         elif self.vcodec in ("h264_videotoolbox", "hevc_videotoolbox"):
             if self.crf is not None:
-                opts["q:v"] = max(1, min(100, 100 - self.crf * 2))
+                set_if("q:v", max(1, min(100, 100 - self.crf * 2)))
         elif self.vcodec in ("h264_nvenc", "hevc_nvenc"):
-            opts["rc"] = 0
+            set_if("rc", 0)
             set_if("qp", self.crf)
             set_if("preset", self.preset)
         elif self.vcodec == "h264_vaapi":

src/lerobot/datasets/compute_stats.py

@@ -509,7 +509,7 @@ def compute_episode_stats(
         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``) skip
-        this rescaling and remain in their stored units.
+        this rescaling and remain in their stored units (stored in ``depth_unit``).
     """
     if quantile_list is None:
         quantile_list = DEFAULT_QUANTILES

src/lerobot/datasets/dataset_metadata.py

@@ -26,12 +26,13 @@ import pyarrow as pa
 import pyarrow.parquet as pq
 from huggingface_hub import snapshot_download
 
-from lerobot.configs import VideoEncoderConfig
+from lerobot.configs import DEPTH_METER_UNIT, VideoEncoderConfig
 from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE
 from lerobot.utils.feature_utils import _validate_feature_names
 from lerobot.utils.utils import flatten_dict
 
 from .compute_stats import aggregate_stats
+from .depth_utils import MM_PER_METRE
 from .feature_utils import create_empty_dataset_info
 from .io_utils import (
     get_file_size_in_mb,
@@ -358,6 +359,35 @@ class LeRobotDatasetMetadata:
 
         return [key for key, ft in self.features.items() if _is_depth(ft)]
 
+    def rescale_depth_stats(self, output_unit: str) -> None:
+        """Rescale depth feature stats in place from their recorded unit to ``output_unit``.
+
+        Depth stats are stored in the unit the frames were recorded in
+        (``features[key]["info"]["depth_unit"]``), while frames are returned in
+        ``output_unit`` on read. This converts the unit-bearing stat entries so
+        stats match the frames consumers see.
+        """
+        missing_unit_keys = [
+            key for key in self.depth_keys if (self.features[key].get("info") or {}).get("depth_unit") is None
+        ]
+        if missing_unit_keys:
+            logging.warning(
+                f"Depth feature(s) {missing_unit_keys} have no recorded 'depth_unit' in their info. "
+                f"Depth maps and stats for these keys will be returned AS IS, with no unit conversion "
+                f"to the requested output unit {output_unit!r}. Re-record the dataset or set 'depth_unit' "
+                f"in the feature info (meta/info.json) to enable conversion."
+            )
+        if self.stats is None:
+            return
+        for key in self.depth_keys:
+            stored_unit = (self.features[key].get("info") or {}).get("depth_unit")
+            if stored_unit is None or stored_unit == output_unit or key not in self.stats:
+                continue
+            factor = MM_PER_METRE if stored_unit == DEPTH_METER_UNIT else 1.0 / MM_PER_METRE
+            self.stats[key] = {
+                stat: value if stat == "count" else value * factor for stat, value in self.stats[key].items()
+            }
+
     @property
     def camera_keys(self) -> list[str]:
         """Keys to access visual modalities (regardless of their storage method)."""

src/lerobot/datasets/dataset_reader.py

@@ -22,10 +22,14 @@ from pathlib import Path
 import datasets
 import torch
 
-from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig
+from lerobot.configs import (
+    DEFAULT_DEPTH_UNIT,
+    DEPTH_METER_UNIT,
+    DepthEncoderConfig,
+)
 
 from .dataset_metadata import LeRobotDatasetMetadata
-from .depth_utils import dequantize_depth
+from .depth_utils import MM_PER_METRE, dequantize_depth
 from .feature_utils import (
     check_delta_timestamps,
     get_delta_indices,
@@ -102,6 +106,13 @@ class DatasetReader:
             for vid_key in self._meta.depth_keys
         }
 
+        # Get the input unit of each depth feature stored as raw images.
+        self._image_depth_units: dict[str, str | None] = {
+            key: (self._meta.features[key].get("info") or {}).get("depth_unit")
+            for key in self._meta.depth_keys
+            if key in self._meta.image_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):
@@ -329,6 +340,13 @@ class DatasetReader:
                     continue
                 item[cam] = self._image_transforms(item[cam])
 
+        # Convert depth features to the output unit.
+        for key, stored_unit in self._image_depth_units.items():
+            if key in item and stored_unit is not None and stored_unit != self._depth_output_unit:
+                item[key] = (
+                    item[key] * MM_PER_METRE if stored_unit == DEPTH_METER_UNIT else item[key] / MM_PER_METRE
+                )
+
         # Add task as a string
         task_idx = item["task_index"].item()
         item["task"] = self._meta.tasks.iloc[task_idx].name

src/lerobot/datasets/dataset_writer.py

@@ -36,6 +36,7 @@ from lerobot.configs import (
     RGBEncoderConfig,
     VideoEncoderConfig,
     depth_encoder_defaults,
+    infer_depth_unit,
     rgb_encoder_defaults,
 )
 
@@ -209,6 +210,15 @@ class DatasetWriter:
         self.episode_buffer["timestamp"].append(timestamp)
         self.episode_buffer["task"].append(frame.pop("task"))
 
+        # Record each depth feature's input unit once, inferred from the first frame's dtype.
+        if frame_index == 0:
+            for depth_key in self._meta.depth_keys:
+                if depth_key not in frame:
+                    continue
+                info = self._meta.features[depth_key].setdefault("info", {})
+                if info.get("depth_unit") is None:
+                    info["depth_unit"] = infer_depth_unit(np.asarray(frame[depth_key]).dtype)
+
         # Start streaming encoder on first frame of episode
         if frame_index == 0 and self._streaming_encoder is not None:
             self._streaming_encoder.start_episode(

src/lerobot/datasets/depth_utils.py

@@ -34,12 +34,13 @@ from lerobot.configs.video import (
     DEPTH_METER_UNIT,
     DEPTH_MILLIMETER_UNIT,
     DEPTH_QMAX,
+    infer_depth_unit,
 )
 
 from .image_writer import squeeze_single_channel
 from .pyav_utils import write_u16_plane
 
-_MM_PER_METRE = 1000.0
+MM_PER_METRE = 1000.0
 _UINT16_MAX = 65535
 
 
@@ -57,11 +58,7 @@ def _depth_input_to_float32_and_unit(
     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
-    )
+    resolved_unit = infer_depth_unit(depth.dtype) if input_unit == "auto" else input_unit
     return depth.astype(np.float32, order="K"), resolved_unit
 
 
@@ -126,12 +123,12 @@ def quantize_depth(
 
     # 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)
+        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)
+        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)
+    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:
@@ -236,7 +233,7 @@ def dequantize_depth(
 
         # 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)
+        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)
@@ -259,7 +256,7 @@ def dequantize_depth(
     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.multiply(buf, MM_PER_METRE, out=buf)
     np.rint(buf, out=buf)
     np.clip(buf, 0.0, _UINT16_MAX, out=buf)
     if output_tensor:

src/lerobot/datasets/language.py

@@ -43,10 +43,10 @@ CORE_STYLES = {
 # validation. Empty by default — populate from a downstream module that
 # also extends ``PERSISTENT_STYLES`` or ``EVENT_ONLY_STYLES`` to declare
 # the new style's column.
-EXTENDED_STYLES: set[str] = set()
+EXTENDED_STYLES: set[str] = {"advantage"}
 STYLE_REGISTRY = CORE_STYLES | EXTENDED_STYLES
 
-PERSISTENT_STYLES = {"subtask", "plan", "memory", "motion", "task_aug"}
+PERSISTENT_STYLES = {"subtask", "plan", "memory", "motion", "task_aug", "advantage"}
 EVENT_ONLY_STYLES = {"interjection", "vqa", "trace"}
 
 # Styles whose ``content`` is grounded in a specific camera view. Rows of these

src/lerobot/datasets/lerobot_dataset.py

@@ -224,6 +224,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         )
         self.root = self.meta.root
         self.revision = self.meta.revision
+        self.meta.rescale_depth_stats(self._depth_output_unit)
 
         if episodes is not None and any(
             episode >= self.meta.total_episodes or episode < 0 for episode in episodes
@@ -350,6 +351,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
         """Frames per second used during data collection."""
         return self.meta.fps
 
+    @property
+    def depth_output_unit(self) -> str:
+        """Physical unit (``"m"`` or ``"mm"``) depth maps and statistics are returned in on read."""
+        return self._depth_output_unit
+
     @property
     def num_frames(self) -> int:
         """Number of frames in selected episodes."""

src/lerobot/datasets/streaming_dataset.py

@@ -22,11 +22,11 @@ import numpy as np
 import torch
 from datasets import load_dataset
 
-from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig
+from lerobot.configs import DEFAULT_DEPTH_UNIT, DEPTH_METER_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 .depth_utils import MM_PER_METRE, dequantize_depth
 from .feature_utils import get_delta_indices
 from .io_utils import item_to_torch
 from .utils import (
@@ -310,6 +310,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         )
         self.root = self.meta.root
         self.revision = self.meta.revision
+        self.meta.rescale_depth_stats(self._depth_output_unit)
         # Check version
         check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION)
 
@@ -318,6 +319,13 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             for vid_key in self.meta.depth_keys
         }
 
+        # Input unit of each depth feature stored as raw images (dequantized separately from videos).
+        self._image_depth_units: dict[str, str | None] = {
+            key: (self.meta.features[key].get("info") or {}).get("depth_unit")
+            for key in self.meta.depth_keys
+            if key in self.meta.image_keys
+        }
+
         self.delta_timestamps = None
         self.delta_indices = None
 
@@ -348,6 +356,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
     def fps(self):
         return self.meta.fps
 
+    @property
+    def depth_output_unit(self) -> str:
+        """Physical unit (``"m"`` or ``"mm"``) depth maps are returned in on read."""
+        return self._depth_output_unit
+
     @staticmethod
     def _iter_random_indices(
         rng: np.random.Generator, buffer_size: int, random_batch_size=100
@@ -530,6 +543,15 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         for update in updates:
             result.update(update)
 
+        # Convert raw-image depth features to the output unit (video depth is already converted).
+        for key, stored_unit in self._image_depth_units.items():
+            if key in result and stored_unit is not None and stored_unit != self._depth_output_unit:
+                result[key] = (
+                    result[key] * MM_PER_METRE
+                    if stored_unit == DEPTH_METER_UNIT
+                    else result[key] / MM_PER_METRE
+                )
+
         result["task"] = self.meta.tasks.iloc[item["task_index"]].name
 
         yield result

src/lerobot/envs/configs.py

@@ -757,7 +757,7 @@ class RoboTwinEnvConfig(EnvConfig):
 
     task: str = "beat_block_hammer"  # single task or comma-separated list
     fps: int = 25
-    episode_length: int = 300
+    episode_length: int = 1200
     obs_type: str = "pixels_agent_pos"
     render_mode: str = "rgb_array"
     # Available cameras from RoboTwin's aloha-agilex embodiment: head_camera
@@ -768,6 +768,9 @@ class RoboTwinEnvConfig(EnvConfig):
     # must equal what SAPIEN actually renders.
     observation_height: int = 240
     observation_width: int = 320
+    # "joint": 14-d joint-space control. "ee": 16-d end-effector-pose deltas executed via CuRobo IK
+    # (for world-model policies like LingBot-VA that predict per-arm xyz+quaternion+gripper poses).
+    action_mode: str = "joint"
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
             ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(14,)),
@@ -784,6 +787,8 @@ class RoboTwinEnvConfig(EnvConfig):
     )
 
     def __post_init__(self):
+        if self.action_mode == "ee":
+            self.features[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(16,))
         cam_list = [c.strip() for c in self.camera_names.split(",") if c.strip()]
         for cam in cam_list:
             self.features[f"pixels/{cam}"] = PolicyFeature(
@@ -826,6 +831,7 @@ class RoboTwinEnvConfig(EnvConfig):
             observation_height=self.observation_height,
             observation_width=self.observation_width,
             episode_length=self.episode_length,
+            action_mode=self.action_mode,
         )
 
 

src/lerobot/envs/robotwin.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 
 import importlib
 import logging
+import os
 from collections import defaultdict
 from collections.abc import Callable, Sequence
 from functools import partial
@@ -28,9 +29,17 @@ import torch
 from gymnasium import spaces
 
 from lerobot.types import RobotObservation
+from lerobot.utils.import_utils import _scipy_available
 
 from .utils import _LazyAsyncVectorEnv
 
+# scipy is only used for end-effector-pose composition (``--env.action_mode=ee``); guard it so this
+# module (and its base-env unit tests, which mock the RoboTwin runtime) imports without scipy installed.
+if _scipy_available:
+    from scipy.spatial.transform import Rotation
+else:
+    Rotation = None
+
 logger = logging.getLogger(__name__)
 
 # Camera names as used by RoboTwin 2.0. The wrapper appends "_rgb" when looking
@@ -41,10 +50,124 @@ ROBOTWIN_CAMERA_NAMES: tuple[str, ...] = (
     "right_camera",
 )
 
-ACTION_DIM = 14  # 7 DOF × 2 arms
+ACTION_DIM = 14  # 7 DOF × 2 arms (joint-space control mode)
+# End-effector-pose control mode: per arm [x, y, z, qx, qy, qz, qw, gripper] = 8, dual-arm = 16.
+# Used by world-model policies (e.g. LingBot-VA) that predict eef-pose deltas executed via CuRobo IK.
+EEF_ACTION_DIM = 16
 ACTION_LOW = -1.0
 ACTION_HIGH = 1.0
-DEFAULT_EPISODE_LENGTH = 300
+DEFAULT_EPISODE_LENGTH = 1200
+OFFICIAL_INSTRUCTION_ENV = "LEROBOT_ROBOTWIN_OFFICIAL_INSTRUCTION"
+OFFICIAL_INSTRUCTION_TYPE_ENV = "LEROBOT_ROBOTWIN_INSTRUCTION_TYPE"
+OFFICIAL_INSTRUCTION_MAX_ENV = "LEROBOT_ROBOTWIN_INSTRUCTION_MAX"
+
+
+def _compose_eef_pose(new_pose: np.ndarray, init_pose: np.ndarray) -> np.ndarray:
+    """Compose a single-arm predicted delta pose onto the initial pose.
+
+    ``new_pose`` / ``init_pose`` are 8-vectors ``[x, y, z, qx, qy, qz, qw, gripper]``. Translation
+    is added, rotation is composed (``init_R * new_R``), and the gripper is taken from the
+    prediction. Mirrors ``add_eef_pose`` in the upstream LingBot-VA RoboTwin client.
+    """
+    new_r = Rotation.from_quat(new_pose[3:7])
+    init_r = Rotation.from_quat(init_pose[3:7])
+    out_rot = (init_r * new_r).as_quat()
+    out_trans = new_pose[:3] + init_pose[:3]
+    return np.concatenate([out_trans, out_rot, new_pose[7:8]])
+
+
+def _add_init_eef_pose(delta_pose: np.ndarray, init_pose: np.ndarray) -> np.ndarray:
+    """Compose a dual-arm (16-d) predicted delta pose onto the initial eef pose, normalizing quats."""
+    left = _compose_eef_pose(delta_pose[:8], init_pose[:8])
+    right = _compose_eef_pose(delta_pose[8:], init_pose[8:])
+    out = np.concatenate([left, right])
+    # Normalize the two quaternions (indices 3:7 and 11:15) as the upstream client does.
+    out[3:7] = out[3:7] / (np.linalg.norm(out[3:7]) + 1e-8)
+    out[11:15] = out[11:15] / (np.linalg.norm(out[11:15]) + 1e-8)
+    return out
+
+
+def _env_flag(name: str, default: bool = False) -> bool:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    return raw.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _arm_for_block(block: Any) -> str:
+    return "left" if float(block.get_pose().p[0]) < 0 else "right"
+
+
+def _robotwin_blocks_episode_info(task_name: str, env: Any) -> dict[str, str] | None:
+    """Infer the episode-info dict used by RoboTwin's official instruction generator for block ranking."""
+    if task_name == "blocks_ranking_rgb":
+        return {
+            "{A}": "red block",
+            "{B}": "green block",
+            "{C}": "blue block",
+            "{a}": _arm_for_block(env.block1),
+            "{b}": _arm_for_block(env.block2),
+            "{c}": _arm_for_block(env.block3),
+        }
+    if task_name == "blocks_ranking_size":
+        return {
+            "{A}": "large block",
+            "{B}": "medium block",
+            "{C}": "small block",
+            "{a}": _arm_for_block(env.block1),
+            "{b}": _arm_for_block(env.block2),
+            "{c}": _arm_for_block(env.block3),
+        }
+    return None
+
+
+def _generate_robotwin_official_instruction(task_name: str, env: Any) -> str:
+    """Generate language with RoboTwin's official task templates, matching its eval client."""
+    fallback = task_name.replace("_", " ")
+    episode_info = _robotwin_blocks_episode_info(task_name, env)
+    if episode_info is None:
+        logger.warning(
+            "Official RoboTwin instruction is not implemented for task=%s; using %r.", task_name, fallback
+        )
+        return fallback
+
+    try:
+        # Part of the robotwin simulator repo, this is being pulled by the docker image running robotwin
+        # see https://github.com/RoboTwin-Platform/RoboTwin/tree/main/description
+        # Used to generate the official instructions
+        from description.utils.generate_episode_instructions import generate_episode_descriptions
+    except Exception:
+        logger.warning(
+            "Failed to import RoboTwin official instruction generator; using %r.", fallback, exc_info=True
+        )
+        return fallback
+
+    instruction_type = os.environ.get(OFFICIAL_INSTRUCTION_TYPE_ENV, "seen")
+    try:
+        max_descriptions = int(os.environ.get(OFFICIAL_INSTRUCTION_MAX_ENV, "1000000"))
+    except ValueError:
+        max_descriptions = 1000000
+
+    results = generate_episode_descriptions(task_name, [episode_info], max_descriptions=max_descriptions)
+    if not results:
+        logger.warning(
+            "RoboTwin generated no official instructions for task=%s; using %r.", task_name, fallback
+        )
+        return fallback
+
+    options = results[0].get(instruction_type) or results[0].get("seen") or results[0].get("unseen")
+    if not options:
+        logger.warning(
+            "RoboTwin generated no %s official instructions for task=%s; using %r.",
+            instruction_type,
+            task_name,
+            fallback,
+        )
+        return fallback
+
+    return str(np.random.choice(options))
+
+
 # D435 dims from task_config/_camera_config.yml (what demo_clean.yml selects).
 DEFAULT_CAMERA_H = 240
 DEFAULT_CAMERA_W = 320
@@ -234,6 +357,7 @@ class RoboTwinEnv(gym.Env):
         observation_width: int | None = None,
         episode_length: int = DEFAULT_EPISODE_LENGTH,
         render_mode: str = "rgb_array",
+        action_mode: str = "joint",
     ):
         super().__init__()
         self.task_name = task_name
@@ -241,6 +365,13 @@ class RoboTwinEnv(gym.Env):
         self.task_description = task_name.replace("_", " ")
         self.episode_index = episode_index
         self._reset_stride = n_envs
+        # "joint": 14-d joint-space actions via take_action(action). "ee": 16-d end-effector-pose
+        # deltas (added onto the episode's initial eef pose) executed via take_action(.., "ee") + IK.
+        if action_mode not in ("joint", "ee"):
+            raise ValueError(f"action_mode must be 'joint' or 'ee'; got {action_mode!r}")
+        self.action_mode = action_mode
+        self._action_dim = EEF_ACTION_DIM if action_mode == "ee" else ACTION_DIM
+        self._init_eef_pose: np.ndarray | None = None
         self.camera_names = list(camera_names)
         # Default to D435 dims (the camera type baked into task_config/demo_clean.yml).
         # The YAML-driven lookup is deferred to reset() so construction doesn't
@@ -271,7 +402,7 @@ class RoboTwinEnv(gym.Env):
             }
         )
         self.action_space = spaces.Box(
-            low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32
+            low=ACTION_LOW, high=ACTION_HIGH, shape=(self._action_dim,), dtype=np.float32
         )
 
     def _ensure_env(self) -> None:
@@ -317,6 +448,18 @@ class RoboTwinEnv(gym.Env):
 
         return {"pixels": images, "agent_pos": joint_state}
 
+    def _read_eef_pose(self) -> np.ndarray:
+        """Read the current 16-d dual-arm eef pose [left(xyz+quat)+grip, right(xyz+quat)+grip]."""
+        assert self._env is not None, "_read_eef_pose called before _ensure_env()"
+        ep = self._env.get_obs()["endpose"]
+        pose = (
+            list(ep["left_endpose"])
+            + [ep["left_gripper"]]
+            + list(ep["right_endpose"])
+            + [ep["right_gripper"]]
+        )
+        return np.asarray(pose, dtype=np.float64)
+
     def reset(self, seed: int | None = None, **kwargs) -> tuple[RobotObservation, dict]:
         self._ensure_env()
         super().reset(seed=seed)
@@ -330,16 +473,32 @@ class RoboTwinEnv(gym.Env):
         self.episode_index += self._reset_stride
         self._step_count = 0
 
+        use_official_instruction = self.task_name in {"blocks_ranking_rgb", "blocks_ranking_size"}
+        if _env_flag(OFFICIAL_INSTRUCTION_ENV, default=use_official_instruction):
+            self.task_description = _generate_robotwin_official_instruction(self.task_name, self._env)
+            if hasattr(self._env, "set_instruction"):
+                self._env.set_instruction(instruction=self.task_description)
+            logger.info("RoboTwin official instruction | task=%s | %s", self.task_name, self.task_description)
+        else:
+            self.task_description = self.task_name.replace("_", " ")
+
+        # In eef mode the policy predicts pose deltas relative to the initial eef pose.
+        if self.action_mode == "ee":
+            self._init_eef_pose = self._read_eef_pose()
+
         obs = self._get_obs()
         return obs, {"is_success": False, "task": self.task_name}
 
     def step(self, action: np.ndarray) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]:
         assert self._env is not None, "step() called before reset()"
-        if action.ndim != 1 or action.shape[0] != ACTION_DIM:
-            raise ValueError(f"Expected 1-D action of shape ({ACTION_DIM},), got {action.shape}")
+        if action.ndim != 1 or action.shape[0] != self._action_dim:
+            raise ValueError(f"Expected 1-D action of shape ({self._action_dim},), got {action.shape}")
 
         with torch.enable_grad():
-            if hasattr(self._env, "take_action"):
+            if self.action_mode == "ee":
+                ee_action = _add_init_eef_pose(np.asarray(action, dtype=np.float64), self._init_eef_pose)
+                self._env.take_action(ee_action, action_type="ee")
+            elif hasattr(self._env, "take_action"):
                 self._env.take_action(action)
             else:
                 self._env.step(action)
@@ -398,6 +557,7 @@ def _make_env_fns(
     observation_height: int,
     observation_width: int,
     episode_length: int,
+    action_mode: str = "joint",
 ) -> list[Callable[[], RoboTwinEnv]]:
     """Return n_envs factory callables for a single task."""
 
@@ -410,6 +570,7 @@ def _make_env_fns(
             observation_height=observation_height,
             observation_width=observation_width,
             episode_length=episode_length,
+            action_mode=action_mode,
         )
 
     return [partial(_make_one, i) for i in range(n_envs)]
@@ -423,6 +584,7 @@ def create_robotwin_envs(
     observation_height: int = DEFAULT_CAMERA_H,
     observation_width: int = DEFAULT_CAMERA_W,
     episode_length: int = DEFAULT_EPISODE_LENGTH,
+    action_mode: str = "joint",
 ) -> dict[str, dict[int, Any]]:
     """Create vectorized RoboTwin 2.0 environments.
 
@@ -473,6 +635,7 @@ def create_robotwin_envs(
             observation_height=observation_height,
             observation_width=observation_width,
             episode_length=episode_length,
+            action_mode=action_mode,
         )
         if is_async:
             lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)

src/lerobot/jobs/__init__.py

@@ -0,0 +1,23 @@
+# 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"]

src/lerobot/jobs/dataset.py

@@ -0,0 +1,53 @@
+# 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.")

src/lerobot/jobs/hf.py

@@ -0,0 +1,425 @@
+# 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}"
+        )

src/lerobot/optim/schedulers.py

@@ -83,6 +83,28 @@ class VQBeTSchedulerConfig(LRSchedulerConfig):
         return LambdaLR(optimizer, lr_lambda, -1)
 
 
+@LRSchedulerConfig.register_subclass("constant_with_warmup")
+@dataclass
+class ConstantWithWarmupSchedulerConfig(LRSchedulerConfig):
+    """Linear warmup followed by a constant learning rate.
+
+    Mirrors the ``warmup_constant_lambda`` used by LingBot-VA (upstream ``wan_va/train.py``):
+    the LR ramps linearly from 0 to the peak over ``num_warmup_steps`` steps, then stays flat.
+    """
+
+    num_warmup_steps: int = 1000
+
+    def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
+        warmup_steps = self.num_warmup_steps or 0
+
+        def lr_lambda(current_step):
+            if current_step < warmup_steps:
+                return float(current_step) / float(max(1, warmup_steps))
+            return 1.0
+
+        return LambdaLR(optimizer, lr_lambda, -1)
+
+
 @LRSchedulerConfig.register_subclass("cosine_decay_with_warmup")
 @dataclass
 class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):

src/lerobot/policies/__init__.py

@@ -18,8 +18,10 @@ 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 .lingbot_va.configuration_lingbot_va import LingBotVAConfig as LingBotVAConfig
 from .molmoact2.configuration_molmoact2 import MolmoAct2Config as MolmoAct2Config
 from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig as MultiTaskDiTConfig
 from .pi0.configuration_pi0 import PI0Config as PI0Config
@@ -42,8 +44,10 @@ __all__ = [
     "ACTConfig",
     "DiffusionConfig",
     "EO1Config",
+    "FastWAMConfig",
     "GaussianActorConfig",
     "GrootConfig",
+    "LingBotVAConfig",
     "MolmoAct2Config",
     "MultiTaskDiTConfig",
     "PI0Config",

src/lerobot/policies/factory.py

@@ -47,8 +47,10 @@ 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 .lingbot_va.configuration_lingbot_va import LingBotVAConfig
 from .molmoact2.configuration_molmoact2 import MolmoAct2Config
 from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
 from .pi0.configuration_pi0 import PI0Config
@@ -162,6 +164,14 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
 
         return VLAJEPAPolicy
+    elif name == "lingbot_va":
+        from .lingbot_va.modeling_lingbot_va import LingBotVAPolicy
+
+        return LingBotVAPolicy
+    elif name == "fastwam":
+        from .fastwam.modeling_fastwam import FastWAMPolicy
+
+        return FastWAMPolicy
     else:
         try:
             return _get_policy_cls_from_policy_name(name=name)
@@ -218,6 +228,10 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return MolmoAct2Config(**kwargs)
     elif policy_type == "vla_jepa":
         return VLAJEPAConfig(**kwargs)
+    elif policy_type == "lingbot_va":
+        return LingBotVAConfig(**kwargs)
+    elif policy_type == "fastwam":
+        return FastWAMConfig(**kwargs)
     else:
         try:
             config_cls = PreTrainedConfig.get_choice_class(policy_type)
@@ -451,6 +465,22 @@ def make_pre_post_processors(
             dataset_stats=kwargs.get("dataset_stats"),
         )
 
+    elif isinstance(policy_cfg, LingBotVAConfig):
+        from .lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
+
+        processors = make_lingbot_va_pre_post_processors(
+            config=policy_cfg,
+            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(

src/lerobot/policies/fastwam/README.md

@@ -0,0 +1 @@
+../../../../docs/source/policy_fastwam_README.md

src/lerobot/policies/fastwam/__init__.py

@@ -0,0 +1,23 @@
+# 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",
+]

src/lerobot/policies/fastwam/configuration_fastwam.py

@@ -0,0 +1,399 @@
+# 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

src/lerobot/policies/fastwam/modeling_fastwam.py

@@ -0,0 +1,440 @@
+# 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))

src/lerobot/policies/fastwam/processor_fastwam.py

@@ -0,0 +1,142 @@
+# 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,
+        ),
+    )

src/lerobot/policies/fastwam/wan/README.md

@@ -0,0 +1,34 @@
+# 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`, …).

src/lerobot/policies/fastwam/wan/__init__.py

@@ -0,0 +1,33 @@
+# 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",
+]

src/lerobot/policies/fastwam/wan/adapters.py

@@ -0,0 +1,108 @@
+# 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)

src/lerobot/policies/fastwam/wan/components.py

@@ -0,0 +1,175 @@
+# 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

src/lerobot/policies/fastwam/wan/model.py

@@ -0,0 +1,341 @@
+# 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)

src/lerobot/policies/fastwam/wan/video_dit.py

@@ -0,0 +1,800 @@
+# 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)

src/lerobot/policies/lingbot_va/README.md

@@ -0,0 +1 @@
+../../../../docs/source/lingbot_va.mdx

src/lerobot/policies/lingbot_va/__init__.py

@@ -0,0 +1,21 @@
+#!/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.
+
+from .configuration_lingbot_va import LingBotVAConfig
+from .modeling_lingbot_va import LingBotVAPolicy
+from .processor_lingbot_va import make_lingbot_va_pre_post_processors
+
+__all__ = ["LingBotVAConfig", "LingBotVAPolicy", "make_lingbot_va_pre_post_processors"]

src/lerobot/policies/lingbot_va/configuration_lingbot_va.py

@@ -0,0 +1,168 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# 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.
+
+"""Configuration for the LingBot-VA policy.
+
+LingBot-VA is an autoregressive video-action world-model policy built on the Wan2.2
+video-diffusion stack. It interleaves prediction of future video latents and robot
+actions in a single dual-stream transformer. See ``docs/source/lingbot_va.mdx`` and the
+upstream repository (https://github.com/Robbyant/lingbot-va).
+
+Defaults below match the upstream LIBERO configuration (``wan_va/configs/va_libero_cfg.py``)
+and the ``transformer/config.json`` of the released checkpoints.
+"""
+
+from dataclasses import dataclass, field
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim.optimizers import AdamWConfig
+from lerobot.optim.schedulers import ConstantWithWarmupSchedulerConfig, LRSchedulerConfig
+from lerobot.utils.constants import ACTION
+
+
+@PreTrainedConfig.register_subclass("lingbot_va")
+@dataclass
+class LingBotVAConfig(PreTrainedConfig):
+    """Configuration for the native LingBot-VA policy integration in LeRobot."""
+
+    # Wan transformer architecture
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    num_attention_heads: int = 24
+    attention_head_dim: int = 128
+    in_channels: int = 48
+    out_channels: int = 48
+    action_dim: int = 30
+    text_dim: int = 4096
+    freq_dim: int = 256
+    ffn_dim: int = 14336
+    num_layers: int = 30
+    cross_attn_norm: bool = True
+    eps: float = 1e-6
+    rope_max_seq_len: int = 1024
+    # "flex" = training only (needs recent torch); inference uses "torch" SDPA or "flashattn".
+    attn_mode: str = "torch"
+
+    # Frozen sub-models (VAE + UMT5 text encoder + tokenizer)
+    # ~20 GB of frozen weights, NOT bundled in the checkpoint; lazily pulled from this HF repo /
+    # local dir (must hold diffusers-style ``vae/``, ``text_encoder/``, ``tokenizer/`` sub-folders).
+    wan_pretrained_path: str = "robbyant/lingbot-va-base"
+    dtype: str = "bfloat16"  # transformer / VAE / text-encoder dtype: "bfloat16", "float16", "float32"
+    # Frozen UMT5-XXL encoder device; "cpu" frees ~11 GB VRAM (it runs once per episode).
+    text_encoder_device: str = "cpu"
+
+    # Observation cameras (order matters: latents are concatenated on width; LIBERO defaults)
+    obs_cam_keys: list[str] = field(
+        default_factory=lambda: ["observation.images.image", "observation.images.image2"]
+    )
+    # Undo the LIBERO env processor's extra horizontal flip to match the model's training orientation.
+    image_hflip: bool = False
+    # Camera latent layout: "width_concat" (cameras concatenated on width; LIBERO) or
+    # "robotwin_tshape" (full-res head + half-res wrists in a "T"; RoboTwin).
+    camera_layout: str = "width_concat"
+
+    # Inference hyperparameters (LIBERO defaults)
+    n_obs_steps: int = 1
+    height: int = 128
+    width: int = 128
+    action_per_frame: int = 4
+    frame_chunk_size: int = 4
+    attn_window: int = 30
+    num_inference_steps: int = 20
+    video_exec_step: int = -1
+    action_num_inference_steps: int = 50
+    guidance_scale: float = 5.0
+    action_guidance_scale: float = 1.0
+    snr_shift: float = 5.0
+    action_snr_shift: float = 0.05
+    max_sequence_length: int = 512  # UMT5 prompt length
+
+    # Subset of the 30-d action space used by the benchmark (LIBERO = 7-DoF). The action
+    # (un)normalization quantiles live in the checkpoint's ``policy_postprocessor.json``, not here.
+    used_action_channel_ids: list[int] = field(default_factory=lambda: list(range(7)))
+
+    # Opt-in: VAE-decode predicted video latents to ``self.last_predicted_frames`` for saving MP4s.
+    save_predicted_video: bool = False
+
+    # Normalization: IDENTITY here; images are scaled + VAE-encoded and actions are
+    # quantile-(un)normalized inside the policy / dedicated processor steps.
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.IDENTITY,
+            "ACTION": NormalizationMode.IDENTITY,
+        }
+    )
+
+    # Optimizer / scheduler (training; AdamW + warmup-constant per upstream train.py)
+    optimizer_lr: float = 1e-5
+    optimizer_betas: tuple[float, float] = (0.9, 0.95)
+    optimizer_eps: float = 1e-8
+    optimizer_weight_decay: float = 1e-4
+    optimizer_grad_clip_norm: float = 1.0
+    scheduler_warmup_steps: int = 1000
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.attn_mode not in ("torch", "flashattn", "flex"):
+            raise ValueError(f"attn_mode must be one of 'torch', 'flashattn', 'flex'; got {self.attn_mode!r}")
+
+    @property
+    def chunk_size(self) -> int:
+        """Number of single-step actions produced per autoregressive chunk."""
+        return self.frame_chunk_size * self.action_per_frame
+
+    @property
+    def n_action_steps(self) -> int:
+        """Number of actions executed before refilling (the whole chunk)."""
+        return self.chunk_size
+
+    def validate_features(self) -> None:
+        image_features = [key for key, feat in self.input_features.items() if feat.type == FeatureType.VISUAL]
+        if not image_features:
+            raise ValueError(
+                "LingBot-VA requires at least one visual input feature. "
+                "No features of type FeatureType.VISUAL found in input_features."
+            )
+        if ACTION not in self.output_features:
+            self.output_features[ACTION] = PolicyFeature(
+                type=FeatureType.ACTION, shape=(len(self.used_action_channel_ids),)
+            )
+
+    def get_optimizer_preset(self) -> AdamWConfig:
+        return AdamWConfig(
+            lr=self.optimizer_lr,
+            betas=self.optimizer_betas,
+            eps=self.optimizer_eps,
+            weight_decay=self.optimizer_weight_decay,
+            grad_clip_norm=self.optimizer_grad_clip_norm,
+        )
+
+    def get_scheduler_preset(self) -> LRSchedulerConfig | None:
+        # Upstream uses a linear warmup followed by a constant LR (warmup_constant_lambda).
+        return ConstantWithWarmupSchedulerConfig(num_warmup_steps=self.scheduler_warmup_steps)
+
+    @property
+    def observation_delta_indices(self) -> list[int]:
+        temporal_downsample = 4
+        stride = max(1, self.action_per_frame // temporal_downsample)
+        return list(range(0, self.frame_chunk_size * temporal_downsample * stride, stride))
+
+    @property
+    def action_delta_indices(self) -> list[int]:
+        return list(range(self.chunk_size))
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None

src/lerobot/policies/lingbot_va/modeling_lingbot_va.py

@@ -0,0 +1,853 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+# 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.
+
+"""LingBot-VA policy: an autoregressive video-action world model on the Wan2.2 stack.
+
+The sampling loop is a faithful re-implementation of the upstream streaming server
+(``wan_va/wan_va_server.py``) and LIBERO client (``evaluation/libero/client.py``), adapted
+to LeRobot's ``select_action`` interface:
+
+  * the trainable dual-stream transformer is owned as a sub-module and round-trips in the
+    single ``model.safetensors`` checkpoint;
+  * the frozen Wan VAE + UMT5 text encoder + tokenizer are *lazily pulled* from
+    ``config.wan_pretrained_path`` (not bundled), so the LeRobot checkpoint stays small;
+  * ``predict_action_chunk`` runs one autoregressive chunk (video stream then action
+    stream, each with CFG and its own flow-matching scheduler) and updates the KV cache;
+  * ``select_action`` drains a per-step action queue and records the real observed
+    keyframes that are fed back into the KV cache when the queue is refilled.
+
+NOTE: The streaming path is written for single-environment eval (``--eval.batch_size=1``).
+"""
+
+from collections import deque
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from einops import rearrange
+from torch import Tensor
+
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.utils.constants import ACTION
+from lerobot.utils.import_utils import require_package
+
+from .configuration_lingbot_va import LingBotVAConfig
+from .utils import (
+    FlowMatchScheduler,
+    WanTransformer3DModel,
+    WanVAEStreamingWrapper,
+    _sample_timestep_id,
+    _torch_dtype,
+    clean_prompt,
+    data_seq_to_patch,
+    denormalize_latents,
+    get_mesh_id,
+    load_text_encoder,
+    load_tokenizer,
+    load_vae,
+)
+
+
+class LingBotVAPolicy(PreTrainedPolicy):
+    """LeRobot wrapper for the LingBot-VA autoregressive video-action world model."""
+
+    config_class = LingBotVAConfig
+    name = "lingbot_va"
+
+    def __init__(self, config: LingBotVAConfig, **kwargs):
+        require_package("diffusers", extra="lingbot_va")
+        require_package("transformers", extra="lingbot_va")
+        super().__init__(config)
+        config.validate_features()
+        self.config = config
+
+        self.dtype = _torch_dtype(config.dtype)
+
+        # Trainable dual-stream transformer (the only sub-module saved in the LeRobot checkpoint).
+        self.transformer = WanTransformer3DModel(
+            patch_size=tuple(config.patch_size),
+            num_attention_heads=config.num_attention_heads,
+            attention_head_dim=config.attention_head_dim,
+            in_channels=config.in_channels,
+            out_channels=config.out_channels,
+            action_dim=config.action_dim,
+            text_dim=config.text_dim,
+            freq_dim=config.freq_dim,
+            ffn_dim=config.ffn_dim,
+            num_layers=config.num_layers,
+            cross_attn_norm=config.cross_attn_norm,
+            eps=config.eps,
+            rope_max_seq_len=config.rope_max_seq_len,
+            attn_mode=config.attn_mode,
+        )
+        # Run the transformer in config.dtype (bf16); norm/modulation paths upcast to fp32 internally.
+        self.transformer = self.transformer.to(self.dtype)
+
+        # Frozen modules are stored OUTSIDE the nn.Module registry (plain dict) so they are
+        # neither saved into model.safetensors nor moved by ``.to()``. They are lazily loaded
+        # from ``config.wan_pretrained_path`` the first time inference runs.
+        self._frozen: dict = {}
+
+        self.last_predicted_frames: Tensor | None = None
+        self.last_predicted_latents: Tensor | None = None
+        self.reset()
+
+    # Frozen-module lazy loading (VAE + UMT5 + tokenizer)
+    def _ensure_frozen_modules(self):
+        if self._frozen:
+            return
+        path = self.config.wan_pretrained_path
+        device = self.config.device
+
+        # The frozen modules always live in ``vae/``, ``text_encoder/`` and ``tokenizer/``
+        # sub-folders -- both in the released diffusers-style HF repos and in the local
+        # ``--bundle-frozen`` output dir. ``from_pretrained(path, subfolder=...)`` resolves
+        # them for either a HF repo id or a local directory.
+        vae = load_vae(path, torch_dtype=self.dtype, torch_device=device, subfolder="vae")
+        # The UMT5-XXL text encoder (~11 GB) runs once per episode; keep it on its own
+        # (CPU by default) device so the 5B transformer + VAE fit on a single GPU.
+        text_encoder = load_text_encoder(
+            path,
+            torch_dtype=self.dtype,
+            torch_device=self.config.text_encoder_device,
+            subfolder="text_encoder",
+        )
+        tokenizer = load_tokenizer(path, subfolder="tokenizer")
+        self._frozen = {
+            "vae": vae.eval(),
+            "streaming_vae": WanVAEStreamingWrapper(vae),
+            "text_encoder": text_encoder.eval(),
+            "tokenizer": tokenizer,
+        }
+        # RoboTwin's T-shape layout encodes the half-resolution wrist cameras through a second
+        # streaming VAE (separate causal cache) alongside the full-res head camera.
+        if self.config.camera_layout == "robotwin_tshape":
+            vae_half = load_vae(path, torch_dtype=self.dtype, torch_device=device, subfolder="vae")
+            self._frozen["streaming_vae_half"] = WanVAEStreamingWrapper(vae_half.eval())
+
+    @property
+    def _vae(self):
+        return self._frozen["vae"]
+
+    @property
+    def _streaming_vae(self):
+        return self._frozen["streaming_vae"]
+
+    # PreTrainedPolicy API
+    def get_optim_params(self) -> dict:
+        # Only the transformer is trainable; the VAE / text encoder stay frozen (kept outside the
+        # nn.Module registry). With PEFT/LoRA this naturally returns just the adapter params.
+        return [p for p in self.transformer.parameters() if p.requires_grad]
+
+    def reset(self):
+        """Reset all per-episode streaming state (KV cache, queues, frame counter)."""
+        cfg = self.config
+        self._action_queue: deque = deque(maxlen=cfg.n_action_steps)
+        self._obs_buffer: list = []  # raw keyframe obs (one per env substep) observed this chunk
+        self._executed_actions: Tensor | None = (
+            None  # last chunk's actions (model-normalized) for KV feedback
+        )
+        self._started = False  # first select_action call uses the obs as the conditioning frame
+        self._exec_step = 0  # index of the action being executed within the current chunk
+        self._prev_j = 0  # sub-step index (within a predicted frame) of the last executed action
+        # Sample one keyframe every ``action_per_frame / temporal_downsample`` executed sub-steps so
+        # that exactly ``frame_chunk_size * temporal_downsample`` frames are VAE-encoded per chunk
+        # (the Wan2.2 VAE temporal downsample is 4 -> ``frame_chunk_size`` latent frames).
+        self._keyframe_stride = max(1, cfg.action_per_frame // 4)
+        self._frame_st_id = 0
+        self._first_chunk = True
+        self._prompt: str | None = None
+        self._prompt_embeds = None
+        self._negative_prompt_embeds = None
+        self.last_predicted_frames = None
+        self.last_predicted_latents = None
+        self._use_cfg = (cfg.guidance_scale > 1) or (cfg.action_guidance_scale > 1)
+        # Two independent flow-matching schedulers (video latent + action streams).
+        self._scheduler = FlowMatchScheduler(shift=cfg.snr_shift, sigma_min=0.0, extra_one_step=True)
+        self._action_scheduler = FlowMatchScheduler(
+            shift=cfg.action_snr_shift, sigma_min=0.0, extra_one_step=True
+        )
+        self._scheduler.set_timesteps(1000, training=True)
+        self._action_scheduler.set_timesteps(1000, training=True)
+        self._cache_initialised = False
+        # Clear KV cache on the (already-built) transformer, if present.
+        if hasattr(self, "transformer"):
+            self.transformer.clear_cache("pos")
+        # Reset the causal streaming-VAE feat cache between episodes (mirrors upstream ``_reset``).
+        # Without this the encoder carries over the previous episode's temporal state, corrupting the
+        # latent frame counts on the next episode's first encode.
+        if self._frozen:
+            self._frozen["streaming_vae"].clear_cache()
+            if "streaming_vae_half" in self._frozen:
+                self._frozen["streaming_vae_half"].clear_cache()
+
+    # Training (flow-matching dual-stream loss). Requires attn_mode="flex".
+    def _ensure_train_schedulers(self):
+        if getattr(self, "_train_sched_latent", None) is None:
+            cfg = self.config
+            self._train_sched_latent = FlowMatchScheduler(
+                shift=cfg.snr_shift, sigma_min=0.0, extra_one_step=True
+            )
+            self._train_sched_latent.set_timesteps(1000, training=True)
+            self._train_sched_action = FlowMatchScheduler(
+                shift=cfg.action_snr_shift, sigma_min=0.0, extra_one_step=True
+            )
+            self._train_sched_action.set_timesteps(1000, training=True)
+
+    @torch.no_grad()
+    def _add_noise_stream(self, latent, scheduler, action_mask, action_mode, noisy_cond_prob):
+        """Flow-matching noising of one stream (port of upstream ``Trainer._add_noise``)."""
+        device = latent.device
+        b, _c, f, _h, _w = latent.shape
+        p = self.config.patch_size
+        patch_f, patch_h, patch_w = (1, 1, 1) if action_mode else (p[0], p[1], p[2])
+
+        ts_ids = _sample_timestep_id(f, num_train_timesteps=scheduler.num_train_timesteps)
+        noise = torch.zeros_like(latent).normal_()
+        timesteps = scheduler.timesteps[ts_ids].to(device)
+        noisy_latents = scheduler.add_noise(latent, noise, timesteps, t_dim=2)
+        targets = scheduler.training_target(latent, noise, timesteps)
+
+        grid_id = (
+            get_mesh_id(
+                latent.shape[-3] // patch_f,
+                latent.shape[-2] // patch_h,
+                latent.shape[-1] // patch_w,
+                t=1 if action_mode else 0,
+                f_w=1,
+                f_shift=0,
+                action=action_mode,
+            )
+            .to(device)[None]
+            .repeat(b, 1, 1)
+        )
+
+        if torch.rand(1).item() < noisy_cond_prob:
+            cond_ids = _sample_timestep_id(
+                f, min_timestep_bd=0.5, max_timestep_bd=1.0, num_train_timesteps=scheduler.num_train_timesteps
+            )
+            cond_noise = torch.zeros_like(latent).normal_()
+            cond_timesteps = scheduler.timesteps[cond_ids].to(device)
+            latent = scheduler.add_noise(latent, cond_noise, cond_timesteps, t_dim=2)
+        else:
+            cond_timesteps = torch.zeros_like(timesteps)
+
+        if action_mask is not None:
+            noisy_latents = noisy_latents * action_mask.float()
+            targets = targets * action_mask.float()
+            latent = latent * action_mask.float()
+
+        return {
+            "timesteps": timesteps[None].repeat(b, 1),
+            "noisy_latents": noisy_latents,
+            "targets": targets,
+            "latent": latent,
+            "cond_timesteps": cond_timesteps[None].repeat(b, 1),
+            "grid_id": grid_id,
+        }
+
+    def _flow_matching_loss(self, input_dict, pred):
+        """Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``)."""
+        latent_pred, action_pred = pred
+        ld, ad = input_dict["latent_dict"], input_dict["action_dict"]
+        action_pred = rearrange(action_pred, "b (f n) c -> b c f n 1", f=ad["targets"].shape[-3])
+        latent_pred = data_seq_to_patch(
+            self.config.patch_size,
+            latent_pred,
+            ld["targets"].shape[-3],
+            ld["targets"].shape[-2],
+            ld["targets"].shape[-1],
+            batch_size=latent_pred.shape[0],
+        )
+        bn, fn = ld["timesteps"].shape
+        lw = self._train_sched_latent.training_weight(ld["timesteps"].flatten()).reshape(bn, fn)
+        aw = self._train_sched_action.training_weight(ad["timesteps"].flatten()).reshape(bn, fn)
+
+        latent_loss = F.mse_loss(latent_pred.float(), ld["targets"].float().detach(), reduction="none")
+        latent_loss = (
+            (latent_loss * lw[:, None, :, None, None]).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
+        )
+        latent_loss = (latent_loss.sum(dim=1) / (torch.ones_like(latent_loss).sum(dim=1) + 1e-6)).mean()
+
+        amask = ad["actions_mask"].float()
+        action_loss = F.mse_loss(action_pred.float(), ad["targets"].float().detach(), reduction="none")
+        action_loss = (
+            (action_loss * aw[:, None, :, None, None] * amask).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
+        )
+        amask_f = amask.permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
+        action_loss = (action_loss.sum(dim=1) / (amask_f.sum(dim=1) + 1e-6)).mean()
+        return latent_loss, action_loss
+
+    def training_loss_from_streams(self, latents, actions, actions_mask, text_emb):
+        """Core dual-stream training loss given prepared latents / actions / text embeddings.
+
+        ``latents``: ``[B, in_channels, F, h, w]`` (normalized video latents).
+        ``actions`` / ``actions_mask``: ``[B, action_dim, F, action_per_frame, 1]``.
+        ``text_emb``: ``[B, seq_len, text_dim]``. Returns ``(loss, {latent_loss, action_loss})``.
+        """
+        if self.config.attn_mode != "flex":
+            raise ValueError(
+                "LingBot-VA training requires attn_mode='flex' (block-causal flow-matching masks). "
+                "Load/convert the policy with --policy.attn_mode=flex for training/fine-tuning."
+            )
+        self._ensure_train_schedulers()
+        latent_dict = self._add_noise_stream(
+            latents, self._train_sched_latent, action_mask=None, action_mode=False, noisy_cond_prob=0.5
+        )
+        action_dict = self._add_noise_stream(
+            actions, self._train_sched_action, action_mask=actions_mask, action_mode=True, noisy_cond_prob=0.0
+        )
+        latent_dict["text_emb"] = text_emb
+        action_dict["text_emb"] = text_emb
+        action_dict["actions_mask"] = actions_mask
+        input_dict = {
+            "latent_dict": latent_dict,
+            "action_dict": action_dict,
+            "chunk_size": int(torch.randint(1, 5, (1,)).item()),
+            "window_size": int(torch.randint(4, 65, (1,)).item()),
+        }
+        pred = self.transformer(input_dict, train_mode=True)
+        latent_loss, action_loss = self._flow_matching_loss(input_dict, pred)
+        loss = latent_loss + action_loss
+        return loss, {"latent_loss": latent_loss.item(), "action_loss": action_loss.item()}
+
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
+        """Training forward: dual-stream flow-matching loss.
+
+        Builds the (video-latent, action, text) training streams from a LeRobot batch
+        (VAE-encoding the camera frames and UMT5-encoding the task), then runs the flow-matching
+        dual-stream loss. Requires the policy to be built with ``attn_mode='flex'``.
+        """
+        self._ensure_frozen_modules()
+        latents, actions, actions_mask, text_emb = self._build_training_streams(batch)
+        return self.training_loss_from_streams(latents, actions, actions_mask, text_emb)
+
+    @torch.no_grad()
+    def _build_training_streams(self, batch):
+        """Build (latents, actions, actions_mask, text_emb) from a LeRobot training batch.
+
+        Camera frames per ``obs_cam_keys`` are expected as a temporal clip ``[B, C, T, H, W]`` (or
+        ``[B, T, C, H, W]``); they are VAE-encoded into ``F = T / temporal_downsample`` latent frames.
+        Actions ``[B, F*action_per_frame, n_used]`` are scattered into the model's ``action_dim`` space.
+        """
+        cfg = self.config
+        device = cfg.device
+        # text embeddings
+        task = batch.get("task")
+        if isinstance(task, str):
+            task = [task]
+        text_emb = self._get_t5_prompt_embeds(list(task), cfg.max_sequence_length)
+
+        # video latents (VAE-encode the camera clips)
+        latents = self._encode_training_latents(batch)
+
+        # actions -> [B, action_dim, F, action_per_frame, 1]
+        act = batch[ACTION].to(device)  # [B, F*apf, n_used]
+        b = act.shape[0]
+        used = cfg.used_action_channel_ids
+        apf, fc = cfg.action_per_frame, cfg.frame_chunk_size
+        act = act[:, : fc * apf].reshape(b, fc, apf, len(used)).permute(0, 3, 1, 2)  # [B, n_used, F, apf]
+        full = act.new_zeros(b, cfg.action_dim, fc, apf)
+        idx = torch.as_tensor(used, device=device)
+        full[:, idx] = act
+        actions = full.unsqueeze(-1).to(self.dtype)  # [B, action_dim, F, apf, 1]
+        mask = torch.zeros(cfg.action_dim, device=device, dtype=self.dtype)
+        mask[idx] = 1.0
+        actions_mask = mask.view(1, -1, 1, 1, 1).expand_as(actions)
+        return latents, actions, actions_mask, text_emb
+
+    @torch.no_grad()
+    def _encode_training_latents(self, batch) -> Tensor:
+        """VAE-encode the per-camera training clips into normalized video latents [B, C, F, h, w]."""
+        vae_device = next(self._vae.parameters()).device
+
+        def _clip(key):
+            x = batch[key].to(vae_device)
+            if x.dim() == 4:  # [B, C, H, W] -> single frame clip
+                x = x.unsqueeze(2)
+            elif x.shape[1] not in (1, 3) and x.shape[2] in (1, 3):  # [B, T, C, H, W] -> [B, C, T, H, W]
+                x = x.permute(0, 2, 1, 3, 4)
+            return x.contiguous()
+
+        def _encode(x, size):
+            b, c, t = x.shape[:3]
+            x = F.interpolate(x.flatten(0, 1).float(), size=size, mode="bilinear", align_corners=False)
+            x = (x.view(b, c, t, *size) * 2.0 - 1.0).to(self.dtype)
+            mu = self._vae.encode(x).latent_dist.mode()  # [B, z_dim, F, h, w]
+            mean = torch.tensor(self._vae.config.latents_mean).view(1, -1, 1, 1, 1).to(mu.device)
+            inv_std = (1.0 / torch.tensor(self._vae.config.latents_std)).view(1, -1, 1, 1, 1).to(mu.device)
+            return ((mu.float() - mean) * inv_std).to(mu)
+
+        keys = self.config.obs_cam_keys
+        if self.config.camera_layout == "robotwin_tshape":
+            h, w = self.config.height, self.config.width
+            head = _encode(_clip(keys[0]), (h, w))
+            left = _encode(_clip(keys[1]), (h // 2, w // 2))
+            right = _encode(_clip(keys[2]), (h // 2, w // 2))
+            return torch.cat([torch.cat([left, right], dim=-1), head], dim=-2).to(self.config.device)
+        per_cam = [_encode(_clip(k), (self.config.height, self.config.width)) for k in keys]
+        return torch.cat(per_cam, dim=-1).to(self.config.device)
+
+    @torch.no_grad()
+    def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
+        """Return one action, refilling the chunk (and feeding back observed keyframes) as needed.
+
+        Mirrors the upstream LIBERO client loop (``evaluation/libero/client.py``): the first obs is
+        the conditioning frame; every observation produced afterwards is buffered as a keyframe and,
+        once the chunk's actions are exhausted, the buffered frames + executed actions are fed back
+        into the KV cache before the next chunk is predicted.
+        """
+        self.eval()
+        self._ensure_frozen_modules()
+        self._maybe_init_prompt(batch)
+
+        if not self._started:
+            # First call: this observation conditions the first chunk (it is *not* a keyframe).
+            self._started = True
+            actions = self.predict_action_chunk(batch)  # [B, chunk_size, n_used]
+            self._action_queue.extend(actions.transpose(0, 1))  # [chunk_size, B, n_used]
+            self._obs_buffer = []
+            self._exec_step = 0
+        else:
+            # This observation is the result of the previously executed action -> a candidate
+            # keyframe. Buffer it on the sub-step boundary the upstream client samples on.
+            if (self._prev_j + 1) % self._keyframe_stride == 0:
+                self._obs_buffer.append(self._extract_raw_obs(batch))
+            if len(self._action_queue) == 0:
+                # All actions for the current chunk have been executed; feed the observed
+                # keyframes + executed actions back and predict the next chunk.
+                actions = self.predict_action_chunk(None)
+                self._action_queue.extend(actions.transpose(0, 1))
+                self._exec_step = 0
+
+        self._prev_j = self._exec_step % self.config.action_per_frame
+        self._exec_step += 1
+        return self._action_queue.popleft()
+
+    @torch.no_grad()
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
+        """Run one autoregressive chunk and return actions ``[B, chunk_size, n_used]`` (normalized)."""
+        self.eval()
+        self._ensure_frozen_modules()
+        self._maybe_init_prompt(batch)
+
+        is_first = self._first_chunk
+        if is_first:
+            init_latent = self._encode_frames([self._extract_raw_obs(batch)])
+            self._init_latent = init_latent
+            self._init_streaming_cache(init_latent)
+            self._obs_buffer = []  # frame 0 (the init obs) conditions the chunk; it is not fed back
+            actions, latents = self._infer(init_latent, frame_st_id=0)
+            self._first_chunk = False
+        else:
+            # Feed the real observed keyframes + the executed actions back into the KV cache.
+            self._compute_kv_cache(self._obs_buffer, self._executed_actions)
+            self._obs_buffer = []
+            actions, latents = self._infer(None, frame_st_id=self._frame_st_id)
+
+        # actions: [B, action_dim, F, action_per_frame, 1] (model-normalized). Keep for KV feedback.
+        self._executed_actions = actions
+
+        if self.config.save_predicted_video:
+            # Match upstream LingBot-VA visualization: collect chunk latents and decode the
+            # concatenated latent sequence once after the rollout finishes.
+            self.last_predicted_frames = None
+            self.last_predicted_latents = latents.detach().to("cpu")
+
+        # On the first chunk, frame 0 is the conditioning frame (already "known"): the upstream
+        # LIBERO client skips it (start_idx=1), so we drop the first frame's actions here.
+        used = self.config.used_action_channel_ids
+        a = actions[:, used]  # [B, n_used, F, action_per_frame, 1]
+        if is_first:
+            a = a[:, :, 1:]  # drop frame 0 -> (F-1) frames of actions
+        a = a.squeeze(-1).flatten(2)  # [B, n_used, n_steps]
+        a = a.transpose(1, 2).contiguous()  # [B, n_steps, n_used]
+        return a.to(torch.float32)
+
+    # Prompt / text encoding
+    def _maybe_init_prompt(self, batch):
+        if self._prompt_embeds is not None or batch is None:
+            return
+        task = batch.get("task")
+        prompt = task[0] if isinstance(task, list | tuple) else task
+        self._prompt = prompt or ""
+        self._prompt_embeds, self._negative_prompt_embeds = self._encode_prompt(self._prompt)
+
+    def _get_t5_prompt_embeds(self, prompt, max_sequence_length):
+        tokenizer = self._frozen["tokenizer"]
+        text_encoder = self._frozen["text_encoder"]
+        device = self.config.device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        prompt = [clean_prompt(u) for u in prompt]
+
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+        text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
+        seq_lens = mask.gt(0).sum(dim=1).long()
+
+        te_device = next(text_encoder.parameters()).device
+        prompt_embeds = text_encoder(text_input_ids.to(te_device), mask.to(te_device)).last_hidden_state
+        prompt_embeds = prompt_embeds.to(dtype=self.dtype, device=device)
+        prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens, strict=False)]
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds],
+            dim=0,
+        )
+        return prompt_embeds.to(device)
+
+    def _encode_prompt(self, prompt):
+        max_len = self.config.max_sequence_length
+        prompt_embeds = self._get_t5_prompt_embeds(prompt, max_len)
+        negative_prompt_embeds = None
+        if self._use_cfg:
+            negative_prompt_embeds = self._get_t5_prompt_embeds("", max_len)
+        return prompt_embeds, negative_prompt_embeds
+
+    # Observation (image) encoding -> normalized video latents
+    def _extract_raw_obs(self, batch) -> dict[str, Tensor]:
+        """Snapshot the configured camera images from a batch (kept raw for later VAE encoding)."""
+        return {k: batch[k].detach() for k in self.config.obs_cam_keys}
+
+    def _camera_frame(self, raw_obs, key, size=None) -> Tensor:
+        """Return a single-frame camera tensor [1, C, 1, H, W] resized + scaled to [-1, 1]."""
+        img = raw_obs[key]
+        if img.dim() == 3:  # [C, H, W]
+            img = img.unsqueeze(0)
+        # LeRobot images arrive as float in [0, 1], shape [B, C, H, W].
+        img = img.to(self.config.device, torch.float32)
+        if self.config.image_hflip:
+            img = torch.flip(img, dims=[-1])  # undo the env processor's horizontal flip
+        if size is None:
+            size = (self.config.height, self.config.width)
+        img = F.interpolate(img, size=size, mode="bilinear", align_corners=False)
+        img = img * 2.0 - 1.0
+        return img.unsqueeze(2).to(self.dtype)  # [1, C, F=1, H, W]
+
+    def _normalize_vae_latent(self, enc_out: Tensor) -> Tensor:
+        """Take the mean of a VAE encoder output and channel-normalize it (matches upstream)."""
+        mu, _logvar = torch.chunk(enc_out, 2, dim=1)
+        latents_mean = torch.tensor(self._vae.config.latents_mean).to(mu.device)
+        latents_std = torch.tensor(self._vae.config.latents_std).to(mu.device)
+        mean = latents_mean.view(1, -1, 1, 1, 1)
+        inv_std = (1.0 / latents_std).view(1, -1, 1, 1, 1)
+        return ((mu.float() - mean) * inv_std).to(mu)
+
+    @torch.no_grad()
+    def _encode_frames(self, raw_frames: list) -> Tensor:
+        """VAE-encode a temporal clip of observed frames and concat the per-camera latents on width.
+
+        ``raw_frames`` is a list of per-frame obs dicts (one per env sub-step). Each configured
+        camera is stacked along the temporal axis into a ``[1, C, F, H, W]`` clip and encoded in a
+        single streaming ``encode_chunk`` call so the VAE temporal downsample (x4) collapses the F
+        input frames into ``F / 4`` latent frames, with the causal ``feat_cache`` carried across
+        chunks (mirrors upstream ``_encode_obs``).
+        """
+        vae_device = next(self._vae.parameters()).device
+        if self.config.camera_layout == "robotwin_tshape":
+            return self._encode_frames_tshape(raw_frames, vae_device)
+        per_cam_videos = []
+        for k in self.config.obs_cam_keys:
+            frames = [self._camera_frame(fb, k) for fb in raw_frames]
+            per_cam_videos.append(torch.cat(frames, dim=2))  # [1, C, F, H, W]
+        videos = torch.cat(per_cam_videos, dim=0)  # [num_cam, C, F, H, W]
+        enc_out = self._streaming_vae.encode_chunk(videos.to(vae_device).to(self.dtype))
+        mu_norm = self._normalize_vae_latent(enc_out)
+        # Concatenate the per-camera latents along width.
+        video_latent = torch.cat(mu_norm.split(1, dim=0), dim=-1)
+        return video_latent.to(self.config.device)
+
+    @torch.no_grad()
+    def _encode_frames_tshape(self, raw_frames: list, vae_device) -> Tensor:
+        """RoboTwin T-shape latent assembly: full-res head + half-res wrists (second streaming VAE).
+
+        The two wrist latents are concatenated on width and stacked (on the height axis) on top of
+        the head latent, mirroring upstream ``_encode_obs`` for ``env_type='robotwin_tshape'``.
+        """
+        cfg = self.config
+        h, w = cfg.height, cfg.width
+        head_key, left_key, right_key = cfg.obs_cam_keys[0], cfg.obs_cam_keys[1], cfg.obs_cam_keys[2]
+        head = torch.cat([self._camera_frame(fb, head_key, size=(h, w)) for fb in raw_frames], dim=2)
+        left = torch.cat(
+            [self._camera_frame(fb, left_key, size=(h // 2, w // 2)) for fb in raw_frames], dim=2
+        )
+        right = torch.cat(
+            [self._camera_frame(fb, right_key, size=(h // 2, w // 2)) for fb in raw_frames], dim=2
+        )
+        wrists = torch.cat([left, right], dim=0)  # [2, C, F, H/2, W/2]
+        enc_high = self._streaming_vae.encode_chunk(head.to(vae_device).to(self.dtype))
+        enc_lr = self._frozen["streaming_vae_half"].encode_chunk(wrists.to(vae_device).to(self.dtype))
+        # wrists side-by-side on width, then stacked on top of the head latent on the height axis.
+        enc_out = torch.cat([torch.cat(enc_lr.split(1, dim=0), dim=-1), enc_high], dim=-2)
+        video_latent = self._normalize_vae_latent(enc_out)
+        return video_latent.to(self.config.device)
+
+    # KV cache management
+    @property
+    def _latent_hw(self):
+        if self.config.camera_layout == "robotwin_tshape":
+            # head (full) on the bottom, two half-res wrists side-by-side on top -> 1.5x height.
+            return ((self.config.height // 16) * 3) // 2, self.config.width // 16
+        h = self.config.height // 16
+        w = (self.config.width // 16) * len(self.config.obs_cam_keys)
+        return h, w
+
+    def _init_streaming_cache(self, init_latent):
+        cfg = self.config
+        latent_h, latent_w = self._latent_hw
+        p = cfg.patch_size
+        latent_token_per_chunk = (cfg.frame_chunk_size * latent_h * latent_w) // (p[0] * p[1] * p[2])
+        action_token_per_chunk = cfg.frame_chunk_size * cfg.action_per_frame
+        self.transformer.create_empty_cache(
+            "pos",
+            cfg.attn_window,
+            latent_token_per_chunk,
+            action_token_per_chunk,
+            device=self.config.device,
+            dtype=self.dtype,
+            batch_size=2 if self._use_cfg else 1,
+        )
+        self._cache_initialised = True
+
+    def _repeat_input_for_cfg(self, input_dict):
+        if self._use_cfg:
+            input_dict["noisy_latents"] = input_dict["noisy_latents"].repeat(2, 1, 1, 1, 1)
+            input_dict["text_emb"] = torch.cat(
+                [
+                    self._prompt_embeds.to(self.dtype).clone(),
+                    self._negative_prompt_embeds.to(self.dtype).clone(),
+                ],
+                dim=0,
+            )
+            input_dict["grid_id"] = input_dict["grid_id"][None].repeat(2, 1, 1)
+            input_dict["timesteps"] = input_dict["timesteps"][None].repeat(2, 1)
+        else:
+            input_dict["grid_id"] = input_dict["grid_id"][None]
+            input_dict["timesteps"] = input_dict["timesteps"][None]
+        return input_dict
+
+    def _prepare_latent_input(
+        self,
+        latent_model_input,
+        action_model_input,
+        latent_t=0,
+        action_t=0,
+        latent_cond=None,
+        action_cond=None,
+        frame_st_id=0,
+    ):
+        cfg = self.config
+        device = self.config.device
+        p = cfg.patch_size
+        out = {}
+        if latent_model_input is not None:
+            out["latent_res_lst"] = {
+                "noisy_latents": latent_model_input,
+                "timesteps": torch.ones([latent_model_input.shape[2]], dtype=torch.float32, device=device)
+                * latent_t,
+                "grid_id": get_mesh_id(
+                    latent_model_input.shape[-3] // p[0],
+                    latent_model_input.shape[-2] // p[1],
+                    latent_model_input.shape[-1] // p[2],
+                    0,
+                    1,
+                    frame_st_id,
+                ).to(device),
+                "text_emb": self._prompt_embeds.to(self.dtype).clone(),
+            }
+            if latent_cond is not None:
+                out["latent_res_lst"]["noisy_latents"][:, :, 0:1] = latent_cond[:, :, 0:1]
+                out["latent_res_lst"]["timesteps"][0:1] *= 0
+        if action_model_input is not None:
+            out["action_res_lst"] = {
+                "noisy_latents": action_model_input,
+                "timesteps": torch.ones([action_model_input.shape[2]], dtype=torch.float32, device=device)
+                * action_t,
+                "grid_id": get_mesh_id(
+                    action_model_input.shape[-3],
+                    action_model_input.shape[-2],
+                    action_model_input.shape[-1],
+                    1,
+                    1,
+                    frame_st_id,
+                    action=True,
+                ).to(device),
+                "text_emb": self._prompt_embeds.to(self.dtype).clone(),
+            }
+            if action_cond is not None:
+                out["action_res_lst"]["noisy_latents"][:, :, 0:1] = action_cond[:, :, 0:1]
+                out["action_res_lst"]["timesteps"][0:1] *= 0
+            out["action_res_lst"]["noisy_latents"][:, ~self._action_mask] *= 0
+        return out
+
+    @property
+    def _action_mask(self):
+        mask = torch.zeros([self.config.action_dim], dtype=torch.bool)
+        mask[self.config.used_action_channel_ids] = True
+        return mask
+
+    # Action conditioning (executed action history) (de)normalization
+    def _preprocess_action_state(self, action_norm: Tensor) -> Tensor:
+        """Build the action-conditioning tensor from the already-normalized executed actions.
+
+        ``action_norm`` is the model-space action chunk ``[B, action_dim, F, action_per_frame, 1]``.
+        Upstream re-derives the conditioning from the raw executed action via quantile norm; here
+        the executed actions are already in the model-normalized space, so we pass them through.
+        """
+        return action_norm.to(self.config.device, self.dtype)
+
+    def _compute_kv_cache(self, obs_buffer, executed_actions):
+        """Feed real observed keyframes + executed actions back into the KV cache."""
+        if not obs_buffer or executed_actions is None:
+            return
+        self.transformer.clear_pred_cache("pos")
+        # Encode the buffered keyframe clip in one streaming call (carries the causal VAE cache).
+        latent_model_input = self._encode_frames(obs_buffer)
+        # On the first feedback, prepend the init latent so the latent/action frame counts align
+        # (upstream prepends ``init_latent`` to the observed keyframes when frame_st_id == 0).
+        if self._frame_st_id == 0 and getattr(self, "_init_latent", None) is not None:
+            latent_model_input = torch.cat([self._init_latent, latent_model_input], dim=2)
+        action_model_input = self._preprocess_action_state(executed_actions)
+        action_model_input = action_model_input.to(latent_model_input)
+        input_dict = self._prepare_latent_input(
+            latent_model_input, action_model_input, frame_st_id=self._frame_st_id
+        )
+        with torch.no_grad():
+            self.transformer(
+                self._repeat_input_for_cfg(input_dict["latent_res_lst"]),
+                update_cache=2,
+                cache_name="pos",
+                action_mode=False,
+            )
+            self.transformer(
+                self._repeat_input_for_cfg(input_dict["action_res_lst"]),
+                update_cache=2,
+                cache_name="pos",
+                action_mode=True,
+            )
+        self._frame_st_id += latent_model_input.shape[2]
+
+    # The core dual-stream denoising loop (one chunk)
+    @torch.no_grad()
+    def _infer(self, init_latent, frame_st_id=0):
+        cfg = self.config
+        device = self.config.device
+        latent_h, latent_w = self._latent_hw
+        frame_chunk_size = cfg.frame_chunk_size
+
+        latents = torch.randn(1, 48, frame_chunk_size, latent_h, latent_w, device=device, dtype=self.dtype)
+        actions = torch.randn(
+            1, cfg.action_dim, frame_chunk_size, cfg.action_per_frame, 1, device=device, dtype=self.dtype
+        )
+
+        self._scheduler.set_timesteps(cfg.num_inference_steps)
+        self._action_scheduler.set_timesteps(cfg.action_num_inference_steps)
+        timesteps = F.pad(self._scheduler.timesteps, (0, 1), mode="constant", value=0)
+        if cfg.video_exec_step != -1:
+            timesteps = timesteps[: cfg.video_exec_step]
+        action_timesteps = F.pad(self._action_scheduler.timesteps, (0, 1), mode="constant", value=0)
+
+        # 1. Video-latent denoising loop
+        for i, t in enumerate(timesteps):
+            last_step = i == len(timesteps) - 1
+            latent_cond = (
+                init_latent[:, :, 0:1].to(self.dtype)
+                if frame_st_id == 0 and init_latent is not None
+                else None
+            )
+            input_dict = self._prepare_latent_input(
+                latents, None, t, t, latent_cond, None, frame_st_id=frame_st_id
+            )
+            video_noise_pred = self.transformer(
+                self._repeat_input_for_cfg(input_dict["latent_res_lst"]),
+                update_cache=1 if last_step else 0,
+                cache_name="pos",
+                action_mode=False,
+            )
+            if not last_step or cfg.video_exec_step != -1:
+                video_noise_pred = data_seq_to_patch(
+                    cfg.patch_size,
+                    video_noise_pred,
+                    frame_chunk_size,
+                    latent_h,
+                    latent_w,
+                    batch_size=2 if self._use_cfg else 1,
+                )
+                if cfg.guidance_scale > 1:
+                    video_noise_pred = video_noise_pred[1:] + cfg.guidance_scale * (
+                        video_noise_pred[:1] - video_noise_pred[1:]
+                    )
+                else:
+                    video_noise_pred = video_noise_pred[:1]
+                latents = self._scheduler.step(video_noise_pred, t, latents, return_dict=False)
+            if frame_st_id == 0 and latent_cond is not None:
+                latents[:, :, 0:1] = latent_cond
+
+        # 2. Action denoising loop
+        for i, t in enumerate(action_timesteps):
+            last_step = i == len(action_timesteps) - 1
+            action_cond = (
+                torch.zeros([1, cfg.action_dim, 1, cfg.action_per_frame, 1], device=device, dtype=self.dtype)
+                if frame_st_id == 0
+                else None
+            )
+            input_dict = self._prepare_latent_input(
+                None, actions, t, t, None, action_cond, frame_st_id=frame_st_id
+            )
+            action_noise_pred = self.transformer(
+                self._repeat_input_for_cfg(input_dict["action_res_lst"]),
+                update_cache=1 if last_step else 0,
+                cache_name="pos",
+                action_mode=True,
+            )
+            if not last_step:
+                action_noise_pred = rearrange(action_noise_pred, "b (f n) c -> b c f n 1", f=frame_chunk_size)
+                if cfg.action_guidance_scale > 1:
+                    action_noise_pred = action_noise_pred[1:] + cfg.action_guidance_scale * (
+                        action_noise_pred[:1] - action_noise_pred[1:]
+                    )
+                else:
+                    action_noise_pred = action_noise_pred[:1]
+                actions = self._action_scheduler.step(action_noise_pred, t, actions, return_dict=False)
+            if frame_st_id == 0 and action_cond is not None:
+                actions[:, :, 0:1] = action_cond
+
+        actions[:, ~self._action_mask] *= 0
+        return actions, latents
+
+    # Predicted-video decoding (opt-in)
+    @torch.no_grad()
+    def decode_predicted_latents(self, latents) -> Tensor:
+        """Decode a concatenated predicted-latent sequence into ``[T, H, W, 3]`` uint8 frames."""
+        return self._decode_predicted_video(latents)
+
+    @torch.no_grad()
+    def _decode_predicted_video(self, latents) -> Tensor:
+        """VAE-decode predicted latents into a uint8 frame stack ``[T, H, W, 3]`` on CPU."""
+        vae = self._vae
+        z_dim = vae.config.z_dim
+        vae_device = next(vae.parameters()).device
+        latents = latents.to(device=vae_device, dtype=vae.dtype)
+        latents = denormalize_latents(latents, vae.config.latents_mean, vae.config.latents_std, z_dim)
+        video = vae.decode(latents, return_dict=False)[0]  # [B, C, F, H, W] in [-1, 1]
+        video = (video.float().clamp(-1, 1) + 1.0) / 2.0
+        video = (video[0].permute(1, 2, 3, 0) * 255.0).round().to(torch.uint8)  # [F, H, W, C]
+        return video.cpu()

src/lerobot/policies/lingbot_va/processor_lingbot_va.py

@@ -0,0 +1,87 @@
+# 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.
+
+"""Pre/post-processor pipelines for the LingBot-VA policy.
+
+The preprocessor passes inputs through (IDENTITY) and the postprocessor maps the policy's
+``[-1, 1]`` actions back to physical units with the built-in ``UnnormalizerProcessorStep``
+(QUANTILES) using per-channel q01/q99 restored from the checkpoint.
+"""
+
+from typing import Any
+
+import torch
+
+from lerobot.configs.types import FeatureType, NormalizationMode
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    ProcessorStep,
+    RenameObservationsProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+from lerobot.utils.constants import (
+    POLICY_POSTPROCESSOR_DEFAULT_NAME,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
+
+from .configuration_lingbot_va import LingBotVAConfig
+
+
+def make_lingbot_va_pre_post_processors(
+    config: LingBotVAConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """Build the pre/post processor pipelines for LingBot-VA."""
+
+    input_steps: list[ProcessorStep] = [
+        RenameObservationsProcessorStep(rename_map={}),
+        AddBatchDimensionProcessorStep(),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+        DeviceProcessorStep(device=config.device),
+    ]
+
+    # Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint.
+    output_steps: list[ProcessorStep] = [
+        UnnormalizerProcessorStep(
+            features=config.output_features,
+            norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES},
+            stats=dataset_stats,
+        ),
+        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,
+        ),
+    )

src/lerobot/policies/molmoact2/configuration_molmoact2.py

@@ -73,12 +73,34 @@ class MolmoAct2Config(PreTrainedConfig):
     num_inference_steps: int | None = None
     mask_action_dim_padding: bool = True
     enable_inference_cuda_graph: bool = True
+
+    # Language conditioning (e.g. RECAP advantage). When set, RenderMessagesStep
+    # resolves language_persistent rows via the recipe YAML. Same mechanism as PI05.
+    recipe_path: str | None = None
+
+    # Inference-time advantage indicator (e.g. "Advantage: positive. ").
+    # Used during rollout when no language_persistent data is available.
+    # Placed after the user prompt, before action tokens.
+    advantage_prefix: str = ""
+
+    # Classifier-Free Guidance (CFG) scale for inference (RECAP Eq. 13).
+    # 1.0 = no guidance. >1.0 = dual-path: v = v_uncond + beta * (v_cond - v_uncond)
+    cfg_beta: float = 1.0
     # MolmoAct2-local eval option. When enabled, stochastic continuous action
     # generation uses a rollout-local generator derived from eval_seed.
     per_episode_seed: bool = False
     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
@@ -123,6 +145,10 @@ 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}. "

src/lerobot/policies/molmoact2/processor_molmoact2.py

@@ -359,6 +359,7 @@ def _build_robot_text(
     add_setup_tokens: bool,
     add_control_tokens: bool,
     num_images: int,
+    advantage: str = "",
 ) -> str:
     setup_text = _wrap_setup_text(setup_type, add_setup_tokens=add_setup_tokens)
     control_text = _wrap_control_text(control_mode, add_control_tokens=add_control_tokens)
@@ -375,7 +376,10 @@ def _build_robot_text(
         image_prefix = "<|image|>"
     else:
         image_prefix = "".join(f"Image {idx + 1}<|image|>" for idx in range(num_images))
-    return f"{image_prefix}<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n{ACTION_OUTPUT_TOKEN}"
+    # Per RECAP paper (Section V-B): advantage indicator goes after context,
+    # before actions, so only action log-likelihoods are affected.
+    advantage_clause = f"Advantage: {advantage}. " if advantage else ""
+    return f"{image_prefix}<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n{advantage_clause}{ACTION_OUTPUT_TOKEN}"
 
 
 def _as_text_list(value: Any, batch_size: int) -> list[str]:
@@ -695,6 +699,39 @@ class MolmoAct2ClampNormalizedProcessorStep(ProcessorStep):
         return features
 
 
+@ProcessorStepRegistry.register(name="molmoact2_normalize_task")
+@dataclass
+class MolmoAct2NormalizeTaskStep(ProcessorStep):
+    """Normalize the task text in complementary_data before recipe rendering.
+
+    Ensures ${task} in recipe templates gets the same normalized form that
+    MolmoAct2PackInputsProcessorStep would produce, so training prompts
+    match inference prompts.
+    """
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        complementary = transition.get(TransitionKey.COMPLEMENTARY_DATA)
+        if not isinstance(complementary, dict):
+            return transition
+        task = complementary.get("task")
+        if task is None:
+            return transition
+
+        transition = transition.copy()
+        complementary = dict(complementary)
+        if isinstance(task, str):
+            complementary["task"] = _normalize_question_text(task)
+        elif isinstance(task, list):
+            complementary["task"] = [_normalize_question_text(t) for t in task]
+        transition[TransitionKey.COMPLEMENTARY_DATA] = complementary
+        return transition
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
 @ProcessorStepRegistry.register(name="molmoact2_pack_inputs")
 @dataclass
 class MolmoAct2PackInputsProcessorStep(ProcessorStep):
@@ -715,6 +752,8 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep):
     chunk_size: int = 30
     max_action_dim: int = 32
     env_action_dim: int | None = None
+    # RECAP: advantage indicator for inference (e.g. "Advantage: positive. ")
+    advantage_prefix: str = ""
 
     def __post_init__(self) -> None:
         require_package("transformers", extra="molmoact2")
@@ -757,6 +796,7 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep):
             "chunk_size": self.chunk_size,
             "max_action_dim": self.max_action_dim,
             "env_action_dim": self.env_action_dim,
+            "advantage_prefix": self.advantage_prefix,
         }
 
     def _resolve_max_sequence_length(
@@ -919,8 +959,40 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep):
         if task_source is None:
             task_source = complementary.get("language_instruction")
         tasks = _as_text_list(task_source, batch_size)
-        if self.normalize_language:
-            tasks = [_normalize_question_text(task) for task in tasks]
+
+        # Resolve the advantage indicator. Per RECAP paper (Section V-B), it goes
+        # after all context but before actions — handled by _build_robot_text.
+        # Source priority: recipe-rendered "advantage" key > config advantage_prefix.
+        advantages: list[str] = []
+        recipe_rendered = "base_task" in complementary
+        if recipe_rendered:
+            # Recipe rendered the task as "<task> Advantage: <value>".
+            # Extract the advantage value and restore the clean task.
+            clean_tasks: list[str] = []
+            for t in tasks:
+                if " Advantage: " in t:
+                    split_idx = t.rindex(" Advantage: ")
+                    clean_task = t[:split_idx]
+                    adv = t[split_idx + len(" Advantage: ") :]
+                    advantages.append(adv)
+                    clean_tasks.append(clean_task)
+                else:
+                    advantages.append("")
+                    clean_tasks.append(t)
+            tasks = clean_tasks
+        else:
+            if self.normalize_language:
+                tasks = [_normalize_question_text(task) for task in tasks]
+            if self.advantage_prefix:
+                # Extract just the value from prefix like "Advantage: positive. "
+                prefix = self.advantage_prefix.strip()
+                if prefix.startswith("Advantage:"):
+                    adv_val = prefix[len("Advantage:") :].strip().rstrip(".")
+                else:
+                    adv_val = prefix
+                advantages = [adv_val] * batch_size
+            else:
+                advantages = [""] * batch_size
         complementary["task"] = tasks
 
         action_padded = None
@@ -953,6 +1025,7 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep):
                 add_setup_tokens=self.add_setup_tokens,
                 add_control_tokens=self.add_control_tokens,
                 num_images=len(images),
+                advantage=advantages[batch_idx],
             )
             prompt_texts.append(prompt)
             if build_action_labels:
@@ -1005,6 +1078,93 @@ 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):
@@ -1067,34 +1227,58 @@ 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,
             stats=masked_dataset_stats,
         ),
         MolmoAct2ClampNormalizedProcessorStep(normalization_masks=normalization_masks),
-        MolmoAct2PackInputsProcessorStep(
-            checkpoint_path=config.checkpoint_path,
-            checkpoint_revision=config.checkpoint_revision,
-            checkpoint_force_download=config.checkpoint_force_download,
-            action_mode=config.action_mode,
-            discrete_action_tokenizer=config.discrete_action_tokenizer,
-            image_keys=image_keys,
-            allow_image_key_fallback=not bool(config.image_keys),
-            setup_type=setup_type,
-            control_mode=control_mode,
-            normalize_language=config.normalize_language,
-            add_setup_tokens=config.add_setup_tokens,
-            add_control_tokens=config.add_control_tokens,
-            num_state_tokens=config.num_state_tokens,
-            max_sequence_length=config.max_sequence_length,
-            chunk_size=chunk_size,
-            max_action_dim=config.expected_max_action_dim,
-            env_action_dim=env_action_dim,
-        ),
-        DeviceProcessorStep(device=config.device),
     ]
 
+    # Insert language rendering steps when a recipe is configured (e.g. RECAP advantage)
+    if config.recipe_path is not None:
+        from lerobot.configs.recipe import load_recipe
+        from lerobot.processor.render_messages_processor import RenderMessagesStep
+        from lerobot.processor.rendered_messages_to_task import RenderedMessagesToTaskStep
+
+        recipe = load_recipe(config.recipe_path)
+        # Normalize task text before recipe uses ${task}, ensuring consistency
+        # between training (recipe-rendered) and inference (advantage_prefix).
+        if config.normalize_language:
+            input_steps.append(MolmoAct2NormalizeTaskStep())
+        input_steps.append(RenderMessagesStep(recipe=recipe))
+        input_steps.append(RenderedMessagesToTaskStep())
+
+    input_steps.extend(
+        [
+            MolmoAct2PackInputsProcessorStep(
+                checkpoint_path=config.checkpoint_path,
+                checkpoint_revision=config.checkpoint_revision,
+                checkpoint_force_download=config.checkpoint_force_download,
+                action_mode=config.action_mode,
+                discrete_action_tokenizer=config.discrete_action_tokenizer,
+                image_keys=image_keys,
+                allow_image_key_fallback=not bool(config.image_keys),
+                setup_type=setup_type,
+                control_mode=control_mode,
+                normalize_language=config.normalize_language,
+                add_setup_tokens=config.add_setup_tokens,
+                add_control_tokens=config.add_control_tokens,
+                num_state_tokens=config.num_state_tokens,
+                max_sequence_length=config.max_sequence_length,
+                chunk_size=chunk_size,
+                max_action_dim=config.expected_max_action_dim,
+                env_action_dim=env_action_dim,
+                advantage_prefix=config.advantage_prefix,
+            ),
+            DeviceProcessorStep(device=config.device),
+        ]
+    )
+
     output_steps: list[ProcessorStep] = [
         MolmoAct2ClampActionProcessorStep(),
         MolmoAct2MaskedUnnormalizerProcessorStep(
@@ -1102,6 +1286,10 @@ 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"),
     ]
 

src/lerobot/policies/pi05/configuration_pi05.py

@@ -87,6 +87,17 @@ class PI05Config(PreTrainedConfig):
     freeze_vision_encoder: bool = False  # Freeze only the vision encoder
     train_expert_only: bool = False  # Freeze entire VLM, train only action expert and projections
 
+    # Language conditioning (e.g. RECAP advantage). When set, RenderMessagesStep
+    # is inserted into the preprocessor to resolve language_persistent rows via
+    # the recipe YAML before prompt construction.
+    recipe_path: str | None = None
+
+    # Classifier-Free Guidance (CFG) scale for inference (Eq. 13 in RECAP paper).
+    # 1.0 = no guidance (default). >1.0 enables dual-path denoising where:
+    #   v = v_uncond + cfg_beta * (v_cond - v_uncond)
+    # VLM runs twice (cond + uncond prompts), action expert runs 2x per step.
+    cfg_beta: float = 1.0
+
     # Optimizer settings: see openpi `AdamW`
     optimizer_lr: float = 2.5e-5  # see openpi `CosineDecaySchedule: peak_lr`
     optimizer_betas: tuple[float, float] = (0.9, 0.95)

src/lerobot/policies/pi05/modeling_pi05.py

@@ -52,6 +52,8 @@ from lerobot.utils.constants import (
     ACTION,
     OBS_LANGUAGE_ATTENTION_MASK,
     OBS_LANGUAGE_TOKENS,
+    OBS_LANGUAGE_UNCOND_ATTENTION_MASK,
+    OBS_LANGUAGE_UNCOND_TOKENS,
     OPENPI_ATTENTION_MASK_VALUE,
 )
 
@@ -148,6 +150,20 @@ def clone_past_key_values(past_key_values):
     )
 
 
+def cat_past_key_values(kv_a, kv_b):
+    """Concatenate two DynamicCaches along the batch dimension for batched CFG."""
+    return DynamicCache(
+        tuple(
+            (
+                torch.cat([ka, kb], dim=0),
+                torch.cat([va, vb], dim=0),
+                sw_a,
+            )
+            for (ka, va, sw_a), (kb, vb, _sw_b) in zip(kv_a, kv_b, strict=True)
+        )
+    )
+
+
 def pad_vector(vector, new_dim):
     """Pad the last dimension of a vector to new_dim with zeros.
 
@@ -797,9 +813,17 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         masks,
         noise=None,
         num_steps=None,
+        uncond_tokens=None,
+        uncond_masks=None,
         **kwargs: Unpack[ActionSelectKwargs],
     ) -> Tensor:
-        """Do a full inference forward and compute the action."""
+        """Do a full inference forward and compute the action.
+
+        When cfg_beta > 1.0 and uncond_tokens/uncond_masks are provided, performs
+        Classifier-Free Guidance: VLM runs twice (conditioned + unconditional), action
+        expert runs twice per denoising step, and velocities are interpolated via
+        v = v_uncond + cfg_beta * (v_cond - v_uncond).
+        """
         if num_steps is None:
             num_steps = self.config.num_inference_steps
 
@@ -815,6 +839,9 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             )  # Use config max_action_dim for internal processing
             noise = self.sample_noise(actions_shape, device)
 
+        cfg_enabled = self.config.cfg_beta > 1.0 and uncond_tokens is not None and uncond_masks is not None
+
+        # Prefill VLM for conditioned prompt
         prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, tokens, masks)
         prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
         prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
@@ -830,6 +857,23 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             use_cache=True,
         )
 
+        # Prefill VLM for unconditional prompt (CFG)
+        if cfg_enabled:
+            uncond_prefix_embs, uncond_prefix_pad_masks, uncond_prefix_att_masks = self.embed_prefix(
+                images, img_masks, uncond_tokens, uncond_masks
+            )
+            uncond_prefix_att_2d_masks = make_att_2d_masks(uncond_prefix_pad_masks, uncond_prefix_att_masks)
+            uncond_prefix_position_ids = torch.cumsum(uncond_prefix_pad_masks, dim=1) - 1
+            uncond_prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(uncond_prefix_att_2d_masks)
+
+            _, uncond_past_key_values = self.paligemma_with_expert.forward(
+                attention_mask=uncond_prefix_att_2d_masks_4d,
+                position_ids=uncond_prefix_position_ids,
+                past_key_values=None,
+                inputs_embeds=[uncond_prefix_embs, None],
+                use_cache=True,
+            )
+
         dt = -1.0 / num_steps
 
         x_t = noise
@@ -838,6 +882,15 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
 
             def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
+                if cfg_enabled:
+                    return self.denoise_step_cfg_batched(
+                        cond_prefix_pad_masks=prefix_pad_masks,
+                        cond_past_key_values=past_key_values,
+                        uncond_prefix_pad_masks=uncond_prefix_pad_masks,
+                        uncond_past_key_values=uncond_past_key_values,
+                        x_t=input_x_t,
+                        timestep=current_timestep,
+                    )
                 return self.denoise_step(
                     prefix_pad_masks=prefix_pad_masks,
                     past_key_values=past_key_values,
@@ -907,6 +960,80 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         suffix_out = suffix_out.to(dtype=torch.float32)
         return self.action_out_proj(suffix_out)
 
+    def denoise_step_cfg_batched(
+        self,
+        cond_prefix_pad_masks,
+        cond_past_key_values,
+        uncond_prefix_pad_masks,
+        uncond_past_key_values,
+        x_t,
+        timestep,
+    ):
+        """Batched CFG denoising: runs cond + uncond in a single forward pass.
+
+        Concatenates cond and uncond inputs along the batch dimension, runs one
+        action expert forward (2x batch), then splits and applies CFG interpolation.
+        This is ~1.5x faster than two sequential denoise_step calls due to better
+        GPU utilization (inspired by Qwen2.5-Omni DiT / diffusers batched CFG).
+        """
+        # Embed suffix once (same x_t and timestep for both branches)
+        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, timestep)
+
+        bsize = cond_prefix_pad_masks.shape[0]
+        suffix_len = suffix_pad_masks.shape[1]
+        cond_prefix_len = cond_prefix_pad_masks.shape[1]
+        uncond_prefix_len = uncond_prefix_pad_masks.shape[1]
+
+        # Build attention masks for cond branch
+        cond_prefix_2d = cond_prefix_pad_masks[:, None, :].expand(bsize, suffix_len, cond_prefix_len)
+        cond_suffix_att_2d = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        cond_full_att = torch.cat([cond_prefix_2d, cond_suffix_att_2d], dim=2)
+        cond_prefix_offsets = torch.sum(cond_prefix_pad_masks, dim=-1)[:, None]
+        cond_position_ids = cond_prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
+
+        # Build attention masks for uncond branch
+        uncond_prefix_2d = uncond_prefix_pad_masks[:, None, :].expand(bsize, suffix_len, uncond_prefix_len)
+        uncond_suffix_att_2d = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        uncond_full_att = torch.cat([uncond_prefix_2d, uncond_suffix_att_2d], dim=2)
+        uncond_prefix_offsets = torch.sum(uncond_prefix_pad_masks, dim=-1)[:, None]
+        uncond_position_ids = uncond_prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
+
+        # Concatenate on batch dim: [cond_batch; uncond_batch]
+        batched_full_att = torch.cat([cond_full_att, uncond_full_att], dim=0)
+        batched_full_att_4d = self._prepare_attention_masks_4d(batched_full_att)
+        batched_position_ids = torch.cat([cond_position_ids, uncond_position_ids], dim=0)
+        batched_suffix_embs = torch.cat([suffix_embs, suffix_embs], dim=0)
+        batched_adarms_cond = torch.cat([adarms_cond, adarms_cond], dim=0)
+
+        # Concatenate KV caches on batch dim
+        batched_past_kv = cat_past_key_values(
+            clone_past_key_values(cond_past_key_values),
+            clone_past_key_values(uncond_past_key_values),
+        )
+
+        self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"  # noqa: SLF001
+
+        # Single forward pass for both branches
+        outputs_embeds, _ = self.paligemma_with_expert.forward(
+            attention_mask=batched_full_att_4d,
+            position_ids=batched_position_ids,
+            past_key_values=batched_past_kv,
+            inputs_embeds=[None, batched_suffix_embs],
+            use_cache=False,
+            adarms_cond=[None, batched_adarms_cond],
+        )
+
+        suffix_out = outputs_embeds[1]
+        suffix_out = suffix_out[:, -self.config.chunk_size :]
+        suffix_out = suffix_out.to(dtype=torch.float32)
+        v_all = self.action_out_proj(suffix_out)
+
+        # Split: first half = cond, second half = uncond
+        v_cond, v_uncond = v_all.chunk(2, dim=0)
+
+        # CFG interpolation: v = v_uncond + beta * (v_cond - v_uncond)
+        return v_uncond + self.config.cfg_beta * (v_cond - v_uncond)
+
 
 class PI05Policy(PreTrainedPolicy):
     """PI05 Policy for LeRobot."""
@@ -1243,8 +1370,20 @@ class PI05Policy(PreTrainedPolicy):
         images, img_masks = self._preprocess_images(batch)
         tokens, masks = batch[f"{OBS_LANGUAGE_TOKENS}"], batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
 
+        # CFG: pass unconditional tokens if available
+        uncond_tokens = batch.get(f"{OBS_LANGUAGE_UNCOND_TOKENS}")
+        uncond_masks = batch.get(f"{OBS_LANGUAGE_UNCOND_ATTENTION_MASK}")
+
         # Sample actions using the model (pass through RTC kwargs, no separate state needed for PI05)
-        actions = self.model.sample_actions(images, img_masks, tokens, masks, **kwargs)
+        actions = self.model.sample_actions(
+            images,
+            img_masks,
+            tokens,
+            masks,
+            uncond_tokens=uncond_tokens,
+            uncond_masks=uncond_masks,
+            **kwargs,
+        )
 
         # Unpad actions to actual action dimension
         original_action_dim = self.config.output_features[ACTION].shape[0]

src/lerobot/policies/pi05/processor_pi05.py

@@ -40,6 +40,8 @@ from lerobot.processor import (
 )
 from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
+    OBS_LANGUAGE_UNCOND_ATTENTION_MASK,
+    OBS_LANGUAGE_UNCOND_TOKENS,
     OBS_STATE,
     POLICY_POSTPROCESSOR_DEFAULT_NAME,
     POLICY_PREPROCESSOR_DEFAULT_NAME,
@@ -57,6 +59,7 @@ class Pi05PrepareStateTokenizerProcessorStep(ProcessorStep):
 
     max_state_dim: int = 32
     task_key: str = "task"
+    cfg_enabled: bool = False
 
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         transition = transition.copy()
@@ -84,8 +87,25 @@ class Pi05PrepareStateTokenizerProcessorStep(ProcessorStep):
             full_prompts.append(full_prompt)
 
         transition[TransitionKey.COMPLEMENTARY_DATA][self.task_key] = full_prompts
-        # Normalize state to [-1, 1] range if needed (assuming it's already normalized by normalizer processor step!!)
-        # Discretize into 256 bins (see openpi `PaligemmaTokenizer.tokenize()`)
+
+        # Build unconditional prompts for CFG (same state but original task without advantage)
+        if self.cfg_enabled:
+            base_tasks = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}).get("base_task")
+            if base_tasks is None:
+                base_tasks = tasks
+
+            if isinstance(base_tasks, str):
+                base_tasks = [base_tasks] * len(tasks)
+
+            uncond_prompts = []
+            for i, base_task in enumerate(base_tasks):
+                cleaned_text = base_task.strip().replace("_", " ").replace("\n", " ")
+                state_str = " ".join(map(str, discretized_states[i]))
+                uncond_prompt = f"Task: {cleaned_text}, State: {state_str};\nAction: "
+                uncond_prompts.append(uncond_prompt)
+
+            transition[TransitionKey.COMPLEMENTARY_DATA]["uncond_task"] = uncond_prompts
+
         return transition
 
     def transform_features(
@@ -111,9 +131,10 @@ def make_pi05_pre_post_processors(
     1. Renaming features to match pretrained configurations.
     2. Normalizing input and output features based on dataset statistics.
     3. Adding a batch dimension.
-    4. Appending a newline character to the task description for tokenizer compatibility.
-    5. Tokenizing the text prompt using the PaliGemma tokenizer.
-    6. Moving all data to the specified device.
+    4. (Optional) Rendering language annotations via recipe YAML.
+    5. (Optional) Flattening rendered messages into the task string.
+    6. Tokenizing the text prompt using the PaliGemma tokenizer.
+    7. Moving all data to the specified device.
 
     The post-processing pipeline handles the model's output by:
     1. Moving data to the CPU.
@@ -122,8 +143,6 @@ def make_pi05_pre_post_processors(
     Args:
         config: The configuration object for the PI0 policy.
         dataset_stats: A dictionary of statistics for normalization.
-        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
-        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
 
     Returns:
         A tuple containing the configured pre-processor and post-processor pipelines.
@@ -147,16 +166,51 @@ def make_pi05_pre_post_processors(
             norm_map=config.normalization_mapping,
             stats=dataset_stats,
         ),
-        Pi05PrepareStateTokenizerProcessorStep(max_state_dim=config.max_state_dim),
-        TokenizerProcessorStep(
-            tokenizer_name="google/paligemma-3b-pt-224",
-            max_length=config.tokenizer_max_length,
-            padding_side="right",
-            padding="max_length",
-        ),
-        DeviceProcessorStep(device=config.device),
     ]
 
+    # Insert language rendering steps when a recipe is configured (e.g. RECAP advantage)
+    if config.recipe_path is not None:
+        from lerobot.configs.recipe import load_recipe
+        from lerobot.processor.render_messages_processor import RenderMessagesStep
+        from lerobot.processor.rendered_messages_to_task import RenderedMessagesToTaskStep
+
+        recipe = load_recipe(config.recipe_path)
+        input_steps.append(RenderMessagesStep(recipe=recipe))
+        input_steps.append(RenderedMessagesToTaskStep())
+
+    cfg_enabled = config.cfg_beta > 1.0
+
+    input_steps.extend(
+        [
+            Pi05PrepareStateTokenizerProcessorStep(
+                max_state_dim=config.max_state_dim,
+                cfg_enabled=cfg_enabled,
+            ),
+            TokenizerProcessorStep(
+                tokenizer_name="google/paligemma-3b-pt-224",
+                max_length=config.tokenizer_max_length,
+                padding_side="right",
+                padding="max_length",
+            ),
+        ]
+    )
+
+    # Add unconditional prompt tokenizer for CFG inference
+    if cfg_enabled:
+        input_steps.append(
+            TokenizerProcessorStep(
+                tokenizer_name="google/paligemma-3b-pt-224",
+                max_length=config.tokenizer_max_length,
+                padding_side="right",
+                padding="max_length",
+                task_key="uncond_task",
+                output_tokens_key=OBS_LANGUAGE_UNCOND_TOKENS,
+                output_mask_key=OBS_LANGUAGE_UNCOND_ATTENTION_MASK,
+            )
+        )
+
+    input_steps.append(DeviceProcessorStep(device=config.device))
+
     output_steps: list[ProcessorStep] = [
         UnnormalizerProcessorStep(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats

src/lerobot/policies/pretrained.py

@@ -11,6 +11,8 @@
 # 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
@@ -19,7 +21,7 @@ import os
 from importlib.resources import files
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TypedDict, TypeVar, Unpack
+from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack
 
 import packaging
 import safetensors
@@ -38,10 +40,13 @@ 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_repo_id: str | None,
+    dataset_meta: LeRobotDatasetMetadata | None,
     input_features: dict | None,
     output_features: dict | None,
 ) -> dict:
@@ -72,30 +77,16 @@ def _build_card_context(
             "lerobot_version": __version__,
         }
 
-    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})"
-            )
+    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]
 
     return context
 
@@ -304,6 +295,7 @@ 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(
@@ -325,7 +317,12 @@ 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
+                cfg.dataset.repo_id,
+                self.config.type,
+                self.config.license,
+                self.config.tags,
+                cfg=cfg,
+                dataset_meta=dataset_meta,
             )
             card.save(str(saved_path / "README.md"))
 
@@ -340,6 +337,9 @@ 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,6 +349,7 @@ 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",
@@ -369,7 +370,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         )
 
         context = _build_card_context(
-            cfg, dataset_repo_id, self.config.input_features, self.config.output_features
+            cfg, dataset_meta, 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)
@@ -386,7 +387,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.
 

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

@@ -17,12 +17,10 @@ from __future__ import annotations
 import logging
 from collections import deque
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
-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
@@ -55,12 +53,13 @@ class VLAJEPAModel(nn.Module):
       - DiT-B: flow-matching action head for future action prediction
       - V-JEPA: world model for video frame prediction
 
-    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)
+    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)
     """
 
     def __init__(self, config: VLAJEPAConfig) -> None:
@@ -75,6 +74,11 @@ 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)
@@ -161,166 +165,123 @@ class VLAJEPAModel(nn.Module):
 
     # ---- Native VLA-JEPA forward (follows original VLA_JEPA.py) ----
 
-    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) ----
+    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."""
         qwen_inputs = self.qwen.build_inputs(
-            images=batch_images,
+            images=images,
             instructions=instructions,
             action_prompt=self.replace_prompt,
             embodied_prompt=self.embodied_replace_prompt,
         )
-
-        # 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)
+        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)
 
         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
 
-            if self.config.enable_world_model:
-                action_tokens = last_hidden[action_indices[0], action_indices[1], :].view(b, -1, h)
+    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]
 
-            embodied_action_tokens = last_hidden[embodied_indices[0], embodied_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]
 
-        # ---- 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)
+        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)
         else:
-            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)
+            wm_loss = torch.zeros((), device=embodied_action_tokens.device)
 
-            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:
+        if actions is None:
             return {"wm_loss": wm_loss}
 
-        # ---- 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
-            )
-
+        action_loss = self._action_loss(embodied_action_tokens, actions, state, action_is_pad)
         return {"action_loss": action_loss, "wm_loss": wm_loss * self.config.world_model_loss_weight}
 
     # ---- Native predict_action (follows original VLA_JEPA.predict_action) ----
@@ -328,58 +289,23 @@ class VLAJEPAModel(nn.Module):
     @torch.no_grad()
     def predict_action(
         self,
-        batch_images: list[list[Image.Image]],
+        images: list[list[Tensor]],
         instructions: list[str],
-        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.
-        """
+        state: Tensor | None = None,
+    ) -> Tensor:
+        """Predict an action chunk. `images` is per-sample, per-view float [0,1] [C, H, W] tensors."""
         if self.config.resize_images_to is not None:
             height, width = self.config.resize_images_to
-            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
+            images = [
+                [F.interpolate(img[None], size=(height, width), mode="area")[0] for img in views]
+                for views in images
             ]
 
-        qwen_inputs = self.qwen.build_inputs(
-            images=batch_images,
-            instructions=instructions,
-            action_prompt=self.replace_prompt,
-            embodied_prompt=self.embodied_replace_prompt,
+        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
         )
 
-        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
@@ -390,9 +316,9 @@ class VLAJEPAPolicy(PreTrainedPolicy):
     """
     LeRobot adapter for VLA-JEPA.
 
-    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.
+    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.
     """
 
     config_class = VLAJEPAConfig
@@ -419,9 +345,8 @@ class VLAJEPAPolicy(PreTrainedPolicy):
 
     # ---- Format Conversion: LeRobot → Native ----
 
-    def _prepare_model_inputs(self, batch: dict[str, Tensor]) -> list[dict]:
-        """
-        Convert LeRobot batch format to native VLA-JEPA examples format.
+    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.
 
         LeRobot format:
             batch = {
@@ -431,65 +356,25 @@ class VLAJEPAPolicy(PreTrainedPolicy):
                 "task": str | List[str],  (optional instruction)
             }
 
-        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
-            }
+        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.
         """
-        # 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.")
-        first_key = image_keys[0]
-        first_tensor = batch[first_key]
-        batch_size = first_tensor.shape[0]
+        batch_size = batch[image_keys[0]].shape[0]
 
-        # ---- 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)]
+        # Current-frame image per view ([B, C, H, W]); regroup per sample for Qwen messages.
+        frames = []
         for key in image_keys:
-            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]))
+            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)]
 
-        # ---- 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
@@ -498,52 +383,32 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         else:
             instructions = list(tasks)
 
-        # ---- 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)]
+        inputs: dict[str, Any] = {"images": images, "instructions": instructions}
 
-        # ---- 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)]
+        # 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))
 
-        # ---- 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)
+        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
 
-        return examples
+        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
 
     # ---- LeRobot Policy Interface ----
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """LeRobot train forward: convert → native forward → aggregate losses."""
-        examples = self._prepare_model_inputs(batch)
-        native_output = self.model.forward(examples)
+        native_output = self.model.forward(**self._prepare_model_inputs(batch, training=True))
 
         ref = next(iter(native_output.values()))
         zero = torch.zeros((), device=ref.device, dtype=ref.dtype)
@@ -561,16 +426,9 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         self.eval()
         self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
-        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)
+        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)
 
     @torch.no_grad()
     def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:

src/lerobot/policies/vla_jepa/qwen_interface.py

@@ -17,9 +17,7 @@ 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
 
@@ -78,7 +76,7 @@ class Qwen3VLInterface(torch.nn.Module):
 
     def build_inputs(
         self,
-        images: Sequence[Sequence[Image.Image]],
+        images: Sequence[Sequence[torch.Tensor]],
         instructions: Sequence[str],
         action_prompt: str,
         embodied_prompt: str,
@@ -94,24 +92,42 @@ 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"},
+            processor_kwargs={
+                "padding": True,
+                "return_tensors": "pt",
+                "device": self.model.device,
+                "do_rescale": False,
+            },
         )
         return batch_inputs.to(self.model.device)
 
     @staticmethod
-    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)
+    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

src/lerobot/processor/rendered_messages_to_task.py

@@ -0,0 +1,86 @@
+#!/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.
+
+"""Adapter step that flattens rendered chat messages back into a task string.
+
+Bridges RenderMessagesStep (which outputs structured messages) to policies
+that expect a plain task string in complementary_data["task"] (e.g. PI05).
+"""
+
+from __future__ import annotations
+
+from lerobot.configs import PipelineFeatureType, PolicyFeature
+
+from .pipeline import ComplementaryDataProcessorStep, ProcessorStepRegistry
+
+
+@ProcessorStepRegistry.register(name="rendered_messages_to_task")
+class RenderedMessagesToTaskStep(ComplementaryDataProcessorStep):
+    """Extract user-role message content from rendered messages into the task string.
+
+    After RenderMessagesStep renders a recipe into structured messages, this
+    step extracts content from all user-role messages, joins them, and writes
+    the result to complementary_data["task"]. This allows downstream steps
+    (like Pi05PrepareStateTokenizerProcessorStep) to consume the
+    advantage-conditioned prompt without modification.
+
+    No-ops when the "messages" key is absent (backward compatible with
+    pipelines that don't use language annotations).
+    """
+
+    def complementary_data(self, complementary_data: dict) -> dict:
+        messages = complementary_data.get("messages")
+        if messages is None:
+            return complementary_data
+
+        user_parts = []
+        for msg in messages:
+            if msg.get("role") == "user":
+                content = msg.get("content", "")
+                if isinstance(content, str) and content:
+                    user_parts.append(content)
+                elif isinstance(content, list):
+                    # HF multimodal blocks: extract text blocks
+                    for block in content:
+                        if isinstance(block, dict) and block.get("type") == "text":
+                            text = block.get("text", "")
+                            if text:
+                                user_parts.append(text)
+
+        new_complementary_data = dict(complementary_data)
+
+        if user_parts:
+            task = complementary_data.get("task")
+            # Preserve the original task for CFG unconditional prompt
+            new_complementary_data["base_task"] = task
+            # Wrap in list if the original task was a list (batched)
+            joined = "\n".join(user_parts)
+            if isinstance(task, list):
+                new_complementary_data["task"] = [joined] * len(task)
+            else:
+                new_complementary_data["task"] = joined
+
+        # Remove consumed rendering outputs
+        new_complementary_data.pop("messages", None)
+        new_complementary_data.pop("message_streams", None)
+        new_complementary_data.pop("target_message_indices", None)
+
+        return new_complementary_data
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features

src/lerobot/processor/tokenizer_processor.py

@@ -81,6 +81,8 @@ class TokenizerProcessorStep(ObservationProcessorStep):
     padding_side: str = "right"
     padding: str = "max_length"
     truncation: bool = True
+    output_tokens_key: str = OBS_LANGUAGE_TOKENS
+    output_mask_key: str = OBS_LANGUAGE_ATTENTION_MASK
 
     # Internal tokenizer instance (not part of the config)
     input_tokenizer: Any = field(default=None, init=False, repr=False)
@@ -201,8 +203,8 @@ class TokenizerProcessorStep(ObservationProcessorStep):
         new_observation = dict(observation)
 
         # Add tokenized data to the observation
-        new_observation[OBS_LANGUAGE_TOKENS] = tokenized_prompt["input_ids"]
-        new_observation[OBS_LANGUAGE_ATTENTION_MASK] = tokenized_prompt["attention_mask"].to(dtype=torch.bool)
+        new_observation[self.output_tokens_key] = tokenized_prompt["input_ids"]
+        new_observation[self.output_mask_key] = tokenized_prompt["attention_mask"].to(dtype=torch.bool)
 
         # Tokenize subtask if available
         subtask = self.get_subtask(self.transition)
@@ -309,14 +311,14 @@ class TokenizerProcessorStep(ObservationProcessorStep):
             The updated dictionary of policy features.
         """
         # Add a feature for the token IDs if it doesn't already exist
-        if OBS_LANGUAGE_TOKENS not in features[PipelineFeatureType.OBSERVATION]:
-            features[PipelineFeatureType.OBSERVATION][OBS_LANGUAGE_TOKENS] = PolicyFeature(
+        if self.output_tokens_key not in features[PipelineFeatureType.OBSERVATION]:
+            features[PipelineFeatureType.OBSERVATION][self.output_tokens_key] = PolicyFeature(
                 type=FeatureType.LANGUAGE, shape=(self.max_length,)
             )
 
         # Add a feature for the attention mask if it doesn't already exist
-        if OBS_LANGUAGE_ATTENTION_MASK not in features[PipelineFeatureType.OBSERVATION]:
-            features[PipelineFeatureType.OBSERVATION][OBS_LANGUAGE_ATTENTION_MASK] = PolicyFeature(
+        if self.output_mask_key not in features[PipelineFeatureType.OBSERVATION]:
+            features[PipelineFeatureType.OBSERVATION][self.output_mask_key] = PolicyFeature(
                 type=FeatureType.LANGUAGE, shape=(self.max_length,)
             )
 

src/lerobot/rewards/__init__.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 
 from .classifier.configuration_classifier import RewardClassifierConfig as RewardClassifierConfig
+from .distributional_value_function.configuration_distributional_value_function import (
+    DistributionalVFConfig as DistributionalVFConfig,
+)
 from .factory import (
     get_reward_model_class as get_reward_model_class,
     make_reward_model as make_reward_model,
@@ -26,6 +29,7 @@ from .topreward.configuration_topreward import TOPRewardConfig as TOPRewardConfi
 
 __all__ = [
     # Configuration classes
+    "DistributionalVFConfig",
     "RewardClassifierConfig",
     "RobometerConfig",
     "SARMConfig",

src/lerobot/rewards/distributional_value_function/__init__.py

@@ -0,0 +1,23 @@
+# 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 .configuration_distributional_value_function import DistributionalVFConfig
+from .modeling_distributional_value_function import DistributionalVFRewardModel
+from .processor_distributional_value_function import make_distributional_vf_pre_post_processors
+
+__all__ = [
+    "DistributionalVFConfig",
+    "DistributionalVFRewardModel",
+    "make_distributional_vf_pre_post_processors",
+]

src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py

@@ -0,0 +1,108 @@
+# 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.
+
+"""Configuration for RECAP's distributional value function.
+
+Paper: "π*0.6: a VLA That Learns From Experience" (Physical Intelligence, 2025)
+       https://pi.website/blog/pistar06
+
+Implements the distributional value function V^{pi_ref}(o_t, l) from Section IV-A.
+Architecture: the paper uses a 670M-parameter Gemma 3 VLM (the actor is 4B Gemma 3).
+We match that scale on PaliGemma (PI05's Gemma 2B backbone) by truncating to 6 Gemma
+LM layers and 13 SigLIP vision layers (~670M params), with a [CLS] token and linear
+head predicting a categorical distribution over B=201 discrete value bins in [-1, 0].
+
+Training: cross-entropy on HL-Gauss soft targets (or Dirac delta projection),
+with optional one-hot targets for terminal states; MC returns normalized per task.
+Weights initialized from a pre-trained PI05 actor checkpoint.
+"""
+
+from dataclasses import dataclass, field
+
+from lerobot.configs import FeatureType, NormalizationMode
+from lerobot.configs.rewards import RewardModelConfig
+from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
+
+
+@RewardModelConfig.register_subclass("distributional_value_function")
+@dataclass
+class DistributionalVFConfig(RewardModelConfig):
+    """Configuration for RECAP's distributional value function.
+
+    The value function predicts V^{pi_ref}(o_t, l) as a distribution over B discrete
+    bins spanning [value_support_min, value_support_max]. It is trained with cross-entropy
+    on HL-Gauss soft targets or Dirac delta projection, derived from Monte Carlo returns
+    (Eq. 1 in the paper).
+
+    Architecture: the paper value function is a 670M Gemma 3 VLM; the actor is 4B Gemma 3.
+    We use truncated PaliGemma (``num_hidden_layers=6``, ``num_vision_layers=13``) to reach
+    about 670M params and initialize from the PI05 actor checkpoint.
+    """
+
+    # Backbone
+    paligemma_variant: str = "gemma_2b"
+    num_hidden_layers: int = 6
+    num_vision_layers: int = 13
+
+    # Distributional head
+    num_value_bins: int = 201
+    value_support_min: float = -1.0
+    value_support_max: float = 0.0
+    hl_gauss_sigma_ratio: float = 5.0
+
+    # Target distribution method: "hl_gauss" (default, soft) or "dirac_delta" (C51, hard)
+    target_method: str = "hl_gauss"
+
+    # Whether to use one-hot targets for terminal states (exact return, no smoothing).
+    # When False, terminal states use the same target method as non-terminal states.
+    use_one_hot_terminal: bool = True
+
+    # Image
+    image_resolution: tuple[int, int] = (224, 224)
+
+    # Tokenizer
+    tokenizer_max_length: int = 64
+
+    # Init from actor (required for first training: provides SigLIP vision tower + Gemma embeddings).
+    # Pass a PI05 checkpoint path or Hub repo_id here.
+    # After training, load the value function with RewardModel.from_pretrained() instead.
+    init_from_actor_path: str = ""
+
+    # Normalization
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.IDENTITY,
+        }
+    )
+
+    def get_optimizer_preset(self) -> AdamWConfig:
+        return AdamWConfig(
+            lr=3e-4,
+            weight_decay=1e-4,
+            grad_clip_norm=1.0,
+        )
+
+    def get_scheduler_preset(self) -> CosineDecayWithWarmupSchedulerConfig:
+        return CosineDecayWithWarmupSchedulerConfig(
+            num_warmup_steps=500,
+            num_decay_steps=50000,
+        )
+
+    def validate_features(self) -> None:
+        if not self.input_features:
+            return
+        has_image = any(ft.type == FeatureType.VISUAL for ft in self.input_features.values())
+        if not has_image:
+            raise ValueError("DistributionalVFConfig requires at least one VISUAL input feature.")

src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py

@@ -0,0 +1,567 @@
+# 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.
+
+"""Modeling for RECAP's distributional value function.
+
+Paper: "π*0.6: a VLA That Learns From Experience" (Physical Intelligence, 2025)
+       https://pi.website/blog/pistar06
+
+Implements the distributional value function V^{pi_ref}(o_t, l) from Section IV-A.
+Architecture: the paper uses a 670M-parameter Gemma 3 VLM (the actor is 4B Gemma 3).
+We match that scale on PaliGemma (PI05's Gemma 2B backbone) by truncating to 6 Gemma
+LM layers and 13 SigLIP vision layers (~670M params), with a [CLS] token and linear
+head predicting a categorical distribution over B=201 discrete value bins in [-1, 0].
+
+Inputs: single image observation + task text prompt ("Task: {task}.")
+Outputs: softmax distribution over value bins; expected value E[V] for inference.
+Training: cross-entropy on HL-Gauss soft targets (or Dirac delta projection),
+with optional one-hot targets for terminal states; MC returns normalized per task.
+
+Weight initialization: vision tower, multi-modal projector, token embeddings, and
+the first N transformer layers are copied from a pre-trained PI05 actor checkpoint.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import Tensor, nn
+
+from lerobot.rewards.pretrained import PreTrainedRewardModel
+from lerobot.utils.import_utils import _transformers_available, require_package
+
+from .configuration_distributional_value_function import DistributionalVFConfig
+
+if TYPE_CHECKING or _transformers_available:
+    from transformers.models.auto import CONFIG_MAPPING
+    from transformers.models.gemma import modeling_gemma
+
+    from lerobot.policies.pi_gemma import (
+        PaliGemmaForConditionalGenerationWithPiGemma,
+        PiGemmaRMSNorm,
+        _gated_residual,
+        _get_pi_gemma_decoder_layer_base,
+    )
+else:
+    CONFIG_MAPPING = None
+    modeling_gemma = None
+    PaliGemmaForConditionalGenerationWithPiGemma = None
+    PiGemmaRMSNorm = None
+    _gated_residual = None
+    _get_pi_gemma_decoder_layer_base = None
+
+PALIGEMMA_VOCAB_SIZE = 257152
+
+
+class DistributionalVFRewardModel(PreTrainedRewardModel):
+    """Distributional value function model for RECAP.
+
+    Predicts V^{pi_ref}(o_t, l) as a categorical distribution over B bins (default 201).
+    Trained with cross-entropy on HL-Gauss or Dirac delta targets centered on
+    per-task normalized Monte Carlo returns.
+
+    Architecture: truncated PaliGemma (``num_hidden_layers=6``, ``num_vision_layers=13``),
+    causal attention, [CLS] token, and Linear(D, num_bins) value head.
+    The expected value is E[V] = sum(softmax(logits) * bin_centers).
+    """
+
+    name = "distributional_value_function"
+    config_class = DistributionalVFConfig
+
+    def __init__(self, config: DistributionalVFConfig, **kwargs) -> None:
+        require_package("transformers", extra="recap")
+        super().__init__(config)
+        self.config = config
+
+        from transformers.models.gemma.modeling_gemma import GemmaRotaryEmbedding
+
+        from lerobot.policies.pi05.modeling_pi05 import get_gemma_config
+
+        # Get base dimensions from the paligemma variant (OpenPI config format)
+        base_config = get_gemma_config(config.paligemma_variant)
+        hidden_dim = base_config.width
+        mlp_dim = base_config.mlp_dim
+        num_layers = config.num_hidden_layers
+
+        # HuggingFace GemmaConfig for transformer layers
+        gemma_config = CONFIG_MAPPING["gemma"](
+            head_dim=base_config.head_dim,
+            hidden_size=hidden_dim,
+            intermediate_size=mlp_dim,
+            num_attention_heads=base_config.num_heads,
+            num_hidden_layers=num_layers,
+            num_key_value_heads=base_config.num_kv_heads,
+            vocab_size=PALIGEMMA_VOCAB_SIZE,
+            hidden_activation="gelu_pytorch_tanh",
+        )
+        self.gemma_config = gemma_config
+        self.hidden_dim = hidden_dim
+        self.num_value_bins = config.num_value_bins
+
+        # Single learned [CLS] token for value prediction
+        self.cls_embedding = nn.Parameter(torch.randn(1, 1, hidden_dim) * 0.02)
+
+        # Value projection head: Linear(hidden_dim, num_bins)
+        self.value_head = nn.Linear(in_features=hidden_dim, out_features=config.num_value_bins)
+
+        # Transformer layers (overwritten by _initialize_from_actor on first run)
+        self.rotary_emb = GemmaRotaryEmbedding(gemma_config)
+        pi_gemma_decoder_layer_base = _get_pi_gemma_decoder_layer_base()
+        self.layers = nn.ModuleList(
+            [pi_gemma_decoder_layer_base(gemma_config, layer_idx=i) for i in range(num_layers)]
+        )
+        self.norm = PiGemmaRMSNorm(hidden_dim, eps=gemma_config.rms_norm_eps)
+
+        # Vision tower + projector + token embedding (overwritten by _initialize_from_actor on first run)
+        # PaliGemmaConfig wraps both vision and text configs into a single model
+        paligemma_config = CONFIG_MAPPING["paligemma"]()
+        paligemma_config.text_config = gemma_config
+        paligemma_config.vision_config.image_size = config.image_resolution[0]
+        paligemma_config.vision_config.intermediate_size = 4304
+        paligemma_config.vision_config.projection_dim = 2048
+        paligemma_config.vision_config.projector_hidden_act = "gelu_fast"
+
+        paligemma_full = PaliGemmaForConditionalGenerationWithPiGemma(config=paligemma_config)
+        self.vision_tower = paligemma_full.model.vision_tower
+        self.multi_modal_projector = paligemma_full.model.multi_modal_projector
+        self.token_embedding = paligemma_full.model.language_model.embed_tokens
+        del paligemma_full
+
+        # Truncate vision tower to num_vision_layers
+        if hasattr(self.vision_tower, "vision_model") and hasattr(self.vision_tower.vision_model, "encoder"):
+            vision_encoder = self.vision_tower.vision_model.encoder
+            vision_encoder.layers = vision_encoder.layers[: config.num_vision_layers]
+
+        # Bin support: evenly spaced centers from value_support_min to value_support_max
+        bin_centers = torch.linspace(config.value_support_min, config.value_support_max, self.num_value_bins)
+        self.register_buffer("bin_centers", bin_centers, persistent=False)
+        bin_width = (config.value_support_max - config.value_support_min) / (self.num_value_bins - 1)
+        self.hl_gauss_sigma = float(config.hl_gauss_sigma_ratio * bin_width)
+
+        # Overwrite with pre-trained PI05 actor weights (first training run only)
+        if config.init_from_actor_path:
+            self._initialize_from_actor()
+
+    def _initialize_from_actor(self) -> None:
+        """Overwrite weights from a pre-trained PI05 actor checkpoint.
+
+        Called on first training run only (when init_from_actor_path is set).
+        """
+        from lerobot.policies.pi05.modeling_pi05 import PI05Policy
+
+        actor_policy = PI05Policy.from_pretrained(self.config.init_from_actor_path)
+        actor_model = actor_policy.model
+
+        paligemma_model = actor_model.paligemma_with_expert.paligemma
+        source_language_model = paligemma_model.model.language_model
+
+        # Transformer components
+        self.rotary_emb.load_state_dict(source_language_model.rotary_emb.state_dict())
+        num_layers = self.gemma_config.num_hidden_layers
+        for i in range(num_layers):
+            self.layers[i].load_state_dict(source_language_model.layers[i].state_dict())
+        self.norm.load_state_dict(source_language_model.norm.state_dict())
+
+        # Vision tower (truncate source first, then copy)
+        source_vision_tower = paligemma_model.model.vision_tower
+        if hasattr(source_vision_tower, "vision_model") and hasattr(
+            source_vision_tower.vision_model, "encoder"
+        ):
+            source_encoder = source_vision_tower.vision_model.encoder
+            source_encoder.layers = source_encoder.layers[: self.config.num_vision_layers]
+        self.vision_tower.load_state_dict(source_vision_tower.state_dict())
+
+        # Multi-modal projector
+        self.multi_modal_projector.load_state_dict(paligemma_model.model.multi_modal_projector.state_dict())
+
+        # Token embedding table
+        self.token_embedding.load_state_dict(paligemma_model.model.language_model.embed_tokens.state_dict())
+
+        del actor_policy
+
+    def embed_image(self, image: Tensor) -> Tensor:
+        """Embed images using the value function's SigLIP vision tower.
+
+        Args:
+            image: [batch_size, channels, height, width] preprocessed images in [-1, 1].
+
+        Returns:
+            [batch_size, num_patches, hidden_dim] projected image features.
+        """
+        out_dtype = image.dtype
+        if image.dtype != torch.float32:
+            image = image.to(torch.float32)
+
+        image_outputs = self.vision_tower(image, return_dict=True)
+        image_features = self.multi_modal_projector(image_outputs.last_hidden_state)
+        image_features = image_features / (self.hidden_dim**0.5)
+
+        if image_features.dtype != out_dtype:
+            image_features = image_features.to(out_dtype)
+        return image_features
+
+    def embed_text(self, token_ids: Tensor) -> Tensor:
+        """Embed text token IDs using the value function's token embedding table.
+
+        Args:
+            token_ids: [batch_size, seq_len] integer token IDs
+
+        Returns:
+            [batch_size, seq_len, hidden_dim] text embeddings
+        """
+        return self.token_embedding(token_ids)
+
+    def _get_cls_embedding(self, batch_size: int) -> Tensor:
+        """Get [CLS] token embedding expanded to batch size.
+
+        Args:
+            batch_size: number of samples in the batch.
+
+        Returns:
+            [batch_size, 1, hidden_dim] learned [CLS] embedding.
+        """
+        return self.cls_embedding.expand(batch_size, -1, -1)
+
+    def forward_value(
+        self, vision_features: Tensor, text_embeddings: Tensor, text_padding_mask: Tensor
+    ) -> dict[str, Tensor]:
+        """Core forward pass through the distributional value function.
+
+        Args:
+            vision_features: [batch_size, num_patches, hidden_dim]
+            text_embeddings: [batch_size, seq_len, hidden_dim]
+            text_padding_mask: [batch_size, seq_len] boolean mask for text tokens
+
+        Returns:
+            logits: [batch_size, num_value_bins]
+            probs: [batch_size, num_value_bins]
+            value: [batch_size, 1]
+        """
+        from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE
+
+        batch_size = text_embeddings.shape[0]
+        device = text_embeddings.device
+
+        # Build sequence: [vision, text, CLS]
+        cls_embedding = self._get_cls_embedding(batch_size)
+        hidden_states = torch.cat([vision_features, text_embeddings, cls_embedding], dim=1)
+
+        # Build causal attention mask
+        vision_len = vision_features.shape[1]
+        vision_padding_mask = torch.ones(batch_size, vision_len, dtype=torch.bool, device=device)
+        cls_padding_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=device)
+        full_padding_mask = torch.cat([vision_padding_mask, text_padding_mask, cls_padding_mask], dim=1)
+
+        full_seq_len = full_padding_mask.shape[1]
+
+        # Causal mask
+        causal_mask = torch.tril(torch.ones(full_seq_len, full_seq_len, device=device, dtype=torch.bool))
+        # Combine causal mask with padding mask
+        padding_mask_4d = full_padding_mask[:, None, None, :].expand(
+            batch_size, 1, full_seq_len, full_seq_len
+        )
+        attention_mask = causal_mask[None, None, :, :] & padding_mask_4d
+        attention_mask = torch.where(attention_mask, 0.0, OPENPI_ATTENTION_MASK_VALUE)
+
+        position_ids = torch.cumsum(full_padding_mask.long(), dim=1) - 1
+        cos, sin = self.rotary_emb(hidden_states, position_ids)
+
+        for layer in self.layers:
+            norm_output = layer.input_layernorm(hidden_states, cond=None)
+            if isinstance(norm_output, tuple):
+                hidden_states_normed, gate = norm_output
+            else:
+                hidden_states_normed, gate = norm_output, None
+
+            input_shape = hidden_states_normed.shape[:-1]
+            hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
+
+            query_states = layer.self_attn.q_proj(hidden_states_normed).view(hidden_shape).transpose(1, 2)
+            key_states = layer.self_attn.k_proj(hidden_states_normed).view(hidden_shape).transpose(1, 2)
+            value_states = layer.self_attn.v_proj(hidden_states_normed).view(hidden_shape).transpose(1, 2)
+
+            query_states, key_states = modeling_gemma.apply_rotary_pos_emb(
+                query_states, key_states, cos, sin, unsqueeze_dim=1
+            )
+
+            attention_output, _ = modeling_gemma.eager_attention_forward(
+                layer.self_attn,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                layer.self_attn.scaling,
+            )
+
+            attention_output = attention_output.reshape(batch_size, -1, self.gemma_config.hidden_size)
+            if attention_output.dtype != layer.self_attn.o_proj.weight.dtype:
+                attention_output = attention_output.to(layer.self_attn.o_proj.weight.dtype)
+            projected_attention = layer.self_attn.o_proj(attention_output)
+
+            if gate is not None:
+                projected_attention = _gated_residual(hidden_states, projected_attention, gate)
+            else:
+                projected_attention = hidden_states + projected_attention
+
+            after_attention_residual = projected_attention.clone()
+
+            norm_output = layer.post_attention_layernorm(projected_attention, cond=None)
+            if isinstance(norm_output, tuple):
+                mlp_input, gate = norm_output
+            else:
+                mlp_input, gate = norm_output, None
+
+            mlp_output = layer.mlp(mlp_input)
+
+            if gate is not None:
+                hidden_states = _gated_residual(after_attention_residual, mlp_output, gate)
+            else:
+                hidden_states = after_attention_residual + mlp_output
+
+        hidden_states = self.norm(hidden_states)
+        if isinstance(hidden_states, tuple):
+            hidden_states = hidden_states[0]
+
+        # Extract [CLS] token (last position in the sequence)
+        cls_hidden_state = hidden_states[:, -1, :]  # [batch_size, hidden_dim]
+
+        # Value head: Linear(hidden_dim, num_bins) -> logits
+        value_logits = self.value_head(cls_hidden_state)  # [batch_size, num_value_bins]
+        value_probs = F.softmax(value_logits, dim=-1)
+        predicted_value = (value_probs * self.bin_centers.to(dtype=value_probs.dtype)).sum(
+            dim=-1, keepdim=True
+        )
+
+        return {"logits": value_logits, "probs": value_probs, "value": predicted_value}
+
+    def hl_gauss_target(self, target_value: Tensor) -> Tensor:
+        """HL-Gauss soft target distribution.
+
+        Places a Gaussian N(target, sigma^2) over the bin support and computes
+        per-bin probabilities as CDF differences at bin edges, normalized to sum to 1.
+
+        Reference: Farebrother et al. 2024, "Stop Regressing: Training Value
+        Functions via Classification for Scalable Deep RL", Section 3.1.
+        arXiv:2403.03950
+
+        Args:
+            target_value: [batch_size] or [batch_size, 1] target values.
+
+        Returns:
+            [batch_size, num_value_bins] target probability distribution.
+        """
+        if target_value.ndim == 2:
+            target_value = target_value.squeeze(-1)
+        target_value = target_value.to(dtype=self.bin_centers.dtype)
+
+        # Bin edges: half a bin-width outside the first/last center
+        bin_width = (self.config.value_support_max - self.config.value_support_min) / (
+            self.num_value_bins - 1
+        )
+        support_edges = torch.linspace(
+            self.config.value_support_min - bin_width / 2,
+            self.config.value_support_max + bin_width / 2,
+            self.num_value_bins + 1,
+            device=target_value.device,
+            dtype=target_value.dtype,
+        )
+
+        # CDF of N(target, sigma^2) evaluated at each edge
+        cdf_at_edges = 0.5 * (
+            1.0
+            + torch.erf(
+                (support_edges.unsqueeze(0) - target_value.unsqueeze(-1))
+                / (self.hl_gauss_sigma * math.sqrt(2))
+            )
+        )  # [batch_size, num_bins + 1]
+
+        # Normalize: z = cdf(max_edge) - cdf(min_edge)
+        normalization_constant = (cdf_at_edges[:, -1] - cdf_at_edges[:, 0]).unsqueeze(-1).clamp(min=1e-10)
+
+        # Bin probabilities = differences of consecutive CDF values, normalized
+        bin_probabilities = (cdf_at_edges[:, 1:] - cdf_at_edges[:, :-1]) / normalization_constant
+
+        return bin_probabilities
+
+    def dirac_delta_target(self, target_value: Tensor) -> Tensor:
+        """Dirac delta (C51) projection: split probability between two nearest bins.
+
+        Standard distributional RL projection from Bellemare et al. 2017.
+        "A Distributional Perspective on Reinforcement Learning"
+        arXiv:1707.06887
+
+        Args:
+            target_value: [batch_size] or [batch_size, 1] target values.
+
+        Returns:
+            [batch_size, num_value_bins] target probability distribution.
+        """
+        if target_value.ndim == 2:
+            target_value = target_value.squeeze(-1)
+        target_value = target_value.clamp(self.config.value_support_min, self.config.value_support_max)
+        target_value = target_value.to(dtype=self.bin_centers.dtype)
+
+        bin_width = self.bin_centers[1] - self.bin_centers[0]
+        normalized_position = (target_value - self.config.value_support_min) / bin_width
+        lower_bin_idx = normalized_position.floor().long().clamp(0, self.num_value_bins - 1)
+        upper_bin_idx = normalized_position.ceil().long().clamp(0, self.num_value_bins - 1)
+
+        weight_upper = normalized_position - lower_bin_idx.float()
+        weight_lower = upper_bin_idx.float() - normalized_position
+
+        same_bin = lower_bin_idx == upper_bin_idx
+        weight_upper = torch.where(same_bin, torch.zeros_like(weight_upper), weight_upper)
+        weight_lower = torch.where(same_bin, torch.ones_like(weight_lower), weight_lower)
+
+        batch_size = target_value.shape[0]
+        target_distribution = torch.zeros(batch_size, self.num_value_bins, device=target_value.device)
+        batch_indices = torch.arange(batch_size, device=target_value.device)
+        target_distribution[batch_indices, lower_bin_idx] += weight_lower
+        target_distribution[batch_indices, upper_bin_idx] += weight_upper
+
+        return target_distribution
+
+    def one_hot_target(self, target_value: Tensor) -> Tensor:
+        """One-hot target for terminal states (exact return, no smoothing).
+
+        Args:
+            target_value: [batch_size] or [batch_size, 1] target values.
+
+        Returns:
+            [batch_size, num_value_bins] one-hot distribution at the nearest bin.
+        """
+        if target_value.ndim == 2:
+            target_value = target_value.squeeze(-1)
+        target_value = target_value.to(dtype=self.bin_centers.dtype)
+        nearest_bin_idx = torch.argmin(
+            torch.abs(self.bin_centers.unsqueeze(0) - target_value.unsqueeze(-1)), dim=-1
+        )
+        return F.one_hot(nearest_bin_idx, num_classes=self.num_value_bins).to(dtype=self.bin_centers.dtype)
+
+    def compute_target_distribution(
+        self,
+        target_value: Tensor,
+        is_terminal: Tensor,
+        method: str = "hl_gauss",
+        use_one_hot_terminal: bool = True,
+    ) -> Tensor:
+        """Compute target distribution using configured method.
+
+        Args:
+            target_value: [batch_size] scalar return targets
+            is_terminal: [batch_size] boolean terminal flags
+            method: "hl_gauss" or "dirac_delta"
+            use_one_hot_terminal: if True, terminal states get one-hot targets
+                (exact return, no smoothing). If False, all states use the same method.
+
+        Returns:
+            [batch_size, num_value_bins] target probability distribution
+        """
+        if method == "hl_gauss":
+            base_distribution = self.hl_gauss_target(target_value)
+        elif method == "dirac_delta":
+            base_distribution = self.dirac_delta_target(target_value)
+        else:
+            raise ValueError(f"Unknown target method: {method}. Use 'hl_gauss' or 'dirac_delta'.")
+
+        if not use_one_hot_terminal:
+            return base_distribution
+
+        terminal_distribution = self.one_hot_target(target_value)
+
+        return torch.where(is_terminal[:, None].bool(), terminal_distribution, base_distribution)
+
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]:
+        """Training forward pass — computes cross-entropy loss against MC return targets.
+
+        The batch is expected to be preprocessed by the processor pipeline.
+        Keys expected in batch:
+            - observation.images.*: [B, C, H, W] preprocessed images
+            - observation.language_tokens: [B, seq_len] tokenized task prompt
+            - observation.language_attention_mask: [B, seq_len] padding mask
+            - mc_return: [B] normalized Monte Carlo return targets in (-1, 0)
+            - is_terminal: [B] boolean terminal flags
+
+        Returns:
+            (loss, output_dict) where loss is scalar cross-entropy
+        """
+        from lerobot.utils.constants import OBS_IMAGES, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
+
+        # Get first image key from batch
+        image_keys = [k for k in batch if k.startswith(f"{OBS_IMAGES}.") or k == OBS_IMAGES]
+        if not image_keys:
+            raise KeyError(f"No image keys found in batch. Expected keys starting with '{OBS_IMAGES}.'")
+        images = batch[image_keys[0]]
+
+        token_ids = batch[OBS_LANGUAGE_TOKENS]
+        text_padding_mask = batch[OBS_LANGUAGE_ATTENTION_MASK].bool()
+        mc_return = batch["mc_return"]
+        is_terminal = batch["is_terminal"]
+
+        # Embed observations
+        vision_features = self.embed_image(images)
+        text_embeddings = self.embed_text(token_ids)
+
+        # Forward through value function transformer
+        vf_output = self.forward_value(vision_features, text_embeddings, text_padding_mask)
+        value_logits = vf_output["logits"]
+        predicted_value = vf_output["value"]
+
+        # Compute target distribution
+        target_distribution = self.compute_target_distribution(
+            mc_return,
+            is_terminal,
+            method=self.config.target_method,
+            use_one_hot_terminal=self.config.use_one_hot_terminal,
+        )
+
+        # Cross-entropy loss (Eq. 1 in pi*0.6 paper)
+        log_probs = F.log_softmax(value_logits, dim=-1)
+        loss = -(target_distribution * log_probs).sum(dim=-1).mean()
+
+        output_dict = {
+            "loss": loss.item(),
+            "predicted_value_mean": predicted_value.mean().item(),
+            "mc_return_mean": mc_return.mean().item(),
+        }
+
+        return loss, output_dict
+
+    def compute_reward(self, batch: dict[str, Tensor]) -> Tensor:
+        """Compute V(s) for a batch of observations. Used for advantage scoring.
+
+        Args:
+            batch: preprocessed batch with images and tokenized text
+
+        Returns:
+            [batch_size] tensor of predicted values V(s)
+        """
+        from lerobot.utils.constants import OBS_IMAGES, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
+
+        image_keys = [k for k in batch if k.startswith(f"{OBS_IMAGES}.") or k == OBS_IMAGES]
+        if not image_keys:
+            raise KeyError(f"No image keys found in batch. Expected keys starting with '{OBS_IMAGES}.'")
+        images = batch[image_keys[0]]
+
+        token_ids = batch[OBS_LANGUAGE_TOKENS]
+        text_padding_mask = batch[OBS_LANGUAGE_ATTENTION_MASK].bool()
+
+        vision_features = self.embed_image(images)
+        text_embeddings = self.embed_text(token_ids)
+
+        vf_output = self.forward_value(vision_features, text_embeddings, text_padding_mask)
+        return vf_output["value"].squeeze(-1)  # [batch_size]

src/lerobot/rewards/distributional_value_function/processor_distributional_value_function.py

@@ -0,0 +1,235 @@
+# 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.
+
+"""Processor for RECAP's distributional value function.
+
+Paper: "π*0.6: a VLA That Learns From Experience" (Physical Intelligence, 2025)
+       https://pi.website/blog/pistar06
+
+Prepares inputs for V^{pi_ref}(o_t, l): single image observation and task text only.
+1. Image preprocessing (resize-with-pad + normalize to [-1, 1]) for SigLIP
+2. Task prompt formatting ("Task: {task}.") and tokenization via PaliGemma tokenizer
+
+Training targets (mc_return, is_terminal) are NOT routed through the processor.
+They are dataset columns read directly from the batch in the model's forward().
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import Tensor
+
+from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    ProcessorStep,
+    ProcessorStepRegistry,
+    RenameObservationsProcessorStep,
+    TokenizerProcessorStep,
+    batch_to_transition,
+    policy_action_to_transition,
+    transition_to_batch,
+)
+from lerobot.processor.converters import to_tensor
+from lerobot.types import EnvTransition, TransitionKey
+from lerobot.utils.constants import (
+    OBS_IMAGES,
+    POLICY_POSTPROCESSOR_DEFAULT_NAME,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
+
+from .configuration_distributional_value_function import DistributionalVFConfig
+
+PALIGEMMA_TOKENIZER_NAME = "google/paligemma-3b-pt-224"
+
+
+@ProcessorStepRegistry.register(name="distributional_vf_prepare_task_prompt")
+@dataclass
+class DistributionalVFPrepareTaskPromptStep(ProcessorStep):
+    """Format the task string for the distributional value function.
+
+    The value function receives only visual observations and task text.
+    Builds prompt: "Task: {task}."
+    """
+
+    task_key: str = "task"
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        transition = transition.copy()
+
+        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
+        tasks = complementary_data.get(self.task_key)
+        if tasks is None:
+            raise ValueError("No task found in complementary data")
+
+        if isinstance(tasks, str):
+            tasks = [tasks]
+
+        full_prompts = []
+        for task in tasks:
+            cleaned_text = task.strip().replace("_", " ").replace("\n", " ")
+            full_prompts.append(f"Task: {cleaned_text}.")
+
+        new_complementary_data = dict(complementary_data)
+        new_complementary_data[self.task_key] = full_prompts
+        transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
+        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 {"task_key": self.task_key}
+
+
+@ProcessorStepRegistry.register(name="distributional_vf_image_preprocessor")
+@dataclass
+class DistributionalVFImagePreprocessorStep(ProcessorStep):
+    """Resize and normalize images for the value function's SigLIP vision tower.
+
+    Expects float images in [0, 1].
+    - Resize-with-pad to ``image_resolution`` (preserves aspect ratio)
+    - Scale to [-1, 1] for SigLIP
+    """
+
+    image_resolution: tuple[int, int] = (224, 224)
+    image_keys: tuple[str, ...] | None = None
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        from lerobot.policies.pi05.modeling_pi05 import resize_with_pad_torch
+
+        observation = transition.get(TransitionKey.OBSERVATION)
+        if not isinstance(observation, dict):
+            raise ValueError("DistributionalVFImagePreprocessorStep requires an observation dict")
+
+        image_keys = self.image_keys or tuple(
+            key for key in observation if key == OBS_IMAGES or key.startswith(f"{OBS_IMAGES}.")
+        )
+        if not image_keys:
+            raise KeyError(
+                f"Distributional value function expected image keys under {OBS_IMAGES!r} in observation"
+            )
+
+        new_observation = dict(observation)
+        for image_key in image_keys:
+            image = new_observation[image_key]
+            if not isinstance(image, Tensor):
+                image = to_tensor(image)
+            if image.dtype != torch.float32:
+                image = image.to(torch.float32)
+
+            is_channels_first = image.ndim == 4 and image.shape[1] == 3
+            if is_channels_first:
+                image = image.permute(0, 2, 3, 1)
+
+            if image.shape[1:3] != self.image_resolution:
+                image = resize_with_pad_torch(image, *self.image_resolution)
+
+            image = image * 2.0 - 1.0
+
+            if is_channels_first:
+                image = image.permute(0, 3, 1, 2)
+
+            new_observation[image_key] = image
+
+        new_transition = transition.copy()
+        new_transition[TransitionKey.OBSERVATION] = new_observation
+        return new_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 {
+            "image_resolution": self.image_resolution,
+            "image_keys": list(self.image_keys) if self.image_keys is not None else None,
+        }
+
+
+def _visual_image_keys(config: DistributionalVFConfig) -> tuple[str, ...]:
+    return tuple(
+        feature_name
+        for feature_name, feature in config.input_features.items()
+        if feature.type == FeatureType.VISUAL
+    )
+
+
+def make_distributional_vf_pre_post_processors(
+    config: DistributionalVFConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """Create pre/post processors for the distributional value function.
+
+    Preprocessor steps:
+        1. Rename observations (no-op by default)
+        2. Add a batch dimension
+        3. Normalize features (images use identity, so they stay in [0, 1])
+        4. Format task prompt: "Task: {task}."
+        5. Tokenize with the PaliGemma tokenizer
+        6. Resize-with-pad and scale images to [-1, 1] for SigLIP
+        7. Move tensors to the configured device
+
+    Training targets (mc_return, is_terminal) are not processed here.
+    The model reads them directly from the batch in forward().
+
+    The postprocessor is a no-op because the value function does not need
+    action postprocessing.
+    """
+    image_keys = _visual_image_keys(config)
+
+    preprocessor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+        steps=[
+            RenameObservationsProcessorStep(rename_map={}),
+            AddBatchDimensionProcessorStep(),
+            NormalizerProcessorStep(
+                features={**config.input_features, **config.output_features},
+                norm_map=config.normalization_mapping,
+                stats=dataset_stats,
+            ),
+            DistributionalVFPrepareTaskPromptStep(),
+            TokenizerProcessorStep(
+                tokenizer_name=PALIGEMMA_TOKENIZER_NAME,
+                max_length=config.tokenizer_max_length,
+                padding_side="right",
+                padding="max_length",
+            ),
+            DistributionalVFImagePreprocessorStep(
+                image_resolution=config.image_resolution,
+                image_keys=image_keys or None,
+            ),
+            DeviceProcessorStep(device=config.device or "cpu"),
+        ],
+        name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        to_transition=batch_to_transition,
+        to_output=transition_to_batch,
+    )
+    postprocessor = PolicyProcessorPipeline(
+        name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+        to_transition=policy_action_to_transition,
+    )
+    return preprocessor, postprocessor

src/lerobot/rewards/factory.py

@@ -24,6 +24,7 @@ from lerobot.configs.rewards import RewardModelConfig
 from lerobot.processor import PolicyAction, PolicyProcessorPipeline
 
 from .classifier.configuration_classifier import RewardClassifierConfig
+from .distributional_value_function.configuration_distributional_value_function import DistributionalVFConfig
 from .pretrained import PreTrainedRewardModel
 from .robometer.configuration_robometer import RobometerConfig
 from .sarm.configuration_sarm import SARMConfig
@@ -63,6 +64,12 @@ def get_reward_model_class(name: str) -> type[PreTrainedRewardModel]:
         from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
 
         return TOPRewardModel
+    elif name == "distributional_value_function":
+        from lerobot.rewards.distributional_value_function.modeling_distributional_value_function import (
+            DistributionalVFRewardModel,
+        )
+
+        return DistributionalVFRewardModel
     else:
         try:
             return _get_reward_model_cls_from_name(name=name)
@@ -96,6 +103,8 @@ def make_reward_model_config(reward_type: str, **kwargs) -> RewardModelConfig:
         return RobometerConfig(**kwargs)
     elif reward_type == "topreward":
         return TOPRewardConfig(**kwargs)
+    elif reward_type == "distributional_value_function":
+        return DistributionalVFConfig(**kwargs)
     else:
         try:
             config_cls = RewardModelConfig.get_choice_class(reward_type)
@@ -192,6 +201,16 @@ def make_reward_pre_post_processors(
             dataset_stats=kwargs.get("dataset_stats"),
         )
 
+    elif isinstance(reward_cfg, DistributionalVFConfig):
+        from lerobot.rewards.distributional_value_function.processor_distributional_value_function import (
+            make_distributional_vf_pre_post_processors,
+        )
+
+        return make_distributional_vf_pre_post_processors(
+            config=reward_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
     else:
         try:
             processors = _make_processors_from_reward_model_config(

src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py

@@ -65,7 +65,13 @@ 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,
         )
 
@@ -80,7 +86,13 @@ 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,
         )
 

src/lerobot/robots/rebot_b601_follower/config_rebot_b601_follower.py

@@ -65,18 +65,33 @@ class RebotB601FollowerConfig:
         }
     )
 
-    # 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)
+    # 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]
+    )
 
-    # Torque/current ratio for the gripper's FORCE_POS mode, in range [0, 1].
-    gripper_torque_ratio: float = 0.1
+    # 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
 
     # Soft joint limits (degrees). These are clipped against on every action.
     joint_limits: dict[str, tuple[float, float]] = field(
         default_factory=lambda: {
-            "shoulder_pan": (-145.0, 145.0),
-            "shoulder_lift": (-170.0, 1.0),
+            "shoulder_pan": (-150.0, 150.0),
+            "shoulder_lift": (-200.0, 1.0),
             "elbow_flex": (-200.0, 1.0),
             "wrist_flex": (-80.0, 90.0),
             "wrist_yaw": (-90.0, 90.0),

src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py

@@ -174,11 +174,25 @@ 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():
-            target_mode = (
-                MotorBridgeMode.FORCE_POS if motor_name == GRIPPER_MOTOR else MotorBridgeMode.POS_VEL
-            )
+            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
             for attempt in range(_ENSURE_MODE_RETRIES + 1):
                 try:
                     motor.ensure_mode(target_mode)
@@ -264,22 +278,34 @@ 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:
-                motor.send_force_pos(pos_rad, vel_rad, self.config.gripper_torque_ratio)
+                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)
             else:
-                motor.send_pos_vel(pos_rad, vel_rad)
+                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))
 
         return {f"{motor}.pos": val for motor, val in goal_pos.items()}
 

src/lerobot/rollout/configs.py

@@ -106,6 +106,8 @@ class DAggerKeyboardConfig:
     pause_resume: str = "space"
     correction: str = "tab"
     upload: str = "enter"
+    success: str = "s"
+    failure: str = "f"
 
 
 @dataclass
@@ -119,6 +121,8 @@ class DAggerPedalConfig:
     pause_resume: str = "KEY_A"
     correction: str = "KEY_B"
     upload: str = "KEY_C"
+    success: str = "KEY_D"
+    failure: str = "KEY_E"
 
 
 @RolloutStrategyConfig.register_subclass("episodic")
@@ -165,6 +169,10 @@ class DAggerStrategyConfig(RolloutStrategyConfig):
     2. **correction** — toggle human correction recording.
     3. **upload** — push dataset to hub on demand (corrections-only mode).
 
+    Episode success labeling:
+    4. **success** — mark current episode as successful.
+    5. **failure** — mark current episode as failed.
+
     When ``record_autonomous=False`` (default) only human-correction windows
     are recorded — each correction becomes its own episode.  Set to ``True``
     to record both autonomous and correction frames with size-based episode
@@ -226,11 +234,14 @@ class RolloutConfig:
     device: str | None = None
     task: str = ""
     display_data: bool = False
-    # Display data on a remote Rerun server
+    # Visualization backend used when display_data is True: "rerun" or "foxglove".
+    display_mode: str = "rerun"
+    # For "rerun": IP of a remote server to send to. For "foxglove": interface to bind the WebSocket
+    # server to (127.0.0.1 for local only, 0.0.0.0 for all interfaces).
     display_ip: str | None = None
-    # Port of the remote Rerun server
+    # For "rerun": port of the remote server. For "foxglove": port to bind the WebSocket server to.
     display_port: int | None = None
-    # Whether to display compressed images in Rerun
+    # Whether to display compressed (JPEG) images instead of raw frames
     display_compressed_images: bool = False
     # Use vocal synthesis to read events
     play_sounds: bool = True

src/lerobot/rollout/context.py

@@ -320,7 +320,9 @@ def build_rollout_context(
             raise ValueError(
                 f"Visual feature mismatch between policy and robot hardware.\n"
                 f"Policy expects: {expected_visuals}\n"
-                f"Robot provides: {provided_visuals}"
+                f"Robot provides: {provided_visuals}\n"
+                f"Use --rename_map to map camera names, e.g. "
+                f"""--rename_map='{{"observation.images.top": "observation.images.cam0"}}'"""
             )
 
     # --- 5. Dataset -------------
@@ -348,6 +350,11 @@ def build_rollout_context(
                     "shape": (1,),
                     "names": None,
                 }
+                dataset_features["next.success"] = {
+                    "dtype": "bool",
+                    "shape": (1,),
+                    "names": None,
+                }
 
             repo_name = cfg.dataset.repo_id.split("/", 1)[-1]
             if not repo_name.startswith("rollout_"):

src/lerobot/rollout/strategies/core.py

@@ -26,7 +26,7 @@ from lerobot.utils.action_interpolator import ActionInterpolator
 from lerobot.utils.constants import OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame
 from lerobot.utils.robot_utils import precise_sleep
-from lerobot.utils.visualization_utils import log_rerun_data
+from lerobot.utils.visualization_utils import log_visualization_data
 
 from ..inference import InferenceEngine
 
@@ -162,11 +162,12 @@ class RolloutStrategy(abc.ABC):
         action_dict: dict | None,
         runtime_ctx: RuntimeContext,
     ) -> None:
-        """Log observation/action telemetry to Rerun if display_data is enabled."""
+        """Log observation/action telemetry to the visualization backend if display_data is enabled."""
         cfg = runtime_ctx.cfg
         if not cfg.display_data:
             return
-        log_rerun_data(
+        log_visualization_data(
+            cfg.display_mode,
             observation=obs_processed,
             action=action_dict,
             compress_images=cfg.display_compressed_images,

src/lerobot/rollout/strategies/dagger.py

@@ -112,6 +112,14 @@ class DAggerEvents:
         # Session-level flags
         self.stop_recording = Event()
         self.upload_requested = Event()
+        # Set when operator presses success/failure key to end the current episode.
+        self.save_episode_requested = Event()
+
+        # Episode success labeling
+        self._episode_success: bool | None = None
+
+        # Episode success labeling
+        self._episode_success: bool | None = None
 
     # -- Thread-safe phase access ------------------------------------------
 
@@ -155,7 +163,43 @@ class DAggerEvents:
         with self._lock:
             self._phase = DAggerPhase.AUTONOMOUS
             self._pending_transition = None
+            self._episode_success = None
         self.upload_requested.clear()
+        self.save_episode_requested.clear()
+
+    def mark_success(self) -> None:
+        """Mark the current episode as successful (called from input threads)."""
+        with self._lock:
+            self._episode_success = True
+
+    def mark_failure(self) -> None:
+        """Mark the current episode as failed (called from input threads)."""
+        with self._lock:
+            self._episode_success = False
+
+    def consume_episode_success(self) -> bool | None:
+        """Consume and reset the episode success label. Returns None if unlabeled."""
+        with self._lock:
+            result = self._episode_success
+            self._episode_success = None
+            return result
+
+    def mark_success(self) -> None:
+        """Mark the current episode as successful (called from input threads)."""
+        with self._lock:
+            self._episode_success = True
+
+    def mark_failure(self) -> None:
+        """Mark the current episode as failed (called from input threads)."""
+        with self._lock:
+            self._episode_success = False
+
+    def consume_episode_success(self) -> bool | None:
+        """Consume and reset the episode success label. Returns None if unlabeled."""
+        with self._lock:
+            result = self._episode_success
+            self._episode_success = None
+            return result
 
 
 # ---------------------------------------------------------------------------
@@ -186,12 +230,20 @@ def _init_dagger_keyboard(events: DAggerEvents, cfg: DAggerKeyboardConfig):
             events.request_transition(key_to_event[name])
         if name == cfg.upload:
             events.upload_requested.set()
+        if name == cfg.success:
+            events.mark_success()
+            events.save_episode_requested.set()
+            logger.info("Episode marked as SUCCESS — saving")
+        if name == cfg.failure:
+            events.mark_failure()
+            events.save_episode_requested.set()
+            logger.info("Episode marked as FAILURE — saving")
 
     return create_key_listener(
         dispatch,
         controls_help=(
             f"pause_resume='{cfg.pause_resume}', correction='{cfg.correction}', "
-            f"upload='{cfg.upload}', ESC=stop"
+            f"upload='{cfg.upload}', success='{cfg.success}', failure='{cfg.failure}', ESC=stop"
         ),
     )
 
@@ -211,6 +263,12 @@ def _init_dagger_pedal(events: DAggerEvents, cfg: DAggerPedalConfig):
             events.request_transition(code_to_event[code])
         if code == cfg.upload:
             events.upload_requested.set()
+        if code == cfg.success:
+            events.mark_success()
+            logger.info("Episode marked as SUCCESS (pedal)")
+        if code == cfg.failure:
+            events.mark_failure()
+            logger.info("Episode marked as FAILURE (pedal)")
 
     logger.info("Initializing DAgger foot pedal listener (device=%s)", cfg.device_path)
     return start_pedal_listener(on_press, device_path=cfg.device_path)
@@ -313,6 +371,31 @@ class DAggerStrategy(RolloutStrategy):
         )
         logger.info("DAgger strategy teardown complete")
 
+    # ------------------------------------------------------------------
+    # Episode success labeling
+    # ------------------------------------------------------------------
+
+    def _stamp_episode_success(self, dataset) -> None:
+        """Set next.success on the terminal frame based on operator label.
+
+        Called just before save_episode(). If the operator pressed the success
+        key during this episode, the last frame's next.success is set to True.
+        Otherwise all frames remain False (unlabeled = assumed failure).
+        """
+        buf = dataset.writer.episode_buffer
+        if buf is None:
+            return
+
+        success_buf = buf.get("next.success")
+        if not success_buf:
+            return
+
+        label = self._events.consume_episode_success()
+
+        if label:
+            success_buf[-1] = np.array([True], dtype=bool)
+            logger.info("Terminal frame stamped next.success=True")
+
     # ------------------------------------------------------------------
     # Continuous recording mode (record_autonomous=True)
     # ------------------------------------------------------------------
@@ -350,7 +433,12 @@ class DAggerStrategy(RolloutStrategy):
         episode_start = time.perf_counter()
         episodes_since_push = 0
         episode_duration_s = self._episode_duration_s
-        logger.info("DAgger continuous recording started (episode_duration=%.0fs)", episode_duration_s)
+        num_episodes = self.config.num_episodes
+        logger.info(
+            "DAgger continuous recording started (episode_duration=%.0fs, target=%s eps)",
+            episode_duration_s,
+            num_episodes if num_episodes is not None else "∞",
+        )
 
         with VideoEncodingManager(dataset):
             try:
@@ -399,6 +487,7 @@ class DAggerStrategy(RolloutStrategy):
                                 **action_frame,
                                 "task": task_str,
                                 "intervention": np.array([True], dtype=bool),
+                                "next.success": np.array([False], dtype=bool),
                             }
                             dataset.add_frame(frame)
                         record_tick += 1
@@ -427,23 +516,32 @@ class DAggerStrategy(RolloutStrategy):
                                     **action_frame,
                                     "task": task_str,
                                     "intervention": np.array([False], dtype=bool),
+                                    "next.success": np.array([False], dtype=bool),
                                 }
                                 dataset.add_frame(frame)
                             record_tick += 1
 
-                    # Episode rotation derived from the video file-size target.
-                    # Saving is deferred while a correction is ongoing so the
-                    # episode boundary lands on a clean autonomous frame.
+                    # Episode rotation: either the operator pressed success/failure,
+                    # or the video file-size target was reached.
+                    # Defer the save while a correction is ongoing so the episode
+                    # boundary lands on a clean autonomous frame. The event stays
+                    # set until we actually save, so it won't be lost.
+                    manual_save = events.save_episode_requested.is_set()
+
                     elapsed = time.perf_counter() - episode_start
-                    if elapsed >= episode_duration_s and phase != DAggerPhase.CORRECTING:
+                    if (manual_save or elapsed >= episode_duration_s) and phase != DAggerPhase.CORRECTING:
+                        if manual_save:
+                            events.save_episode_requested.clear()
                         with self._episode_lock:
+                            self._stamp_episode_success(dataset)
                             dataset.save_episode()
                         episodes_since_push += 1
                         self._needs_push.set()
+                        save_reason = "manual save" if manual_save else f"elapsed {elapsed:.1f}s"
                         logger.info(
-                            "Episode saved (total: %d, elapsed: %.1fs)",
+                            "Episode saved (%s, total: %d)",
+                            save_reason,
                             dataset.num_episodes,
-                            elapsed,
                         )
                         log_say(f"Episode {dataset.num_episodes} saved", play_sounds)
 
@@ -451,6 +549,25 @@ class DAggerStrategy(RolloutStrategy):
                             self._background_push(dataset, cfg)
                             episodes_since_push = 0
 
+                        if num_episodes is not None and dataset.num_episodes >= num_episodes:
+                            logger.info("Target episode count reached (%d), stopping session", num_episodes)
+                            log_say(f"All {num_episodes} episodes collected", play_sounds)
+                            events.stop_recording.set()
+                            break
+
+                        # Pause after manual save: stop the policy, return robot to
+                        # initial position, and wait for the operator to reset the
+                        # environment and press SPACE.
+                        if manual_save:
+                            engine.pause()
+                            events.phase = DAggerPhase.PAUSED
+                            self._return_to_initial_position(ctx.hardware)
+                            last_action = None
+                            logger.info(
+                                "Episode saved — paused for environment reset. Press SPACE to start next episode."
+                            )
+                            log_say("Reset the environment, then press space", play_sounds)
+
                         episode_start = time.perf_counter()
 
                     dt = time.perf_counter() - loop_start
@@ -465,10 +582,13 @@ class DAggerStrategy(RolloutStrategy):
                 logger.info("DAgger continuous control loop ended — pausing engine")
                 engine.pause()
                 with contextlib.suppress(Exception):
-                    with self._episode_lock:
-                        dataset.save_episode()
-                    self._needs_push.set()
-                    logger.info("Final in-progress episode saved")
+                    buf = dataset.writer.episode_buffer
+                    if buf and any(len(v) > 0 for v in buf.values() if isinstance(v, list)):
+                        with self._episode_lock:
+                            self._stamp_episode_success(dataset)
+                            dataset.save_episode()
+                        self._needs_push.set()
+                        logger.info("Final in-progress episode saved")
 
     # ------------------------------------------------------------------
     # Corrections-only mode (record_autonomous=False)
@@ -540,6 +660,7 @@ class DAggerStrategy(RolloutStrategy):
                         # Correction ended -> save episode (blocking if not streaming)
                         if old_phase == DAggerPhase.CORRECTING and new_phase == DAggerPhase.PAUSED:
                             with self._episode_lock:
+                                self._stamp_episode_success(dataset)
                                 dataset.save_episode()
                             recorded += 1
                             self._needs_push.set()
@@ -581,6 +702,7 @@ class DAggerStrategy(RolloutStrategy):
                                     **action_frame,
                                     "task": task_str,
                                     "intervention": np.array([True], dtype=bool),
+                                    "next.success": np.array([False], dtype=bool),
                                 }
                             )
                         record_tick += 1
@@ -614,10 +736,13 @@ class DAggerStrategy(RolloutStrategy):
                 logger.info("DAgger corrections-only loop ended — pausing engine")
                 engine.pause()
                 with contextlib.suppress(Exception):
-                    with self._episode_lock:
-                        dataset.save_episode()
-                    self._needs_push.set()
-                    logger.info("Final in-progress episode saved")
+                    buf = dataset.writer.episode_buffer
+                    if buf and any(len(v) > 0 for v in buf.values() if isinstance(v, list)):
+                        with self._episode_lock:
+                            self._stamp_episode_success(dataset)
+                            dataset.save_episode()
+                        self._needs_push.set()
+                        logger.info("Final in-progress episode saved")
 
     # ------------------------------------------------------------------
     # State-machine transition side-effects

src/lerobot/rollout/strategies/episodic.py

@@ -44,7 +44,7 @@ from lerobot.utils.feature_utils import build_dataset_frame
 from lerobot.utils.keyboard_input import init_keyboard_listener
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import log_rerun_data
+from lerobot.utils.visualization_utils import log_visualization_data
 
 from ..configs import EpisodicStrategyConfig
 from ..context import RolloutContext
@@ -171,6 +171,7 @@ class EpisodicStrategy(RolloutStrategy):
                             fps=fps,
                             control_time_s=reset_time_s,
                             display_data=cfg.display_data,
+                            display_mode=cfg.display_mode,
                             display_compressed=display_compressed,
                         )
 
@@ -259,6 +260,7 @@ class EpisodicStrategy(RolloutStrategy):
         fps: float,
         control_time_s: float,
         display_data: bool,
+        display_mode: str,
         display_compressed: bool,
     ) -> None:
         """Reset-phase loop: teleop drives the robot if available, no recording."""
@@ -288,7 +290,8 @@ class EpisodicStrategy(RolloutStrategy):
 
                 if display_data:
                     obs_processed = processors.robot_observation_processor(obs)
-                    log_rerun_data(
+                    log_visualization_data(
+                        display_mode,
                         observation=obs_processed,
                         action=act_teleop,
                         compress_images=display_compressed,

src/lerobot/scripts/lerobot_annotate.py

@@ -34,6 +34,7 @@ from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConf
 from lerobot.annotations.steerable_pipeline.executor import Executor
 from lerobot.annotations.steerable_pipeline.frames import make_frame_provider
 from lerobot.annotations.steerable_pipeline.modules import (
+    AdvantageModule,
     GeneralVqaModule,
     InterjectionsAndSpeechModule,
     PlanSubtasksMemoryModule,
@@ -63,17 +64,23 @@ def annotate(cfg: AnnotationPipelineConfig) -> None:
     root = _resolve_root(cfg)
     logger.info("annotate: root=%s", root)
 
-    vlm = make_vlm_client(cfg.vlm)
-    frame_provider = make_frame_provider(root, camera_key=cfg.vlm.camera_key, video_backend=cfg.video_backend)
+    needs_vlm = cfg.plan.enabled or cfg.interjections.enabled or cfg.vqa.enabled
+    vlm = make_vlm_client(cfg.vlm) if needs_vlm else None
+    frame_provider = (
+        make_frame_provider(root, camera_key=cfg.vlm.camera_key, video_backend=cfg.video_backend)
+        if needs_vlm
+        else None
+    )
     # Surface the resolved cameras up front so a silent vqa-module no-op
     # is obvious in job output rather than discovered post-hoc by counting
     # parquet rows.
-    cam_keys = list(getattr(frame_provider, "camera_keys", []) or [])
-    logger.info(
-        "annotate: frame_provider default camera=%r, all cameras=%s",
-        getattr(frame_provider, "camera_key", None),
-        cam_keys,
-    )
+    cam_keys = list(getattr(frame_provider, "camera_keys", []) or []) if frame_provider else []
+    if frame_provider:
+        logger.info(
+            "annotate: frame_provider default camera=%r, all cameras=%s",
+            getattr(frame_provider, "camera_key", None),
+            cam_keys,
+        )
     if cfg.vqa.enabled and not cam_keys:
         logger.warning(
             "annotate: the vqa module is enabled but no cameras were "
@@ -81,14 +88,27 @@ def annotate(cfg: AnnotationPipelineConfig) -> None:
             "meta/info.json for observation.images.* features, or pass "
             "--vlm.camera_key=<key> to seed the cameras list."
         )
-    plan = PlanSubtasksMemoryModule(vlm=vlm, config=cfg.plan, frame_provider=frame_provider)
-    interjections = InterjectionsAndSpeechModule(
-        vlm=vlm, config=cfg.interjections, seed=cfg.seed, frame_provider=frame_provider
+    plan = (
+        PlanSubtasksMemoryModule(vlm=vlm, config=cfg.plan, frame_provider=frame_provider)
+        if needs_vlm
+        else None
     )
-    vqa = GeneralVqaModule(vlm=vlm, config=cfg.vqa, seed=cfg.seed, frame_provider=frame_provider)
+    interjections = (
+        InterjectionsAndSpeechModule(
+            vlm=vlm, config=cfg.interjections, seed=cfg.seed, frame_provider=frame_provider
+        )
+        if needs_vlm
+        else None
+    )
+    vqa = (
+        GeneralVqaModule(vlm=vlm, config=cfg.vqa, seed=cfg.seed, frame_provider=frame_provider)
+        if needs_vlm
+        else None
+    )
+    advantage = AdvantageModule(config=cfg.advantage)
     writer = LanguageColumnsWriter()
     validator = StagingValidator(
-        dataset_camera_keys=tuple(getattr(frame_provider, "camera_keys", []) or []) or None,
+        dataset_camera_keys=tuple(cam_keys) or None,
     )
 
     executor = Executor(
@@ -96,6 +116,7 @@ def annotate(cfg: AnnotationPipelineConfig) -> None:
         plan=plan,
         interjections=interjections,
         vqa=vqa,
+        advantage=advantage,
         writer=writer,
         validator=validator,
     )

src/lerobot/scripts/lerobot_compute_returns.py

@@ -0,0 +1,382 @@
+#!/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.
+
+"""Compute per-frame ``is_terminal`` and ``mc_return`` for a LeRobot dataset.
+
+Implements the sparse reward function from pi*0.6 / RECAP (Eq. 5):
+
+    r_t = -1          for non-terminal steps
+    r_T = 0           for terminal success
+    r_T = -C_fail     for terminal failure
+
+Monte Carlo returns are the cumulative sum from each step to the end of
+the episode, normalized by ``max_episode_length`` so that values are bounded
+to approximately (-1, 0).
+
+The columns are written directly into the dataset's parquet data shards as
+flat per-frame scalars. These serve as training targets for the distributional
+value function.
+
+Usage:
+    # Compute returns using the default "next.success" column (from lerobot-eval/rollout)
+    lerobot-compute-returns \\
+        --dataset-repo-id lerobot/aloha_sim_insertion_human_image
+
+    # Override: treat all episodes as successful (demo-only datasets)
+    lerobot-compute-returns \\
+        --dataset-repo-id lerobot/aloha_sim_insertion_human_image \\
+        --default-success true
+
+    # Custom success key, failure penalty, and discount
+    lerobot-compute-returns \\
+        --dataset-repo-id my_org/my_dataset \\
+        --success-key episode_success \\
+        --c-fail 100 \\
+        --gamma 0.99
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+IS_TERMINAL_COL = "is_terminal"
+MC_RETURN_COL = "mc_return"
+
+
+@dataclass
+class ComputeReturnsConfig:
+    """Configuration for the returns computation script."""
+
+    dataset_repo_id: str = ""
+    root: str | None = None
+    success_key: str = "next.success"
+    default_success: bool | None = None
+    max_episode_length: int | None = None
+    c_fail: float = 50.0
+    gamma: float = 1.0
+    episodes: list[int] = field(default_factory=list)
+    force: bool = False
+
+
+def _get_episode_success(
+    episode_table: pa.Table,
+    success_key: str,
+    default_success: bool | None,
+) -> bool:
+    """Determine whether an episode was successful.
+
+    Priority:
+    1. If ``default_success`` is set, use it unconditionally.
+    2. Look for ``success_key`` in the parquet columns and reduce with any().
+    3. Fall back to True (assume success for demo datasets).
+    """
+    if default_success is not None:
+        return default_success
+
+    if success_key in episode_table.column_names:
+        col = episode_table.column(success_key)
+        for val in col:
+            py_val = val.as_py()
+            if isinstance(py_val, bool) and py_val:
+                return True
+            if isinstance(py_val, (int, float)) and py_val:
+                return True
+        return False
+
+    return True
+
+
+def compute_episode_returns(
+    num_frames: int,
+    success: bool,
+    c_fail: float,
+    gamma: float,
+    max_episode_length: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Compute is_terminal and mc_return arrays for a single episode.
+
+    Args:
+        num_frames: Number of frames in the episode.
+        success: Whether the episode ended successfully.
+        c_fail: Failure penalty constant.
+        gamma: Discount factor (1.0 = undiscounted).
+        max_episode_length: Normalization horizon H.
+
+    Returns:
+        Tuple of (is_terminal, mc_return) arrays, each of length num_frames.
+    """
+    horizon = max_episode_length
+
+    rewards = np.full(num_frames, -1.0 / horizon, dtype=np.float64)
+
+    if success:
+        rewards[-1] = 0.0
+    else:
+        rewards[-1] = -c_fail / horizon
+
+    is_terminal = np.zeros(num_frames, dtype=bool)
+    is_terminal[-1] = True
+
+    if gamma == 1.0:
+        # Reverse cumulative sum
+        mc_return = np.cumsum(rewards[::-1])[::-1].astype(np.float32)
+    else:
+        mc_return = np.zeros(num_frames, dtype=np.float64)
+        mc_return[-1] = rewards[-1]
+        for t in range(num_frames - 2, -1, -1):
+            mc_return[t] = rewards[t] + gamma * mc_return[t + 1]
+        mc_return = mc_return.astype(np.float32)
+
+    return is_terminal, mc_return
+
+
+def compute_returns(config: ComputeReturnsConfig) -> Path:
+    """Compute returns and write them into parquet shards."""
+    from lerobot.datasets import LeRobotDataset
+
+    logger.info(f"Loading dataset: {config.dataset_repo_id}")
+    kwargs = {"repo_id": config.dataset_repo_id, "download_videos": False}
+    if config.root:
+        kwargs["root"] = config.root
+    dataset = LeRobotDataset(**kwargs)
+
+    meta = dataset.meta
+    root = Path(meta.root)
+    logger.info(f"Dataset root: {root}")
+    logger.info(f"Episodes: {meta.total_episodes}, Frames: {meta.total_frames}")
+
+    episode_indices = config.episodes if config.episodes else list(range(meta.total_episodes))
+
+    if config.max_episode_length is not None:
+        max_ep_len = config.max_episode_length
+    else:
+        max_ep_len = max(int(meta.episodes[i]["length"]) for i in episode_indices)
+    logger.info(f"Normalization horizon (max_episode_length): {max_ep_len}")
+
+    parquet_files_to_rewrite: dict[Path, list[int]] = {}
+    for ep_idx in episode_indices:
+        rel_path = meta.get_data_file_path(ep_idx)
+        abs_path = root / rel_path
+        parquet_files_to_rewrite.setdefault(abs_path, []).append(ep_idx)
+
+    logger.info(f"Parquet shards to rewrite: {len(parquet_files_to_rewrite)}")
+
+    for parquet_path, ep_indices_in_file in tqdm(parquet_files_to_rewrite.items(), desc="Processing shards"):
+        table = pq.read_table(parquet_path)
+
+        if not config.force and IS_TERMINAL_COL in table.column_names:
+            logger.info(f"Skipping {parquet_path.name} (already has {IS_TERMINAL_COL})")
+            continue
+
+        all_is_terminal = np.zeros(len(table), dtype=bool)
+        all_mc_return = np.zeros(len(table), dtype=np.float32)
+
+        episode_col = table.column("episode_index").to_pylist()
+
+        for ep_idx in ep_indices_in_file:
+            ep_info = meta.episodes[ep_idx]
+            ep_from = int(ep_info["dataset_from_index"])
+            ep_to = int(ep_info["dataset_to_index"])
+            ep_len = ep_to - ep_from
+
+            mask = np.array([v == ep_idx for v in episode_col], dtype=bool)
+            local_indices = np.where(mask)[0]
+
+            if len(local_indices) != ep_len:
+                logger.warning(
+                    f"Episode {ep_idx}: expected {ep_len} frames in shard, "
+                    f"found {len(local_indices)}. Using found count."
+                )
+                ep_len = len(local_indices)
+
+            if ep_len == 0:
+                continue
+
+            ep_subtable = table.filter(mask)
+            success = _get_episode_success(ep_subtable, config.success_key, config.default_success)
+
+            is_terminal, mc_return = compute_episode_returns(
+                num_frames=ep_len,
+                success=success,
+                c_fail=config.c_fail,
+                gamma=config.gamma,
+                max_episode_length=max_ep_len,
+            )
+
+            all_is_terminal[local_indices] = is_terminal
+            all_mc_return[local_indices] = mc_return
+
+        if IS_TERMINAL_COL in table.column_names:
+            table = table.drop(IS_TERMINAL_COL)
+        if MC_RETURN_COL in table.column_names:
+            table = table.drop(MC_RETURN_COL)
+
+        table = table.append_column(IS_TERMINAL_COL, pa.array(all_is_terminal))
+        table = table.append_column(MC_RETURN_COL, pa.array(all_mc_return))
+
+        pq.write_table(table, parquet_path)
+
+    _update_info_json(root, meta)
+
+    logger.info("Done. Columns written: is_terminal, mc_return")
+    return root
+
+
+def _update_info_json(root: Path, meta) -> None:
+    """Add is_terminal and mc_return to the dataset's info.json features."""
+    info_path = root / "meta" / "info.json"
+    if not info_path.exists():
+        logger.warning(f"info.json not found at {info_path}, skipping metadata update.")
+        return
+
+    info = json.loads(info_path.read_text())
+    features = info.get("features", {})
+    changed = False
+
+    if IS_TERMINAL_COL not in features:
+        features[IS_TERMINAL_COL] = {
+            "dtype": "bool",
+            "shape": [1],
+            "names": None,
+        }
+        changed = True
+
+    if MC_RETURN_COL not in features:
+        features[MC_RETURN_COL] = {
+            "dtype": "float32",
+            "shape": [1],
+            "names": None,
+        }
+        changed = True
+
+    if changed:
+        info["features"] = features
+        info_path.write_text(json.dumps(info, indent=2) + "\n")
+        logger.info("Updated meta/info.json with is_terminal and mc_return features.")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Compute per-frame is_terminal and mc_return for a LeRobot dataset.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    # Use the 'success' column from the dataset
+    lerobot-compute-returns --dataset-repo-id lerobot/aloha_sim_insertion_human_image
+
+    # Override all episodes as successful (demo-only data)
+    lerobot-compute-returns --dataset-repo-id my_org/my_dataset --default-success true
+
+    # Custom failure penalty
+    lerobot-compute-returns --dataset-repo-id my_org/my_dataset --c-fail 100
+        """,
+    )
+    parser.add_argument(
+        "--dataset-repo-id",
+        type=str,
+        required=True,
+        help="HuggingFace dataset repo id or local path.",
+    )
+    parser.add_argument(
+        "--root",
+        type=str,
+        default=None,
+        help="Local root directory override for the dataset.",
+    )
+    parser.add_argument(
+        "--success-key",
+        type=str,
+        default="next.success",
+        help="Column name in parquet that indicates episode success (default: 'next.success').",
+    )
+    parser.add_argument(
+        "--default-success",
+        type=str,
+        default=None,
+        choices=["true", "false"],
+        help="Override success for all episodes ('true' or 'false').",
+    )
+    parser.add_argument(
+        "--max-episode-length",
+        type=int,
+        default=None,
+        help="Normalization horizon H. If not set, uses max episode length in dataset.",
+    )
+    parser.add_argument(
+        "--c-fail",
+        type=float,
+        default=50.0,
+        help="Failure penalty constant (default: 50.0).",
+    )
+    parser.add_argument(
+        "--gamma",
+        type=float,
+        default=1.0,
+        help="Discount factor (default: 1.0, undiscounted).",
+    )
+    parser.add_argument(
+        "--episodes",
+        type=int,
+        nargs="+",
+        default=None,
+        help="Process only these episode indices (default: all).",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Overwrite existing is_terminal/mc_return columns.",
+    )
+
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+
+    default_success = None
+    if args.default_success is not None:
+        default_success = args.default_success.lower() == "true"
+
+    config = ComputeReturnsConfig(
+        dataset_repo_id=args.dataset_repo_id,
+        root=args.root,
+        success_key=args.success_key,
+        default_success=default_success,
+        max_episode_length=args.max_episode_length,
+        c_fail=args.c_fail,
+        gamma=args.gamma,
+        episodes=args.episodes or [],
+        force=args.force,
+    )
+
+    root = compute_returns(config)
+    logger.info(f"Returns computed and written to: {root}")
+    logger.info(f"  Columns added: {IS_TERMINAL_COL}, {MC_RETURN_COL}")
+    logger.info("To train the distributional value function, these columns")
+    logger.info("will be read as flat batch keys during training.")
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/scripts/lerobot_dataset_viz.py

@@ -59,6 +59,18 @@ distant$ lerobot-dataset-viz \
 local$ rerun rerun+http://IP:GRPC_PORT/proxy
 ```
 
+- Visualize data in Foxglove with a seekable, scrubbable timeline:
+```
+local$ lerobot-dataset-viz \
+    --repo-id lerobot/pusht \
+    --episode-index 0 \
+    --display-mode foxglove
+
+# then open the Foxglove app and connect to ws://127.0.0.1:8765
+```
+This starts a Foxglove WebSocket server that serves the episode on demand from the on-disk dataset,
+so you can play/pause and scrub anywhere in the episode using Foxglove's playback controls.
+
 """
 
 import argparse
@@ -72,10 +84,29 @@ import torch
 import torch.utils.data
 import tqdm
 
+from lerobot.configs import DEPTH_MILLIMETER_UNIT
 from lerobot.datasets import LeRobotDataset
-from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD
+from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
 from lerobot.utils.utils import init_logging
 
+DEFAULT_FOXGLOVE_PORT = 8765
+DEFAULT_RERUN_PORT = 9090
+
+
+def get_feature_names(dataset: LeRobotDataset, key: str) -> list[str]:
+    """Return per-dimension names for a feature from the dataset metadata.
+
+    Only flat-list ``names`` metadata is used. Dict-style ``names`` and missing names fall back to ``{key}_{i}`` indices.
+    """
+    feature = dataset.features[key]
+    dim = feature["shape"][-1]
+
+    names = feature.get("names")
+    if isinstance(names, list) and len(names) == dim:
+        return [str(name) for name in names]
+
+    return [f"{key}_{d}" for d in range(dim)]
+
 
 def check_chw_float32(frame: torch.Tensor) -> None:
     """
@@ -93,10 +124,35 @@ def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:
     return hwc_uint8_numpy
 
 
-def to_hwc_uint16_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:
+def to_hwc_float32_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:
     check_chw_float32(chw_float32_torch)
-    hwc_uint16_numpy = chw_float32_torch.round().type(torch.uint16).permute(1, 2, 0).numpy()
-    return hwc_uint16_numpy
+    hwc_float32_numpy = chw_float32_torch.permute(1, 2, 0).numpy()
+    return hwc_float32_numpy
+
+
+def build_blueprint_from_dataset(dataset: LeRobotDataset):
+    """Build a Rerun blueprint laying out camera images and time series for the given dataset.
+
+    Camera images and scalar signals (action, state, reward, done, success) are arranged in a grid.
+    The per-dimension series names for ``action`` and ``state`` are applied directly
+    via blueprint overrides.
+    """
+    import rerun as rr
+    import rerun.blueprint as rrb
+
+    views = [rrb.Spatial2DView(origin=key, name=key) for key in dataset.meta.camera_keys]
+
+    # Style multi-dimensional signals (action, state) with per-dimension names.
+    for origin, key in ((ACTION, ACTION), ("state", OBS_STATE)):
+        if key in dataset.features:
+            names = get_feature_names(dataset, key)
+            styling = rr.SeriesLines(names=names)
+            views.append(rrb.TimeSeriesView(origin=origin, name=origin, overrides={origin: styling}))
+    for key in (DONE, REWARD, SUCCESS):
+        if key in dataset.features:
+            views.append(rrb.TimeSeriesView(origin=key, name=key))
+
+    return rrb.Blueprint(rrb.Grid(*views))
 
 
 def visualize_dataset(
@@ -105,13 +161,30 @@ def visualize_dataset(
     batch_size: int = 32,
     num_workers: int = 0,
     mode: str = "local",
-    web_port: int = 9090,
+    web_port: int | None = None,
     grpc_port: int = 9876,
     save: bool = False,
     output_dir: Path | None = None,
     display_compressed_images: bool = False,
+    display_mode: str = "rerun",
+    host: str = "127.0.0.1",
+    autoplay: bool = True,
     **kwargs,
 ) -> Path | None:
+    if display_mode == "foxglove":
+        from lerobot.utils.foxglove_visualization import serve_foxglove_dataset_playback
+
+        logging.info("Starting Foxglove server")
+        serve_foxglove_dataset_playback(
+            dataset,
+            episode_index,
+            host=host,
+            port=web_port if web_port is not None else DEFAULT_FOXGLOVE_PORT,
+            compress_images=display_compressed_images,
+            autoplay=autoplay,
+        )
+        return None
+
     if save:
         assert output_dir is not None, (
             "Set an output directory where to write .rrd files with `--output-dir path/to/directory`."
@@ -137,7 +210,8 @@ def visualize_dataset(
     import rerun as rr
 
     spawn_local_viewer = mode == "local" and not save
-    rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer)
+    blueprint = build_blueprint_from_dataset(dataset)
+    rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer, default_blueprint=blueprint)
 
     # Manually call python garbage collector after `rr.init` to avoid hanging in a blocking flush
     # when iterating on a dataloader with `num_workers` > 0
@@ -147,14 +221,23 @@ def visualize_dataset(
     if mode == "distant":
         server_uri = rr.serve_grpc(grpc_port=grpc_port)
         logging.info(f"Connect to a Rerun Server: rerun rerun+http://IP:{grpc_port}/proxy")
-        rr.serve_web_viewer(open_browser=False, web_port=web_port, connect_to=server_uri)
+        rr.serve_web_viewer(
+            open_browser=False,
+            web_port=web_port if web_port is not None else DEFAULT_RERUN_PORT,
+            connect_to=server_uri,
+        )
 
     logging.info("Logging to Rerun")
 
+    # Depth frames and stats are dequantized to the dataset's depth_output_unit on load.
+    depth_meter = 1000.0 if dataset.depth_output_unit == DEPTH_MILLIMETER_UNIT else 1.0
+
     # Use the dataset's q01/q99 depth statistics for robust depth range bounds
     depth_ranges = {}
     for key in dataset.meta.depth_keys:
-        stats = dataset.meta.stats[key]
+        stats = (dataset.meta.stats or {}).get(key)
+        if not stats:
+            continue
         lo = stats["q01"] if "q01" in stats else stats["min"]
         hi = stats["q99"] if "q99" in stats else stats["max"]
         depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item()))
@@ -163,19 +246,21 @@ def visualize_dataset(
     for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
         if first_index is None:
             first_index = batch["index"][0].item()
+
         # iterate over the batch
         for i in range(len(batch["index"])):
             rr.set_time("frame_index", sequence=batch["index"][i].item() - first_index)
             rr.set_time("timestamp", timestamp=batch["timestamp"][i].item())
 
-            # display each camera image
+            # display each camera image (or depth map)
             for key in dataset.meta.camera_keys:
                 if key in dataset.meta.depth_keys:
-                    depth = to_hwc_uint16_numpy(batch[key][i])
+                    depth = to_hwc_float32_numpy(batch[key][i])
                     depth_entity = rr.DepthImage(
                         depth,
+                        meter=depth_meter,
                         colormap=rr.components.Colormap.Viridis,
-                        depth_range=depth_ranges[key],
+                        depth_range=depth_ranges.get(key),
                     )
                     rr.log(key, entity=depth_entity)
                 else:
@@ -183,15 +268,13 @@ def visualize_dataset(
                     img_entity = rr.Image(img).compress() if display_compressed_images else rr.Image(img)
                     rr.log(key, entity=img_entity)
 
-            # display each dimension of action space (e.g. actuators command)
+            # display the action space (e.g. actuators command)
             if ACTION in batch:
-                for dim_idx, val in enumerate(batch[ACTION][i]):
-                    rr.log(f"{ACTION}/{dim_idx}", rr.Scalars(val.item()))
+                rr.log(ACTION, rr.Scalars(batch[ACTION][i].numpy()))
 
-            # display each dimension of observed state space (e.g. agent position in joint space)
+            # display the observed state space (e.g. agent position in joint space)
             if OBS_STATE in batch:
-                for dim_idx, val in enumerate(batch[OBS_STATE][i]):
-                    rr.log(f"state/{dim_idx}", rr.Scalars(val.item()))
+                rr.log("state", rr.Scalars(batch[OBS_STATE][i].numpy()))
 
             if DONE in batch:
                 rr.log(DONE, rr.Scalars(batch[DONE][i].item()))
@@ -199,12 +282,11 @@ def visualize_dataset(
             if REWARD in batch:
                 rr.log(REWARD, rr.Scalars(batch[REWARD][i].item()))
 
-            if "next.success" in batch:
-                rr.log("next.success", rr.Scalars(batch["next.success"][i].item()))
+            if SUCCESS in batch:
+                rr.log(SUCCESS, rr.Scalars(batch[SUCCESS][i].item()))
 
+    # save .rrd locally
     if mode == "local" and save:
-        # save .rrd locally
-        output_dir = Path(output_dir)
         output_dir.mkdir(parents=True, exist_ok=True)
         repo_id_str = repo_id.replace("/", "_")
         rrd_path = output_dir / f"{repo_id_str}_episode_{episode_index}.rrd"
@@ -212,7 +294,7 @@ def visualize_dataset(
         return rrd_path
 
     elif mode == "distant":
-        # stop the process from exiting since it is serving the websocket connection
+        # Keep the process alive while it serves the gRPC/web connection.
         try:
             while True:
                 time.sleep(1)
@@ -273,13 +355,11 @@ def main():
     parser.add_argument(
         "--web-port",
         type=int,
-        default=9090,
-        help="Web port for rerun.io when `--mode distant` is set.",
-    )
-    parser.add_argument(
-        "--ws-port",
-        type=int,
-        help="deprecated, please use --grpc-port instead.",
+        default=None,
+        help=(
+            "Web/WebSocket port. For rerun `--mode distant` it is the web viewer port (default 9090); "
+            "for `--display-mode foxglove` it is the server bind port (default 8765)."
+        ),
     )
     parser.add_argument(
         "--grpc-port",
@@ -312,27 +392,61 @@ def main():
     parser.add_argument(
         "--display-compressed-images",
         action="store_true",
-        help="If set, display compressed images in Rerun instead of uncompressed ones.",
+        help="If set, display compressed (JPEG) images instead of uncompressed ones.",
+    )
+
+    parser.add_argument(
+        "--display-mode",
+        type=str,
+        default="rerun",
+        choices=["rerun", "foxglove"],
+        help=(
+            "Visualization backend. 'rerun' uses the Rerun viewer (--mode/--save/--*-port apply). "
+            "'foxglove' starts a Foxglove WebSocket server that serves the episode as a seekable, "
+            "scrubbable timeline; connect the Foxglove app to ws://HOST:PORT (--host/--web-port)."
+        ),
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="127.0.0.1",
+        help=(
+            "Host to bind the Foxglove WebSocket server to when `--display-mode foxglove` is set "
+            "(127.0.0.1 for local only, 0.0.0.0 for all interfaces)."
+        ),
+    )
+    parser.add_argument(
+        "--no-autoplay",
+        dest="autoplay",
+        action="store_false",
+        help=(
+            "For `--display-mode foxglove`: don't start playing automatically when a client "
+            "connects; wait for play to be pressed in the Foxglove app instead."
+        ),
     )
 
     args = parser.parse_args()
+
+    if args.display_mode == "foxglove":
+        rerun_only = ("mode", "save", "output_dir", "grpc_port", "batch_size", "num_workers")
+        ignored = [name for name in rerun_only if getattr(args, name) != parser.get_default(name)]
+        if ignored:
+            logging.warning(
+                "These flags only apply to `--display-mode rerun` and are ignored with "
+                "`--display-mode foxglove`: %s.",
+                ", ".join(f"--{name.replace('_', '-')}" for name in ignored),
+            )
+
     kwargs = vars(args)
     repo_id = kwargs.pop("repo_id")
     root = kwargs.pop("root")
     tolerance_s = kwargs.pop("tolerance_s")
 
-    if kwargs["ws_port"] is not None:
-        logging.warning(
-            "--ws-port is deprecated and will be removed in future versions. Please use --grpc-port instead."
-        )
-        logging.warning("Setting grpc_port to ws_port value.")
-        kwargs["grpc_port"] = kwargs.pop("ws_port")
-
     init_logging()
     logging.info("Loading dataset")
     dataset = LeRobotDataset(repo_id, episodes=[args.episode_index], root=root, tolerance_s=tolerance_s)
 
-    visualize_dataset(dataset, **vars(args))
+    visualize_dataset(dataset, **kwargs)
 
 
 if __name__ == "__main__":

src/lerobot/scripts/lerobot_eval.py

@@ -169,6 +169,7 @@ def rollout(
     env_features: dict | None = None,
     recording_repo_id: str | None = None,
     recording_private: bool = False,
+    predicted_latents_callback: Callable[[PreTrainedPolicy], None] | None = None,
 ) -> dict:
     """Run a batched policy rollout once through a batch of environments.
 
@@ -198,6 +199,9 @@ def rollout(
             are returned optionally because they typically take more memory to cache. Defaults to False.
         render_callback: Optional rendering callback to be used after the environments are reset, and after
             every step.
+        predicted_latents_callback: Optional callback invoked after every ``select_action`` with the policy
+            itself. World-model policies (e.g. LingBot-VA) stash predicted video latents on
+            ``policy.last_predicted_latents``; this lets the caller concatenate chunks and decode once.
     Returns:
         The dictionary described above.
     """
@@ -276,6 +280,8 @@ def rollout(
             observation = preprocessor(observation)
             with torch.inference_mode():
                 action = policy.select_action(observation)
+            if predicted_latents_callback is not None:
+                predicted_latents_callback(policy)
             action = postprocessor(action)
 
             action_transition = {ACTION: action}
@@ -295,12 +301,22 @@ def rollout(
             # available if none of the envs finished.
             if "final_info" in info:
                 final_info = info["final_info"]
-                if not isinstance(final_info, dict):
-                    raise RuntimeError(
-                        "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). "
-                        "You're likely using an older version of gymnasium (< 1.0). Please upgrade."
+                if isinstance(final_info, dict):
+                    is_success = final_info.get("is_success", [False] * env.num_envs)
+                    successes = (
+                        is_success.tolist()
+                        if hasattr(is_success, "tolist")
+                        else [bool(is_success)] * env.num_envs
                     )
-                successes = final_info["is_success"].tolist()
+                else:
+                    # Gymnasium < 1.0 returns final_info as a per-env sequence/object array,
+                    # with entries set to a dict only for envs that just finished.
+                    successes = []
+                    for item in final_info:
+                        if isinstance(item, dict) and "is_success" in item:
+                            successes.append(bool(item["is_success"]))
+                        else:
+                            successes.append(False)
             elif "is_success" in info:
                 is_success = info["is_success"]
                 successes = (
@@ -400,6 +416,7 @@ def eval_policy(
     env_features: dict | None = None,
     recording_repo_id: str | None = None,
     recording_private: bool = False,
+    save_predicted_video: bool = False,
 ) -> dict:
     """
     Args:
@@ -418,6 +435,11 @@ def eval_policy(
     if max_episodes_rendered > 0 and not videos_dir:
         raise ValueError("If max_episodes_rendered > 0, videos_dir must be provided.")
 
+    # World-model policies (e.g. LingBot-VA) opt into predicted-video saving via their config.
+    save_predicted_video = save_predicted_video or bool(
+        getattr(getattr(policy, "config", None), "save_predicted_video", False)
+    )
+
     if not isinstance(policy, PreTrainedPolicy):
         exc = ValueError(
             f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
@@ -461,6 +483,22 @@ def eval_policy(
     if max_episodes_rendered > 0:
         video_paths: list[str] = []
 
+    if save_predicted_video:
+        if not videos_dir:
+            raise ValueError("If save_predicted_video is True, videos_dir must be provided.")
+        predicted_video_paths: list[str] = []
+        n_predicted_rendered = 0
+
+    # Collect predicted-video latents across a rollout (world-model policies only). The latents are
+    # concatenated and decoded once after the rollout, matching upstream LingBot-VA's visualization path.
+    def collect_predicted_latents(policy: PreTrainedPolicy):
+        latents = getattr(policy, "last_predicted_latents", None)
+        if latents is not None:
+            pred_latents.append(
+                latents.detach().to("cpu") if hasattr(latents, "detach") else torch.as_tensor(latents).cpu()
+            )
+            policy.last_predicted_latents = None
+
     if return_episode_data:
         episode_data: dict | None = None
 
@@ -472,6 +510,9 @@ def eval_policy(
         if max_episodes_rendered > 0:
             ep_frames: list[np.ndarray] = []
 
+        if save_predicted_video:
+            pred_latents: list[torch.Tensor] = []
+
         if start_seed is None:
             seeds = None
         else:
@@ -492,6 +533,7 @@ def eval_policy(
             env_features=env_features,
             recording_repo_id=recording_repo_id,
             recording_private=recording_private,
+            predicted_latents_callback=collect_predicted_latents if save_predicted_video else None,
         )
 
         # Figure out where in each rollout sequence the first done condition was encountered (results after
@@ -557,6 +599,35 @@ def eval_policy(
                 threads.append(thread)
                 n_episodes_rendered += 1
 
+        # Maybe save the policy's predicted (imagined) video for this batch's rollout.
+        if save_predicted_video and len(pred_latents) > 0:
+            predicted_latent = torch.cat(pred_latents, dim=2)
+            decoder = getattr(policy, "decode_predicted_latents", None) or getattr(
+                policy, "_decode_predicted_video", None
+            )
+            if decoder is None:
+                raise AttributeError(
+                    "Policy config requested predicted-video saving, but the policy does not expose "
+                    "`decode_predicted_latents` or `_decode_predicted_video`."
+                )
+            predicted_video = decoder(predicted_latent)
+            if hasattr(predicted_video, "detach"):
+                predicted_video = predicted_video.detach().to("cpu").numpy()
+            videos_dir.mkdir(parents=True, exist_ok=True)
+            predicted_video_path = videos_dir / f"pred_episode_{n_predicted_rendered}.mp4"
+            predicted_video_paths.append(str(predicted_video_path))
+            thread = threading.Thread(
+                target=write_video,
+                args=(
+                    str(predicted_video_path),
+                    predicted_video,
+                    env.unwrapped.metadata["render_fps"],
+                ),
+            )
+            thread.start()
+            threads.append(thread)
+            n_predicted_rendered += 1
+
         progbar.set_postfix(
             {"running_success_rate": f"{np.mean(all_successes[:n_episodes]).item() * 100:.1f}%"}
         )
@@ -600,6 +671,9 @@ def eval_policy(
     if max_episodes_rendered > 0:
         info["video_paths"] = video_paths
 
+    if save_predicted_video:
+        info["predicted_video_paths"] = predicted_video_paths
+
     return info
 
 
@@ -740,9 +814,10 @@ class TaskMetrics(TypedDict):
     max_rewards: list[float]
     successes: list[bool]
     video_paths: list[str]
+    predicted_video_paths: list[str]
 
 
-ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths")
+ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths", "predicted_video_paths")
 
 
 def eval_one(
@@ -791,6 +866,7 @@ def eval_one(
         max_rewards=[ep["max_reward"] for ep in per_episode],
         successes=[ep["success"] for ep in per_episode],
         video_paths=task_result.get("video_paths", []),
+        predicted_video_paths=task_result.get("predicted_video_paths", []),
     )
 
 
@@ -851,6 +927,7 @@ def run_one(
 
     if max_episodes_rendered > 0:
         metrics.setdefault("video_paths", [])
+    metrics.setdefault("predicted_video_paths", [])
     return task_group, task_id, metrics
 
 
@@ -908,11 +985,11 @@ def eval_policy_all(
         _append("sum_rewards", metrics.get("sum_rewards"))
         _append("max_rewards", metrics.get("max_rewards"))
         _append("successes", metrics.get("successes"))
-        # video_paths is list-like
-        paths = metrics.get("video_paths", [])
-        if paths:
-            group_acc[group]["video_paths"].extend(paths)
-            overall["video_paths"].extend(paths)
+        for key in ("video_paths", "predicted_video_paths"):
+            paths = metrics.get(key, [])
+            if paths:
+                group_acc[group][key].extend(paths)
+                overall[key].extend(paths)
 
     # Choose runner (sequential vs threaded)
     task_runner = partial(
@@ -984,6 +1061,7 @@ def eval_policy_all(
             "pc_success": _agg_from_list(acc["successes"]) * 100 if acc["successes"] else float("nan"),
             "n_episodes": len(acc["sum_rewards"]),
             "video_paths": list(acc["video_paths"]),
+            "predicted_video_paths": list(acc["predicted_video_paths"]),
         }
 
     # overall aggregates
@@ -995,6 +1073,7 @@ def eval_policy_all(
         "eval_s": time.time() - start_t,
         "eval_ep_s": (time.time() - start_t) / max(1, len(overall["sum_rewards"])),
         "video_paths": list(overall["video_paths"]),
+        "predicted_video_paths": list(overall["predicted_video_paths"]),
     }
 
     return {

src/lerobot/scripts/lerobot_record.py

@@ -38,6 +38,9 @@ lerobot-record \\
     --display_data=true
 ```
 
+To stream the data to Foxglove instead of Rerun, add ``--display_mode=foxglove`` (then connect the
+Foxglove app to ``ws://127.0.0.1:8765``; override the port with ``--display_port=<port>``).
+
 Example recording with bimanual so100:
 ```shell
 lerobot-record \\
@@ -157,7 +160,11 @@ from lerobot.utils.utils import (
     init_logging,
     log_say,
 )
-from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import (
+    init_visualization,
+    log_visualization_data,
+    shutdown_visualization,
+)
 
 
 @dataclass
@@ -168,11 +175,14 @@ class RecordConfig:
     teleop: TeleoperatorConfig | None = None
     # Display all cameras on screen
     display_data: bool = False
-    # Display data on a remote Rerun server
+    # Visualization backend used when display_data is True: "rerun" or "foxglove".
+    display_mode: str = "rerun"
+    # For "rerun": IP of a remote server to send to. For "foxglove": interface to bind the WebSocket
+    # server to (127.0.0.1 for local only, 0.0.0.0 for all interfaces).
     display_ip: str | None = None
-    # Port of the remote Rerun server
+    # For "rerun": port of the remote server. For "foxglove": port to bind the WebSocket server to.
     display_port: int | None = None
-    # Whether to  display compressed images in Rerun
+    # Whether to display compressed (JPEG) images instead of raw frames
     display_compressed_images: bool = False
     # Use vocal synthesis to read events.
     play_sounds: bool = True
@@ -233,6 +243,7 @@ def record_loop(
     control_time_s: int | None = None,
     single_task: str | None = None,
     display_data: bool = False,
+    display_mode: str = "rerun",
     display_compressed_images: bool = False,
 ):
     if dataset is not None and dataset.fps != fps:
@@ -327,8 +338,11 @@ def record_loop(
             dataset.add_frame(frame)
 
         if display_data:
-            log_rerun_data(
-                observation=obs_processed, action=action_values, compress_images=display_compressed_images
+            log_visualization_data(
+                display_mode,
+                observation=obs_processed,
+                action=action_values,
+                compress_images=display_compressed_images,
             )
 
         dt_s = time.perf_counter() - start_loop_t
@@ -354,7 +368,9 @@ def record(
     init_logging()
     logging.info(pformat(asdict(cfg)))
     if cfg.display_data:
-        init_rerun(session_name="recording", ip=cfg.display_ip, port=cfg.display_port)
+        init_visualization(
+            cfg.display_mode, session_name="recording", ip=cfg.display_ip, port=cfg.display_port
+        )
     display_compressed_images = (
         True
         if (cfg.display_data and cfg.display_ip is not None and cfg.display_port is not None)
@@ -464,6 +480,7 @@ def record(
                     control_time_s=cfg.dataset.episode_time_s,
                     single_task=cfg.dataset.single_task,
                     display_data=cfg.display_data,
+                    display_mode=cfg.display_mode,
                     display_compressed_images=display_compressed_images,
                 )
 
@@ -485,6 +502,7 @@ def record(
                         control_time_s=cfg.dataset.reset_time_s,
                         single_task=cfg.dataset.single_task,
                         display_data=cfg.display_data,
+                        display_mode=cfg.display_mode,
                     )
 
                 if events["rerecord_episode"]:
@@ -510,6 +528,9 @@ def record(
         if listener is not None:
             listener.stop()
 
+        if cfg.display_data:
+            shutdown_visualization(cfg.display_mode)
+
         if cfg.dataset.push_to_hub:
             if dataset and dataset.num_episodes > 0:
                 dataset.push_to_hub(tags=cfg.dataset.tags, private=cfg.dataset.private)

src/lerobot/scripts/lerobot_rollout.py

@@ -145,6 +145,9 @@ Usage examples
         --dataset.rgb_encoder.vcodec=h264 \\
         --dataset.rgb_encoder.preset=fast \\
         --dataset.rgb_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2}
+
+    # Stream to Foxglove instead of Rerun:
+    # add --display_mode=foxglove, then connect the Foxglove app to ws://127.0.0.1:8765.
 """
 
 import logging
@@ -190,7 +193,7 @@ from lerobot.teleoperators import (  # noqa: F401
 from lerobot.utils.import_utils import register_third_party_plugins
 from lerobot.utils.process import ProcessSignalHandler
 from lerobot.utils.utils import init_logging
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import init_visualization, shutdown_visualization
 
 logger = logging.getLogger(__name__)
 
@@ -201,8 +204,13 @@ def rollout(cfg: RolloutConfig):
     init_logging()
 
     if cfg.display_data:
-        logger.info("Initializing Rerun visualization (ip=%s, port=%s)", cfg.display_ip, cfg.display_port)
-        init_rerun(session_name="rollout", ip=cfg.display_ip, port=cfg.display_port)
+        logger.info(
+            "Initializing %s visualization (ip=%s, port=%s)",
+            cfg.display_mode,
+            cfg.display_ip,
+            cfg.display_port,
+        )
+        init_visualization(cfg.display_mode, session_name="rollout", ip=cfg.display_ip, port=cfg.display_port)
 
     signal_handler = ProcessSignalHandler(use_threads=True, display_pid=False)
     shutdown_event = signal_handler.shutdown_event
@@ -227,6 +235,8 @@ def rollout(cfg: RolloutConfig):
         logger.info("Interrupted by user")
     finally:
         strategy.teardown(ctx)
+        if cfg.display_data:
+            shutdown_visualization(cfg.display_mode)
 
     logger.info("Rollout finished")
 

src/lerobot/scripts/lerobot_teleoperate.py

@@ -31,6 +31,22 @@ lerobot-teleoperate \
     --display_data=true
 ```
 
+To stream the data to Foxglove instead of Rerun, add ``--display_mode=foxglove``
+(then connect the Foxglove app to ``ws://127.0.0.1:8765``; override the port with ``--display_port=<port>``):
+
+```shell
+lerobot-teleoperate \
+    --robot.type=so101_follower \
+    --robot.port=/dev/tty.usbmodem58760431541 \
+    --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}}" \
+    --robot.id=black \
+    --teleop.type=so101_leader \
+    --teleop.port=/dev/tty.usbmodem58760431551 \
+    --teleop.id=blue \
+    --display_data=true \
+    --display_mode=foxglove
+```
+
 Example teleoperation with bimanual so100:
 
 ```shell
@@ -108,7 +124,11 @@ from lerobot.teleoperators import (  # noqa: F401
 from lerobot.utils.import_utils import register_third_party_plugins
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import init_logging, move_cursor_up
-from lerobot.utils.visualization_utils import init_rerun, log_rerun_data, shutdown_rerun
+from lerobot.utils.visualization_utils import (
+    init_visualization,
+    log_visualization_data,
+    shutdown_visualization,
+)
 
 
 @dataclass
@@ -121,11 +141,14 @@ class TeleoperateConfig:
     teleop_time_s: float | None = None
     # Display all cameras on screen
     display_data: bool = False
-    # Display data on a remote Rerun server
+    # Visualization backend used when display_data is True: "rerun" or "foxglove".
+    display_mode: str = "rerun"
+    # For "rerun": IP of a remote server to send to. For "foxglove": interface to bind the WebSocket
+    # server to (127.0.0.1 for local only, 0.0.0.0 for all interfaces).
     display_ip: str | None = None
-    # Port of the remote Rerun server
+    # For "rerun": port of the remote server. For "foxglove": port to bind the WebSocket server to.
     display_port: int | None = None
-    # Whether to  display compressed images in Rerun
+    # Whether to display compressed (JPEG) images instead of raw frames
     display_compressed_images: bool = False
 
 
@@ -137,6 +160,7 @@ def teleop_loop(
     robot_action_processor: RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction],
     robot_observation_processor: RobotProcessorPipeline[RobotObservation, RobotObservation],
     display_data: bool = False,
+    display_mode: str = "rerun",
     duration: float | None = None,
     display_compressed_images: bool = False,
 ):
@@ -149,8 +173,10 @@ def teleop_loop(
         teleop: The teleoperator device instance providing control actions.
         robot: The robot instance being controlled.
         fps: The target frequency for the control loop in frames per second.
-        display_data: If True, fetches robot observations and displays them in the console and Rerun.
-        display_compressed_images: If True, compresses images before sending them to Rerun for display.
+        display_data: If True, fetches robot observations and displays them in the console and the
+            visualization backend.
+        display_mode: Visualization backend to use when display_data is True ("rerun" or "foxglove").
+        display_compressed_images: If True, compresses images before sending them to the backend for display.
         duration: The maximum duration of the teleoperation loop in seconds. If None, the loop runs indefinitely.
         teleop_action_processor: An optional pipeline to process raw actions from the teleoperator.
         robot_action_processor: An optional pipeline to process actions before they are sent to the robot.
@@ -187,7 +213,8 @@ def teleop_loop(
             # Process robot observation through pipeline
             obs_transition = robot_observation_processor(obs)
 
-            log_rerun_data(
+            log_visualization_data(
+                display_mode,
                 observation=obs_transition,
                 action=teleop_action,
                 compress_images=display_compressed_images,
@@ -215,7 +242,9 @@ def teleoperate(cfg: TeleoperateConfig):
     init_logging()
     logging.info(pformat(asdict(cfg)))
     if cfg.display_data:
-        init_rerun(session_name="teleoperation", ip=cfg.display_ip, port=cfg.display_port)
+        init_visualization(
+            cfg.display_mode, session_name="teleoperation", ip=cfg.display_ip, port=cfg.display_port
+        )
     display_compressed_images = (
         True
         if (cfg.display_data and cfg.display_ip is not None and cfg.display_port is not None)
@@ -235,6 +264,7 @@ def teleoperate(cfg: TeleoperateConfig):
             robot=robot,
             fps=cfg.fps,
             display_data=cfg.display_data,
+            display_mode=cfg.display_mode,
             duration=cfg.teleop_time_s,
             teleop_action_processor=teleop_action_processor,
             robot_action_processor=robot_action_processor,
@@ -245,7 +275,7 @@ def teleoperate(cfg: TeleoperateConfig):
         pass
     finally:
         if cfg.display_data:
-            shutdown_rerun()
+            shutdown_visualization(cfg.display_mode)
         teleop.disconnect()
         robot.disconnect()
 

src/lerobot/scripts/lerobot_train.py

@@ -20,6 +20,7 @@ Requires: pip install 'lerobot[training]'  (includes dataset + accelerate + wand
 
 import dataclasses
 import logging
+import sys
 import time
 from contextlib import nullcontext
 from pprint import pformat
@@ -41,15 +42,17 @@ from lerobot.common.train_utils import (
     load_training_batch_size,
     load_training_num_processes,
     load_training_state,
+    push_checkpoint_to_hub,
     save_checkpoint,
     update_last_checkpoint,
 )
 from lerobot.common.wandb_utils import WandBLogger
-from lerobot.configs import parser
+from lerobot.configs import JobConfig, parser
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.datasets import EpisodeAwareSampler, compute_sampler_state
 from lerobot.datasets.factory import make_train_eval_datasets
 from lerobot.envs import close_envs, make_env, make_env_pre_post_processors
+from lerobot.jobs import submit_to_hf
 from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
 from lerobot.rewards import make_reward_pre_post_processors
@@ -188,6 +191,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         cfg: A `TrainPipelineConfig` object containing all training configurations.
         accelerator: Optional Accelerator instance. If None, one will be created automatically.
     """
+    if cfg.job.is_remote:
+        return submit_to_hf(cfg)
+
     from lerobot.utils.import_utils import require_package
 
     require_package("accelerate", extra="training")
@@ -205,8 +211,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting.
         # Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training).
         force_cpu = cfg.trainable_config.device == "cpu"
+        # Drive Accelerate's autocast from policy.dtype (bf16/fp16 activate it; float32/absent -> launcher default).
+        policy_dtype = getattr(cfg.trainable_config, "dtype", None)
+        mixed_precision = {"bfloat16": "bf16", "float16": "fp16", "float32": "no"}.get(policy_dtype)
         accelerator = Accelerator(
             step_scheduler_with_optimizer=False,
+            mixed_precision=mixed_precision,
             kwargs_handlers=[ddp_kwargs],
             cpu=force_cpu,
         )
@@ -655,6 +665,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     optim_state_dict=optim_state_dict,
                 )
                 update_last_checkpoint(checkpoint_dir)
+                if cfg.save_checkpoint_to_hub:
+                    push_checkpoint_to_hub(
+                        checkpoint_dir,
+                        cfg.policy.repo_id,
+                        private=cfg.policy.private,
+                    )
                 if wandb_logger:
                     wandb_logger.log_policy(checkpoint_dir)
 
@@ -724,9 +740,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             unwrapped_model = accelerator.unwrap_model(policy)
             # PEFT only applies when training a policy — reward models use the plain path.
             if not cfg.is_reward_model_training and cfg.policy.use_peft:
-                unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model)
+                unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model, dataset_meta=dataset.meta)
             else:
-                unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict)
+                unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict, dataset_meta=dataset.meta)
             preprocessor.push_to_hub(active_cfg.repo_id)
             postprocessor.push_to_hub(active_cfg.repo_id)
 
@@ -735,8 +751,25 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     accelerator.end_training()
 
 
+def _remote_target_in_argv() -> bool:
+    """True when the CLI requests a remote HF Jobs run (--job.target=<non-local>)."""
+    target = None
+    args = sys.argv[1:]
+    for i, tok in enumerate(args):
+        if tok == "--job.target" and i + 1 < len(args):
+            target = args[i + 1]
+        elif tok.startswith("--job.target="):
+            target = tok.split("=", 1)[1]
+    return JobConfig.is_remote_target(target)
+
+
 def main():
     register_third_party_plugins()
+    if _remote_target_in_argv():
+        # The policy device is resolved on the remote pod, not here, so silence the
+        # client-side "Device '...' is not available" warning PreTrainedConfig emits
+        # while parsing the config (it fires before train() can dispatch remotely).
+        logging.getLogger("lerobot.configs.policies").setLevel(logging.ERROR)
     train()
 
 

src/lerobot/teleoperators/rebot_102_leader/config_rebot_102_leader.py

@@ -65,7 +65,7 @@ class RebotArm102LeaderConfig:
     joint_ranges: dict[str, list[int]] = field(
         default_factory=lambda: {
             "shoulder_pan": [-150, 150],
-            "shoulder_lift": [-170, 1],
+            "shoulder_lift": [-200, 1],
             "elbow_flex": [-200, 1],
             "wrist_flex": [-80, 90],
             "wrist_yaw": [-90, 90],

src/lerobot/utils/constants.py

@@ -26,6 +26,9 @@ OBS_IMAGES = OBS_IMAGE + "s"
 OBS_LANGUAGE = OBS_STR + ".language"
 OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
 OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
+OBS_LANGUAGE_UNCOND = OBS_STR + ".language_uncond"
+OBS_LANGUAGE_UNCOND_TOKENS = OBS_LANGUAGE_UNCOND + ".tokens"
+OBS_LANGUAGE_UNCOND_ATTENTION_MASK = OBS_LANGUAGE_UNCOND + ".attention_mask"
 OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask"
 OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens"
 OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask"
@@ -37,6 +40,7 @@ ACTION_TOKEN_MASK = ACTION + ".token_mask"
 REWARD = "next.reward"
 TRUNCATED = "next.truncated"
 DONE = "next.done"
+SUCCESS = "next.success"
 INFO = "info"
 
 ROBOTS = "robots"

src/lerobot/utils/foxglove_visualization.py

@@ -0,0 +1,651 @@
+# 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.
+
+"""Foxglove visualization backend.
+
+Live control-loop streaming (:func:`log_foxglove_data`) and seekable dataset playback
+(:func:`serve_foxglove_dataset_playback`) over a Foxglove WebSocket server. Callers usually select a
+backend at runtime through the dispatch in :mod:`lerobot.utils.visualization_utils` rather than
+importing from here directly. Requires the ``viz`` extra (``pip install 'lerobot[viz]'``).
+"""
+
+import logging
+import numbers
+import time
+
+import cv2
+import numpy as np
+
+from lerobot.types import RobotAction, RobotObservation
+
+from .constants import (
+    ACTION,
+    ACTION_PREFIX,
+    DONE,
+    OBS_IMAGES,
+    OBS_PREFIX,
+    OBS_STATE,
+    OBS_STR,
+    REWARD,
+    SUCCESS,
+    TRUNCATED,
+)
+from .import_utils import require_package
+
+# Static schema shared by all scalar topics. Each message carries a flat list of ``{label, value}``
+# pairs rather than one field per feature, so the same schema fits any robot regardless of which
+# observation/action features it reports. The ``label`` field name is what Foxglove looks for to name
+# each series automatically, so a single filtered path plots every feature, e.g.
+# ``/observation/state.scalars[:]``.
+_SCALARS_SCHEMA = {
+    "type": "object",
+    "title": "lerobot.Scalars",
+    "properties": {
+        "scalars": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "properties": {
+                    "label": {"type": "string"},
+                    "value": {"type": "number"},
+                },
+            },
+        }
+    },
+}
+
+
+def _is_scalar(x):
+    return isinstance(x, (float | numbers.Real | np.integer | np.floating)) or (
+        isinstance(x, np.ndarray) and x.ndim == 0
+    )
+
+
+def init_foxglove(host: str = "127.0.0.1", port: int | None = 8765) -> None:
+    """
+    Starts a Foxglove WebSocket server for visualizing the control loop.
+
+    Connect to it from the Foxglove app at ``ws://<host>:<port>``. Calling this
+    more than once is a no-op while a server is already running.
+
+    Args:
+        host: Host interface to bind the WebSocket server to.
+        port: Port to bind the WebSocket server to (defaults to 8765).
+    """
+
+    require_package("foxglove-sdk", extra="viz", import_name="foxglove")
+    import foxglove
+
+    # Live-stream state lives as attributes on ``log_foxglove_data``:
+    # ``.server`` is the shared WebSocket server and
+    # ``.channels`` caches one Foxglove channel per topic
+    if getattr(log_foxglove_data, "server", None) is not None:
+        return
+    log_foxglove_data.server = foxglove.start_server(host=host, port=port or 8765)
+    log_foxglove_data.channels = {}
+
+
+def shutdown_foxglove() -> None:
+    """Stops the Foxglove WebSocket server and clears cached channels."""
+
+    server = getattr(log_foxglove_data, "server", None)
+    if server is not None:
+        server.stop()
+    log_foxglove_data.server = None
+    log_foxglove_data.channels = {}
+
+
+def _foxglove_safe_name(name: str) -> str:
+    """Replace ``.`` with ``_`` so a feature name is a single Foxglove topic-path segment.
+
+    Foxglove treats ``.`` as a path separator, so an unsanitized name like ``observation.images.front``
+    would split into nested segments instead of naming one topic.
+    """
+
+    return name.replace(".", "_")
+
+
+def _foxglove_topic(key: str, *, is_image: bool = False) -> str:
+    """Build the Foxglove topic for a feature ``key``.
+
+    Camera features map to a per-source image topic (``/observation/images/<name>``); scalar features
+    share one aggregate topic per source: ``/observation/state`` for observations, ``/action/state``
+    for actions.
+    """
+
+    if is_image:
+        name = str(key)
+        for prefix in (f"{OBS_IMAGES}.", OBS_PREFIX):
+            if name.startswith(prefix):
+                name = name[len(prefix) :]
+                break
+        return f"/{OBS_STR}/images/{_foxglove_safe_name(name)}"
+    source = ACTION if (str(key).startswith(ACTION_PREFIX) or str(key) == ACTION) else OBS_STR
+    return f"/{source}/state"
+
+
+def _log_foxglove_scalars(
+    topic: str, values: dict[str, float], *, channels: dict | None = None, log_time: int | None = None
+) -> None:
+    """Log scalars on a typed JSON channel using the static :data:`_SCALARS_SCHEMA`.
+
+    ``values`` is an ordered mapping of feature name to value; it is emitted as a ``scalars`` array of
+    ``{label, value}`` objects. Insertion order is preserved so series stay stable across messages.
+
+    ``channels`` is the per-topic channel cache to reuse (defaults to the live-stream cache on
+    :func:`log_foxglove_data`; dataset playback passes its own local cache to stay self-contained).
+    ``log_time`` is the message time in nanoseconds; when ``None`` the server's receive time is used.
+    """
+
+    if not values:
+        return
+
+    import foxglove
+
+    if channels is None:
+        channels = log_foxglove_data.channels
+    channel = channels.get(topic)
+    if channel is None:
+        channel = channels[topic] = foxglove.Channel(topic, schema=_SCALARS_SCHEMA, message_encoding="json")
+    msg = {"scalars": [{"label": label, "value": value} for label, value in values.items()]}
+    if log_time is None:
+        channel.log(msg)
+    else:
+        channel.log(msg, log_time=log_time)
+
+
+def _labeled_scalars(name: str, values, labels: list[str] | None = None) -> dict[str, float]:
+    """Expand a 1D sequence into ``{label: value}`` entries with a consistent fallback."""
+
+    flat = [float(v) for v in values]
+    if labels is None or len(labels) != len(flat):
+        labels = [f"{name}_{i}" for i in range(len(flat))]
+    return dict(zip(labels, flat, strict=True))
+
+
+def _log_foxglove_image(
+    topic: str,
+    frame_id: str,
+    arr: np.ndarray,
+    *,
+    compress_images: bool,
+    channels: dict | None = None,
+    log_time: int | None = None,
+    depth_range: tuple[float, float] | None = None,
+    raw_depth_values: bool = False,
+) -> None:
+    """Log an image on a cached per-topic channel.
+
+    The encoding is chosen from the channel count and dtype: a single-channel ``float`` or ``uint16``
+    frame is a depth map (``32FC1``/``16UC1``), single-channel ``uint8`` is ``mono8``, 3 => ``rgb8``
+    (float input assumed in [0, 1], cast to uint8), 4 => ``rgba8``; other counts are skipped with a
+    warning. When ``compress_images`` is set, ``rgb8`` is JPEG-encoded instead.
+
+    Args:
+        topic: Foxglove topic to log on.
+        frame_id: Frame id stamped on the message.
+        arr: Image as HWC or CHW (CHW is transposed to HWC), any dtype.
+        compress_images: JPEG-encode ``rgb8`` frames; ignored for other encodings.
+        channels: Per-topic channel cache to reuse (see :func:`_log_foxglove_scalars`).
+        log_time: Message time in nanoseconds, also written to the header timestamp; when ``None``
+            the server's receive time is used.
+        depth_range: ``(lo, hi)`` clip bounds in a depth frame's own input units. Depth frames
+            (``32FC1``/``16UC1``) are rescaled onto Foxglove's default display max for their encoding
+            (``1.0`` / ``10000``) so they show with sensible contrast; ``depth_range`` sets the source
+            range, else the frame's own min/max is used. Ignored for ``mono8``/``rgb8``/``rgba8``.
+        raw_depth_values: If True, depth values are not rescaled and are logged as is.
+    """
+
+    from foxglove.channels import CompressedImageChannel, RawImageChannel
+    from foxglove.messages import CompressedImage, RawImage, Timestamp
+
+    if channels is None:
+        channels = log_foxglove_data.channels
+    time_ns = time.time_ns() if log_time is None else log_time
+    timestamp = Timestamp(sec=time_ns // 1_000_000_000, nsec=time_ns % 1_000_000_000)
+    log_kwargs = {} if log_time is None else {"log_time": log_time}
+
+    # Convert CHW -> HWC when needed (mirrors log_rerun_data).
+    if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
+        arr = np.transpose(arr, (1, 2, 0))
+    height, width = arr.shape[0], arr.shape[1]
+    n_channels = 1 if arr.ndim == 2 else arr.shape[2]
+
+    if n_channels == 1 and arr.dtype != np.uint8:
+        # Depth map: infer the encoding from the dtype.
+        encoding, target_dtype, value_max = (
+            ("32FC1", np.float32, 1.0)
+            if np.issubdtype(arr.dtype, np.floating)
+            else ("16UC1", np.uint16, 10000.0)
+        )
+        if not raw_depth_values:
+            # Rescale onto the encoding's display max with respect to the given depth_range.
+            lo, hi = depth_range if depth_range is not None else (float(arr.min()), float(arr.max()))
+            arr = arr.clip(lo, hi).astype(np.float32)
+            arr = (arr - lo) / ((hi - lo) if hi > lo else 1.0) * value_max
+        arr = np.ascontiguousarray(arr, dtype=target_dtype)
+    else:
+        if n_channels == 3 and np.issubdtype(arr.dtype, np.floating):
+            arr = (arr * 255.0).clip(0, 255)
+        arr = np.ascontiguousarray(arr, dtype=np.uint8)
+
+        if compress_images and n_channels == 3:
+            buf_src = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
+            _, buf = cv2.imencode(".jpg", buf_src)
+            channel = channels.get(topic)
+            if channel is None:
+                channel = channels[topic] = CompressedImageChannel(topic=topic)
+            channel.log(
+                CompressedImage(timestamp=timestamp, frame_id=frame_id, data=buf.tobytes(), format="jpeg"),
+                **log_kwargs,
+            )
+            return
+
+        encoding = {1: "mono8", 3: "rgb8", 4: "rgba8"}.get(n_channels)
+        if encoding is None:
+            logging.warning(
+                "Foxglove: skipping image on topic '%s' with unsupported shape %s (%d channels); "
+                "expected 1 (mono8/16UC1/32FC1), 3 (rgb8), or 4 (rgba8) channels.",
+                topic,
+                tuple(arr.shape),
+                n_channels,
+            )
+            return
+
+    channel = channels.get(topic)
+    if channel is None:
+        channel = channels[topic] = RawImageChannel(topic=topic)
+    channel.log(
+        RawImage(
+            timestamp=timestamp,
+            frame_id=frame_id,
+            width=width,
+            height=height,
+            encoding=encoding,
+            step=width * n_channels * arr.itemsize,
+            data=arr.tobytes(),
+        ),
+        **log_kwargs,
+    )
+
+
+def log_foxglove_data(
+    observation: RobotObservation | None = None,
+    action: RobotAction | None = None,
+    compress_images: bool = False,
+) -> None:
+    """
+    Logs observation and action data to a Foxglove WebSocket server for real-time visualization.
+
+    Mirrors ``log_rerun_data`` but emits Foxglove messages over the server started by
+    :func:`init_foxglove`. Data is mapped as follows:
+    - Scalars (and elements of 1D arrays) are accumulated per source and logged on the
+      ``/observation/state`` and ``/action/state`` topics as typed JSON messages using the static
+      ``lerobot.Scalars`` schema: a ``scalars`` array of ``{label, value}`` objects (see
+      :data:`_SCALARS_SCHEMA`). The ``label`` field lets Foxglove name each series automatically, so
+      ``/observation/state.scalars[:].value`` plots every feature at once.
+    - 3D NumPy arrays that resemble images are transposed from CHW to HWC when needed and logged on a
+      per-source topic (e.g. ``/observation/images/front``) as a ``RawImage`` (or a JPEG
+      ``CompressedImage`` when ``compress_images`` is True).
+
+    Args:
+        observation: An optional dictionary containing observation data to log.
+        action: An optional dictionary containing action data to log.
+        compress_images: Whether to JPEG-compress images before logging to save bandwidth in exchange
+            for CPU and quality.
+    """
+
+    require_package("foxglove-sdk", extra="viz", import_name="foxglove")
+
+    if getattr(log_foxglove_data, "server", None) is None:
+        raise RuntimeError("init_foxglove() must be called before log_foxglove_data().")
+
+    now = time.time_ns()
+
+    if observation:
+        obs_scalars: dict[str, float] = {}
+        for k, v in observation.items():
+            if v is None:
+                continue
+            key = k[len(OBS_PREFIX) :] if str(k).startswith(OBS_PREFIX) else str(k)
+            if _is_scalar(v):
+                obs_scalars[key] = float(v)
+            elif isinstance(v, np.ndarray):
+                if v.ndim == 1:
+                    obs_scalars.update(_labeled_scalars(key, v))
+                else:
+                    _log_foxglove_image(
+                        _foxglove_topic(k, is_image=True),
+                        key,
+                        v,
+                        compress_images=compress_images,
+                        log_time=now,
+                    )
+        _log_foxglove_scalars(_foxglove_topic(OBS_STATE), obs_scalars, log_time=now)
+
+    if action:
+        action_scalars: dict[str, float] = {}
+        for k, v in action.items():
+            if v is None:
+                continue
+            key = k[len(ACTION_PREFIX) :] if str(k).startswith(ACTION_PREFIX) else str(k)
+            if _is_scalar(v):
+                action_scalars[key] = float(v)
+            elif isinstance(v, np.ndarray):
+                action_scalars.update(_labeled_scalars(key, v.flatten()))
+        _log_foxglove_scalars(_foxglove_topic(ACTION), action_scalars, log_time=now)
+
+
+# ── Dataset playback over a Foxglove WebSocket server ─────────────────────
+# A LeRobotDataset is random-access on disk, so rather than fire-and-forget a forward stream we
+# advertise a seekable timeline and serve frames on demand for whatever time the user scrubs/plays
+# to in the Foxglove app. This relies on the SDK's PlaybackControl capability.
+
+
+def _feature_dim_names(feature: dict | None) -> list[str] | None:
+    """Best-effort per-dimension series labels for a 1D feature, or ``None`` to fall back to indices.
+
+    LeRobot records a feature's ``names`` inconsistently: a flat list (``["x", "y"]``), a category
+    mapping (``{"motors": ["motor_0", "motor_1"]}``), or a name->index mapping
+    (``{"delta_x": 0, "delta_y": 1}``). Each is handled, but labels are only returned when their count
+    matches the feature's 1D shape, so a malformed/mismatched ``names`` can't silently mislabel series.
+    """
+
+    if not feature:
+        return None
+    shape = feature.get("shape")
+    dim = shape[0] if shape and len(shape) == 1 else None
+    names = feature.get("names")
+    labels: list[str] | None = None
+    if isinstance(names, dict):
+        values = list(names.values())
+        if values and all(isinstance(v, (list, tuple)) for v in values):
+            labels = [str(n) for group in values for n in group]
+        elif values and all(isinstance(v, int) and not isinstance(v, bool) for v in values):
+            labels = [name for name, _ in sorted(names.items(), key=lambda kv: kv[1])]
+    elif isinstance(names, (list, tuple)):
+        labels = [str(n) for n in names]
+    if labels is not None and dim is not None and len(labels) == dim:
+        return labels
+    return None
+
+
+def _frame_to_scalars(sample: dict, key: str, labels: list[str] | None = None) -> dict[str, float]:
+    """Flatten a frame's vector/scalar feature ``key`` into ``{label: value}`` entries.
+
+    ``labels`` provides one name per dimension (from the dataset's feature metadata); when absent or
+    the wrong length, dimensions fall back to ``{name}_{i}`` (the short feature name), matching the
+    live stream so series names agree. A scalar feature becomes a single entry. Missing or ``None``
+    features yield an empty mapping.
+    """
+
+    v = sample.get(key)
+    if v is None:
+        return {}
+    arr = v.numpy() if hasattr(v, "numpy") else np.asarray(v)
+    if key.startswith(OBS_PREFIX):
+        name = key[len(OBS_PREFIX) :]
+    elif key.startswith(ACTION_PREFIX):
+        name = key[len(ACTION_PREFIX) :]
+    else:
+        name = key
+    if arr.ndim == 0:
+        return {name: float(arr)}
+    return _labeled_scalars(name, arr.flatten(), labels)
+
+
+def serve_foxglove_dataset_playback(
+    dataset,
+    episode_index: int,
+    *,
+    host: str = "127.0.0.1",
+    port: int = 8765,
+    compress_images: bool = False,
+    autoplay: bool = True,
+) -> None:
+    """Serve a single dataset episode to Foxglove as a seekable, scrubbable timeline.
+
+    Starts a Foxglove WebSocket server advertising the ``PlaybackControl`` capability over the
+    episode's time range. The Foxglove app drives play/pause/seek/speed; a background thread and a
+    ``ServerListener`` read frames from the on-disk ``dataset`` on demand and log them stamped at
+    their dataset timestamps, so the user can scrub anywhere in the episode. Blocks until interrupted.
+
+    Args:
+        dataset: A ``LeRobotDataset`` loaded for the single episode to visualize.
+        episode_index: Index of the episode being visualized (used only for the session name).
+        host: Host interface to bind the WebSocket server to.
+        port: Port to bind the WebSocket server to.
+        compress_images: Whether to JPEG-compress camera frames before logging.
+        autoplay: If True, start playing automatically as soon as a client connects, instead of
+            waiting for the user to press play in the Foxglove app.
+    """
+
+    require_package("foxglove-sdk", extra="viz", import_name="foxglove")
+    import bisect
+    import threading
+
+    import foxglove
+    from foxglove.websocket import (
+        Capability,
+        PlaybackCommand,
+        PlaybackControlRequest,
+        PlaybackState,
+        PlaybackStatus,
+        ServerListener,
+    )
+
+    # Per-frame timestamps in nanoseconds (read straight from the table, no video decode).
+    times_ns = [int(round(float(t) * 1e9)) for t in dataset.hf_dataset["timestamp"]]
+    n_frames = len(times_ns)
+    if n_frames == 0:
+        raise ValueError("Cannot visualize an empty episode.")
+    first_ns, last_ns = times_ns[0], times_ns[-1]
+    camera_keys = list(dataset.meta.camera_keys)
+    # Dataset-wide q01/q99 depth bounds (fallback min/max) used to normalize depth to [0, 1].
+    depth_ranges: dict[str, tuple[float, float]] = {}
+    for key in dataset.meta.depth_keys:
+        stats = (dataset.meta.stats or {}).get(key)
+        if not stats:
+            continue
+        lo = stats["q01"] if "q01" in stats else stats["min"]
+        hi = stats["q99"] if "q99" in stats else stats["max"]
+        depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item()))
+    # Per-dimension series labels from the dataset metadata (e.g. joint names), computed once.
+    scalar_labels = {
+        OBS_STATE: _feature_dim_names(dataset.meta.features.get(OBS_STATE)),
+        ACTION: _feature_dim_names(dataset.meta.features.get(ACTION)),
+    }
+    # Local channel cache so the playback server is self-contained and doesn't touch the live-stream cache.
+    channels: dict = {}
+
+    def emit_frame(i: int) -> None:
+        """Log every channel for frame ``i`` stamped at its dataset timestamp."""
+        sample = dataset[i]
+        log_time = times_ns[i]
+        for key in camera_keys:
+            arr = sample.get(key)
+            if arr is None:
+                continue
+            arr = arr.numpy() if hasattr(arr, "numpy") else np.asarray(arr)
+            _log_foxglove_image(
+                _foxglove_topic(key, is_image=True),
+                key,
+                arr,
+                compress_images=compress_images,
+                channels=channels,
+                log_time=log_time,
+                depth_range=depth_ranges.get(key),
+                raw_depth_values=True,
+            )
+        _log_foxglove_scalars(
+            _foxglove_topic(OBS_STATE),
+            _frame_to_scalars(sample, OBS_STATE, scalar_labels[OBS_STATE]),
+            channels=channels,
+            log_time=log_time,
+        )
+        _log_foxglove_scalars(
+            _foxglove_topic(ACTION),
+            _frame_to_scalars(sample, ACTION, scalar_labels[ACTION]),
+            channels=channels,
+            log_time=log_time,
+        )
+        episode_scalars = {}
+        for feat, label in (
+            (DONE, "done"),
+            (TRUNCATED, "truncated"),
+            (REWARD, "reward"),
+            (SUCCESS, "success"),
+        ):
+            v = sample.get(feat)
+            if v is not None:
+                episode_scalars[label] = float(v)
+        _log_foxglove_scalars("/episode/state", episode_scalars, channels=channels, log_time=log_time)
+
+    lock = threading.Lock()
+    stop_event = threading.Event()
+    # Shared playback state, guarded by ``lock``. ``seek_idx`` is a one-shot request set by the
+    # listener and serviced by the playback loop, which is the *only* thread that emits frames (so
+    # concurrent random access into the on-disk dataset / video decoder never overlaps).
+    state = {
+        "status": PlaybackStatus.Paused,
+        "cursor": first_ns,
+        "speed": 1.0,
+        "last_idx": -1,
+        "seek_idx": None,
+    }
+
+    def index_at(t_ns: int) -> int:
+        return max(0, min(n_frames - 1, bisect.bisect_right(times_ns, t_ns) - 1))
+
+    # One-shot latch so autoplay fires only on the first client subscription.
+    autoplay_started = threading.Event()
+
+    class _PlaybackListener(ServerListener):
+        def on_subscribe(self, client, channel):
+            # Start playing automatically once a client actually connects (subscribes). Using the
+            # subscribe hook, rather than starting in Playing up front, means the timeline doesn't
+            # advance before anyone is watching. Fires once; the user can still pause/seek after.
+            if not autoplay:
+                return
+            with lock:
+                if autoplay_started.is_set() or state["status"] != PlaybackStatus.Paused:
+                    return
+                autoplay_started.set()
+                state["status"] = PlaybackStatus.Playing
+                cursor, speed = state["cursor"], state["speed"]
+            server.broadcast_playback_state(PlaybackState(PlaybackStatus.Playing, cursor, speed, False, ""))
+
+        def on_playback_control_request(self, req: PlaybackControlRequest):
+            # Only mutate state here; the playback loop performs all frame emission.
+            with lock:
+                did_seek = False
+                if req.seek_time is not None:
+                    cursor = max(first_ns, min(last_ns, req.seek_time))
+                    state["cursor"] = cursor
+                    state["last_idx"] = state["seek_idx"] = index_at(cursor)
+                    did_seek = True
+                if req.playback_speed and req.playback_speed > 0:
+                    state["speed"] = req.playback_speed
+                if req.playback_command == PlaybackCommand.Play:
+                    # Restarting from the end replays from the beginning.
+                    if state["cursor"] >= last_ns:
+                        state["cursor"] = first_ns
+                        state["last_idx"] = state["seek_idx"] = 0
+                        did_seek = True
+                    state["status"] = PlaybackStatus.Playing
+                elif req.playback_command == PlaybackCommand.Pause:
+                    state["status"] = PlaybackStatus.Paused
+                status, cursor, speed = state["status"], state["cursor"], state["speed"]
+                request_id = req.request_id or ""
+            return PlaybackState(status, cursor, speed, did_seek, request_id)
+
+    server = foxglove.start_server(
+        name=f"{dataset.repo_id}/episode_{episode_index}",
+        host=host,
+        port=port,
+        capabilities=[Capability.PlaybackControl, Capability.Time],
+        server_listener=_PlaybackListener(),
+        playback_time_range=(first_ns, last_ns),
+    )
+
+    def playback_loop() -> None:
+        # Cap how far the cursor may advance in a single tick. A slow frame decode (or any stall)
+        # would otherwise make ``dt`` huge and produce one enormous catch-up batch; clamping it makes
+        # playback trail wall-clock under a slow decoder while each tick emits a bounded frame range.
+        max_tick_dt_s = 0.25
+        prev = time.monotonic()
+        while not stop_event.is_set():
+            time.sleep(1.0 / 60.0)
+            ended = False
+            speed = 1.0
+            with lock:
+                now = time.monotonic()
+                dt = min(now - prev, max_tick_dt_s)
+                prev = now
+                # A queued seek is always serviced, even while paused, so scrubbing updates the view.
+                work = []
+                seek_idx = state["seek_idx"]
+                if seek_idx is not None:
+                    state["seek_idx"] = None
+                    work.append(seek_idx)
+                if state["status"] == PlaybackStatus.Playing:
+                    cursor = state["cursor"] + int(dt * 1e9 * state["speed"])
+                    start_idx = state["last_idx"] + 1
+                    if cursor >= last_ns:
+                        cursor, target, ended = last_ns, n_frames - 1, True
+                    else:
+                        target = index_at(cursor)
+                    state["cursor"] = cursor
+                    work.extend(range(start_idx, target + 1))
+                    # cursor only grows while playing (seeks reset last_idx in the listener), so
+                    # target >= last_idx here; a plain assignment is correct and clearer than max().
+                    state["last_idx"] = target
+                    if ended:
+                        state["status"] = PlaybackStatus.Ended
+                if not work:
+                    continue
+                cursor, speed = state["cursor"], state["speed"]
+            # Emit outside the lock; this is the only thread that calls emit_frame. Re-check
+            # stop_event between frames so shutdown stays responsive even mid-batch.
+            for i in work:
+                if stop_event.is_set():
+                    break
+                emit_frame(i)
+            server.broadcast_time(cursor)
+            if ended:
+                server.broadcast_playback_state(PlaybackState(PlaybackStatus.Ended, cursor, speed, False, ""))
+
+    # Emit the first frame so channels are advertised (done before the loop starts, so emission stays
+    # single-threaded). Late-connecting clients re-receive frames once they seek/play.
+    emit_frame(0)
+    with lock:
+        state["last_idx"] = 0
+    server.broadcast_time(first_ns)
+    server.broadcast_playback_state(PlaybackState(PlaybackStatus.Paused, first_ns, 1.0, True, ""))
+
+    thread = threading.Thread(target=playback_loop, name="foxglove-playback", daemon=True)
+    thread.start()
+
+    print(f"Foxglove server running. Connect the Foxglove app to ws://{host}:{port}")
+    print("Use the playback controls in Foxglove to play/pause and scrub the episode. Ctrl-C to exit.")
+    try:
+        while not stop_event.is_set():
+            time.sleep(0.5)
+    except KeyboardInterrupt:
+        print("Ctrl-C received. Exiting.")
+    finally:
+        stop_event.set()
+        thread.join(timeout=2.0)
+        server.stop()
+        channels.clear()

src/lerobot/utils/hub.py

@@ -20,9 +20,33 @@ from typing import Any, TypeVar
 from huggingface_hub import HfApi
 from huggingface_hub.utils import validate_hf_hub_args
 
+from .constants import CHECKPOINTS_DIR
+
 T = TypeVar("T", bound="HubMixin")
 
 
+def find_latest_hub_checkpoint(
+    repo_id: str,
+    *,
+    token: str | bool | None = None,
+    revision: str | None = None,
+) -> str | None:
+    """Repo-relative path of the most recent checkpoint in a training repo.
+
+    Training runs push checkpoints to ``checkpoints/<step>/`` (see
+    ``push_checkpoint_to_hub``). This lists those step dirs and returns
+    ``checkpoints/<highest-step>``, or ``None`` if the repo has no checkpoints.
+    """
+    files = HfApi().list_repo_files(repo_id=repo_id, repo_type="model", revision=revision, token=token)
+    prefix = f"{CHECKPOINTS_DIR}/"
+    steps = {
+        name for f in files if f.startswith(prefix) and (name := f[len(prefix) :].split("/", 1)[0]).isdigit()
+    }
+    if not steps:
+        return None
+    return f"{CHECKPOINTS_DIR}/{max(steps, key=int)}"
+
+
 class HubMixin:
     """
     A Mixin containing the functionality to push an object to the hub.

src/lerobot/utils/rerun_visualization.py

@@ -0,0 +1,191 @@
+# 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.
+
+"""Rerun visualization backend.
+
+Live control-loop streaming to the Rerun viewer (:func:`log_rerun_data`). Callers usually select a
+backend at runtime through the dispatch in :mod:`lerobot.utils.visualization_utils` rather than
+importing from here directly. Requires the ``viz`` extra (``pip install 'lerobot[viz]'``).
+"""
+
+import numbers
+import os
+
+import numpy as np
+
+from lerobot.configs import DEPTH_MILLIMETER_UNIT, infer_depth_unit
+from lerobot.types import RobotAction, RobotObservation
+
+from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR
+from .import_utils import require_package
+
+
+def _is_scalar(x):
+    return isinstance(x, (float | numbers.Real | np.integer | np.floating)) or (
+        isinstance(x, np.ndarray) and x.ndim == 0
+    )
+
+
+def init_rerun(
+    session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
+) -> None:
+    """
+    Initializes the Rerun SDK for visualizing the control loop.
+
+    Args:
+        session_name: Name of the Rerun session.
+        ip: Optional IP for connecting to a Rerun server.
+        port: Optional port for connecting to a Rerun server.
+    """
+
+    require_package("rerun-sdk", extra="viz", import_name="rerun")
+    import rerun as rr
+
+    log_rerun_data.blueprint = None  # Reset blueprint cache for new session
+
+    batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000")
+    os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size
+    rr.init(session_name)
+    memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%")
+    if ip and port:
+        rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy")
+    else:
+        rr.spawn(memory_limit=memory_limit)
+
+
+def shutdown_rerun() -> None:
+    """Shuts down the Rerun SDK gracefully."""
+
+    require_package("rerun-sdk", extra="viz", import_name="rerun")
+    import rerun as rr
+
+    rr.rerun_shutdown()
+
+
+def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]):
+    """Build a Rerun blueprint laying out camera images, observation and action scalars in separate views.
+
+    Camera images, observation and action scalars are arranged in a grid.
+    """
+
+    # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
+    import rerun.blueprint as rrb
+
+    views = [rrb.Spatial2DView(origin=path, name=path) for path in sorted(image_paths)]
+
+    if observation_paths:
+        views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths)))
+    if action_paths:
+        views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths)))
+
+    return rrb.Blueprint(rrb.Grid(*views))
+
+
+def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]) -> None:
+    """Build and send the blueprint once, from the first observation and action data."""
+    if getattr(log_rerun_data, "blueprint", None) is not None:
+        return
+
+    if not (observation_paths or action_paths or image_paths):
+        return
+
+    # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
+    import rerun as rr
+
+    blueprint = _build_blueprint(observation_paths, action_paths, image_paths)
+    log_rerun_data.blueprint = blueprint
+    rr.send_blueprint(blueprint)
+
+
+def log_rerun_data(
+    observation: RobotObservation | None = None,
+    action: RobotAction | None = None,
+    compress_images: bool = False,
+) -> None:
+    """
+    Logs observation and action data to Rerun for real-time visualization.
+
+    This function iterates through the provided observation and action dictionaries and sends their contents
+    to the Rerun viewer. It handles different data types appropriately:
+    - Scalars values (floats, ints) are logged as `rr.Scalars`.
+    - 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed
+      from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`.
+    - 1D NumPy arrays are logged as a single `rr.Scalars` batch under one entity path, so that every
+      dimension shares the same view instead of being split across one view per element.
+    - Multi-dimensional **action** arrays are flattened and logged as a single `rr.Scalars` batch.
+
+    Keys are automatically namespaced with "observation." or "action." if not already present.
+
+    On the first call, a blueprint is built and sent so observation and action scalars get separate
+    time-series views and each image gets its own spatial view.
+
+    Args:
+        observation: An optional dictionary containing observation data to log.
+        action: An optional dictionary containing action data to log.
+        compress_images: Whether to compress images before logging to save bandwidth & memory in exchange for cpu and quality.
+    """
+
+    require_package("rerun-sdk", extra="viz", import_name="rerun")
+    import rerun as rr
+
+    observation_paths: set[str] = set()
+    action_paths: set[str] = set()
+    image_paths: set[str] = set()
+
+    if observation:
+        for k, v in observation.items():
+            if v is None:
+                continue
+            key = k if str(k).startswith(OBS_PREFIX) else f"{OBS_STR}.{k}"
+
+            if _is_scalar(v):
+                rr.log(key, rr.Scalars(float(v)))
+                observation_paths.add(key)
+            elif isinstance(v, np.ndarray):
+                arr = v
+                # Convert CHW -> HWC when needed
+                if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
+                    arr = np.transpose(arr, (1, 2, 0))
+                if arr.ndim == 1:
+                    rr.log(key, rr.Scalars(arr.astype(float)))
+                    observation_paths.add(key)
+                else:
+                    if arr.shape[-1] == 1:
+                        # At record time, the depth unit is inferred from the frame type.
+                        depth_unit = infer_depth_unit(arr.dtype)
+                        img_entity = rr.DepthImage(
+                            arr,
+                            meter=1000.0 if depth_unit == DEPTH_MILLIMETER_UNIT else 1.0,
+                            colormap=rr.components.Colormap.Viridis,
+                        )
+                    else:
+                        img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
+                    rr.log(key, entity=img_entity, static=True)
+                    image_paths.add(key)
+
+    if action:
+        for k, v in action.items():
+            if v is None:
+                continue
+            key = k if str(k).startswith(ACTION_PREFIX) else f"{ACTION}.{k}"
+
+            if _is_scalar(v):
+                rr.log(key, rr.Scalars(float(v)))
+                action_paths.add(key)
+            elif isinstance(v, np.ndarray):
+                # Flatten any (incl. higher-dimensional) array into a single batched Scalars
+                rr.log(key, rr.Scalars(v.reshape(-1).astype(float)))
+                action_paths.add(key)
+
+    _ensure_blueprint(observation_paths, action_paths, image_paths)

src/lerobot/utils/visualization_utils.py

@@ -12,121 +12,68 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numbers
-import os
+"""Backend-agnostic visualization dispatch.
 
-import numpy as np
+Selects a visualization backend at runtime via a display-mode string (e.g. a ``--display_mode`` CLI
+flag) so callers never branch on the backend. The concrete implementations live in
+:mod:`lerobot.utils.rerun_visualization` and :mod:`lerobot.utils.foxglove_visualization`; importing
+this module does not import ``rerun`` or ``foxglove`` (each backend imports its SDK lazily behind a
+``require_package`` guard).
+"""
 
 from lerobot.types import RobotAction, RobotObservation
 
-from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR
-from .import_utils import require_package
+from .foxglove_visualization import init_foxglove, log_foxglove_data, shutdown_foxglove
+from .rerun_visualization import init_rerun, log_rerun_data, shutdown_rerun
+
+# Visualization backends selectable at runtime via a display-mode string (e.g. a --display_mode flag).
+VISUALIZATION_MODES = ("rerun", "foxglove")
 
 
-def init_rerun(
-    session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
+def init_visualization(
+    display_mode: str,
+    *,
+    session_name: str = "lerobot_control_loop",
+    ip: str | None = None,
+    port: int | None = None,
 ) -> None:
-    """
-    Initializes the Rerun SDK for visualizing the control loop.
+    """Initializes the visualization backend selected by ``display_mode``.
 
-    Args:
-        session_name: Name of the Rerun session.
-        ip: Optional IP for connecting to a Rerun server.
-        port: Optional port for connecting to a Rerun server.
+    For ``"rerun"``, ``ip``/``port`` point at an optional remote Rerun server. For ``"foxglove"``,
+    ``ip`` is the interface to bind the WebSocket server to (``127.0.0.1`` for local only, ``0.0.0.0``
+    for all interfaces) and ``port`` is its port.
     """
 
-    require_package("rerun-sdk", extra="viz", import_name="rerun")
-    import rerun as rr
-
-    batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000")
-    os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size
-    rr.init(session_name)
-    memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%")
-    if ip and port:
-        rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy")
+    if display_mode == "rerun":
+        init_rerun(session_name=session_name, ip=ip, port=port)
+    elif display_mode == "foxglove":
+        init_foxglove(host=ip or "127.0.0.1", port=port)
     else:
-        rr.spawn(memory_limit=memory_limit)
+        raise ValueError(f"Unknown display_mode '{display_mode}'. Expected one of {VISUALIZATION_MODES}.")
 
 
-def shutdown_rerun() -> None:
-    """Shuts down the Rerun SDK gracefully."""
-
-    require_package("rerun-sdk", extra="viz", import_name="rerun")
-    import rerun as rr
-
-    rr.rerun_shutdown()
-
-
-def _is_scalar(x):
-    return isinstance(x, (float | numbers.Real | np.integer | np.floating)) or (
-        isinstance(x, np.ndarray) and x.ndim == 0
-    )
-
-
-def log_rerun_data(
+def log_visualization_data(
+    display_mode: str,
     observation: RobotObservation | None = None,
     action: RobotAction | None = None,
     compress_images: bool = False,
 ) -> None:
-    """
-    Logs observation and action data to Rerun for real-time visualization.
+    """Logs observation/action data to the backend selected by ``display_mode``."""
 
-    This function iterates through the provided observation and action dictionaries and sends their contents
-    to the Rerun viewer. It handles different data types appropriately:
-    - Scalars values (floats, ints) are logged as `rr.Scalars`.
-    - 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed
-      from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`.
-    - 1D NumPy arrays are logged as a series of individual scalars, with each element indexed.
-    - Other multi-dimensional arrays are flattened and logged as individual scalars.
+    if display_mode == "rerun":
+        log_rerun_data(observation=observation, action=action, compress_images=compress_images)
+    elif display_mode == "foxglove":
+        log_foxglove_data(observation=observation, action=action, compress_images=compress_images)
+    else:
+        raise ValueError(f"Unknown display_mode '{display_mode}'. Expected one of {VISUALIZATION_MODES}.")
 
-    Keys are automatically namespaced with "observation." or "action." if not already present.
 
-    Args:
-        observation: An optional dictionary containing observation data to log.
-        action: An optional dictionary containing action data to log.
-        compress_images: Whether to compress images before logging to save bandwidth & memory in exchange for cpu and quality.
-    """
+def shutdown_visualization(display_mode: str) -> None:
+    """Shuts down the backend selected by ``display_mode``."""
 
-    require_package("rerun-sdk", extra="viz", import_name="rerun")
-    import rerun as rr
-
-    if observation:
-        for k, v in observation.items():
-            if v is None:
-                continue
-            key = k if str(k).startswith(OBS_PREFIX) else f"{OBS_STR}.{k}"
-
-            if _is_scalar(v):
-                rr.log(key, rr.Scalars(float(v)))
-            elif isinstance(v, np.ndarray):
-                arr = v
-                # Convert CHW -> HWC when needed
-                if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
-                    arr = np.transpose(arr, (1, 2, 0))
-                if arr.ndim == 1:
-                    for i, vi in enumerate(arr):
-                        rr.log(f"{key}_{i}", rr.Scalars(float(vi)))
-                else:
-                    if arr.shape[-1] == 1:
-                        img_entity = rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis)
-                    else:
-                        img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
-                    rr.log(key, entity=img_entity, static=True)
-
-    if action:
-        for k, v in action.items():
-            if v is None:
-                continue
-            key = k if str(k).startswith(ACTION_PREFIX) else f"{ACTION}.{k}"
-
-            if _is_scalar(v):
-                rr.log(key, rr.Scalars(float(v)))
-            elif isinstance(v, np.ndarray):
-                if v.ndim == 1:
-                    for i, vi in enumerate(v):
-                        rr.log(f"{key}_{i}", rr.Scalars(float(vi)))
-                else:
-                    # Fall back to flattening higher-dimensional arrays
-                    flat = v.flatten()
-                    for i, vi in enumerate(flat):
-                        rr.log(f"{key}_{i}", rr.Scalars(float(vi)))
+    if display_mode == "rerun":
+        shutdown_rerun()
+    elif display_mode == "foxglove":
+        shutdown_foxglove()
+    else:
+        raise ValueError(f"Unknown display_mode '{display_mode}'. Expected one of {VISUALIZATION_MODES}.")

tests/annotations/run_e2e_smoke.py

@@ -28,9 +28,10 @@ import sys
 import tempfile
 from pathlib import Path
 
-from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
+from lerobot.annotations.steerable_pipeline.config import AdvantageConfig, AnnotationPipelineConfig
 from lerobot.annotations.steerable_pipeline.executor import Executor
 from lerobot.annotations.steerable_pipeline.modules import (
+    AdvantageModule,
     GeneralVqaModule,
     InterjectionsAndSpeechModule,
     PlanSubtasksMemoryModule,
@@ -85,6 +86,7 @@ def main() -> int:
             plan=PlanSubtasksMemoryModule(vlm=vlm, config=cfg.plan),
             interjections=InterjectionsAndSpeechModule(vlm=vlm, config=cfg.interjections, seed=cfg.seed),
             vqa=GeneralVqaModule(vlm=vlm, config=cfg.vqa, seed=cfg.seed),
+            advantage=AdvantageModule(config=AdvantageConfig(enabled=False)),
             writer=LanguageColumnsWriter(),
             validator=StagingValidator(),
         )

tests/annotations/test_advantage.py

@@ -0,0 +1,305 @@
+#!/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.
+
+"""Tests for the advantage scoring annotation module."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+
+from lerobot.annotations.steerable_pipeline.config import AdvantageConfig
+from lerobot.annotations.steerable_pipeline.modules.advantage import AdvantageModule
+from lerobot.annotations.steerable_pipeline.reader import EpisodeRecord
+from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging
+
+
+def _make_record(
+    episode_index: int = 0,
+    num_frames: int = 20,
+    task: str = "pick up the cup",
+    mc_returns: np.ndarray | None = None,
+    intervention_mask: np.ndarray | None = None,
+    fps: float = 10.0,
+) -> EpisodeRecord:
+    """Build a minimal EpisodeRecord with a mocked frames_df."""
+    import pandas as pd
+
+    timestamps = tuple(round(i / fps, 6) for i in range(num_frames))
+    frame_indices = tuple(range(num_frames))
+
+    if mc_returns is None:
+        mc_returns = np.linspace(-0.9, -0.1, num_frames).astype(np.float32)
+
+    data = {
+        "episode_index": [episode_index] * num_frames,
+        "frame_index": list(range(num_frames)),
+        "timestamp": list(timestamps),
+        "mc_return": mc_returns,
+    }
+
+    if intervention_mask is not None:
+        data["intervention"] = intervention_mask.astype(bool)
+
+    df = pd.DataFrame(data)
+
+    record = EpisodeRecord(
+        episode_index=episode_index,
+        episode_task=task,
+        frame_timestamps=timestamps,
+        frame_indices=frame_indices,
+        data_path=Path("/fake/data.parquet"),
+        row_offset=0,
+        row_count=num_frames,
+    )
+    record._frames_df_cache = df
+    return record
+
+
+@pytest.fixture
+def staging(tmp_path: Path) -> EpisodeStaging:
+    return EpisodeStaging(tmp_path, episode_index=0)
+
+
+def test_advantage_module_disabled():
+    """Disabled module has enabled=False."""
+    cfg = AdvantageConfig(enabled=False)
+    module = AdvantageModule(config=cfg)
+    assert not module.enabled
+
+
+def test_advantage_module_enabled_by_default():
+    """Module is enabled by default."""
+    cfg = AdvantageConfig()
+    module = AdvantageModule(config=cfg)
+    assert module.enabled
+
+
+def test_run_episode_skips_without_value_function_path(staging: EpisodeStaging):
+    """Module gracefully returns when no value_function_path is configured."""
+    cfg = AdvantageConfig(value_function_path="")
+    module = AdvantageModule(config=cfg)
+    record = _make_record()
+
+    module.run_episode(record, staging)
+
+    rows = staging.read("advantage")
+    assert rows == []
+
+
+def test_binarization_with_mock_values(staging: EpisodeStaging):
+    """Advantage binarization produces positive/negative labels based on threshold."""
+    num_frames = 10
+    mc_returns = np.array([-0.5, -0.4, -0.3, -0.2, -0.1, -0.5, -0.6, -0.7, -0.8, -0.9], dtype=np.float32)
+    mock_values = np.array([-0.4, -0.4, -0.4, -0.4, -0.4, -0.4, -0.4, -0.4, -0.4, -0.4], dtype=np.float32)
+
+    cfg = AdvantageConfig(
+        value_function_path="/fake/vf",
+        dropout_rate=0.0,
+        threshold_percentile=0.5,
+    )
+    module = AdvantageModule(config=cfg)
+    record = _make_record(num_frames=num_frames, mc_returns=mc_returns)
+
+    with (
+        patch.object(module, "_ensure_model_loaded"),
+        patch.object(module, "_compute_values", return_value=mock_values),
+    ):
+        module.run_episode(record, staging)
+
+    rows = staging.read("advantage")
+    assert len(rows) == num_frames
+
+    # A_t = mc_returns - values
+    # advantages = [-0.1, 0.0, 0.1, 0.2, 0.3, -0.1, -0.2, -0.3, -0.4, -0.5]
+    # Median (50th pctile) = -0.1
+    # positive: advantage > -0.1 → indices 1,2,3,4
+    # negative: advantage <= -0.1 → indices 0,5,6,7,8,9
+    positives = [r for r in rows if r["content"] == "positive"]
+    negatives = [r for r in rows if r["content"] == "negative"]
+    assert len(positives) == 4
+    assert len(negatives) == 6
+
+
+def test_intervention_frames_forced_positive(staging: EpisodeStaging):
+    """Intervention frames are always scored as positive regardless of advantage value."""
+    num_frames = 5
+    mc_returns = np.array([-0.9, -0.9, -0.9, -0.9, -0.9], dtype=np.float32)
+    mock_values = np.array([-0.1, -0.1, -0.1, -0.1, -0.1], dtype=np.float32)
+    intervention = np.array([False, False, True, False, False])
+
+    cfg = AdvantageConfig(
+        value_function_path="/fake/vf",
+        dropout_rate=0.0,
+        force_positive_on_intervention=True,
+    )
+    module = AdvantageModule(config=cfg)
+    record = _make_record(num_frames=num_frames, mc_returns=mc_returns, intervention_mask=intervention)
+
+    with (
+        patch.object(module, "_ensure_model_loaded"),
+        patch.object(module, "_compute_values", return_value=mock_values),
+    ):
+        module.run_episode(record, staging)
+
+    rows = staging.read("advantage")
+    # Frame 2 (intervention) should be positive despite negative advantage
+    assert rows[2]["content"] == "positive"
+
+
+def test_dropout_reduces_output_rows(staging: EpisodeStaging):
+    """Non-zero dropout rate omits some frames."""
+    num_frames = 100
+    mc_returns = np.linspace(-0.9, -0.1, num_frames).astype(np.float32)
+    mock_values = np.full(num_frames, -0.5, dtype=np.float32)
+
+    cfg = AdvantageConfig(
+        value_function_path="/fake/vf",
+        dropout_rate=0.3,
+    )
+    module = AdvantageModule(config=cfg)
+    record = _make_record(num_frames=num_frames, mc_returns=mc_returns)
+
+    with (
+        patch.object(module, "_ensure_model_loaded"),
+        patch.object(module, "_compute_values", return_value=mock_values),
+    ):
+        module.run_episode(record, staging)
+
+    rows = staging.read("advantage")
+    # With 30% dropout on 100 frames, expect ~70 rows (with some variance)
+    assert 50 < len(rows) < 90
+
+
+def test_staged_row_format(staging: EpisodeStaging):
+    """Staged rows have the correct schema for language_persistent."""
+    num_frames = 5
+    mc_returns = np.array([-0.5, -0.4, -0.3, -0.2, -0.1], dtype=np.float32)
+    mock_values = np.full(5, -0.3, dtype=np.float32)
+
+    cfg = AdvantageConfig(
+        value_function_path="/fake/vf",
+        dropout_rate=0.0,
+    )
+    module = AdvantageModule(config=cfg)
+    record = _make_record(num_frames=num_frames, mc_returns=mc_returns)
+
+    with (
+        patch.object(module, "_ensure_model_loaded"),
+        patch.object(module, "_compute_values", return_value=mock_values),
+    ):
+        module.run_episode(record, staging)
+
+    rows = staging.read("advantage")
+    for row in rows:
+        assert row["role"] == "user"
+        assert row["content"] in ("positive", "negative")
+        assert row["style"] == "advantage"
+        assert isinstance(row["timestamp"], float)
+        assert row["camera"] is None
+        assert row["tool_calls"] is None
+
+
+def test_n_step_advantage():
+    """N-step advantage uses partial returns + bootstrapped value."""
+    num_frames = 10
+    mc_returns = np.linspace(-0.9, 0.0, num_frames).astype(np.float32)
+    mock_values = np.full(num_frames, -0.45, dtype=np.float32)
+
+    cfg = AdvantageConfig(
+        value_function_path="/fake/vf",
+        n_step=3,
+        dropout_rate=0.0,
+    )
+    module = AdvantageModule(config=cfg)
+    record = _make_record(num_frames=num_frames, mc_returns=mc_returns)
+
+    with patch.object(module, "_ensure_model_loaded"):
+        advantages, _ = (
+            module.compute_advantages_for_episode.__wrapped__(module, record)
+            if hasattr(module.compute_advantages_for_episode, "__wrapped__")
+            else (None, None)
+        )
+
+    # Just verify computation works - use the internal method directly
+    module._model = MagicMock()
+    module._preprocessor = MagicMock()
+    with patch.object(module, "_compute_values", return_value=mock_values):
+        advantages, _ = module.compute_advantages_for_episode(record)
+
+    # For t where t+n < num_frames: A = mc_return[t] - mc_return[t+n] + values[t+n] - values[t]
+    # Since values are constant: A = mc_return[t] - mc_return[t+n]
+    # For t where t+n >= num_frames: A = mc_return[t] - values[t]
+    for t in range(num_frames):
+        if t + 3 < num_frames:
+            expected = mc_returns[t] - mc_returns[t + 3] + mock_values[t + 3] - mock_values[t]
+        else:
+            expected = mc_returns[t] - mock_values[t]
+        np.testing.assert_almost_equal(advantages[t], expected, decimal=5)
+
+
+def test_compute_threshold():
+    """Threshold is computed as configured percentile of non-intervention advantages."""
+    cfg = AdvantageConfig(threshold_percentile=0.3)
+    module = AdvantageModule(config=cfg)
+
+    advantages = np.array([-1.0, -0.5, 0.0, 0.5, 1.0], dtype=np.float32)
+    intervention_mask = np.array([False, False, False, False, False])
+
+    threshold = module._compute_threshold(advantages, intervention_mask)
+    expected = float(np.percentile(advantages, 30))
+    assert abs(threshold - expected) < 1e-6
+
+
+def test_compute_threshold_excludes_intervention():
+    """Threshold computation excludes intervention frames."""
+    cfg = AdvantageConfig(threshold_percentile=0.5)
+    module = AdvantageModule(config=cfg)
+
+    advantages = np.array([100.0, -1.0, 0.0, 1.0, 100.0], dtype=np.float32)
+    intervention_mask = np.array([True, False, False, False, True])
+
+    threshold = module._compute_threshold(advantages, intervention_mask)
+    # Only non-intervention: [-1.0, 0.0, 1.0], median = 0.0
+    expected = float(np.percentile([-1.0, 0.0, 1.0], 50))
+    assert abs(threshold - expected) < 1e-6
+
+
+def test_missing_mc_return_raises():
+    """Module raises if mc_return column is missing from dataset."""
+    import pandas as pd
+
+    cfg = AdvantageConfig(value_function_path="/fake/vf")
+    module = AdvantageModule(config=cfg)
+    module._model = MagicMock()
+    module._preprocessor = MagicMock()
+
+    record = EpisodeRecord(
+        episode_index=0,
+        episode_task="test",
+        frame_timestamps=(0.0, 0.1),
+        frame_indices=(0, 1),
+        data_path=Path("/fake/data.parquet"),
+        row_offset=0,
+        row_count=2,
+    )
+    record._frames_df_cache = pd.DataFrame({"episode_index": [0, 0], "frame_index": [0, 1]})
+
+    with pytest.raises(KeyError, match="mc_return"):
+        module.compute_advantages_for_episode(record)