fix(edit-dataset): align root with LeRobotDataset v2 semantics

* chore(docstrings): fixing deprecated `root` argument docstrings in LeRobotDataset class

* chore(draccus): updating draccus CLI help

* chore(revert): reverting changes in lerobot_dataset_viz.py

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

AI_POLICY.md

@@ -0,0 +1,25 @@
+# AI Usage Policy
+
+The LeRobot project welcomes contributions from everyone, and we have a few guidelines regarding AI usage to ensure high code quality, clear communication, and a healthy open-source ecosystem:
+
+- **Please disclose significant AI assistance.** If you used AI tools (e.g., Copilot, Claude, Cursor, ChatGPT) to generate a substantial portion of your code or text, let us know in your PR description. Transparency helps us review your changes more effectively.
+- **Own your code (The Human-in-the-Loop).** You must fully understand all the changes you are proposing. If you cannot explain what your AI-assisted code does or how it interacts with LeRobot's broader architecture, please take the time to learn and test it before submitting.
+- **Keep issues and discussions focused.** You are welcome to use AI to help draft issues or PR descriptions, but please review and edit them carefully before posting. AI can often be overly verbose; trimming the noise and getting straight to the point helps our maintainers address your needs faster.
+
+Our core maintainers also use AI tools to aid their workflows, but they do so while bringing deep contextual knowledge of the LeRobot codebase to validate the output. We ask all contributors to apply that same level of rigor.
+
+## Remember the Human Maintainers
+
+Please remember that LeRobot is maintained by a dedicated team of humans.
+
+Every discussion, issue, and pull request is read and reviewed by real people. While AI tools can generate thousands of lines of code in seconds, reviewing that code still takes human time and energy. Submitting unverified or low-effort AI output puts an unfair burden on our maintainers.
+
+Today, the quality of the AI output still heavily depends on the developer driving the tool. We ask that you respect our maintainers' time by thoroughly vetting, testing, and refining your submissions.
+
+## AI is Welcome Here
+
+LeRobot operates at the cutting edge of AI and robotics, and many of our maintainers actively embrace AI coding assistants as valuable productivity tools. We are a pro-AI project!
+
+Our reason for having an AI policy is not an anti-AI stance. Rather, it exists to ensure that AI is used to enhance human contributions, not replace them with unverified noise. It's about how the tools are used, not the tools themselves.
+
+We value the unique human insight you bring to the LeRobot community. Let AI empower your workflow, but always let your own judgment take the wheel.

CONTRIBUTING.md

@@ -2,7 +2,7 @@
 
 Everyone is welcome to contribute, and we value everybody's contribution. Code is not the only way to help the community. Answering questions, helping others, reaching out, and improving the documentation are immensely valuable.
 
-Whichever way you choose to contribute, please be mindful to respect our [code of conduct](./CODE_OF_CONDUCT.md).
+Whichever way you choose to contribute, please be mindful to respect our [code of conduct](./CODE_OF_CONDUCT.md) and our [AI policy](./AI_POLICY.md).
 
 ## Ways to Contribute
 

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').
+    # Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
     root: str | Path | None = None
     # Limit the frames per second. By default, uses the policy fps.
     fps: int = 30

pyproject.toml

@@ -25,7 +25,7 @@ discord = "https://discord.gg/s3KuuzsPFb"
 
 [project]
 name = "lerobot"
-version = "0.4.4"
+version = "0.4.5"
 description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
 dynamic = ["readme"]
 license = { text = "Apache-2.0" }

src/lerobot/async_inference/robot_client.py

@@ -49,23 +49,18 @@ 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 (  # noqa: F401
-    Robot,
-    RobotConfig,
-    bi_so_follower,
-    koch_follower,
+from lerobot.robots import (
+    RobotConfig,  # noqa: F401
     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,
@@ -485,8 +480,9 @@ class RobotClient:
 def async_client(cfg: RobotClientConfig):
     logging.info(pformat(asdict(cfg)))
 
-    if cfg.robot.type not in SUPPORTED_ROBOTS:
-        raise ValueError(f"Robot {cfg.robot.type} not yet supported!")
+    # 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!")
 
     client = RobotClient(cfg)
 
@@ -512,4 +508,5 @@ 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').
+    # Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
     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. 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'.
+            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).
             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,11 +1771,12 @@ 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)
-            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)
+            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)
 
         self.image_transforms = image_transforms
         self.delta_timestamps = delta_timestamps

src/lerobot/policies/diffusion/configuration_diffusion.py

@@ -55,10 +55,16 @@ 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.
-        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).
+        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).
         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.
@@ -114,7 +120,9 @@ class DiffusionConfig(PreTrainedConfig):
     # Architecture / modeling.
     # Vision backbone.
     vision_backbone: str = "resnet18"
-    crop_shape: tuple[int, int] | None = (84, 84)
+    resize_shape: tuple[int, int] | None = None
+    crop_ratio: float = 1.0
+    crop_shape: tuple[int, int] | None = None
     crop_is_random: bool = True
     pretrained_backbone_weights: str | None = None
     use_group_norm: bool = True
@@ -175,6 +183,25 @@ 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)
@@ -202,13 +229,12 @@ 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.crop_shape is not None:
+        if self.resize_shape is None and 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 images shapes. Got {self.crop_shape} "
-                        f"for `crop_shape` and {image_ft.shape} for "
-                        f"`{key}`."
+                        f"`crop_shape` should fit within the image shapes. Got {self.crop_shape} "
+                        f"for `crop_shape` and {image_ft.shape} for `{key}`."
                     )
 
         # Check that all input images have the same shape.

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -454,12 +454,18 @@ class DiffusionRgbEncoder(nn.Module):
     def __init__(self, config: DiffusionConfig):
         super().__init__()
         # Set up optional preprocessing.
-        if config.crop_shape is not None:
+        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:
             self.do_crop = True
             # Always use center crop for eval
-            self.center_crop = torchvision.transforms.CenterCrop(config.crop_shape)
+            self.center_crop = torchvision.transforms.CenterCrop(crop_shape)
             if config.crop_is_random:
-                self.maybe_random_crop = torchvision.transforms.RandomCrop(config.crop_shape)
+                self.maybe_random_crop = torchvision.transforms.RandomCrop(crop_shape)
             else:
                 self.maybe_random_crop = self.center_crop
         else:
@@ -485,13 +491,16 @@ class DiffusionRgbEncoder(nn.Module):
 
         # Set up pooling and final layers.
         # Use a dry run to get the feature map shape.
-        # 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`.
+        # The dummy shape mirrors the runtime preprocessing order: resize -> crop.
 
         # 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
-        dummy_shape_h_w = config.crop_shape if config.crop_shape is not None else images_shape[1:]
+        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 = (1, images_shape[0], *dummy_shape_h_w)
         feature_map_shape = get_output_shape(self.backbone, dummy_shape)[1:]
 
@@ -507,7 +516,10 @@ class DiffusionRgbEncoder(nn.Module):
         Returns:
             (B, D) image feature.
         """
-        # Preprocess: maybe crop (if it was set up in the __init__).
+        # Preprocess: resize if configured, then crop if configured.
+
+        if self.resize is not None:
+            x = self.resize(x)
         if self.do_crop:
             if self.training:  # noqa: SIM108
                 x = self.maybe_random_crop(x)

src/lerobot/policies/smolvla/configuration_smolvla.py

@@ -106,6 +106,9 @@ 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,6 +593,12 @@ 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,7 +77,6 @@ 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,6 +56,7 @@ from lerobot.teleoperators import (  # noqa: F401
     make_teleoperator_from_config,
     omx_leader,
     openarm_leader,
+    openarm_mini,
     so_leader,
     unitree_g1,
 )

src/lerobot/scripts/lerobot_edit_dataset.py

@@ -19,7 +19,9 @@ Edit LeRobot datasets using various transformation tools.
 
 This script allows you to delete episodes, split datasets, merge datasets,
 remove features, modify tasks, and convert image datasets to video format.
-When new_repo_id is specified, creates a new dataset.
+
+Path semantics (v2): --root and --new_root are exact dataset folders containing
+meta/, data/, videos/. When omitted, defaults to $HF_LEROBOT_HOME/{repo_id}.
 
 Usage Examples:
 
@@ -29,10 +31,18 @@ Delete episodes 0, 2, and 5 from a dataset:
         --operation.type delete_episodes \
         --operation.episode_indices "[0, 2, 5]"
 
-Delete episodes and save to a new dataset:
+Delete episodes from a local dataset at a specific path:
+    lerobot-edit-dataset \
+        --repo_id lerobot/pusht \
+        --root /data/lerobot/pusht \
+        --operation.type delete_episodes \
+        --operation.episode_indices "[0, 2, 5]"
+
+Delete episodes and save to a new dataset at a specific path:
     lerobot-edit-dataset \
         --repo_id lerobot/pusht \
         --new_repo_id lerobot/pusht_filtered \
+        --new_root /data/lerobot/pusht_filtered \
         --operation.type delete_episodes \
         --operation.episode_indices "[0, 2, 5]"
 
@@ -42,6 +52,13 @@ Split dataset by fractions:
         --operation.type split \
         --operation.splits '{"train": 0.8, "val": 0.2}'
 
+Split dataset by fractions with custom output base directory:
+    lerobot-edit-dataset \
+        --repo_id lerobot/pusht \
+        --new_root /data/splits \
+        --operation.type split \
+        --operation.splits '{"train": 0.8, "val": 0.2}'
+
 Split dataset by episode indices:
     lerobot-edit-dataset \
         --repo_id lerobot/pusht \
@@ -54,12 +71,27 @@ Split into more than two splits:
         --operation.type split \
         --operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'
 
-Merge multiple datasets:
+Merge multiple datasets (from Hub or default cache):
     lerobot-edit-dataset \
         --repo_id lerobot/pusht_merged \
         --operation.type merge \
         --operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
 
+Merge local datasets with explicit paths:
+    lerobot-edit-dataset \
+        --repo_id lerobot/pusht_merged \
+        --operation.type merge \
+        --operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']" \
+        --operation.roots "['/data/pusht_train', '/data/pusht_val']"
+
+Merge local datasets with custom output path:
+    lerobot-edit-dataset \
+        --repo_id lerobot/pusht_merged \
+        --new_root /output/merged \
+        --operation.type merge \
+        --operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']" \
+        --operation.roots "['/data/pusht_train', '/data/pusht_val']"
+
 Remove camera feature:
     lerobot-edit-dataset \
         --repo_id lerobot/pusht \
@@ -88,8 +120,8 @@ Modify tasks - set default task with overrides for specific episodes (WARNING: m
 Convert image dataset to video format and save locally:
     lerobot-edit-dataset \
         --repo_id lerobot/pusht_image \
-        --operation.type convert_image_to_video \
-        --operation.output_dir /path/to/output/pusht_video
+        --new_root /path/to/output/pusht_video \
+        --operation.type convert_image_to_video
 
 Convert image dataset to video format and save with new repo_id:
     lerobot-edit-dataset \
@@ -167,6 +199,8 @@ class SplitConfig(OperationConfig):
 @dataclass
 class MergeConfig(OperationConfig):
     repo_ids: list[str] | None = None
+    # Per-dataset exact folder paths. Must match repo_ids length when provided.
+    roots: list[str] | None = None
 
 
 @OperationConfig.register_subclass("remove_feature")
@@ -208,26 +242,35 @@ class InfoConfig(OperationConfig):
 class EditDatasetConfig:
     repo_id: str
     operation: OperationConfig
+    # Exact dataset folder (containing meta/, data/, videos/). If None, defaults to HF_LEROBOT_HOME/{repo_id}.
     root: str | None = None
     new_repo_id: str | None = None
+    # Exact output dataset folder. If None, derived from new_repo_id or in-place.
+    new_root: str | None = None
     push_to_hub: bool = False
 
 
-def get_output_path(repo_id: str, new_repo_id: str | None, root: Path | None) -> tuple[str, Path]:
-    if new_repo_id:
+def get_output_path(
+    input_root: Path,
+    new_repo_id: str | None,
+    new_root: Path | None,
+    default_repo_id: str,
+) -> tuple[str, Path]:
+    if new_root:
+        output_repo_id = new_repo_id or default_repo_id
+        output_dir = new_root
+    elif new_repo_id:
         output_repo_id = new_repo_id
-        output_dir = root / new_repo_id if root else HF_LEROBOT_HOME / new_repo_id
+        output_dir = HF_LEROBOT_HOME / new_repo_id
     else:
-        output_repo_id = repo_id
-        dataset_path = root / repo_id if root else HF_LEROBOT_HOME / repo_id
-        old_path = Path(str(dataset_path) + "_old")
+        output_repo_id = default_repo_id
+        output_dir = input_root
+        backup_path = input_root.with_name(input_root.name + "_old")
 
-        if dataset_path.exists():
-            if old_path.exists():
-                shutil.rmtree(old_path)
-            shutil.move(str(dataset_path), str(old_path))
-
-        output_dir = dataset_path
+        if input_root.exists():
+            if backup_path.exists():
+                shutil.rmtree(backup_path)
+            shutil.move(str(input_root), str(backup_path))
 
     return output_repo_id, output_dir
 
@@ -241,11 +284,14 @@ def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
 
     dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
     output_repo_id, output_dir = get_output_path(
-        cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
+        input_root=dataset.root,
+        new_repo_id=cfg.new_repo_id,
+        new_root=Path(cfg.new_root) if cfg.new_root else None,
+        default_repo_id=cfg.repo_id,
     )
 
-    if cfg.new_repo_id is None:
-        dataset.root = Path(str(dataset.root) + "_old")
+    if output_dir == dataset.root:
+        dataset.root = dataset.root.with_name(dataset.root.name + "_old")
 
     logging.info(f"Deleting episodes {cfg.operation.episode_indices} from {cfg.repo_id}")
     new_dataset = delete_episodes(
@@ -275,16 +321,20 @@ def handle_split(cfg: EditDatasetConfig) -> None:
     dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
 
     logging.info(f"Splitting dataset {cfg.repo_id} with splits: {cfg.operation.splits}")
-    split_datasets = split_dataset(dataset, splits=cfg.operation.splits)
+    split_output_dir = Path(cfg.new_root) if cfg.new_root else None
+    split_datasets = split_dataset(
+        dataset,
+        splits=cfg.operation.splits,
+        output_dir=split_output_dir,
+    )
 
     for split_name, split_ds in split_datasets.items():
-        split_repo_id = f"{cfg.repo_id}_{split_name}"
         logging.info(
             f"{split_name}: {split_ds.meta.total_episodes} episodes, {split_ds.meta.total_frames} frames"
         )
 
         if cfg.push_to_hub:
-            logging.info(f"Pushing {split_name} split to hub as {split_repo_id}")
+            logging.info(f"Pushing {split_name} split to hub as {split_ds.repo_id}")
             LeRobotDataset(split_ds.repo_id, root=split_ds.root).push_to_hub()
 
 
@@ -298,10 +348,19 @@ def handle_merge(cfg: EditDatasetConfig) -> None:
     if not cfg.repo_id:
         raise ValueError("repo_id must be specified as the output repository for merged dataset")
 
-    logging.info(f"Loading {len(cfg.operation.repo_ids)} datasets to merge")
-    datasets = [LeRobotDataset(repo_id, root=cfg.root) for repo_id in cfg.operation.repo_ids]
+    if cfg.operation.roots:
+        if len(cfg.operation.roots) != len(cfg.operation.repo_ids):
+            raise ValueError("operation.roots and operation.repo_ids must have the same length")
+        logging.info(f"Loading {len(cfg.operation.roots)} local datasets to merge")
+        datasets = [
+            LeRobotDataset(repo_id, root=root)
+            for repo_id, root in zip(cfg.operation.repo_ids, cfg.operation.roots, strict=True)
+        ]
+    else:
+        logging.info(f"Loading {len(cfg.operation.repo_ids)} datasets to merge")
+        datasets = [LeRobotDataset(repo_id) for repo_id in cfg.operation.repo_ids]
 
-    output_dir = Path(cfg.root) / cfg.repo_id if cfg.root else HF_LEROBOT_HOME / cfg.repo_id
+    output_dir = Path(cfg.new_root) if cfg.new_root else HF_LEROBOT_HOME / cfg.repo_id
 
     logging.info(f"Merging datasets into {cfg.repo_id}")
     merged_dataset = merge_datasets(
@@ -329,11 +388,14 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
 
     dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
     output_repo_id, output_dir = get_output_path(
-        cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
+        input_root=dataset.root,
+        new_repo_id=cfg.new_repo_id,
+        new_root=Path(cfg.new_root) if cfg.new_root else None,
+        default_repo_id=cfg.repo_id,
     )
 
-    if cfg.new_repo_id is None:
-        dataset.root = Path(str(dataset.root) + "_old")
+    if output_dir == dataset.root:
+        dataset.root = dataset.root.with_name(dataset.root.name + "_old")
 
     logging.info(f"Removing features {cfg.operation.feature_names} from {cfg.repo_id}")
     new_dataset = remove_feature(
@@ -361,9 +423,10 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
     if new_task is None and episode_tasks_raw is None:
         raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
 
-    # Warn about in-place modification behavior
-    if cfg.new_repo_id is not None:
-        logging.warning("modify_tasks modifies datasets in-place. The --new_repo_id parameter is ignored.")
+    if cfg.new_repo_id is not None or cfg.new_root is not None:
+        logging.warning(
+            "modify_tasks modifies datasets in-place. The --new_repo_id and --new_root parameters are ignored."
+        )
 
     dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
     logging.warning(f"Modifying dataset in-place at {dataset.root}. Original data will be overwritten.")
@@ -399,31 +462,24 @@ def handle_convert_image_to_video(cfg: EditDatasetConfig) -> None:
     dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
 
     # Determine output directory and repo_id
-    # Priority: 1) new_repo_id, 2) operation.output_dir, 3) auto-generated name
+    # Priority: 1) new_root, 2) new_repo_id, 3) operation.output_dir, 4) auto-generated name
     output_dir_config = getattr(cfg.operation, "output_dir", None)
 
-    if cfg.new_repo_id:
-        # Use new_repo_id for both local storage and hub push
+    if cfg.new_root:
+        output_dir = Path(cfg.new_root)
+        output_repo_id = cfg.new_repo_id or f"{cfg.repo_id}_video"
+        logging.info(f"Saving to {output_dir} as {output_repo_id}")
+    elif cfg.new_repo_id:
         output_repo_id = cfg.new_repo_id
-        # Place new dataset as a sibling to the original dataset
-        # Get the parent of the actual dataset root (not cfg.root which might be the lerobot cache dir)
-        # Extract just the dataset name (after last slash) for the local directory
-        local_dir_name = cfg.new_repo_id.split("/")[-1]
-        output_dir = dataset.root.parent / local_dir_name
+        output_dir = HF_LEROBOT_HOME / cfg.new_repo_id
         logging.info(f"Saving to new dataset: {cfg.new_repo_id} at {output_dir}")
     elif output_dir_config:
-        # Use custom output directory for local-only storage
         output_dir = Path(output_dir_config)
-        # Extract repo name from output_dir for the dataset
         output_repo_id = output_dir.name
         logging.info(f"Saving to local directory: {output_dir}")
     else:
-        # Auto-generate name: append "_video" to original repo_id
         output_repo_id = f"{cfg.repo_id}_video"
-        # Place new dataset as a sibling to the original dataset
-        # Extract just the dataset name (after last slash) for the local directory
-        local_dir_name = output_repo_id.split("/")[-1]
-        output_dir = dataset.root.parent / local_dir_name
+        output_dir = HF_LEROBOT_HOME / output_repo_id
         logging.info(f"Saving to auto-generated location: {output_dir}")
 
     logging.info(f"Converting dataset {cfg.repo_id} to video format")

src/lerobot/scripts/lerobot_find_joint_limits.py

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

@@ -1,366 +0,0 @@
-# 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,6 +125,7 @@ from lerobot.teleoperators import (  # noqa: F401
     make_teleoperator_from_config,
     omx_leader,
     openarm_leader,
+    openarm_mini,
     reachy2_teleoperator,
     so_leader,
     unitree_g1,
@@ -154,7 +155,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').
+    # Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
     root: str | Path | None = None
     # Limit the frames per second.
     fps: int = 30
@@ -333,6 +334,7 @@ def record_loop(
         preprocessor.reset()
         postprocessor.reset()
 
+    no_action_count = 0
     timestamp = 0
     start_episode_t = time.perf_counter()
     while timestamp < control_time_s:
@@ -380,11 +382,13 @@ 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:
-            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."
-            )
+            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."
+                )
             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').
+    # Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
     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,6 +94,7 @@ 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,7 +38,6 @@ 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
@@ -381,10 +380,10 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
         "dataloading_s": AverageMeter("data_s", ":.3f"),
     }
 
-    # Use effective batch size for proper epoch calculation in distributed training
+    # Keep global batch size for logging; MetricsTracker handles world size internally.
     effective_batch_size = cfg.batch_size * accelerator.num_processes
     train_tracker = MetricsTracker(
-        effective_batch_size,
+        cfg.batch_size,
         dataset.num_frames,
         dataset.num_episodes,
         train_metrics,

src/lerobot/utils/logging_utils.py

@@ -104,9 +104,10 @@ 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
+        self.samples = self.steps * self._batch_size * world_size
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
         self.accelerator = accelerator
@@ -132,7 +133,8 @@ class MetricsTracker:
         Updates metrics that depend on 'step' for one step.
         """
         self.steps += 1
-        self.samples += self._batch_size * (self.accelerator.num_processes if self.accelerator else 1)
+        world_size = self.accelerator.num_processes if self.accelerator else 1
+        self.samples += self._batch_size * world_size
         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:19eaaa85f66ba4aa6388dbb83819ffad6ea4363247208f871a8dc385689f6fc8
+oid sha256:54aecbc1af72a4cd5e9261492f5e7601890517516257aacdf2a0ffb3ce281f1b
 size 992

tests/artifacts/policies/pusht_diffusion_/grad_stats.safetensors

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

tests/artifacts/policies/pusht_diffusion_/output_dict.safetensors

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

tests/artifacts/policies/pusht_diffusion_/param_stats.safetensors

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

tests/utils/test_logging_utils.py

@@ -24,6 +24,11 @@ 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"
@@ -82,6 +87,37 @@ 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"]