feat(policiy): fine tune vision encoder refactor completness

Fixes #1774

AGENT_GUIDE.md

@@ -61,20 +61,16 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
 **4.1 Install**
 
 ```bash
-# uv (recommended — see AGENTS.md and CLAUDE.md)
-uv sync --locked --extra feetech          # SO-100/SO-101 motor stack
-# uv sync --locked --extra all            # everything
-# uv sync --locked --extra smolvla        # add SmolVLA deps
-
-# pip (alternative, e.g. when not working from source)
-# pip install 'lerobot[feetech]'
-# pip install 'lerobot[all]'
-# pip install 'lerobot[smolvla]'
-
+pip install 'lerobot[feetech]'              # SO-100/SO-101 motor stack
+# pip install 'lerobot[all]'                # everything
+# pip install 'lerobot[aloha,pusht]'        # specific features
+# pip install 'lerobot[smolvla]'            # add SmolVLA deps
 git lfs install && git lfs pull
-hf auth login                             # required to push datasets/policies
+hf auth login                               # required to push datasets/policies
 ```
 
+Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
+
 **4.2 Find USB ports** — run once per arm, unplug when prompted.
 
 ```bash
@@ -325,10 +321,11 @@ SmolVLA ships with `freeze_vision_encoder=True`. Unfreezing usually **improves p
 
 ```bash
 lerobot-train ... --policy.type=smolvla \
-  --policy.freeze_vision_encoder=false \
-  --policy.train_expert_only=false
+  --policy.fine_tune_vision_encoder=true
 ```
 
+This selectively trains the vision encoder and connector while leaving the language model frozen. Their learning rate defaults to `0.1 × optimizer_lr`; adjust it with `--policy.vision_encoder_lr_multiplier` if needed.
+
 ### 7.7 Signals to stop / keep going
 
 - Train loss plateaus → stop, save a Hub checkpoint.

docs/source/env_processor.mdx

@@ -88,6 +88,20 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
 
 The same policy can work with different environment processors, and the same environment processor can work with different policies:
 
+````python
+# Use SmolVLA policy with LIBERO environment
+# Use SmolVLA policy with LIBERO environment
+libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
+    env_cfg=libero_cfg,
+    policy_cfg=smolvla_cfg,
+)
+smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
+# Or use ACT policy with the same LIBERO environment
+libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
+    env_cfg=libero_cfg,
+    policy_cfg=act_cfg,
+)
+act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
 ```python
 # Use SmolVLA policy with LIBERO environment
 libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -102,7 +116,6 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
     policy_cfg=act_cfg,
 )
 act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
-```
 
 ### 3. **Easier Experimentation**
 
@@ -132,7 +145,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
         state = torch.cat([eef_pos, eef_axisangle, eef_vel,
                           gripper_pos, gripper_vel], dim=-1)  # 14D
         return state
-```
+````
 
 ### 4. **Cleaner Environment Code**
 

docs/source/hope_jr.mdx

@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
 
 ### Record
 
-This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
+This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
 
 ```bash
 lerobot-record \

docs/source/integrate_hardware.mdx

@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
 - [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
 
 Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
-For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
+For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
 
 Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
 

docs/source/lekiwi.mdx

@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
 pip install -e ".[lekiwi]"
 ```
 
-Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
+Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
 Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
 
 # Step-by-Step Assembly Instructions

docs/source/pi0fast.mdx

@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
 
 ## Reproducing π₀Fast results
 
-We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
+We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
 
 The finetuned model can be found here:
 

docs/source/smolvla.mdx

@@ -70,6 +70,19 @@ cd lerobot && lerobot-train \
   GPU allows it, as long as loading times remain short.
 </Tip>
 
+For tasks that require adapting visual features, such as distinguishing new colors or shapes, selectively
+fine-tune the vision encoder and its connector:
+
+```bash
+lerobot-train ... \
+  --policy.path=lerobot/smolvla_base \
+  --policy.fine_tune_vision_encoder=true
+```
+
+This keeps the language model frozen with the default `train_expert_only=true` setting and trains the vision
+path at `0.1` times the main learning rate by default. Fine-tuning the vision encoder increases memory use and
+can reduce the model's general visual knowledge, so enable it only when the frozen encoder is insufficient.
+
 Fine-tuning is an art. For a complete overview of the options for finetuning, run
 
 ```bash
@@ -93,7 +106,7 @@ lerobot-train --help
 
 ## Evaluate the finetuned model and run it in real-time
 
-Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
+Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
 Once you are logged in, you can run inference in your setup by doing:
 
 ```bash

src/lerobot/policies/smolvla/configuration_smolvla.py

@@ -67,6 +67,8 @@ class SmolVLAConfig(PreTrainedConfig):
 
     # Finetuning settings
     freeze_vision_encoder: bool = True
+    fine_tune_vision_encoder: bool = False  # Fine-tune vision + connector; takes priority over freezing.
+    vision_encoder_lr_multiplier: float = 0.1
     train_expert_only: bool = True
     train_state_proj: bool = True
 
@@ -110,6 +112,12 @@ class SmolVLAConfig(PreTrainedConfig):
         super().__post_init__()
 
         """Input validation (not exhaustive)."""
+        if self.fine_tune_vision_encoder:
+            self.freeze_vision_encoder = False
+        if self.vision_encoder_lr_multiplier <= 0:
+            raise ValueError(
+                f"`vision_encoder_lr_multiplier` must be positive, got {self.vision_encoder_lr_multiplier}."
+            )
         if self.n_action_steps > self.chunk_size:
             raise ValueError(
                 f"The chunk size is the upper bound for the number of action steps per model invocation. Got "

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -186,8 +186,27 @@ class SmolVLAPolicy(PreTrainedPolicy):
             if model_value is not None:
                 model_value.rtc_processor = self.rtc_processor
 
-    def get_optim_params(self) -> dict:
-        return self.parameters()
+    def get_optim_params(self):
+        if not self.config.fine_tune_vision_encoder:
+            return self.parameters()
+
+        vision_params = []
+        other_params = []
+        for name, param in self.named_parameters():
+            if not param.requires_grad:
+                continue
+            if ".vision_model." in name or ".connector." in name:
+                vision_params.append(param)
+            else:
+                other_params.append(param)
+
+        return [
+            {"params": other_params},
+            {
+                "params": vision_params,
+                "lr": self.config.optimizer_lr * self.config.vision_encoder_lr_multiplier,
+            },
+        ]
 
     def _get_action_chunk(
         self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
@@ -493,6 +512,7 @@ class VLAFlowMatching(nn.Module):
         self.vlm_with_expert = SmolVLMWithExpertModel(
             model_id=self.config.vlm_model_name,
             freeze_vision_encoder=self.config.freeze_vision_encoder,
+            fine_tune_vision_encoder=self.config.fine_tune_vision_encoder,
             train_expert_only=self.config.train_expert_only,
             load_vlm_weights=self.config.load_vlm_weights,
             attention_mode=self.config.attention_mode,

src/lerobot/policies/smolvla/smolvlm_with_expert.py

@@ -78,6 +78,7 @@ class SmolVLMWithExpertModel(nn.Module):
         load_vlm_weights: bool = True,
         train_expert_only: bool = True,
         freeze_vision_encoder: bool = False,
+        fine_tune_vision_encoder: bool = False,
         attention_mode: str = "self_attn",
         num_expert_layers: int = -1,
         num_vlm_layers: int = -1,
@@ -141,6 +142,7 @@ class SmolVLMWithExpertModel(nn.Module):
         self.num_key_value_heads = self.config.text_config.num_key_value_heads
 
         self.freeze_vision_encoder = freeze_vision_encoder
+        self.fine_tune_vision_encoder = fine_tune_vision_encoder
         self.train_expert_only = train_expert_only
         self.attention_mode = attention_mode
         self.expert_hidden_size = lm_expert_config.hidden_size
@@ -150,10 +152,6 @@ class SmolVLMWithExpertModel(nn.Module):
         return self.vlm.model
 
     def set_requires_grad(self):
-        if self.freeze_vision_encoder:
-            self.get_vlm_model().vision_model.eval()
-            for params in self.get_vlm_model().vision_model.parameters():
-                params.requires_grad = False
         if self.train_expert_only:
             self.vlm.eval()
             for params in self.vlm.parameters():
@@ -176,6 +174,18 @@ class SmolVLMWithExpertModel(nn.Module):
             for name, params in self.vlm.named_parameters():
                 if any(k in name for k in frozen_layers):
                     params.requires_grad = False
+
+        if self.freeze_vision_encoder:
+            self.get_vlm_model().vision_model.eval()
+            for params in self.get_vlm_model().vision_model.parameters():
+                params.requires_grad = False
+
+        if self.fine_tune_vision_encoder:
+            for params in self.get_vlm_model().vision_model.parameters():
+                params.requires_grad = True
+            for params in self.get_vlm_model().connector.parameters():
+                params.requires_grad = True
+
         # To avoid unused params issue with distributed training
         for name, params in self.lm_expert.named_parameters():
             if "lm_head" in name:
@@ -184,11 +194,15 @@ class SmolVLMWithExpertModel(nn.Module):
     def train(self, mode: bool = True):
         super().train(mode)
 
+        if self.train_expert_only:
+            self.vlm.eval()
+
         if self.freeze_vision_encoder:
             self.get_vlm_model().vision_model.eval()
 
-        if self.train_expert_only:
-            self.vlm.eval()
+        if self.fine_tune_vision_encoder:
+            self.get_vlm_model().vision_model.train(mode)
+            self.get_vlm_model().connector.train(mode)
 
     def embed_image(self, image: torch.Tensor):
         patch_attention_mask = None

src/lerobot/rl/buffer.py

@@ -18,7 +18,7 @@ import functools
 import threading
 from collections.abc import Callable, Sequence
 from contextlib import suppress
-from typing import NotRequired, TypedDict
+from typing import TypedDict
 
 import torch
 import torch.nn.functional as F  # noqa: N812
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
     next_state: dict[str, torch.Tensor]
     done: torch.Tensor
     truncated: torch.Tensor
-    complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
+    complementary_info: dict[str, torch.Tensor | float | int] | None = None
 
 
 def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:

src/lerobot/scripts/lerobot_dataset_viz.py

@@ -87,11 +87,6 @@ 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.dataset_visualization_utils import (
-    get_extra_scalar_keys,
-    is_scalar_like,
-    scalar_to_float,
-)
 from lerobot.utils.utils import init_logging
 
 logger = logging.getLogger(__name__)
@@ -158,8 +153,6 @@ def build_blueprint_from_dataset(dataset: LeRobotDataset):
     for key in (DONE, REWARD, SUCCESS):
         if key in dataset.features:
             views.append(rrb.TimeSeriesView(origin=key, name=key))
-    for key in get_extra_scalar_keys(dataset):
-        views.append(rrb.TimeSeriesView(origin=key, name=key))
 
     return rrb.Blueprint(rrb.Grid(*views))
 
@@ -251,8 +244,6 @@ def visualize_dataset(
         hi = stats["q99"] if "q99" in stats else stats["max"]
         depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item()))
 
-    extra_scalar_keys = get_extra_scalar_keys(dataset)
-
     first_index = None
     for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
         if first_index is None:
@@ -296,10 +287,6 @@ def visualize_dataset(
             if SUCCESS in batch:
                 rr.log(SUCCESS, rr.Scalars(batch[SUCCESS][i].item()))
 
-            for key in extra_scalar_keys:
-                if key in batch and is_scalar_like(batch[key][i]):
-                    rr.log(key, rr.Scalars(scalar_to_float(batch[key][i])))
-
     # save .rrd locally
     if mode == "local" and save:
         output_dir.mkdir(parents=True, exist_ok=True)

src/lerobot/scripts/lerobot_find_cameras.py

@@ -28,6 +28,7 @@ lerobot-find-cameras
 # NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
 
 import argparse
+import concurrent.futures
 import logging
 import time
 from pathlib import Path
@@ -132,7 +133,7 @@ def save_image(
     camera_identifier: str | int,
     images_dir: Path,
     camera_type: str,
-) -> None:
+):
     """
     Saves a single image to disk using Pillow. Handles color conversion if necessary.
     """
@@ -151,7 +152,7 @@ def save_image(
         logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
 
 
-def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None:
+def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
     """Create and connect to a camera instance based on metadata."""
     cam_type = cam_meta.get("type")
     cam_id = cam_meta.get("id")
@@ -164,14 +165,12 @@ def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> di
             cv_config = OpenCVCameraConfig(
                 index_or_path=cam_id,
                 color_mode=ColorMode.RGB,
-                warmup_s=warmup_s,
             )
             instance = OpenCVCamera(cv_config)
         elif cam_type == "RealSense":
             rs_config = RealSenseCameraConfig(
                 serial_number_or_name=cam_id,
                 color_mode=ColorMode.RGB,
-                warmup_s=warmup_s,
             )
             instance = RealSenseCamera(rs_config)
         else:
@@ -189,7 +188,9 @@ def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> di
         return None
 
 
-def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None:
+def process_camera_image(
+    cam_dict: dict[str, Any], output_dir: Path, current_time: float
+) -> concurrent.futures.Future | None:
     """Capture and process an image from a single camera."""
     cam = cam_dict["instance"]
     meta = cam_dict["meta"]
@@ -199,7 +200,7 @@ def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_tim
     try:
         image_data = cam.read()
 
-        save_image(
+        return save_image(
             image_data,
             cam_id_str,
             output_dir,
@@ -214,21 +215,21 @@ def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_tim
     return None
 
 
-def cleanup_camera(cam_dict: dict[str, Any]) -> None:
+def cleanup_cameras(cameras_to_use: list[dict[str, Any]]):
     """Disconnect all cameras."""
-    logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...")
-    try:
-        if cam_dict["instance"] and cam_dict["instance"].is_connected:
-            cam_dict["instance"].disconnect()
-    except Exception as e:
-        logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
+    logger.info(f"Disconnecting {len(cameras_to_use)} cameras...")
+    for cam_dict in cameras_to_use:
+        try:
+            if cam_dict["instance"] and cam_dict["instance"].is_connected:
+                cam_dict["instance"].disconnect()
+        except Exception as e:
+            logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
 
 
 def save_images_from_all_cameras(
     output_dir: Path,
     record_time_s: float = 2.0,
     camera_type: str | None = None,
-    warmup_s: int = 1,
 ):
     """
     Connects to detected cameras (optionally filtered by type) and saves images from each.
@@ -239,7 +240,6 @@ def save_images_from_all_cameras(
         record_time_s: Duration in seconds to record images.
         camera_type: Optional string to filter cameras ("realsense" or "opencv").
                             If None, uses all detected cameras.
-        warmup_s: Duration in seconds to warmup camera before recording images.
     """
     output_dir.mkdir(parents=True, exist_ok=True)
     logger.info(f"Saving images to {output_dir}")
@@ -249,24 +249,40 @@ def save_images_from_all_cameras(
         logger.warning("No cameras detected matching the criteria. Cannot save images.")
         return
 
-    logger.info(
-        f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras."
-    )
+    cameras_to_use = []
+    for cam_meta in all_camera_metadata:
+        camera_instance = create_camera_instance(cam_meta)
+        if camera_instance:
+            cameras_to_use.append(camera_instance)
 
-    try:
-        for cam_meta in all_camera_metadata:
-            cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s)
-            if cam_dict is None:
-                continue
-            start_time = time.perf_counter()
+    if not cameras_to_use:
+        logger.warning("No cameras could be connected. Aborting image save.")
+        return
+
+    logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
+    start_time = time.perf_counter()
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
+        try:
             while time.perf_counter() - start_time < record_time_s:
+                futures = []
                 current_capture_time = time.perf_counter()
-                process_camera_image(cam_dict, output_dir, current_capture_time)
-            cleanup_camera(cam_dict)
-    except KeyboardInterrupt:
-        logger.info("Capture interrupted by user.")
-    finally:
-        print(f"Image capture finished. Images saved to {output_dir}")
+
+                for cam_dict in cameras_to_use:
+                    future = process_camera_image(cam_dict, output_dir, current_capture_time)
+                    if future:
+                        futures.append(future)
+
+                if futures:
+                    concurrent.futures.wait(futures)
+
+        except KeyboardInterrupt:
+            logger.info("Capture interrupted by user.")
+        finally:
+            print("\nFinalizing image saving...")
+            executor.shutdown(wait=True)
+            cleanup_cameras(cameras_to_use)
+            print(f"Image capture finished. Images saved to {output_dir}")
 
 
 def main():
@@ -275,6 +291,7 @@ def main():
     parser = argparse.ArgumentParser(
         description="Unified camera utility script for listing cameras and capturing images."
     )
+
     parser.add_argument(
         "camera_type",
         type=str,
@@ -292,14 +309,8 @@ def main():
     parser.add_argument(
         "--record-time-s",
         type=float,
-        default=2.0,
-        help="Time duration to attempt capturing frames. Default: 2 seconds.",
-    )
-    parser.add_argument(
-        "--warmup-s",
-        type=int,
-        default=1,
-        help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
+        default=6.0,
+        help="Time duration to attempt capturing frames. Default: 6 seconds.",
     )
     args = parser.parse_args()
     save_images_from_all_cameras(**vars(args))

src/lerobot/teleoperators/phone/teleop_phone.py

@@ -171,13 +171,7 @@ class IOSPhone(BasePhone, Teleoperator):
         # HEBI provides orientation in w, x, y, z format.
         # Scipy's Rotation expects x, y, z, w.
         quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]]))  # wxyz to xyzw
-        # ARKit can emit zero/NaN quaternions before tracking is ready or on a
-        # dropped packet. Rotation.from_quat now rejects those; degrade the same
-        # way as a missing pose so teleop stays alive mid-session.
-        try:
-            rot = Rotation.from_quat(quat_xyzw)
-        except ValueError:
-            return False, None, None, None
+        rot = Rotation.from_quat(quat_xyzw)
         pos = ar_pos - rot.apply(self.config.camera_offset)
         return True, pos, rot, pose
 

src/lerobot/utils/dataset_visualization_utils.py

@@ -1,99 +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.
-
-"""Shared helpers for visualizing scalar features from a LeRobot dataset."""
-
-from collections.abc import Iterable, Mapping
-from typing import TYPE_CHECKING
-
-import numpy as np
-import torch
-
-from .constants import ACTION, DEFAULT_FEATURES, DONE, OBS_STATE, REWARD, SUCCESS
-
-if TYPE_CHECKING:
-    from lerobot.datasets import LeRobotDataset
-
-
-METADATA_KEYS = {*DEFAULT_FEATURES, "task"}
-KNOWN_SCALAR_KEYS = {DONE, REWARD, SUCCESS}
-SCALAR_DTYPE_KINDS = {"b", "i", "u", "f"}
-
-
-def is_scalar_feature(feature: Mapping) -> bool:
-    """Return whether a feature schema describes a numeric or boolean scalar."""
-
-    dtype = feature.get("dtype")
-    if not isinstance(dtype, str):
-        return False
-    try:
-        dtype_kind = np.dtype(dtype).kind
-    except (TypeError, ValueError):
-        return False
-    if dtype_kind not in SCALAR_DTYPE_KINDS:
-        return False
-
-    shape = feature.get("shape")
-    if shape is None:
-        return True
-    if isinstance(shape, int):
-        return shape == 1
-    if not isinstance(shape, (list, tuple)):
-        return False
-    return len(shape) == 0 or (len(shape) == 1 and shape[0] == 1)
-
-
-def get_extra_scalar_keys(dataset: "LeRobotDataset", additional_known_keys: Iterable[str] = ()) -> list[str]:
-    """Return scalar feature keys not handled by the visualizer's standard paths."""
-
-    known_keys = {
-        ACTION,
-        OBS_STATE,
-        *KNOWN_SCALAR_KEYS,
-        *METADATA_KEYS,
-        *additional_known_keys,
-        *dataset.meta.camera_keys,
-    }
-    return [
-        key
-        for key, feature in dataset.features.items()
-        if key not in known_keys and is_scalar_feature(feature)
-    ]
-
-
-def is_scalar_like(value: object) -> bool:
-    """Return whether a runtime value contains exactly one numeric or boolean scalar."""
-
-    if isinstance(value, torch.Tensor):
-        return value.numel() == 1 and not value.is_complex()
-    if isinstance(value, np.ndarray):
-        return value.size == 1 and value.dtype.kind in SCALAR_DTYPE_KINDS
-    return np.isscalar(value) and np.asarray(value).dtype.kind in SCALAR_DTYPE_KINDS
-
-
-def scalar_to_float(value: object) -> float:
-    """Convert a scalar-like tensor, array, or Python value to ``float``."""
-
-    return float(value.item() if hasattr(value, "item") else value)
-
-
-def get_scalar_values(sample: Mapping, keys: Iterable[str]) -> dict[str, float]:
-    """Select and convert scalar-like values from ``sample`` for the requested keys."""
-
-    values = {}
-    for key in keys:
-        value = sample.get(key)
-        if value is not None and is_scalar_like(value):
-            values[key] = scalar_to_float(value)
-    return values

src/lerobot/utils/device_utils.py

@@ -37,25 +37,16 @@ def auto_select_torch_device() -> torch.device:
 
 # TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
 def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
-    """Given a string, return a torch.device with checks on whether the device is available.
-
-    Raises:
-        ValueError: If the requested device family is known but not available on
-            this machine (``AssertionError`` was previously used and is easy to
-            mistake for a programmer bug under ``python -O`` where asserts vanish).
-    """
+    """Given a string, return a torch.device with checks on whether the device is available."""
     try_device = str(try_device)
     if try_device.startswith("cuda"):
-        if not torch.cuda.is_available():
-            raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
+        assert torch.cuda.is_available()
         device = torch.device(try_device)
     elif try_device == "mps":
-        if not torch.backends.mps.is_available():
-            raise ValueError("Requested device 'mps' but MPS is not available.")
+        assert torch.backends.mps.is_available()
         device = torch.device("mps")
     elif try_device == "xpu":
-        if not torch.xpu.is_available():
-            raise ValueError("Requested device 'xpu' but XPU is not available.")
+        assert torch.xpu.is_available()
         device = torch.device("xpu")
     elif try_device == "cpu":
         device = torch.device("cpu")

src/lerobot/utils/foxglove_visualization.py

@@ -23,7 +23,6 @@ importing from here directly. Requires the ``viz`` extra (``pip install 'lerobot
 import logging
 import numbers
 import time
-from collections.abc import Iterable, Mapping
 
 import cv2
 import numpy as np
@@ -42,7 +41,6 @@ from .constants import (
     SUCCESS,
     TRUNCATED,
 )
-from .dataset_visualization_utils import get_extra_scalar_keys, get_scalar_values
 from .import_utils import require_package
 
 # Static schema shared by all scalar topics. Each message carries a flat list of ``{label, value}``
@@ -407,24 +405,6 @@ def _frame_to_scalars(sample: dict, key: str, labels: list[str] | None = None) -
     return _labeled_scalars(name, arr.flatten(), labels)
 
 
-def _dataset_frame_scalar_groups(
-    sample: Mapping, extra_scalar_keys: Iterable[str]
-) -> tuple[dict[str, float], dict[str, float]]:
-    """Return standard episode scalars and custom scalar features for one dataset frame."""
-
-    episode_scalars = {}
-    for feature, label in (
-        (DONE, "done"),
-        (TRUNCATED, "truncated"),
-        (REWARD, "reward"),
-        (SUCCESS, "success"),
-    ):
-        value = sample.get(feature)
-        if value is not None:
-            episode_scalars[label] = float(value)
-    return episode_scalars, get_scalar_values(sample, extra_scalar_keys)
-
-
 def serve_foxglove_dataset_playback(
     dataset,
     episode_index: int,
@@ -472,7 +452,6 @@ def serve_foxglove_dataset_playback(
         raise ValueError("Cannot visualize an empty episode.")
     first_ns, last_ns = times_ns[0], times_ns[-1]
     camera_keys = list(dataset.meta.camera_keys)
-    extra_scalar_keys = get_extra_scalar_keys(dataset, additional_known_keys=(TRUNCATED,))
     # 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:
@@ -521,14 +500,17 @@ def serve_foxglove_dataset_playback(
             channels=channels,
             log_time=log_time,
         )
-        episode_scalars, extra_scalars = _dataset_frame_scalar_groups(sample, extra_scalar_keys)
+        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)
-        _log_foxglove_scalars(
-            "/episode/extras",
-            extra_scalars,
-            channels=channels,
-            log_time=log_time,
-        )
 
     lock = threading.Lock()
     stop_event = threading.Event()

src/lerobot/utils/io_utils.py

@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
     Returns:
         Any: The data loaded from the JSON file.
     """
-    with open(fpath, encoding="utf-8") as f:
+    with open(fpath) as f:
         return json.load(f)
 
 
-def write_json(data: JsonLike, fpath: Path) -> None:
-    """Write JSON-serializable data to a file.
+def write_json(data: dict, fpath: Path) -> None:
+    """Write data to a JSON file.
 
     Creates parent directories if they don't exist.
 
     Args:
-        data: JSON-serializable data to write.
+        data (dict): The dictionary to write.
         fpath (Path): The path to the output JSON file.
     """
     fpath.parent.mkdir(exist_ok=True, parents=True)
-    with open(fpath, "w", encoding="utf-8") as f:
+    with open(fpath, "w") as f:
         json.dump(data, f, indent=4, ensure_ascii=False)
 
 

src/lerobot/utils/robot_utils.py

@@ -30,10 +30,6 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
     """
     if seconds <= 0:
         return
-    if spin_threshold < 0:
-        raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
-    if sleep_margin < 0:
-        raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")
 
     system = platform.system()
     # On macOS and Windows the scheduler / sleep granularity can make

src/lerobot/utils/rotation.py

@@ -29,13 +29,10 @@ class Rotation:
     def __init__(self, quat: np.ndarray) -> None:
         """Initialize rotation from quaternion [x, y, z, w]."""
         self._quat = np.asarray(quat, dtype=float)
-        if self._quat.shape != (4,):
-            raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
-        # Normalize quaternion. Reject the zero vector — it has no orientation.
+        # Normalize quaternion
         norm = np.linalg.norm(self._quat)
-        if norm <= 0.0 or not np.isfinite(norm):
-            raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})")
-        self._quat = self._quat / norm
+        if norm > 0:
+            self._quat = self._quat / norm
 
     @classmethod
     def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":

src/lerobot/utils/transition.py

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import NotRequired, TypedDict
+from typing import TypedDict
 
 import torch
 
@@ -28,7 +28,7 @@ class Transition(TypedDict):
     next_state: dict[str, torch.Tensor]
     done: bool
     truncated: bool
-    complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
+    complementary_info: dict[str, torch.Tensor | float | int] | None = None
 
 
 def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:

src/lerobot/utils/utils.py

@@ -24,6 +24,7 @@ import sys
 import time
 from collections.abc import Iterator
 from copy import copy, deepcopy
+from datetime import datetime
 from pathlib import Path
 from statistics import mean
 from typing import TYPE_CHECKING, Any
@@ -60,16 +61,14 @@ def init_logging(
         accelerator: Optional Accelerator instance (for multi-GPU detection)
     """
 
-    class LeRobotFormatter(logging.Formatter):
-        def format(self, record: logging.LogRecord) -> str:
-            record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:]
-            record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else ""
-            return super().format(record)
+    def custom_format(record: logging.LogRecord) -> str:
+        dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        fnameline = f"{record.pathname}:{record.lineno}"
+        pid_str = f"[PID: {os.getpid()}] " if display_pid else ""
+        return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}"
 
-    formatter = LeRobotFormatter(
-        "%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-    )
+    formatter = logging.Formatter()
+    formatter.format = custom_format
 
     logger = logging.getLogger()
     logger.setLevel(logging.NOTSET)

tests/datasets/test_visualize_dataset.py

@@ -13,78 +13,11 @@
 # 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 sys
-from types import SimpleNamespace
-
-import numpy as np
 import pytest
-import torch
 
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 
 from lerobot.scripts.lerobot_dataset_viz import visualize_dataset
-from lerobot.utils import import_utils
-from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS, TRUNCATED
-from lerobot.utils.dataset_visualization_utils import (
-    get_extra_scalar_keys,
-    is_scalar_feature,
-    is_scalar_like,
-    scalar_to_float,
-)
-
-
-class DummyMeta:
-    camera_keys = ["observation.images.front"]
-
-
-class DummyFeatureDataset:
-    meta = DummyMeta()
-    features = {
-        "index": {"dtype": "int64", "shape": [1]},
-        "timestamp": {"dtype": "float32", "shape": [1]},
-        "episode_index": {"dtype": "int64", "shape": [1]},
-        "frame_index": {"dtype": "int64", "shape": [1]},
-        "task_index": {"dtype": "int64", "shape": [1]},
-        ACTION: {"dtype": "float32", "shape": [6]},
-        OBS_STATE: {"dtype": "float32", "shape": [6]},
-        DONE: {"dtype": "bool", "shape": [1]},
-        REWARD: {"dtype": "float32", "shape": [1]},
-        SUCCESS: {"dtype": "bool", "shape": [1]},
-        TRUNCATED: {"dtype": "bool", "shape": [1]},
-        "observation.images.front": {"dtype": "video", "shape": [3, 480, 640]},
-        "q_target": {"dtype": "float32", "shape": [1]},
-        "quality": {"dtype": "double", "shape": [1]},
-        "intervention": {"dtype": "bool", "shape": []},
-        "embedding": {"dtype": "float32", "shape": [32]},
-        "comment": {"dtype": "string", "shape": [1]},
-    }
-
-
-class DummyVizDataset:
-    repo_id = "dummy/custom-scalars"
-    depth_output_unit = "m"
-    meta = SimpleNamespace(camera_keys=[], depth_keys=[], stats=None)
-    features = {
-        "index": {"dtype": "int64", "shape": [1]},
-        "timestamp": {"dtype": "float32", "shape": [1]},
-        ACTION: {"dtype": "float32", "shape": [2], "names": ["x", "y"]},
-        "q_target": {"dtype": "float32", "shape": [1]},
-        "intervention": {"dtype": "bool", "shape": [1]},
-        "embedding": {"dtype": "float32", "shape": [2]},
-    }
-
-    def __len__(self):
-        return 2
-
-    def __getitem__(self, index):
-        return {
-            "index": torch.tensor(index),
-            "timestamp": torch.tensor(index / 10),
-            ACTION: torch.tensor([index, index + 1], dtype=torch.float32),
-            "q_target": torch.tensor([index + 0.5]),
-            "intervention": torch.tensor([index == 0]),
-            "embedding": torch.tensor([index, index + 1], dtype=torch.float32),
-        }
 
 
 @pytest.mark.skip("TODO: add dummy videos")
@@ -100,93 +33,3 @@ def test_visualize_local_dataset(tmp_path, lerobot_dataset_factory):
         output_dir=output_dir,
     )
     assert rrd_path.exists()
-
-
-def test_get_extra_scalar_keys_skips_known_metadata_and_non_scalars():
-    dataset = DummyFeatureDataset()
-
-    assert get_extra_scalar_keys(dataset) == [TRUNCATED, "q_target", "quality", "intervention"]
-    assert get_extra_scalar_keys(dataset, additional_known_keys=(TRUNCATED,)) == [
-        "q_target",
-        "quality",
-        "intervention",
-    ]
-
-
-@pytest.mark.parametrize(
-    ("feature", "expected"),
-    [
-        ({"dtype": "float32", "shape": (1,)}, True),
-        ({"dtype": "double", "shape": [1]}, True),
-        ({"dtype": "bool", "shape": []}, True),
-        ({"dtype": "int64", "shape": 1}, True),
-        ({"dtype": "float32", "shape": (3,)}, False),
-        ({"dtype": "complex64", "shape": [1]}, False),
-        ({"dtype": "datetime64[ns]", "shape": [1]}, False),
-        ({"dtype": "string", "shape": [1]}, False),
-        ({"shape": [1]}, False),
-    ],
-)
-def test_is_scalar_feature(feature, expected):
-    assert is_scalar_feature(feature) is expected
-
-
-@pytest.mark.parametrize(
-    ("value", "expected"),
-    [
-        (torch.tensor(1.5), True),
-        (torch.tensor([1.5]), True),
-        (torch.tensor([1.5, 2.5]), False),
-        (np.array(2.0), True),
-        (np.array([2.0]), True),
-        (np.array([2.0, 3.0]), False),
-        (True, True),
-        ("not numeric", False),
-    ],
-)
-def test_is_scalar_like(value, expected):
-    assert is_scalar_like(value) is expected
-
-
-def test_scalar_to_float():
-    assert scalar_to_float(torch.tensor(3.0)) == 3.0
-    assert scalar_to_float(np.array([4.0])) == 4.0
-
-
-def test_visualize_dataset_logs_extra_scalars(monkeypatch):
-    logged = []
-    initialized = []
-
-    dummy_rrb = SimpleNamespace(
-        Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
-            kind="spatial", origin=origin, name=name
-        ),
-        TimeSeriesView=lambda origin=None, name=None, overrides=None: SimpleNamespace(
-            kind="time_series", origin=origin, name=name, overrides=overrides
-        ),
-        Grid=lambda *views: SimpleNamespace(views=views),
-        Blueprint=lambda root: SimpleNamespace(root=root),
-    )
-    dummy_rr = SimpleNamespace(
-        SeriesLines=lambda names=None: SimpleNamespace(names=names),
-        Scalars=lambda value: SimpleNamespace(value=value),
-        init=lambda *args, **kwargs: initialized.append((args, kwargs)),
-        log=lambda key, entity: logged.append((key, entity)),
-        set_time=lambda *args, **kwargs: None,
-        blueprint=dummy_rrb,
-    )
-    monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
-    monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
-    monkeypatch.setattr(import_utils, "require_package", lambda *args, **kwargs: None)
-
-    visualize_dataset(DummyVizDataset(), episode_index=0, batch_size=2)
-
-    logged_keys = [key for key, _ in logged]
-    assert logged_keys.count("q_target") == 2
-    assert logged_keys.count("intervention") == 2
-    assert "embedding" not in logged_keys
-
-    blueprint = initialized[0][1]["default_blueprint"]
-    time_series_origins = {view.origin for view in blueprint.root.views if view.kind == "time_series"}
-    assert {"q_target", "intervention"} <= time_series_origins
-    assert "embedding" not in time_series_origins

tests/utils/test_device_utils.py

@@ -1,36 +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 unittest.mock import patch
-
-import pytest
-import torch
-
-from lerobot.utils.device_utils import get_safe_torch_device, is_torch_device_available
-
-
-def test_cpu_always_available():
-    assert get_safe_torch_device("cpu") == torch.device("cpu")
-    assert is_torch_device_available("cpu")
-
-
-def test_missing_cuda_raises_valueerror():
-    with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"):
-        get_safe_torch_device("cuda")
-
-
-def test_missing_mps_raises_valueerror():
-    with patch("torch.backends.mps.is_available", return_value=False), pytest.raises(ValueError, match="MPS"):
-        get_safe_torch_device("mps")

tests/utils/test_foxglove_visualization.py

@@ -24,7 +24,7 @@ the functions that talk to the server, so the helpers below run in the base test
 import numpy as np
 
 from lerobot.utils import foxglove_visualization as fv
-from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
+from lerobot.utils.constants import ACTION, OBS_STATE
 
 
 def test_foxglove_safe_name_collapses_dots():
@@ -93,24 +93,6 @@ def test_feature_dim_names_formats():
     assert fv._feature_dim_names({"shape": [2]}) is None
 
 
-def test_dataset_frame_scalar_groups_include_custom_scalars():
-    sample = {
-        DONE: np.array(True),
-        REWARD: np.array(0.5),
-        SUCCESS: np.array(False),
-        "q_target": np.array([0.75]),
-        "intervention": np.array([True]),
-        "embedding": np.array([1.0, 2.0]),
-    }
-
-    episode_scalars, extra_scalars = fv._dataset_frame_scalar_groups(
-        sample, ("q_target", "intervention", "embedding")
-    )
-
-    assert episode_scalars == {"done": 1.0, "reward": 0.5, "success": 0.0}
-    assert extra_scalars == {"q_target": 0.75, "intervention": 1.0}
-
-
 def test_is_scalar():
     assert fv._is_scalar(1.0)
     assert fv._is_scalar(np.float32(2.0))

tests/utils/test_rotation.py

@@ -1,46 +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.
-
-import numpy as np
-import pytest
-
-from lerobot.utils.rotation import Rotation
-
-
-def test_zero_quaternion_rejected():
-    with pytest.raises(ValueError, match="non-zero"):
-        Rotation(np.zeros(4))
-
-
-def test_non_finite_quaternion_rejected():
-    with pytest.raises(ValueError, match="non-zero|finite"):
-        Rotation(np.array([np.nan, 0.0, 0.0, 1.0]))
-
-
-def test_wrong_shape_rejected():
-    with pytest.raises(ValueError, match="shape"):
-        Rotation(np.array([1.0, 0.0, 0.0]))
-
-
-def test_identity_roundtrip():
-    r = Rotation.from_rotvec(np.zeros(3))
-    assert np.allclose(r.as_rotvec(), 0.0)
-    assert np.allclose(r.as_matrix(), np.eye(3))
-
-
-def test_rotvec_roundtrip():
-    rotvec = np.array([0.1, -0.2, 0.3])
-    r = Rotation.from_rotvec(rotvec)
-    assert np.allclose(r.as_rotvec(), rotvec, atol=1e-6)