more changres

docs/source/_toctree.yml

@@ -92,6 +92,10 @@
   - local: phone_teleop
     title: Phone
   title: "Teleoperators"
+- sections:
+  - local: torch_accelerators
+    title: PyTorch accelerators
+  title: "Supported Hardware"
 - sections:
   - local: notebooks
     title: Notebooks

docs/source/torch_accelerators.mdx

@@ -0,0 +1,42 @@
+# PyTorch accelerators
+
+LeRobot supports multiple hardware acceleration options for both training and inference.
+
+These options include:
+
+- **CPU**: CPU executes all computations, no dedicated accelerator is used
+- **CUDA**: acceleration with NVIDIA & AMD GPUs
+- **MPS**: acceleration with Apple Silicon GPUs
+- **XPU**: acceleration with Intel integrated and discrete GPUs
+
+## Getting Started
+
+To use particular accelerator, a suitable version of PyTorch should be installed.
+
+For CPU, CUDA, and MPS backends follow instructions provided on [PyTorch installation page](https://pytorch.org/get-started/locally).
+For XPU backend, follow instructions from [PyTorch documentation](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html).
+
+### Verifying the installation
+
+After installation, accelerator availability can be verified by running
+
+```python
+import torch
+print(torch.<backend_name>.is_available())  # <backend_name> is cuda, mps, or xpu
+```
+
+## How to run training or evaluation
+
+To select the desired accelerator, use the `--policy.device` flag when running `lerobot-train` or `lerobot-eval`. For example, to use MPS on Apple Silicon, run:
+
+```bash
+lerobot-train
+    --policy.device=mps ...
+```
+
+```bash
+lerobot-eval \
+    --policy.device=mps ...
+```
+
+However, in most cases, presence of an accelerator is detected automatically and `policy.device` parameter can be omitted from CLI commands.

docs/source/using_dataset_tools.mdx

@@ -11,13 +11,14 @@ LeRobot provides several utilities for manipulating datasets:
 3. **Merge Datasets** - Combine multiple datasets into one. The datasets must have identical features, and episodes are concatenated in the order specified in `repo_ids`
 4. **Add Features** - Add new features to a dataset
 5. **Remove Features** - Remove features from a dataset
+6. **Convert to Video** - Convert image-based datasets to video format for efficient storage
 
 The core implementation is in `lerobot.datasets.dataset_tools`.
 An example script detailing how to use the tools API is available in `examples/dataset/use_dataset_tools.py`.
 
 ## Command-Line Tool: lerobot-edit-dataset
 
-`lerobot-edit-dataset` is a command-line script for editing datasets. It can be used to delete episodes, split datasets, merge datasets, add features, and remove features.
+`lerobot-edit-dataset` is a command-line script for editing datasets. It can be used to delete episodes, split datasets, merge datasets, add features, remove features, and convert image datasets to video format.
 
 Run `lerobot-edit-dataset --help` for more information on the configuration of each operation.
 
@@ -86,9 +87,71 @@ lerobot-edit-dataset \
     --operation.feature_names "['observation.images.top']"
 ```
 
+#### Convert to Video
+
+Convert an image-based dataset to video format, creating a new LeRobotDataset where images are stored as videos. This is useful for reducing storage requirements and improving data loading performance. The new dataset will have the exact same structure as the original, but with images encoded as MP4 videos in the proper LeRobot format.
+
+```bash
+# Local-only: Save to a custom output directory (no hub push)
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht_image \
+    --operation.type convert_to_video \
+    --operation.output_dir /path/to/output/pusht_video
+
+# Save with new repo_id (local storage)
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht_image \
+    --new_repo_id lerobot/pusht_video \
+    --operation.type convert_to_video
+
+# Convert and push to Hugging Face Hub
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht_image \
+    --new_repo_id lerobot/pusht_video \
+    --operation.type convert_to_video \
+    --push_to_hub true
+
+# Convert with custom video codec and quality settings
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht_image \
+    --operation.type convert_to_video \
+    --operation.output_dir outputs/pusht_video \
+    --operation.vcodec libsvtav1 \
+    --operation.pix_fmt yuv420p \
+    --operation.g 2 \
+    --operation.crf 30
+
+# Convert only specific episodes
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht_image \
+    --operation.type convert_to_video \
+    --operation.output_dir outputs/pusht_video \
+    --operation.episode_indices "[0, 1, 2, 5, 10]"
+
+# Convert with multiple workers for parallel processing
+lerobot-edit-dataset \
+    --repo_id lerobot/pusht_image \
+    --operation.type convert_to_video \
+    --operation.output_dir outputs/pusht_video \
+    --operation.num_workers 8
+```
+
+**Parameters:**
+
+- `output_dir`: Custom output directory (optional - by default uses `new_repo_id` or `{repo_id}_video`)
+- `vcodec`: Video codec to use - options: `h264`, `hevc`, `libsvtav1` (default: `libsvtav1`)
+- `pix_fmt`: Pixel format - options: `yuv420p`, `yuv444p` (default: `yuv420p`)
+- `g`: Group of pictures (GOP) size - lower values give better quality but larger files (default: 2)
+- `crf`: Constant rate factor - lower values give better quality but larger files, 0 is lossless (default: 30)
+- `fast_decode`: Fast decode tuning option (default: 0)
+- `episode_indices`: List of specific episodes to convert (default: all episodes)
+- `num_workers`: Number of parallel workers for processing (default: 4)
+
+**Note:** The resulting dataset will be a proper LeRobotDataset with all cameras encoded as videos in the `videos/` directory, with parquet files containing only metadata (no raw image data). All episodes, stats, and tasks are preserved.
+
 ### Push to Hub
 
-Add the `--push_to_hub` flag to any command to automatically upload the resulting dataset to the Hugging Face Hub:
+Add the `--push_to_hub true` flag to any command to automatically upload the resulting dataset to the Hugging Face Hub:
 
 ```bash
 lerobot-edit-dataset \
@@ -96,7 +159,7 @@ lerobot-edit-dataset \
     --new_repo_id lerobot/pusht_after_deletion \
     --operation.type delete_episodes \
     --operation.episode_indices "[0, 2, 5]" \
-    --push_to_hub
+    --push_to_hub true
 ```
 
 There is also a tool for adding features to a dataset that is not yet covered in `lerobot-edit-dataset`.

examples/behavior_1k/behavior_lerobot_dataset_v3.py

@@ -1,464 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-BehaviorLeRobotDatasetV3: A wrapper around LeRobotDataset v3.0 for loading BEHAVIOR-1K data.
-
-This wrapper extends LeRobotDataset to support BEHAVIOR-1K specific features:
-- Modality and camera selection (rgb, depth, seg_instance_id)
-- Efficient chunk streaming mode with keyframe access
-- Additional BEHAVIOR-1K metadata (cam_rel_poses, task_info, etc.)
-"""
-
-import logging
-from collections.abc import Callable
-from pathlib import Path
-
-import datasets
-import numpy as np
-from behaviour_1k_constants import ROBOT_CAMERA_NAMES, ROBOT_TYPE
-from torch.utils.data import Dataset, get_worker_info
-
-from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset, LeRobotDatasetMetadata
-from lerobot.datasets.utils import (
-    check_delta_timestamps,
-    get_delta_indices,
-    get_safe_version,
-    hf_transform_to_torch,
-)
-from lerobot.datasets.video_utils import decode_video_frames, get_safe_default_codec
-from lerobot.utils.constants import HF_LEROBOT_HOME
-
-logger = logging.getLogger(__name__)
-
-
-class BehaviorLeRobotDatasetMetadata(LeRobotDatasetMetadata):
-    """
-    Extended metadata class for BEHAVIOR-1K datasets.
-
-    Adds support for:
-    - Modality and camera filtering
-    - Custom metainfo and annotation paths
-    """
-
-    def __init__(
-        self,
-        repo_id: str,
-        root: str | Path | None = None,
-        revision: str | None = None,
-        force_cache_sync: bool = False,
-        metadata_buffer_size: int = 10,
-        modalities: set[str] | None = None,
-        cameras: set[str] | None = None,
-    ):
-        self.modalities = set(modalities) if modalities else {"rgb", "depth", "seg_instance_id"}
-        self.camera_names = set(cameras) if cameras else {"head", "left_wrist", "right_wrist"}
-
-        assert self.modalities.issubset({"rgb", "depth", "seg_instance_id"}), (
-            f"Modalities must be subset of ['rgb', 'depth', 'seg_instance_id'], got {self.modalities}"
-        )
-
-        assert self.camera_names.issubset(set(ROBOT_CAMERA_NAMES[ROBOT_TYPE])), (
-            f"Camera names must be subset of {list(ROBOT_CAMERA_NAMES[ROBOT_TYPE])}, got {self.camera_names}"
-        )
-
-        super().__init__(repo_id, root, revision, force_cache_sync, metadata_buffer_size)
-
-    @property
-    def filtered_features(self) -> dict[str, dict]:
-        """Return only features matching selected modalities and cameras."""
-        features = {}
-        for name, feature_info in self.features.items():
-            if not name.startswith("observation.images."):
-                features[name] = feature_info
-                continue
-
-            parts = name.split(".")
-            if len(parts) >= 4:
-                modality = parts[2]
-                camera = parts[3]
-                if modality in self.modalities and camera in self.camera_names:
-                    features[name] = feature_info
-
-        return features
-
-    @property
-    def video_keys(self) -> list[str]:
-        """Return only video keys for selected modalities and cameras."""
-        all_video_keys = super().video_keys
-
-        filtered_keys = []
-        for key in all_video_keys:
-            parts = key.split(".")
-            if len(parts) >= 4:
-                modality = parts[2]
-                camera = parts[3]
-                if modality in self.modalities and camera in self.camera_names:
-                    filtered_keys.append(key)
-
-        return filtered_keys
-
-    def get_metainfo_path(self, ep_index: int) -> Path:
-        """Get path to episode metainfo file."""
-        if "metainfo_path" in self.info:
-            fpath = self.info["metainfo_path"].format(episode_index=ep_index)
-            return Path(fpath)
-        return None
-
-    def get_annotation_path(self, ep_index: int) -> Path:
-        """Get path to episode annotation file."""
-        if "annotation_path" in self.info:
-            fpath = self.info["annotation_path"].format(episode_index=ep_index)
-            return Path(fpath)
-        return None
-
-
-class BehaviorLeRobotDatasetV3(LeRobotDataset):
-    """
-    BEHAVIOR-1K wrapper for LeRobotDataset v3.0.
-
-    Each BEHAVIOR-1K dataset contains a single task (e.g., behavior1k-task0000).
-    See https://huggingface.co/collections/lerobot/behavior-1k for all available tasks.
-
-    Key features:
-    - Modality and camera selection
-    - Efficient chunk streaming with keyframe access (recommended for B1K with GOP=250)
-    - Support for BEHAVIOR-1K specific observations (cam_rel_poses, task_info, task_index)
-    """
-
-    def __init__(
-        self,
-        repo_id: str,
-        root: str | Path | None = None,
-        episodes: list[int] | None = None,
-        image_transforms: Callable | None = None,
-        delta_timestamps: dict[list[float]] | None = None,
-        tolerance_s: float = 1e-4,
-        revision: str | None = None,
-        force_cache_sync: bool = False,
-        download_videos: bool = True,
-        video_backend: str | None = None,
-        batch_encoding_size: int = 1,
-        # BEHAVIOR-1K specific arguments
-        modalities: list[str] | None = None,
-        cameras: list[str] | None = None,
-        check_timestamp_sync: bool = True,
-        chunk_streaming_using_keyframe: bool = True,
-        shuffle: bool = True,
-        seed: int = 42,
-    ):
-        """
-        Initialize BEHAVIOR-1K dataset.
-
-        Args:
-            repo_id: HuggingFace repository ID (e.g., "lerobot/behavior1k-task0000")
-            root: Local directory for dataset storage
-            episodes: List of episode indices to load (for train/val split)
-            image_transforms: Torchvision v2 transforms for images
-            delta_timestamps: Temporal offsets for history/future frames
-            tolerance_s: Tolerance for timestamp synchronization
-            revision: Git revision/branch to load
-            force_cache_sync: Force re-download from hub
-            download_videos: Whether to download video files
-            video_backend: Video decoder ('pyav' or 'torchcodec')
-            batch_encoding_size: Batch size for video encoding
-            modalities: List of modalities to load (None = all: rgb, depth, seg_instance_id)
-            cameras: List of cameras to load (None = all: head, left_wrist, right_wrist)
-            check_timestamp_sync: Verify timestamp synchronization (can be slow)
-            chunk_streaming_using_keyframe: Use keyframe-based streaming (STRONGLY RECOMMENDED for B1K)
-            shuffle: Shuffle chunks in streaming mode
-            seed: Random seed for shuffling
-        """
-        Dataset.__init__(self)
-
-        self.repo_id = repo_id
-        if root:
-            self.root = Path(root)
-        else:
-            dataset_name = repo_id.split("/")[-1] if "/" in repo_id else repo_id
-            self.root = HF_LEROBOT_HOME / dataset_name
-
-        self.image_transforms = image_transforms
-        self.delta_timestamps = delta_timestamps
-        self.tolerance_s = tolerance_s
-        self.revision = revision if revision else CODEBASE_VERSION
-        self.video_backend = video_backend if video_backend else get_safe_default_codec()
-        self.delta_indices = None
-        self.batch_encoding_size = batch_encoding_size
-        self.episodes_since_last_encoding = 0
-        self.seed = seed
-
-        self.image_writer = None
-        self.episode_buffer = None
-        self.writer = None
-        self.latest_episode = None
-        self._current_file_start_frame = None
-
-        self.root.mkdir(exist_ok=True, parents=True)
-
-        if modalities is None:
-            modalities = ["rgb", "depth", "seg_instance_id"]
-        if "seg_instance_id" in modalities:
-            assert chunk_streaming_using_keyframe, (
-                "For performance, seg_instance_id requires chunk_streaming_using_keyframe=True"
-            )
-        if "depth" in modalities:
-            assert self.video_backend == "pyav", "Depth videos require video_backend='pyav'"
-        if cameras is None:
-            cameras = ["head", "left_wrist", "right_wrist"]
-
-        self.meta = BehaviorLeRobotDatasetMetadata(
-            repo_id=self.repo_id,
-            root=self.root,
-            revision=self.revision,
-            force_cache_sync=force_cache_sync,
-            modalities=modalities,
-            cameras=cameras,
-        )
-
-        if episodes is not None:
-            self.episodes = sorted([i for i in episodes if i < len(self.meta.episodes)])
-        else:
-            self.episodes = list(range(len(self.meta.episodes)))
-
-        logger.info(f"Total episodes: {len(self.episodes)}")
-
-        self._chunk_streaming_using_keyframe = chunk_streaming_using_keyframe
-        if self._chunk_streaming_using_keyframe:
-            if not shuffle:
-                logger.warning("Chunk streaming enabled but shuffle=False. This may reduce randomness.")
-            self.chunks = self._get_keyframe_chunk_indices()
-            self.current_streaming_chunk_idx = None if shuffle else 0
-            self.current_streaming_frame_idx = None if shuffle else self.chunks[0][0] if self.chunks else 0
-            self.obs_loaders = {}
-            self._should_obs_loaders_reload = True
-
-        self._lazy_loading = False
-        self._recorded_frames = self.meta.total_frames
-        self._writer_closed_for_reading = False
-
-        try:
-            if force_cache_sync:
-                raise FileNotFoundError
-            self.hf_dataset = self.load_hf_dataset()
-        except (AssertionError, FileNotFoundError, NotADirectoryError):
-            self.revision = get_safe_version(self.repo_id, self.revision)
-            self.download_episodes(download_videos)
-            self.hf_dataset = self.load_hf_dataset()
-
-        if self.delta_timestamps is not None:
-            check_delta_timestamps(self.delta_timestamps, self.meta.fps, self.tolerance_s)
-            self.delta_indices = get_delta_indices(self.delta_timestamps, self.meta.fps)
-
-    @property
-    def fps(self) -> int:
-        """Frames per second."""
-        return self.meta.fps
-
-    @property
-    def features(self) -> dict:
-        """Dataset features (filtered by modalities/cameras)."""
-        return self.meta.filtered_features
-
-    @property
-    def num_episodes(self) -> int:
-        """Number of episodes."""
-        return len(self.episodes)
-
-    @property
-    def num_frames(self) -> int:
-        """Total number of frames."""
-        return len(self.hf_dataset)
-
-    def get_episodes_file_paths(self) -> list[str]:
-        """
-        Get download patterns for requested episodes.
-
-        Returns glob patterns for download rather than specific file paths.
-
-        Note: Unlike the base LeRobotDataset, this method cannot filter downloads to only
-        requested episodes because:
-        1. BEHAVIOR-1K episode indices are encoded (e.g., 10010 for task 1, episode 10)
-        2. Episodes are chunked across multiple parquet/video files
-        3. The parquet files are organized by chunk, not by episode
-
-        Therefore, we download full data/meta/video directories and rely on
-        `self.load_hf_dataset()` to filter to requested episodes from the loaded data.
-        """
-        allow_patterns = ["data/**", "meta/**"]
-
-        # Filter by modalities and cameras for video patterns
-        if len(self.meta.video_keys) > 0:
-            if len(self.meta.modalities) != 3 or len(self.meta.camera_names) != 3:
-                # Only download specific modality/camera combinations
-                for modality in self.meta.modalities:
-                    for camera in self.meta.camera_names:
-                        allow_patterns.append(f"**/observation.images.{modality}.{camera}/**")
-            else:
-                # Download all videos (no filtering needed)
-                allow_patterns.append("videos/**")
-
-        return allow_patterns
-
-    def download_episodes(self, download_videos: bool = True) -> None:
-        """
-        Download episodes with modality/camera filtering.
-
-        Follows the same pattern as base LeRobotDataset.download() but uses
-        get_episodes_file_paths() which returns patterns for modality/camera filtering.
-        """
-        ignore_patterns = None if download_videos else "videos/"
-        files = self.get_episodes_file_paths()
-        self.pull_from_repo(allow_patterns=files, ignore_patterns=ignore_patterns)
-
-    def pull_from_repo(
-        self,
-        allow_patterns: list[str] | str | None = None,
-        ignore_patterns: list[str] | str | None = None,
-    ) -> None:
-        """Pull dataset from HuggingFace Hub."""
-
-        from huggingface_hub import snapshot_download
-
-        logger.info(f"Pulling dataset {self.repo_id} from HuggingFace Hub...")
-        snapshot_download(
-            self.repo_id,
-            repo_type="dataset",
-            revision=self.revision,
-            local_dir=self.root,
-            allow_patterns=allow_patterns,
-            ignore_patterns=ignore_patterns,
-        )
-
-    def load_hf_dataset(self) -> datasets.Dataset:
-        """Load dataset from parquet files."""
-        from datasets import load_dataset
-
-        path = str(self.root / "data")
-        hf_dataset = load_dataset("parquet", data_dir=path, split="train")
-
-        hf_dataset.set_transform(hf_transform_to_torch)
-        return hf_dataset
-
-    def _get_keyframe_chunk_indices(self, chunk_size: int = 250) -> list[tuple[int, int, int]]:
-        """
-        Divide episodes into chunks based on GOP size (keyframe interval).
-
-        For BEHAVIOR-1K, GOP size is 250 frames for efficient storage.
-
-        Returns:
-            List of (start_index, end_index, local_start_index) tuples
-        """
-        chunks = []
-        offset = 0
-
-        for ep_array_idx in self.episodes:
-            # self.episodes contains array indices, so access directly
-            ep = self.meta.episodes[ep_array_idx]
-            length = ep["length"]
-            local_starts = list(range(0, length, chunk_size))
-            local_ends = local_starts[1:] + [length]
-
-            for local_start, local_end in zip(local_starts, local_ends, strict=True):
-                chunks.append((offset + local_start, offset + local_end, local_start))
-            offset += length
-
-        return chunks
-
-    def __getitem__(self, idx: int) -> dict:
-        """Get item by index, with optional chunk streaming."""
-        if not self._chunk_streaming_using_keyframe:
-            item = self.hf_dataset[idx]
-
-            for key in self.meta.video_keys:
-                if key in self.features:
-                    ep_idx = item["episode_index"].item()
-                    timestamp = item["timestamp"].item()
-                    video_path = self.root / self.meta.get_video_file_path(ep_idx, key)
-                    frames = decode_video_frames(
-                        video_path, [timestamp], self.tolerance_s, self.video_backend
-                    )
-                    item[key] = frames.squeeze(0)
-
-            if self.image_transforms is not None:
-                for key in self.features:
-                    if key.startswith("observation.images."):
-                        item[key] = self.image_transforms(item[key])
-
-            if "task_index" in item:
-                task_idx = item["task_index"].item()
-                try:
-                    item["task"] = self.meta.tasks.iloc[task_idx].name
-                except (IndexError, AttributeError):
-                    item["task"] = f"task_{task_idx}"
-
-            return item
-
-        return self._get_item_streaming(idx)
-
-    def _get_item_streaming(self, idx: int) -> dict:
-        """Get item in chunk streaming mode."""
-        if self.current_streaming_chunk_idx is None:
-            worker_info = get_worker_info()
-            worker_id = 0 if worker_info is None else worker_info.id
-            rng = np.random.default_rng(self.seed + worker_id)
-            rng.shuffle(self.chunks)
-            self.current_streaming_chunk_idx = rng.integers(0, len(self.chunks)).item()
-            self.current_streaming_frame_idx = self.chunks[self.current_streaming_chunk_idx][0]
-
-        if self.current_streaming_frame_idx >= self.chunks[self.current_streaming_chunk_idx][1]:
-            self.current_streaming_chunk_idx += 1
-            if self.current_streaming_chunk_idx >= len(self.chunks):
-                self.current_streaming_chunk_idx = 0
-            self.current_streaming_frame_idx = self.chunks[self.current_streaming_chunk_idx][0]
-            self._should_obs_loaders_reload = True
-
-        item = self.hf_dataset[self.current_streaming_frame_idx]
-        ep_idx = item["episode_index"].item()
-
-        if self._should_obs_loaders_reload:
-            for loader in self.obs_loaders.values():
-                if hasattr(loader, "close"):
-                    loader.close()
-            self.obs_loaders = {}
-            self.current_streaming_episode_idx = ep_idx
-            self._should_obs_loaders_reload = False
-
-        for key in self.meta.video_keys:
-            if key in self.features:
-                timestamp = item["timestamp"].item()
-                video_path = self.root / self.meta.get_video_file_path(ep_idx, key)
-                frames = decode_video_frames(video_path, [timestamp], self.tolerance_s, self.video_backend)
-                item[key] = frames.squeeze(0)
-
-        if self.image_transforms is not None:
-            for key in self.features:
-                if key.startswith("observation.images."):
-                    item[key] = self.image_transforms(item[key])
-
-        if "task_index" in item:
-            task_idx = item["task_index"].item()
-            try:
-                item["task"] = self.meta.tasks.iloc[task_idx].name
-            except (IndexError, AttributeError):
-                item["task"] = f"task_{task_idx}"
-
-        self.current_streaming_frame_idx += 1
-        return item
-
-    def __len__(self) -> int:
-        """Total number of frames."""
-        return len(self.hf_dataset)

examples/behavior_1k/behaviour_1k_constants.py

@@ -1,350 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from collections import OrderedDict
-
-import numpy as np
-import torch as th
-
-ROBOT_TYPE = "R1Pro"
-FPS = 30
-
-ROBOT_CAMERA_NAMES = {
-    "A1": {
-        "external": "external::external_camera",
-        "wrist": "external::wrist_camera",
-    },
-    "R1Pro": {
-        "left_wrist": "robot_r1::robot_r1:left_realsense_link:Camera:0",
-        "right_wrist": "robot_r1::robot_r1:right_realsense_link:Camera:0",
-        "head": "robot_r1::robot_r1:zed_link:Camera:0",
-    },
-}
-
-# Camera resolutions and corresponding intrinstics
-HEAD_RESOLUTION = (720, 720)
-WRIST_RESOLUTION = (480, 480)
-# TODO: Fix A1
-CAMERA_INTRINSICS = {
-    "A1": {
-        "external": np.array(
-            [[306.0, 0.0, 360.0], [0.0, 306.0, 360.0], [0.0, 0.0, 1.0]], dtype=np.float32
-        ),  # 240x240
-        "wrist": np.array(
-            [[388.6639, 0.0, 240.0], [0.0, 388.6639, 240.0], [0.0, 0.0, 1.0]], dtype=np.float32
-        ),  # 240x240
-    },
-    "R1Pro": {
-        "head": np.array(
-            [[306.0, 0.0, 360.0], [0.0, 306.0, 360.0], [0.0, 0.0, 1.0]], dtype=np.float32
-        ),  # 720x720
-        "left_wrist": np.array(
-            [[388.6639, 0.0, 240.0], [0.0, 388.6639, 240.0], [0.0, 0.0, 1.0]], dtype=np.float32
-        ),  # 480x480
-        "right_wrist": np.array(
-            [[388.6639, 0.0, 240.0], [0.0, 388.6639, 240.0], [0.0, 0.0, 1.0]], dtype=np.float32
-        ),  # 480x480
-    },
-}
-
-
-# Dataset features for BEHAVIOR-1K LeRobotDataset v3.0
-BEHAVIOR_DATASET_FEATURES = {
-    # Actions
-    "action": {
-        "dtype": "float32",
-        "shape": (23,),  # 23-dimensional action space for R1Pro
-        "names": None,
-    },
-    # Proprioception
-    "observation.state": {
-        "dtype": "float32",
-        "shape": (256,),  # Full proprioception state
-        "names": None,
-    },
-    # Camera relative poses
-    "observation.cam_rel_poses": {
-        "dtype": "float32",
-        "shape": (21,),  # 3 cameras * 7 (pos + quat)
-        "names": None,
-    },
-    # Task information
-    "observation.task_info": {
-        "dtype": "float32",
-        "shape": (None,),  # Variable size
-        "names": None,
-    },
-    # RGB images
-    "observation.images.rgb.head": {
-        "dtype": "video",
-        "shape": [720, 720, 3],
-        "names": ["height", "width", "channels"],
-    },
-    "observation.images.rgb.left_wrist": {
-        "dtype": "video",
-        "shape": [480, 480, 3],
-        "names": ["height", "width", "channels"],
-    },
-    "observation.images.rgb.right_wrist": {
-        "dtype": "video",
-        "shape": [480, 480, 3],
-        "names": ["height", "width", "channels"],
-    },
-    # Depth images
-    "observation.images.depth.head": {
-        "dtype": "video",
-        "shape": [720, 720, 1],
-        "names": ["height", "width", "channels"],
-    },
-    "observation.images.depth.left_wrist": {
-        "dtype": "video",
-        "shape": [480, 480, 1],
-        "names": ["height", "width", "channels"],
-    },
-    "observation.images.depth.right_wrist": {
-        "dtype": "video",
-        "shape": [480, 480, 1],
-        "names": ["height", "width", "channels"],
-    },
-    # Segmentation instance ID images
-    "observation.images.seg_instance_id.head": {
-        "dtype": "video",
-        "shape": [720, 720, 1],
-        "names": ["height", "width", "channels"],
-    },
-    "observation.images.seg_instance_id.left_wrist": {
-        "dtype": "video",
-        "shape": [480, 480, 1],
-        "names": ["height", "width", "channels"],
-    },
-    "observation.images.seg_instance_id.right_wrist": {
-        "dtype": "video",
-        "shape": [480, 480, 1],
-        "names": ["height", "width", "channels"],
-    },
-}
-
-
-# Action indices
-ACTION_QPOS_INDICES = {
-    "A1": OrderedDict(
-        {
-            "arm": np.s_[0:6],
-            "gripper": np.s_[6:7],
-        }
-    ),
-    "R1Pro": OrderedDict(
-        {
-            "base": np.s_[0:3],
-            "torso": np.s_[3:7],
-            "left_arm": np.s_[7:14],
-            "left_gripper": np.s_[14:15],
-            "right_arm": np.s_[15:22],
-            "right_gripper": np.s_[22:23],
-        }
-    ),
-}
-
-
-# Proprioception configuration
-PROPRIOCEPTION_INDICES = {
-    "A1": OrderedDict(
-        {
-            "joint_qpos": np.s_[0:8],
-            "joint_qpos_sin": np.s_[8:16],
-            "joint_qpos_cos": np.s_[16:24],
-            "joint_qvel": np.s_[24:32],
-            "joint_qeffort": np.s_[32:40],
-            "eef_0_pos": np.s_[40:43],
-            "eef_0_quat": np.s_[43:47],
-            "grasp_0": np.s_[47:48],
-            "gripper_0_qpos": np.s_[48:50],
-            "gripper_0_qvel": np.s_[50:52],
-        }
-    ),
-    "R1Pro": OrderedDict(
-        {
-            "joint_qpos": np.s_[
-                0:28
-            ],  # Full robot joint positions, the first 6 are base joints, which is NOT allowed in standard track
-            "joint_qpos_sin": np.s_[
-                28:56
-            ],  # Full robot joint positions, the first 6 are base joints, which is NOT allowed in standard track
-            "joint_qpos_cos": np.s_[
-                56:84
-            ],  # Full robot joint positions, the first 6 are base joints, which is NOT allowed in standard track
-            "joint_qvel": np.s_[84:112],
-            "joint_qeffort": np.s_[112:140],
-            "robot_pos": np.s_[140:143],  # Global pos, this is NOT allowed in standard track
-            "robot_ori_cos": np.s_[143:146],  # Global ori, this is NOT allowed in standard track
-            "robot_ori_sin": np.s_[146:149],  # Global ori, this is NOT allowed in standard track
-            "robot_2d_ori": np.s_[149:150],  # 2D global ori, this is NOT allowed in standard track
-            "robot_2d_ori_cos": np.s_[150:151],  # 2D global ori, this is NOT allowed in standard track
-            "robot_2d_ori_sin": np.s_[151:152],  # 2D global ori, this is NOT allowed in standard track
-            "robot_lin_vel": np.s_[152:155],
-            "robot_ang_vel": np.s_[155:158],
-            "arm_left_qpos": np.s_[158:165],
-            "arm_left_qpos_sin": np.s_[165:172],
-            "arm_left_qpos_cos": np.s_[172:179],
-            "arm_left_qvel": np.s_[179:186],
-            "eef_left_pos": np.s_[186:189],
-            "eef_left_quat": np.s_[189:193],
-            "gripper_left_qpos": np.s_[193:195],
-            "gripper_left_qvel": np.s_[195:197],
-            "arm_right_qpos": np.s_[197:204],
-            "arm_right_qpos_sin": np.s_[204:211],
-            "arm_right_qpos_cos": np.s_[211:218],
-            "arm_right_qvel": np.s_[218:225],
-            "eef_right_pos": np.s_[225:228],
-            "eef_right_quat": np.s_[228:232],
-            "gripper_right_qpos": np.s_[232:234],
-            "gripper_right_qvel": np.s_[234:236],
-            "trunk_qpos": np.s_[236:240],
-            "trunk_qvel": np.s_[240:244],
-            "base_qpos": np.s_[244:247],  # Base joint position, this is NOT allowed in standard track
-            "base_qpos_sin": np.s_[247:250],  # Base joint position, this is NOT allowed in standard track
-            "base_qpos_cos": np.s_[250:253],  # Base joint position, this is NOT allowed in standard track
-            "base_qvel": np.s_[253:256],
-        }
-    ),
-}
-
-# Proprioception indices
-PROPRIO_QPOS_INDICES = {
-    "A1": OrderedDict(
-        {
-            "arm": np.s_[0:6],
-            "gripper": np.s_[6:8],
-        }
-    ),
-    "R1Pro": OrderedDict(
-        {
-            "torso": np.s_[6:10],
-            "left_arm": np.s_[10:24:2],
-            "right_arm": np.s_[11:24:2],
-            "left_gripper": np.s_[24:26],
-            "right_gripper": np.s_[26:28],
-        }
-    ),
-}
-
-
-# Joint limits (lower, upper)
-JOINT_RANGE = {
-    "A1": {
-        "arm": (
-            th.tensor([-2.8798, 0.0, -3.3161, -2.8798, -1.6581, -2.8798], dtype=th.float32),
-            th.tensor([2.8798, 3.1415, 0.0, 2.8798, 1.6581, 2.8798], dtype=th.float32),
-        ),
-        "gripper": (th.tensor([0.00], dtype=th.float32), th.tensor([0.03], dtype=th.float32)),
-    },
-    "R1Pro": {
-        "base": (
-            th.tensor([-0.75, -0.75, -1.0], dtype=th.float32),
-            th.tensor([0.75, 0.75, 1.0], dtype=th.float32),
-        ),
-        "torso": (
-            th.tensor([-1.1345, -2.7925, -1.8326, -3.0543], dtype=th.float32),
-            th.tensor([1.8326, 2.5307, 1.5708, 3.0543], dtype=th.float32),
-        ),
-        "left_arm": (
-            th.tensor([-4.4506, -0.1745, -2.3562, -2.0944, -2.3562, -1.0472, -1.5708], dtype=th.float32),
-            th.tensor([1.3090, 3.1416, 2.3562, 0.3491, 2.3562, 1.0472, 1.5708], dtype=th.float32),
-        ),
-        "left_gripper": (th.tensor([0.00], dtype=th.float32), th.tensor([0.05], dtype=th.float32)),
-        "right_arm": (
-            th.tensor([-4.4506, -3.1416, -2.3562, -2.0944, -2.3562, -1.0472, -1.5708], dtype=th.float32),
-            th.tensor([1.3090, 0.1745, 2.3562, 0.3491, 2.3562, 1.0472, 1.5708], dtype=th.float32),
-        ),
-        "right_gripper": (th.tensor([0.00], dtype=th.float32), th.tensor([0.05], dtype=th.float32)),
-    },
-}
-
-
-EEF_POSITION_RANGE = {
-    "A1": {
-        "0": (th.tensor([0.0, -0.7, 0.0], dtype=th.float32), th.tensor([0.7, 0.7, 0.7], dtype=th.float32)),
-    },
-    "R1Pro": {
-        "left": (
-            th.tensor([0.0, -0.65, 0.0], dtype=th.float32),
-            th.tensor([0.65, 0.65, 2.5], dtype=th.float32),
-        ),
-        "right": (
-            th.tensor([0.0, -0.65, 0.0], dtype=th.float32),
-            th.tensor([0.65, 0.65, 2.5], dtype=th.float32),
-        ),
-    },
-}
-
-
-TASK_NAMES_TO_INDICES = {
-    # B10
-    "turning_on_radio": 0,
-    "picking_up_trash": 1,
-    "putting_away_Halloween_decorations": 2,
-    "cleaning_up_plates_and_food": 3,
-    "can_meat": 4,
-    "setting_mousetraps": 5,
-    "hiding_Easter_eggs": 6,
-    "picking_up_toys": 7,
-    "rearranging_kitchen_furniture": 8,
-    "putting_up_Christmas_decorations_inside": 9,
-    # B20
-    "set_up_a_coffee_station_in_your_kitchen": 10,
-    "putting_dishes_away_after_cleaning": 11,
-    "preparing_lunch_box": 12,
-    "loading_the_car": 13,
-    "carrying_in_groceries": 14,
-    "bringing_in_wood": 15,
-    "moving_boxes_to_storage": 16,
-    "bringing_water": 17,
-    "tidying_bedroom": 18,
-    "outfit_a_basic_toolbox": 19,
-    # B30
-    "sorting_vegetables": 20,
-    "collecting_childrens_toys": 21,
-    "putting_shoes_on_rack": 22,
-    "boxing_books_up_for_storage": 23,
-    "storing_food": 24,
-    "clearing_food_from_table_into_fridge": 25,
-    "assembling_gift_baskets": 26,
-    "sorting_household_items": 27,
-    "getting_organized_for_work": 28,
-    "clean_up_your_desk": 29,
-    # B40
-    "setting_the_fire": 30,
-    "clean_boxing_gloves": 31,
-    "wash_a_baseball_cap": 32,
-    "wash_dog_toys": 33,
-    "hanging_pictures": 34,
-    "attach_a_camera_to_a_tripod": 35,
-    "clean_a_patio": 36,
-    "clean_a_trumpet": 37,
-    "spraying_for_bugs": 38,
-    "spraying_fruit_trees": 39,
-    # B50
-    "make_microwave_popcorn": 40,
-    "cook_cabbage": 41,
-    "chop_an_onion": 42,
-    "slicing_vegetables": 43,
-    "chopping_wood": 44,
-    "cook_hot_dogs": 45,
-    "cook_bacon": 46,
-    "freeze_pies": 47,
-    "canning_food": 48,
-    "make_pizza": 49,
-}
-TASK_INDICES_TO_NAMES = {v: k for k, v in TASK_NAMES_TO_INDICES.items()}

examples/behavior_1k/convert_to_lerobot_v3.py

@@ -1,605 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert Behavior Dataset to LeRobotDataset v3.0 format"""
-
-import argparse
-import json
-import logging
-import shutil
-from pathlib import Path
-
-import jsonlines
-import pandas as pd
-import pyarrow as pa
-import tqdm
-from datasets import Dataset, Features, Image
-
-from lerobot.datasets.compute_stats import aggregate_stats
-from lerobot.datasets.utils import (
-    DEFAULT_CHUNK_SIZE,
-    DEFAULT_DATA_FILE_SIZE_IN_MB,
-    DEFAULT_DATA_PATH,
-    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
-    DEFAULT_VIDEO_PATH,
-    LEGACY_EPISODES_PATH,
-    LEGACY_EPISODES_STATS_PATH,
-    LEGACY_TASKS_PATH,
-    cast_stats_to_numpy,
-    flatten_dict,
-    get_file_size_in_mb,
-    get_parquet_file_size_in_mb,
-    get_parquet_num_frames,
-    load_info,
-    update_chunk_file_indices,
-    write_episodes,
-    write_info,
-    write_stats,
-    write_tasks,
-)
-from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
-from lerobot.utils.utils import init_logging
-
-# script to convert one single task to v3.1
-# TASK = 1
-NEW_ROOT = Path("/fsx/jade_choghari/tmp/bb")
-
-
-def get_total_episodes_task(local_dir: Path, task_id: int, task_ranges: dict, step) -> int:
-    """
-    Calculates the total number of episodes for a single, specified task.
-    """
-    # Simply load the episodes for the task and count them.
-    episodes = legacy_load_episodes_task(
-        local_dir=local_dir, task_id=task_id, task_ranges=task_ranges, step=step
-    )
-    return len(episodes)
-
-
-NUM_CAMERAS = 9
-
-
-def get_total_frames_task(local_dir, meta_path, task_id: int, task_ranges: dict, step: int) -> int:
-    episodes_metadata = legacy_load_episodes_task(
-        local_dir=local_dir, task_id=task_id, task_ranges=task_ranges, step=step
-    )
-    total_frames = 0
-    # like 'duration'
-    for ep in episodes_metadata.values():
-        duration_s = ep["length"]
-        total_frames += int(duration_s)
-    return total_frames
-
-
-def convert_info(
-    root, new_root, data_file_size_in_mb, video_file_size_in_mb, meta_path, task_id: int, task_ranges, step
-):
-    info = load_info(root)
-    info["codebase_version"] = "v3.0"
-    del info["total_videos"]
-    info["data_files_size_in_mb"] = data_file_size_in_mb
-    info["video_files_size_in_mb"] = video_file_size_in_mb
-    info["data_path"] = DEFAULT_DATA_PATH
-    info["video_path"] = DEFAULT_VIDEO_PATH if info["video_path"] is not None else None
-    info["fps"] = int(info["fps"])
-    for key in info["features"]:
-        if info["features"][key]["dtype"] == "video":
-            # already has fps in video_info
-            continue
-        info["features"][key]["fps"] = info["fps"]
-
-    info["total_episodes"] = get_total_episodes_task(root, task_id, task_ranges, step)
-    info["total_videos"] = info["total_episodes"] * NUM_CAMERAS
-    info["total_frames"] = get_total_frames_task(root, meta_path, task_id, task_ranges, step)
-    info["total_tasks"] = 1
-    write_info(info, new_root)
-
-
-def load_jsonlines(fpath: Path) -> list[any]:
-    with jsonlines.open(fpath, "r") as reader:
-        return list(reader)
-
-
-def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]:
-    tasks = load_jsonlines(local_dir / LEGACY_TASKS_PATH)
-    # return tasks dict such that
-    tasks = {item["task_index"]: item["task"] for item in sorted(tasks, key=lambda x: x["task_index"])}
-    task_to_task_index = {task: task_index for task_index, task in tasks.items()}
-    return tasks, task_to_task_index
-
-
-def convert_tasks(root, new_root, task_id: int):
-    tasks, _ = legacy_load_tasks(root)
-    if task_id not in tasks:
-        raise ValueError(f"Task ID {task_id} not found in tasks (available: {list(tasks.keys())})")
-    tasks = {task_id: tasks[task_id]}
-    task_indices = tasks.keys()
-    task_strings = tasks.values()
-    df_tasks = pd.DataFrame({"task_index": task_indices}, index=task_strings)
-    write_tasks(df_tasks, new_root)
-
-
-def concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys):
-    # TODO(rcadene): to save RAM use Dataset.from_parquet(file) and concatenate_datasets
-    dataframes = [pd.read_parquet(file) for file in paths_to_cat]
-    # Concatenate all DataFrames along rows
-    concatenated_df = pd.concat(dataframes, ignore_index=True)
-
-    path = new_root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    if len(image_keys) > 0:
-        schema = pa.Schema.from_pandas(concatenated_df)
-        features = Features.from_arrow_schema(schema)
-        for key in image_keys:
-            features[key] = Image()
-        schema = features.arrow_schema
-    else:
-        schema = None
-
-    concatenated_df.to_parquet(path, index=False, schema=schema)
-
-
-def get_image_keys(root):
-    info = load_info(root)
-    features = info["features"]
-    image_keys = [key for key, ft in features.items() if ft["dtype"] == "image"]
-    return image_keys
-
-
-def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int, task_index: int):
-    task_dir_name = f"task-00{task_index}"
-    data_dir = root / "data" / task_dir_name
-    ep_paths = sorted(data_dir.glob("*.parquet"))
-    image_keys = get_image_keys(root)
-
-    ep_idx = 0
-    chunk_idx = 0
-    file_idx = 0
-    size_in_mb = 0
-    num_frames = 0
-    paths_to_cat = []
-    episodes_metadata = []
-
-    logging.info(f"Converting data files from {len(ep_paths)} episodes")
-
-    for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
-        ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
-        ep_num_frames = get_parquet_num_frames(ep_path)
-        ep_metadata = {
-            "episode_index": ep_idx,
-            "data/chunk_index": chunk_idx,
-            "data/file_index": file_idx,
-            "dataset_from_index": num_frames,
-            "dataset_to_index": num_frames + ep_num_frames,
-        }
-        size_in_mb += ep_size_in_mb
-        num_frames += ep_num_frames
-        episodes_metadata.append(ep_metadata)
-        ep_idx += 1
-
-        if size_in_mb < data_file_size_in_mb:
-            paths_to_cat.append(ep_path)
-            continue
-
-        if paths_to_cat:
-            concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
-
-        # Reset for the next file
-        size_in_mb = ep_size_in_mb
-        paths_to_cat = [ep_path]
-
-        chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
-
-    # Write remaining data if any
-    if paths_to_cat:
-        concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
-
-    return episodes_metadata
-
-
-def convert_videos_of_camera(
-    root: Path, new_root: Path, video_key: str, video_file_size_in_mb: int, task_index: int
-):
-    # Access old paths to mp4
-    # videos_dir = root / "videos"
-    # ep_paths = sorted(videos_dir.glob(f"*/{video_key}/*.mp4"))
-    task_dir_name = f"task-00{task_index}"
-    videos_dir = root / "videos" / task_dir_name / video_key
-    ep_paths = sorted(videos_dir.glob("*.mp4"))
-    print("ep_paths", ep_paths)
-    ep_idx = 0
-    chunk_idx = 0
-    file_idx = 0
-    size_in_mb = 0
-    duration_in_s = 0.0
-    paths_to_cat = []
-    episodes_metadata = []
-
-    for ep_path in tqdm.tqdm(ep_paths, desc=f"convert videos of {video_key}"):
-        ep_size_in_mb = get_file_size_in_mb(ep_path)
-        ep_duration_in_s = get_video_duration_in_s(ep_path)
-
-        # Check if adding this episode would exceed the limit
-        if size_in_mb + ep_size_in_mb >= video_file_size_in_mb and len(paths_to_cat) > 0:
-            # Size limit would be exceeded, save current accumulation WITHOUT this episode
-            concatenate_video_files(
-                paths_to_cat,
-                new_root
-                / DEFAULT_VIDEO_PATH.format(video_key=video_key, chunk_index=chunk_idx, file_index=file_idx),
-            )
-
-            # Update episodes metadata for the file we just saved
-            for i, _ in enumerate(paths_to_cat):
-                past_ep_idx = ep_idx - len(paths_to_cat) + i
-                episodes_metadata[past_ep_idx][f"videos/{video_key}/chunk_index"] = chunk_idx
-                episodes_metadata[past_ep_idx][f"videos/{video_key}/file_index"] = file_idx
-
-            # Move to next file and start fresh with current episode
-            chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
-            size_in_mb = 0
-            duration_in_s = 0.0
-            paths_to_cat = []
-
-        # Add current episode metadata
-        ep_metadata = {
-            "episode_index": ep_idx,
-            f"videos/{video_key}/chunk_index": chunk_idx,  # Will be updated when file is saved
-            f"videos/{video_key}/file_index": file_idx,  # Will be updated when file is saved
-            f"videos/{video_key}/from_timestamp": duration_in_s,
-            f"videos/{video_key}/to_timestamp": duration_in_s + ep_duration_in_s,
-        }
-        episodes_metadata.append(ep_metadata)
-
-        # Add current episode to accumulation
-        paths_to_cat.append(ep_path)
-        size_in_mb += ep_size_in_mb
-        duration_in_s += ep_duration_in_s
-        ep_idx += 1
-
-    # Write remaining videos if any
-    if paths_to_cat:
-        concatenate_video_files(
-            paths_to_cat,
-            new_root
-            / DEFAULT_VIDEO_PATH.format(video_key=video_key, chunk_index=chunk_idx, file_index=file_idx),
-        )
-
-        # Update episodes metadata for the final file
-        for i, _ in enumerate(paths_to_cat):
-            past_ep_idx = ep_idx - len(paths_to_cat) + i
-            episodes_metadata[past_ep_idx][f"videos/{video_key}/chunk_index"] = chunk_idx
-            episodes_metadata[past_ep_idx][f"videos/{video_key}/file_index"] = file_idx
-
-    return episodes_metadata
-
-
-def get_video_keys(root):
-    info = load_info(root)
-    features = info["features"]
-    video_keys = [key for key, ft in features.items() if ft["dtype"] == "video"]
-    return video_keys
-
-
-def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int, task_id: int):
-    logging.info(f"Converting videos from {root} to {new_root}")
-
-    video_keys = get_video_keys(root)
-    if len(video_keys) == 0:
-        return None
-
-    video_keys = sorted(video_keys)
-
-    eps_metadata_per_cam = []
-    for camera in video_keys:
-        eps_metadata = convert_videos_of_camera(root, new_root, camera, video_file_size_in_mb, task_id)
-        eps_metadata_per_cam.append(eps_metadata)
-
-    num_eps_per_cam = [len(eps_cam_map) for eps_cam_map in eps_metadata_per_cam]
-    if len(set(num_eps_per_cam)) != 1:
-        raise ValueError(f"All cams dont have same number of episodes ({num_eps_per_cam}).")
-
-    episods_metadata = []
-    num_cameras = len(video_keys)
-    num_episodes = num_eps_per_cam[0]
-    for ep_idx in tqdm.tqdm(range(num_episodes), desc="convert videos"):
-        # Sanity check
-        ep_ids = [eps_metadata_per_cam[cam_idx][ep_idx]["episode_index"] for cam_idx in range(num_cameras)]
-        ep_ids += [ep_idx]
-        if len(set(ep_ids)) != 1:
-            raise ValueError(f"All episode indices need to match ({ep_ids}).")
-
-        ep_dict = {}
-        for cam_idx in range(num_cameras):
-            ep_dict.update(eps_metadata_per_cam[cam_idx][ep_idx])
-        episods_metadata.append(ep_dict)
-
-    return episods_metadata
-
-
-def infer_task_episode_ranges(episodes_jsonl_path: Path) -> dict:
-    """
-    Parse the Behavior-1K episodes.jsonl metadata and infer contiguous episode ranges per unique task.
-    Returns a dict:
-      { task_id: { "task_string": ..., "ep_start": ..., "ep_end": ... } }
-    """
-    task_ranges = {}
-    task_id = 0
-    current_task_str = None
-    ep_start = None
-    ep_end = None
-
-    with open(episodes_jsonl_path) as f:
-        for line in f:
-            if not line.strip():
-                continue
-            ep = json.loads(line)
-            ep_idx = ep["episode_index"]
-            task_str = ep["tasks"][0] if ep["tasks"] else "UNKNOWN"
-
-            if current_task_str is None:
-                current_task_str = task_str
-                ep_start = ep_idx
-                ep_end = ep_idx
-            elif task_str == current_task_str:
-                ep_end = ep_idx
-            else:
-                # close previous task group
-                task_ranges[task_id] = {
-                    "task_string": current_task_str,
-                    "ep_start": ep_start,
-                    "ep_end": ep_end,
-                }
-                task_id += 1
-                # start new one
-                current_task_str = task_str
-                ep_start = ep_idx
-                ep_end = ep_idx
-
-    # store last task
-    if current_task_str is not None:
-        task_ranges[task_id] = {
-            "task_string": current_task_str,
-            "ep_start": ep_start,
-            "ep_end": ep_end,
-        }
-
-    return task_ranges
-
-
-def legacy_load_episodes_task(local_dir: Path, task_id: int, task_ranges: dict, step: int = 10) -> dict:
-    """
-    Load only the episodes belonging to a specific task, inferred automatically from episode ranges.
-
-    Args:
-        local_dir (Path): Root path containing legacy meta/episodes.jsonl
-        task_id (int): Which task to load (key from the inferred task_ranges dict)
-        task_ranges (dict): Mapping from infer_task_episode_ranges()
-        step (int): Episode index step (Behavior-1K = 10)
-    """
-    all_episodes = legacy_load_episodes(local_dir)
-
-    # get the range for this task
-    if task_id not in task_ranges:
-        raise ValueError(f"Task id {task_id} not found in task_ranges")
-
-    ep_start = task_ranges[task_id]["ep_start"]
-    ep_end = task_ranges[task_id]["ep_end"]
-
-    task_episode_indices = range(ep_start, ep_end + step, step)
-    return {i: all_episodes[i] for i in task_episode_indices if i in all_episodes}
-
-
-def legacy_load_episodes(local_dir: Path) -> dict:
-    episodes = load_jsonlines(local_dir / LEGACY_EPISODES_PATH)
-    return {item["episode_index"]: item for item in sorted(episodes, key=lambda x: x["episode_index"])}
-
-
-def legacy_load_episodes_stats(local_dir: Path) -> dict:
-    episodes_stats = load_jsonlines(local_dir / LEGACY_EPISODES_STATS_PATH)
-    return {
-        item["episode_index"]: cast_stats_to_numpy(item["stats"])
-        for item in sorted(episodes_stats, key=lambda x: x["episode_index"])
-    }
-
-
-def legacy_load_episodes_stats_task(local_dir: Path, task_id: int, task_ranges: dict, step: int = 10) -> dict:
-    all_stats = legacy_load_episodes_stats(local_dir)
-
-    if task_id not in task_ranges:
-        raise ValueError(f"Task id {task_id} not found in task_ranges")
-
-    ep_start = task_ranges[task_id]["ep_start"]
-    ep_end = task_ranges[task_id]["ep_end"]
-
-    task_episode_indices = range(ep_start, ep_end + step, step)
-    return {i: all_stats[i] for i in task_episode_indices if i in all_stats}
-
-
-def generate_episode_metadata_dict(
-    episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_videos=None
-):
-    num_episodes = len(episodes_metadata)
-    episodes_legacy_metadata_vals = list(episodes_legacy_metadata.values())
-    episodes_stats_vals = list(episodes_stats.values())
-    episodes_stats_keys = list(episodes_stats.keys())
-
-    for i in range(num_episodes):
-        ep_legacy_metadata = episodes_legacy_metadata_vals[i]
-        ep_metadata = episodes_metadata[i]
-        ep_stats = episodes_stats_vals[i]
-
-        ep_ids_set = {
-            ep_legacy_metadata["episode_index"],
-            ep_metadata["episode_index"],
-            episodes_stats_keys[i],
-        }
-
-        if episodes_videos is None:
-            ep_video = {}
-        else:
-            ep_video = episodes_videos[i]
-            ep_ids_set.add(ep_video["episode_index"])
-        # we skip this check because ep_ids have a step of 10, whereas we convert with a step of 1
-        # if len(ep_ids_set) != 1:
-        #     raise ValueError(f"Number of episodes is not the same ({ep_ids_set}).")
-
-        ep_dict = {**ep_metadata, **ep_video, **ep_legacy_metadata, **flatten_dict({"stats": ep_stats})}
-        ep_dict["meta/episodes/chunk_index"] = 0
-        ep_dict["meta/episodes/file_index"] = 0
-        yield ep_dict
-
-
-def convert_episodes_metadata(
-    root, new_root, episodes_metadata, task_id: int, task_ranges, episodes_video_metadata=None
-):
-    logging.info(f"Converting episodes metadata from {root} to {new_root}")
-
-    # filter by task
-    episodes_legacy_metadata = legacy_load_episodes_task(root, task_id=task_id, task_ranges=task_ranges)
-    episodes_stats = legacy_load_episodes_stats_task(root, task_id=task_id, task_ranges=task_ranges)
-
-    num_eps_set = {len(episodes_legacy_metadata), len(episodes_metadata)}
-    if episodes_video_metadata is not None:
-        num_eps_set.add(len(episodes_video_metadata))
-
-    if len(num_eps_set) != 1:
-        raise ValueError(f"Number of episodes is not the same ({num_eps_set}).")
-
-    ds_episodes = Dataset.from_generator(
-        lambda: generate_episode_metadata_dict(
-            episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_video_metadata
-        )
-    )
-    write_episodes(ds_episodes, new_root)
-
-    stats = aggregate_stats(list(episodes_stats.values()))
-    write_stats(stats, new_root)
-
-
-def convert_dataset_local(
-    data_path: Path,
-    new_repo: Path,
-    task_id: int,
-    data_file_size_in_mb: int = DEFAULT_DATA_FILE_SIZE_IN_MB,
-    video_file_size_in_mb: int = DEFAULT_VIDEO_FILE_SIZE_IN_MB,
-    force_conversion: bool = False,
-):
-    """
-    Convert a local dataset to v3.x format, task-by-task, without using the Hugging Face Hub.
-
-    Args:
-        data_path (Path): path to local dataset root (e.g. /fsx/.../2025-challenge-demos)
-        new_repo (Path): path where converted dataset will be written (e.g. /fsx/.../behavior1k_v3)
-        task_id (int): which task to convert (index)
-        data_file_size_in_mb (int): max size per data chunk
-        video_file_size_in_mb (int): max size per video chunk
-        force_conversion (bool): overwrite existing conversion if True
-    """
-
-    root = Path(data_path)
-    new_root = Path(new_repo)
-
-    # Clean up if needed
-    if new_root.exists() and force_conversion:
-        shutil.rmtree(new_root)
-    new_root.mkdir(parents=True, exist_ok=True)
-
-    print(f"🔹 Starting conversion for task {task_id}")
-    print(f"Input root: {root}")
-    print(f"Output root: {new_root}")
-    # Infer task episode ranges
-    episodes_meta_path = root / "meta" / "episodes.jsonl"
-    task_ranges = infer_task_episode_ranges(episodes_meta_path)
-    convert_info(
-        root,
-        new_root,
-        data_file_size_in_mb,
-        video_file_size_in_mb,
-        episodes_meta_path,
-        task_id,
-        task_ranges,
-        step=10,
-    )
-    convert_tasks(root, new_root, task_id)
-    episodes_metadata = convert_data(root, new_root, data_file_size_in_mb, task_index=task_id)
-    episodes_videos_metadata = convert_videos(root, new_root, video_file_size_in_mb, task_id=task_id)
-    convert_episodes_metadata(
-        root,
-        new_root,
-        episodes_metadata,
-        task_id=task_id,
-        task_ranges=task_ranges,
-        episodes_video_metadata=episodes_videos_metadata,
-    )
-
-    print(f"✅ Conversion complete for task {task_id}")
-    print(f"Converted dataset written to: {new_root}")
-
-
-if __name__ == "__main__":
-    import argparse
-    from pathlib import Path
-
-    init_logging()
-
-    parser = argparse.ArgumentParser(
-        description="Convert Behavior-1K tasks to LeRobot v3 format (local only)"
-    )
-    parser.add_argument(
-        "--data-path",
-        type=str,
-        required=True,
-        help="Path to the local Behavior-1K dataset (e.g. /fsx/francesco_capuano/.cache/behavior-1k/2025-challenge-demos)",
-    )
-    parser.add_argument(
-        "--new-repo",
-        type=str,
-        required=True,
-        help="Path to the output directory for the converted dataset",
-    )
-    parser.add_argument(
-        "--task-id",
-        type=int,
-        required=True,
-        help="Task index to convert (e.g. 0, 1, 2, ...)",
-    )
-    parser.add_argument(
-        "--data-file-size-in-mb",
-        type=int,
-        default=DEFAULT_DATA_FILE_SIZE_IN_MB,
-        help=f"Maximum size per data chunk (default: {DEFAULT_DATA_FILE_SIZE_IN_MB})",
-    )
-    parser.add_argument(
-        "--video-file-size-in-mb",
-        type=int,
-        default=DEFAULT_VIDEO_FILE_SIZE_IN_MB,
-        help=f"Maximum size per video chunk (default: {DEFAULT_VIDEO_FILE_SIZE_IN_MB})",
-    )
-    parser.add_argument(
-        "--force-conversion",
-        action="store_true",
-        help="Force overwrite of existing conversion output if present.",
-    )
-
-    args = parser.parse_args()
-
-    convert_dataset_local(
-        data_path=Path(args.data_path),
-        new_repo=Path(args.new_repo),
-        task_id=args.task_id,
-        data_file_size_in_mb=args.data_file_size_in_mb,
-        video_file_size_in_mb=args.video_file_size_in_mb,
-        force_conversion=args.force_conversion,
-    )

examples/behavior_1k/load_behavior_1k_dataset.py

@@ -1,130 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Test script to verify BEHAVIOR-1K dataset loading with v3.0 wrapper.
-"""
-
-import argparse
-import logging
-
-from behavior_lerobot_dataset_v3 import BehaviorLeRobotDatasetV3
-
-from lerobot.utils.utils import init_logging
-
-init_logging()
-
-
-def load_behavior1k_dataset(repo_id, root):
-    """Test basic dataset loading."""
-    logging.info("=" * 80)
-    logging.info("Testing BEHAVIOR-1K dataset loading")
-    logging.info("=" * 80)
-
-    logging.info(f"\n1. Loading dataset with repo_id: {repo_id}")
-    dataset = BehaviorLeRobotDatasetV3(
-        repo_id=repo_id,
-        root=root,
-        modalities=["rgb"],
-        cameras=["head"],
-        chunk_streaming_using_keyframe=False,
-        check_timestamp_sync=False,
-    )
-
-    logging.info("\n2. Dataset loaded successfully!")
-    logging.info(f"   - Number of episodes: {dataset.num_episodes}")
-    logging.info(f"   - Number of frames: {dataset.num_frames}")
-    logging.info(f"   - FPS: {dataset.fps}")
-    logging.info(f"   - Features: {list(dataset.features)}")
-
-    return dataset
-
-
-def load_behavior1k_dataset_with_multiple_modalities(repo_id, root):
-    """Test loading multiple modalities and cameras."""
-    logging.info("\n" + "=" * 80)
-    logging.info("Testing multi-modality loading with repo_id: {repo_id}")
-    logging.info("=" * 80)
-
-    logging.info(f"\n1. Loading dataset with RGB + Depth with repo_id: {repo_id}")
-    dataset = BehaviorLeRobotDatasetV3(
-        repo_id=repo_id,
-        root=root,
-        modalities=["rgb", "depth"],
-        cameras=["head", "left_wrist", "right_wrist"],
-        chunk_streaming_using_keyframe=False,
-        check_timestamp_sync=False,
-        video_backend="pyav",
-    )
-
-    logging.info(f"\n2. Dataset loaded with modalities: {list(dataset.features)}")
-    logging.info(f"   - Total features: {len(dataset.features)}")
-
-    rgb_keys = [k for k in dataset.features if "rgb" in k]
-    depth_keys = [k for k in dataset.features if "depth" in k]
-    logging.info(f"   - RGB features: {rgb_keys}")
-    logging.info(f"   - Depth features: {depth_keys}")
-
-    logging.info("\n3. SUCCESS! Multi-modality loading works.")
-
-    return dataset
-
-
-def stream_behavior1k_dataset(repo_id, root):
-    """Test chunk streaming mode."""
-    logging.info("\n" + "=" * 80)
-    logging.info("Testing chunk streaming mode")
-    logging.info("=" * 80)
-
-    logging.info("\n1. Loading dataset with chunk streaming...")
-    dataset = BehaviorLeRobotDatasetV3(
-        repo_id=repo_id,
-        root=root,
-        modalities=["rgb"],
-        cameras=["head"],
-        chunk_streaming_using_keyframe=True,
-        shuffle=True,
-        seed=42,
-        check_timestamp_sync=False,
-    )
-
-    logging.info("\n2. Dataset loaded in streaming mode")
-    logging.info(f"   - Number of chunks: {len(dataset.chunks)}")
-    logging.info(f"   - First chunk range: {dataset.chunks[0]}")
-
-    logging.info("\n3. Testing frame access in streaming mode...")
-    for i in range(min(3, len(dataset))):
-        frame = dataset[i]
-        logging.info(
-            f"   - Frame {i}: episode_index={frame['episode_index'].item()}, "
-            f"task_index={frame['task_index'].item()}"
-        )
-
-    logging.info("\n4. SUCCESS! Chunk streaming works.")
-
-    return dataset
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--repo-id", type=str, default=None)
-    parser.add_argument("--root", type=str, default=None)
-
-    args = parser.parse_args()
-
-    load_behavior1k_dataset(args.repo_id, args.root)
-    load_behavior1k_dataset_with_multiple_modalities(args.repo_id, args.root)
-    stream_behavior1k_dataset(args.repo_id, args.root)

examples/dataset/prompt.txt

@@ -0,0 +1,20 @@
+Goal:
+Create a high and low policy, the high policy take obs + text and return text, its a VLM
+The low policy is a VLA that take text returned by the high and ouptut actions
+
+
+High policy is ran every one second or when user send prompt
+
+
+Synthetic data generation:
+D demo -> teleop data with a global task annotation
+D label -> segment data into short skills (one to three seconds)
+D syn: p-gen will create the high level prompt user might gave to p-hi 
+Given D label prompt p-gen to imagine appropriate action, take images, ALL PRIOR skill labels in the episode: ℓ̂₀, …, ℓ̂ₜ₋₁
+
+“Given the scene + all previous steps + current needed skill ℓ̂₅,
+generate a user request that logically leads to ℓ̂₅.”
+
+Train:
+Phi(lt| images, global label) cross entropy - next token predictions
+Plow(At| images, qt, lt)

examples/port_datasets/convert_to_video_parallel.sh

@@ -0,0 +1,141 @@
+#!/bin/bash
+#SBATCH --time=24:00:00
+#SBATCH --partition=hopper-cpu
+#SBATCH --cpus-per-task=96
+#SBATCH --output=/fsx/jade_choghari/logs/launcher_%j.out
+#SBATCH --error=/fsx/jade_choghari/logs/launcher_%j.err
+# Activate conda environment
+# Load conda
+source /fsx/jade_choghari/miniforge3/etc/profile.d/conda.sh
+conda activate lerobot
+set -e  # Exit on error
+
+# Input dataset
+REPO_ID="lerobot"
+ROOT="/fsx/jade_choghari/vlabench-primitive/"
+
+# Output paths
+OUTPUT_DIR="/fsx/jade_choghari/vlabench-primitive-encoded"
+OUTPUT_REPO_ID="vlabench-primitive-encoded"
+LOGS_DIR="/fsx/jade_choghari/logs/convert_video"
+
+# Video encoding settings
+VCODEC="libsvtav1"
+PIX_FMT="yuv420p"
+GOP_SIZE=2
+CRF=30
+FAST_DECODE=0
+
+# Parallelization settings
+NUM_WORKERS=24           # Number of parallel SLURM workers
+NUM_IMAGE_WORKERS=4       # Threads per worker for image saving
+
+# SLURM settings
+PARTITION="hopper-cpu"  # Change to your CPU partition name
+CPUS_PER_TASK=32           # CPUs per worker
+MEM_PER_CPU="4G"          # Memory per CPU
+TIME_LIMIT="24:00:00"     # Time limit per job
+
+###############################################################################
+# STEP 1: Parallel Video Conversion
+###############################################################################
+rm -rf "${OUTPUT_DIR}"
+mkdir -p "${OUTPUT_DIR}"
+
+echo "=============================================="
+echo "STEP 1: Starting parallel video conversion"
+echo "  Workers: ${NUM_WORKERS}"
+echo "  Input: ${REPO_ID} (root: ${ROOT})"
+echo "  Output: ${OUTPUT_DIR}"
+echo "=============================================="
+
+python /admin/home/jade_choghari/lerobot/examples/port_datasets/slurm_convert_to_video.py\
+    --repo-id "${REPO_ID}" \
+    --root "${ROOT}" \
+    --output-dir "${OUTPUT_DIR}" \
+    --output-repo-id "${OUTPUT_REPO_ID}" \
+    --vcodec "${VCODEC}" \
+    --pix-fmt "${PIX_FMT}" \
+    --g ${GOP_SIZE} \
+    --crf ${CRF} \
+    --fast-decode ${FAST_DECODE} \
+    --num-image-workers ${NUM_IMAGE_WORKERS} \
+    --logs-dir "${LOGS_DIR}" \
+    --job-name "convert_video" \
+    --slurm 1 \
+    --workers ${NUM_WORKERS} \
+    --partition "${PARTITION}" \
+    --cpus-per-task ${CPUS_PER_TASK} \
+    --mem-per-cpu "${MEM_PER_CPU}" \
+    --time-limit "${TIME_LIMIT}"
+
+echo ""
+echo "✓ Parallel conversion jobs submitted!"
+echo "  Monitor with: squeue -u \$USER"
+echo "  Check logs in: ${LOGS_DIR}/convert_video"
+echo ""
+echo "Wait for all jobs to complete before running Step 2."
+echo "You can check completion with: squeue -u \$USER | grep convert_video"
+echo ""
+echo "After all jobs complete, run Step 2 to aggregate shards:"
+echo "  bash convert_to_video_parallel.sh aggregate"
+
+###############################################################################
+# STEP 2: Aggregate Shards (run this after Step 1 completes)
+###############################################################################
+
+if [ "$1" == "aggregate" ]; then
+    echo ""
+    echo "=============================================="
+    echo "STEP 2: Aggregating video shards"
+    echo "  Shards: ${NUM_WORKERS}"
+    echo "  Input: ${OUTPUT_DIR}/shard_XXXX"
+    echo "  Output: ${OUTPUT_DIR}_final"
+    echo "=============================================="
+
+    python slurm_aggregate_video_shards.py \
+        --shards-dir "${OUTPUT_DIR}" \
+        --output-dir "${OUTPUT_DIR}_final" \
+        --output-repo-id "${OUTPUT_REPO_ID}" \
+        --num-shards ${NUM_WORKERS} \
+        --logs-dir "${LOGS_DIR}" \
+        --job-name "aggregate_video" \
+        --slurm 1 \
+        --partition "${PARTITION}" \
+        --cpus-per-task 16 \
+        --mem-per-cpu "8G" \
+        --time-limit "08:00:00"
+
+    echo ""
+    echo "✓ Aggregation job submitted!"
+    echo "  Monitor with: squeue -u \$USER | grep aggregate_video"
+    echo "  Check logs in: ${LOGS_DIR}/aggregate_video"
+    echo ""
+    echo "After completion, your final dataset will be in:"
+    echo "  ${OUTPUT_DIR}_final"
+fi
+
+###############################################################################
+# Helpful information
+###############################################################################
+
+if [ "$1" != "aggregate" ]; then
+    echo ""
+    echo "=============================================="
+    echo "WORKFLOW SUMMARY"
+    echo "=============================================="
+    echo ""
+    echo "1. Step 1 is now running - it will:"
+    echo "   - Split episodes across ${NUM_WORKERS} workers"
+    echo "   - Each worker converts its episodes to video"
+    echo "   - Creates shard datasets in ${OUTPUT_DIR}/shard_XXXX"
+    echo ""
+    echo "2. After Step 1 completes, run Step 2:"
+    echo "   bash convert_to_video_parallel.sh aggregate"
+    echo ""
+    echo "3. Step 2 will merge all shards into a single dataset"
+    echo ""
+    echo "=============================================="
+fi
+
+

examples/port_datasets/slurm_aggregate_video_shards.py

@@ -0,0 +1,266 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Aggregate video dataset shards into a single dataset.
+
+After parallel conversion using slurm_convert_to_video.py, this script merges
+all the shard datasets into one final dataset.
+
+Example usage:
+    python slurm_aggregate_video_shards.py \
+        --shards-dir /fsx/jade_choghari/libero_video \
+        --output-dir /fsx/jade_choghari/libero_video_final \
+        --output-repo-id lerobot_video \
+        --num-workers 100 \
+        --partition cpu_partition \
+        --cpus-per-task 16
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+from datatrove.executor import LocalPipelineExecutor
+from datatrove.executor.slurm import SlurmPipelineExecutor
+from datatrove.pipeline.base import PipelineStep
+
+
+class AggregateVideoShards(PipelineStep):
+    """Pipeline step that aggregates video dataset shards."""
+
+    def __init__(
+        self,
+        shards_dir: str | Path,
+        output_dir: str | Path,
+        output_repo_id: str,
+        num_shards: int,
+    ):
+        super().__init__()
+        self.shards_dir = Path(shards_dir)
+        self.output_dir = Path(output_dir)
+        self.output_repo_id = output_repo_id
+        self.num_shards = num_shards
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        """Aggregate all shards into a single dataset."""
+        from lerobot.datasets.dataset_tools import merge_datasets
+        from lerobot.datasets.lerobot_dataset import LeRobotDataset
+        from lerobot.utils.utils import init_logging
+
+        init_logging()
+
+        # Only worker 0 performs aggregation
+        if rank != 0:
+            logging.info(f"Worker {rank} skipping - only worker 0 performs aggregation")
+            return
+
+        logging.info(f"Starting aggregation of {self.num_shards} shards")
+
+        # Collect all shard datasets
+        shard_datasets = []
+        for shard_idx in range(self.num_shards):
+            shard_dir = self.shards_dir / f"shard_{shard_idx:04d}"
+            if not shard_dir.exists():
+                logging.warning(f"Shard directory not found: {shard_dir}")
+                continue
+
+            # Find the repo_id for this shard
+            shard_repo_id = f"{self.output_repo_id}_shard_{shard_idx:04d}"
+            try:
+                shard_dataset = LeRobotDataset(shard_repo_id, root=shard_dir)
+                shard_datasets.append(shard_dataset)
+                logging.info(
+                    f"Loaded shard {shard_idx}: {shard_dataset.meta.total_episodes} episodes, "
+                    f"{shard_dataset.meta.total_frames} frames"
+                )
+            except Exception as e:
+                logging.error(f"Failed to load shard {shard_idx}: {e}")
+                continue
+
+        if len(shard_datasets) == 0:
+            raise ValueError(f"No valid shards found in {self.shards_dir}")
+
+        logging.info(f"Successfully loaded {len(shard_datasets)} shards, starting merge")
+
+        # Merge all shards
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        merged_dataset = merge_datasets(
+            shard_datasets,
+            output_repo_id=self.output_repo_id,
+            output_dir=self.output_dir,
+        )
+
+        logging.info("✓ Aggregation complete!")
+        logging.info(f"Merged dataset saved to: {self.output_dir}")
+        logging.info(f"Total episodes: {merged_dataset.meta.total_episodes}")
+        logging.info(f"Total frames: {merged_dataset.meta.total_frames}")
+
+
+def make_aggregate_executor(
+    shards_dir,
+    output_dir,
+    output_repo_id,
+    num_shards,
+    job_name,
+    logs_dir,
+    partition,
+    cpus_per_task,
+    mem_per_cpu,
+    time_limit,
+    slurm=True,
+):
+    """Create executor for shard aggregation."""
+    kwargs = {
+        "pipeline": [
+            AggregateVideoShards(
+                shards_dir=shards_dir,
+                output_dir=output_dir,
+                output_repo_id=output_repo_id,
+                num_shards=num_shards,
+            ),
+        ],
+        "logging_dir": str(logs_dir / job_name),
+    }
+
+    if slurm:
+        # Only need 1 worker for aggregation
+        kwargs.update(
+            {
+                "job_name": job_name,
+                "tasks": 1,
+                "workers": 1,
+                "time": time_limit,
+                "partition": partition,
+                "cpus_per_task": cpus_per_task,
+                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
+            }
+        )
+        executor = SlurmPipelineExecutor(**kwargs)
+    else:
+        kwargs.update(
+            {
+                "tasks": 1,
+                "workers": 1,
+            }
+        )
+        executor = LocalPipelineExecutor(**kwargs)
+
+    return executor
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Aggregate video dataset shards into a single dataset",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--shards-dir",
+        type=Path,
+        required=True,
+        help="Directory containing shard_XXXX subdirectories",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        required=True,
+        help="Output directory for the aggregated dataset",
+    )
+    parser.add_argument(
+        "--output-repo-id",
+        type=str,
+        required=True,
+        help="Repository ID for the aggregated dataset",
+    )
+    parser.add_argument(
+        "--num-shards",
+        type=int,
+        required=True,
+        help="Number of shards to aggregate (should match --workers from conversion)",
+    )
+    parser.add_argument(
+        "--logs-dir",
+        type=Path,
+        required=True,
+        help="Path to logs directory for datatrove",
+    )
+    parser.add_argument(
+        "--job-name",
+        type=str,
+        default="aggregate_video_shards",
+        help="Job name for SLURM",
+    )
+    parser.add_argument(
+        "--slurm",
+        type=int,
+        default=1,
+        help="Launch over SLURM (1) or locally (0)",
+    )
+    parser.add_argument(
+        "--partition",
+        type=str,
+        required=True,
+        help="SLURM partition (use CPU partition)",
+    )
+    parser.add_argument(
+        "--cpus-per-task",
+        type=int,
+        default=16,
+        help="Number of CPUs per task (aggregation can use more CPUs)",
+    )
+    parser.add_argument(
+        "--mem-per-cpu",
+        type=str,
+        default="8G",
+        help="Memory per CPU",
+    )
+    parser.add_argument(
+        "--time-limit",
+        type=str,
+        default="08:00:00",
+        help="Time limit for SLURM job",
+    )
+
+    args = parser.parse_args()
+
+    # Convert slurm flag to boolean
+    slurm = args.slurm == 1
+
+    # Create and run executor
+    executor = make_aggregate_executor(
+        shards_dir=args.shards_dir,
+        output_dir=args.output_dir,
+        output_repo_id=args.output_repo_id,
+        num_shards=args.num_shards,
+        job_name=args.job_name,
+        logs_dir=args.logs_dir,
+        partition=args.partition,
+        cpus_per_task=args.cpus_per_task,
+        mem_per_cpu=args.mem_per_cpu,
+        time_limit=args.time_limit,
+        slurm=slurm,
+    )
+
+    logging.info("Starting shard aggregation")
+    executor.run()
+    logging.info("Aggregation job submitted/completed")
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/port_datasets/slurm_convert_to_video.py

@@ -0,0 +1,539 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Parallelize video conversion using SLURM and Datatrove.
+
+This script converts an image-based LeRobot dataset to video format by distributing 
+episodes across multiple workers for parallel processing.
+
+Example usage:
+    python slurm_convert_to_video.py \
+        --repo-id lerobot \
+        --root /fsx/jade_choghari/libero/ \
+        --output-dir /fsx/jade_choghari/libero_video \
+        --output-repo-id lerobot_video \
+        --workers 100 \
+        --partition cpu_partition \
+        --cpus-per-task 8 \
+        --logs-dir ./logs
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+from datatrove.executor import LocalPipelineExecutor
+from datatrove.executor.slurm import SlurmPipelineExecutor
+from datatrove.pipeline.base import PipelineStep
+
+
+class ConvertEpisodesToVideo(PipelineStep):
+    """Pipeline step that converts episodes to videos in parallel."""
+
+    def __init__(
+        self,
+        repo_id: str,
+        root: str | None,
+        output_dir: str,
+        output_repo_id: str | None,
+        vcodec: str = "libsvtav1",
+        pix_fmt: str = "yuv420p",
+        g: int = 2,
+        crf: int = 30,
+        fast_decode: int = 0,
+        num_image_workers: int = 4,
+    ):
+        super().__init__()
+        self.repo_id = repo_id
+        self.root = root
+        self.output_dir = Path(output_dir)
+        self.output_repo_id = output_repo_id or f"{repo_id}_video"
+        self.vcodec = vcodec
+        self.pix_fmt = pix_fmt
+        self.g = g
+        self.crf = crf
+        self.fast_decode = fast_decode
+        self.num_image_workers = num_image_workers
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        """Process a shard of episodes."""
+        from datasets.utils.tqdm import disable_progress_bars
+
+        from lerobot.datasets.lerobot_dataset import LeRobotDataset
+        from lerobot.scripts.lerobot_edit_dataset import (
+            encode_episode_videos,
+        )
+        from lerobot.utils.utils import init_logging
+        import logging
+        init_logging()
+        disable_progress_bars()
+
+        logging.info(f"Worker {rank}/{world_size} starting video conversion")
+
+        # Load source dataset
+        dataset = LeRobotDataset(self.repo_id, root=self.root)
+        total_episodes = dataset.meta.total_episodes
+
+        # Determine which episodes this worker processes
+        episodes_per_worker = total_episodes // world_size
+        remainder = total_episodes % world_size
+
+        if rank < remainder:
+            # First 'remainder' workers get one extra episode
+            start_ep = rank * (episodes_per_worker + 1)
+            end_ep = start_ep + episodes_per_worker + 1
+        else:
+            start_ep = remainder * (episodes_per_worker + 1) + (rank - remainder) * episodes_per_worker
+            end_ep = start_ep + episodes_per_worker
+
+        episode_indices = list(range(start_ep, end_ep))
+
+        logging.info(
+            f"Worker {rank} processing episodes {start_ep} to {end_ep-1} ({len(episode_indices)} episodes)"
+        )
+
+        if len(episode_indices) == 0:
+            logging.info(f"Worker {rank} has no episodes to process")
+            return
+
+        # Create shard-specific output directory
+        import shutil
+        shard_output_dir = self.output_dir / f"shard_{rank:04d}"
+        
+        # Remove existing directory to avoid conflicts with LeRobotDatasetMetadata.create
+        if shard_output_dir.exists():
+            logging.warning(f"Shard directory {shard_output_dir} already exists, removing it")
+            shutil.rmtree(shard_output_dir)
+
+        # Import conversion function
+        from lerobot.scripts.lerobot_edit_dataset import convert_dataset_to_videos
+
+        logging.info(
+            f"Worker {rank} converting {len(episode_indices)} episodes with codec {self.vcodec}, CRF {self.crf}"
+        )
+
+        # Convert this shard's episodes with remapped indices
+        # We need to remap episode indices to start from 0 for proper file structure
+        self._convert_shard_to_videos(
+            dataset=dataset,
+            shard_output_dir=shard_output_dir,
+            shard_repo_id=f"{self.output_repo_id}_shard_{rank:04d}",
+            episode_indices=episode_indices,
+        )
+
+        logging.info(f"Worker {rank} completed successfully")
+
+    def _convert_shard_to_videos(
+        self,
+        dataset,
+        shard_output_dir,
+        shard_repo_id,
+        episode_indices,
+    ):
+        """Convert a shard's episodes to videos with proper index remapping."""
+        import shutil
+        import pandas as pd
+        from tqdm import tqdm
+
+        from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
+        from lerobot.datasets.utils import write_info, write_stats, write_tasks
+        from lerobot.datasets.video_utils import encode_video_frames, get_video_info
+        from lerobot.scripts.lerobot_edit_dataset import (
+            save_episode_images_for_video,
+            _copy_data_without_images,
+        )
+        from lerobot.utils.constants import OBS_IMAGE
+
+        # Check that it's an image dataset
+        if len(dataset.meta.video_keys) > 0:
+            raise ValueError(
+                f"This operation is for image datasets only. Video dataset provided: {dataset.repo_id}"
+            )
+
+        # Get all image keys
+        hf_dataset = dataset.hf_dataset.with_format(None)
+        img_keys = [key for key in hf_dataset.features if key.startswith(OBS_IMAGE)]
+
+        if len(img_keys) == 0:
+            raise ValueError(f"No image keys found in dataset {dataset.repo_id}")
+
+        logging.info(f"Converting {len(episode_indices)} episodes with {len(img_keys)} cameras")
+
+        # Create new features dict, converting image features to video features
+        new_features = {}
+        for key, value in dataset.meta.features.items():
+            if key not in img_keys:
+                new_features[key] = value
+            else:
+                # Convert image key to video format
+                new_features[key] = value.copy()
+                new_features[key]["dtype"] = "video"
+
+        # Create new metadata for video dataset
+        new_meta = LeRobotDatasetMetadata.create(
+            repo_id=shard_repo_id,
+            fps=dataset.meta.fps,
+            features=new_features,
+            robot_type=dataset.meta.robot_type,
+            root=shard_output_dir,
+            use_videos=True,
+            chunks_size=dataset.meta.chunks_size,
+            data_files_size_in_mb=dataset.meta.data_files_size_in_mb,
+            video_files_size_in_mb=dataset.meta.video_files_size_in_mb,
+        )
+
+        # Create temporary directory for image extraction
+        temp_dir = shard_output_dir / "temp_images"
+        temp_dir.mkdir(parents=True, exist_ok=True)
+
+        # Process each episode with REMAPPED indices (0, 1, 2, ...)
+        all_episode_metadata = []
+        fps = int(dataset.fps)
+
+        try:
+            for new_ep_idx, orig_ep_idx in enumerate(tqdm(episode_indices, desc="Converting episodes")):
+                # Get episode metadata from source using ORIGINAL index
+                src_episode = dataset.meta.episodes[orig_ep_idx]
+                episode_length = src_episode["length"]
+                episode_duration = episode_length / dataset.fps
+
+                video_metadata = {}
+
+                # Encode videos for each camera
+                for img_key in img_keys:
+                    # Save images temporarily using ORIGINAL episode index
+                    imgs_dir = temp_dir / f"episode_{orig_ep_idx:06d}" / img_key
+                    save_episode_images_for_video(
+                        dataset, imgs_dir, img_key, orig_ep_idx, self.num_image_workers
+                    )
+
+                    # Determine chunk and file indices using NEW (remapped) episode index
+                    chunk_idx = new_ep_idx // new_meta.chunks_size
+                    file_idx = new_ep_idx % new_meta.chunks_size
+
+                    # Create video path in the new dataset structure
+                    video_path = new_meta.root / new_meta.video_path.format(
+                        video_key=img_key, chunk_index=chunk_idx, file_index=file_idx
+                    )
+                    video_path.parent.mkdir(parents=True, exist_ok=True)
+
+                    # Encode video
+                    encode_video_frames(
+                        imgs_dir=imgs_dir,
+                        video_path=video_path,
+                        fps=fps,
+                        vcodec=self.vcodec,
+                        pix_fmt=self.pix_fmt,
+                        g=self.g,
+                        crf=self.crf,
+                        fast_decode=self.fast_decode,
+                        overwrite=True,
+                    )
+
+                    # Clean up temporary images
+                    shutil.rmtree(imgs_dir)
+
+                    # Store video metadata
+                    video_metadata[img_key] = {
+                        f"videos/{img_key}/chunk_index": chunk_idx,
+                        f"videos/{img_key}/file_index": file_idx,
+                        f"videos/{img_key}/from_timestamp": 0.0,
+                        f"videos/{img_key}/to_timestamp": episode_duration,
+                    }
+
+                # Build episode metadata using NEW index
+                episode_meta = {
+                    "episode_index": new_ep_idx,
+                    "length": episode_length,
+                    "dataset_from_index": new_ep_idx * episode_length,
+                    "dataset_to_index": (new_ep_idx + 1) * episode_length,
+                }
+
+                # Add video metadata
+                for img_key in img_keys:
+                    episode_meta.update(video_metadata[img_key])
+
+                # Add data chunk/file info
+                if "data/chunk_index" in src_episode:
+                    episode_meta["data/chunk_index"] = src_episode["data/chunk_index"]
+                    episode_meta["data/file_index"] = src_episode["data/file_index"]
+
+                all_episode_metadata.append(episode_meta)
+
+            # Copy and transform data files (removing image columns)
+            _copy_data_without_images(dataset, new_meta, episode_indices, img_keys)
+
+            # Save episode metadata
+            episodes_df = pd.DataFrame(all_episode_metadata)
+            episodes_path = new_meta.root / "meta" / "episodes" / "chunk-000" / "file-000.parquet"
+            episodes_path.parent.mkdir(parents=True, exist_ok=True)
+            episodes_df.to_parquet(episodes_path, index=False)
+
+            # Update metadata info
+            new_meta.info["total_episodes"] = len(episode_indices)
+            new_meta.info["total_frames"] = sum(ep["length"] for ep in all_episode_metadata)
+            new_meta.info["total_tasks"] = dataset.meta.total_tasks
+            new_meta.info["splits"] = {"train": f"0:{len(episode_indices)}"}
+
+            # Update video info for all image keys using the actual first video file
+            for img_key in img_keys:
+                if not new_meta.features[img_key].get("info", None):
+                    # Use the first actually created video file
+                    chunk_idx = all_episode_metadata[0][f"videos/{img_key}/chunk_index"]
+                    file_idx = all_episode_metadata[0][f"videos/{img_key}/file_index"]
+                    video_path = new_meta.root / new_meta.video_path.format(
+                        video_key=img_key, chunk_index=chunk_idx, file_index=file_idx
+                    )
+                    new_meta.info["features"][img_key]["info"] = get_video_info(video_path)
+
+            write_info(new_meta.info, new_meta.root)
+
+            # Copy stats and tasks
+            if dataset.meta.stats is not None:
+                # Remove image stats
+                new_stats = {k: v for k, v in dataset.meta.stats.items() if k not in img_keys}
+                write_stats(new_stats, new_meta.root)
+
+            if dataset.meta.tasks is not None:
+                write_tasks(dataset.meta.tasks, new_meta.root)
+
+        finally:
+            # Clean up temporary directory
+            if temp_dir.exists():
+                shutil.rmtree(temp_dir)
+
+
+def make_convert_executor(
+    repo_id,
+    root,
+    output_dir,
+    output_repo_id,
+    vcodec,
+    pix_fmt,
+    g,
+    crf,
+    fast_decode,
+    num_image_workers,
+    job_name,
+    logs_dir,
+    workers,
+    partition,
+    cpus_per_task,
+    mem_per_cpu,
+    time_limit,
+    slurm=True,
+):
+    """Create executor for parallel video conversion."""
+    kwargs = {
+        "pipeline": [
+            ConvertEpisodesToVideo(
+                repo_id=repo_id,
+                root=root,
+                output_dir=output_dir,
+                output_repo_id=output_repo_id,
+                vcodec=vcodec,
+                pix_fmt=pix_fmt,
+                g=g,
+                crf=crf,
+                fast_decode=fast_decode,
+                num_image_workers=num_image_workers,
+            ),
+        ],
+        "logging_dir": str(logs_dir / job_name),
+    }
+
+    if slurm:
+        kwargs.update(
+            {
+                "job_name": job_name,
+                "tasks": workers,  # Number of parallel tasks = number of workers
+                "workers": workers,
+                "time": time_limit,
+                "partition": partition,
+                "cpus_per_task": cpus_per_task,
+                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
+            }
+        )
+        executor = SlurmPipelineExecutor(**kwargs)
+    else:
+        kwargs.update(
+            {
+                "tasks": workers,
+                "workers": 1,  # Local mode: sequential
+            }
+        )
+        executor = LocalPipelineExecutor(**kwargs)
+
+    return executor
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Parallelize video conversion across SLURM workers",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    # Input/Output paths
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        required=True,
+        help="Repository ID of the source image dataset",
+    )
+    parser.add_argument(
+        "--root",
+        type=str,
+        default=None,
+        help="Root directory containing the source dataset (default: HF cache)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        required=True,
+        help="Output directory for the video dataset",
+    )
+    parser.add_argument(
+        "--output-repo-id",
+        type=str,
+        default=None,
+        help="Repository ID for output dataset (default: <repo_id>_video)",
+    )
+
+    # Video encoding parameters
+    parser.add_argument(
+        "--vcodec",
+        type=str,
+        default="libsvtav1",
+        help="Video codec",
+    )
+    parser.add_argument(
+        "--pix-fmt",
+        type=str,
+        default="yuv420p",
+        help="Pixel format",
+    )
+    parser.add_argument(
+        "--g",
+        type=int,
+        default=2,
+        help="GOP size (group of pictures)",
+    )
+    parser.add_argument(
+        "--crf",
+        type=int,
+        default=30,
+        help="Constant rate factor (quality)",
+    )
+    parser.add_argument(
+        "--fast-decode",
+        type=int,
+        default=0,
+        help="Fast decode tuning",
+    )
+    parser.add_argument(
+        "--num-image-workers",
+        type=int,
+        default=4,
+        help="Number of threads per worker for saving images",
+    )
+
+    # SLURM parameters
+    parser.add_argument(
+        "--logs-dir",
+        type=Path,
+        required=True,
+        help="Path to logs directory for datatrove",
+    )
+    parser.add_argument(
+        "--job-name",
+        type=str,
+        default="convert_to_video",
+        help="Job name for SLURM",
+    )
+    parser.add_argument(
+        "--slurm",
+        type=int,
+        default=1,
+        help="Launch over SLURM (1) or locally (0)",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=100,
+        help="Number of parallel workers (each processes different episodes)",
+    )
+    parser.add_argument(
+        "--partition",
+        type=str,
+        required=True,
+        help="SLURM partition (use CPU partition)",
+    )
+    parser.add_argument(
+        "--cpus-per-task",
+        type=int,
+        default=8,
+        help="Number of CPUs per SLURM task",
+    )
+    parser.add_argument(
+        "--mem-per-cpu",
+        type=str,
+        default="4G",
+        help="Memory per CPU",
+    )
+    parser.add_argument(
+        "--time-limit",
+        type=str,
+        default="08:00:00",
+        help="Time limit for SLURM job",
+    )
+
+    args = parser.parse_args()
+
+    # Convert slurm flag to boolean
+    slurm = args.slurm == 1
+
+    # Create and run executor
+    executor = make_convert_executor(
+        repo_id=args.repo_id,
+        root=args.root,
+        output_dir=args.output_dir,
+        output_repo_id=args.output_repo_id,
+        vcodec=args.vcodec,
+        pix_fmt=args.pix_fmt,
+        g=args.g,
+        crf=args.crf,
+        fast_decode=args.fast_decode,
+        num_image_workers=args.num_image_workers,
+        job_name=args.job_name,
+        logs_dir=args.logs_dir,
+        workers=args.workers,
+        partition=args.partition,
+        cpus_per_task=args.cpus_per_task,
+        mem_per_cpu=args.mem_per_cpu,
+        time_limit=args.time_limit,
+        slurm=slurm,
+    )
+
+    logging.info(f"Starting video conversion with {args.workers} workers")
+    executor.run()
+    logging.info("All workers submitted/completed")
+
+
+if __name__ == "__main__":
+    main()
+
+

slurm-21724787.out

@@ -0,0 +1,7 @@
+==============================================
+STEP 1: Starting parallel video conversion
+  Workers: 2
+  Input: lerobot (root: /fsx/jade_choghari/vlabench-primitive/)
+  Output: /fsx/jade_choghari/vlabench-primitive-encoded
+==============================================
+python: can't open file '/admin/home/jade_choghari/lerobot/slurm_convert_to_video.py': [Errno 2] No such file or directory

src/lerobot/datasets/aggregate.py

@@ -136,21 +136,40 @@ def update_meta_data(
         df["_orig_chunk"] = df[orig_chunk_col].copy()
         df["_orig_file"] = df[orig_file_col].copy()
 
-        # Update chunk and file indices to point to destination
-        df[orig_chunk_col] = video_idx["chunk"]
-        df[orig_file_col] = video_idx["file"]
-
-        # Apply per-source-file timestamp offsets
+        # Get mappings for this video key
         src_to_offset = video_idx.get("src_to_offset", {})
-        if src_to_offset:
-            # Apply offset based on original source file
+        src_to_dst = video_idx.get("src_to_dst", {})
+
+        # Apply per-source-file mappings
+        if src_to_dst:
+            # Map each episode to its correct destination file and apply offset
             for idx in df.index:
-                src_key = (df.at[idx, "_orig_chunk"], df.at[idx, "_orig_file"])
+                # Convert to Python int to avoid numpy type mismatch in dict lookup
+                src_key = (int(df.at[idx, "_orig_chunk"]), int(df.at[idx, "_orig_file"]))
+
+                # Get destination chunk/file for this source file
+                dst_chunk, dst_file = src_to_dst.get(src_key, (video_idx["chunk"], video_idx["file"]))
+                df.at[idx, orig_chunk_col] = dst_chunk
+                df.at[idx, orig_file_col] = dst_file
+
+                # Apply timestamp offset
+                offset = src_to_offset.get(src_key, 0)
+                df.at[idx, f"videos/{key}/from_timestamp"] += offset
+                df.at[idx, f"videos/{key}/to_timestamp"] += offset
+        elif src_to_offset:
+            # Fallback: use same destination for all, but apply per-file offsets
+            df[orig_chunk_col] = video_idx["chunk"]
+            df[orig_file_col] = video_idx["file"]
+            for idx in df.index:
+                # Convert to Python int to avoid numpy type mismatch in dict lookup
+                src_key = (int(df.at[idx, "_orig_chunk"]), int(df.at[idx, "_orig_file"]))
                 offset = src_to_offset.get(src_key, 0)
                 df.at[idx, f"videos/{key}/from_timestamp"] += offset
                 df.at[idx, f"videos/{key}/to_timestamp"] += offset
         else:
             # Fallback to simple offset (for backward compatibility)
+            df[orig_chunk_col] = video_idx["chunk"]
+            df[orig_file_col] = video_idx["file"]
             df[f"videos/{key}/from_timestamp"] = (
                 df[f"videos/{key}/from_timestamp"] + video_idx["latest_duration"]
             )
@@ -268,6 +287,12 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
         videos_idx[key]["episode_duration"] = 0
         # Track offset for each source (chunk, file) pair
         videos_idx[key]["src_to_offset"] = {}
+        # Track destination (chunk, file) for each source (chunk, file) pair
+        videos_idx[key]["src_to_dst"] = {}
+        # Initialize dst_file_durations if not present
+        # dst_file_durations tracks duration of each destination file
+        if "dst_file_durations" not in videos_idx[key]:
+            videos_idx[key]["dst_file_durations"] = {}
 
     for key, video_idx in videos_idx.items():
         unique_chunk_file_pairs = {
@@ -282,9 +307,13 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
 
         chunk_idx = video_idx["chunk"]
         file_idx = video_idx["file"]
-        current_offset = video_idx["latest_duration"]
+        dst_file_durations = video_idx["dst_file_durations"]
 
         for src_chunk_idx, src_file_idx in unique_chunk_file_pairs:
+            # Convert to Python int to ensure consistent dict keys
+            src_chunk_idx = int(src_chunk_idx)
+            src_file_idx = int(src_file_idx)
+
             src_path = src_meta.root / DEFAULT_VIDEO_PATH.format(
                 video_key=key,
                 chunk_index=src_chunk_idx,
@@ -298,14 +327,17 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
             )
 
             src_duration = get_video_duration_in_s(src_path)
+            dst_key = (chunk_idx, file_idx)
 
             if not dst_path.exists():
-                # Store offset before incrementing
-                videos_idx[key]["src_to_offset"][(src_chunk_idx, src_file_idx)] = current_offset
+                # New destination file: offset is 0
+                videos_idx[key]["src_to_offset"][(src_chunk_idx, src_file_idx)] = 0
+                videos_idx[key]["src_to_dst"][(src_chunk_idx, src_file_idx)] = dst_key
                 dst_path.parent.mkdir(parents=True, exist_ok=True)
                 shutil.copy(str(src_path), str(dst_path))
+                # Track duration of this destination file
+                dst_file_durations[dst_key] = src_duration
                 videos_idx[key]["episode_duration"] += src_duration
-                current_offset += src_duration
                 continue
 
             # Check file sizes before appending
@@ -313,10 +345,11 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
             dst_size = get_file_size_in_mb(dst_path)
 
             if dst_size + src_size >= video_files_size_in_mb:
-                # Rotate to a new file, this source becomes start of new destination
-                # So its offset should be 0
-                videos_idx[key]["src_to_offset"][(src_chunk_idx, src_file_idx)] = 0
+                # Rotate to a new file - offset is 0
                 chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, chunk_size)
+                dst_key = (chunk_idx, file_idx)
+                videos_idx[key]["src_to_offset"][(src_chunk_idx, src_file_idx)] = 0
+                videos_idx[key]["src_to_dst"][(src_chunk_idx, src_file_idx)] = dst_key
                 dst_path = dst_meta.root / DEFAULT_VIDEO_PATH.format(
                     video_key=key,
                     chunk_index=chunk_idx,
@@ -324,16 +357,20 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
                 )
                 dst_path.parent.mkdir(parents=True, exist_ok=True)
                 shutil.copy(str(src_path), str(dst_path))
-                # Reset offset for next file
-                current_offset = src_duration
+                # Track duration of this new destination file
+                dst_file_durations[dst_key] = src_duration
             else:
-                # Append to existing video file - use current accumulated offset
-                videos_idx[key]["src_to_offset"][(src_chunk_idx, src_file_idx)] = current_offset
+                # Append to existing destination file
+                # Offset is the current duration of this destination file
+                current_dst_duration = dst_file_durations.get(dst_key, 0)
+                videos_idx[key]["src_to_offset"][(src_chunk_idx, src_file_idx)] = current_dst_duration
+                videos_idx[key]["src_to_dst"][(src_chunk_idx, src_file_idx)] = dst_key
                 concatenate_video_files(
                     [dst_path, src_path],
                     dst_path,
                 )
-                current_offset += src_duration
+                # Update duration of this destination file
+                dst_file_durations[dst_key] = current_dst_duration + src_duration
 
             videos_idx[key]["episode_duration"] += src_duration
 

src/lerobot/scripts/lerobot_edit_dataset.py

@@ -18,7 +18,8 @@
 Edit LeRobot datasets using various transformation tools.
 
 This script allows you to delete episodes, split datasets, merge datasets,
-and remove features. When new_repo_id is specified, creates a new dataset.
+remove features, and convert image datasets to video format.
+When new_repo_id is specified, creates a new dataset.
 
 Usage Examples:
 
@@ -65,6 +66,25 @@ Remove camera feature:
         --operation.type remove_feature \
         --operation.feature_names "['observation.images.top']"
 
+Convert image dataset to video format (saves locally):
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht_image \
+        --operation.type convert_to_video \
+        --operation.output_dir /path/to/output/pusht_video
+
+Convert image dataset and save with new repo_id:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht_image \
+        --new_repo_id lerobot/pusht_video \
+        --operation.type convert_to_video
+
+Convert and push to hub:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht_image \
+        --new_repo_id lerobot/pusht_video \
+        --operation.type convert_to_video \
+        --push_to_hub true
+
 Using JSON config file:
     python -m lerobot.scripts.lerobot_edit_dataset \
         --config_path path/to/edit_config.json
@@ -72,9 +92,13 @@ Using JSON config file:
 
 import logging
 import shutil
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
 from pathlib import Path
 
+import pandas as pd
+from tqdm import tqdm
+
 from lerobot.configs import parser
 from lerobot.datasets.dataset_tools import (
     delete_episodes,
@@ -82,8 +106,10 @@ from lerobot.datasets.dataset_tools import (
     remove_feature,
     split_dataset,
 )
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.utils.constants import HF_LEROBOT_HOME
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.datasets.utils import write_stats, write_tasks
+from lerobot.datasets.video_utils import encode_video_frames, get_video_info
+from lerobot.utils.constants import HF_LEROBOT_HOME, OBS_IMAGE
 from lerobot.utils.utils import init_logging
 
 
@@ -111,10 +137,23 @@ class RemoveFeatureConfig:
     feature_names: list[str] | None = None
 
 
+@dataclass
+class ConvertToVideoConfig:
+    type: str = "convert_to_video"
+    output_dir: str | None = None
+    vcodec: str = "libsvtav1"
+    pix_fmt: str = "yuv420p"
+    g: int = 2
+    crf: int = 30
+    fast_decode: int = 0
+    episode_indices: list[int] | None = None
+    num_workers: int = 4
+
+
 @dataclass
 class EditDatasetConfig:
     repo_id: str
-    operation: DeleteEpisodesConfig | SplitConfig | MergeConfig | RemoveFeatureConfig
+    operation: DeleteEpisodesConfig | SplitConfig | MergeConfig | RemoveFeatureConfig | ConvertToVideoConfig
     root: str | None = None
     new_repo_id: str | None = None
     push_to_hub: bool = False
@@ -258,6 +297,415 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
         LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
 
 
+def save_episode_images_for_video(
+    dataset: LeRobotDataset,
+    imgs_dir: Path,
+    img_key: str,
+    episode_index: int,
+    num_workers: int = 4,
+) -> None:
+    """Save images from a specific episode and camera to disk for video encoding.
+
+    Args:
+        dataset: The LeRobot dataset to extract images from
+        imgs_dir: Directory to save images to
+        img_key: The image key (camera) to extract
+        episode_index: Index of the episode to save
+        num_workers: Number of threads for parallel image saving
+    """
+    # Create directory
+    imgs_dir.mkdir(parents=True, exist_ok=True)
+
+    # Get dataset without torch format for PIL image access
+    hf_dataset = dataset.hf_dataset.with_format(None)
+
+    # Select only this camera's images
+    imgs_dataset = hf_dataset.select_columns(img_key)
+
+    # Get episode start and end indices
+    from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
+    to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
+
+    # Get all items for this episode
+    episode_dataset = imgs_dataset.select(range(from_idx, to_idx))
+
+    # Define function to save a single image
+    def save_single_image(i_item_tuple):
+        i, item = i_item_tuple
+        img = item[img_key]
+        # Use frame-XXXXXX.png format to match encode_video_frames expectations
+        img.save(str(imgs_dir / f"frame-{i:06d}.png"), quality=100)
+        return i
+
+    # Save images with proper naming convention for encode_video_frames (frame-XXXXXX.png)
+    items = list(enumerate(episode_dataset))
+
+    with ThreadPoolExecutor(max_workers=num_workers) as executor:
+        futures = [executor.submit(save_single_image, item) for item in items]
+        for future in as_completed(futures):
+            future.result()  # This will raise any exceptions that occurred
+
+
+def encode_episode_videos(
+    dataset: LeRobotDataset,
+    new_meta: LeRobotDatasetMetadata,
+    episode_index: int,
+    vcodec: str,
+    pix_fmt: str,
+    g: int,
+    crf: int,
+    fast_decode: int,
+    temp_dir: Path,
+    num_image_workers: int = 4,
+) -> dict[str, dict]:
+    """Encode videos for a single episode and return video metadata.
+
+    Args:
+        dataset: Source dataset with images
+        new_meta: Metadata object for the new video dataset
+        episode_index: Episode index to process
+        vcodec: Video codec
+        pix_fmt: Pixel format
+        g: Group of pictures size
+        crf: Constant rate factor
+        fast_decode: Fast decode tuning
+        temp_dir: Temporary directory for images
+        num_image_workers: Number of workers for saving images
+
+    Returns:
+        Dictionary mapping video keys to their metadata (chunk_index, file_index, timestamps)
+    """
+    hf_dataset = dataset.hf_dataset.with_format(None)
+    img_keys = [key for key in hf_dataset.features if key.startswith(OBS_IMAGE)]
+
+    video_metadata = {}
+    fps = int(dataset.fps)  # Convert to int for PyAV compatibility
+    episode_length = dataset.meta.episodes["length"][episode_index]
+    episode_duration = episode_length / dataset.fps  # Use original fps for duration calculation
+
+    for img_key in img_keys:
+        # Save images temporarily
+        imgs_dir = temp_dir / f"episode_{episode_index:06d}" / img_key
+        save_episode_images_for_video(dataset, imgs_dir, img_key, episode_index, num_image_workers)
+
+        # Determine chunk and file indices
+        # For simplicity, we'll put each episode in its own file
+        chunk_idx = episode_index // new_meta.chunks_size
+        file_idx = episode_index % new_meta.chunks_size
+
+        # Create video path in the new dataset structure
+        video_path = new_meta.root / new_meta.video_path.format(
+            video_key=img_key, chunk_index=chunk_idx, file_index=file_idx
+        )
+        video_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Encode video
+        encode_video_frames(
+            imgs_dir=imgs_dir,
+            video_path=video_path,
+            fps=fps,
+            vcodec=vcodec,
+            pix_fmt=pix_fmt,
+            g=g,
+            crf=crf,
+            fast_decode=fast_decode,
+            overwrite=True,
+        )
+
+        # Clean up temporary images
+        shutil.rmtree(imgs_dir)
+
+        # Store video metadata
+        video_metadata[img_key] = {
+            f"videos/{img_key}/chunk_index": chunk_idx,
+            f"videos/{img_key}/file_index": file_idx,
+            f"videos/{img_key}/from_timestamp": 0.0,
+            f"videos/{img_key}/to_timestamp": episode_duration,
+        }
+
+    return video_metadata
+
+
+def convert_dataset_to_videos(
+    dataset: LeRobotDataset,
+    output_dir: Path,
+    repo_id: str | None = None,
+    vcodec: str = "libsvtav1",
+    pix_fmt: str = "yuv420p",
+    g: int = 2,
+    crf: int = 30,
+    fast_decode: int = 0,
+    episode_indices: list[int] | None = None,
+    num_workers: int = 4,
+) -> LeRobotDataset:
+    """Convert image-based dataset to video-based dataset.
+
+    Creates a new LeRobotDataset with videos instead of images, following the proper
+    LeRobot dataset structure with videos stored in chunked MP4 files.
+
+    Args:
+        dataset: The source LeRobot dataset with images
+        output_dir: Directory to save the new video dataset
+        repo_id: Repository ID for the new dataset (default: original_id + "_video")
+        vcodec: Video codec (default: libsvtav1)
+        pix_fmt: Pixel format (default: yuv420p)
+        g: Group of pictures size (default: 2)
+        crf: Constant rate factor (default: 30)
+        fast_decode: Fast decode tuning (default: 0)
+        episode_indices: List of episode indices to convert (None = all episodes)
+        num_workers: Number of threads for parallel processing (default: 4)
+
+    Returns:
+        New LeRobotDataset with videos
+    """
+    # Check that it's an image dataset
+    if len(dataset.meta.video_keys) > 0:
+        raise ValueError(
+            f"This operation is for image datasets only. Video dataset provided: {dataset.repo_id}"
+        )
+
+    # Get all image keys
+    hf_dataset = dataset.hf_dataset.with_format(None)
+    img_keys = [key for key in hf_dataset.features if key.startswith(OBS_IMAGE)]
+
+    if len(img_keys) == 0:
+        raise ValueError(f"No image keys found in dataset {dataset.repo_id}")
+
+    # Determine which episodes to process
+    if episode_indices is None:
+        episode_indices = list(range(dataset.meta.total_episodes))
+
+    if repo_id is None:
+        repo_id = f"{dataset.repo_id}_video"
+
+    logging.info(
+        f"Converting {len(episode_indices)} episodes with {len(img_keys)} cameras from {dataset.repo_id}"
+    )
+    logging.info(f"Video codec: {vcodec}, pixel format: {pix_fmt}, GOP: {g}, CRF: {crf}")
+
+    # Create new features dict, converting image features to video features
+    new_features = {}
+    for key, value in dataset.meta.features.items():
+        if key not in img_keys:
+            new_features[key] = value
+        else:
+            # Convert image key to video format
+            new_features[key] = value.copy()
+            new_features[key]["dtype"] = "video"  # Change dtype from "image" to "video"
+            # Video info will be updated after episodes are encoded
+
+    # Create new metadata for video dataset
+    new_meta = LeRobotDatasetMetadata.create(
+        repo_id=repo_id,
+        fps=dataset.meta.fps,
+        features=new_features,
+        robot_type=dataset.meta.robot_type,
+        root=output_dir,
+        use_videos=True,
+        chunks_size=dataset.meta.chunks_size,
+        data_files_size_in_mb=dataset.meta.data_files_size_in_mb,
+        video_files_size_in_mb=dataset.meta.video_files_size_in_mb,
+    )
+
+    # Create temporary directory for image extraction
+    temp_dir = output_dir / "temp_images"
+    temp_dir.mkdir(parents=True, exist_ok=True)
+
+    # Process each episode
+    all_episode_metadata = []
+
+    try:
+        for ep_idx in tqdm(episode_indices, desc="Converting episodes to videos"):
+            # Get episode metadata from source
+            src_episode = dataset.meta.episodes[ep_idx]
+
+            # Encode videos for this episode
+            video_metadata = encode_episode_videos(
+                dataset=dataset,
+                new_meta=new_meta,
+                episode_index=ep_idx,
+                vcodec=vcodec,
+                pix_fmt=pix_fmt,
+                g=g,
+                crf=crf,
+                fast_decode=fast_decode,
+                temp_dir=temp_dir,
+                num_image_workers=num_workers,
+            )
+
+            # Build episode metadata
+            episode_meta = {
+                "episode_index": ep_idx,
+                "length": src_episode["length"],
+                "dataset_from_index": ep_idx * src_episode["length"],
+                "dataset_to_index": (ep_idx + 1) * src_episode["length"],
+            }
+
+            # Add video metadata
+            for img_key in img_keys:
+                episode_meta.update(video_metadata[img_key])
+
+            # Add data chunk/file info (using same structure as source)
+            if "data/chunk_index" in src_episode:
+                episode_meta["data/chunk_index"] = src_episode["data/chunk_index"]
+                episode_meta["data/file_index"] = src_episode["data/file_index"]
+
+            all_episode_metadata.append(episode_meta)
+
+        # Copy and transform data files (removing image columns)
+        _copy_data_without_images(dataset, new_meta, episode_indices, img_keys)
+
+        # Save episode metadata
+        episodes_df = pd.DataFrame(all_episode_metadata)
+        episodes_path = new_meta.root / "meta" / "episodes" / "chunk-000" / "file-000.parquet"
+        episodes_path.parent.mkdir(parents=True, exist_ok=True)
+        episodes_df.to_parquet(episodes_path, index=False)
+
+        # Update metadata info
+        new_meta.info["total_episodes"] = len(episode_indices)
+        new_meta.info["total_frames"] = sum(ep["length"] for ep in all_episode_metadata)
+        new_meta.info["total_tasks"] = dataset.meta.total_tasks
+        new_meta.info["splits"] = {"train": f"0:{len(episode_indices)}"}
+
+        # Update video info for all image keys (now videos)
+        # We need to manually set video info since update_video_info() checks video_keys first
+        for img_key in img_keys:
+            if not new_meta.features[img_key].get("info", None):
+                video_path = new_meta.root / new_meta.video_path.format(
+                    video_key=img_key, chunk_index=0, file_index=0
+                )
+                new_meta.info["features"][img_key]["info"] = get_video_info(video_path)
+
+        from lerobot.datasets.utils import write_info
+
+        write_info(new_meta.info, new_meta.root)
+
+        # Copy stats and tasks
+        if dataset.meta.stats is not None:
+            # Remove image stats
+            new_stats = {k: v for k, v in dataset.meta.stats.items() if k not in img_keys}
+            write_stats(new_stats, new_meta.root)
+
+        if dataset.meta.tasks is not None:
+            write_tasks(dataset.meta.tasks, new_meta.root)
+
+    finally:
+        # Clean up temporary directory
+        if temp_dir.exists():
+            shutil.rmtree(temp_dir)
+
+    logging.info(f"✓ Completed converting {dataset.repo_id} to video format")
+    logging.info(f"New dataset saved to: {output_dir}")
+
+    # Return new dataset
+    return LeRobotDataset(repo_id=repo_id, root=output_dir)
+
+
+def _copy_data_without_images(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    episode_indices: list[int],
+    img_keys: list[str],
+) -> None:
+    """Copy data files without image columns.
+
+    Args:
+        src_dataset: Source dataset
+        dst_meta: Destination metadata
+        episode_indices: Episodes to include
+        img_keys: Image keys to remove
+    """
+    from lerobot.datasets.utils import DATA_DIR
+
+    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}")
+
+    episode_set = set(episode_indices)
+
+    for src_path in tqdm(parquet_files, desc="Processing data files"):
+        df = pd.read_parquet(src_path).reset_index(drop=True)
+
+        # Filter to only include selected episodes
+        df = df[df["episode_index"].isin(episode_set)].copy()
+
+        if len(df) == 0:
+            continue
+
+        # Remove image columns
+        columns_to_drop = [col for col in img_keys if col in df.columns]
+        if columns_to_drop:
+            df = df.drop(columns=columns_to_drop)
+
+        # Get chunk and file indices from path
+        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])
+
+        # Write to destination without pandas index
+        dst_path = dst_meta.root / f"data/chunk-{chunk_idx:03d}/file-{file_idx:03d}.parquet"
+        dst_path.parent.mkdir(parents=True, exist_ok=True)
+        df.to_parquet(dst_path, index=False)
+
+
+def handle_convert_to_video(cfg: EditDatasetConfig) -> None:
+    # Note: Parser may create any config type with the right fields, so we access fields directly
+    # instead of checking isinstance()
+    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
+    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
+        output_repo_id = cfg.new_repo_id
+        output_dir = Path(cfg.root) / cfg.new_repo_id if cfg.root else HF_LEROBOT_HOME / cfg.new_repo_id
+        logging.info(f"Saving to new dataset: {cfg.new_repo_id}")
+    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"
+        output_dir = Path(cfg.root) / output_repo_id if cfg.root else 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")
+
+    new_dataset = convert_dataset_to_videos(
+        dataset=dataset,
+        output_dir=output_dir,
+        repo_id=output_repo_id,
+        vcodec=getattr(cfg.operation, "vcodec", "libsvtav1"),
+        pix_fmt=getattr(cfg.operation, "pix_fmt", "yuv420p"),
+        g=getattr(cfg.operation, "g", 2),
+        crf=getattr(cfg.operation, "crf", 30),
+        fast_decode=getattr(cfg.operation, "fast_decode", 0),
+        episode_indices=getattr(cfg.operation, "episode_indices", None),
+        num_workers=getattr(cfg.operation, "num_workers", 4),
+    )
+
+    logging.info("Video dataset created successfully!")
+    logging.info(f"Location: {output_dir}")
+    logging.info(f"Episodes: {new_dataset.meta.total_episodes}")
+    logging.info(f"Frames: {new_dataset.meta.total_frames}")
+
+    if cfg.push_to_hub:
+        logging.info(f"Pushing to hub as {output_repo_id}...")
+        new_dataset.push_to_hub()
+        logging.info("✓ Successfully pushed to hub!")
+    else:
+        logging.info("Dataset saved locally (not pushed to hub)")
+
+
 @parser.wrap()
 def edit_dataset(cfg: EditDatasetConfig) -> None:
     operation_type = cfg.operation.type
@@ -270,10 +718,12 @@ def edit_dataset(cfg: EditDatasetConfig) -> None:
         handle_merge(cfg)
     elif operation_type == "remove_feature":
         handle_remove_feature(cfg)
+    elif operation_type == "convert_to_video":
+        handle_convert_to_video(cfg)
     else:
         raise ValueError(
             f"Unknown operation type: {operation_type}\n"
-            f"Available operations: delete_episodes, split, merge, remove_feature"
+            f"Available operations: delete_episodes, split, merge, remove_feature, convert_to_video"
         )
 
 

tests/datasets/test_dataset_tools.py

@@ -29,6 +29,7 @@ from lerobot.datasets.dataset_tools import (
     remove_feature,
     split_dataset,
 )
+from lerobot.scripts.lerobot_edit_dataset import convert_dataset_to_videos
 
 
 @pytest.fixture
@@ -1047,3 +1048,107 @@ def test_modify_features_preserves_file_structure(sample_dataset, tmp_path):
         assert new_chunk_indices == original_chunk_indices, "Chunk indices should be preserved"
         assert new_file_indices == original_file_indices, "File indices should be preserved"
         assert "reward" in modified_dataset.meta.features
+
+
+def test_convert_dataset_to_videos(tmp_path):
+    """Test converting lerobot/pusht_image dataset to video format."""
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+    # Load the actual lerobot/pusht_image dataset (only first 2 episodes for speed)
+    source_dataset = LeRobotDataset("lerobot/pusht_image", episodes=[0, 1])
+
+    output_dir = tmp_path / "pusht_video"
+
+    with (
+        patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
+        patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
+    ):
+        mock_get_safe_version.return_value = "v3.0"
+        mock_snapshot_download.return_value = str(output_dir)
+
+        # Verify source dataset has images, not videos
+        assert len(source_dataset.meta.video_keys) == 0
+        assert "observation.image" in source_dataset.meta.features
+
+        # Convert to video dataset (only first 2 episodes for speed)
+        video_dataset = convert_dataset_to_videos(
+            dataset=source_dataset,
+            output_dir=output_dir,
+            repo_id="lerobot/pusht_video",
+            vcodec="libsvtav1",
+            pix_fmt="yuv420p",
+            g=2,
+            crf=30,
+            episode_indices=[0, 1],
+            num_workers=2,
+        )
+
+        # Verify new dataset has videos
+        assert len(video_dataset.meta.video_keys) > 0
+        assert "observation.image" in video_dataset.meta.video_keys
+
+        # Verify correct number of episodes and frames (2 episodes)
+        assert video_dataset.meta.total_episodes == 2
+        # Compare against the actual number of frames in the loaded episodes, not metadata total
+        assert len(video_dataset) == len(source_dataset)
+
+        # Verify video files exist
+        for ep_idx in range(video_dataset.meta.total_episodes):
+            for video_key in video_dataset.meta.video_keys:
+                video_path = video_dataset.root / video_dataset.meta.get_video_file_path(ep_idx, video_key)
+                assert video_path.exists(), f"Video file should exist: {video_path}"
+
+        # Verify we can load the dataset and access it
+        assert len(video_dataset) == video_dataset.meta.total_frames
+
+        # Test that we can actually get an item from the video dataset
+        item = video_dataset[0]
+        assert "observation.image" in item
+        assert "action" in item
+
+        # Cleanup
+        import shutil
+
+        if output_dir.exists():
+            shutil.rmtree(output_dir)
+
+
+def test_convert_dataset_to_videos_subset_episodes(tmp_path):
+    """Test converting only specific episodes from lerobot/pusht_image to video format."""
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+    # Load the actual lerobot/pusht_image dataset (only first 3 episodes)
+    source_dataset = LeRobotDataset("lerobot/pusht_image", episodes=[0, 1, 2])
+
+    output_dir = tmp_path / "pusht_video_subset"
+
+    with (
+        patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
+        patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
+    ):
+        mock_get_safe_version.return_value = "v3.0"
+        mock_snapshot_download.return_value = str(output_dir)
+
+        # Convert only episode 0 to video (subset of loaded episodes)
+        episode_indices = [0]
+
+        video_dataset = convert_dataset_to_videos(
+            dataset=source_dataset,
+            output_dir=output_dir,
+            repo_id="lerobot/pusht_video_subset",
+            episode_indices=episode_indices,
+            num_workers=2,
+        )
+
+        # Verify correct number of episodes
+        assert video_dataset.meta.total_episodes == len(episode_indices)
+
+        # Verify video files exist for selected episodes
+        assert len(video_dataset.meta.video_keys) > 0
+        assert "observation.image" in video_dataset.meta.video_keys
+
+        # Cleanup
+        import shutil
+
+        if output_dir.exists():
+            shutil.rmtree(output_dir)