adding lerobot-train requirement inside PR checklist

.dockerignore

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

AGENTS.md

@@ -51,6 +51,7 @@ pre-commit run --all-files                           # Lint + format (ruff, typo
 ## Notes
 
 - **Mypy is gradual**: strict only for `lerobot.envs`, `lerobot.configs`, `lerobot.optim`, `lerobot.model`, `lerobot.cameras`, `lerobot.motors`, `lerobot.transport`. Add type annotations when modifying these modules.
-- **Optional dependencies**: many policies, envs, and robots are behind extras (e.g., `lerobot[aloha]`). New imports for optional packages must be guarded or lazy. See `pyproject.toml [project.optional-dependencies]`.
+- **Imports**: prefer top-level imports; relative (`from .sibling import X`) across sibling files within a module, absolute (`from lerobot.module import X`) across modules.
+- **Optional dependencies**: many policies, envs, and robots are behind extras (e.g., `lerobot[aloha]`, see `pyproject.toml`). Guard optional imports with `TYPE_CHECKING or _foo_available` at module top + a `require_package(...)` check at use time. Reuse the `_foo_available` flags in `utils/import_utils.py`; don't call `is_package_available`.
 - **Video decoding**: datasets can store observations as video files. `LeRobotDataset` handles frame extraction, but tests need ffmpeg installed.
 - **Prioritize use of `uv run`** to execute Python commands (not raw `python` or `pip`).

docs/source/_toctree.yml

@@ -69,8 +69,6 @@
     title: VLA-JEPA
   - local: eo1
     title: EO-1
-  - local: lingbot_va
-    title: LingBot-VA
   - local: fastwam
     title: FastWAM
   - local: groot

docs/source/bring_your_own_policies.mdx

@@ -165,6 +165,8 @@ Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constant
 
 LeRobot uses `PolicyProcessorPipeline`s to normalize inputs and de-normalize outputs around your policy. For a concrete reference, see [`processor_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/processor_act.py) or [`processor_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/processor_diffusion.py).
 
+Pay close attention here: processors are the most common reproducibility pain point. A mismatch in normalization mode (`IDENTITY` vs `MEAN_STD` vs `MIN_MAX` vs `QUANTILES`/`QUANTILE10`) or in which features get normalized will train and eval without erroring, yet silently wreck results. Make sure the modes match how the checkpoint was trained, that the required stats exist (e.g. `QUANTILES` needs `q01`/`q99`), and that the pre- and post-processors stay consistent.
+
 ```python
 # processor_my_policy.py
 from typing import Any
@@ -371,6 +373,7 @@ The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingfa
 - [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard.
 - [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests.
 - [ ] `src/lerobot/policies/<name>/README.md` symlinked into `docs/source/policy_<name>_README.md`; user-facing `docs/source/<name>.mdx` written and added to `_toctree.yml`.
+- [ ] `lerobot-train --policy.type my_policy ...` runs end-to-end for at least a few steps + save a checkpoint that can be loaded and run by `lerobot-eval` or `lerobot-rollout`.
 - [ ] At least one reproducible benchmark eval in the policy MDX with a published checkpoint (sim benchmark, or real-robot dataset + checkpoint).
 
 The fastest way to get a clean PR is to copy the directory of the existing policy closest to yours, rename, and replace contents method by method. Don't wait until everything is polished — open a draft PR early and iterate with us; reviewers would much rather give feedback on a half-finished branch than a fully-merged one.

docs/source/hardware_guide.mdx

@@ -82,18 +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, without owning a GPU. `lerobot-train` submits and streams the job for you — just add `--job.target=<flavor>` to a normal training command:
+[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.
 
 ```bash
-lerobot-train \
-  --policy.type=act --dataset.repo_id=<USER>/<DATASET> \
-  --policy.repo_id=<USER>/act_<task> \
-  --job.target=a10g-large
+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"
 ```
 
 Notes:
 
-- 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).
+- The leading `nvidia-smi` is a quick sanity check that CUDA is visible inside the container — useful to fail fast if the flavor or driver mismatched.
+- The default Job timeout is 30 minutes; pass `--timeout 4h` (or longer) for real training.
+- `--flavor` maps onto the table above: `t4-small`/`t4-medium` (T4, ACT only), `l4x1`/`l4x4` (L4 24 GB), `a10g-small/large/largex2/largex4` (A10G 24 GB scaled out), `a100-large` (A100). For the current full catalogue + pricing see [https://huggingface.co/docs/hub/jobs](https://huggingface.co/docs/hub/jobs).
+- Prefer not to write the `hf jobs run` wrapper yourself? `lerobot-train` can submit the job for you: just add `--job.target=<flavor>` to a normal training command and it handles dataset upload, log streaming, and the final model push. See the [imitation-learning training guide](./il_robots).

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_visualization, log_visualization_data, shutdown_visualization
+from lerobot.utils.visualization_utils import init_rerun, log_rerun_data, shutdown_rerun
 
 robot_config = SO101FollowerConfig(
     port="/dev/tty.usbmodem5AB90687491",
@@ -142,7 +142,7 @@ teleop_config = SO101LeaderConfig(
     id="my_leader_arm",
 )
 
-init_visualization("rerun", session_name="teleoperation")  # pass "foxglove" to stream to Foxglove instead
+init_rerun(session_name="teleoperation")
 
 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_visualization_data("rerun", observation=observation, action=action)
+    log_rerun_data(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_visualization
+from lerobot.utils.visualization_utils import init_rerun
 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_visualization("rerun", session_name="recording")
+    init_rerun(session_name="recording")
 
     # Connect the robot and teleoperator
     robot.connect()
@@ -532,7 +532,84 @@ 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).
 
-`lerobot-train` runs locally by default. To run on a HuggingFace GPU, pass `--job.target` with a hardware flavor name:
+> **Tip:** if you just want to launch a standard training run, you can skip building the command below and use the integrated **Train on HF Jobs via `--job.target`** flow described further down — `lerobot-train` then submits the job, uploads a local-only dataset for you, and streams the logs.
+
+To run the training manually use this command:
+
+<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
+```
+</hfoption>
+<hfoption id="API example">
+
+<!-- prettier-ignore-start -->
+```python
+from huggingface_hub import run_job, get_token
+
+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}")
+```
+<!-- prettier-ignore-end -->
+
+</hfoption>
+</hfoptions>
+
+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.
+
+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.
+
+After training the model will be pushed to hub and you can use it as any other model with LeRobot.
+
+#### Train on HF Jobs via `--job.target` (integrated CLI)
+
+`lerobot-train` runs locally by default. To run on a HuggingFace GPU without constructing the Docker command yourself, pass `--job.target` with a hardware flavor name:
 
 ```bash
 lerobot-train \

docs/source/lingbot_va.mdx

@@ -1,187 +0,0 @@
-# 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/using_dataset_tools.mdx

@@ -265,8 +265,6 @@ 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

@@ -125,7 +125,6 @@ hardware = [
 ]
 viz = [
     "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.
@@ -236,7 +235,6 @@ fastwam = [
 ]
 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]"]
@@ -319,7 +317,6 @@ all = [
     "lerobot[xvla]",
     "lerobot[hilserl]",
     "lerobot[vla_jepa]",
-    "lerobot[lingbot_va]",
     "lerobot[async]",
     "lerobot[dev]",
     "lerobot[test]",

src/lerobot/configs/__init__.py

@@ -34,8 +34,6 @@ from .types import (
 )
 from .video import (
     DEFAULT_DEPTH_UNIT,
-    DEPTH_METER_UNIT,
-    DEPTH_MILLIMETER_UNIT,
     VALID_VIDEO_CODECS,
     VIDEO_ENCODER_INFO_KEYS,
     DepthEncoderConfig,
@@ -43,7 +41,6 @@ from .video import (
     VideoEncoderConfig,
     depth_encoder_defaults,
     encoder_config_from_video_info,
-    infer_depth_unit,
     rgb_encoder_defaults,
 )
 
@@ -73,11 +70,8 @@ __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/video.py

@@ -22,8 +22,6 @@ 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__)
@@ -38,9 +36,7 @@ HW_VIDEO_CODECS = [
     "h264_vaapi",  # Linux Intel/AMD
     "h264_qsv",  # Intel Quick Sync
 ]
-VALID_VIDEO_CODECS: frozenset[str] = frozenset(
-    {"h264", "hevc", "libsvtav1", "libaom-av1", "auto", *HW_VIDEO_CODECS}
-)
+VALID_VIDEO_CODECS: frozenset[str] = frozenset({"h264", "hevc", "libsvtav1", "auto", *HW_VIDEO_CODECS})
 # Aliases for legacy video codec names.
 VIDEO_CODECS_ALIASES: dict[str, str] = {"av1": "libsvtav1"}
 
@@ -69,15 +65,6 @@ 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"})
 
@@ -226,24 +213,18 @@ class VideoEncoderConfig:
             if encoder_threads is not None:
                 svtav1_parts.append(f"lp={encoder_threads}")
             if svtav1_parts:
-                set_if("svtav1-params", ":".join(svtav1_parts))
+                opts["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:
-                set_if("tune", "fastdecode")
+                opts["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:
-                set_if("q:v", max(1, min(100, 100 - self.crf * 2)))
+                opts["q:v"] = max(1, min(100, 100 - self.crf * 2))
         elif self.vcodec in ("h264_nvenc", "hevc_nvenc"):
-            set_if("rc", 0)
+            opts["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 (stored in ``depth_unit``).
+        this rescaling and remain in their stored units.
     """
     if quantile_list is None:
         quantile_list = DEFAULT_QUANTILES

src/lerobot/datasets/dataset_metadata.py

@@ -26,13 +26,12 @@ import pyarrow as pa
 import pyarrow.parquet as pq
 from huggingface_hub import snapshot_download
 
-from lerobot.configs import DEPTH_METER_UNIT, VideoEncoderConfig
+from lerobot.configs import 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,
@@ -359,35 +358,6 @@ 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,14 +22,10 @@ from pathlib import Path
 import datasets
 import torch
 
-from lerobot.configs import (
-    DEFAULT_DEPTH_UNIT,
-    DEPTH_METER_UNIT,
-    DepthEncoderConfig,
-)
+from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig
 
 from .dataset_metadata import LeRobotDatasetMetadata
-from .depth_utils import MM_PER_METRE, dequantize_depth
+from .depth_utils import dequantize_depth
 from .feature_utils import (
     check_delta_timestamps,
     get_delta_indices,
@@ -106,13 +102,6 @@ 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):
@@ -340,13 +329,6 @@ 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,7 +36,6 @@ from lerobot.configs import (
     RGBEncoderConfig,
     VideoEncoderConfig,
     depth_encoder_defaults,
-    infer_depth_unit,
     rgb_encoder_defaults,
 )
 
@@ -210,15 +209,6 @@ 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,13 +34,12 @@ 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
 
 
@@ -58,7 +57,11 @@ 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 = infer_depth_unit(depth.dtype) if input_unit == "auto" else input_unit
+    resolved_unit = (
+        (DEPTH_METER_UNIT if np.issubdtype(depth.dtype, np.floating) else DEPTH_MILLIMETER_UNIT)
+        if input_unit == "auto"
+        else input_unit
+    )
     return depth.astype(np.float32, order="K"), resolved_unit
 
 
@@ -123,12 +126,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:
@@ -233,7 +236,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)
@@ -256,7 +259,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/lerobot_dataset.py

@@ -224,7 +224,6 @@ 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
@@ -351,11 +350,6 @@ 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, DEPTH_METER_UNIT, DepthEncoderConfig
+from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig
 from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
 
 from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
-from .depth_utils import MM_PER_METRE, dequantize_depth
+from .depth_utils import dequantize_depth
 from .feature_utils import get_delta_indices
 from .io_utils import item_to_torch
 from .utils import (
@@ -310,7 +310,6 @@ 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)
 
@@ -319,13 +318,6 @@ 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
 
@@ -356,11 +348,6 @@ 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
@@ -543,15 +530,6 @@ 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 = 1200
+    episode_length: int = 300
     obs_type: str = "pixels_agent_pos"
     render_mode: str = "rgb_array"
     # Available cameras from RoboTwin's aloha-agilex embodiment: head_camera
@@ -768,9 +768,6 @@ 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,)),
@@ -787,8 +784,6 @@ 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(
@@ -831,7 +826,6 @@ 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,7 +17,6 @@ from __future__ import annotations
 
 import importlib
 import logging
-import os
 from collections import defaultdict
 from collections.abc import Callable, Sequence
 from functools import partial
@@ -29,17 +28,9 @@ 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
@@ -50,124 +41,10 @@ ROBOTWIN_CAMERA_NAMES: tuple[str, ...] = (
     "right_camera",
 )
 
-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_DIM = 14  # 7 DOF × 2 arms
 ACTION_LOW = -1.0
 ACTION_HIGH = 1.0
-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))
-
-
+DEFAULT_EPISODE_LENGTH = 300
 # D435 dims from task_config/_camera_config.yml (what demo_clean.yml selects).
 DEFAULT_CAMERA_H = 240
 DEFAULT_CAMERA_W = 320
@@ -357,7 +234,6 @@ 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
@@ -365,13 +241,6 @@ 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
@@ -402,7 +271,7 @@ class RoboTwinEnv(gym.Env):
             }
         )
         self.action_space = spaces.Box(
-            low=ACTION_LOW, high=ACTION_HIGH, shape=(self._action_dim,), dtype=np.float32
+            low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32
         )
 
     def _ensure_env(self) -> None:
@@ -448,18 +317,6 @@ 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)
@@ -473,32 +330,16 @@ 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] != self._action_dim:
-            raise ValueError(f"Expected 1-D action of shape ({self._action_dim},), got {action.shape}")
+        if action.ndim != 1 or action.shape[0] != ACTION_DIM:
+            raise ValueError(f"Expected 1-D action of shape ({ACTION_DIM},), got {action.shape}")
 
         with torch.enable_grad():
-            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"):
+            if hasattr(self._env, "take_action"):
                 self._env.take_action(action)
             else:
                 self._env.step(action)
@@ -557,7 +398,6 @@ 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."""
 
@@ -570,7 +410,6 @@ 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)]
@@ -584,7 +423,6 @@ 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.
 
@@ -635,7 +473,6 @@ 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/optim/schedulers.py

@@ -83,28 +83,6 @@ 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

@@ -21,7 +21,6 @@ from .factory import get_policy_class, make_policy, make_policy_config, make_pre
 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
@@ -47,7 +46,6 @@ __all__ = [
     "FastWAMConfig",
     "GaussianActorConfig",
     "GrootConfig",
-    "LingBotVAConfig",
     "MolmoAct2Config",
     "MultiTaskDiTConfig",
     "PI0Config",

src/lerobot/policies/factory.py

@@ -50,7 +50,6 @@ 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
@@ -164,10 +163,6 @@ 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
 
@@ -228,8 +223,6 @@ 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:
@@ -465,14 +458,6 @@ 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
 

src/lerobot/policies/lingbot_va/README.md

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

src/lerobot/policies/lingbot_va/__init__.py

@@ -1,21 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-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

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

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

@@ -1,87 +0,0 @@
-# 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/rollout/configs.py

@@ -106,8 +106,6 @@ class DAggerKeyboardConfig:
     pause_resume: str = "space"
     correction: str = "tab"
     upload: str = "enter"
-    success: str = "s"
-    failure: str = "f"
 
 
 @dataclass
@@ -167,10 +165,6 @@ 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
@@ -232,14 +226,11 @@ class RolloutConfig:
     device: str | None = None
     task: str = ""
     display_data: bool = False
-    # 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 data on a remote Rerun server
     display_ip: str | None = None
-    # For "rerun": port of the remote server. For "foxglove": port to bind the WebSocket server to.
+    # Port of the remote Rerun server
     display_port: int | None = None
-    # Whether to display compressed (JPEG) images instead of raw frames
+    # Whether to display compressed images in Rerun
     display_compressed_images: bool = False
     # Use vocal synthesis to read events
     play_sounds: bool = True

src/lerobot/rollout/context.py

@@ -350,11 +350,6 @@ 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_visualization_data
+from lerobot.utils.visualization_utils import log_rerun_data
 
 from ..inference import InferenceEngine
 
@@ -162,12 +162,11 @@ class RolloutStrategy(abc.ABC):
         action_dict: dict | None,
         runtime_ctx: RuntimeContext,
     ) -> None:
-        """Log observation/action telemetry to the visualization backend if display_data is enabled."""
+        """Log observation/action telemetry to Rerun if display_data is enabled."""
         cfg = runtime_ctx.cfg
         if not cfg.display_data:
             return
-        log_visualization_data(
-            cfg.display_mode,
+        log_rerun_data(
             observation=obs_processed,
             action=action_dict,
             compress_images=cfg.display_compressed_images,

src/lerobot/rollout/strategies/dagger.py

@@ -112,11 +112,6 @@ 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
 
     # -- Thread-safe phase access ------------------------------------------
 
@@ -160,26 +155,7 @@ 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
 
 
 # ---------------------------------------------------------------------------
@@ -210,20 +186,12 @@ 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}', success='{cfg.success}', failure='{cfg.failure}', ESC=stop"
+            f"upload='{cfg.upload}', ESC=stop"
         ),
     )
 
@@ -345,32 +313,6 @@ 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()
-        logger.info("_stamp_episode_success: label=%s, buffer_len=%d", label, len(success_buf))
-
-        if label:
-            success_buf[-1] = np.array([True], dtype=bool)
-            logger.info("Terminal frame stamped next.success=True")
-
     # ------------------------------------------------------------------
     # Continuous recording mode (record_autonomous=True)
     # ------------------------------------------------------------------
@@ -408,12 +350,7 @@ class DAggerStrategy(RolloutStrategy):
         episode_start = time.perf_counter()
         episodes_since_push = 0
         episode_duration_s = self._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 "∞",
-        )
+        logger.info("DAgger continuous recording started (episode_duration=%.0fs)", episode_duration_s)
 
         with VideoEncodingManager(dataset):
             try:
@@ -462,7 +399,6 @@ 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
@@ -491,32 +427,23 @@ 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: 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()
-
+                    # 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.
                     elapsed = time.perf_counter() - episode_start
-                    if (manual_save or elapsed >= episode_duration_s) and phase != DAggerPhase.CORRECTING:
-                        if manual_save:
-                            events.save_episode_requested.clear()
+                    if elapsed >= episode_duration_s and phase != DAggerPhase.CORRECTING:
                         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 (%s, total: %d)",
-                            save_reason,
+                            "Episode saved (total: %d, elapsed: %.1fs)",
                             dataset.num_episodes,
+                            elapsed,
                         )
                         log_say(f"Episode {dataset.num_episodes} saved", play_sounds)
 
@@ -524,25 +451,6 @@ 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
@@ -557,13 +465,10 @@ class DAggerStrategy(RolloutStrategy):
                 logger.info("DAgger continuous control loop ended — pausing engine")
                 engine.pause()
                 with contextlib.suppress(Exception):
-                    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")
+                    with self._episode_lock:
+                        dataset.save_episode()
+                    self._needs_push.set()
+                    logger.info("Final in-progress episode saved")
 
     # ------------------------------------------------------------------
     # Corrections-only mode (record_autonomous=False)
@@ -635,7 +540,6 @@ 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()
@@ -677,7 +581,6 @@ class DAggerStrategy(RolloutStrategy):
                                     **action_frame,
                                     "task": task_str,
                                     "intervention": np.array([True], dtype=bool),
-                                    "next.success": np.array([False], dtype=bool),
                                 }
                             )
                         record_tick += 1
@@ -711,13 +614,10 @@ class DAggerStrategy(RolloutStrategy):
                 logger.info("DAgger corrections-only loop ended — pausing engine")
                 engine.pause()
                 with contextlib.suppress(Exception):
-                    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")
+                    with self._episode_lock:
+                        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_visualization_data
+from lerobot.utils.visualization_utils import log_rerun_data
 
 from ..configs import EpisodicStrategyConfig
 from ..context import RolloutContext
@@ -171,7 +171,6 @@ 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,
                         )
 
@@ -260,7 +259,6 @@ 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."""
@@ -290,8 +288,7 @@ class EpisodicStrategy(RolloutStrategy):
 
                 if display_data:
                     obs_processed = processors.robot_observation_processor(obs)
-                    log_visualization_data(
-                        display_mode,
+                    log_rerun_data(
                         observation=obs_processed,
                         action=act_teleop,
                         compress_images=display_compressed,

src/lerobot/scripts/lerobot_dataset_viz.py

@@ -59,18 +59,6 @@ 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
@@ -84,14 +72,10 @@ 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, SUCCESS
+from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD
 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.
@@ -124,12 +108,6 @@ def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:
     return hwc_uint8_numpy
 
 
-def to_hwc_float32_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:
-    check_chw_float32(chw_float32_torch)
-    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.
 
@@ -148,43 +126,32 @@ def build_blueprint_from_dataset(dataset: LeRobotDataset):
             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):
+    for key in (DONE, REWARD, "next.success"):
         if key in dataset.features:
             views.append(rrb.TimeSeriesView(origin=key, name=key))
 
     return rrb.Blueprint(rrb.Grid(*views))
 
 
+def to_hwc_uint16_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
+
+
 def visualize_dataset(
     dataset: LeRobotDataset,
     episode_index: int,
     batch_size: int = 32,
     num_workers: int = 0,
     mode: str = "local",
-    web_port: int | None = None,
+    web_port: int = 9090,
     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`."
@@ -221,23 +188,14 @@ 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 if web_port is not None else DEFAULT_RERUN_PORT,
-            connect_to=server_uri,
-        )
+        rr.serve_web_viewer(open_browser=False, web_port=web_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 or {}).get(key)
-        if not stats:
-            continue
+        stats = dataset.meta.stats[key]
         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()))
@@ -255,12 +213,11 @@ def visualize_dataset(
             # display each camera image (or depth map)
             for key in dataset.meta.camera_keys:
                 if key in dataset.meta.depth_keys:
-                    depth = to_hwc_float32_numpy(batch[key][i])
+                    depth = to_hwc_uint16_numpy(batch[key][i])
                     depth_entity = rr.DepthImage(
                         depth,
-                        meter=depth_meter,
                         colormap=rr.components.Colormap.Viridis,
-                        depth_range=depth_ranges.get(key),
+                        depth_range=depth_ranges[key],
                     )
                     rr.log(key, entity=depth_entity)
                 else:
@@ -282,8 +239,8 @@ def visualize_dataset(
             if REWARD in batch:
                 rr.log(REWARD, rr.Scalars(batch[REWARD][i].item()))
 
-            if SUCCESS in batch:
-                rr.log(SUCCESS, rr.Scalars(batch[SUCCESS][i].item()))
+            if "next.success" in batch:
+                rr.log("next.success", rr.Scalars(batch["next.success"][i].item()))
 
     # save .rrd locally
     if mode == "local" and save:
@@ -355,11 +312,13 @@ def main():
     parser.add_argument(
         "--web-port",
         type=int,
-        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)."
-        ),
+        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.",
     )
     parser.add_argument(
         "--grpc-port",
@@ -392,56 +351,24 @@ def main():
     parser.add_argument(
         "--display-compressed-images",
         action="store_true",
-        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."
-        ),
+        help="If set, display compressed images in Rerun instead of uncompressed ones.",
     )
 
     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")
+    else:
+        kwargs.pop("ws_port")  # Always remove ws_port from kwargs
+
     init_logging()
     logging.info("Loading dataset")
     dataset = LeRobotDataset(repo_id, episodes=[args.episode_index], root=root, tolerance_s=tolerance_s)

src/lerobot/scripts/lerobot_eval.py

@@ -169,7 +169,6 @@ 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.
 
@@ -199,9 +198,6 @@ 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.
     """
@@ -280,8 +276,6 @@ 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}
@@ -301,22 +295,12 @@ def rollout(
             # available if none of the envs finished.
             if "final_info" in info:
                 final_info = info["final_info"]
-                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
+                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."
                     )
-                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)
+                successes = final_info["is_success"].tolist()
             elif "is_success" in info:
                 is_success = info["is_success"]
                 successes = (
@@ -416,7 +400,6 @@ 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:
@@ -435,11 +418,6 @@ 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."
@@ -483,22 +461,6 @@ 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
 
@@ -510,9 +472,6 @@ 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:
@@ -533,7 +492,6 @@ 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
@@ -599,35 +557,6 @@ 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}%"}
         )
@@ -671,9 +600,6 @@ 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
 
 
@@ -814,10 +740,9 @@ 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", "predicted_video_paths")
+ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths")
 
 
 def eval_one(
@@ -866,7 +791,6 @@ 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", []),
     )
 
 
@@ -927,7 +851,6 @@ def run_one(
 
     if max_episodes_rendered > 0:
         metrics.setdefault("video_paths", [])
-    metrics.setdefault("predicted_video_paths", [])
     return task_group, task_id, metrics
 
 
@@ -985,11 +908,11 @@ def eval_policy_all(
         _append("sum_rewards", metrics.get("sum_rewards"))
         _append("max_rewards", metrics.get("max_rewards"))
         _append("successes", metrics.get("successes"))
-        for key in ("video_paths", "predicted_video_paths"):
-            paths = metrics.get(key, [])
-            if paths:
-                group_acc[group][key].extend(paths)
-                overall[key].extend(paths)
+        # video_paths is list-like
+        paths = metrics.get("video_paths", [])
+        if paths:
+            group_acc[group]["video_paths"].extend(paths)
+            overall["video_paths"].extend(paths)
 
     # Choose runner (sequential vs threaded)
     task_runner = partial(
@@ -1061,7 +984,6 @@ 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
@@ -1073,7 +995,6 @@ 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,9 +38,6 @@ 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 \\
@@ -160,11 +157,7 @@ from lerobot.utils.utils import (
     init_logging,
     log_say,
 )
-from lerobot.utils.visualization_utils import (
-    init_visualization,
-    log_visualization_data,
-    shutdown_visualization,
-)
+from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 
 
 @dataclass
@@ -175,14 +168,11 @@ class RecordConfig:
     teleop: TeleoperatorConfig | None = None
     # Display all cameras on screen
     display_data: bool = False
-    # 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 data on a remote Rerun server
     display_ip: str | None = None
-    # For "rerun": port of the remote server. For "foxglove": port to bind the WebSocket server to.
+    # Port of the remote Rerun server
     display_port: int | None = None
-    # Whether to display compressed (JPEG) images instead of raw frames
+    # Whether to  display compressed images in Rerun
     display_compressed_images: bool = False
     # Use vocal synthesis to read events.
     play_sounds: bool = True
@@ -243,7 +233,6 @@ 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:
@@ -338,11 +327,8 @@ def record_loop(
             dataset.add_frame(frame)
 
         if display_data:
-            log_visualization_data(
-                display_mode,
-                observation=obs_processed,
-                action=action_values,
-                compress_images=display_compressed_images,
+            log_rerun_data(
+                observation=obs_processed, action=action_values, compress_images=display_compressed_images
             )
 
         dt_s = time.perf_counter() - start_loop_t
@@ -368,9 +354,7 @@ def record(
     init_logging()
     logging.info(pformat(asdict(cfg)))
     if cfg.display_data:
-        init_visualization(
-            cfg.display_mode, session_name="recording", ip=cfg.display_ip, port=cfg.display_port
-        )
+        init_rerun(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)
@@ -480,7 +464,6 @@ 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,
                 )
 
@@ -502,7 +485,6 @@ 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"]:
@@ -528,9 +510,6 @@ 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,9 +145,6 @@ 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
@@ -193,7 +190,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_visualization, shutdown_visualization
+from lerobot.utils.visualization_utils import init_rerun
 
 logger = logging.getLogger(__name__)
 
@@ -204,13 +201,8 @@ def rollout(cfg: RolloutConfig):
     init_logging()
 
     if cfg.display_data:
-        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)
+        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)
 
     signal_handler = ProcessSignalHandler(use_threads=True, display_pid=False)
     shutdown_event = signal_handler.shutdown_event
@@ -235,8 +227,6 @@ 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,22 +31,6 @@ 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
@@ -124,11 +108,7 @@ 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_visualization,
-    log_visualization_data,
-    shutdown_visualization,
-)
+from lerobot.utils.visualization_utils import init_rerun, log_rerun_data, shutdown_rerun
 
 
 @dataclass
@@ -141,14 +121,11 @@ class TeleoperateConfig:
     teleop_time_s: float | None = None
     # Display all cameras on screen
     display_data: bool = False
-    # 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 data on a remote Rerun server
     display_ip: str | None = None
-    # For "rerun": port of the remote server. For "foxglove": port to bind the WebSocket server to.
+    # Port of the remote Rerun server
     display_port: int | None = None
-    # Whether to display compressed (JPEG) images instead of raw frames
+    # Whether to  display compressed images in Rerun
     display_compressed_images: bool = False
 
 
@@ -160,7 +137,6 @@ 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,
 ):
@@ -173,10 +149,8 @@ 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 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.
+        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.
         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.
@@ -213,8 +187,7 @@ def teleop_loop(
             # Process robot observation through pipeline
             obs_transition = robot_observation_processor(obs)
 
-            log_visualization_data(
-                display_mode,
+            log_rerun_data(
                 observation=obs_transition,
                 action=teleop_action,
                 compress_images=display_compressed_images,
@@ -242,9 +215,7 @@ def teleoperate(cfg: TeleoperateConfig):
     init_logging()
     logging.info(pformat(asdict(cfg)))
     if cfg.display_data:
-        init_visualization(
-            cfg.display_mode, session_name="teleoperation", ip=cfg.display_ip, port=cfg.display_port
-        )
+        init_rerun(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)
@@ -264,7 +235,6 @@ 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,
@@ -275,7 +245,7 @@ def teleoperate(cfg: TeleoperateConfig):
         pass
     finally:
         if cfg.display_data:
-            shutdown_visualization(cfg.display_mode)
+            shutdown_rerun()
         teleop.disconnect()
         robot.disconnect()
 

src/lerobot/scripts/lerobot_train.py

@@ -211,12 +211,8 @@ 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,
         )

src/lerobot/utils/constants.py

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

@@ -1,651 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""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/rerun_visualization.py

@@ -1,191 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""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,68 +12,166 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Backend-agnostic visualization dispatch.
+import numbers
+import os
 
-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).
-"""
+import numpy as np
 
 from lerobot.types import RobotAction, RobotObservation
 
-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")
+from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR
+from .import_utils import require_package
 
 
-def init_visualization(
-    display_mode: str,
-    *,
-    session_name: str = "lerobot_control_loop",
-    ip: str | None = None,
-    port: int | None = None,
+def init_rerun(
+    session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
 ) -> None:
-    """Initializes the visualization backend selected by ``display_mode``.
+    """
+    Initializes the Rerun SDK for visualizing the control loop.
 
-    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.
+    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.
     """
 
-    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)
+    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:
-        raise ValueError(f"Unknown display_mode '{display_mode}'. Expected one of {VISUALIZATION_MODES}.")
+        rr.spawn(memory_limit=memory_limit)
 
 
-def log_visualization_data(
-    display_mode: str,
+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 _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/action data to the backend selected by ``display_mode``."""
+    """
+    Logs observation and action data to Rerun for real-time visualization.
 
-    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}.")
+    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.
 
-def shutdown_visualization(display_mode: str) -> None:
-    """Shuts down the backend selected by ``display_mode``."""
+    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.
 
-    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}.")
+    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:
+                        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)
+                    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)

tests/conftest.py

@@ -14,23 +14,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import functools
 import traceback
 
-import draccus.wrappers.docstring as _draccus_docstring
 import pytest
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.utils.import_utils import is_package_available
 from tests.utils import DEVICE
 
-# On every `draccus.parse()`, draccus rebuilds each dataclass field's help text by
-# re-reading and re-parsing the class source (draccus.wrappers.docstring). For a config
-# as large as TrainPipelineConfig this costs ~2.5s per parse — negligible for the single
-# parse a CLI does, but tests parse configs hundreds of times. The source can't change
-# within a run, so memoize it for the whole test session.
-_draccus_docstring.get_attribute_docstring = functools.cache(_draccus_docstring.get_attribute_docstring)
-
 # Import fixture modules as plugins.
 # Fixtures that depend on optional packages are only registered when those packages are available,
 # so that tests can be collected and run even with a minimal install.

tests/datasets/test_datasets.py

@@ -1531,7 +1531,6 @@ def test_valid_video_codecs_constant():
     assert "h264" in VALID_VIDEO_CODECS
     assert "hevc" in VALID_VIDEO_CODECS
     assert "libsvtav1" in VALID_VIDEO_CODECS
-    assert "libaom-av1" in VALID_VIDEO_CODECS
     assert "auto" in VALID_VIDEO_CODECS
     assert "h264_videotoolbox" in VALID_VIDEO_CODECS
     assert "h264_nvenc" in VALID_VIDEO_CODECS
@@ -1539,7 +1538,7 @@ def test_valid_video_codecs_constant():
     assert "h264_qsv" in VALID_VIDEO_CODECS
     assert "hevc_videotoolbox" in VALID_VIDEO_CODECS
     assert "hevc_nvenc" in VALID_VIDEO_CODECS
-    assert len(VALID_VIDEO_CODECS) == 11
+    assert len(VALID_VIDEO_CODECS) == 10
 
 
 def test_delta_timestamps_with_episodes_filter(tmp_path, empty_lerobot_dataset_factory):

tests/datasets/test_depth.py

@@ -32,7 +32,6 @@ from lerobot.configs.video import (
 )
 from lerobot.datasets.depth_utils import dequantize_depth, quantize_depth
 from lerobot.datasets.image_writer import image_array_to_pil_image, write_image
-from lerobot.utils.constants import DEFAULT_FEATURES
 from tests.fixtures.constants import (
     DEFAULT_FPS,
     DUMMY_CAMERA_FEATURES,
@@ -246,91 +245,3 @@ class TestFeatureFileRouting:
 
         dataset.save_episode()
         dataset.finalize()
-
-
-class TestDepthUnitMetadata:
-    """The depth unit is inferred once from dtype, stored in ``info``, and drives stats + reads."""
-
-    NUM_FRAMES = 4
-
-    def _record(self, root, features_factory, depth_dtype, value, use_videos):
-        from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-        features = features_factory(camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, use_videos=use_videos)
-        dataset = LeRobotDataset.create(
-            repo_id=DUMMY_REPO_ID,
-            fps=DEFAULT_FPS,
-            features=features,
-            root=root,
-            use_videos=use_videos,
-            streaming_encoding=use_videos,
-        )
-        for _ in range(self.NUM_FRAMES):
-            frame: dict = {"task": "test"}
-            for key, ft in dataset.meta.features.items():
-                if key in DEFAULT_FEATURES:
-                    continue
-                if key in dataset.meta.depth_keys:
-                    frame[key] = np.full(ft["shape"], value, dtype=depth_dtype)
-                elif key in dataset.meta.camera_keys:
-                    frame[key] = np.random.randint(0, 256, ft["shape"], dtype=np.uint8)
-                else:
-                    frame[key] = np.zeros(ft["shape"], dtype=np.float32)
-            dataset.add_frame(frame)
-        return dataset
-
-    @pytest.mark.parametrize("use_videos", [False, True])
-    @pytest.mark.parametrize(
-        ("depth_dtype", "value", "expected_unit"),
-        [(np.float32, 2.0, DEPTH_METER_UNIT), (np.uint16, 2000, DEPTH_MILLIMETER_UNIT)],
-    )
-    def test_recorded_unit_inferred_persisted_and_kept_in_stats(
-        self, tmp_path, features_factory, use_videos, depth_dtype, value, expected_unit
-    ):
-        """Unit is inferred from the first frame's dtype, drives stats (raw, never canonicalized), and survives a reload."""
-        from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-        dataset = self._record(tmp_path / "ds", features_factory, depth_dtype, value, use_videos)
-        assert dataset.meta.features[DEPTH_KEY]["info"]["depth_unit"] == expected_unit
-        dataset.save_episode()
-        mean = float(np.asarray(dataset.meta.stats[DEPTH_KEY]["mean"]).reshape(-1)[0])
-        np.testing.assert_allclose(mean, value, rtol=0.05)
-        dataset.finalize()
-
-        reloaded = LeRobotDataset(repo_id=DUMMY_REPO_ID, root=tmp_path / "ds")
-        assert reloaded.meta.features[DEPTH_KEY]["info"]["depth_unit"] == expected_unit
-
-    @pytest.mark.parametrize("use_videos", [False, True])
-    @pytest.mark.parametrize(
-        ("output_unit", "expected"),
-        [(DEPTH_MILLIMETER_UNIT, 2000.0), (DEPTH_METER_UNIT, 2.0)],
-    )
-    def test_read_honors_output_unit_for_frames_and_stats(
-        self, tmp_path, features_factory, use_videos, output_unit, expected
-    ):
-        """Reloading with a ``depth_output_unit`` converts metre frames (image mode) and rescales stats while preserving count."""
-        from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-        dataset = self._record(tmp_path / "ds", features_factory, np.float32, 2.0, use_videos=use_videos)
-        dataset.save_episode()
-        count = float(np.asarray(dataset.meta.stats[DEPTH_KEY]["count"]).reshape(-1)[0])
-        dataset.finalize()
-
-        read_dataset = LeRobotDataset(
-            repo_id=DUMMY_REPO_ID, root=tmp_path / "ds", depth_output_unit=output_unit
-        )
-        stats = read_dataset.meta.stats[DEPTH_KEY]
-        np.testing.assert_allclose(float(np.asarray(stats["mean"]).reshape(-1)[0]), expected, rtol=0.05)
-        np.testing.assert_allclose(float(np.asarray(stats["count"]).reshape(-1)[0]), count)
-
-        if not use_videos:
-            depth = read_dataset[0][DEPTH_KEY]
-            assert torch.allclose(depth, torch.full_like(depth, expected))
-
-            from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
-
-            stream_dataset = StreamingLeRobotDataset(
-                repo_id=DUMMY_REPO_ID, root=tmp_path / "ds", depth_output_unit=output_unit
-            )
-            stream_depth = next(iter(stream_dataset))[DEPTH_KEY]
-            assert torch.allclose(stream_depth, torch.full_like(stream_depth, expected))

tests/datasets/test_video_encoding.py

@@ -345,9 +345,7 @@ class TestExtraOptions:
         opts = cfg.get_codec_options()
         assert opts["qp"] == 20
         assert isinstance(opts["qp"], int)
-        str_opts = cfg.get_codec_options(as_strings=True)
-        assert str_opts["qp"] == "20"
-        assert all(isinstance(v, str) for v in str_opts.values())
+        assert cfg.get_codec_options(as_strings=True)["qp"] == "20"
 
     @require_libsvtav1
     def test_structured_fields_win_on_collision(self):

tests/fixtures/dataset_factories.py

@@ -26,7 +26,6 @@ import pytest
 import torch
 from datasets import Dataset
 
-from lerobot.configs.video import infer_depth_unit
 from lerobot.datasets.dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
 from lerobot.datasets.feature_utils import get_hf_features_from_features
 from lerobot.datasets.io_utils import flatten_dict, hf_transform_to_torch
@@ -536,13 +535,6 @@ def lerobot_dataset_factory(
                 chunks_size=chunks_size,
                 **info_kwargs,
             )
-            # This synthetic path skips add_frame, so record the depth unit the writer would
-            # have stored (dummy depth is uint16) to keep ``depth_unit`` present in info.json.
-            # Reassign a fresh info dict to avoid mutating the shared feature constants.
-            for ft in info.features.values():
-                ft_info = ft.get("info")
-                if ft_info is not None and ft_info.get("is_depth_map") and "depth_unit" not in ft_info:
-                    ft["info"] = {**ft_info, "depth_unit": infer_depth_unit(np.uint16)}
         if stats is None:
             stats = stats_factory(features=info.features)
         if tasks is None:

tests/policies/lingbot_va/test_configuration_lingbot_va.py

@@ -1,78 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import pytest
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import FeatureType, PolicyFeature
-from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
-from lerobot.utils.constants import ACTION, OBS_IMAGES
-
-
-def make_config(**overrides) -> LingBotVAConfig:
-    kwargs = {"device": "cpu"}
-    kwargs.update(overrides)
-    return LingBotVAConfig(**kwargs)
-
-
-def test_registered_in_choice_registry() -> None:
-    assert "lingbot_va" in PreTrainedConfig.get_known_choices()
-    assert PreTrainedConfig.get_choice_class("lingbot_va") is LingBotVAConfig
-
-
-def test_type_property() -> None:
-    assert make_config().type == "lingbot_va"
-
-
-def test_chunk_size_and_action_steps() -> None:
-    cfg = make_config(frame_chunk_size=4, action_per_frame=4)
-    assert cfg.chunk_size == 16
-    assert cfg.n_action_steps == 16
-    assert cfg.action_delta_indices == list(range(16))
-    assert cfg.observation_delta_indices == list(range(16))
-    assert cfg.reward_delta_indices is None
-
-
-def test_optimizer_and_scheduler_presets() -> None:
-    cfg = make_config()
-    opt = cfg.get_optimizer_preset()
-    assert opt.lr == cfg.optimizer_lr
-    sched = cfg.get_scheduler_preset()
-    assert sched.num_warmup_steps == cfg.scheduler_warmup_steps
-
-
-def test_validate_features_sets_action_feature() -> None:
-    cfg = make_config()
-    cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))}
-    cfg.output_features = {}
-    cfg.validate_features()
-    assert ACTION in cfg.output_features
-    assert cfg.output_features[ACTION].shape == (len(cfg.used_action_channel_ids),)
-
-
-def test_validate_features_no_visual_raises() -> None:
-    cfg = make_config()
-    cfg.input_features = {}
-    cfg.output_features = {}
-    with pytest.raises(ValueError, match="at least one visual input feature"):
-        cfg.validate_features()
-
-
-def test_invalid_attn_mode_raises() -> None:
-    with pytest.raises(ValueError, match="attn_mode"):
-        make_config(attn_mode="banana")

tests/policies/lingbot_va/test_factory.py

@@ -1,38 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import pytest
-
-from lerobot.policies.factory import make_policy_config
-from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
-
-
-def test_make_policy_config_returns_lingbot_va() -> None:
-    cfg = make_policy_config("lingbot_va", device="cpu")
-    assert isinstance(cfg, LingBotVAConfig)
-
-
-def test_get_policy_class_resolves_lazily() -> None:
-    # Importing the policy class pulls in diffusers (Wan2.2 stack); skip if unavailable.
-    pytest.importorskip("diffusers")
-    pytest.importorskip("transformers")
-    from lerobot.policies.factory import get_policy_class
-
-    cls = get_policy_class("lingbot_va")
-    assert cls.name == "lingbot_va"
-    assert cls.config_class is LingBotVAConfig

tests/policies/lingbot_va/test_modules.py

@@ -1,128 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Unit tests for the vendored LingBot-VA helper code (scheduler + grid utilities)."""
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-pytest.importorskip("diffusers")  # the model code lives in modeling_lingbot_va, which imports diffusers
-
-from lerobot.policies.lingbot_va.modeling_lingbot_va import FlowMatchScheduler
-from lerobot.policies.lingbot_va.utils import data_seq_to_patch, get_mesh_id
-
-
-def test_flow_match_scheduler_timesteps_monotone_decreasing() -> None:
-    sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True)
-    sch.set_timesteps(20)
-    assert sch.timesteps.shape == (20,)
-    diffs = sch.timesteps[1:] - sch.timesteps[:-1]
-    assert torch.all(diffs <= 0)  # decreasing
-
-
-def test_flow_match_scheduler_step_preserves_shape() -> None:
-    sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True)
-    sch.set_timesteps(20)
-    sample = torch.zeros(1, 48, 4, 8, 16)
-    out = sch.step(torch.ones_like(sample), sch.timesteps[0], sample)
-    assert out.shape == sample.shape
-
-
-def test_flow_match_scheduler_add_noise() -> None:
-    sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True)
-    sch.set_timesteps(20)
-    sample = torch.randn(1, 48, 4, 8, 16)
-    noise = torch.randn_like(sample)
-    noisy = sch.add_noise(sample, noise, sch.timesteps[:4], t_dim=2)
-    assert noisy.shape == sample.shape
-
-
-def test_get_mesh_id_latent_shape() -> None:
-    grid = get_mesh_id(4, 8, 16, 0, 1, 0)
-    assert grid.shape == (4, 4 * 8 * 16)  # (f, h, w, stream) x tokens
-
-
-def test_get_mesh_id_action_shape() -> None:
-    grid = get_mesh_id(4, 4, 1, 1, 1, 0, action=True)
-    assert grid.shape == (4, 4 * 4 * 1)
-    # Action rows for h/w are sentinel -1.
-    assert torch.all(grid[1] < 0)
-    assert torch.all(grid[2] < 0)
-
-
-def test_data_seq_to_patch_roundtrip_shape() -> None:
-    b, f, h, w, c = 1, 4, 8, 16, 48
-    seq = torch.arange(b * f * h * w * c, dtype=torch.float32).reshape(b, f * h * w, c)
-    out = data_seq_to_patch((1, 2, 2), seq, f, h, w, batch_size=b)
-    assert out.shape == (b, c, f, h, w)
-
-
-def test_training_step_reduces_loss_tiny_flex() -> None:
-    """End-to-end single training step (flow-matching loss -> backward -> AdamW) on a tiny config.
-
-    Exercises the flex-attention training path; requires a CUDA GPU with flex-attention support.
-    """
-    if not torch.cuda.is_available():
-        import pytest
-
-        pytest.skip("training step test requires a CUDA GPU (flex-attention)")
-
-    from lerobot.configs.types import FeatureType, PolicyFeature
-    from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
-    from lerobot.policies.lingbot_va.modeling_lingbot_va import LingBotVAPolicy
-    from lerobot.utils.constants import ACTION, OBS_IMAGES
-
-    cfg = LingBotVAConfig(
-        attn_mode="flex",
-        dtype="bfloat16",
-        in_channels=16,
-        out_channels=16,
-        action_dim=8,
-        text_dim=32,
-        freq_dim=64,
-        ffn_dim=64,
-        num_attention_heads=2,
-        attention_head_dim=24,
-        num_layers=2,
-        frame_chunk_size=2,
-        action_per_frame=4,
-        used_action_channel_ids=[0, 1, 2, 3],
-        obs_cam_keys=[f"{OBS_IMAGES}.image"],
-        device="cuda",
-    )
-    cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 64, 64))}
-    cfg.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(4,))}
-    cfg.validate_features()
-
-    policy = LingBotVAPolicy(cfg).to("cuda")
-    policy.train()
-    opt = torch.optim.AdamW(policy.get_optim_params(), lr=1e-4)
-
-    b, fc, apf = 1, cfg.frame_chunk_size, cfg.action_per_frame
-    latents = torch.randn(b, cfg.in_channels, fc, 4, 4, device="cuda", dtype=torch.bfloat16)
-    actions = torch.randn(b, cfg.action_dim, fc, apf, 1, device="cuda", dtype=torch.bfloat16)
-    amask = torch.zeros(cfg.action_dim, device="cuda")
-    amask[cfg.used_action_channel_ids] = 1.0
-    actions_mask = amask.view(1, -1, 1, 1, 1).expand_as(actions)
-    text_emb = torch.randn(b, cfg.max_sequence_length, cfg.text_dim, device="cuda", dtype=torch.bfloat16)
-
-    loss, metrics = policy.training_loss_from_streams(latents, actions, actions_mask, text_emb)
-    assert torch.isfinite(loss) and {"latent_loss", "action_loss"} <= set(metrics)
-    loss.backward()
-    assert any(p.grad is not None and torch.isfinite(p.grad).all() for p in policy.get_optim_params())
-    opt.step()

tests/policies/lingbot_va/test_processor.py

@@ -1,88 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import torch
-
-from lerobot.configs.types import FeatureType, PolicyFeature
-from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
-from lerobot.policies.lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
-from lerobot.processor import PolicyProcessorPipeline, UnnormalizerProcessorStep
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import (
-    ACTION,
-    OBS_IMAGES,
-    POLICY_POSTPROCESSOR_DEFAULT_NAME,
-    POLICY_PREPROCESSOR_DEFAULT_NAME,
-)
-
-
-def _make_config() -> LingBotVAConfig:
-    cfg = LingBotVAConfig(device="cpu")
-    cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))}
-    cfg.output_features = {}
-    cfg.validate_features()
-    return cfg
-
-
-def test_make_pre_post_processors_names_and_steps() -> None:
-    cfg = _make_config()
-    pre, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=None)
-    assert pre.name == POLICY_PREPROCESSOR_DEFAULT_NAME
-    assert post.name == POLICY_POSTPROCESSOR_DEFAULT_NAME
-    # Actions are unnormalized by the standard built-in quantile unnormalizer.
-    assert any(isinstance(s, UnnormalizerProcessorStep) for s in post.steps)
-
-
-def test_freshly_built_postprocessor_is_identity() -> None:
-    # Without action stats the quantile unnormalizer is a no-op (identity passthrough): the real
-    # per-benchmark q01/q99 are restored from the saved checkpoint on load, not hardcoded here.
-    cfg = _make_config()
-    _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=None)
-    normed = torch.tensor([[0.3, -0.5, 1.0, -1.0, 0.0, 0.7, -0.2]])
-    assert torch.allclose(post(normed), normed, atol=1e-6)
-
-
-def test_postprocessor_quantile_unnormalization() -> None:
-    # QUANTILES unnormalize maps [-1, 1] -> [q01, q99]: -1 -> q01, +1 -> q99.
-    cfg = _make_config()
-    q01 = [-1.0, -0.5, 0.0, -1.0, -1.0, -1.0, -1.0]
-    q99 = [1.0, 0.5, 2.0, 1.0, 1.0, 1.0, 1.0]
-    stats = {ACTION: {"q01": q01, "q99": q99}}
-    _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=stats)
-    out_lo = post(torch.full((1, 7), -1.0))
-    out_hi = post(torch.full((1, 7), 1.0))
-    assert torch.allclose(out_lo, torch.tensor(q01).unsqueeze(0), atol=1e-4)
-    assert torch.allclose(out_hi, torch.tensor(q99).unsqueeze(0), atol=1e-4)
-
-
-def test_postprocessor_stats_survive_save_load(tmp_path) -> None:
-    # Regression guard for the Hub mechanism: the q01/q99 stats live in the saved post-processor
-    # state and must round-trip through save_pretrained / from_pretrained.
-    cfg = _make_config()
-    q01 = [-0.6, -0.8, -0.9, -0.1, -0.15, -0.25, -1.0]
-    q99 = [0.9, 0.85, 0.9, 0.17, 0.18, 0.34, 1.0]
-    _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats={ACTION: {"q01": q01, "q99": q99}})
-    post.save_pretrained(tmp_path)
-    loaded = PolicyProcessorPipeline.from_pretrained(
-        tmp_path,
-        config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
-        to_transition=policy_action_to_transition,
-        to_output=transition_to_policy_action,
-    )
-    out = loaded(torch.full((1, 7), -1.0))
-    assert torch.allclose(out, torch.tensor(q01).unsqueeze(0), atol=1e-4)

tests/utils/test_foxglove_visualization.py

@@ -1,101 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for the Foxglove backend's pure helpers.
-
-These cover topic naming, series labelling and feature-name parsing. They import
-``foxglove_visualization`` directly and need NO ``foxglove`` extra: the SDK is imported lazily inside
-the functions that talk to the server, so the helpers below run in the base test tier.
-"""
-
-import numpy as np
-
-from lerobot.utils import foxglove_visualization as fv
-from lerobot.utils.constants import ACTION, OBS_STATE
-
-
-def test_foxglove_safe_name_collapses_dots():
-    assert fv._foxglove_safe_name("observation.images.front") == "observation_images_front"
-    assert fv._foxglove_safe_name("plain") == "plain"
-
-
-def test_foxglove_topic_image_strips_prefix_without_doubling_images():
-    # Fully-qualified camera key -> single clean segment (no doubled "images").
-    assert fv._foxglove_topic("observation.images.front", is_image=True) == "/observation/images/front"
-    # A nested camera name keeps its structure via safe-name collapsing.
-    assert (
-        fv._foxglove_topic("observation.images.wrist.left", is_image=True) == "/observation/images/wrist_left"
-    )
-    # Bare camera name (as real robots emit).
-    assert fv._foxglove_topic("front", is_image=True) == "/observation/images/front"
-
-
-def test_foxglove_topic_scalar_sources():
-    assert fv._foxglove_topic(OBS_STATE) == "/observation/state"
-    assert fv._foxglove_topic("observation.environment_state") == "/observation/state"
-    assert fv._foxglove_topic(ACTION) == "/action/state"
-    assert fv._foxglove_topic("action.delta") == "/action/state"
-
-
-def test_labeled_scalars_uses_labels_then_index_fallback():
-    assert fv._labeled_scalars("state", np.array([1.0, 2.0, 3.0])) == {
-        "state_0": 1.0,
-        "state_1": 2.0,
-        "state_2": 3.0,
-    }
-    assert fv._labeled_scalars("state", [1.0, 2.0], ["pan", "lift"]) == {"pan": 1.0, "lift": 2.0}
-    # Wrong-length labels fall back to index naming (never silently mislabels).
-    assert fv._labeled_scalars("q", [1.0, 2.0], ["only_one"]) == {"q_0": 1.0, "q_1": 2.0}
-
-
-def test_frame_to_scalars_matches_live_labeling_and_handles_scalar():
-    frame = {OBS_STATE: np.array([1.0, 2.0])}
-    # No metadata -> {short_name}_{i}, identical to the live-stream fallback.
-    assert fv._frame_to_scalars(frame, OBS_STATE) == fv._labeled_scalars("state", np.array([1.0, 2.0]))
-    assert fv._frame_to_scalars(frame, OBS_STATE) == {"state_0": 1.0, "state_1": 2.0}
-    # Metadata labels are honored.
-    assert fv._frame_to_scalars(frame, OBS_STATE, ["pan", "lift"]) == {"pan": 1.0, "lift": 2.0}
-    # A 0-d scalar becomes a single entry named by the short feature name.
-    assert fv._frame_to_scalars({ACTION: np.array(5.0)}, ACTION) == {"action": 5.0}
-    # A missing feature yields an empty mapping.
-    assert fv._frame_to_scalars({}, OBS_STATE) == {}
-
-
-def test_feature_dim_names_formats():
-    # Flat list of names.
-    assert fv._feature_dim_names({"shape": [2], "names": ["x", "y"]}) == ["x", "y"]
-    # Category mapping (dict of lists).
-    assert fv._feature_dim_names({"shape": [2], "names": {"motors": ["m0", "m1"]}}) == ["m0", "m1"]
-    # name -> index mapping (returned sorted by index).
-    assert fv._feature_dim_names({"shape": [2], "names": {"delta_x": 0, "delta_y": 1}}) == [
-        "delta_x",
-        "delta_y",
-    ]
-    # Bool values must NOT be treated as an index map (bool is a subclass of int).
-    assert fv._feature_dim_names({"shape": [2], "names": {"a": True, "b": False}}) is None
-    # Mismatched length -> None (won't silently mislabel).
-    assert fv._feature_dim_names({"shape": [3], "names": ["x", "y"]}) is None
-    # Missing / absent names -> None.
-    assert fv._feature_dim_names(None) is None
-    assert fv._feature_dim_names({"shape": [2]}) is None
-
-
-def test_is_scalar():
-    assert fv._is_scalar(1.0)
-    assert fv._is_scalar(np.float32(2.0))
-    assert fv._is_scalar(np.array(3.0))  # 0-d array
-    assert not fv._is_scalar(np.array([1.0, 2.0]))
-    assert not fv._is_scalar("x")

tests/utils/test_rerun_visualization.py

@@ -1,311 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import importlib
-import sys
-from types import SimpleNamespace
-
-import numpy as np
-import pytest
-
-pytest.importorskip("rerun", reason="rerun-sdk is required (install lerobot[viz])")
-
-from lerobot.types import TransitionKey
-from lerobot.utils.constants import OBS_STATE
-
-
-@pytest.fixture
-def mock_rerun(monkeypatch):
-    """
-    Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't
-    depend on the real library. Also reload the module-under-test so it binds to
-    this mock `rr`.
-    """
-    calls = []
-    blueprints = []
-
-    class DummyScalar:
-        def __init__(self, value):
-            # Scalars may be built from a single float or from a 1D array batch.
-            self.value = value
-
-    class DummyImage:
-        def __init__(self, arr):
-            self.arr = arr
-
-        def compress(self, *a, **k):
-            return self
-
-    class DummyDepthImage:
-        def __init__(self, arr, meter=None, colormap=None):
-            self.arr = arr
-            self.meter = meter
-            self.colormap = colormap
-
-    def dummy_log(key, obj=None, **kwargs):
-        # Accept either positional `obj` or keyword `entity` and record remaining kwargs.
-        if obj is None and "entity" in kwargs:
-            obj = kwargs.pop("entity")
-        calls.append((key, obj, kwargs))
-
-    def dummy_send_blueprint(blueprint, *a, **k):
-        blueprints.append(blueprint)
-
-    # Mock the `rerun.blueprint` submodule used to build the layout.
-    dummy_rrb = SimpleNamespace(
-        Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
-            kind="Spatial2DView", origin=origin, name=name
-        ),
-        TimeSeriesView=lambda name=None, contents=None: SimpleNamespace(
-            kind="TimeSeriesView", name=name, contents=contents
-        ),
-        Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)),
-        Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root),
-    )
-
-    dummy_rr = SimpleNamespace(
-        __name__="rerun",
-        __package__="rerun",
-        __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None),
-        Scalars=DummyScalar,
-        Image=DummyImage,
-        DepthImage=DummyDepthImage,
-        components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")),
-        log=dummy_log,
-        send_blueprint=dummy_send_blueprint,
-        init=lambda *a, **k: None,
-        spawn=lambda *a, **k: None,
-        blueprint=dummy_rrb,
-    )
-
-    # Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`).
-    monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
-    monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
-
-    # Now import and reload the module under test, to bind to our rerun mock
-    import lerobot.utils.rerun_visualization as rv
-
-    importlib.reload(rv)
-
-    # Expose the reloaded module, the call recorder and the captured blueprints
-    yield rv, calls, blueprints
-
-
-def _keys(calls):
-    """Helper to extract just the keys logged to rr.log"""
-    return [k for (k, _obj, _kw) in calls]
-
-
-def _obj_for(calls, key):
-    """Find the first object logged under a given key."""
-    for k, obj, _kw in calls:
-        if k == key:
-            return obj
-    raise KeyError(f"Key {key} not found in calls: {calls}")
-
-
-def _kwargs_for(calls, key):
-    for k, _obj, kw in calls:
-        if k == key:
-            return kw
-    raise KeyError(f"Key {key} not found in calls: {calls}")
-
-
-def _views_by_kind(blueprint, kind):
-    """Return the views of a given kind from the (single) blueprint's grid."""
-    return [v for v in blueprint.root.views if v.kind == kind]
-
-
-def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
-    rv, calls, blueprints = mock_rerun
-
-    # Build EnvTransition dict
-    obs = {
-        f"{OBS_STATE}.temperature": np.float32(25.0),
-        # CHW image should be converted to HWC for rr.Image
-        "observation.camera": np.zeros((3, 10, 20), dtype=np.uint8),
-    }
-    act = {
-        "action.throttle": 0.7,
-        # 1D array should be logged as a single Scalars batch under one entity path
-        "action.vector": np.array([1.0, 2.0], dtype=np.float32),
-    }
-    transition = {
-        TransitionKey.OBSERVATION: obs,
-        TransitionKey.ACTION: act,
-    }
-
-    # Extract observation and action data from transition like in the real call sites
-    obs_data = transition.get(TransitionKey.OBSERVATION, {})
-    action_data = transition.get(TransitionKey.ACTION, {})
-    rv.log_rerun_data(observation=obs_data, action=action_data)
-
-    # We expect:
-    # - observation.state.temperature -> Scalars
-    # - observation.camera -> Image (HWC) with static=True
-    # - action.throttle -> Scalars
-    # - action.vector -> single Scalars batch (no per-element suffix)
-    expected_keys = {
-        f"{OBS_STATE}.temperature",
-        "observation.camera",
-        "action.throttle",
-        "action.vector",
-    }
-    assert set(_keys(calls)) == expected_keys
-
-    # Check scalar types and values
-    temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature")
-    assert type(temp_obj).__name__ == "DummyScalar"
-    assert float(temp_obj.value) == pytest.approx(25.0)
-
-    throttle_obj = _obj_for(calls, "action.throttle")
-    assert type(throttle_obj).__name__ == "DummyScalar"
-    assert float(throttle_obj.value) == pytest.approx(0.7)
-
-    # 1D vector logged as a single batched Scalars under one entity path
-    vec = _obj_for(calls, "action.vector")
-    assert type(vec).__name__ == "DummyScalar"
-    np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0])
-
-    # Check image handling: CHW -> HWC
-    img_obj = _obj_for(calls, "observation.camera")
-    assert type(img_obj).__name__ == "DummyImage"
-    assert img_obj.arr.shape == (10, 20, 3)  # transposed
-    assert _kwargs_for(calls, "observation.camera").get("static", False) is True  # static=True for images
-
-    # A blueprint should have been built and sent exactly once, and cached on the function.
-    assert len(blueprints) == 1
-    assert rv.log_rerun_data.blueprint is blueprints[0]
-
-    bp = blueprints[0]
-    # One spatial view per image path
-    spatial_views = _views_by_kind(bp, "Spatial2DView")
-    assert {v.origin for v in spatial_views} == {"observation.camera"}
-
-    # One time-series view each for observation and action scalars
-    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
-    assert set(ts_views) == {"observation", "action"}
-    assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"]
-    assert ts_views["action"].contents == ["action.throttle", "action.vector"]
-
-
-def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
-    rv, calls, blueprints = mock_rerun
-
-    # First dict without prefixes treated as observation
-    # Second dict without prefixes treated as action
-    obs_plain = {
-        "temp": 1.5,
-        # Already HWC image => should stay as-is
-        "img": np.zeros((5, 6, 3), dtype=np.uint8),
-        "none": None,  # should be skipped
-    }
-    act_plain = {
-        "throttle": 0.3,
-        "vec": np.array([9, 8, 7], dtype=np.float32),
-    }
-
-    # Extract observation and action data from list like the old function logic did
-    # First dict was treated as observation, second as action
-    rv.log_rerun_data(observation=obs_plain, action=act_plain)
-
-    # Expected keys with auto-prefixes. The 1D vector is a single batched Scalars.
-    expected = {
-        "observation.temp",
-        "observation.img",
-        "action.throttle",
-        "action.vec",
-    }
-    logged = set(_keys(calls))
-    assert logged == expected
-
-    # Scalars
-    t = _obj_for(calls, "observation.temp")
-    assert type(t).__name__ == "DummyScalar"
-    assert float(t.value) == pytest.approx(1.5)
-
-    throttle = _obj_for(calls, "action.throttle")
-    assert type(throttle).__name__ == "DummyScalar"
-    assert float(throttle.value) == pytest.approx(0.3)
-
-    # Image stays HWC
-    img = _obj_for(calls, "observation.img")
-    assert type(img).__name__ == "DummyImage"
-    assert img.arr.shape == (5, 6, 3)
-    assert _kwargs_for(calls, "observation.img").get("static", False) is True
-
-    # Vector logged as a single batched Scalars under one entity path
-    vec = _obj_for(calls, "action.vec")
-    assert type(vec).__name__ == "DummyScalar"
-    np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7])
-
-    # Blueprint sent once with the expected view layout
-    assert len(blueprints) == 1
-    bp = blueprints[0]
-    spatial_views = _views_by_kind(bp, "Spatial2DView")
-    assert {v.origin for v in spatial_views} == {"observation.img"}
-    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
-    assert ts_views["observation"].contents == ["observation.temp"]
-    assert ts_views["action"].contents == ["action.throttle", "action.vec"]
-
-
-def test_log_rerun_data_kwargs_only(mock_rerun):
-    rv, calls, blueprints = mock_rerun
-
-    rv.log_rerun_data(
-        observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)},
-        action={"action.a": 1.0},
-    )
-
-    keys = set(_keys(calls))
-    assert "observation.temp" in keys
-    assert "observation.gray" in keys
-    assert "action.a" in keys
-
-    temp = _obj_for(calls, "observation.temp")
-    assert type(temp).__name__ == "DummyScalar"
-    assert float(temp.value) == pytest.approx(10.0)
-
-    img = _obj_for(calls, "observation.gray")
-    assert type(img).__name__ == "DummyDepthImage"  # single-channel -> DepthImage
-    assert img.arr.shape == (8, 8, 1)  # remains HWC
-    assert _kwargs_for(calls, "observation.gray").get("static", False) is True
-
-    a = _obj_for(calls, "action.a")
-    assert type(a).__name__ == "DummyScalar"
-    assert float(a.value) == pytest.approx(1.0)
-
-    # Blueprint sent once, with a spatial view for the image and time-series views for scalars
-    assert len(blueprints) == 1
-    bp = blueprints[0]
-    assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"}
-    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
-    assert ts_views["observation"].contents == ["observation.temp"]
-    assert ts_views["action"].contents == ["action.a"]
-
-
-def test_log_rerun_data_blueprint_sent_only_once(mock_rerun):
-    """The blueprint is built from the first call and not resent on subsequent calls."""
-    rv, calls, blueprints = mock_rerun
-
-    rv.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0})
-    assert len(blueprints) == 1
-    first_blueprint = rv.log_rerun_data.blueprint
-
-    rv.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0})
-    # Still only one blueprint, and the cached one is unchanged.
-    assert len(blueprints) == 1
-    assert rv.log_rerun_data.blueprint is first_blueprint

tests/utils/test_visualization_utils.py

@@ -14,23 +14,297 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for the backend-agnostic visualization dispatch.
-
-These exercise the display-mode routing/validation only; they need neither ``rerun`` nor
-``foxglove`` installed since the unknown-mode branch raises before touching any backend. Backend
-behavior is covered in ``test_rerun_visualization.py`` and ``test_foxglove_visualization.py``.
-"""
+import importlib
+import sys
+from types import SimpleNamespace
 
+import numpy as np
 import pytest
 
-from lerobot.utils import visualization_utils as vu
+pytest.importorskip("rerun", reason="rerun-sdk is required (install lerobot[viz])")
+
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import OBS_STATE
 
 
-def test_visualization_modes():
-    assert vu.VISUALIZATION_MODES == ("rerun", "foxglove")
+@pytest.fixture
+def mock_rerun(monkeypatch):
+    """
+    Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't
+    depend on the real library. Also reload the module-under-test so it binds to
+    this mock `rr`.
+    """
+    calls = []
+    blueprints = []
+
+    class DummyScalar:
+        def __init__(self, value):
+            # Scalars may be built from a single float or from a 1D array batch.
+            self.value = value
+
+    class DummyImage:
+        def __init__(self, arr):
+            self.arr = arr
+
+        def compress(self, *a, **k):
+            return self
+
+    class DummyDepthImage:
+        def __init__(self, arr, colormap=None):
+            self.arr = arr
+            self.colormap = colormap
+
+    def dummy_log(key, obj=None, **kwargs):
+        # Accept either positional `obj` or keyword `entity` and record remaining kwargs.
+        if obj is None and "entity" in kwargs:
+            obj = kwargs.pop("entity")
+        calls.append((key, obj, kwargs))
+
+    def dummy_send_blueprint(blueprint, *a, **k):
+        blueprints.append(blueprint)
+
+    # Mock the `rerun.blueprint` submodule used to build the layout.
+    dummy_rrb = SimpleNamespace(
+        Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
+            kind="Spatial2DView", origin=origin, name=name
+        ),
+        TimeSeriesView=lambda name=None, contents=None: SimpleNamespace(
+            kind="TimeSeriesView", name=name, contents=contents
+        ),
+        Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)),
+        Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root),
+    )
+
+    dummy_rr = SimpleNamespace(
+        __name__="rerun",
+        __package__="rerun",
+        __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None),
+        Scalars=DummyScalar,
+        Image=DummyImage,
+        DepthImage=DummyDepthImage,
+        components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")),
+        log=dummy_log,
+        send_blueprint=dummy_send_blueprint,
+        init=lambda *a, **k: None,
+        spawn=lambda *a, **k: None,
+        blueprint=dummy_rrb,
+    )
+
+    # Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`).
+    monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
+    monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
+
+    # Now import and reload the module under test, to bind to our rerun mock
+    import lerobot.utils.visualization_utils as vu
+
+    importlib.reload(vu)
+
+    # Expose the reloaded module, the call recorder and the captured blueprints
+    yield vu, calls, blueprints
 
 
-@pytest.mark.parametrize("func", ["init_visualization", "log_visualization_data", "shutdown_visualization"])
-def test_dispatch_rejects_unknown_mode(func):
-    with pytest.raises(ValueError, match="Unknown display_mode"):
-        getattr(vu, func)("bogus")
+def _keys(calls):
+    """Helper to extract just the keys logged to rr.log"""
+    return [k for (k, _obj, _kw) in calls]
+
+
+def _obj_for(calls, key):
+    """Find the first object logged under a given key."""
+    for k, obj, _kw in calls:
+        if k == key:
+            return obj
+    raise KeyError(f"Key {key} not found in calls: {calls}")
+
+
+def _kwargs_for(calls, key):
+    for k, _obj, kw in calls:
+        if k == key:
+            return kw
+    raise KeyError(f"Key {key} not found in calls: {calls}")
+
+
+def _views_by_kind(blueprint, kind):
+    """Return the views of a given kind from the (single) blueprint's grid."""
+    return [v for v in blueprint.root.views if v.kind == kind]
+
+
+def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
+    vu, calls, blueprints = mock_rerun
+
+    # Build EnvTransition dict
+    obs = {
+        f"{OBS_STATE}.temperature": np.float32(25.0),
+        # CHW image should be converted to HWC for rr.Image
+        "observation.camera": np.zeros((3, 10, 20), dtype=np.uint8),
+    }
+    act = {
+        "action.throttle": 0.7,
+        # 1D array should be logged as a single Scalars batch under one entity path
+        "action.vector": np.array([1.0, 2.0], dtype=np.float32),
+    }
+    transition = {
+        TransitionKey.OBSERVATION: obs,
+        TransitionKey.ACTION: act,
+    }
+
+    # Extract observation and action data from transition like in the real call sites
+    obs_data = transition.get(TransitionKey.OBSERVATION, {})
+    action_data = transition.get(TransitionKey.ACTION, {})
+    vu.log_rerun_data(observation=obs_data, action=action_data)
+
+    # We expect:
+    # - observation.state.temperature -> Scalars
+    # - observation.camera -> Image (HWC) with static=True
+    # - action.throttle -> Scalars
+    # - action.vector -> single Scalars batch (no per-element suffix)
+    expected_keys = {
+        f"{OBS_STATE}.temperature",
+        "observation.camera",
+        "action.throttle",
+        "action.vector",
+    }
+    assert set(_keys(calls)) == expected_keys
+
+    # Check scalar types and values
+    temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature")
+    assert type(temp_obj).__name__ == "DummyScalar"
+    assert float(temp_obj.value) == pytest.approx(25.0)
+
+    throttle_obj = _obj_for(calls, "action.throttle")
+    assert type(throttle_obj).__name__ == "DummyScalar"
+    assert float(throttle_obj.value) == pytest.approx(0.7)
+
+    # 1D vector logged as a single batched Scalars under one entity path
+    vec = _obj_for(calls, "action.vector")
+    assert type(vec).__name__ == "DummyScalar"
+    np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0])
+
+    # Check image handling: CHW -> HWC
+    img_obj = _obj_for(calls, "observation.camera")
+    assert type(img_obj).__name__ == "DummyImage"
+    assert img_obj.arr.shape == (10, 20, 3)  # transposed
+    assert _kwargs_for(calls, "observation.camera").get("static", False) is True  # static=True for images
+
+    # A blueprint should have been built and sent exactly once, and cached on the function.
+    assert len(blueprints) == 1
+    assert vu.log_rerun_data.blueprint is blueprints[0]
+
+    bp = blueprints[0]
+    # One spatial view per image path
+    spatial_views = _views_by_kind(bp, "Spatial2DView")
+    assert {v.origin for v in spatial_views} == {"observation.camera"}
+
+    # One time-series view each for observation and action scalars
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert set(ts_views) == {"observation", "action"}
+    assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"]
+    assert ts_views["action"].contents == ["action.throttle", "action.vector"]
+
+
+def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
+    vu, calls, blueprints = mock_rerun
+
+    # First dict without prefixes treated as observation
+    # Second dict without prefixes treated as action
+    obs_plain = {
+        "temp": 1.5,
+        # Already HWC image => should stay as-is
+        "img": np.zeros((5, 6, 3), dtype=np.uint8),
+        "none": None,  # should be skipped
+    }
+    act_plain = {
+        "throttle": 0.3,
+        "vec": np.array([9, 8, 7], dtype=np.float32),
+    }
+
+    # Extract observation and action data from list like the old function logic did
+    # First dict was treated as observation, second as action
+    vu.log_rerun_data(observation=obs_plain, action=act_plain)
+
+    # Expected keys with auto-prefixes. The 1D vector is a single batched Scalars.
+    expected = {
+        "observation.temp",
+        "observation.img",
+        "action.throttle",
+        "action.vec",
+    }
+    logged = set(_keys(calls))
+    assert logged == expected
+
+    # Scalars
+    t = _obj_for(calls, "observation.temp")
+    assert type(t).__name__ == "DummyScalar"
+    assert float(t.value) == pytest.approx(1.5)
+
+    throttle = _obj_for(calls, "action.throttle")
+    assert type(throttle).__name__ == "DummyScalar"
+    assert float(throttle.value) == pytest.approx(0.3)
+
+    # Image stays HWC
+    img = _obj_for(calls, "observation.img")
+    assert type(img).__name__ == "DummyImage"
+    assert img.arr.shape == (5, 6, 3)
+    assert _kwargs_for(calls, "observation.img").get("static", False) is True
+
+    # Vector logged as a single batched Scalars under one entity path
+    vec = _obj_for(calls, "action.vec")
+    assert type(vec).__name__ == "DummyScalar"
+    np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7])
+
+    # Blueprint sent once with the expected view layout
+    assert len(blueprints) == 1
+    bp = blueprints[0]
+    spatial_views = _views_by_kind(bp, "Spatial2DView")
+    assert {v.origin for v in spatial_views} == {"observation.img"}
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert ts_views["observation"].contents == ["observation.temp"]
+    assert ts_views["action"].contents == ["action.throttle", "action.vec"]
+
+
+def test_log_rerun_data_kwargs_only(mock_rerun):
+    vu, calls, blueprints = mock_rerun
+
+    vu.log_rerun_data(
+        observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)},
+        action={"action.a": 1.0},
+    )
+
+    keys = set(_keys(calls))
+    assert "observation.temp" in keys
+    assert "observation.gray" in keys
+    assert "action.a" in keys
+
+    temp = _obj_for(calls, "observation.temp")
+    assert type(temp).__name__ == "DummyScalar"
+    assert float(temp.value) == pytest.approx(10.0)
+
+    img = _obj_for(calls, "observation.gray")
+    assert type(img).__name__ == "DummyDepthImage"  # single-channel -> DepthImage
+    assert img.arr.shape == (8, 8, 1)  # remains HWC
+    assert _kwargs_for(calls, "observation.gray").get("static", False) is True
+
+    a = _obj_for(calls, "action.a")
+    assert type(a).__name__ == "DummyScalar"
+    assert float(a.value) == pytest.approx(1.0)
+
+    # Blueprint sent once, with a spatial view for the image and time-series views for scalars
+    assert len(blueprints) == 1
+    bp = blueprints[0]
+    assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"}
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert ts_views["observation"].contents == ["observation.temp"]
+    assert ts_views["action"].contents == ["action.a"]
+
+
+def test_log_rerun_data_blueprint_sent_only_once(mock_rerun):
+    """The blueprint is built from the first call and not resent on subsequent calls."""
+    vu, calls, blueprints = mock_rerun
+
+    vu.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0})
+    assert len(blueprints) == 1
+    first_blueprint = vu.log_rerun_data.blueprint
+
+    vu.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0})
+    # Still only one blueprint, and the cached one is unchanged.
+    assert len(blueprints) == 1
+    assert vu.log_rerun_data.blueprint is first_blueprint

uv.lock

@@ -1550,26 +1550,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/2c/47/c99d5268f354002ce80f8d029cd9d7d872969da1de8b93d32de4dc56d6f4/fonttools-4.63.0-py3-none-any.whl", hash = "sha256:445af2eab030a16b9171ea8bdda7ebf7d96bda2df88ee182a464252f6e05e20d", size = 1164562, upload-time = "2026-05-14T12:04:29.092Z" },
 ]
 
-[[package]]
-name = "foxglove-sdk"
-version = "0.25.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/c1/a7/86a252782ea0d9baf1357369ad1bbf1ed644768702b0266a3fa3a05361d0/foxglove_sdk-0.25.1.tar.gz", hash = "sha256:8230f3c32ea3ab715818687377491594ec9c7e58e6b0ed8ed91aadf937ce706b", size = 547778, upload-time = "2026-06-02T03:13:18.942Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/58/15/59f02e8201b8da09ce05d8774820c29efc9149862b70ee6b3a27968e791a/foxglove_sdk-0.25.1-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:5af9f9a691eefbe6e0a47875ff2f7d0fc36607f0920e8690bbdc2dfd4fb22451", size = 17911538, upload-time = "2026-06-02T03:13:12.493Z" },
-    { url = "https://files.pythonhosted.org/packages/27/ed/16d809fab24cbfdf97c15c9cdd80eabfeb447ca545ede426950d62bac848/foxglove_sdk-0.25.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:3e908bd87d1926a05c785779d8252db6b87eef685f284ec1cf46ee501645d08e", size = 16452309, upload-time = "2026-06-02T03:13:10.607Z" },
-    { url = "https://files.pythonhosted.org/packages/d6/c3/f95874935a3436841487df1f0202de4d20eabc0adb6b79c94c531bbe7eb3/foxglove_sdk-0.25.1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:968e32c8668d172f6b546c8e7af658ed35a21ec165adc3bacf53a04dda159f12", size = 2355680, upload-time = "2026-06-02T02:34:01.668Z" },
-    { url = "https://files.pythonhosted.org/packages/38/da/ad22d8d6e3fedde9fc0c49aa8b20394e5e0bc44ab3fba564c77a64ddc7e2/foxglove_sdk-0.25.1-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3f75374fedafe259c40b19bc645589d9453708eab679a5b07c603035f936d29a", size = 2274075, upload-time = "2026-06-02T02:34:07.212Z" },
-    { url = "https://files.pythonhosted.org/packages/a3/fa/1254adb5e72eff507695473e9c82d0e90395b61463e5353762250db30d3d/foxglove_sdk-0.25.1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7b6d3af517a00342bf7b08a4a65b043f3eafaa197138752b6fbd704fb91043fa", size = 2282160, upload-time = "2026-06-02T02:34:08.812Z" },
-    { url = "https://files.pythonhosted.org/packages/7c/e4/2b22ef06ba4058494c7aa35974d138f8f1ae4cf5273f77d69c9dc3a99b45/foxglove_sdk-0.25.1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:aed27c0f03a45fd6abdd566498bfee2672391602bcff32c827b8e3a6d8f67ab1", size = 22685338, upload-time = "2026-06-02T02:34:04.688Z" },
-    { url = "https://files.pythonhosted.org/packages/35/7c/58324c99b80eef0b674c8d4f5c2e07c66fd1480a27a8f0d4d79371805111/foxglove_sdk-0.25.1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:419dd8308e3f91e2ae487b727f1bf1804642990876163b2a353db4a1b1de1425", size = 19326096, upload-time = "2026-06-02T02:34:10.939Z" },
-    { url = "https://files.pythonhosted.org/packages/fe/9c/3452d92959e05fc6b1c1e5f032605d55623aeb6704357d20408f8781bc84/foxglove_sdk-0.25.1-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0fcb36e628ab3d9043e193f12ad4dbbb955fe18616aac7ef5bca82c52910f108", size = 2539020, upload-time = "2026-06-02T03:13:14.365Z" },
-    { url = "https://files.pythonhosted.org/packages/b5/af/57fa58525d3acb5c5480a6f0ef86450b1a0ccae2b21248edb1376073ce55/foxglove_sdk-0.25.1-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:7909fd9f94935935dd8813702d84ffdbfebeb3866673c618ce35e8cfedd03029", size = 2550999, upload-time = "2026-06-02T03:13:15.715Z" },
-    { url = "https://files.pythonhosted.org/packages/90/78/f74bb167186c965d475ff360fa6eb7441d5ac6c6239d60f542f63984f849/foxglove_sdk-0.25.1-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:69d5966213b5212b8841b4004fe582db924a74f1610d8452ad890f6931702926", size = 2560166, upload-time = "2026-06-02T03:13:17.254Z" },
-    { url = "https://files.pythonhosted.org/packages/81/83/1c4c6d04fbd4784fe44fb2da021db1adf1f03a371f1e5679a383c1173235/foxglove_sdk-0.25.1-cp310-abi3-win32.whl", hash = "sha256:2a1121a5c74590ff6e61628c4a46dc57d392d290b4beeb29d6852933da56224a", size = 1618124, upload-time = "2026-06-02T03:13:20.158Z" },
-    { url = "https://files.pythonhosted.org/packages/5f/4d/bdb9e252a41a951eb53908ac9cb965b7480c3ba649174f5398d4fcf0ca1d/foxglove_sdk-0.25.1-cp310-abi3-win_amd64.whl", hash = "sha256:6ed3ad0d3e72cd7875e7e293709c5ff90494fe14f1b48a336baffc313a7272cc", size = 16588452, upload-time = "2026-06-02T03:13:21.636Z" },
-]
-
 [[package]]
 name = "fqdn"
 version = "1.5.1"
@@ -2831,7 +2811,6 @@ all = [
     { name = "faker" },
     { name = "fastapi" },
     { name = "feetech-servo-sdk" },
-    { name = "foxglove-sdk" },
     { name = "grpcio" },
     { name = "grpcio-tools" },
     { name = "gym-aloha" },
@@ -2916,7 +2895,6 @@ core-scripts = [
     { name = "av" },
     { name = "datasets" },
     { name = "deepdiff" },
-    { name = "foxglove-sdk" },
     { name = "jsonlines" },
     { name = "pandas" },
     { name = "pyarrow" },
@@ -2939,7 +2917,6 @@ dataset = [
 dataset-viz = [
     { name = "av" },
     { name = "datasets" },
-    { name = "foxglove-sdk" },
     { name = "jsonlines" },
     { name = "pandas" },
     { name = "pyarrow" },
@@ -3057,11 +3034,6 @@ libero = [
     { name = "torchcodec", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'arm64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" },
     { name = "transformers" },
 ]
-lingbot-va = [
-    { name = "accelerate" },
-    { name = "diffusers" },
-    { name = "transformers" },
-]
 matplotlib-dep = [
     { name = "contourpy" },
     { name = "matplotlib" },
@@ -3215,7 +3187,6 @@ video-benchmark = [
     { name = "scikit-image" },
 ]
 viz = [
-    { name = "foxglove-sdk" },
     { name = "rerun-sdk" },
 ]
 vla-jepa = [
@@ -3255,7 +3226,6 @@ requires-dist = [
     { name = "fastapi", marker = "extra == 'phone'", specifier = "<1.0" },
     { name = "feetech-servo-sdk", marker = "extra == 'feetech'", specifier = ">=1.0.0,<2.0.0" },
     { name = "flash-attn", marker = "sys_platform != 'darwin' and extra == 'groot'", specifier = ">=2.5.9,<3.0.0" },
-    { name = "foxglove-sdk", marker = "extra == 'viz'", specifier = ">=0.25.1,<0.26.0" },
     { name = "grpcio", marker = "extra == 'grpcio-dep'", specifier = ">=1.73.1,<2.0.0" },
     { name = "grpcio", marker = "extra == 'reachy2'", specifier = "<=1.73.1" },
     { name = "grpcio-tools", marker = "extra == 'dev'", specifier = ">=1.73.1,<2.0.0" },
@@ -3270,7 +3240,6 @@ requires-dist = [
     { name = "ipykernel", marker = "extra == 'notebook'", specifier = ">=6.0.0,<7.0.0" },
     { name = "jsonlines", marker = "extra == 'dataset'", specifier = ">=4.0.0,<5.0.0" },
     { name = "jupyter", marker = "extra == 'notebook'", specifier = ">=1.0.0,<2.0.0" },
-    { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'lingbot-va'" },
     { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'smolvla'" },
     { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'training'" },
     { name = "lerobot", extras = ["aloha"], marker = "extra == 'all'" },
@@ -3298,7 +3267,6 @@ requires-dist = [
     { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'diffusion'" },
     { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'fastwam'" },
     { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'groot'" },
-    { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'lingbot-va'" },
     { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'multi-task-dit'" },
     { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'vla-jepa'" },
     { name = "lerobot", extras = ["diffusion"], marker = "extra == 'all'" },
@@ -3319,7 +3287,6 @@ requires-dist = [
     { name = "lerobot", extras = ["kinematics"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["lekiwi"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["libero"], marker = "sys_platform == 'linux' and extra == 'all'" },
-    { name = "lerobot", extras = ["lingbot-va"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'async'" },
     { name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'sarm'" },
     { name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'unitree-g1'" },
@@ -3378,7 +3345,6 @@ requires-dist = [
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'groot'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'hilserl'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'libero'" },
-    { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'lingbot-va'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'molmoact2'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'multi-task-dit'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'peft'" },
@@ -3458,7 +3424,7 @@ requires-dist = [
     { name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" },
     { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" },
 ]
-provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "hilserl", "vla-jepa", "lingbot-va", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
+provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
 
 [[package]]
 name = "librt"