feat(scripts): add dataset mirroring utility

Add the mirror-dataset script from feat/mirror so bimanual datasets can be mirrored by swapping left/right features, applying joint masks, and flipping video frames.

Made-with: Cursor

examples/backward_compatibility/replay.py

@@ -57,7 +57,7 @@ class DatasetReplayConfig:
     repo_id: str
     # Episode to replay.
     episode: int
-    # Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
+    # Root directory where the dataset will be stored (e.g. 'dataset/path').
     root: str | Path | None = None
     # Limit the frames per second. By default, uses the policy fps.
     fps: int = 30

src/lerobot/async_inference/robot_client.py

@@ -49,18 +49,23 @@ import torch
 
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
-from lerobot.robots import (
-    RobotConfig,  # noqa: F401
+from lerobot.robots import (  # noqa: F401
+    Robot,
+    RobotConfig,
+    bi_so_follower,
+    koch_follower,
     make_robot_from_config,
+    omx_follower,
+    so_follower,
 )
 from lerobot.transport import (
     services_pb2,  # type: ignore
     services_pb2_grpc,  # type: ignore
 )
 from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks
-from lerobot.utils.import_utils import register_third_party_plugins
 
 from .configs import RobotClientConfig
+from .constants import SUPPORTED_ROBOTS
 from .helpers import (
     Action,
     FPSTracker,
@@ -480,9 +485,8 @@ class RobotClient:
 def async_client(cfg: RobotClientConfig):
     logging.info(pformat(asdict(cfg)))
 
-    # TODO: Assert if checking robot support is still needed with the plugin system
-    # if cfg.robot.type not in SUPPORTED_ROBOTS:
-    #     raise ValueError(f"Robot {cfg.robot.type} not yet supported!")
+    if cfg.robot.type not in SUPPORTED_ROBOTS:
+        raise ValueError(f"Robot {cfg.robot.type} not yet supported!")
 
     client = RobotClient(cfg)
 
@@ -508,5 +512,4 @@ def async_client(cfg: RobotClientConfig):
 
 
 if __name__ == "__main__":
-    register_third_party_plugins()
     async_client()  # run the client

src/lerobot/configs/default.py

@@ -27,7 +27,7 @@ class DatasetConfig:
     # "dataset_index" into the returned item. The index mapping is made according to the order in which the
     # datasets are provided.
     repo_id: str
-    # Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
+    # Root directory where the dataset will be stored (e.g. 'dataset/path').
     root: str | None = None
     episodes: list[int] | None = None
     image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)

src/lerobot/datasets/lerobot_dataset.py

@@ -664,11 +664,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
         for the README).
 
         Args:
-            repo_id (str): This is the repo id that will be used to fetch the dataset.
-            root (Path | None, optional): Local directory where the dataset will be downloaded and
-                stored. If set, all dataset files will be stored directly under this path. If not set, the
-                dataset files will be stored under $HF_LEROBOT_HOME/repo_id (configurable via the
-                HF_LEROBOT_HOME environment variable).
+            repo_id (str): This is the repo id that will be used to fetch the dataset. Locally, the dataset
+                will be stored under root/repo_id.
+            root (Path | None, optional): Local directory to use for downloading/writing files. You can also
+                set the HF_LEROBOT_HOME environment variable to point to a different location. Defaults to
+                '~/.cache/huggingface/lerobot'.
             episodes (list[int] | None, optional): If specified, this will only load episodes specified by
                 their episode_index in this list. Defaults to None.
             image_transforms (Callable | None, optional): You can pass standard v2 image transforms from
@@ -1771,12 +1771,11 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
             )
         for repo_id, ds in zip(self.repo_ids, self._datasets, strict=True):
             extra_keys = set(ds.features).difference(intersection_features)
-            if extra_keys:
-                logging.warning(
-                    f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the "
-                    "other datasets."
-                )
-                self.disabled_features.update(extra_keys)
+            logging.warning(
+                f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the "
+                "other datasets."
+            )
+            self.disabled_features.update(extra_keys)
 
         self.image_transforms = image_transforms
         self.delta_timestamps = delta_timestamps

src/lerobot/policies/diffusion/configuration_diffusion.py

@@ -55,16 +55,10 @@ class DiffusionConfig(PreTrainedConfig):
         normalization_mapping: A dictionary that maps from a str value of FeatureType (e.g., "STATE", "VISUAL") to
             a corresponding NormalizationMode (e.g., NormalizationMode.MIN_MAX)
         vision_backbone: Name of the torchvision resnet backbone to use for encoding images.
-        resize_shape: (H, W) shape to resize images to as a preprocessing step for the vision
-            backbone. If None, no resizing is done and the original image resolution is used.
-        crop_ratio: Ratio in (0, 1] used to derive the crop size from resize_shape
-            (crop_h = int(resize_shape[0] * crop_ratio), likewise for width).
-            Set to 1.0 to disable cropping. Only takes effect when resize_shape is not None.
-        crop_shape: (H, W) shape to crop images to. When resize_shape is set and crop_ratio < 1.0,
-            this is computed automatically. Can also be set directly for legacy configs that use
-            crop-only (without resize). If None and no derivation applies, no cropping is done.
-        crop_is_random: Whether the crop should be random at training time (it's always a center
-            crop in eval mode).
+        crop_shape: (H, W) shape to crop images to as a preprocessing step for the vision backbone. Must fit
+            within the image size. If None, no cropping is done.
+        crop_is_random: Whether the crop should be random at training time (it's always a center crop in eval
+            mode).
         pretrained_backbone_weights: Pretrained weights from torchvision to initialize the backbone.
             `None` means no pretrained weights.
         use_group_norm: Whether to replace batch normalization with group normalization in the backbone.
@@ -120,9 +114,7 @@ class DiffusionConfig(PreTrainedConfig):
     # Architecture / modeling.
     # Vision backbone.
     vision_backbone: str = "resnet18"
-    resize_shape: tuple[int, int] | None = None
-    crop_ratio: float = 1.0
-    crop_shape: tuple[int, int] | None = None
+    crop_shape: tuple[int, int] | None = (84, 84)
     crop_is_random: bool = True
     pretrained_backbone_weights: str | None = None
     use_group_norm: bool = True
@@ -183,25 +175,6 @@ class DiffusionConfig(PreTrainedConfig):
                 f"Got {self.noise_scheduler_type}."
             )
 
-        if self.resize_shape is not None and (
-            len(self.resize_shape) != 2 or any(d <= 0 for d in self.resize_shape)
-        ):
-            raise ValueError(f"`resize_shape` must be a pair of positive integers. Got {self.resize_shape}.")
-        if not (0 < self.crop_ratio <= 1.0):
-            raise ValueError(f"`crop_ratio` must be in (0, 1]. Got {self.crop_ratio}.")
-
-        if self.resize_shape is not None:
-            if self.crop_ratio < 1.0:
-                self.crop_shape = (
-                    int(self.resize_shape[0] * self.crop_ratio),
-                    int(self.resize_shape[1] * self.crop_ratio),
-                )
-            else:
-                # Explicitly disable cropping for resize+ratio path when crop_ratio == 1.0.
-                self.crop_shape = None
-        if self.crop_shape is not None and (self.crop_shape[0] <= 0 or self.crop_shape[1] <= 0):
-            raise ValueError(f"`crop_shape` must have positive dimensions. Got {self.crop_shape}.")
-
         # Check that the horizon size and U-Net downsampling is compatible.
         # U-Net downsamples by 2 with each stage.
         downsampling_factor = 2 ** len(self.down_dims)
@@ -229,12 +202,13 @@ class DiffusionConfig(PreTrainedConfig):
         if len(self.image_features) == 0 and self.env_state_feature is None:
             raise ValueError("You must provide at least one image or the environment state among the inputs.")
 
-        if self.resize_shape is None and self.crop_shape is not None:
+        if self.crop_shape is not None:
             for key, image_ft in self.image_features.items():
                 if self.crop_shape[0] > image_ft.shape[1] or self.crop_shape[1] > image_ft.shape[2]:
                     raise ValueError(
-                        f"`crop_shape` should fit within the image shapes. Got {self.crop_shape} "
-                        f"for `crop_shape` and {image_ft.shape} for `{key}`."
+                        f"`crop_shape` should fit within the images shapes. Got {self.crop_shape} "
+                        f"for `crop_shape` and {image_ft.shape} for "
+                        f"`{key}`."
                     )
 
         # Check that all input images have the same shape.

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -454,18 +454,12 @@ class DiffusionRgbEncoder(nn.Module):
     def __init__(self, config: DiffusionConfig):
         super().__init__()
         # Set up optional preprocessing.
-        if config.resize_shape is not None:
-            self.resize = torchvision.transforms.Resize(config.resize_shape)
-        else:
-            self.resize = None
-
-        crop_shape = config.crop_shape
-        if crop_shape is not None:
+        if config.crop_shape is not None:
             self.do_crop = True
             # Always use center crop for eval
-            self.center_crop = torchvision.transforms.CenterCrop(crop_shape)
+            self.center_crop = torchvision.transforms.CenterCrop(config.crop_shape)
             if config.crop_is_random:
-                self.maybe_random_crop = torchvision.transforms.RandomCrop(crop_shape)
+                self.maybe_random_crop = torchvision.transforms.RandomCrop(config.crop_shape)
             else:
                 self.maybe_random_crop = self.center_crop
         else:
@@ -491,16 +485,13 @@ class DiffusionRgbEncoder(nn.Module):
 
         # Set up pooling and final layers.
         # Use a dry run to get the feature map shape.
-        # The dummy shape mirrors the runtime preprocessing order: resize -> crop.
+        # The dummy input should take the number of image channels from `config.image_features` and it should
+        # use the height and width from `config.crop_shape` if it is provided, otherwise it should use the
+        # height and width from `config.image_features`.
 
         # Note: we have a check in the config class to make sure all images have the same shape.
         images_shape = next(iter(config.image_features.values())).shape
-        if config.crop_shape is not None:
-            dummy_shape_h_w = config.crop_shape
-        elif config.resize_shape is not None:
-            dummy_shape_h_w = config.resize_shape
-        else:
-            dummy_shape_h_w = images_shape[1:]
+        dummy_shape_h_w = config.crop_shape if config.crop_shape is not None else images_shape[1:]
         dummy_shape = (1, images_shape[0], *dummy_shape_h_w)
         feature_map_shape = get_output_shape(self.backbone, dummy_shape)[1:]
 
@@ -516,10 +507,7 @@ class DiffusionRgbEncoder(nn.Module):
         Returns:
             (B, D) image feature.
         """
-        # Preprocess: resize if configured, then crop if configured.
-
-        if self.resize is not None:
-            x = self.resize(x)
+        # Preprocess: maybe crop (if it was set up in the __init__).
         if self.do_crop:
             if self.training:  # noqa: SIM108
                 x = self.maybe_random_crop(x)

src/lerobot/policies/smolvla/configuration_smolvla.py

@@ -106,9 +106,6 @@ class SmolVLAConfig(PreTrainedConfig):
     # Real-Time Chunking (RTC) configuration
     rtc_config: RTCConfig | None = None
 
-    compile_model: bool = False  # Whether to use torch.compile for model optimization
-    compile_mode: str = "max-autotune"  # Torch compile mode
-
     def __post_init__(self):
         super().__post_init__()
 

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -593,12 +593,6 @@ class VLAFlowMatching(nn.Module):
         self.prefix_length = self.config.prefix_length
         self.rtc_processor = rtc_processor
 
-        # Compile model if requested
-        if config.compile_model:
-            torch.set_float32_matmul_precision("high")
-            self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
-            self.forward = torch.compile(self.forward, mode=config.compile_mode)
-
     def _rtc_enabled(self):
         return self.config.rtc_config is not None and self.config.rtc_config.enabled
 

src/lerobot/policies/smolvla/smolvlm_with_expert.py

@@ -77,6 +77,7 @@ class SmolVLMWithExpertModel(nn.Module):
             print(f"Loading  {model_id} weights ...")
             self.vlm = AutoModelForImageTextToText.from_pretrained(
                 model_id,
+                device_map=device,
                 torch_dtype="bfloat16",
                 low_cpu_mem_usage=True,
             )

src/lerobot/scripts/lerobot_calibrate.py

@@ -56,7 +56,6 @@ from lerobot.teleoperators import (  # noqa: F401
     make_teleoperator_from_config,
     omx_leader,
     openarm_leader,
-    openarm_mini,
     so_leader,
     unitree_g1,
 )

src/lerobot/scripts/lerobot_find_joint_limits.py

@@ -61,7 +61,6 @@ from lerobot.teleoperators import (  # noqa: F401
     make_teleoperator_from_config,
     omx_leader,
     openarm_leader,
-    openarm_mini,
     so_leader,
 )
 from lerobot.utils.robot_utils import precise_sleep

src/lerobot/scripts/lerobot_mirror_dataset.py

@@ -0,0 +1,366 @@
+# 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.
+
+"""
+Mirror a bimanual robot dataset by swapping left/right arms and inverting joint values.
+
+This script creates a mirrored version of a dataset where:
+1. Left and right arm observations/actions are swapped
+2. Joint values are inverted according to a mirroring mask
+3. Video frames are horizontally flipped
+
+Example usage:
+```shell
+python -m lerobot.scripts.lerobot_mirror_dataset \
+    --repo_id=pepijn/openarm_bimanual \
+    --output_repo_id=pepijn/openarm_bimanual_mirrored
+```
+"""
+
+import argparse
+import logging
+import os
+import subprocess
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.datasets.utils import (
+    DATA_DIR,
+    DEFAULT_DATA_PATH,
+    write_info,
+    write_stats,
+    write_tasks,
+)
+from lerobot.utils.constants import HF_LEROBOT_HOME
+
+logger = logging.getLogger(__name__)
+
+OPENARM_MIRRORING_MASK = {
+    "joint_1": -1,  # Pan - invert
+    "joint_2": -1,  # Lift - invert
+    "joint_3": -1,  # Roll - invert
+    "joint_4": 1,   # Elbow - no invert
+    "joint_5": -1,  # W-Roll - invert
+    "joint_6": -1,  # W-Pitch - invert
+    "joint_7": -1,  # W-Yaw - invert
+    "gripper": 1,   # Gripper - no invert
+}
+
+
+def get_mirroring_mask(robot_type: str) -> dict[str, int]:
+    """Get the mirroring mask for a given robot type."""
+    if robot_type in ["bi_openarm_follower", "openarm_follower", "bi_openarms_follower", "openarms_follower"]:
+        return OPENARM_MIRRORING_MASK
+    raise ValueError(f"Unknown robot type: {robot_type}. Add a mirroring mask for this robot.")
+
+
+def swap_left_right_name(name: str) -> str:
+    """Swap 'left' and 'right' in a feature name."""
+    # Use placeholder to avoid double-swap
+    result = name.replace("left_", "LEFT_PLACEHOLDER_")
+    result = result.replace("right_", "left_")
+    result = result.replace("LEFT_PLACEHOLDER_", "right_")
+    return result
+
+
+def mirror_feature_names(names: list[str]) -> tuple[list[str], dict[int, int]]:
+    """Mirror feature names by swapping left/right and return the new names and index mapping."""
+    mirrored_names = [swap_left_right_name(n) for n in names]
+    old_to_new_idx = {}
+    for old_idx, old_name in enumerate(names):
+        new_name = swap_left_right_name(old_name)
+        new_idx = mirrored_names.index(new_name)
+        old_to_new_idx[old_idx] = new_idx
+    return mirrored_names, old_to_new_idx
+
+
+def apply_mirroring_mask(
+    value: float,
+    feature_name: str,
+    mirroring_mask: dict[str, int],
+) -> float:
+    """Apply mirroring mask to a joint value."""
+    name_without_prefix = feature_name.split("_", 1)[1] if "_" in feature_name else feature_name
+    joint_name = name_without_prefix.split(".")[0]
+    if joint_name in mirroring_mask:
+        return value * mirroring_mask[joint_name]
+    return value
+
+
+def mirror_array(
+    array: np.ndarray,
+    names: list[str],
+    mirroring_mask: dict[str, int],
+) -> np.ndarray:
+    """Mirror an array of values (action or state) by swapping left/right and applying mask."""
+    mirrored_names, idx_mapping = mirror_feature_names(names)
+    result = np.zeros_like(array)
+    for old_idx, new_idx in idx_mapping.items():
+        old_name = names[old_idx]
+        new_name = mirrored_names[new_idx]
+        value = array[old_idx]
+        mirrored_value = apply_mirroring_mask(value, new_name, mirroring_mask)
+        result[new_idx] = mirrored_value
+    return result
+
+
+def flip_video_frames(
+    input_path: Path,
+    output_path: Path,
+    fps: float,
+    vcodec: str = "libsvtav1",
+):
+    """Flip video frames horizontally using FFmpeg with same settings as encode_video_frames."""
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    cmd = [
+        "ffmpeg", "-y", "-i", str(input_path),
+        "-vf", "hflip",
+        "-c:v", vcodec,
+        "-g", "2",
+        "-crf", "30",
+        "-r", str(int(fps)),
+        "-pix_fmt", "yuv420p",
+        "-loglevel", "error",
+    ]
+    if vcodec == "libsvtav1":
+        cmd.extend(["-preset", "12"])
+    cmd.append(str(output_path))
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"FFmpeg failed: {result.stderr}")
+
+
+def mirror_dataset(
+    repo_id: str,
+    output_repo_id: str,
+    root: str | Path | None = None,
+    output_root: str | Path | None = None,
+    mirroring_mask: dict[str, int] | None = None,
+    vcodec: str = "libsvtav1",
+    num_workers: int | None = None,
+) -> LeRobotDataset:
+    """Mirror a bimanual robot dataset."""
+    logger.info(f"Loading dataset: {repo_id}")
+    dataset = LeRobotDataset(repo_id, root=root)
+
+    if mirroring_mask is None:
+        robot_type = dataset.meta.robot_type or "bi_openarms_follower"
+        mirroring_mask = get_mirroring_mask(robot_type)
+        logger.info(f"Using mirroring mask for robot type: {robot_type}")
+
+    output_root = Path(output_root) if output_root else HF_LEROBOT_HOME / output_repo_id
+
+    mirrored_features = {}
+    for key, feat in dataset.meta.features.items():
+        new_key = swap_left_right_name(key)
+        new_feat = feat.copy()
+        if "names" in new_feat and new_feat["names"]:
+            new_feat["names"] = [swap_left_right_name(n) for n in new_feat["names"]]
+        mirrored_features[new_key] = new_feat
+
+    logger.info("Creating mirrored dataset metadata...")
+    new_meta = LeRobotDatasetMetadata.create(
+        repo_id=output_repo_id,
+        fps=dataset.meta.fps,
+        features=mirrored_features,
+        robot_type=dataset.meta.robot_type,
+        root=output_root,
+        use_videos=len(dataset.meta.video_keys) > 0,
+    )
+
+    if dataset.meta.tasks is not None:
+        write_tasks(dataset.meta.tasks, new_meta.root)
+        new_meta.tasks = dataset.meta.tasks.copy()
+
+    _mirror_data(dataset, new_meta, mirroring_mask)
+    _mirror_videos(dataset, new_meta, vcodec, num_workers)
+    _copy_episodes_metadata(dataset, new_meta)
+
+    logger.info(f"Mirrored dataset saved to: {output_root}")
+    return LeRobotDataset(output_repo_id, root=output_root)
+
+
+def _mirror_data(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    mirroring_mask: dict[str, int],
+) -> None:
+    """Mirror parquet data files."""
+    data_dir = src_dataset.root / DATA_DIR
+    parquet_files = sorted(data_dir.glob("*/*.parquet"))
+
+    if not parquet_files:
+        raise ValueError(f"No parquet files found in {data_dir}")
+
+    action_names = src_dataset.meta.features.get("action", {}).get("names", [])
+    state_names = src_dataset.meta.features.get("observation.state", {}).get("names", [])
+
+    for src_path in tqdm(parquet_files, desc="Mirroring data files"):
+        df = pd.read_parquet(src_path).reset_index(drop=True)
+        relative_path = src_path.relative_to(src_dataset.root)
+        chunk_dir = relative_path.parts[1]
+        file_name = relative_path.parts[2]
+        chunk_idx = int(chunk_dir.split("-")[1])
+        file_idx = int(file_name.split("-")[1].split(".")[0])
+
+        if "action" in df.columns and action_names:
+            actions = np.stack(df["action"].values)
+            mirrored_actions = np.array([
+                mirror_array(row, action_names, mirroring_mask) for row in actions
+            ])
+            df["action"] = list(mirrored_actions)
+
+        if "observation.state" in df.columns and state_names:
+            states = np.stack(df["observation.state"].values)
+            mirrored_states = np.array([
+                mirror_array(row, state_names, mirroring_mask) for row in states
+            ])
+            df["observation.state"] = list(mirrored_states)
+
+        dst_path = dst_meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
+        dst_path.parent.mkdir(parents=True, exist_ok=True)
+        df.to_parquet(dst_path, index=False)
+
+
+def _mirror_videos(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    vcodec: str,
+    num_workers: int | None = None,
+) -> None:
+    """Mirror video files by flipping horizontally and swapping left/right names."""
+    if not src_dataset.meta.video_keys:
+        return
+
+    video_tasks = []
+    for old_video_key in src_dataset.meta.video_keys:
+        new_video_key = swap_left_right_name(old_video_key)
+        for ep_idx in range(src_dataset.meta.total_episodes):
+            try:
+                src_path = src_dataset.root / src_dataset.meta.get_video_file_path(ep_idx, old_video_key)
+                dst_relative = src_dataset.meta.get_video_file_path(ep_idx, old_video_key)
+                dst_relative_str = str(dst_relative).replace(old_video_key, new_video_key)
+                dst_path = dst_meta.root / dst_relative_str
+                if src_path.exists():
+                    video_tasks.append((src_path, dst_path))
+            except KeyError:
+                continue
+
+    def process_video(task, pbar):
+        src_path, dst_path = task
+        pbar.set_postfix_str(src_path.name)
+        flip_video_frames(src_path, dst_path, src_dataset.meta.fps, vcodec)
+        return src_path
+
+    if num_workers is None:
+        num_workers = os.cpu_count() or 16
+    num_workers = min(len(video_tasks), num_workers)
+    logger.info(f"Processing {len(video_tasks)} videos with {num_workers} workers")
+    with tqdm(total=len(video_tasks), desc="Mirroring videos") as pbar:
+        with ThreadPoolExecutor(max_workers=num_workers) as executor:
+            futures = {executor.submit(process_video, t, pbar): t for t in video_tasks}
+            for future in as_completed(futures):
+                task = futures[future]
+                future.result()
+                pbar.set_postfix_str(f"done: {task[0].name}")
+                pbar.update(1)
+
+
+def _copy_episodes_metadata(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+) -> None:
+    """Copy episodes metadata with swapped video keys."""
+    episodes_dir = src_dataset.root / "meta/episodes"
+    dst_episodes_dir = dst_meta.root / "meta/episodes"
+
+    if episodes_dir.exists():
+        dst_episodes_dir.mkdir(parents=True, exist_ok=True)
+        for src_parquet in episodes_dir.glob("*/*.parquet"):
+            df = pd.read_parquet(src_parquet)
+            columns_to_rename = {}
+            for col in df.columns:
+                if col.startswith("videos/"):
+                    parts = col.split("/")
+                    if len(parts) >= 2:
+                        video_key = parts[1]
+                        new_video_key = swap_left_right_name(video_key)
+                        new_col = col.replace(f"videos/{video_key}/", f"videos/{new_video_key}/")
+                        columns_to_rename[col] = new_col
+            if columns_to_rename:
+                df = df.rename(columns=columns_to_rename)
+            dst_parquet = dst_episodes_dir / src_parquet.relative_to(episodes_dir)
+            dst_parquet.parent.mkdir(parents=True, exist_ok=True)
+            df.to_parquet(dst_parquet, index=False)
+
+    dst_meta.info.update({
+        "total_episodes": src_dataset.meta.total_episodes,
+        "total_frames": src_dataset.meta.total_frames,
+        "total_tasks": src_dataset.meta.total_tasks,
+        "total_videos": src_dataset.meta.total_videos,
+        "total_chunks": src_dataset.meta.total_chunks,
+    })
+    write_info(dst_meta.info, dst_meta.root)
+
+    if src_dataset.meta.stats is not None:
+        mirrored_stats = _mirror_stats(src_dataset.meta.stats)
+        write_stats(mirrored_stats, dst_meta.root)
+
+
+def _mirror_stats(stats: dict) -> dict:
+    """Mirror stats by swapping left/right in feature names."""
+    mirrored = {}
+    for key, value in stats.items():
+        new_key = swap_left_right_name(key)
+        if isinstance(value, dict):
+            mirrored[new_key] = _mirror_stats(value)
+        else:
+            mirrored[new_key] = value
+    return mirrored
+
+
+def main():
+    logging.basicConfig(level=logging.INFO)
+    parser = argparse.ArgumentParser(description="Mirror a bimanual robot dataset")
+    parser.add_argument("--repo_id", type=str, required=True, help="Source dataset repo_id")
+    parser.add_argument("--output_repo_id", type=str, required=True, help="Output dataset repo_id")
+    parser.add_argument("--root", type=str, default=None, help="Source dataset root directory")
+    parser.add_argument("--output_root", type=str, default=None, help="Output dataset root directory")
+    parser.add_argument("--vcodec", type=str, default="libsvtav1", help="Video codec (libsvtav1, h264, hevc)")
+    parser.add_argument("--num_workers", type=int, default=None, help="Number of parallel workers for video processing")
+    parser.add_argument("--push-to-hub", action="store_true", help="Push mirrored dataset to HuggingFace Hub")
+    args = parser.parse_args()
+
+    dataset = mirror_dataset(
+        repo_id=args.repo_id,
+        output_repo_id=args.output_repo_id,
+        root=args.root,
+        output_root=args.output_root,
+        vcodec=args.vcodec,
+        num_workers=args.num_workers,
+    )
+
+    if getattr(args, "push_to_hub", False):
+        logger.info(f"Pushing dataset to HuggingFace Hub: {args.output_repo_id}")
+        dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()
+

src/lerobot/scripts/lerobot_record.py

@@ -125,7 +125,6 @@ from lerobot.teleoperators import (  # noqa: F401
     make_teleoperator_from_config,
     omx_leader,
     openarm_leader,
-    openarm_mini,
     reachy2_teleoperator,
     so_leader,
     unitree_g1,
@@ -155,7 +154,7 @@ class DatasetRecordConfig:
     repo_id: str
     # A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.")
     single_task: str
-    # Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
+    # Root directory where the dataset will be stored (e.g. 'dataset/path').
     root: str | Path | None = None
     # Limit the frames per second.
     fps: int = 30
@@ -334,7 +333,6 @@ def record_loop(
         preprocessor.reset()
         postprocessor.reset()
 
-    no_action_count = 0
     timestamp = 0
     start_episode_t = time.perf_counter()
     while timestamp < control_time_s:
@@ -382,13 +380,11 @@ def record_loop(
             act = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
             act_processed_teleop = teleop_action_processor((act, obs))
         else:
-            no_action_count += 1
-            if no_action_count == 1 or no_action_count % 10 == 0:
-                logging.warning(
-                    "No policy or teleoperator provided, skipping action generation. "
-                    "This is likely to happen when resetting the environment without a teleop device. "
-                    "The robot won't be at its rest position at the start of the next episode."
-                )
+            logging.info(
+                "No policy or teleoperator provided, skipping action generation."
+                "This is likely to happen when resetting the environment without a teleop device."
+                "The robot won't be at its rest position at the start of the next episode."
+            )
             continue
 
         # Applies a pipeline to the action, default is IdentityProcessor

src/lerobot/scripts/lerobot_replay.py

@@ -80,7 +80,7 @@ class DatasetReplayConfig:
     repo_id: str
     # Episode to replay.
     episode: int
-    # Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
+    # Root directory where the dataset will be stored (e.g. 'dataset/path').
     root: str | Path | None = None
     # Limit the frames per second. By default, uses the policy fps.
     fps: int = 30

src/lerobot/scripts/lerobot_teleoperate.py

@@ -94,7 +94,6 @@ from lerobot.teleoperators import (  # noqa: F401
     make_teleoperator_from_config,
     omx_leader,
     openarm_leader,
-    openarm_mini,
     reachy2_teleoperator,
     so_leader,
     unitree_g1,

src/lerobot/scripts/lerobot_train.py

@@ -38,6 +38,7 @@ from lerobot.policies.factory import make_policy, make_pre_post_processors
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.rl.wandb_utils import WandBLogger
 from lerobot.scripts.lerobot_eval import eval_policy_all
+from lerobot.teleoperators import openarm_mini  # noqa: F401
 from lerobot.utils.import_utils import register_third_party_plugins
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 from lerobot.utils.random_utils import set_seed
@@ -380,10 +381,10 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
         "dataloading_s": AverageMeter("data_s", ":.3f"),
     }
 
-    # Keep global batch size for logging; MetricsTracker handles world size internally.
+    # Use effective batch size for proper epoch calculation in distributed training
     effective_batch_size = cfg.batch_size * accelerator.num_processes
     train_tracker = MetricsTracker(
-        cfg.batch_size,
+        effective_batch_size,
         dataset.num_frames,
         dataset.num_episodes,
         train_metrics,

src/lerobot/utils/logging_utils.py

@@ -104,10 +104,9 @@ class MetricsTracker:
         self.metrics = metrics
 
         self.steps = initial_step
-        world_size = accelerator.num_processes if accelerator else 1
         # A sample is an (observation,action) pair, where observation and action
         # can be on multiple timestamps. In a batch, we have `batch_size` number of samples.
-        self.samples = self.steps * self._batch_size * world_size
+        self.samples = self.steps * self._batch_size
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
         self.accelerator = accelerator
@@ -133,8 +132,7 @@ class MetricsTracker:
         Updates metrics that depend on 'step' for one step.
         """
         self.steps += 1
-        world_size = self.accelerator.num_processes if self.accelerator else 1
-        self.samples += self._batch_size * world_size
+        self.samples += self._batch_size * (self.accelerator.num_processes if self.accelerator else 1)
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
 

tests/artifacts/policies/pusht_diffusion_/actions.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:54aecbc1af72a4cd5e9261492f5e7601890517516257aacdf2a0ffb3ce281f1b
+oid sha256:19eaaa85f66ba4aa6388dbb83819ffad6ea4363247208f871a8dc385689f6fc8
 size 992

tests/artifacts/policies/pusht_diffusion_/grad_stats.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:88a9c3775a2aa1e90a08850521970070a4fcf0f6b82aab43cd8ccc5cf77e0013
+oid sha256:227296eaeeb54acdc3dae2eb8af3d4d08fb87e245337624447140b1e91cfd002
 size 47424

tests/artifacts/policies/pusht_diffusion_/output_dict.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:91a2635e05a75fe187a5081504c5f35ce3417378813fa2deaf9ca4e8200e1819
+oid sha256:271b00cb2f0cd5fd26b1d53463638e3d1a6e92692ec625fcffb420ca190869e5
 size 68

tests/artifacts/policies/pusht_diffusion_/param_stats.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:645bff922ac7bea63ad018ebf77c303c0e4cd2c1c0dc5ef3192865281bef3dc6
+oid sha256:778fddbbaa64248cee35cb377c02cc2b6076f7ce5855146de677128900617ddf
 size 47424

tests/utils/test_logging_utils.py

@@ -24,11 +24,6 @@ def mock_metrics():
     return {"loss": AverageMeter("loss", ":.3f"), "accuracy": AverageMeter("accuracy", ":.2f")}
 
 
-class MockAccelerator:
-    def __init__(self, num_processes: int):
-        self.num_processes = num_processes
-
-
 def test_average_meter_initialization():
     meter = AverageMeter("loss", ":.2f")
     assert meter.name == "loss"
@@ -87,37 +82,6 @@ def test_metrics_tracker_step(mock_metrics):
     assert tracker.epochs == tracker.samples / 1000
 
 
-def test_metrics_tracker_initialization_with_accelerator(mock_metrics):
-    tracker = MetricsTracker(
-        batch_size=32,
-        num_frames=1000,
-        num_episodes=50,
-        metrics=mock_metrics,
-        initial_step=10,
-        accelerator=MockAccelerator(num_processes=2),
-    )
-    assert tracker.steps == 10
-    assert tracker.samples == 10 * 32 * 2
-    assert tracker.episodes == tracker.samples / (1000 / 50)
-    assert tracker.epochs == tracker.samples / 1000
-
-
-def test_metrics_tracker_step_with_accelerator(mock_metrics):
-    tracker = MetricsTracker(
-        batch_size=32,
-        num_frames=1000,
-        num_episodes=50,
-        metrics=mock_metrics,
-        initial_step=5,
-        accelerator=MockAccelerator(num_processes=2),
-    )
-    tracker.step()
-    assert tracker.steps == 6
-    assert tracker.samples == (5 * 32 * 2) + (32 * 2)
-    assert tracker.episodes == tracker.samples / (1000 / 50)
-    assert tracker.epochs == tracker.samples / 1000
-
-
 def test_metrics_tracker_getattr(mock_metrics):
     tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=mock_metrics)
     assert tracker.loss == mock_metrics["loss"]