fix port

fix import
add download script
2026-07-29 04:36:04 +00:00 · 2025-09-08 18:58:42 +02:00 · 2025-09-08 18:49:26 +02:00 · 2025-09-08 17:57:04 +02:00 · 2025-09-08 13:40:47 +02:00 · 2025-09-08 12:44:43 +02:00
41 changed files with 2986 additions and 5552 deletions
@@ -29,7 +29,7 @@ ENV DEBIAN_FRONTEND=noninteractive \

 # Install system dependencies and uv (as root)
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    build-essential git curl libglib2.0-0 libegl1-mesa ffmpeg \
+    build-essential git curl libglib2.0-0 libegl1-mesa-dev ffmpeg \
    libusb-1.0-0-dev speech-dispatcher libgeos-dev portaudio19-dev \
    && curl -LsSf https://astral.sh/uv/install.sh | sh \
    && mv /root/.local/bin/uv /usr/local/bin/uv \
@@ -150,7 +150,7 @@ gsutil -m cp -r gs://gresearch/robotics/droid_100 /your/data/
 ### Step 3: Port the Dataset

 ```bash
-python examples/port_datasets/droid_rlds/port_droid.py \
+python examples/port_datasets/port_droid_rlds.py \
    --raw-dir /your/data/droid/1.0.1 \
    --repo-id your_id/droid_1.0.1 \
    --push-to-hub
@@ -161,7 +161,7 @@ python examples/port_datasets/droid_rlds/port_droid.py \
 For development, you can port a single shard:

 ```bash
-python examples/port_datasets/droid_rlds/port_droid.py \
+python examples/port_datasets/port_droid_rlds.py \
    --raw-dir /your/data/droid/1.0.1 \
    --repo-id your_id/droid_1.0.1_test \
    --num-shards 2048 \
@@ -194,7 +194,7 @@ Choose a **CPU partition** - no GPU needed for dataset porting.
 ### Step 3: Launch Parallel Porting Jobs

 ```bash
-python examples/port_datasets/droid_rlds/slurm_port_shards.py \
+python examples/port_datasets/slurm_port_shards.py \
    --raw-dir /your/data/droid/1.0.1 \
    --repo-id your_id/droid_1.0.1 \
    --logs-dir /your/logs \
@@ -245,7 +245,7 @@ failed_logs /your/logs/port_droid
 Once all porting jobs complete:

 ```bash
-python examples/port_datasets/droid_rlds/slurm_aggregate_shards.py \
+python examples/port_datasets/slurm_aggregate_shards.py \
    --repo-id your_id/droid_1.0.1 \
    --logs-dir /your/logs \
    --job-name aggr_droid \
@@ -258,7 +258,7 @@ python examples/port_datasets/droid_rlds/slurm_aggregate_shards.py \
 ### Step 6: Upload to Hub

 ```bash
-python examples/port_datasets/droid_rlds/slurm_upload.py \
+python examples/port_datasets/slurm_upload.py \
    --repo-id your_id/droid_1.0.1 \
    --logs-dir /your/logs \
    --job-name upload_droid \
@@ -1,503 +0,0 @@
-import json
-import logging
-import shutil
-import time
-from pathlib import Path
-
-import h5py
-import numpy as np
-import pandas as pd
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.utils import (
-    DEFAULT_CHUNK_SIZE,
-    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
-    DEFAULT_VIDEO_PATH,
-    EPISODES_DIR,
-    get_video_duration_in_s,
-    get_video_size_in_mb,
-    update_chunk_file_indices,
-    write_info,
-)
-from lerobot.datasets.video_utils import concat_video_files
-from lerobot.utils.utils import get_elapsed_time_in_days_hours_minutes_seconds
-
-AGIBOT_FPS = 30
-AGIBOT_ROBOT_TYPE = "AgiBot_A2D"
-AGIBOT_FEATURES = {
-    # gripper open range in mm (0 for pull open, 1 for full close)
-    "observation.state.effector.position": {
-        "dtype": "float32",
-        "shape": (2,),
-        "names": {
-            "axes": ["left_gripper", "right_gripper"],
-        },
-    },
-    # flange xyz in meters
-    "observation.state.end.position": {
-        "dtype": "float32",
-        "shape": (6,),
-        "names": {
-            "axes": ["left_x", "left_y", "left_z", "right_x", "right_y", "right_z"],
-        },
-    },
-    # flange quaternion with xyzw
-    "observation.state.end.orientation": {
-        "dtype": "float32",
-        "shape": (8,),
-        "names": {
-            "axes": ["left_x", "left_y", "left_z", "left_w", "right_x", "right_y", "right_z", "right_w"],
-        },
-    },
-    # in radians
-    "observation.state.head.position": {
-        "dtype": "float32",
-        "shape": (2,),
-        "names": {
-            "axes": ["yaw", "pitch"],
-        },
-    },
-    # in motor steps
-    "observation.state.joint.current_value": {
-        "dtype": "float32",
-        "shape": (14,),
-        "names": {
-            "axes": [f"left_joint_{i}" for i in range(7)] + [f"right_joint_{i}" for i in range(7)],
-        },
-    },
-    # same as current_value but in radians
-    "observation.state.joint.position": {
-        "dtype": "float32",
-        "shape": (14,),
-        "names": {
-            "axes": [f"left_joint_{i}" for i in range(7)] + [f"right_joint_{i}" for i in range(7)],
-        },
-    },
-    # pitch in radians, lift in meters
-    "observation.state.waist.position": {
-        "dtype": "float32",
-        "shape": (2,),
-        "names": {
-            "axes": ["pitch", "lift"],
-        },
-    },
-    # concatenation of head.position, joint.position, effector.position, waist.position
-    "observation.state": {
-        "dtype": "float32",
-        "shape": (20,),
-        "names": {
-            "axes": ["head_yaw", "head_pitch"]
-            + [f"left_joint_{i}" for i in range(7)]
-            + ["left_gripper"]
-            + [f"right_joint_{i}" for i in range(7)]
-            + ["right_gripper"]
-            + ["waist_pitch", "waist_lift"],
-        },
-    },
-    # gripper open range in mm (0 for pull open, 1 for full close)
-    "action.effector.position": {
-        "dtype": "float32",
-        "shape": (2,),
-        "names": {
-            "axes": ["left_gripper", "right_gripper"],
-        },
-    },
-    # flange xyz in meters
-    "action.end.position": {
-        "dtype": "float32",
-        "shape": (6,),
-        "names": {
-            "axes": ["left_x", "left_y", "left_z", "right_x", "right_y", "right_z"],
-        },
-    },
-    # flange quaternion with xyzw
-    "action.end.orientation": {
-        "dtype": "float32",
-        "shape": (8,),
-        "names": {
-            "axes": ["left_x", "left_y", "left_z", "left_w", "right_x", "right_y", "right_z", "right_w"],
-        },
-    },
-    # in radians
-    "action.head.position": {
-        "dtype": "float32",
-        "shape": (2,),
-        "names": {
-            "axes": ["yaw", "pitch"],
-        },
-    },
-    # goal joint position in radians
-    "action.joint.position": {
-        "dtype": "float32",
-        "shape": (14,),
-        "names": {
-            "axes": [f"left_joint_{i}" for i in range(7)] + [f"right_joint_{i}" for i in range(7)],
-        },
-    },
-    "action.robot.velocity": {
-        "dtype": "float32",
-        "shape": (2,),
-        "names": {
-            "axes": ["velocity_x", "yaw_rate"],
-        },
-    },
-    # pitch in radians, lift in meters
-    "action.waist.position": {
-        "dtype": "float32",
-        "shape": (2,),
-        "names": {
-            "axes": ["pitch", "lift"],
-        },
-    },
-    # concatenation of head.position, joint.position, effector.position, waist.position, robot.velocity
-    "action": {
-        "dtype": "float32",
-        "shape": (22,),
-        "names": {
-            "axes": ["head_yaw", "head_pitch"]
-            + [f"left_joint_{i}" for i in range(7)]
-            + ["left_gripper"]
-            + [f"right_joint_{i}" for i in range(7)]
-            + ["right_gripper"]
-            + ["waist_pitch", "waist_lift"]
-            + ["velocity_x", "yaw_rate"],
-        },
-    },
-    # episode level annotation
-    "init_scene_text": {
-        "dtype": "string",
-        "shape": (1,),
-        "names": None,
-    },
-    # frame level annotation
-    "action_text": {
-        "dtype": "string",
-        "shape": (1,),
-        "names": None,
-    },
-    # frame level annotation
-    "skill": {
-        "dtype": "string",
-        "shape": (1,),
-        "names": None,
-    },
-}
-
-AGIBOT_IMAGES_FEATURES = {
-    "observation.images.top_head": {
-        "dtype": "video",
-        "shape": (480, 640, 3),
-        "names": ["height", "width", "channel"],
-    },
-    "observation.images.hand_left": {
-        "dtype": "video",
-        "shape": (480, 640, 3),
-        "names": ["height", "width", "channel"],
-    },
-    "observation.images.hand_right": {
-        "dtype": "video",
-        "shape": (480, 640, 3),
-        "names": ["height", "width", "channel"],
-    },
-    "observation.images.head_center_fisheye": {
-        "dtype": "video",
-        "shape": (748, 960, 3),
-        "names": ["height", "width", "channel"],
-    },
-    "observation.images.head_left_fisheye": {
-        "dtype": "video",
-        "shape": (748, 960, 3),
-        "names": ["height", "width", "channel"],
-    },
-    "observation.images.head_right_fisheye": {
-        "dtype": "video",
-        "shape": (748, 960, 3),
-        "names": ["height", "width", "channel"],
-    },
-    "observation.images.back_left_fisheye": {
-        "dtype": "video",
-        "shape": (748, 960, 3),
-        "names": ["height", "width", "channel"],
-    },
-    "observation.images.back_right_fisheye": {
-        "dtype": "video",
-        "shape": (748, 960, 3),
-        "names": ["height", "width", "channel"],
-    },
-}
-
-
-def load_info_per_task(raw_dir):
-    info_per_task = {}
-    task_info_dir = raw_dir / "task_info"
-    for path in task_info_dir.glob("task_*.json"):
-        task_index = int(path.name.replace("task_", "").replace(".json", ""))
-        with open(path) as f:
-            task_info = json.load(f)
-
-        task_info = {ep["episode_id"]: ep for ep in task_info}
-        info_per_task[task_index] = task_info
-
-    return info_per_task
-
-
-def create_frame_idx_to_frames_label_idx(ep_info):
-    frame_idx_to_frames_label_idx = {}
-    for label_idx, frames_label in enumerate(ep_info["label_info"]["action_config"]):
-        for frame_idx in range(frames_label["start_frame"], frames_label["end_frame"]):
-            frame_idx_to_frames_label_idx[frame_idx] = label_idx
-    return frame_idx_to_frames_label_idx
-
-
-def generate_lerobot_frames(raw_dir: Path, task_index: int, episode_index: int):
-    r"""/!\ The frames dont contain observation.cameras.*"""
-    info_per_task = load_info_per_task(raw_dir)
-    ep_info = info_per_task[task_index][episode_index]
-    frame_idx_to_frames_label_idx = create_frame_idx_to_frames_label_idx(ep_info)
-
-    # Empty features are commented out.
-    keys_mapping = {
-        # STATE
-        # "observation.state.effector.force": "state/effector/force",
-        "observation.state.effector.position": "state/effector/position",
-        # "observation.state.end.angular": "state/end/angular",
-        "observation.state.end.position": "state/end/position",
-        "observation.state.end.orientation": "state/end/orientation",
-        # "observation.state.end.velocity": "state/end/velocity",
-        # "observation.state.end.wrench": "state/end/wrench",
-        # "observation.state.head.effort": "state/head/effort",
-        "observation.state.head.position": "state/head/position",
-        # "observation.state.head.velocity": "state/head/velocity",
-        "observation.state.joint.current_value": "state/joint/current_value",
-        # "observation.state.joint.effort": "state/joint/effort",
-        "observation.state.joint.position": "state/joint/position",
-        # "observation.state.joint.velocity": "state/joint/velocity",
-        # "observation.state.robot.orientation": "state/robot/orientation",
-        # "observation.state.robot.orientation_drift": "state/robot/orientation_drift",
-        # "observation.state.robot.position": "state/robot/position",
-        # "observation.state.robot.position_drift": "state/robot/position_drift",
-        # "observation.state.waist.effort": "state/waist/effort",
-        "observation.state.waist.position": "state/waist/position",
-        # "observation.state.waist.velocity": "state/waist/velocity",
-        # ----- ACTION (index are also commented out) -----
-        # "action.effector.index": "action/effector/index",
-        "action.effector.position": "action/effector/position",
-        # "action.effector.force": "action/effector/force",
-        # "action.end.index": "action/end/index",
-        "action.end.position": "action/end/position",
-        "action.end.orientation": "action/end/orientation",
-        # "action.head.index": "action/head/index",
-        "action.head.position": "action/head/position",
-        # "action.joint.index": "action/joint/index",
-        "action.joint.position": "action/joint/position",
-        # "action.joint.effort": "action/joint/effort",
-        # "action.joint.velocity": "action/joint/velocity",
-        # "action.robot.index": "action/robot/index",
-        # "action.robot.position": "action/robot/position",
-        # "action.robot.orientation": "action/robot/orientation",
-        # "action.robot.angular": "action/robot/angular",
-        "action.robot.velocity": "action/robot/velocity",
-        # "action.waist.index": "action/waist/index",
-        "action.waist.position": "action/waist/position",
-    }
-
-    h5_path = raw_dir / f"proprio_stats/{task_index}/{episode_index}/proprio_stats.h5"
-    with h5py.File(h5_path) as h5:
-        num_frames = len(h5["state/joint/position"])
-
-        for h5_key in keys_mapping.values():
-            col_num_frames = h5[h5_key].shape[0]
-            if col_num_frames != num_frames:
-                raise ValueError(
-                    f"HDF5 column '{h5_key}' is expected to have {num_frames} but has {col_num_frames}' frames instead."
-                )
-
-        for i in range(num_frames):
-            # Create frame
-            f = {new_key: h5[h5_key][i] for new_key, h5_key in keys_mapping.items()}
-
-            for key in f:
-                f[key] = np.array(f[key]).astype(np.float32)
-
-            f["observation.state.end.position"] = f["observation.state.end.position"].reshape(6)
-            f["observation.state.end.orientation"] = f["observation.state.end.orientation"].reshape(8)
-            f["observation.state"] = np.concatenate(
-                [
-                    f["observation.state.head.position"],
-                    f["observation.state.joint.position"][:7],  # left
-                    f["observation.state.effector.position"][[0]],  # left
-                    f["observation.state.joint.position"][7:],  # right
-                    f["observation.state.effector.position"][[1]],  # right
-                    f["observation.state.waist.position"],
-                ]
-            )
-
-            f["action.end.position"] = f["action.end.position"].reshape(6)
-            f["action.end.orientation"] = f["action.end.orientation"].reshape(8)
-            f["action"] = np.concatenate(
-                [
-                    f["action.head.position"],
-                    f["action.joint.position"][:7],  # left
-                    f["action.effector.position"][[0]],  # left
-                    f["action.joint.position"][7:],  # right
-                    f["action.effector.position"][[1]],  # right
-                    f["action.waist.position"],
-                    f["action.robot.velocity"],
-                ]
-            )
-
-            # episode level annotation
-            f["task"] = ep_info["task_name"]
-            f["init_scene_text"] = ep_info["init_scene_text"]
-
-            # frame level annotation
-            if i in frame_idx_to_frames_label_idx:
-                frames_label_idx = frame_idx_to_frames_label_idx[i]
-                frames_label = ep_info["label_info"]["action_config"][frames_label_idx]
-                f["action_text"] = frames_label["action_text"]
-                f["skill"] = frames_label["skill"]
-            else:
-                f["action_text"] = ""
-                f["skill"] = ""
-
-            yield f
-
-
-def update_meta_data(
-    df,
-    ep_to_meta,
-):
-    def _update(row):
-        ep_idx = row["episode_index"]
-        for key, meta in ep_to_meta[ep_idx].items():
-            row[f"videos/{key}/chunk_index"] = meta["chunk_index"]
-            row[f"videos/{key}/file_index"] = meta["file_index"]
-            row[f"videos/{key}/from_timestamp"] = meta["from_timestamp"]
-            row[f"videos/{key}/to_timestamp"] = meta["to_timestamp"]
-        return row
-
-    return df.apply(_update, axis=1)
-
-
-def move_videos_to_lerobot_directory(lerobot_dataset, raw_dir, task_index, episode_names):
-    keys_mapping = {
-        "observation.images.top_head": "head_color",
-        "observation.images.hand_left": "hand_left_color",
-        "observation.images.hand_right": "hand_right_color",
-        "observation.images.head_center_fisheye": "head_center_fisheye_color",
-        "observation.images.head_left_fisheye": "head_left_fisheye_color",
-        "observation.images.head_right_fisheye": "head_right_fisheye_color",
-        "observation.images.back_left_fisheye": "back_left_fisheye_color",
-        "observation.images.back_right_fisheye": "back_right_fisheye_color",
-    }
-
-    # sanity check
-    for key in keys_mapping:
-        if key not in lerobot_dataset.meta.info["features"]:
-            raise ValueError(f"Key '{key}' not found in features.")
-
-    video_keys = keys_mapping.keys()
-    chunk_idx = dict.fromkeys(video_keys, 0)
-    file_idx = dict.fromkeys(video_keys, 0)
-    latest_duration_in_s = dict.fromkeys(video_keys, 0)
-    ep_to_meta = {}
-    for ep_idx, ep_name in enumerate(episode_names):
-        for key in video_keys:
-            raw_videos_dir = raw_dir / f"observations/{task_index}/{ep_name}/videos"
-            old_key = keys_mapping[key]
-            ep_path = raw_videos_dir / f"{old_key}.mp4"
-            ep_duration_in_s = get_video_duration_in_s(ep_path)
-
-            aggr_path = lerobot_dataset.root / DEFAULT_VIDEO_PATH.format(
-                video_key=key,
-                chunk_index=chunk_idx[key],
-                file_index=file_idx[key],
-            )
-            if not aggr_path.exists():
-                # First video
-                aggr_path.parent.mkdir(parents=True, exist_ok=True)
-                shutil.copy(str(ep_path), str(aggr_path))
-            else:
-                size_in_mb = get_video_size_in_mb(ep_path)
-                aggr_size_in_mb = get_video_size_in_mb(aggr_path)
-
-                if aggr_size_in_mb + size_in_mb >= DEFAULT_VIDEO_FILE_SIZE_IN_MB:
-                    # Size limit is reached, prepare new parquet file
-                    chunk_idx[key], file_idx[key] = update_chunk_file_indices(
-                        chunk_idx[key], file_idx[key], DEFAULT_CHUNK_SIZE
-                    )
-                    aggr_path = lerobot_dataset.root / DEFAULT_VIDEO_PATH.format(
-                        video_key=key,
-                        chunk_index=chunk_idx[key],
-                        file_index=file_idx[key],
-                    )
-                    aggr_path.parent.mkdir(parents=True, exist_ok=True)
-                    shutil.copy(str(ep_path), str(aggr_path))
-                    latest_duration_in_s[key] = 0
-                else:
-                    # Update the existing parquet file with new rows
-                    concat_video_files(
-                        [aggr_path, ep_path],
-                        lerobot_dataset.root,
-                        key,
-                        chunk_idx[key],
-                        file_idx[key],
-                    )
-
-            if ep_idx not in ep_to_meta:
-                ep_to_meta[ep_idx] = {}
-            ep_to_meta[ep_idx][key] = {
-                "chunk_index": chunk_idx[key],
-                "file_index": file_idx[key],
-                "from_timestamp": latest_duration_in_s[key],
-                "to_timestamp": latest_duration_in_s[key] + ep_duration_in_s,
-            }
-            latest_duration_in_s[key] += ep_duration_in_s
-
-    # Update episodes meta data
-    for meta_path in (lerobot_dataset.root / EPISODES_DIR).glob("chunk-*/file-*.parquet"):
-        df = pd.read_parquet(meta_path)
-        df = update_meta_data(df, ep_to_meta)
-        df.to_parquet(meta_path)
-
-
-def port_agibot(
-    raw_dir: Path, repo_id: str, task_index: int, episode_indices: list[int], push_to_hub: bool = False
-):
-    lerobot_dataset = LeRobotDataset.create(
-        repo_id=repo_id,
-        robot_type=AGIBOT_ROBOT_TYPE,
-        fps=AGIBOT_FPS,
-        features=AGIBOT_FEATURES,
-    )
-
-    start_time = time.time()
-    num_episodes = len(episode_indices)
-    logging.info(f"Number of episodes {num_episodes}")
-
-    for i, episode_index in enumerate(episode_indices):
-        elapsed_time = time.time() - start_time
-        d, h, m, s = get_elapsed_time_in_days_hours_minutes_seconds(elapsed_time)
-
-        logging.info(
-            f"{i} / {num_episodes} episodes processed (after {d} days, {h} hours, {m} minutes, {s:.3f} seconds)"
-        )
-
-        for frame in generate_lerobot_frames(raw_dir, task_index, episode_index):
-            lerobot_dataset.add_frame(frame)
-
-        lerobot_dataset.save_episode()
-        logging.info("Save_episode")
-
-    # Videos have already been encoded with the proper format, so we rely on hacks
-    # HACK: Add extra images features
-    lerobot_dataset.meta.info["features"].update(AGIBOT_IMAGES_FEATURES)
-    write_info(lerobot_dataset.meta.info, lerobot_dataset.meta.root)
-    move_videos_to_lerobot_directory(lerobot_dataset, raw_dir, task_index, episode_indices)
-
-    if push_to_hub:
-        lerobot_dataset.push_to_hub(
-            # Add agibot tag, since it belongs to the agibot collection of datasets
-            tags=["agibot"],
-            private=False,
-        )
@@ -1,198 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import logging
-import tarfile
-from pathlib import Path
-
-from datatrove.executor import LocalPipelineExecutor
-from datatrove.executor.slurm import SlurmPipelineExecutor
-from datatrove.pipeline.base import PipelineStep
-from port_datasets.agibot_hdf5.download import (
-    RAW_REPO_ID,
-    download_meta_data,
-    get_observations_files,
-)
-
-
-class PortAgiBotShards(PipelineStep):
-    def __init__(
-        self,
-        raw_dir: Path | str,
-        repo_id: str = None,
-    ):
-        super().__init__()
-        self.raw_dir = Path(raw_dir)
-        self.repo_id = repo_id
-
-    def run(self, data=None, rank: int = 0, world_size: int = 1):
-        import shutil
-
-        from datasets.utils.tqdm import disable_progress_bars
-        from port_datasets.agibot_hdf5.download import (
-            RAW_REPO_ID,
-            download,
-            get_observations_files,
-            no_depth,
-        )
-        from port_datasets.agibot_hdf5.port_agibot import port_agibot
-        from port_datasets.droid_rlds.port_droid import validate_dataset
-
-        from lerobot.constants import HF_LEROBOT_HOME
-        from lerobot.utils.utils import init_logging
-
-        init_logging()
-        disable_progress_bars()
-
-        shard_repo_id = f"{self.repo_id}_world_{world_size}_rank_{rank}"
-
-        dataset_dir = HF_LEROBOT_HOME / shard_repo_id
-        if dataset_dir.exists():
-            shutil.rmtree(dataset_dir)
-
-        obs_files, _ = get_observations_files(self.raw_dir, RAW_REPO_ID)
-        obs_file = obs_files[rank]
-
-        # Download subset
-        download(self.raw_dir, allow_patterns=obs_file)
-
-        tar_path = self.raw_dir / obs_file
-        with tarfile.open(tar_path, "r") as tar:
-            extracted_files = tar.getnames()
-
-        task_index = int(tar_path.parent.name)
-        episode_names = [int(p) for p in extracted_files if "/" not in p]
-
-        # Untar if needed
-        if not all((tar_path.parent / f"{ep_name}").exists() for ep_name in episode_names):
-            logging.info(f"Untar-ing {tar_path}...")
-            with tarfile.open(tar_path, "r") as tar:
-                tar.extractall(path=tar_path.parent, filter=no_depth)  # nosec B202
-
-        port_agibot(self.raw_dir, shard_repo_id, task_index, episode_names, push_to_hub=False)
-
-        for ep_name in episode_names:
-            shutil.rmtree(str(tar_path.parent / f"{ep_name}"))
-
-        tar_path.unlink()
-
-        validate_dataset(shard_repo_id)
-
-
-def make_port_executor(
-    raw_dir, repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
-):
-    download_meta_data(raw_dir)
-    obs_files, _ = get_observations_files(raw_dir, RAW_REPO_ID)
-    num_shards = len(obs_files)
-
-    kwargs = {
-        "pipeline": [
-            PortAgiBotShards(raw_dir, repo_id),
-        ],
-        "logging_dir": str(logs_dir / job_name),
-    }
-
-    if slurm:
-        kwargs.update(
-            {
-                "job_name": job_name,
-                "tasks": num_shards,
-                "workers": workers,
-                "time": "08:00:00",
-                "partition": partition,
-                "cpus_per_task": cpus_per_task,
-                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
-            }
-        )
-        executor = SlurmPipelineExecutor(**kwargs)
-    else:
-        kwargs.update(
-            {
-                "tasks": num_shards,
-                "workers": 1,
-            }
-        )
-        executor = LocalPipelineExecutor(**kwargs)
-
-    return executor
-
-
-def main():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--raw-dir",
-        type=Path,
-        required=True,
-        help="Directory containing input raw datasets (e.g. `path/to/dataset` or `path/to/dataset/version).",
-    )
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True.",
-    )
-    parser.add_argument(
-        "--logs-dir",
-        type=Path,
-        help="Path to logs directory for `datatrove`.",
-    )
-    parser.add_argument(
-        "--job-name",
-        type=str,
-        default="port_droid",
-        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
-    )
-    parser.add_argument(
-        "--slurm",
-        type=int,
-        default=1,
-        help="Launch over slurm. Use `--slurm 0` to launch sequentially (useful to debug).",
-    )
-    parser.add_argument(
-        "--workers",
-        type=int,
-        default=2048,
-        help="Number of slurm workers. It should be less than the maximum number of shards.",
-    )
-    parser.add_argument(
-        "--partition",
-        type=str,
-        help="Slurm partition. Ideally a CPU partition. No need for GPU partition.",
-    )
-    parser.add_argument(
-        "--cpus-per-task",
-        type=int,
-        default=8,
-        help="Number of cpus that each slurm worker will use.",
-    )
-    parser.add_argument(
-        "--mem-per-cpu",
-        type=str,
-        default="1950M",
-        help="Memory per cpu that each worker will use.",
-    )
-
-    args = parser.parse_args()
-    kwargs = vars(args)
-    kwargs["slurm"] = kwargs.pop("slurm") == 1
-    port_executor = make_port_executor(**kwargs)
-    port_executor.run()
-
-
-if __name__ == "__main__":
-    main()
@@ -0,0 +1,47 @@
+#!/bin/bash
+
+# Example script for converting RT-1 dataset using SLURM
+# Make sure to modify the paths and parameters according to your setup
+
+# Configuration
+RAW_DIR="/path/to/datasets/fractal20220817_data/0.1.0"
+REPO_ID="your_username/rt1_lerobot"
+LOGS_DIR="/path/to/logs"
+PARTITION="cpu"  # Your SLURM partition name
+
+# Step 1: Convert dataset using distributed processing
+echo "Starting RT-1 dataset conversion..."
+python examples/port_datasets/slurm_port_shards.py \
+    --raw-dir "$RAW_DIR" \
+    --repo-id "$REPO_ID" \
+    --dataset-type rlds \
+    --logs-dir "$LOGS_DIR" \
+    --job-name rt1_conversion \
+    --workers 32 \
+    --num-shards 32 \
+    --partition "$PARTITION" \
+    --cpus-per-task 4 \
+    --mem-per-cpu 2G \
+    --slurm 1
+
+# Step 2: Wait for jobs to complete (you can monitor with squeue)
+echo "Conversion jobs submitted. Monitor with 'squeue -u \$USER'"
+echo "Once all jobs complete, run the aggregation step:"
+echo ""
+echo "python examples/port_datasets/slurm_aggregate_shards.py \\"
+echo "    --repo-id $REPO_ID \\"
+echo "    --push-to-hub"
+
+# Uncomment the following lines if you want to automatically aggregate
+# (but make sure all shards are complete first)
+
+# echo "Waiting for jobs to complete..."
+# while [ $(squeue -u $USER -h | wc -l) -gt 0 ]; do
+#     echo "Jobs still running, waiting 60 seconds..."
+#     sleep 60
+# done
+
+# echo "All jobs completed. Starting aggregation..."
+# python examples/port_datasets/slurm_aggregate_shards.py \
+#     --repo-id "$REPO_ID" \
+#     --push-to-hub
@@ -0,0 +1,15 @@
+# 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.
+
+"""Open X-Embodiment utilities for dataset conversion."""
@@ -0,0 +1,854 @@
+"""
+Adapt from https://github.com/openvla/openvla/blob/main/prismatic/vla/datasets/rlds/oxe/configs.py
+configs.py
+
+Defines per-dataset configuration (kwargs) for each dataset in Open-X Embodiment.
+
+Configuration adopts the following structure:
+    image_obs_keys:
+        primary: primary external RGB
+        secondary: secondary external RGB
+        wrist: wrist RGB
+
+    depth_obs_keys:
+        primary: primary external depth
+        secondary: secondary external depth
+        wrist: wrist depth
+
+    # Always 8-dim =>> changes based on `StateEncoding`
+    state_obs_keys:
+        StateEncoding.POS_EULER:    EEF XYZ (3) + Roll-Pitch-Yaw (3) + <PAD> (1) + Gripper Open/Close (1)
+        StateEncoding.POS_QUAT:     EEF XYZ (3) + Quaternion (4) + Gripper Open/Close (1)
+        StateEncoding.JOINT:        Joint Angles (7, <PAD> if fewer) + Gripper Open/Close (1)
+
+    state_encoding: Type of `StateEncoding`
+    action_encoding: Type of action encoding (e.g., EEF Position vs. Joint Position)
+"""
+
+from enum import IntEnum
+
+import tensorflow as tf
+
+
+def zero_action_filter(traj: dict) -> bool:
+    """
+    Filters transitions whose actions are all-0 (only relative actions, no gripper action).
+    Note: this filter is applied *after* action normalization, so need to compare to "normalized 0".
+    """
+    DROID_Q01 = tf.convert_to_tensor(  # NOQA: N806
+        [
+            -0.7776297926902771,
+            -0.5803514122962952,
+            -0.5795090794563293,
+            -0.6464047729969025,
+            -0.7041108310222626,
+            -0.8895104378461838,
+        ]
+    )
+    DROID_Q99 = tf.convert_to_tensor(  # NOQA: N806
+        [
+            0.7597932070493698,
+            0.5726242214441299,
+            0.7351000607013702,
+            0.6705610305070877,
+            0.6464948207139969,
+            0.8897542208433151,
+        ]
+    )
+    DROID_NORM_0_ACT = (  # NOQA: N806
+        2 * (tf.zeros_like(traj["action"][:, :6]) - DROID_Q01) / (DROID_Q99 - DROID_Q01 + 1e-8) - 1
+    )
+
+    return tf.reduce_any(tf.math.abs(traj["action"][:, :6] - DROID_NORM_0_ACT) > 1e-5)
+
+
+# Defines Proprioceptive State Encoding Schemes
+class StateEncoding(IntEnum):
+    # fmt: off
+    NONE = -1               # No Proprioceptive State
+    POS_EULER = 1           # EEF XYZ (3) + Roll-Pitch-Yaw (3) + <PAD> (1) + Gripper Open/Close (1)
+    POS_QUAT = 2            # EEF XYZ (3) + Quaternion (4) + Gripper Open/Close (1)
+    JOINT = 3               # Joint Angles (7, <PAD> if fewer) + Gripper Open/Close (1)
+    JOINT_BIMANUAL = 4      # Joint Angles (2 x [ Joint Angles (6) + Gripper Open/Close (1) ])
+    # fmt: on
+
+
+# Defines Action Encoding Schemes
+class ActionEncoding(IntEnum):
+    # fmt: off
+    EEF_POS = 1             # EEF Delta XYZ (3) + Roll-Pitch-Yaw (3) + Gripper Open/Close (1)
+    EEF_POS_QUAT = 5        # EEF Delta XYZ (3) + Quaternion (4) + Gripper Open/Close (1)
+    JOINT_POS = 2           # Joint Delta Position (7) + Gripper Open/Close (1)
+    JOINT_POS_BIMANUAL = 3  # Joint Delta Position (2 x [ Joint Delta Position (6) + Gripper Open/Close (1) ])
+    EEF_R6 = 4              # EEF Delta XYZ (3) + R6 (6) + Gripper Open/Close (1)
+    # fmt: on
+
+
+# === Individual Dataset Configs ===
+OXE_DATASET_CONFIGS = {
+    "fractal20220817_data": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["base_pose_tool_reached", "gripper_closed"],
+        "state_encoding": StateEncoding.POS_QUAT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 3,
+        "robot_type": "Google Robot",
+    },
+    "kuka": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": [
+            "clip_function_input/base_pose_tool_reached",
+            "gripper_closed",
+        ],
+        "state_encoding": StateEncoding.POS_QUAT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Kuka iiwa",
+    },
+    "bridge_oxe": {  # Version of Bridge V2 in Open X-Embodiment mixture
+        "image_obs_keys": {"primary": "image", "secondary": "image_1", "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "WidowX",
+    },
+    "bridge_orig": {  # Original version of Bridge V2 from project website
+        "image_obs_keys": {"primary": "image_0", "secondary": "image_1", "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "WidowX",
+    },
+    "bridge_dataset": {  # Original version of Bridge V2 from project website
+        "image_obs_keys": {"primary": "image_0", "secondary": "image_1", "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "WidowX",
+    },
+    "taco_play": {
+        "image_obs_keys": {
+            "primary": "rgb_static",
+            "secondary": None,
+            "wrist": "rgb_gripper",
+        },
+        "depth_obs_keys": {
+            "primary": "depth_static",
+            "secondary": None,
+            "wrist": "depth_gripper",
+        },
+        "state_obs_keys": ["state_eef", None, "state_gripper"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 15,
+        "robot_type": "Franka",
+    },
+    "jaco_play": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "image_wrist",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state_eef", None, "state_gripper"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Jaco 2",
+    },
+    "berkeley_cable_routing": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": "top_image",
+            "wrist": "wrist45_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["robot_state", None],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Franka",
+    },
+    "roboturk": {
+        "image_obs_keys": {"primary": "front_rgb", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": [None, None, None, None, None, None, None, None],
+        "state_encoding": StateEncoding.NONE,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Sawyer",
+    },
+    "nyu_door_opening_surprising_effectiveness": {
+        "image_obs_keys": {"primary": None, "secondary": None, "wrist": "image"},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": [None, None, None, None, None, None, None, None],
+        "state_encoding": StateEncoding.NONE,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 3,
+        "robot_type": "Hello Stretch",
+    },
+    "viola": {
+        "image_obs_keys": {
+            "primary": "agentview_rgb",
+            "secondary": None,
+            "wrist": "eye_in_hand_rgb",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["joint_states", "gripper_states"],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+    "berkeley_autolab_ur5": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "hand_image",
+        },
+        "depth_obs_keys": {"primary": "depth", "secondary": None, "wrist": None},
+        "state_obs_keys": ["state"],
+        "state_encoding": StateEncoding.POS_QUAT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "UR5",
+    },
+    "toto": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state", None],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 30,
+        "robot_type": "Franka",
+    },
+    "language_table": {
+        "image_obs_keys": {"primary": "rgb", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["effector_translation", None, None, None, None, None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "xArm",
+    },
+    "columbia_cairlab_pusht_real": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["robot_state", None, None, None, None, None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "UR5",
+    },
+    "stanford_kuka_multimodal_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": "depth_image", "secondary": None, "wrist": None},
+        "state_obs_keys": ["ee_position", "ee_orientation", None],
+        "state_encoding": StateEncoding.POS_QUAT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Kuka iiwa",
+    },
+    "nyu_rot_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 3,
+        "robot_type": "xArm",
+    },
+    "stanford_hydra_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Franka",
+    },
+    "austin_buds_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state"],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+    "nyu_franka_play_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": "image_additional_view",
+            "wrist": None,
+        },
+        "depth_obs_keys": {
+            "primary": "depth",
+            "secondary": "depth_additional_view",
+            "wrist": None,
+        },
+        "state_obs_keys": ["eef_state", None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 3,
+        "robot_type": "Franka",
+    },
+    "maniskill_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {
+            "primary": "depth",
+            "secondary": None,
+            "wrist": "wrist_depth",
+        },
+        "state_obs_keys": ["tcp_pose", "gripper_state"],
+        "state_encoding": StateEncoding.POS_QUAT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+    "furniture_bench_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state"],
+        "state_encoding": StateEncoding.POS_QUAT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Franka",
+    },
+    "cmu_franka_exploration_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {
+            "primary": "highres_image",
+            "secondary": None,
+            "wrist": None,
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": [None, None, None, None, None, None, None, None],
+        "state_encoding": StateEncoding.NONE,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Franka",
+    },
+    "ucsd_kitchen_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["joint_state", None],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 2,
+        "robot_type": "xArm",
+    },
+    "ucsd_pick_and_place_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 3,
+        "robot_type": "xArm",
+    },
+    "austin_sailor_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state"],
+        "state_encoding": StateEncoding.POS_QUAT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+    "austin_sirius_dataset_converted_externally_to_rlds": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state"],
+        "state_encoding": StateEncoding.POS_QUAT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+    "bc_z": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": [
+            "present/xyz",
+            "present/axis_angle",
+            None,
+            "present/sensed_close",
+        ],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Google Robot",
+    },
+    "utokyo_pr2_opening_fridge_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "PR2",
+    },
+    "utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "PR2",
+    },
+    "utokyo_xarm_pick_and_place_converted_externally_to_rlds": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": "image2",
+            "wrist": "hand_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["end_effector_pose", None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "xArm",
+    },
+    "utokyo_xarm_bimanual_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["pose_r", None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "xArm Bimanual",
+    },
+    "robo_net": {
+        "image_obs_keys": {"primary": "image", "secondary": "image1", "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 1,
+        "robot_type": "Multi-Robot",
+    },
+    "berkeley_mvp_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": None, "secondary": None, "wrist": "hand_image"},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["pose", "gripper"],
+        "state_encoding": StateEncoding.POS_QUAT,
+        "action_encoding": ActionEncoding.JOINT_POS,
+        "control_frequency": 5,
+        "robot_type": "xArm",
+    },
+    "berkeley_rpt_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": None, "secondary": None, "wrist": "hand_image"},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["joint_pos", "gripper"],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.JOINT_POS,
+        "control_frequency": 30,
+        "robot_type": "Franka",
+    },
+    "kaist_nonprehensile_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state", None],
+        "state_encoding": StateEncoding.POS_QUAT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Franka",
+    },
+    "stanford_mask_vit_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": None,
+        "robot_type": "Sawyer",
+    },
+    "tokyo_u_lsmo_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Cobotta",
+    },
+    "dlr_sara_pour_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state", None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "DLR SARA",
+    },
+    "dlr_sara_grid_clamp_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state", None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "DLR SARA",
+    },
+    "dlr_edan_shared_control_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state", None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "DLR EDAN",
+    },
+    "asu_table_top_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 12.5,
+        "robot_type": "UR5",
+    },
+    "stanford_robocook_converted_externally_to_rlds": {
+        "image_obs_keys": {"primary": "image_1", "secondary": "image_2", "wrist": None},
+        "depth_obs_keys": {"primary": "depth_1", "secondary": "depth_2", "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "Franka",
+    },
+    "imperialcollege_sawyer_wrist_cam": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": [None, None, None, None, None, None, None, "state"],
+        "state_encoding": StateEncoding.NONE,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Sawyer",
+    },
+    "iamlab_cmu_pickup_insert_converted_externally_to_rlds": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["joint_state", "gripper_state"],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+    "uiuc_d3field": {
+        "image_obs_keys": {"primary": "image_1", "secondary": "image_2", "wrist": None},
+        "depth_obs_keys": {"primary": "depth_1", "secondary": "depth_2", "wrist": None},
+        "state_obs_keys": [None, None, None, None, None, None, None, None],
+        "state_encoding": StateEncoding.NONE,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 1,
+        "robot_type": "Kinova Gen3",
+    },
+    "utaustin_mutex": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state"],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+    "berkeley_fanuc_manipulation": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "wrist_image",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["joint_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Fanuc Mate",
+    },
+    "cmu_playing_with_food": {
+        "image_obs_keys": {
+            "primary": "image",
+            "secondary": None,
+            "wrist": "finger_vision_1",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state", None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Franka",
+    },
+    "cmu_play_fusion": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state"],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "Franka",
+    },
+    "cmu_stretch": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["eef_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Hello Stretch",
+    },
+    "berkeley_gnm_recon": {
+        "image_obs_keys": {"primary": None, "secondary": None, "wrist": "image"},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state", None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 3,
+        "robot_type": "Jackal",
+    },
+    "berkeley_gnm_cory_hall": {
+        "image_obs_keys": {"primary": None, "secondary": None, "wrist": "image"},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state", None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "RC Car",
+    },
+    "berkeley_gnm_sac_son": {
+        "image_obs_keys": {"primary": None, "secondary": None, "wrist": "image"},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["state", None, None],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "TurtleBot 2",
+    },
+    # NOTE: modified
+    "droid": {
+        "image_obs_keys": {
+            "primary": "exterior_image_1_left",
+            "secondary": "exterior_image_2_left",
+            "wrist": "wrist_image_left",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 15,
+        "robot_type": "Franka",
+        "aux_kwargs": {
+            "dataset_frame_transform_kwargs": {
+                "chunk_filter_fn": zero_action_filter,
+            },
+        },
+    },
+    "fmb_dataset": {
+        "image_obs_keys": {
+            "primary": "image_side_1",
+            "secondary": "image_side_2",
+            "wrist": "image_wrist_1",
+        },
+        "depth_obs_keys": {
+            "primary": "image_side_1_depth",
+            "secondary": "image_side_2_depth",
+            "wrist": "image_wrist_1_depth",
+        },
+        "state_obs_keys": ["proprio"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Franka",
+    },
+    # NOTE: modified
+    "dobbe": {
+        "image_obs_keys": {"primary": "wrist_image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 3.75,
+        "robot_type": "Hello Stretch",
+    },
+    "roboset": {
+        "image_obs_keys": {
+            "primary": "image_left",
+            "secondary": "image_right",
+            "wrist": "image_wrist",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["proprio"],
+        "state_encoding": StateEncoding.JOINT,
+        "action_encoding": ActionEncoding.JOINT_POS,
+        "control_frequency": 5,
+        "robot_type": "Franka",
+    },
+    "rh20t": {
+        "image_obs_keys": {
+            "primary": "image_front",
+            "secondary": "image_side_right",
+            "wrist": "image_wrist",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["proprio"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 10,
+        "robot_type": "Flexiv",
+    },
+    ### T-DROID datasets
+    "tdroid_carrot_in_bowl": {  # "put carrot in bowl" task, 50 demos @ 5 Hz control
+        "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "Franka",
+    },
+    "tdroid_pour_corn_in_pot": {  # "pour corn from red bonawl into steel pot" task, 50 demos @ 5 Hz control
+        "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "Franka",
+    },
+    "tdroid_flip_pot_upright": {  # "flip pot upright" task, 10 demos @ 5 Hz control
+        "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "Franka",
+    },
+    "tdroid_move_object_onto_plate": {  # "move <object> onto plate" task, 150 demos @ 5 Hz control
+        "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "Franka",
+    },
+    "tdroid_knock_object_over": {  # "knock <object> over" task, 70 demos @ 5 Hz control
+        "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "Franka",
+    },
+    "tdroid_cover_object_with_towel": {  # "cover <object> with towel" task, 45 demos @ 5 Hz control
+        "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
+        "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", None, "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 5,
+        "robot_type": "Franka",
+    },
+    ### DROID Finetuning datasets
+    "droid_wipe": {
+        "image_obs_keys": {
+            "primary": "exterior_image_2_left",
+            "secondary": None,
+            "wrist": "wrist_image_left",
+        },
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["proprio"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 15,
+        "robot_type": "Franka",
+    },
+    # NOTE: modified
+    ### LIBERO datasets (modified versions)
+    "libero_spatial_no_noops": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+    "libero_object_no_noops": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+    "libero_goal_no_noops": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+    "libero_10_no_noops": {
+        "image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"},
+        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
+        "state_obs_keys": ["EEF_state", "gripper_state"],
+        "state_encoding": StateEncoding.POS_EULER,
+        "action_encoding": ActionEncoding.EEF_POS,
+        "control_frequency": 20,
+        "robot_type": "Franka",
+    },
+}
@@ -0,0 +1,76 @@
+"""
+Copied from https://github.com/openvla/openvla/blob/main/prismatic/vla/datasets/rlds/utils/data_utils.py
+"""
+
+from typing import Any
+
+import tensorflow as tf
+
+
+def binarize_gripper_actions(actions: tf.Tensor) -> tf.Tensor:
+    """
+    Converts gripper actions from continuous to binary values (0 and 1).
+
+    We exploit that fact that most of the time, the gripper is fully open (near 1.0) or fully closed (near 0.0). As it
+    transitions between the two, it sometimes passes through a few intermediate values. We relabel those intermediate
+    values based on the state that is reached _after_ those intermediate values.
+
+    In the edge case that the trajectory ends with an intermediate value, we give up on binarizing and relabel that
+    chunk of intermediate values as the last action in the trajectory.
+
+    The `scan_fn` implements the following logic:
+        new_actions = np.empty_like(actions)
+        carry = actions[-1]
+        for i in reversed(range(actions.shape[0])):
+            if in_between_mask[i]:
+                carry = carry
+            else:
+                carry = float(open_mask[i])
+            new_actions[i] = carry
+    """
+    open_mask, closed_mask = actions > 0.95, actions < 0.05
+    in_between_mask = tf.logical_not(tf.logical_or(open_mask, closed_mask))
+    is_open_float = tf.cast(open_mask, tf.float32)
+
+    def scan_fn(carry, i):
+        return tf.cond(in_between_mask[i], lambda: tf.cast(carry, tf.float32), lambda: is_open_float[i])
+
+    return tf.scan(scan_fn, tf.range(tf.shape(actions)[0]), actions[-1], reverse=True)
+
+
+def invert_gripper_actions(actions: tf.Tensor) -> tf.Tensor:
+    return 1 - actions
+
+
+def rel2abs_gripper_actions(actions: tf.Tensor) -> tf.Tensor:
+    """
+    Converts relative gripper actions (+1 for closing, -1 for opening) to absolute actions (0 = closed; 1 = open).
+
+    Assumes that the first relative gripper is not redundant (i.e. close when already closed)!
+    """
+    # Note =>> -1 for closing, 1 for opening, 0 for no change
+    opening_mask, closing_mask = actions < -0.1, actions > 0.1
+    thresholded_actions = tf.where(opening_mask, 1, tf.where(closing_mask, -1, 0))
+
+    def scan_fn(carry, i):
+        return tf.cond(thresholded_actions[i] == 0, lambda: carry, lambda: thresholded_actions[i])
+
+    # If no relative grasp, assumes open for whole trajectory
+    start = -1 * thresholded_actions[tf.argmax(thresholded_actions != 0, axis=0)]
+    start = tf.cond(start == 0, lambda: 1, lambda: start)
+
+    # Note =>> -1 for closed, 1 for open
+    new_actions = tf.scan(scan_fn, tf.range(tf.shape(actions)[0]), start)
+    new_actions = tf.cast(new_actions, tf.float32) / 2 + 0.5
+
+    return new_actions
+
+
+# === Bridge-V2 =>> Dataset-Specific Transform ===
+def relabel_bridge_actions(traj: dict[str, Any]) -> dict[str, Any]:
+    """Relabels actions to use reached proprioceptive state; discards last timestep (no-action)."""
+    movement_actions = traj["observation"]["state"][1:, :6] - traj["observation"]["state"][:-1, :6]
+    traj_truncated = tf.nest.map_structure(lambda x: x[:-1], traj)
+    traj_truncated["action"] = tf.concat([movement_actions, traj["action"][:-1, -1:]], axis=1)
+
+    return traj_truncated
@@ -0,0 +1,359 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import logging
+import re
+import time
+from functools import partial
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_datasets as tfds
+from oxe_utils.configs import OXE_DATASET_CONFIGS, ActionEncoding, StateEncoding
+from oxe_utils.transforms import OXE_STANDARDIZATION_TRANSFORMS
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.utils.utils import get_elapsed_time_in_days_hours_minutes_seconds
+
+# Default FPS for datasets without specific config
+DEFAULT_FPS = 10
+DEFAULT_ROBOT_TYPE = "unknown"
+
+
+def determine_dataset_info(raw_dir: Path):
+    """Determine dataset name and version from directory structure."""
+    last_part = raw_dir.name
+    if re.match(r"^\d+\.\d+\.\d+$", last_part):
+        version = last_part
+        dataset_name = raw_dir.parent.name
+        data_dir = raw_dir.parent.parent
+    else:
+        version = ""
+        dataset_name = last_part
+        data_dir = raw_dir.parent
+    return dataset_name, version, data_dir
+
+
+def generate_features_from_builder(builder, dataset_name: str) -> dict[str, Any]:
+    """Generate LeRobot features schema from TFDS builder and dataset config."""
+
+    # Generate state names based on encoding type
+    state_names = [f"motor_{i}" for i in range(8)]
+    if dataset_name in OXE_DATASET_CONFIGS:
+        state_encoding = OXE_DATASET_CONFIGS[dataset_name]["state_encoding"]
+        if state_encoding == StateEncoding.POS_EULER:
+            state_names = ["x", "y", "z", "roll", "pitch", "yaw", "pad", "gripper"]
+            if "libero" in dataset_name:
+                state_names = [
+                    "x",
+                    "y",
+                    "z",
+                    "roll",
+                    "pitch",
+                    "yaw",
+                    "gripper",
+                    "gripper",
+                ]  # 2D gripper state
+        elif state_encoding == StateEncoding.POS_QUAT:
+            state_names = ["x", "y", "z", "rx", "ry", "rz", "rw", "gripper"]
+        elif state_encoding == StateEncoding.JOINT:
+            state_names = [f"motor_{i}" for i in range(7)] + ["gripper"]
+            state_obs_keys = OXE_DATASET_CONFIGS[dataset_name]["state_obs_keys"]
+            pad_count = state_obs_keys[:-1].count(None)
+            state_names[-pad_count - 1 : -1] = ["pad"] * pad_count
+            state_names[-1] = "pad" if state_obs_keys[-1] is None else state_names[-1]
+
+    # Generate action names based on encoding type
+    action_names = [f"motor_{i}" for i in range(8)]
+    if dataset_name in OXE_DATASET_CONFIGS:
+        action_encoding = OXE_DATASET_CONFIGS[dataset_name]["action_encoding"]
+        if action_encoding == ActionEncoding.EEF_POS:
+            action_names = ["x", "y", "z", "roll", "pitch", "yaw", "gripper"]
+        elif action_encoding == ActionEncoding.JOINT_POS:
+            action_names = [f"motor_{i}" for i in range(7)] + ["gripper"]
+
+    # Base features (state and action)
+    features = {
+        "observation.state": {
+            "dtype": "float32",
+            "shape": (len(state_names),),
+            "names": {"axes": state_names},
+        },
+        "action": {
+            "dtype": "float32",
+            "shape": (len(action_names),),
+            "names": {"axes": action_names},
+        },
+    }
+
+    # Add image features from TFDS builder info
+    obs_features = builder.info.features["steps"]["observation"]
+    for key, value in obs_features.items():
+        # Skip depth images and non-image features
+        if "depth" in key or not any(x in key for x in ["image", "rgb"]):
+            continue
+
+        features[f"observation.images.{key}"] = {
+            "dtype": "video",
+            "shape": tuple(value.shape),
+            "names": ["height", "width", "channels"],
+        }
+
+    return features
+
+
+def transform_raw_dataset(episode, dataset_name: str):
+    """Apply OXE standardization transforms to raw TFDS episode."""
+    # Batch all steps in the episode
+    traj = next(iter(episode["steps"].batch(episode["steps"].cardinality())))
+
+    # Apply dataset-specific transform if available
+    if dataset_name in OXE_STANDARDIZATION_TRANSFORMS:
+        traj = OXE_STANDARDIZATION_TRANSFORMS[dataset_name](traj)
+
+    # Create consolidated state vector
+    if dataset_name in OXE_DATASET_CONFIGS:
+        state_obs_keys = OXE_DATASET_CONFIGS[dataset_name]["state_obs_keys"]
+    else:
+        state_obs_keys = [None for _ in range(8)]
+
+    # Build proprio (proprioceptive state) vector
+    proprio_components = []
+    for key in state_obs_keys:
+        if key is None:
+            # Add padding for missing state components
+            component = tf.zeros((tf.shape(traj["action"])[0], 1), dtype=tf.float32)
+        else:
+            component = tf.cast(traj["observation"][key], tf.float32)
+            # Ensure component has right shape (add dimension if needed)
+            if len(component.shape) == 1:
+                component = component[:, None]
+        proprio_components.append(component)
+
+    proprio = tf.concat(proprio_components, axis=1)
+
+    # Update trajectory with standardized format
+    traj.update(
+        {
+            "proprio": proprio,
+            "task": traj.get("language_instruction", ""),
+            "action": tf.cast(traj["action"], tf.float32),
+        }
+    )
+
+    episode["steps"] = traj
+    return episode
+
+
+def generate_lerobot_frames(tf_episode):
+    """Generate LeRobot frames from transformed TFDS episode."""
+    traj = tf_episode["steps"]
+
+    # Get the task/language instruction
+    if isinstance(traj["task"], tf.Tensor):
+        if traj["task"].dtype == tf.string:
+            task = traj["task"][0].numpy().decode() if len(traj["task"]) > 0 else ""
+        else:
+            task = str(traj["task"][0].numpy()) if len(traj["task"]) > 0 else ""
+    else:
+        task = str(traj["task"]) if traj["task"] else ""
+
+    # Iterate through each timestep
+    num_steps = tf.shape(traj["action"])[0].numpy()
+    for i in range(num_steps):
+        frame = {}
+
+        # Add observation state
+        frame["observation.state"] = traj["proprio"][i].numpy()
+
+        # Add action
+        frame["action"] = traj["action"][i].numpy()
+
+        # Add images
+        for key, value in traj["observation"].items():
+            if any(x in key for x in ["image", "rgb"]) and "depth" not in key:
+                frame[f"observation.images.{key}"] = value[i].numpy()
+
+        # Add task
+        frame["task"] = task
+
+        # Cast fp64 to fp32
+        for key in frame:
+            if isinstance(frame[key], np.ndarray) and frame[key].dtype == np.float64:
+                frame[key] = frame[key].astype(np.float32)
+
+        yield frame
+
+
+def port_rlds(
+    raw_dir: Path,
+    repo_id: str,
+    push_to_hub: bool = False,
+    num_shards: int | None = None,
+    shard_index: int | None = None,
+):
+    """Port RLDS dataset to LeRobot format."""
+
+    # Determine dataset info
+    dataset_name, version, data_dir = determine_dataset_info(raw_dir)
+
+    # Build TFDS dataset
+    builder = tfds.builder(
+        f"{dataset_name}/{version}" if version else dataset_name, data_dir=data_dir, version=version
+    )
+
+    # Handle sharding if specified
+    if num_shards is not None and shard_index is not None:
+        if shard_index >= num_shards:
+            raise ValueError(f"Shard index {shard_index} >= num_shards {num_shards}")
+
+        # Calculate shard splits
+        total_episodes = builder.info.splits["train"].num_examples
+        episodes_per_shard = total_episodes // num_shards
+        start_idx = shard_index * episodes_per_shard
+        if shard_index == num_shards - 1:
+            # Last shard gets remaining episodes
+            end_idx = total_episodes
+        else:
+            end_idx = start_idx + episodes_per_shard
+
+        split_str = f"train[{start_idx}:{end_idx}]"
+        raw_dataset = builder.as_dataset(split=split_str)
+    else:
+        raw_dataset = builder.as_dataset(split="train")
+
+    # Apply filtering (e.g., success filter for kuka)
+    if dataset_name == "kuka":
+        raw_dataset = raw_dataset.filter(lambda e: e["success"])
+
+    # Apply transformations
+    raw_dataset = raw_dataset.map(partial(transform_raw_dataset, dataset_name=dataset_name))
+
+    # Get dataset configuration
+    fps = DEFAULT_FPS
+    robot_type = DEFAULT_ROBOT_TYPE
+
+    if dataset_name in OXE_DATASET_CONFIGS:
+        config = OXE_DATASET_CONFIGS[dataset_name]
+        fps = config.get("control_frequency", DEFAULT_FPS)
+        robot_type = config.get("robot_type", DEFAULT_ROBOT_TYPE)
+        robot_type = robot_type.lower().replace(" ", "_").replace("-", "_")
+
+    # Generate features schema
+    features = generate_features_from_builder(builder, dataset_name)
+
+    # Create LeRobot dataset
+    lerobot_dataset = LeRobotDataset.create(
+        repo_id=repo_id,
+        robot_type=robot_type,
+        fps=int(fps),
+        features=features,
+    )
+
+    # Process episodes
+    start_time = time.time()
+    num_episodes = raw_dataset.cardinality().numpy().item()
+    logging.info(f"Number of episodes: {num_episodes}")
+
+    for episode_index, episode in enumerate(raw_dataset):
+        elapsed_time = time.time() - start_time
+        d, h, m, s = get_elapsed_time_in_days_hours_minutes_seconds(elapsed_time)
+
+        logging.info(
+            f"{episode_index} / {num_episodes} episodes processed "
+            f"(after {d} days, {h} hours, {m} minutes, {s:.3f} seconds)"
+        )
+
+        # Generate and add frames
+        for frame in generate_lerobot_frames(episode):
+            lerobot_dataset.add_frame(frame)
+
+        lerobot_dataset.save_episode()
+        logging.info("Save_episode")
+
+    # Push to hub if requested
+    if push_to_hub:
+        tags = ["openx", dataset_name]
+        if robot_type != "unknown":
+            tags.append(robot_type)
+
+        lerobot_dataset.push_to_hub(
+            tags=tags,
+            private=False,
+        )
+
+
+def validate_dataset(repo_id):
+    """Sanity check that ensures metadata can be loaded and all files are present."""
+    meta = LeRobotDatasetMetadata(repo_id)
+
+    if meta.total_episodes == 0:
+        raise ValueError("Number of episodes is 0.")
+
+    for ep_idx in range(meta.total_episodes):
+        data_path = meta.root / meta.get_data_file_path(ep_idx)
+
+        if not data_path.exists():
+            raise ValueError(f"Parquet file is missing in: {data_path}")
+
+        for vid_key in meta.video_keys:
+            vid_path = meta.root / meta.get_video_file_path(ep_idx, vid_key)
+            if not vid_path.exists():
+                raise ValueError(f"Video file is missing in: {vid_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--raw-dir",
+        type=Path,
+        required=True,
+        help="Directory containing input raw datasets (e.g. `path/to/dataset` or `path/to/dataset/version).",
+    )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        help="Repository identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True",
+    )
+    parser.add_argument(
+        "--push-to-hub",
+        action="store_true",
+        help="Upload to hub.",
+    )
+    parser.add_argument(
+        "--num-shards",
+        type=int,
+        default=None,
+        help="Number of shards to split the dataset into for parallel processing.",
+    )
+    parser.add_argument(
+        "--shard-index",
+        type=int,
+        default=None,
+        help="Index of the shard to process (0-indexed).",
+    )
+
+    args = parser.parse_args()
+
+    port_rlds(**vars(args))
+
+
+if __name__ == "__main__":
+    main()
@@ -20,7 +20,7 @@ from pathlib import Path
 from datatrove.executor import LocalPipelineExecutor
 from datatrove.executor.slurm import SlurmPipelineExecutor
 from datatrove.pipeline.base import PipelineStep
-from port_datasets.droid_rlds.port_droid import DROID_SHARDS
+from port_droid import DROID_SHARDS


 class PortDroidShards(PipelineStep):
@@ -35,7 +35,7 @@ class PortDroidShards(PipelineStep):

    def run(self, data=None, rank: int = 0, world_size: int = 1):
        from datasets.utils.tqdm import disable_progress_bars
-        from port_datasets.droid_rlds.port_droid import port_droid, validate_dataset
+        from port_droid import port_droid, validate_dataset

        from lerobot.utils.utils import init_logging

@@ -61,13 +61,71 @@ class PortDroidShards(PipelineStep):
        validate_dataset(shard_repo_id)


+class PortRLDSShards(PipelineStep):
+    def __init__(
+        self,
+        raw_dir: Path | str,
+        repo_id: str = None,
+        num_shards: int = None,
+    ):
+        super().__init__()
+        self.raw_dir = Path(raw_dir)
+        self.repo_id = repo_id
+        self.num_shards = num_shards
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        from datasets.utils.tqdm import disable_progress_bars
+        from port_rlds import port_rlds, validate_dataset
+
+        from lerobot.utils.utils import init_logging
+
+        init_logging()
+        disable_progress_bars()
+
+        shard_repo_id = f"{self.repo_id}_world_{world_size}_rank_{rank}"
+
+        try:
+            validate_dataset(shard_repo_id)
+            return
+        except Exception:
+            pass  # nosec B110 - Dataset doesn't exist yet, continue with porting
+
+        port_rlds(
+            self.raw_dir,
+            shard_repo_id,
+            push_to_hub=False,
+            num_shards=world_size,
+            shard_index=rank,
+        )
+
+        validate_dataset(shard_repo_id)
+
+
 def make_port_executor(
-    raw_dir, repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
+    raw_dir,
+    repo_id,
+    job_name,
+    logs_dir,
+    workers,
+    partition,
+    cpus_per_task,
+    mem_per_cpu,
+    slurm=True,
+    dataset_type="droid",
+    num_shards=None,
 ):
+    # Select appropriate pipeline step based on dataset type
+    if dataset_type.lower() == "droid":
+        pipeline_step = PortDroidShards(raw_dir, repo_id)
+        default_shards = DROID_SHARDS
+    elif dataset_type.lower() == "rlds":
+        pipeline_step = PortRLDSShards(raw_dir, repo_id, num_shards)
+        default_shards = num_shards or workers  # Use num_shards or fallback to workers
+    else:
+        raise ValueError(f"Unsupported dataset type: {dataset_type}")
+
    kwargs = {
-        "pipeline": [
-            PortDroidShards(raw_dir, repo_id),
-        ],
+        "pipeline": [pipeline_step],
        "logging_dir": str(logs_dir / job_name),
    }

@@ -75,7 +133,7 @@ def make_port_executor(
        kwargs.update(
            {
                "job_name": job_name,
-                "tasks": DROID_SHARDS,
+                "tasks": default_shards,
                "workers": workers,
                "time": "08:00:00",
                "partition": partition,
@@ -113,13 +171,21 @@ def main():
    parser.add_argument(
        "--logs-dir",
        type=Path,
-        help="Path to logs directory for `datatrove`.",
+        default=Path("./logs"),
+        help="Path to logs directory for `datatrove` (default: ./logs).",
+    )
+    parser.add_argument(
+        "--dataset-type",
+        type=str,
+        choices=["droid", "rlds"],
+        default="droid",
+        help="Type of dataset to process: 'droid' for DROID datasets or 'rlds' for RLDS/OpenX datasets.",
    )
    parser.add_argument(
        "--job-name",
        type=str,
-        default="port_droid",
-        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
+        default=None,
+        help="Job name used in slurm, and name of the directory created inside the provided logs directory. Defaults to 'port_{dataset_type}'.",
    )
    parser.add_argument(
        "--slurm",
@@ -130,8 +196,14 @@ def main():
    parser.add_argument(
        "--workers",
        type=int,
-        default=2048,
-        help="Number of slurm workers. It should be less than the maximum number of shards.",
+        default=None,
+        help="Number of slurm workers. Defaults: 2048 for DROID, 64 for RLDS datasets.",
+    )
+    parser.add_argument(
+        "--num-shards",
+        type=int,
+        default=None,
+        help="Number of shards to split the dataset into. For DROID datasets, this is fixed at 2048. For RLDS datasets, defaults to number of workers.",
    )
    parser.add_argument(
        "--partition",
@@ -152,8 +224,21 @@ def main():
    )

    args = parser.parse_args()
+
+    # Set defaults based on dataset type
+    if args.job_name is None:
+        args.job_name = f"port_{args.dataset_type}"
+
+    if args.workers is None:
+        if args.dataset_type == "droid":
+            args.workers = 2048
+        else:  # rlds
+            args.workers = 64
+
+    # Convert args to kwargs and process
    kwargs = vars(args)
    kwargs["slurm"] = kwargs.pop("slurm") == 1
+
    port_executor = make_port_executor(**kwargs)
    port_executor.run()

@@ -1,111 +0,0 @@
-#!/usr/bin/env python
-"""
-Example script demonstrating dataset tools utilities.
-
-This script shows how to:
-1. Delete episodes from a dataset
-2. Split a dataset into train/val sets
-3. Add/remove features
-4. Merge datasets
-
-Usage:
-    python examples/use_dataset_tools.py
-"""
-
-import numpy as np
-
-from lerobot.datasets.dataset_tools import (
-    add_feature,
-    delete_episodes,
-    merge_datasets,
-    remove_feature,
-    split_dataset,
-)
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-
-def main():
-    # Load an existing dataset (replace with your dataset)
-    dataset = LeRobotDataset("lerobot/pusht")
-
-    print(f"Original dataset: {dataset.meta.total_episodes} episodes, {dataset.meta.total_frames} frames")
-    print(f"Features: {list(dataset.meta.features.keys())}")
-
-    # Example 1: Delete episodes
-    print("\n1. Deleting episodes 0 and 2...")
-    filtered_dataset = delete_episodes(dataset, episode_indices=[0, 2], repo_id="pusht_filtered")
-    print(f"Filtered dataset: {filtered_dataset.meta.total_episodes} episodes")
-
-    # Example 2: Split dataset
-    print("\n2. Splitting dataset into train/val...")
-    splits = split_dataset(
-        dataset,
-        splits={"train": 0.8, "val": 0.2},
-    )
-    print(f"Train split: {splits['train'].meta.total_episodes} episodes")
-    print(f"Val split: {splits['val'].meta.total_episodes} episodes")
-
-    # Example 3: Add a feature
-    print("\n3. Adding a reward feature...")
-
-    # Method 1: Pre-computed values
-    reward_values = np.random.randn(dataset.meta.total_frames).astype(np.float32)
-    dataset_with_reward = add_feature(
-        dataset,
-        feature_name="reward",
-        feature_values=reward_values,
-        feature_info={
-            "dtype": "float32",
-            "shape": (1,),
-            "names": None,
-        },
-        repo_id="pusht_with_reward",
-    )
-
-    # Method 2: Using a callable
-    def compute_success(frame_dict, episode_idx, frame_idx):
-        # Example: mark last 10 frames of each episode as successful
-        episode_length = 10  # You'd get this from episode metadata
-        return float(frame_idx >= episode_length - 10)
-
-    dataset_with_success = add_feature(
-        dataset_with_reward,
-        feature_name="success",
-        feature_values=compute_success,
-        feature_info={
-            "dtype": "float32",
-            "shape": (1,),
-            "names": None,
-        },
-        repo_id="pusht_with_reward_and_success",
-    )
-
-    print(f"New features: {list(dataset_with_success.meta.features.keys())}")
-
-    # Example 4: Remove features
-    print("\n4. Removing the success feature...")
-    dataset_cleaned = remove_feature(dataset_with_success, feature_names="success", repo_id="pusht_cleaned")
-    print(f"Features after removal: {list(dataset_cleaned.meta.features.keys())}")
-
-    # Example 5: Merge datasets
-    print("\n5. Merging train and val splits back together...")
-    merged = merge_datasets([splits["train"], splits["val"]], output_repo_id="pusht_merged")
-    print(f"Merged dataset: {merged.meta.total_episodes} episodes")
-
-    # Example 6: Complex workflow
-    print("\n6. Complex workflow example...")
-
-    # Remove a camera if dataset has multiple
-    if len(dataset.meta.camera_keys) > 1:
-        camera_to_remove = dataset.meta.camera_keys[0]
-        print(f"Removing camera: {camera_to_remove}")
-        dataset_no_cam = remove_feature(
-            dataset, feature_names=camera_to_remove, repo_id="pusht_no_first_camera"
-        )
-        print(f"Remaining cameras: {dataset_no_cam.meta.camera_keys}")
-
-    print("\nDone! Check ~/.cache/huggingface/lerobot/ for the created datasets.")
-
-
-if __name__ == "__main__":
-    main()
@@ -84,7 +84,6 @@ dependencies = [

    # Support dependencies
    "deepdiff>=7.0.1,<9.0.0",
-    "flask>=3.0.3,<4.0.0",
    "imageio[ffmpeg]>=2.34.0,<3.0.0",
    "termcolor>=2.4.0,<4.0.0",
 ]
@@ -152,11 +152,11 @@ def update_meta_data(
 def aggregate_datasets(
    repo_ids: list[str],
    aggr_repo_id: str,
-    roots: list[Path] = None,
-    aggr_root: Path = None,
-    data_files_size_in_mb: float = None,
-    video_files_size_in_mb: float = None,
-    chunk_size: int = None,
+    roots: list[Path] | None = None,
+    aggr_root: Path | None = None,
+    data_files_size_in_mb: float | None = None,
+    video_files_size_in_mb: float | None = None,
+    chunk_size: int | None = None,
 ):
    """Aggregates multiple LeRobot datasets into a single unified dataset.

@@ -70,10 +70,8 @@ class CompatibilityError(Exception): ...

 class BackwardCompatibilityError(CompatibilityError):
    def __init__(self, repo_id: str, version: packaging.version.Version):
-        if version.major == 3:
+        if version.major == 2 and version.minor == 1:
            message = V30_MESSAGE.format(repo_id=repo_id, version=version)
-        elif version.major == 2:
-            message = V2_MESSAGE.format(repo_id=repo_id, version=version)
        else:
            raise NotImplementedError(
                "Contact the maintainer on [Discord](https://discord.com/invite/s3KuuzsPFb)."
@@ -1,761 +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.
-"""Dataset tools utilities for LeRobotDataset.
-
-This module provides utilities for:
- Deleting episodes from datasets
- Splitting datasets into multiple smaller datasets
- Adding/removing features from datasets
- Merging datasets (wrapper around aggregate functionality)
-"""
-
-import logging
-import shutil
-from collections.abc import Callable
-from pathlib import Path
-
-import numpy as np
-import pandas as pd
-import torch
-from tqdm import tqdm
-
-from lerobot.constants import HF_LEROBOT_HOME
-from lerobot.datasets.aggregate import aggregate_datasets
-from lerobot.datasets.compute_stats import aggregate_stats
-from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
-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,
-    get_parquet_file_size_in_mb,
-    get_video_size_in_mb,
-    to_parquet_with_hf_images,
-    update_chunk_file_indices,
-    write_info,
-    write_stats,
-    write_tasks,
-)
-
-
-def delete_episodes(
-    dataset: LeRobotDataset,
-    episode_indices: list[int],
-    output_dir: str | Path | None = None,
-    repo_id: str | None = None,
-) -> LeRobotDataset:
-    """Delete episodes from a LeRobotDataset and create a new dataset.
-
-    Args:
-        dataset: The source LeRobotDataset.
-        episode_indices: List of episode indices to delete.
-        output_dir: Directory to save the new dataset. If None, uses default location.
-        repo_id: Repository ID for the new dataset. If None, appends "_filtered" to original.
-
-    Returns:
-        LeRobotDataset: New dataset with episodes removed.
-    """
-    if not episode_indices:
-        raise ValueError("No episodes to delete")
-
-    # Validate episode indices
-    valid_indices = set(range(dataset.meta.total_episodes))
-    invalid = set(episode_indices) - valid_indices
-    if invalid:
-        raise ValueError(f"Invalid episode indices: {invalid}")
-
-    logging.info(f"Deleting {len(episode_indices)} episodes from dataset")
-
-    # Create new dataset metadata
-    if repo_id is None:
-        repo_id = f"{dataset.repo_id}_filtered"
-    if output_dir is None:
-        output_dir = HF_LEROBOT_HOME / repo_id
-    else:
-        output_dir = Path(output_dir)
-
-    # Get episodes to keep
-    episodes_to_keep = [i for i in range(dataset.meta.total_episodes) if i not in episode_indices]
-    if not episodes_to_keep:
-        raise ValueError("Cannot delete all episodes from dataset")
-
-    # Create new dataset
-    new_meta = LeRobotDatasetMetadata.create(
-        repo_id=repo_id,
-        fps=dataset.meta.fps,
-        features=dataset.meta.features,
-        robot_type=dataset.meta.robot_type,
-        root=output_dir,
-        use_videos=len(dataset.meta.video_keys) > 0,
-    )
-
-    # Process episodes
-    episode_mapping = {}  # old_idx -> new_idx
-    new_episode_idx = 0
-
-    for old_idx in tqdm(episodes_to_keep, desc="Processing episodes"):
-        episode_mapping[old_idx] = new_episode_idx
-        new_episode_idx += 1
-
-    # Copy data files and update indices
-    _copy_and_reindex_data(dataset, new_meta, episode_mapping)
-
-    # Copy video files if present
-    if dataset.meta.video_keys:
-        _copy_and_reindex_videos(dataset, new_meta, episode_mapping)
-
-    # Create new dataset instance
-    new_dataset = LeRobotDataset(
-        repo_id=repo_id,
-        root=output_dir,
-        image_transforms=dataset.image_transforms,
-        delta_timestamps=dataset.delta_timestamps,
-        tolerance_s=dataset.tolerance_s,
-    )
-
-    logging.info(f"Created new dataset with {len(episodes_to_keep)} episodes")
-    return new_dataset
-
-
-def split_dataset(
-    dataset: LeRobotDataset,
-    splits: dict[str, list[int]] | dict[str, float],
-    output_dir: str | Path | None = None,
-) -> dict[str, LeRobotDataset]:
-    """Split a LeRobotDataset into multiple smaller datasets.
-
-    Args:
-        dataset: The source LeRobotDataset to split.
-        splits: Either a dict mapping split names to episode indices, or a dict mapping
-                split names to fractions (must sum to <= 1.0).
-        output_dir: Base directory for output datasets. If None, uses default location.
-
-    Returns:
-        dict[str, LeRobotDataset]: Dictionary mapping split names to new datasets.
-
-    Examples:
-        # Split by specific episodes
-        splits = {"train": [0, 1, 2], "val": [3, 4]}
-        datasets = split_dataset(dataset, splits)
-
-        # Split by fractions
-        splits = {"train": 0.8, "val": 0.2}
-        datasets = split_dataset(dataset, splits)
-    """
-    if not splits:
-        raise ValueError("No splits provided")
-
-    # Convert fractions to episode indices if needed
-    if all(isinstance(v, float) for v in splits.values()):
-        splits = _fractions_to_episode_indices(dataset.meta.total_episodes, splits)
-
-    # Validate episodes
-    all_episodes = set()
-    for split_name, episodes in splits.items():
-        if not episodes:
-            raise ValueError(f"Split '{split_name}' has no episodes")
-        episode_set = set(episodes)
-        if episode_set & all_episodes:
-            raise ValueError("Episodes cannot appear in multiple splits")
-        all_episodes.update(episode_set)
-
-    # Validate all episodes are valid
-    valid_indices = set(range(dataset.meta.total_episodes))
-    invalid = all_episodes - valid_indices
-    if invalid:
-        raise ValueError(f"Invalid episode indices: {invalid}")
-
-    if output_dir is None:
-        output_dir = HF_LEROBOT_HOME
-    else:
-        output_dir = Path(output_dir)
-
-    result_datasets = {}
-
-    for split_name, episodes in splits.items():
-        logging.info(f"Creating split '{split_name}' with {len(episodes)} episodes")
-
-        # Create repo_id for split
-        split_repo_id = f"{dataset.repo_id}_{split_name}"
-        split_output_dir = output_dir / split_repo_id
-
-        # Create episode mapping
-        episode_mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(sorted(episodes))}
-
-        # Create new dataset metadata
-        new_meta = LeRobotDatasetMetadata.create(
-            repo_id=split_repo_id,
-            fps=dataset.meta.fps,
-            features=dataset.meta.features,
-            robot_type=dataset.meta.robot_type,
-            root=split_output_dir,
-            use_videos=len(dataset.meta.video_keys) > 0,
-        )
-
-        # Copy data and videos
-        _copy_and_reindex_data(dataset, new_meta, episode_mapping)
-        if dataset.meta.video_keys:
-            _copy_and_reindex_videos(dataset, new_meta, episode_mapping)
-
-        # Create new dataset instance
-        new_dataset = LeRobotDataset(
-            repo_id=split_repo_id,
-            root=split_output_dir,
-            image_transforms=dataset.image_transforms,
-            delta_timestamps=dataset.delta_timestamps,
-            tolerance_s=dataset.tolerance_s,
-        )
-
-        result_datasets[split_name] = new_dataset
-
-    return result_datasets
-
-
-def merge_datasets(
-    datasets: list[LeRobotDataset],
-    output_repo_id: str,
-    output_dir: str | Path | None = None,
-) -> LeRobotDataset:
-    """Merge multiple LeRobotDatasets into a single dataset.
-
-    This is a wrapper around the aggregate_datasets functionality with a cleaner API.
-
-    Args:
-        datasets: List of LeRobotDatasets to merge.
-        output_repo_id: Repository ID for the merged dataset.
-        output_dir: Directory to save the merged dataset. If None, uses default location.
-
-    Returns:
-        LeRobotDataset: The merged dataset.
-    """
-    if not datasets:
-        raise ValueError("No datasets to merge")
-
-    if output_dir is None:
-        output_dir = HF_LEROBOT_HOME / output_repo_id
-    else:
-        output_dir = Path(output_dir)
-
-    # Extract repo_ids and roots
-    repo_ids = [ds.repo_id for ds in datasets]
-    roots = [ds.root for ds in datasets]
-
-    # Call aggregate_datasets
-    aggregate_datasets(
-        repo_ids=repo_ids,
-        aggr_repo_id=output_repo_id,
-        roots=roots,
-        aggr_root=output_dir,
-    )
-
-    # Create and return the merged dataset
-    merged_dataset = LeRobotDataset(
-        repo_id=output_repo_id,
-        root=output_dir,
-        image_transforms=datasets[0].image_transforms,
-        delta_timestamps=datasets[0].delta_timestamps,
-        tolerance_s=datasets[0].tolerance_s,
-    )
-
-    return merged_dataset
-
-
-def add_feature(
-    dataset: LeRobotDataset,
-    feature_name: str,
-    feature_values: np.ndarray | torch.Tensor | Callable,
-    feature_info: dict,
-    output_dir: str | Path | None = None,
-    repo_id: str | None = None,
-) -> LeRobotDataset:
-    """Add a new feature to a LeRobotDataset.
-
-    Args:
-        dataset: The source LeRobotDataset.
-        feature_name: Name of the new feature.
-        feature_values: Either:
-            - Array/tensor of shape (num_frames, ...) with values for each frame
-            - Callable that takes (frame_dict, episode_index, frame_index) and returns feature value
-        feature_info: Dictionary with feature metadata (dtype, shape, names).
-        output_dir: Directory to save the new dataset. If None, uses default location.
-        repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
-
-    Returns:
-        LeRobotDataset: New dataset with the added feature.
-    """
-    if feature_name in dataset.meta.features:
-        raise ValueError(f"Feature '{feature_name}' already exists in dataset")
-
-    if repo_id is None:
-        repo_id = f"{dataset.repo_id}_modified"
-    if output_dir is None:
-        output_dir = HF_LEROBOT_HOME / repo_id
-    else:
-        output_dir = Path(output_dir)
-
-    # Validate feature_info
-    required_keys = {"dtype", "shape"}
-    if not required_keys.issubset(feature_info.keys()):
-        raise ValueError(f"feature_info must contain keys: {required_keys}")
-
-    # Create new features dict
-    new_features = dataset.meta.features.copy()
-    new_features[feature_name] = feature_info
-
-    # Create new dataset metadata
-    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=len(dataset.meta.video_keys) > 0,
-    )
-
-    # Process data with new feature
-    _copy_data_with_feature_changes(
-        dataset=dataset,
-        new_meta=new_meta,
-        add_features={feature_name: (feature_values, feature_info)},
-    )
-
-    # Copy videos if present
-    if dataset.meta.video_keys:
-        _copy_videos(dataset, new_meta)
-
-    # Create new dataset instance
-    new_dataset = LeRobotDataset(
-        repo_id=repo_id,
-        root=output_dir,
-        image_transforms=dataset.image_transforms,
-        delta_timestamps=dataset.delta_timestamps,
-        tolerance_s=dataset.tolerance_s,
-    )
-
-    return new_dataset
-
-
-def remove_feature(
-    dataset: LeRobotDataset,
-    feature_names: str | list[str],
-    output_dir: str | Path | None = None,
-    repo_id: str | None = None,
-) -> LeRobotDataset:
-    """Remove features from a LeRobotDataset.
-
-    Args:
-        dataset: The source LeRobotDataset.
-        feature_names: Name(s) of features to remove. Can be a single string or list.
-        output_dir: Directory to save the new dataset. If None, uses default location.
-        repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
-
-    Returns:
-        LeRobotDataset: New dataset with features removed.
-    """
-    if isinstance(feature_names, str):
-        feature_names = [feature_names]
-
-    # Validate features exist
-    for name in feature_names:
-        if name not in dataset.meta.features:
-            raise ValueError(f"Feature '{name}' not found in dataset")
-
-    # Check if trying to remove required features
-    required_features = {"timestamp", "frame_index", "episode_index", "index", "task_index"}
-    if any(name in required_features for name in feature_names):
-        raise ValueError(f"Cannot remove required features: {required_features}")
-
-    if repo_id is None:
-        repo_id = f"{dataset.repo_id}_modified"
-    if output_dir is None:
-        output_dir = HF_LEROBOT_HOME / repo_id
-    else:
-        output_dir = Path(output_dir)
-
-    # Create new features dict
-    new_features = {k: v for k, v in dataset.meta.features.items() if k not in feature_names}
-
-    # Check if removing video features
-    video_keys_to_remove = [name for name in feature_names if name in dataset.meta.video_keys]
-
-    # Check if videos will remain after removal
-    remaining_video_keys = [k for k in dataset.meta.video_keys if k not in video_keys_to_remove]
-
-    # Create new dataset metadata
-    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=len(remaining_video_keys) > 0,
-    )
-
-    # Process data with removed features
-    _copy_data_with_feature_changes(
-        dataset=dataset,
-        new_meta=new_meta,
-        remove_features=feature_names,
-    )
-
-    # Copy videos (excluding removed ones)
-    if new_meta.video_keys:
-        _copy_videos(dataset, new_meta, exclude_keys=video_keys_to_remove)
-
-    # Create new dataset instance
-    new_dataset = LeRobotDataset(
-        repo_id=repo_id,
-        root=output_dir,
-        image_transforms=dataset.image_transforms,
-        delta_timestamps=dataset.delta_timestamps,
-        tolerance_s=dataset.tolerance_s,
-    )
-
-    return new_dataset
-
-
-# Helper functions
-
-
-def _fractions_to_episode_indices(
-    total_episodes: int,
-    splits: dict[str, float],
-) -> dict[str, list[int]]:
-    """Convert split fractions to episode indices."""
-    if sum(splits.values()) > 1.0:
-        raise ValueError("Split fractions must sum to <= 1.0")
-
-    indices = list(range(total_episodes))
-    result = {}
-    start_idx = 0
-
-    for split_name, fraction in splits.items():
-        num_episodes = int(total_episodes * fraction)
-        end_idx = start_idx + num_episodes
-        if split_name == list(splits.keys())[-1]:  # Last split gets remaining episodes
-            end_idx = total_episodes
-        result[split_name] = indices[start_idx:end_idx]
-        start_idx = end_idx
-
-    return result
-
-
-def _copy_and_reindex_data(
-    src_dataset: LeRobotDataset,
-    dst_meta: LeRobotDatasetMetadata,
-    episode_mapping: dict[int, int],
-) -> None:
-    """Copy data files and reindex episodes."""
-    # Get unique data files from episodes to keep
-    file_paths = set()
-    for old_idx in episode_mapping:
-        file_paths.add(src_dataset.meta.get_data_file_path(old_idx))
-
-    # Track global index
-    global_index = 0
-    chunk_idx, file_idx = 0, 0
-
-    # Process each data file
-    for src_path in tqdm(sorted(file_paths), desc="Processing data files"):
-        df = pd.read_parquet(src_dataset.root / src_path)
-
-        # Filter to keep only mapped episodes
-        mask = df["episode_index"].isin(episode_mapping.keys())
-        df = df[mask].copy()
-
-        if len(df) == 0:
-            continue
-
-        # Update episode indices
-        df["episode_index"] = df["episode_index"].map(episode_mapping)
-
-        # Update global index to be continuous
-        df["index"] = range(global_index, global_index + len(df))
-        global_index += len(df)
-
-        # Update task indices if needed
-        if dst_meta.tasks is None:
-            # Get unique tasks from filtered data
-            task_indices = df["task_index"].unique()
-            tasks = [src_dataset.meta.tasks.iloc[idx].name for idx in task_indices]
-            dst_meta.save_episode_tasks(list(set(tasks)))
-
-        # Remap task indices
-        task_mapping = {}
-        for old_task_idx in df["task_index"].unique():
-            task_name = src_dataset.meta.tasks.iloc[old_task_idx].name
-            new_task_idx = dst_meta.get_task_index(task_name)
-            task_mapping[old_task_idx] = new_task_idx
-        df["task_index"] = df["task_index"].map(task_mapping)
-
-        # Save processed data
-        chunk_idx, file_idx = _save_data_chunk(df, dst_meta, chunk_idx, file_idx)
-
-    # Process episodes metadata
-    _copy_and_reindex_episodes_metadata(src_dataset, dst_meta, episode_mapping)
-
-
-def _copy_and_reindex_videos(
-    src_dataset: LeRobotDataset,
-    dst_meta: LeRobotDatasetMetadata,
-    episode_mapping: dict[int, int],
-) -> None:
-    """Copy video files and update metadata."""
-    for video_key in src_dataset.meta.video_keys:
-        video_files = set()
-        for old_idx in episode_mapping:
-            video_files.add(src_dataset.meta.get_video_file_path(old_idx, video_key))
-
-        chunk_idx, file_idx = 0, 0
-
-        for src_path in tqdm(sorted(video_files), desc=f"Processing {video_key} videos"):
-            dst_path = dst_meta.root / DEFAULT_VIDEO_PATH.format(
-                video_key=video_key,
-                chunk_index=chunk_idx,
-                file_index=file_idx,
-            )
-            dst_path.parent.mkdir(parents=True, exist_ok=True)
-
-            # For simplicity, copy entire video files
-            # In production, you might want to extract only relevant segments
-            shutil.copy(src_dataset.root / src_path, dst_path)
-
-            # Update indices for next file
-            file_size = get_video_size_in_mb(dst_path)
-            if file_size >= DEFAULT_VIDEO_FILE_SIZE_IN_MB * 0.9:  # 90% threshold
-                chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
-
-
-def _copy_and_reindex_episodes_metadata(
-    src_dataset: LeRobotDataset,
-    dst_meta: LeRobotDatasetMetadata,
-    episode_mapping: dict[int, int],
-) -> None:
-    """Copy and reindex episodes metadata."""
-    all_stats = []
-    frame_offset = 0
-
-    for old_idx, new_idx in tqdm(
-        sorted(episode_mapping.items(), key=lambda x: x[1]), desc="Processing episodes metadata"
-    ):
-        # Get episode from source
-        src_episode = src_dataset.meta.episodes[old_idx]
-
-        # Create episode dict
-        episode_dict = {
-            "episode_index": new_idx,
-            "tasks": src_episode["tasks"],  # Already a list of task names
-            "length": src_episode["length"],
-        }
-
-        # Copy other metadata
-        episode_metadata = {
-            "data/chunk_index": 0,  # Will be recalculated when saving
-            "data/file_index": 0,  # Will be recalculated when saving
-            "dataset_from_index": frame_offset,
-            "dataset_to_index": frame_offset + src_episode["length"],
-        }
-
-        # Update frame offset for next episode
-        frame_offset += src_episode["length"]
-
-        # Copy stats metadata
-        for key in src_episode.keys():
-            if key.startswith("stats/"):
-                episode_dict[key] = src_episode[key]
-
-        # Add episode metadata
-        stats_dict = {
-            key.replace("stats/", ""): value
-            for key, value in episode_dict.items()
-            if key.startswith("stats/")
-        }
-        all_stats.append(stats_dict)
-
-        # Calculate stats from dict
-        episode_stats = {}
-        for key in dst_meta.features:
-            if key in stats_dict:
-                episode_stats[key] = stats_dict[key]
-
-        dst_meta.save_episode(
-            new_idx, episode_dict["length"], episode_dict["tasks"], episode_stats, episode_metadata
-        )
-
-    # Aggregate all stats
-    if all_stats:
-        aggregated_stats = aggregate_stats(all_stats)
-        write_stats(aggregated_stats, dst_meta.root)
-
-
-def _save_data_chunk(
-    df: pd.DataFrame,
-    meta: LeRobotDatasetMetadata,
-    chunk_idx: int = 0,
-    file_idx: int = 0,
-) -> tuple[int, int]:
-    """Save a data chunk and return updated indices."""
-    path = meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
-    path.parent.mkdir(parents=True, exist_ok=True)
-
-    if len(meta.image_keys) > 0:
-        to_parquet_with_hf_images(df, path)
-    else:
-        df.to_parquet(path)
-
-    # Check if we need to rotate files
-    file_size = get_parquet_file_size_in_mb(path)
-    if file_size >= DEFAULT_DATA_FILE_SIZE_IN_MB * 0.9:  # 90% threshold
-        chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
-
-    return chunk_idx, file_idx
-
-
-def _copy_data_with_feature_changes(
-    dataset: LeRobotDataset,
-    new_meta: LeRobotDatasetMetadata,
-    add_features: dict[str, tuple] | None = None,
-    remove_features: list[str] | None = None,
-) -> None:
-    """Copy data while adding or removing features."""
-    # Get all unique data files
-    file_paths = set()
-    for ep_idx in range(dataset.meta.total_episodes):
-        file_paths.add(dataset.meta.get_data_file_path(ep_idx))
-
-    frame_idx = 0
-
-    # Process each data file
-    for src_path in tqdm(sorted(file_paths), desc="Processing data files"):
-        df = pd.read_parquet(dataset.root / src_path)
-
-        # Remove features
-        if remove_features:
-            df = df.drop(columns=remove_features, errors="ignore")
-
-        # Add features
-        if add_features:
-            for feature_name, (values, _) in add_features.items():
-                if callable(values):
-                    # Compute values for each frame
-                    feature_values = []
-                    for _, row in df.iterrows():
-                        ep_idx = row["episode_index"]
-                        frame_in_ep = row["frame_index"]
-                        value = values(row.to_dict(), ep_idx, frame_in_ep)
-                        # Convert numpy arrays to scalars for single-element arrays
-                        if isinstance(value, np.ndarray) and value.size == 1:
-                            value = value.item()
-                        feature_values.append(value)
-                    df[feature_name] = feature_values
-                else:
-                    # Use provided values
-                    end_idx = frame_idx + len(df)
-                    # Convert to list to ensure proper shape handling
-                    feature_slice = values[frame_idx:end_idx]
-                    if len(feature_slice.shape) > 1 and feature_slice.shape[1] == 1:
-                        # Flatten single-element arrays to scalars for pandas
-                        df[feature_name] = feature_slice.flatten()
-                    else:
-                        df[feature_name] = feature_slice
-                    frame_idx = end_idx
-
-        # Save chunk
-        _save_data_chunk(df, new_meta)
-
-    # Copy episodes metadata and update stats
-    _copy_episodes_metadata_and_stats(dataset, new_meta)
-
-
-def _copy_videos(
-    src_dataset: LeRobotDataset,
-    dst_meta: LeRobotDatasetMetadata,
-    exclude_keys: list[str] | None = None,
-) -> None:
-    """Copy video files, optionally excluding certain keys."""
-    if exclude_keys is None:
-        exclude_keys = []
-
-    for video_key in src_dataset.meta.video_keys:
-        if video_key in exclude_keys:
-            continue
-
-        # Get all video files for this key
-        video_files = set()
-        for ep_idx in range(src_dataset.meta.total_episodes):
-            video_files.add(src_dataset.meta.get_video_file_path(ep_idx, video_key))
-
-        # Copy video files
-        for src_path in tqdm(sorted(video_files), desc=f"Copying {video_key} videos"):
-            # Maintain same structure
-            rel_path = src_path.relative_to(src_dataset.root)
-            dst_path = dst_meta.root / rel_path
-            dst_path.parent.mkdir(parents=True, exist_ok=True)
-            shutil.copy(src_dataset.root / src_path, dst_path)
-
-
-def _copy_episodes_metadata_and_stats(
-    src_dataset: LeRobotDataset,
-    dst_meta: LeRobotDatasetMetadata,
-) -> None:
-    """Copy episodes metadata and recalculate stats."""
-    # Copy tasks
-    if src_dataset.meta.tasks is not None:
-        write_tasks(src_dataset.meta.tasks, dst_meta.root)
-        dst_meta.tasks = src_dataset.meta.tasks.copy()
-
-    # Copy episodes metadata files
-    episodes_dir = src_dataset.root / "meta/episodes"
-    dst_episodes_dir = dst_meta.root / "meta/episodes"
-    if episodes_dir.exists():
-        shutil.copytree(episodes_dir, dst_episodes_dir, dirs_exist_ok=True)
-
-    # Update info
-    dst_meta.info.update(
-        {
-            "total_episodes": src_dataset.meta.total_episodes,
-            "total_frames": src_dataset.meta.total_frames,
-            "total_tasks": src_dataset.meta.total_tasks,
-            "splits": src_dataset.meta.info.get("splits", {"train": f"0:{src_dataset.meta.total_episodes}"}),
-        }
-    )
-
-    # Update video info if needed
-    if dst_meta.video_keys and src_dataset.meta.video_keys:
-        for key in dst_meta.video_keys:
-            if key in src_dataset.meta.features:
-                dst_meta.info["features"][key]["info"] = src_dataset.meta.info["features"][key].get(
-                    "info", {}
-                )
-
-    write_info(dst_meta.info, dst_meta.root)
-
-    # Recalculate stats if features changed
-    if set(dst_meta.features.keys()) != set(src_dataset.meta.features.keys()):
-        # Need to recalculate stats
-        logging.info("Recalculating dataset statistics...")
-        # This is a simplified version - in production you'd want to properly recalculate
-        if src_dataset.meta.stats:
-            new_stats = {}
-            for key in dst_meta.features:
-                if key in src_dataset.meta.stats:
-                    new_stats[key] = src_dataset.meta.stats[key]
-            write_stats(new_stats, dst_meta.root)
-    else:
-        # Copy existing stats
-        if src_dataset.meta.stats:
-            write_stats(src_dataset.meta.stats, dst_meta.root)
@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import contextlib
+import gc
 import logging
 import shutil
 import tempfile
@@ -27,7 +28,6 @@ import pandas as pd
 import PIL.Image
 import torch
 import torch.utils
-from datasets import Dataset
 from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.constants import REPOCARD_NAME
 from huggingface_hub.errors import RevisionNotFoundError
@@ -48,12 +48,12 @@ from lerobot.datasets.utils import (
    embed_images,
    flatten_dict,
    get_delta_indices,
+    get_hf_dataset_cache_dir,
    get_hf_dataset_size_in_mb,
    get_hf_features_from_features,
    get_parquet_file_size_in_mb,
    get_parquet_num_frames,
    get_safe_version,
-    get_video_duration_in_s,
    get_video_size_in_mb,
    hf_transform_to_torch,
    is_valid_version,
@@ -77,6 +77,7 @@ from lerobot.datasets.video_utils import (
    decode_video_frames,
    encode_video_frames,
    get_safe_default_codec,
+    get_video_duration_in_s,
    get_video_info,
 )

@@ -104,9 +105,6 @@ class LeRobotDatasetMetadata:
                self.revision = get_safe_version(self.repo_id, self.revision)

            (self.root / "meta").mkdir(exist_ok=True, parents=True)
-            # TODO(rcadene): instead of downloading all episodes metadata files,
-            # download only the ones associated to the requested episodes. This would
-            # require adding `episodes: list[int]` as argument.
            self.pull_from_repo(allow_patterns="meta/")
            self.load_metadata()

@@ -273,7 +271,7 @@ class LeRobotDatasetMetadata:
        """
        # Convert buffer into HF Dataset
        episode_dict = {key: [value] for key, value in episode_dict.items()}
-        ep_dataset = Dataset.from_dict(episode_dict)
+        ep_dataset = datasets.Dataset.from_dict(episode_dict)
        ep_size_in_mb = get_hf_dataset_size_in_mb(ep_dataset)
        df = pd.DataFrame(ep_dataset)
        num_frames = episode_dict["length"][0]
@@ -309,14 +307,21 @@ class LeRobotDatasetMetadata:
                latest_df = pd.read_parquet(latest_path)
                df = pd.concat([latest_df, df], ignore_index=True)

+                # Memort optimization
+                del latest_df
+                gc.collect()
+
        # Write the resulting dataframe from RAM to disk
        path = self.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
        path.parent.mkdir(parents=True, exist_ok=True)
        df.to_parquet(path, index=False)

-        # Update the Hugging Face dataset by reloading it.
-        # This process should be fast because only the latest Parquet file has been modified.
-        # Therefore, only this file needs to be converted to PyArrow; the rest is loaded from the PyArrow memory-mapped cache.
+        if self.episodes is not None:
+            # Remove the episodes cache directory, necessary to avoid cache bloat
+            cached_dir = get_hf_dataset_cache_dir(self.episodes)
+            if cached_dir is not None:
+                shutil.rmtree(cached_dir)
+
        self.episodes = load_episodes(self.root)

    def save_episode(
@@ -358,6 +363,53 @@ class LeRobotDatasetMetadata:
                video_path = self.root / self.get_video_file_path(ep_index=0, vid_key=key)
                self.info["features"][key]["info"] = get_video_info(video_path)

+    def update_chunk_settings(
+        self,
+        chunks_size: int | None = None,
+        data_files_size_in_mb: int | None = None,
+        video_files_size_in_mb: int | None = None,
+    ) -> None:
+        """Update chunk and file size settings after dataset creation.
+
+        This allows users to customize storage organization without modifying the constructor.
+        These settings control how episodes are chunked and how large files can grow before
+        creating new ones.
+
+        Args:
+            chunks_size: Maximum number of files per chunk directory. If None, keeps current value.
+            data_files_size_in_mb: Maximum size for data parquet files in MB. If None, keeps current value.
+            video_files_size_in_mb: Maximum size for video files in MB. If None, keeps current value.
+        """
+        if chunks_size is not None:
+            if chunks_size <= 0:
+                raise ValueError(f"chunks_size must be positive, got {chunks_size}")
+            self.info["chunks_size"] = chunks_size
+
+        if data_files_size_in_mb is not None:
+            if data_files_size_in_mb <= 0:
+                raise ValueError(f"data_files_size_in_mb must be positive, got {data_files_size_in_mb}")
+            self.info["data_files_size_in_mb"] = data_files_size_in_mb
+
+        if video_files_size_in_mb is not None:
+            if video_files_size_in_mb <= 0:
+                raise ValueError(f"video_files_size_in_mb must be positive, got {video_files_size_in_mb}")
+            self.info["video_files_size_in_mb"] = video_files_size_in_mb
+
+        # Update the info file on disk
+        write_info(self.info, self.root)
+
+    def get_chunk_settings(self) -> dict[str, int]:
+        """Get current chunk and file size settings.
+
+        Returns:
+            Dict containing chunks_size, data_files_size_in_mb, and video_files_size_in_mb.
+        """
+        return {
+            "chunks_size": self.chunks_size,
+            "data_files_size_in_mb": self.data_files_size_in_mb,
+            "video_files_size_in_mb": self.video_files_size_in_mb,
+        }
+
    def __repr__(self):
        feature_keys = list(self.features)
        return (
@@ -407,7 +459,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
        root: str | Path | None = None,
        episodes: list[int] | None = None,
        image_transforms: Callable | None = None,
-        delta_timestamps: dict[list[float]] | None = None,
+        delta_timestamps: dict[str, list[float]] | None = None,
        tolerance_s: float = 1e-4,
        revision: str | None = None,
        force_cache_sync: bool = False,
@@ -675,7 +727,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
            return False

        # Get available episode indices from cached dataset
-        available_episodes = set(self.hf_dataset["episode_index"])
+        available_episodes = {
+            ep_idx.item() if isinstance(ep_idx, torch.Tensor) else ep_idx
+            for ep_idx in self.hf_dataset["episode_index"]
+        }

        # Determine requested episodes
        if self.episodes is None:
@@ -993,6 +1048,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
                latest_df = pd.read_parquet(latest_path)
                df = pd.concat([latest_df, df], ignore_index=True)

+                # Memort optimization
+                del latest_df
+                gc.collect()
+
        # Write the resulting dataframe from RAM to disk
        path = self.root / self.meta.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
        path.parent.mkdir(parents=True, exist_ok=True)
@@ -1001,9 +1060,12 @@ class LeRobotDataset(torch.utils.data.Dataset):
        else:
            df.to_parquet(path)

-        # Update the Hugging Face dataset by reloading it.
-        # This process should be fast because only the latest Parquet file has been modified.
-        # Therefore, only this file needs to be converted to PyArrow; the rest is loaded from the PyArrow memory-mapped cache.
+        if self.hf_dataset is not None:
+            # Remove hf dataset cache directory, necessary to avoid cache bloat
+            cached_dir = get_hf_dataset_cache_dir(self.hf_dataset)
+            if cached_dir is not None:
+                shutil.rmtree(cached_dir)
+
        self.hf_dataset = self.load_hf_dataset()

        metadata = {
@@ -1023,7 +1085,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
        if self.meta.episodes is None:
            # Initialize indices for a new dataset made of the first episode data
            chunk_idx, file_idx = 0, 0
-            latest_duration_in_s = 0
+            latest_duration_in_s = 0.0
            new_path = self.root / self.meta.video_path.format(
                video_key=video_key, chunk_index=chunk_idx, file_index=file_idx
            )
@@ -1049,6 +1111,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
                )
                new_path.parent.mkdir(parents=True, exist_ok=True)
                shutil.move(str(ep_path), str(new_path))
+                latest_duration_in_s = 0.0
            else:
                # Update latest video file
                concat_video_files([latest_path, ep_path], self.root, video_key, chunk_idx, file_idx)
@@ -1169,7 +1232,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
        root: str | Path | None = None,
        episodes: dict | None = None,
        image_transforms: Callable | None = None,
-        delta_timestamps: dict[list[float]] | None = None,
+        delta_timestamps: dict[str, list[float]] | None = None,
        tolerances_s: dict | None = None,
        download_videos: bool = True,
        video_backend: str | None = None,
@@ -17,11 +17,9 @@ import contextlib
 import importlib.resources
 import json
 import logging
-import subprocess
 from collections.abc import Iterator
 from pathlib import Path
 from pprint import pformat
-from types import SimpleNamespace
 from typing import Any

 import datasets
@@ -40,11 +38,10 @@ from torchvision import transforms

 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.datasets.backward_compatibility import (
-    V21_MESSAGE,
+    FUTURE_MESSAGE,
    BackwardCompatibilityError,
    ForwardCompatibilityError,
 )
-from lerobot.robots import Robot
 from lerobot.utils.utils import is_valid_numpy_dtype_string

 DEFAULT_CHUNK_SIZE = 1000  # Max number of files per chunk
@@ -105,6 +102,12 @@ def get_hf_dataset_size_in_mb(hf_ds: Dataset) -> int:
    return hf_ds.data.nbytes // (1024**2)


+def get_hf_dataset_cache_dir(hf_ds: Dataset) -> Path | None:
+    if hf_ds.cache_files is None or len(hf_ds.cache_files) == 0:
+        return None
+    return Path(hf_ds.cache_files[0]["filename"]).parents[2]
+
+
 def update_chunk_file_indices(chunk_idx: int, file_idx: int, chunks_size: int) -> tuple[int, int]:
    if file_idx == chunks_size - 1:
        file_idx = 0
@@ -143,26 +146,6 @@ def get_video_size_in_mb(mp4_path: Path) -> float:
    return file_size_mb


-def get_video_duration_in_s(mp4_file: Path) -> float:
-    # TODO(rcadene): move to video_utils.py
-    command = [
-        "ffprobe",
-        "-v",
-        "error",
-        "-show_entries",
-        "format=duration",
-        "-of",
-        "default=noprint_wrappers=1:nokey=1",
-        str(mp4_file),
-    ]
-    result = subprocess.run(
-        command,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-    )
-    return float(result.stdout)
-
-
 def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict:
    """Flatten a nested dictionary structure by collapsing nested keys into one key with a separator.

@@ -272,8 +255,21 @@ def load_tasks(local_dir: Path) -> pandas.DataFrame:


 def write_episodes(episodes: Dataset, local_dir: Path) -> None:
-    if get_hf_dataset_size_in_mb(episodes) > DEFAULT_DATA_FILE_SIZE_IN_MB:
-        raise NotImplementedError("Contact a maintainer.")
+    """Write episode metadata to a parquet file in the LeRobot v3.0 format.
+    This function writes episode-level metadata to a single parquet file.
+    Used primarily during dataset conversion (v2.1 → v3.0) and in test fixtures.
+
+    Args:
+        episodes: HuggingFace Dataset containing episode metadata
+        local_dir: Root directory where the dataset will be stored
+    """
+    episode_size_mb = get_hf_dataset_size_in_mb(episodes)
+    if episode_size_mb > DEFAULT_DATA_FILE_SIZE_IN_MB:
+        raise NotImplementedError(
+            f"Episodes dataset is too large ({episode_size_mb} MB) to write to a single file. "
+            f"The current limit is {DEFAULT_DATA_FILE_SIZE_IN_MB} MB. "
+            "This function only supports single-file episode metadata. "
+        )

    fpath = local_dir / DEFAULT_EPISODES_PATH.format(chunk_index=0, file_index=0)
    fpath.parent.mkdir(parents=True, exist_ok=True)
@@ -352,7 +348,7 @@ def check_version_compatibility(
    if v_check.major < v_current.major and enforce_breaking_major:
        raise BackwardCompatibilityError(repo_id, v_check)
    elif v_check.minor < v_current.minor:
-        logging.warning(V21_MESSAGE.format(repo_id=repo_id, version=v_check))
+        logging.warning(FUTURE_MESSAGE.format(repo_id=repo_id, version=v_check))


 def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
@@ -492,17 +488,6 @@ def build_dataset_frame(
    return frame


-def get_features_from_robot(robot: Robot, use_videos: bool = True) -> dict:
-    # TODO(rcadene): add fps for each feature
-    camera_ft = {}
-    if robot.cameras:
-        camera_ft = {
-            key: {"dtype": "video" if use_videos else "image", **ft}
-            for key, ft in robot.camera_features.items()
-        }
-    return {**robot.motor_features, **camera_ft, **DEFAULT_FEATURES}
-
-
 def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFeature]:
    # TODO(aliberts): Implement "type" in dataset features and simplify this
    policy_features = {}
@@ -540,6 +525,9 @@ def create_empty_dataset_info(
    features: dict,
    use_videos: bool,
    robot_type: str | None = None,
+    chunks_size: int | None = None,
+    data_files_size_in_mb: int | None = None,
+    video_files_size_in_mb: int | None = None,
 ) -> dict:
    return {
        "codebase_version": codebase_version,
@@ -547,9 +535,9 @@ def create_empty_dataset_info(
        "total_episodes": 0,
        "total_frames": 0,
        "total_tasks": 0,
-        "chunks_size": DEFAULT_CHUNK_SIZE,
-        "data_files_size_in_mb": DEFAULT_DATA_FILE_SIZE_IN_MB,
-        "video_files_size_in_mb": DEFAULT_VIDEO_FILE_SIZE_IN_MB,
+        "chunks_size": chunks_size or DEFAULT_CHUNK_SIZE,
+        "data_files_size_in_mb": data_files_size_in_mb or DEFAULT_DATA_FILE_SIZE_IN_MB,
+        "video_files_size_in_mb": video_files_size_in_mb or DEFAULT_VIDEO_FILE_SIZE_IN_MB,
        "fps": fps,
        "splits": {},
        "data_path": DEFAULT_DATA_PATH,
@@ -630,7 +618,7 @@ def create_lerobot_dataset_card(
    **kwargs,
 ) -> DatasetCard:
    """
-    Keyword arguments will be used to replace values in ./lerobot/datasets/card_template.md.
+    Keyword arguments will be used to replace values in src/lerobot/datasets/card_template.md.
    Note: If specified, license must be one of https://huggingface.co/docs/hub/repositories-licenses.
    """
    card_tags = ["LeRobot"]
@@ -662,61 +650,6 @@ def create_lerobot_dataset_card(
    )


-class IterableNamespace(SimpleNamespace):
-    """
-    A namespace object that supports both dictionary-like iteration and dot notation access.
-    Automatically converts nested dictionaries into IterableNamespaces.
-
-    This class extends SimpleNamespace to provide:
-    - Dictionary-style iteration over keys
-    - Access to items via both dot notation (obj.key) and brackets (obj["key"])
-    - Dictionary-like methods: items(), keys(), values()
-    - Recursive conversion of nested dictionaries
-
-    Args:
-        dictionary: Optional dictionary to initialize the namespace
-        **kwargs: Additional keyword arguments passed to SimpleNamespace
-
-    Examples:
-        >>> data = {"name": "Alice", "details": {"age": 25}}
-        >>> ns = IterableNamespace(data)
-        >>> ns.name
-        'Alice'
-        >>> ns.details.age
-        25
-        >>> list(ns.keys())
-        ['name', 'details']
-        >>> for key, value in ns.items():
-        ...     print(f"{key}: {value}")
-        name: Alice
-        details: IterableNamespace(age=25)
-    """
-
-    def __init__(self, dictionary: dict[str, Any] = None, **kwargs):
-        super().__init__(**kwargs)
-        if dictionary is not None:
-            for key, value in dictionary.items():
-                if isinstance(value, dict):
-                    setattr(self, key, IterableNamespace(value))
-                else:
-                    setattr(self, key, value)
-
-    def __iter__(self) -> Iterator[str]:
-        return iter(vars(self))
-
-    def __getitem__(self, key: str) -> Any:
-        return vars(self)[key]
-
-    def items(self):
-        return vars(self).items()
-
-    def values(self):
-        return vars(self).values()
-
-    def keys(self):
-        return vars(self).keys()
-
-
 def validate_frame(frame: dict, features: dict) -> None:
    expected_features = set(features) - set(DEFAULT_FEATURES)
    actual_features = set(frame)
@@ -1,884 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-This script is for internal use to convert all datasets under the 'lerobot' hub user account to v2.
-
-Note: Since the original Aloha datasets don't use shadow motors, you need to comment those out in
-lerobot/configs/robot/aloha.yaml before running this script.
-"""
-
-import traceback
-from pathlib import Path
-from textwrap import dedent
-
-from lerobot import available_datasets
-from lerobot.datasets.v2.convert_dataset_v1_to_v2 import convert_dataset
-from lerobot.robots.aloha.configuration_aloha import AlohaRobotConfig
-
-LOCAL_DIR = Path("data/")
-
-# spellchecker:off
-ALOHA_MOBILE_INFO = {
-    "robot_config": AlohaRobotConfig(),
-    "license": "mit",
-    "url": "https://mobile-aloha.github.io/",
-    "paper": "https://huggingface.co/papers/2401.02117",
-    "citation_bibtex": dedent(r"""
-        @inproceedings{fu2024mobile,
-            author    = {Fu, Zipeng and Zhao, Tony Z. and Finn, Chelsea},
-            title     = {Mobile ALOHA: Learning Bimanual Mobile Manipulation with Low-Cost Whole-Body Teleoperation},
-            booktitle = {arXiv},
-            year      = {2024},
-        }""").lstrip(),
-}
-ALOHA_STATIC_INFO = {
-    "robot_config": AlohaRobotConfig(),
-    "license": "mit",
-    "url": "https://tonyzhaozh.github.io/aloha/",
-    "paper": "https://huggingface.co/papers/2304.13705",
-    "citation_bibtex": dedent(r"""
-        @article{Zhao2023LearningFB,
-            title={Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware},
-            author={Tony Zhao and Vikash Kumar and Sergey Levine and Chelsea Finn},
-            journal={RSS},
-            year={2023},
-            volume={abs/2304.13705},
-            url={https://huggingface.co/papers/2304.13705}
-        }""").lstrip(),
-}
-PUSHT_INFO = {
-    "license": "mit",
-    "url": "https://diffusion-policy.cs.columbia.edu/",
-    "paper": "https://huggingface.co/papers/2303.04137",
-    "citation_bibtex": dedent(r"""
-        @article{chi2024diffusionpolicy,
-            author = {Cheng Chi and Zhenjia Xu and Siyuan Feng and Eric Cousineau and Yilun Du and Benjamin Burchfiel and Russ Tedrake and Shuran Song},
-            title ={Diffusion Policy: Visuomotor Policy Learning via Action Diffusion},
-            journal = {The International Journal of Robotics Research},
-            year = {2024},
-        }""").lstrip(),
-}
-XARM_INFO = {
-    "license": "mit",
-    "url": "https://www.nicklashansen.com/td-mpc/",
-    "paper": "https://huggingface.co/papers/2203.04955",
-    "citation_bibtex": dedent(r"""
-        @inproceedings{Hansen2022tdmpc,
-            title={Temporal Difference Learning for Model Predictive Control},
-            author={Nicklas Hansen and Xiaolong Wang and Hao Su},
-            booktitle={ICML},
-            year={2022}
-        }
-    """),
-}
-UNITREEH_INFO = {
-    "license": "apache-2.0",
-}
-
-DATASETS = {
-    "aloha_mobile_cabinet": {
-        "single_task": "Open the top cabinet, store the pot inside it then close the cabinet.",
-        **ALOHA_MOBILE_INFO,
-    },
-    "aloha_mobile_chair": {
-        "single_task": "Push the chairs in front of the desk to place them against it.",
-        **ALOHA_MOBILE_INFO,
-    },
-    "aloha_mobile_elevator": {
-        "single_task": "Take the elevator to the 1st floor.",
-        **ALOHA_MOBILE_INFO,
-    },
-    "aloha_mobile_shrimp": {
-        "single_task": "Sauté the raw shrimp on both sides, then serve it in the bowl.",
-        **ALOHA_MOBILE_INFO,
-    },
-    "aloha_mobile_wash_pan": {
-        "single_task": "Pick up the pan, rinse it in the sink and then place it in the drying rack.",
-        **ALOHA_MOBILE_INFO,
-    },
-    "aloha_mobile_wipe_wine": {
-        "single_task": "Pick up the wet cloth on the faucet and use it to clean the spilled wine on the table and underneath the glass.",
-        **ALOHA_MOBILE_INFO,
-    },
-    "aloha_static_battery": {
-        "single_task": "Place the battery into the slot of the remote controller.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_candy": {"single_task": "Pick up the candy and unwrap it.", **ALOHA_STATIC_INFO},
-    "aloha_static_coffee": {
-        "single_task": "Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray, then push the 'Hot Water' and 'Travel Mug' buttons.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_coffee_new": {
-        "single_task": "Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_cups_open": {
-        "single_task": "Pick up the plastic cup and open its lid.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_fork_pick_up": {
-        "single_task": "Pick up the fork and place it on the plate.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_pingpong_test": {
-        "single_task": "Transfer one of the two balls in the right glass into the left glass, then transfer it back to the right glass.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_pro_pencil": {
-        "single_task": "Pick up the pencil with the right arm, hand it over to the left arm then place it back onto the table.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_screw_driver": {
-        "single_task": "Pick up the screwdriver with the right arm, hand it over to the left arm then place it into the cup.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_tape": {
-        "single_task": "Cut a small piece of tape from the tape dispenser then place it on the cardboard box's edge.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_thread_velcro": {
-        "single_task": "Pick up the velcro cable tie with the left arm, then insert the end of the velcro tie into the other end's loop with the right arm.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_towel": {
-        "single_task": "Pick up a piece of paper towel and place it on the spilled liquid.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_vinh_cup": {
-        "single_task": "Pick up the plastic cup with the right arm, then pop its lid open with the left arm.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_vinh_cup_left": {
-        "single_task": "Pick up the plastic cup with the left arm, then pop its lid open with the right arm.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_static_ziploc_slide": {"single_task": "Slide open the ziploc bag.", **ALOHA_STATIC_INFO},
-    "aloha_sim_insertion_scripted": {"single_task": "Insert the peg into the socket.", **ALOHA_STATIC_INFO},
-    "aloha_sim_insertion_scripted_image": {
-        "single_task": "Insert the peg into the socket.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_sim_insertion_human": {"single_task": "Insert the peg into the socket.", **ALOHA_STATIC_INFO},
-    "aloha_sim_insertion_human_image": {
-        "single_task": "Insert the peg into the socket.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_sim_transfer_cube_scripted": {
-        "single_task": "Pick up the cube with the right arm and transfer it to the left arm.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_sim_transfer_cube_scripted_image": {
-        "single_task": "Pick up the cube with the right arm and transfer it to the left arm.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_sim_transfer_cube_human": {
-        "single_task": "Pick up the cube with the right arm and transfer it to the left arm.",
-        **ALOHA_STATIC_INFO,
-    },
-    "aloha_sim_transfer_cube_human_image": {
-        "single_task": "Pick up the cube with the right arm and transfer it to the left arm.",
-        **ALOHA_STATIC_INFO,
-    },
-    "pusht": {"single_task": "Push the T-shaped block onto the T-shaped target.", **PUSHT_INFO},
-    "pusht_image": {"single_task": "Push the T-shaped block onto the T-shaped target.", **PUSHT_INFO},
-    "unitreeh1_fold_clothes": {"single_task": "Fold the sweatshirt.", **UNITREEH_INFO},
-    "unitreeh1_rearrange_objects": {"single_task": "Put the object into the bin.", **UNITREEH_INFO},
-    "unitreeh1_two_robot_greeting": {
-        "single_task": "Greet the other robot with a high five.",
-        **UNITREEH_INFO,
-    },
-    "unitreeh1_warehouse": {
-        "single_task": "Grab the spray paint on the shelf and place it in the bin on top of the robot dog.",
-        **UNITREEH_INFO,
-    },
-    "xarm_lift_medium": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
-    "xarm_lift_medium_image": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
-    "xarm_lift_medium_replay": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
-    "xarm_lift_medium_replay_image": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
-    "xarm_push_medium": {"single_task": "Push the cube onto the target.", **XARM_INFO},
-    "xarm_push_medium_image": {"single_task": "Push the cube onto the target.", **XARM_INFO},
-    "xarm_push_medium_replay": {"single_task": "Push the cube onto the target.", **XARM_INFO},
-    "xarm_push_medium_replay_image": {"single_task": "Push the cube onto the target.", **XARM_INFO},
-    "umi_cup_in_the_wild": {
-        "single_task": "Put the cup on the plate.",
-        "license": "apache-2.0",
-    },
-    "asu_table_top": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "paper": "https://link.springer.com/article/10.1007/s10514-023-10129-1",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{zhou2023modularity,
-                title={Modularity through Attention: Efficient Training and Transfer of Language-Conditioned Policies for Robot Manipulation},
-                author={Zhou, Yifan and Sonawani, Shubham and Phielipp, Mariano and Stepputtis, Simon and Amor, Heni},
-                booktitle={Conference on Robot Learning},
-                pages={1684--1695},
-                year={2023},
-                organization={PMLR}
-            }
-            @article{zhou2023learning,
-                title={Learning modular language-conditioned robot policies through attention},
-                author={Zhou, Yifan and Sonawani, Shubham and Phielipp, Mariano and Ben Amor, Heni and Stepputtis, Simon},
-                journal={Autonomous Robots},
-                pages={1--21},
-                year={2023},
-                publisher={Springer}
-            }""").lstrip(),
-    },
-    "austin_buds_dataset": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://ut-austin-rpl.github.io/BUDS-website/",
-        "paper": "https://huggingface.co/papers/2109.13841",
-        "citation_bibtex": dedent(r"""
-            @article{zhu2022bottom,
-                title={Bottom-Up Skill Discovery From Unsegmented Demonstrations for Long-Horizon Robot Manipulation},
-                author={Zhu, Yifeng and Stone, Peter and Zhu, Yuke},
-                journal={IEEE Robotics and Automation Letters},
-                volume={7},
-                number={2},
-                pages={4126--4133},
-                year={2022},
-                publisher={IEEE}
-            }""").lstrip(),
-    },
-    "austin_sailor_dataset": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://ut-austin-rpl.github.io/sailor/",
-        "paper": "https://huggingface.co/papers/2210.11435",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{nasiriany2022sailor,
-                title={Learning and Retrieval from Prior Data for Skill-based Imitation Learning},
-                author={Soroush Nasiriany and Tian Gao and Ajay Mandlekar and Yuke Zhu},
-                booktitle={Conference on Robot Learning (CoRL)},
-                year={2022}
-            }""").lstrip(),
-    },
-    "austin_sirius_dataset": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://ut-austin-rpl.github.io/sirius/",
-        "paper": "https://huggingface.co/papers/2211.08416",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{liu2022robot,
-                title = {Robot Learning on the Job: Human-in-the-Loop Autonomy and Learning During Deployment},
-                author = {Huihan Liu and Soroush Nasiriany and Lance Zhang and Zhiyao Bao and Yuke Zhu},
-                booktitle = {Robotics: Science and Systems (RSS)},
-                year = {2023}
-            }""").lstrip(),
-    },
-    "berkeley_autolab_ur5": {
-        "tasks_col": "language_instruction",
-        "license": "cc-by-4.0",
-        "url": "https://sites.google.com/view/berkeley-ur5/home",
-        "citation_bibtex": dedent(r"""
-            @misc{BerkeleyUR5Website,
-                title = {Berkeley {UR5} Demonstration Dataset},
-                author = {Lawrence Yunliang Chen and Simeon Adebola and Ken Goldberg},
-                howpublished = {https://sites.google.com/view/berkeley-ur5/home},
-            }""").lstrip(),
-    },
-    "berkeley_cable_routing": {
-        "tasks_col": "language_instruction",
-        "license": "cc-by-4.0",
-        "url": "https://sites.google.com/view/cablerouting/home",
-        "paper": "https://huggingface.co/papers/2307.08927",
-        "citation_bibtex": dedent(r"""
-            @article{luo2023multistage,
-                author    = {Jianlan Luo and Charles Xu and Xinyang Geng and Gilbert Feng and Kuan Fang and Liam Tan and Stefan Schaal and Sergey Levine},
-                title     = {Multi-Stage Cable Routing through Hierarchical Imitation Learning},
-                journal   = {arXiv pre-print},
-                year      = {2023},
-                url       = {https://huggingface.co/papers/2307.08927},
-            }""").lstrip(),
-    },
-    "berkeley_fanuc_manipulation": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://sites.google.com/berkeley.edu/fanuc-manipulation",
-        "citation_bibtex": dedent(r"""
-            @article{fanuc_manipulation2023,
-                title={Fanuc Manipulation: A Dataset for Learning-based Manipulation with FANUC Mate 200iD Robot},
-                author={Zhu, Xinghao and Tian, Ran and Xu, Chenfeng and Ding, Mingyu and Zhan, Wei and Tomizuka, Masayoshi},
-                year={2023},
-            }""").lstrip(),
-    },
-    "berkeley_gnm_cory_hall": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "paper": "https://huggingface.co/papers/1709.10489",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{kahn2018self,
-                title={Self-supervised deep reinforcement learning with generalized computation graphs for robot navigation},
-                author={Kahn, Gregory and Villaflor, Adam and Ding, Bosen and Abbeel, Pieter and Levine, Sergey},
-                booktitle={2018 IEEE international conference on robotics and automation (ICRA)},
-                pages={5129--5136},
-                year={2018},
-                organization={IEEE}
-            }""").lstrip(),
-    },
-    "berkeley_gnm_recon": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://sites.google.com/view/recon-robot",
-        "paper": "https://huggingface.co/papers/2104.05859",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{shah2021rapid,
-                title={Rapid Exploration for Open-World Navigation with Latent Goal Models},
-                author={Dhruv Shah and Benjamin Eysenbach and Nicholas Rhinehart and Sergey Levine},
-                booktitle={5th Annual Conference on Robot Learning },
-                year={2021},
-                url={https://openreview.net/forum?id=d_SWJhyKfVw}
-            }""").lstrip(),
-    },
-    "berkeley_gnm_sac_son": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://sites.google.com/view/SACSoN-review",
-        "paper": "https://huggingface.co/papers/2306.01874",
-        "citation_bibtex": dedent(r"""
-            @article{hirose2023sacson,
-                title={SACSoN: Scalable Autonomous Data Collection for Social Navigation},
-                author={Hirose, Noriaki and Shah, Dhruv and Sridhar, Ajay and Levine, Sergey},
-                journal={arXiv preprint arXiv:2306.01874},
-                year={2023}
-            }""").lstrip(),
-    },
-    "berkeley_mvp": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "paper": "https://huggingface.co/papers/2203.06173",
-        "citation_bibtex": dedent(r"""
-            @InProceedings{Radosavovic2022,
-                title = {Real-World Robot Learning with Masked Visual Pre-training},
-                author = {Ilija Radosavovic and Tete Xiao and Stephen James and Pieter Abbeel and Jitendra Malik and Trevor Darrell},
-                booktitle = {CoRL},
-                year = {2022}
-            }""").lstrip(),
-    },
-    "berkeley_rpt": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "paper": "https://huggingface.co/papers/2306.10007",
-        "citation_bibtex": dedent(r"""
-            @article{Radosavovic2023,
-                title={Robot Learning with Sensorimotor Pre-training},
-                author={Ilija Radosavovic and Baifeng Shi and Letian Fu and Ken Goldberg and Trevor Darrell and Jitendra Malik},
-                year={2023},
-                journal={arXiv:2306.10007}
-            }""").lstrip(),
-    },
-    "cmu_franka_exploration_dataset": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://human-world-model.github.io/",
-        "paper": "https://huggingface.co/papers/2308.10901",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{mendonca2023structured,
-                title={Structured World Models from Human Videos},
-                author={Mendonca, Russell  and Bahl, Shikhar and Pathak, Deepak},
-                journal={RSS},
-                year={2023}
-            }""").lstrip(),
-    },
-    "cmu_play_fusion": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://play-fusion.github.io/",
-        "paper": "https://huggingface.co/papers/2312.04549",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{chen2023playfusion,
-                title={PlayFusion: Skill Acquisition via Diffusion from Language-Annotated Play},
-                author={Chen, Lili and Bahl, Shikhar and Pathak, Deepak},
-                booktitle={CoRL},
-                year={2023}
-            }""").lstrip(),
-    },
-    "cmu_stretch": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://robo-affordances.github.io/",
-        "paper": "https://huggingface.co/papers/2304.08488",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{bahl2023affordances,
-                title={Affordances from Human Videos as a Versatile Representation for Robotics},
-                author={Bahl, Shikhar and Mendonca, Russell and Chen, Lili and Jain, Unnat and Pathak, Deepak},
-                booktitle={CVPR},
-                year={2023}
-            }
-                @article{mendonca2023structured,
-                title={Structured World Models from Human Videos},
-                author={Mendonca, Russell and Bahl, Shikhar and Pathak, Deepak},
-                journal={CoRL},
-                year={2023}
-            }""").lstrip(),
-    },
-    "columbia_cairlab_pusht_real": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://diffusion-policy.cs.columbia.edu/",
-        "paper": "https://huggingface.co/papers/2303.04137",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{chi2023diffusionpolicy,
-                title={Diffusion Policy: Visuomotor Policy Learning via Action Diffusion},
-                author={Chi, Cheng and Feng, Siyuan and Du, Yilun and Xu, Zhenjia and Cousineau, Eric and Burchfiel, Benjamin and Song, Shuran},
-                booktitle={Proceedings of Robotics: Science and Systems (RSS)},
-                year={2023}
-            }""").lstrip(),
-    },
-    "conq_hose_manipulation": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://sites.google.com/view/conq-hose-manipulation-dataset/home",
-        "citation_bibtex": dedent(r"""
-            @misc{ConqHoseManipData,
-                author={Peter Mitrano and Dmitry Berenson},
-                title={Conq Hose Manipulation Dataset, v1.15.0},
-                year={2024},
-                howpublished={https://sites.google.com/view/conq-hose-manipulation-dataset}
-            }""").lstrip(),
-    },
-    "dlr_edan_shared_control": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "paper": "https://ieeexplore.ieee.org/document/9341156",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{vogel_edan_2020,
-                title = {EDAN - an EMG-Controlled Daily Assistant to Help People with Physical Disabilities},
-                language = {en},
-                booktitle = {2020 {IEEE}/{RSJ} {International} {Conference} on {Intelligent} {Robots} and {Systems} ({IROS})},
-                author = {Vogel, Jörn and Hagengruber, Annette and Iskandar, Maged and Quere, Gabriel and Leipscher, Ulrike and Bustamante, Samuel and Dietrich, Alexander and Hoeppner, Hannes and Leidner, Daniel and Albu-Schäffer, Alin},
-                year = {2020}
-            }
-            @inproceedings{quere_shared_2020,
-                address = {Paris, France},
-                title = {Shared {Control} {Templates} for {Assistive} {Robotics}},
-                language = {en},
-                booktitle = {2020 {IEEE} {International} {Conference} on {Robotics} and {Automation} ({ICRA})},
-                author = {Quere, Gabriel and Hagengruber, Annette and Iskandar, Maged and Bustamante, Samuel and Leidner, Daniel and Stulp, Freek and Vogel, Joern},
-                year = {2020},
-                pages = {7},
-            }""").lstrip(),
-    },
-    "dlr_sara_grid_clamp": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "paper": "https://www.researchsquare.com/article/rs-3289569/v1",
-        "citation_bibtex": dedent(r"""
-            @article{padalkar2023guided,
-                title={A guided reinforcement learning approach using shared control templates for learning manipulation skills in the real world},
-                author={Padalkar, Abhishek and Quere, Gabriel and Raffin, Antonin and Silv{\'e}rio, Jo{\~a}o and Stulp, Freek},
-                journal={Research square preprint rs-3289569/v1},
-                year={2023}
-            }""").lstrip(),
-    },
-    "dlr_sara_pour": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "paper": "https://elib.dlr.de/193739/1/padalkar2023rlsct.pdf",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{padalkar2023guiding,
-                title={Guiding Reinforcement Learning with Shared Control Templates},
-                author={Padalkar, Abhishek and Quere, Gabriel and Steinmetz, Franz and Raffin, Antonin and Nieuwenhuisen, Matthias and Silv{\'e}rio, Jo{\~a}o and Stulp, Freek},
-                booktitle={40th IEEE International Conference on Robotics and Automation, ICRA 2023},
-                year={2023},
-                organization={IEEE}
-            }""").lstrip(),
-    },
-    "droid_100": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://droid-dataset.github.io/",
-        "paper": "https://huggingface.co/papers/2403.12945",
-        "citation_bibtex": dedent(r"""
-            @article{khazatsky2024droid,
-                title   = {DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset},
-                author  = {Alexander Khazatsky and Karl Pertsch and Suraj Nair and Ashwin Balakrishna and Sudeep Dasari and Siddharth Karamcheti and Soroush Nasiriany and Mohan Kumar Srirama and Lawrence Yunliang Chen and Kirsty Ellis and Peter David Fagan and Joey Hejna and Masha Itkina and Marion Lepert and Yecheng Jason Ma and Patrick Tree Miller and Jimmy Wu and Suneel Belkhale and Shivin Dass and Huy Ha and Arhan Jain and Abraham Lee and Youngwoon Lee and Marius Memmel and Sungjae Park and Ilija Radosavovic and Kaiyuan Wang and Albert Zhan and Kevin Black and Cheng Chi and Kyle Beltran Hatch and Shan Lin and Jingpei Lu and Jean Mercat and Abdul Rehman and Pannag R Sanketi and Archit Sharma and Cody Simpson and Quan Vuong and Homer Rich Walke and Blake Wulfe and Ted Xiao and Jonathan Heewon Yang and Arefeh Yavary and Tony Z. Zhao and Christopher Agia and Rohan Baijal and Mateo Guaman Castro and Daphne Chen and Qiuyu Chen and Trinity Chung and Jaimyn Drake and Ethan Paul Foster and Jensen Gao and David Antonio Herrera and Minho Heo and Kyle Hsu and Jiaheng Hu and Donovon Jackson and Charlotte Le and Yunshuang Li and Kevin Lin and Roy Lin and Zehan Ma and Abhiram Maddukuri and Suvir Mirchandani and Daniel Morton and Tony Nguyen and Abigail O'Neill and Rosario Scalise and Derick Seale and Victor Son and Stephen Tian and Emi Tran and Andrew E. Wang and Yilin Wu and Annie Xie and Jingyun Yang and Patrick Yin and Yunchu Zhang and Osbert Bastani and Glen Berseth and Jeannette Bohg and Ken Goldberg and Abhinav Gupta and Abhishek Gupta and Dinesh Jayaraman and Joseph J Lim and Jitendra Malik and Roberto Martín-Martín and Subramanian Ramamoorthy and Dorsa Sadigh and Shuran Song and Jiajun Wu and Michael C. Yip and Yuke Zhu and Thomas Kollar and Sergey Levine and Chelsea Finn},
-                year    = {2024},
-            }""").lstrip(),
-    },
-    "fmb": {
-        "tasks_col": "language_instruction",
-        "license": "cc-by-4.0",
-        "url": "https://functional-manipulation-benchmark.github.io/",
-        "paper": "https://huggingface.co/papers/2401.08553",
-        "citation_bibtex": dedent(r"""
-            @article{luo2024fmb,
-                title={FMB: a Functional Manipulation Benchmark for Generalizable Robotic Learning},
-                author={Luo, Jianlan and Xu, Charles and Liu, Fangchen and Tan, Liam and Lin, Zipeng and Wu, Jeffrey and Abbeel, Pieter and Levine, Sergey},
-                journal={arXiv preprint arXiv:2401.08553},
-                year={2024}
-            }""").lstrip(),
-    },
-    "iamlab_cmu_pickup_insert": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://openreview.net/forum?id=WuBv9-IGDUA",
-        "paper": "https://huggingface.co/papers/2401.14502",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{saxena2023multiresolution,
-                title={Multi-Resolution Sensing for Real-Time Control with Vision-Language Models},
-                author={Saumya Saxena and Mohit Sharma and Oliver Kroemer},
-                booktitle={7th Annual Conference on Robot Learning},
-                year={2023},
-                url={https://openreview.net/forum?id=WuBv9-IGDUA}
-            }""").lstrip(),
-    },
-    "imperialcollege_sawyer_wrist_cam": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-    },
-    "jaco_play": {
-        "tasks_col": "language_instruction",
-        "license": "cc-by-4.0",
-        "url": "https://github.com/clvrai/clvr_jaco_play_dataset",
-        "citation_bibtex": dedent(r"""
-            @software{dass2023jacoplay,
-                author = {Dass, Shivin and Yapeter, Jullian and Zhang, Jesse and Zhang, Jiahui
-                            and Pertsch, Karl and Nikolaidis, Stefanos and Lim, Joseph J.},
-                title = {CLVR Jaco Play Dataset},
-                url = {https://github.com/clvrai/clvr_jaco_play_dataset},
-                version = {1.0.0},
-                year = {2023}
-            }""").lstrip(),
-    },
-    "kaist_nonprehensile": {
-        "tasks_col": "language_instruction",
-        "license": "cc-by-4.0",
-        "url": "https://github.com/JaeHyung-Kim/rlds_dataset_builder",
-        "citation_bibtex": dedent(r"""
-            @article{kimpre,
-                title={Pre-and post-contact policy decomposition for non-prehensile manipulation with zero-shot sim-to-real transfer},
-                author={Kim, Minchan and Han, Junhyek and Kim, Jaehyung and Kim, Beomjoon},
-                booktitle={2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)},
-                year={2023},
-                organization={IEEE}
-            }""").lstrip(),
-    },
-    "nyu_door_opening_surprising_effectiveness": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://jyopari.github.io/VINN/",
-        "paper": "https://huggingface.co/papers/2112.01511",
-        "citation_bibtex": dedent(r"""
-            @misc{pari2021surprising,
-                title={The Surprising Effectiveness of Representation Learning for Visual Imitation},
-                author={Jyothish Pari and Nur Muhammad Shafiullah and Sridhar Pandian Arunachalam and Lerrel Pinto},
-                year={2021},
-                eprint={2112.01511},
-                archivePrefix={arXiv},
-                primaryClass={cs.RO}
-            }""").lstrip(),
-    },
-    "nyu_franka_play_dataset": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://play-to-policy.github.io/",
-        "paper": "https://huggingface.co/papers/2210.10047",
-        "citation_bibtex": dedent(r"""
-            @article{cui2022play,
-                title   = {From Play to Policy: Conditional Behavior Generation from Uncurated Robot Data},
-                author  = {Cui, Zichen Jeff and Wang, Yibin and Shafiullah, Nur Muhammad Mahi and Pinto, Lerrel},
-                journal = {arXiv preprint arXiv:2210.10047},
-                year    = {2022}
-            }""").lstrip(),
-    },
-    "nyu_rot_dataset": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://rot-robot.github.io/",
-        "paper": "https://huggingface.co/papers/2206.15469",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{haldar2023watch,
-                title={Watch and match: Supercharging imitation with regularized optimal transport},
-                author={Haldar, Siddhant and Mathur, Vaibhav and Yarats, Denis and Pinto, Lerrel},
-                booktitle={Conference on Robot Learning},
-                pages={32--43},
-                year={2023},
-                organization={PMLR}
-            }""").lstrip(),
-    },
-    "roboturk": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://roboturk.stanford.edu/dataset_real.html",
-        "paper": "PAPER",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{mandlekar2019scaling,
-                title={Scaling robot supervision to hundreds of hours with roboturk: Robotic manipulation dataset through human reasoning and dexterity},
-                author={Mandlekar, Ajay and Booher, Jonathan and Spero, Max and Tung, Albert and Gupta, Anchit and Zhu, Yuke and Garg, Animesh and Savarese, Silvio and Fei-Fei, Li},
-                booktitle={2019 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)},
-                pages={1048--1055},
-                year={2019},
-                organization={IEEE}
-            }""").lstrip(),
-    },
-    "stanford_hydra_dataset": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://sites.google.com/view/hydra-il-2023",
-        "paper": "https://huggingface.co/papers/2306.17237",
-        "citation_bibtex": dedent(r"""
-            @article{belkhale2023hydra,
-                title={HYDRA: Hybrid Robot Actions for Imitation Learning},
-                author={Belkhale, Suneel and Cui, Yuchen and Sadigh, Dorsa},
-                journal={arxiv},
-                year={2023}
-            }""").lstrip(),
-    },
-    "stanford_kuka_multimodal_dataset": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://sites.google.com/view/visionandtouch",
-        "paper": "https://huggingface.co/papers/1810.10191",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{lee2019icra,
-                title={Making sense of vision and touch: Self-supervised learning of multimodal representations for contact-rich tasks},
-                author={Lee, Michelle A and Zhu, Yuke and Srinivasan, Krishnan and Shah, Parth and Savarese, Silvio and Fei-Fei, Li and  Garg, Animesh and Bohg, Jeannette},
-                booktitle={2019 IEEE International Conference on Robotics and Automation (ICRA)},
-                year={2019},
-                url={https://huggingface.co/papers/1810.10191}
-            }""").lstrip(),
-    },
-    "stanford_robocook": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://hshi74.github.io/robocook/",
-        "paper": "https://huggingface.co/papers/2306.14447",
-        "citation_bibtex": dedent(r"""
-            @article{shi2023robocook,
-                title={RoboCook: Long-Horizon Elasto-Plastic Object Manipulation with Diverse Tools},
-                author={Shi, Haochen and Xu, Huazhe and Clarke, Samuel and Li, Yunzhu and Wu, Jiajun},
-                journal={arXiv preprint arXiv:2306.14447},
-                year={2023}
-            }""").lstrip(),
-    },
-    "taco_play": {
-        "tasks_col": "language_instruction",
-        "license": "cc-by-4.0",
-        "url": "https://www.kaggle.com/datasets/oiermees/taco-robot",
-        "paper": "https://huggingface.co/papers/2209.08959, https://huggingface.co/papers/2210.01911",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{rosete2022tacorl,
-                author = {Erick Rosete-Beas and Oier Mees and Gabriel Kalweit and Joschka Boedecker and Wolfram Burgard},
-                title = {Latent Plans for Task Agnostic Offline Reinforcement Learning},
-                journal = {Proceedings of the 6th Conference on Robot Learning (CoRL)},
-                year = {2022}
-            }
-            @inproceedings{mees23hulc2,
-                title={Grounding  Language  with  Visual  Affordances  over  Unstructured  Data},
-                author={Oier Mees and Jessica Borja-Diaz and Wolfram Burgard},
-                booktitle = {Proceedings of the IEEE International Conference on Robotics and Automation (ICRA)},
-                year={2023},
-                address = {London, UK}
-            }""").lstrip(),
-    },
-    "tokyo_u_lsmo": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "URL",
-        "paper": "https://huggingface.co/papers/2107.05842",
-        "citation_bibtex": dedent(r"""
-            @Article{Osa22,
-                author  = {Takayuki Osa},
-                journal = {The International Journal of Robotics Research},
-                title   = {Motion Planning by Learning the Solution Manifold in Trajectory Optimization},
-                year    = {2022},
-                number  = {3},
-                pages   = {291--311},
-                volume  = {41},
-            }""").lstrip(),
-    },
-    "toto": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://toto-benchmark.org/",
-        "paper": "https://huggingface.co/papers/2306.00942",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{zhou2023train,
-                author={Zhou, Gaoyue and Dean, Victoria and Srirama, Mohan Kumar and Rajeswaran, Aravind and Pari, Jyothish and Hatch, Kyle and Jain, Aryan and Yu, Tianhe and Abbeel, Pieter and Pinto, Lerrel and Finn, Chelsea and Gupta, Abhinav},
-                booktitle={2023 IEEE International Conference on Robotics and Automation (ICRA)},
-                title={Train Offline, Test Online: A Real Robot Learning Benchmark},
-                year={2023},
-            }""").lstrip(),
-    },
-    "ucsd_kitchen_dataset": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "citation_bibtex": dedent(r"""
-            @ARTICLE{ucsd_kitchens,
-                author = {Ge Yan, Kris Wu, and Xiaolong Wang},
-                title = {{ucsd kitchens Dataset}},
-                year = {2023},
-                month = {August}
-            }""").lstrip(),
-    },
-    "ucsd_pick_and_place_dataset": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://owmcorl.github.io/#",
-        "paper": "https://huggingface.co/papers/2310.16029",
-        "citation_bibtex": dedent(r"""
-            @preprint{Feng2023Finetuning,
-                title={Finetuning Offline World Models in the Real World},
-                author={Yunhai Feng, Nicklas Hansen, Ziyan Xiong, Chandramouli Rajagopalan, Xiaolong Wang},
-                year={2023}
-            }""").lstrip(),
-    },
-    "uiuc_d3field": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://robopil.github.io/d3fields/",
-        "paper": "https://huggingface.co/papers/2309.16118",
-        "citation_bibtex": dedent(r"""
-            @article{wang2023d3field,
-                title={D^3Field: Dynamic 3D Descriptor Fields for Generalizable Robotic Manipulation},
-                author={Wang, Yixuan and Li, Zhuoran and Zhang, Mingtong and Driggs-Campbell, Katherine and Wu, Jiajun and Fei-Fei, Li and Li, Yunzhu},
-                journal={arXiv preprint arXiv:},
-                year={2023},
-            }""").lstrip(),
-    },
-    "usc_cloth_sim": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://uscresl.github.io/dmfd/",
-        "paper": "https://huggingface.co/papers/2207.10148",
-        "citation_bibtex": dedent(r"""
-            @article{salhotra2022dmfd,
-                author={Salhotra, Gautam and Liu, I-Chun Arthur and Dominguez-Kuhne, Marcus and Sukhatme, Gaurav S.},
-                journal={IEEE Robotics and Automation Letters},
-                title={Learning Deformable Object Manipulation From Expert Demonstrations},
-                year={2022},
-                volume={7},
-                number={4},
-                pages={8775-8782},
-                doi={10.1109/LRA.2022.3187843}
-            }""").lstrip(),
-    },
-    "utaustin_mutex": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://ut-austin-rpl.github.io/MUTEX/",
-        "paper": "https://huggingface.co/papers/2309.14320",
-        "citation_bibtex": dedent(r"""
-            @inproceedings{shah2023mutex,
-                title={{MUTEX}: Learning Unified Policies from Multimodal Task Specifications},
-                author={Rutav Shah and Roberto Mart{\'\i}n-Mart{\'\i}n and Yuke Zhu},
-                booktitle={7th Annual Conference on Robot Learning},
-                year={2023},
-                url={https://openreview.net/forum?id=PwqiqaaEzJ}
-            }""").lstrip(),
-    },
-    "utokyo_pr2_opening_fridge": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "citation_bibtex": dedent(r"""
-            @misc{oh2023pr2utokyodatasets,
-                author={Jihoon Oh and Naoaki Kanazawa and Kento Kawaharazuka},
-                title={X-Embodiment U-Tokyo PR2 Datasets},
-                year={2023},
-                url={https://github.com/ojh6404/rlds_dataset_builder},
-            }""").lstrip(),
-    },
-    "utokyo_pr2_tabletop_manipulation": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "citation_bibtex": dedent(r"""
-            @misc{oh2023pr2utokyodatasets,
-                author={Jihoon Oh and Naoaki Kanazawa and Kento Kawaharazuka},
-                title={X-Embodiment U-Tokyo PR2 Datasets},
-                year={2023},
-                url={https://github.com/ojh6404/rlds_dataset_builder},
-            }""").lstrip(),
-    },
-    "utokyo_saytap": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://saytap.github.io/",
-        "paper": "https://huggingface.co/papers/2306.07580",
-        "citation_bibtex": dedent(r"""
-            @article{saytap2023,
-                author = {Yujin Tang and Wenhao Yu and Jie Tan and Heiga Zen and Aleksandra Faust and
-                Tatsuya Harada},
-                title  = {SayTap: Language to Quadrupedal Locomotion},
-                eprint = {arXiv:2306.07580},
-                url    = {https://saytap.github.io},
-                note   = {https://saytap.github.io},
-                year   = {2023}
-            }""").lstrip(),
-    },
-    "utokyo_xarm_bimanual": {
-        "tasks_col": "language_instruction",
-        "license": "cc-by-4.0",
-        "citation_bibtex": dedent(r"""
-            @misc{matsushima2023weblab,
-                title={Weblab xArm Dataset},
-                author={Tatsuya Matsushima and Hiroki Furuta and Yusuke Iwasawa and Yutaka Matsuo},
-                year={2023},
-            }""").lstrip(),
-    },
-    "utokyo_xarm_pick_and_place": {
-        "tasks_col": "language_instruction",
-        "license": "cc-by-4.0",
-        "citation_bibtex": dedent(r"""
-            @misc{matsushima2023weblab,
-                title={Weblab xArm Dataset},
-                author={Tatsuya Matsushima and Hiroki Furuta and Yusuke Iwasawa and Yutaka Matsuo},
-                year={2023},
-            }""").lstrip(),
-    },
-    "viola": {
-        "tasks_col": "language_instruction",
-        "license": "mit",
-        "url": "https://ut-austin-rpl.github.io/VIOLA/",
-        "paper": "https://huggingface.co/papers/2210.11339",
-        "citation_bibtex": dedent(r"""
-            @article{zhu2022viola,
-                title={VIOLA: Imitation Learning for Vision-Based Manipulation with Object Proposal Priors},
-                author={Zhu, Yifeng and Joshi, Abhishek and Stone, Peter and Zhu, Yuke},
-                journal={6th Annual Conference on Robot Learning (CoRL)},
-                year={2022}
-            }""").lstrip(),
-    },
-}
-# spellchecker:on
-
-
-def batch_convert():
-    status = {}
-    logfile = LOCAL_DIR / "conversion_log.txt"
-    assert set(DATASETS) == {id_.split("/")[1] for id_ in available_datasets}
-    for num, (name, kwargs) in enumerate(DATASETS.items()):
-        repo_id = f"lerobot/{name}"
-        print(f"\nConverting {repo_id} ({num}/{len(DATASETS)})")
-        print("---------------------------------------------------------")
-        try:
-            convert_dataset(repo_id, LOCAL_DIR, **kwargs)
-            status = f"{repo_id}: success."
-            with open(logfile, "a") as file:
-                file.write(status + "\n")
-        except Exception:
-            status = f"{repo_id}: failed\n    {traceback.format_exc()}"
-            with open(logfile, "a") as file:
-                file.write(status + "\n")
-            continue
-
-
-if __name__ == "__main__":
-    batch_convert()
@@ -1,687 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 1.6 to
-2.0. You will be required to provide the 'tasks', which is a short but accurate description in plain English
-for each of the task performed in the dataset. This will allow to easily train models with task-conditioning.
-
-We support 3 different scenarios for these tasks (see instructions below):
-    1. Single task dataset: all episodes of your dataset have the same single task.
-    2. Single task episodes: the episodes of your dataset each contain a single task but they can differ from
-      one episode to the next.
-    3. Multi task episodes: episodes of your dataset may each contain several different tasks.
-
-
-Can you can also provide a robot config .yaml file (not mandatory) to this script via the option
-'--robot-config' so that it writes information about the robot (robot type, motors names) this dataset was
-recorded with. For now, only Aloha/Koch type robots are supported with this option.
-
-
-# 1. Single task dataset
-If your dataset contains a single task, you can simply provide it directly via the CLI with the
-'--single-task' option.
-
-Examples:
-
-```bash
-python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
-    --repo-id lerobot/aloha_sim_insertion_human_image \
-    --single-task "Insert the peg into the socket." \
-    --robot-config lerobot/configs/robot/aloha.yaml \
-    --local-dir data
-```
-
-```bash
-python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
-    --repo-id aliberts/koch_tutorial \
-    --single-task "Pick the Lego block and drop it in the box on the right." \
-    --robot-config lerobot/configs/robot/koch.yaml \
-    --local-dir data
-```
-
-
-# 2. Single task episodes
-If your dataset is a multi-task dataset, you have two options to provide the tasks to this script:
-
- If your dataset already contains a language instruction column in its parquet file, you can simply provide
-  this column's name with the '--tasks-col' arg.
-
-    Example:
-
-    ```bash
-    python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
-        --repo-id lerobot/stanford_kuka_multimodal_dataset \
-        --tasks-col "language_instruction" \
-        --local-dir data
-    ```
-
- If your dataset doesn't contain a language instruction, you should provide the path to a .json file with the
-  '--tasks-path' arg. This file should have the following structure where keys correspond to each
-  episode_index in the dataset, and values are the language instruction for that episode.
-
-    Example:
-
-    ```json
-    {
-        "0": "Do something",
-        "1": "Do something else",
-        "2": "Do something",
-        "3": "Go there",
-        ...
-    }
-    ```
-
-# 3. Multi task episodes
-If you have multiple tasks per episodes, your dataset should contain a language instruction column in its
-parquet file, and you must provide this column's name with the '--tasks-col' arg.
-
-Example:
-
-```bash
-python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
-    --repo-id lerobot/stanford_kuka_multimodal_dataset \
-    --tasks-col "language_instruction" \
-    --local-dir data
-```
-"""
-
-import argparse
-import contextlib
-import filecmp
-import json
-import logging
-import math
-import shutil
-import subprocess
-import tempfile
-from pathlib import Path
-
-import datasets
-import pyarrow.compute as pc
-import pyarrow.parquet as pq
-import torch
-from datasets import Dataset
-from huggingface_hub import HfApi
-from huggingface_hub.errors import EntryNotFoundError, HfHubHTTPError
-from safetensors.torch import load_file
-
-from lerobot.datasets.utils import (
-    DEFAULT_CHUNK_SIZE,
-    INFO_PATH,
-    LEGACY_DEFAULT_PARQUET_PATH,
-    LEGACY_DEFAULT_VIDEO_PATH,
-    LEGACY_EPISODES_PATH,
-    LEGACY_TASKS_PATH,
-    STATS_PATH,
-    create_branch,
-    create_lerobot_dataset_card,
-    flatten_dict,
-    get_safe_version,
-    load_json,
-    unflatten_dict,
-    write_json,
-    write_jsonlines,
-)
-from lerobot.datasets.video_utils import (
-    VideoFrame,  # noqa: F401
-    get_image_pixel_channels,
-    get_video_info,
-)
-from lerobot.robots import RobotConfig
-
-V16 = "v1.6"
-V20 = "v2.0"
-
-GITATTRIBUTES_REF = "aliberts/gitattributes_reference"
-V1_VIDEO_FILE = "{video_key}_episode_{episode_index:06d}.mp4"
-V1_INFO_PATH = "meta_data/info.json"
-V1_STATS_PATH = "meta_data/stats.safetensors"
-
-
-def parse_robot_config(robot_cfg: RobotConfig) -> tuple[str, dict]:
-    if robot_cfg.type in ["aloha", "koch"]:
-        state_names = [
-            f"{arm}_{motor}" if len(robot_cfg.follower_arms) > 1 else motor
-            for arm in robot_cfg.follower_arms
-            for motor in robot_cfg.follower_arms[arm].motors
-        ]
-        action_names = [
-            # f"{arm}_{motor}" for arm in ["left", "right"] for motor in robot_cfg["leader_arms"][arm]["motors"]
-            f"{arm}_{motor}" if len(robot_cfg.leader_arms) > 1 else motor
-            for arm in robot_cfg.leader_arms
-            for motor in robot_cfg.leader_arms[arm].motors
-        ]
-    # elif robot_cfg["robot_type"] == "stretch3": TODO
-    else:
-        raise NotImplementedError(
-            "Please provide robot_config={'robot_type': ..., 'names': ...} directly to convert_dataset()."
-        )
-
-    return {
-        "robot_type": robot_cfg.type,
-        "names": {
-            "observation.state": state_names,
-            "observation.effort": state_names,
-            "action": action_names,
-        },
-    }
-
-
-def convert_stats_to_json(v1_dir: Path, v2_dir: Path) -> None:
-    safetensor_path = v1_dir / V1_STATS_PATH
-    stats = load_file(safetensor_path)
-    serialized_stats = {key: value.tolist() for key, value in stats.items()}
-    serialized_stats = unflatten_dict(serialized_stats)
-
-    json_path = v2_dir / STATS_PATH
-    json_path.parent.mkdir(exist_ok=True, parents=True)
-    with open(json_path, "w") as f:
-        json.dump(serialized_stats, f, indent=4)
-
-    # Sanity check
-    with open(json_path) as f:
-        stats_json = json.load(f)
-
-    stats_json = flatten_dict(stats_json)
-    stats_json = {key: torch.tensor(value) for key, value in stats_json.items()}
-    for key in stats:
-        torch.testing.assert_close(stats_json[key], stats[key])
-
-
-def get_features_from_hf_dataset(
-    dataset: Dataset, robot_config: RobotConfig | None = None
-) -> dict[str, list]:
-    robot_config = parse_robot_config(robot_config)
-    features = {}
-    for key, ft in dataset.features.items():
-        if isinstance(ft, datasets.Value):
-            dtype = ft.dtype
-            shape = (1,)
-            names = None
-        if isinstance(ft, datasets.Sequence):
-            assert isinstance(ft.feature, datasets.Value)
-            dtype = ft.feature.dtype
-            shape = (ft.length,)
-            motor_names = (
-                robot_config["names"][key] if robot_config else [f"motor_{i}" for i in range(ft.length)]
-            )
-            assert len(motor_names) == shape[0]
-            names = {"motors": motor_names}
-        elif isinstance(ft, datasets.Image):
-            dtype = "image"
-            image = dataset[0][key]  # Assuming first row
-            channels = get_image_pixel_channels(image)
-            shape = (image.height, image.width, channels)
-            names = ["height", "width", "channels"]
-        elif ft._type == "VideoFrame":
-            dtype = "video"
-            shape = None  # Add shape later
-            names = ["height", "width", "channels"]
-
-        features[key] = {
-            "dtype": dtype,
-            "shape": shape,
-            "names": names,
-        }
-
-    return features
-
-
-def add_task_index_by_episodes(dataset: Dataset, tasks_by_episodes: dict) -> tuple[Dataset, list[str]]:
-    df = dataset.to_pandas()
-    tasks = list(set(tasks_by_episodes.values()))
-    tasks_to_task_index = {task: task_idx for task_idx, task in enumerate(tasks)}
-    episodes_to_task_index = {ep_idx: tasks_to_task_index[task] for ep_idx, task in tasks_by_episodes.items()}
-    df["task_index"] = df["episode_index"].map(episodes_to_task_index).astype(int)
-
-    features = dataset.features
-    features["task_index"] = datasets.Value(dtype="int64")
-    dataset = Dataset.from_pandas(df, features=features, split="train")
-    return dataset, tasks
-
-
-def add_task_index_from_tasks_col(
-    dataset: Dataset, tasks_col: str
-) -> tuple[Dataset, dict[str, list[str]], list[str]]:
-    df = dataset.to_pandas()
-
-    # HACK: This is to clean some of the instructions in our version of Open X datasets
-    prefix_to_clean = "tf.Tensor(b'"
-    suffix_to_clean = "', shape=(), dtype=string)"
-    df[tasks_col] = df[tasks_col].str.removeprefix(prefix_to_clean).str.removesuffix(suffix_to_clean)
-
-    # Create task_index col
-    tasks_by_episode = df.groupby("episode_index")[tasks_col].unique().apply(lambda x: x.tolist()).to_dict()
-    tasks = df[tasks_col].unique().tolist()
-    tasks_to_task_index = {task: idx for idx, task in enumerate(tasks)}
-    df["task_index"] = df[tasks_col].map(tasks_to_task_index).astype(int)
-
-    # Build the dataset back from df
-    features = dataset.features
-    features["task_index"] = datasets.Value(dtype="int64")
-    dataset = Dataset.from_pandas(df, features=features, split="train")
-    dataset = dataset.remove_columns(tasks_col)
-
-    return dataset, tasks, tasks_by_episode
-
-
-def split_parquet_by_episodes(
-    dataset: Dataset,
-    total_episodes: int,
-    total_chunks: int,
-    output_dir: Path,
-) -> list:
-    table = dataset.data.table
-    episode_lengths = []
-    for ep_chunk in range(total_chunks):
-        ep_chunk_start = DEFAULT_CHUNK_SIZE * ep_chunk
-        ep_chunk_end = min(DEFAULT_CHUNK_SIZE * (ep_chunk + 1), total_episodes)
-        chunk_dir = "/".join(LEGACY_DEFAULT_PARQUET_PATH.split("/")[:-1]).format(episode_chunk=ep_chunk)
-        (output_dir / chunk_dir).mkdir(parents=True, exist_ok=True)
-        for ep_idx in range(ep_chunk_start, ep_chunk_end):
-            ep_table = table.filter(pc.equal(table["episode_index"], ep_idx))
-            episode_lengths.insert(ep_idx, len(ep_table))
-            output_file = output_dir / LEGACY_DEFAULT_PARQUET_PATH.format(
-                episode_chunk=ep_chunk, episode_index=ep_idx
-            )
-            pq.write_table(ep_table, output_file)
-
-    return episode_lengths
-
-
-def move_videos(
-    repo_id: str,
-    video_keys: list[str],
-    total_episodes: int,
-    total_chunks: int,
-    work_dir: Path,
-    clean_gittatributes: Path,
-    branch: str = "main",
-) -> None:
-    """
-    HACK: Since HfApi() doesn't provide a way to move files directly in a repo, this function will run git
-    commands to fetch git lfs video files references to move them into subdirectories without having to
-    actually download them.
-    """
-    _lfs_clone(repo_id, work_dir, branch)
-
-    videos_moved = False
-    video_files = [str(f.relative_to(work_dir)) for f in work_dir.glob("videos*/*.mp4")]
-    if len(video_files) == 0:
-        video_files = [str(f.relative_to(work_dir)) for f in work_dir.glob("videos*/*/*/*.mp4")]
-        videos_moved = True  # Videos have already been moved
-
-    assert len(video_files) == total_episodes * len(video_keys)
-
-    lfs_untracked_videos = _get_lfs_untracked_videos(work_dir, video_files)
-
-    current_gittatributes = work_dir / ".gitattributes"
-    if not filecmp.cmp(current_gittatributes, clean_gittatributes, shallow=False):
-        fix_gitattributes(work_dir, current_gittatributes, clean_gittatributes)
-
-    if lfs_untracked_videos:
-        fix_lfs_video_files_tracking(work_dir, video_files)
-
-    if videos_moved:
-        return
-
-    video_dirs = sorted(work_dir.glob("videos*/"))
-    for ep_chunk in range(total_chunks):
-        ep_chunk_start = DEFAULT_CHUNK_SIZE * ep_chunk
-        ep_chunk_end = min(DEFAULT_CHUNK_SIZE * (ep_chunk + 1), total_episodes)
-        for vid_key in video_keys:
-            chunk_dir = "/".join(LEGACY_DEFAULT_VIDEO_PATH.split("/")[:-1]).format(
-                episode_chunk=ep_chunk, video_key=vid_key
-            )
-            (work_dir / chunk_dir).mkdir(parents=True, exist_ok=True)
-
-            for ep_idx in range(ep_chunk_start, ep_chunk_end):
-                target_path = LEGACY_DEFAULT_VIDEO_PATH.format(
-                    episode_chunk=ep_chunk, video_key=vid_key, episode_index=ep_idx
-                )
-                video_file = V1_VIDEO_FILE.format(video_key=vid_key, episode_index=ep_idx)
-                if len(video_dirs) == 1:
-                    video_path = video_dirs[0] / video_file
-                else:
-                    for dir in video_dirs:
-                        if (dir / video_file).is_file():
-                            video_path = dir / video_file
-                            break
-
-                video_path.rename(work_dir / target_path)
-
-    commit_message = "Move video files into chunk subdirectories"
-    subprocess.run(["git", "add", "."], cwd=work_dir, check=True)
-    subprocess.run(["git", "commit", "-m", commit_message], cwd=work_dir, check=True)
-    subprocess.run(["git", "push"], cwd=work_dir, check=True)
-
-
-def fix_lfs_video_files_tracking(work_dir: Path, lfs_untracked_videos: list[str]) -> None:
-    """
-    HACK: This function fixes the tracking by git lfs which was not properly set on some repos. In that case,
-    there's no other option than to download the actual files and reupload them with lfs tracking.
-    """
-    for i in range(0, len(lfs_untracked_videos), 100):
-        files = lfs_untracked_videos[i : i + 100]
-        try:
-            subprocess.run(["git", "rm", "--cached", *files], cwd=work_dir, capture_output=True, check=True)
-        except subprocess.CalledProcessError as e:
-            print("git rm --cached ERROR:")
-            print(e.stderr)
-        subprocess.run(["git", "add", *files], cwd=work_dir, check=True)
-
-    commit_message = "Track video files with git lfs"
-    subprocess.run(["git", "commit", "-m", commit_message], cwd=work_dir, check=True)
-    subprocess.run(["git", "push"], cwd=work_dir, check=True)
-
-
-def fix_gitattributes(work_dir: Path, current_gittatributes: Path, clean_gittatributes: Path) -> None:
-    shutil.copyfile(clean_gittatributes, current_gittatributes)
-    subprocess.run(["git", "add", ".gitattributes"], cwd=work_dir, check=True)
-    subprocess.run(["git", "commit", "-m", "Fix .gitattributes"], cwd=work_dir, check=True)
-    subprocess.run(["git", "push"], cwd=work_dir, check=True)
-
-
-def _lfs_clone(repo_id: str, work_dir: Path, branch: str) -> None:
-    subprocess.run(["git", "lfs", "install"], cwd=work_dir, check=True)
-    repo_url = f"https://huggingface.co/datasets/{repo_id}"
-    env = {"GIT_LFS_SKIP_SMUDGE": "1"}  # Prevent downloading LFS files
-    subprocess.run(
-        ["git", "clone", "--branch", branch, "--single-branch", "--depth", "1", repo_url, str(work_dir)],
-        check=True,
-        env=env,
-    )
-
-
-def _get_lfs_untracked_videos(work_dir: Path, video_files: list[str]) -> list[str]:
-    lfs_tracked_files = subprocess.run(
-        ["git", "lfs", "ls-files", "-n"], cwd=work_dir, capture_output=True, text=True, check=True
-    )
-    lfs_tracked_files = set(lfs_tracked_files.stdout.splitlines())
-    return [f for f in video_files if f not in lfs_tracked_files]
-
-
-def get_videos_info(repo_id: str, local_dir: Path, video_keys: list[str], branch: str) -> dict:
-    # Assumes first episode
-    video_files = [
-        LEGACY_DEFAULT_VIDEO_PATH.format(episode_chunk=0, video_key=vid_key, episode_index=0)
-        for vid_key in video_keys
-    ]
-    hub_api = HfApi()
-    hub_api.snapshot_download(
-        repo_id=repo_id, repo_type="dataset", local_dir=local_dir, revision=branch, allow_patterns=video_files
-    )
-    videos_info_dict = {}
-    for vid_key, vid_path in zip(video_keys, video_files, strict=True):
-        videos_info_dict[vid_key] = get_video_info(local_dir / vid_path)
-
-    return videos_info_dict
-
-
-def convert_dataset(
-    repo_id: str,
-    local_dir: Path,
-    single_task: str | None = None,
-    tasks_path: Path | None = None,
-    tasks_col: Path | None = None,
-    robot_config: RobotConfig | None = None,
-    test_branch: str | None = None,
-    **card_kwargs,
-):
-    v1 = get_safe_version(repo_id, V16)
-    v1x_dir = local_dir / V16 / repo_id
-    v20_dir = local_dir / V20 / repo_id
-    v1x_dir.mkdir(parents=True, exist_ok=True)
-    v20_dir.mkdir(parents=True, exist_ok=True)
-
-    hub_api = HfApi()
-    hub_api.snapshot_download(
-        repo_id=repo_id, repo_type="dataset", revision=v1, local_dir=v1x_dir, ignore_patterns="videos*/"
-    )
-    branch = "main"
-    if test_branch:
-        branch = test_branch
-        create_branch(repo_id=repo_id, branch=test_branch, repo_type="dataset")
-
-    metadata_v1 = load_json(v1x_dir / V1_INFO_PATH)
-    dataset = datasets.load_dataset("parquet", data_dir=v1x_dir / "data", split="train")
-    features = get_features_from_hf_dataset(dataset, robot_config)
-    video_keys = [key for key, ft in features.items() if ft["dtype"] == "video"]
-
-    if single_task and "language_instruction" in dataset.column_names:
-        logging.warning(
-            "'single_task' provided but 'language_instruction' tasks_col found. Using 'language_instruction'.",
-        )
-        single_task = None
-        tasks_col = "language_instruction"
-
-    # Episodes & chunks
-    episode_indices = sorted(dataset.unique("episode_index"))
-    total_episodes = len(episode_indices)
-    assert episode_indices == list(range(total_episodes))
-    total_videos = total_episodes * len(video_keys)
-    total_chunks = total_episodes // DEFAULT_CHUNK_SIZE
-    if total_episodes % DEFAULT_CHUNK_SIZE != 0:
-        total_chunks += 1
-
-    # Tasks
-    if single_task:
-        tasks_by_episodes = dict.fromkeys(episode_indices, single_task)
-        dataset, tasks = add_task_index_by_episodes(dataset, tasks_by_episodes)
-        tasks_by_episodes = {ep_idx: [task] for ep_idx, task in tasks_by_episodes.items()}
-    elif tasks_path:
-        tasks_by_episodes = load_json(tasks_path)
-        tasks_by_episodes = {int(ep_idx): task for ep_idx, task in tasks_by_episodes.items()}
-        dataset, tasks = add_task_index_by_episodes(dataset, tasks_by_episodes)
-        tasks_by_episodes = {ep_idx: [task] for ep_idx, task in tasks_by_episodes.items()}
-    elif tasks_col:
-        dataset, tasks, tasks_by_episodes = add_task_index_from_tasks_col(dataset, tasks_col)
-    else:
-        raise ValueError
-
-    assert set(tasks) == {task for ep_tasks in tasks_by_episodes.values() for task in ep_tasks}
-    tasks = [{"task_index": task_idx, "task": task} for task_idx, task in enumerate(tasks)]
-    write_jsonlines(tasks, v20_dir / LEGACY_TASKS_PATH)
-    features["task_index"] = {
-        "dtype": "int64",
-        "shape": (1,),
-        "names": None,
-    }
-
-    # Videos
-    if video_keys:
-        assert metadata_v1.get("video", False)
-        dataset = dataset.remove_columns(video_keys)
-        clean_gitattr = Path(
-            hub_api.hf_hub_download(
-                repo_id=GITATTRIBUTES_REF, repo_type="dataset", local_dir=local_dir, filename=".gitattributes"
-            )
-        ).absolute()
-        with tempfile.TemporaryDirectory() as tmp_video_dir:
-            move_videos(
-                repo_id, video_keys, total_episodes, total_chunks, Path(tmp_video_dir), clean_gitattr, branch
-            )
-        videos_info = get_videos_info(repo_id, v1x_dir, video_keys=video_keys, branch=branch)
-        for key in video_keys:
-            features[key]["shape"] = (
-                videos_info[key].pop("video.height"),
-                videos_info[key].pop("video.width"),
-                videos_info[key].pop("video.channels"),
-            )
-            features[key]["video_info"] = videos_info[key]
-            assert math.isclose(videos_info[key]["video.fps"], metadata_v1["fps"], rel_tol=1e-3)
-            if "encoding" in metadata_v1:
-                assert videos_info[key]["video.pix_fmt"] == metadata_v1["encoding"]["pix_fmt"]
-    else:
-        assert metadata_v1.get("video", 0) == 0
-        videos_info = None
-
-    # Split data into 1 parquet file by episode
-    episode_lengths = split_parquet_by_episodes(dataset, total_episodes, total_chunks, v20_dir)
-
-    if robot_config is not None:
-        robot_type = robot_config.type
-        repo_tags = [robot_type]
-    else:
-        robot_type = "unknown"
-        repo_tags = None
-
-    # Episodes
-    episodes = [
-        {"episode_index": ep_idx, "tasks": tasks_by_episodes[ep_idx], "length": episode_lengths[ep_idx]}
-        for ep_idx in episode_indices
-    ]
-    write_jsonlines(episodes, v20_dir / LEGACY_EPISODES_PATH)
-
-    # Assemble metadata v2.0
-    metadata_v2_0 = {
-        "codebase_version": V20,
-        "robot_type": robot_type,
-        "total_episodes": total_episodes,
-        "total_frames": len(dataset),
-        "total_tasks": len(tasks),
-        "total_videos": total_videos,
-        "total_chunks": total_chunks,
-        "chunks_size": DEFAULT_CHUNK_SIZE,
-        "fps": metadata_v1["fps"],
-        "splits": {"train": f"0:{total_episodes}"},
-        "data_path": LEGACY_DEFAULT_PARQUET_PATH,
-        "video_path": LEGACY_DEFAULT_VIDEO_PATH if video_keys else None,
-        "features": features,
-    }
-    write_json(metadata_v2_0, v20_dir / INFO_PATH)
-    convert_stats_to_json(v1x_dir, v20_dir)
-    card = create_lerobot_dataset_card(tags=repo_tags, dataset_info=metadata_v2_0, **card_kwargs)
-
-    with contextlib.suppress(EntryNotFoundError, HfHubHTTPError):
-        hub_api.delete_folder(repo_id=repo_id, path_in_repo="data", repo_type="dataset", revision=branch)
-
-    with contextlib.suppress(EntryNotFoundError, HfHubHTTPError):
-        hub_api.delete_folder(repo_id=repo_id, path_in_repo="meta_data", repo_type="dataset", revision=branch)
-
-    with contextlib.suppress(EntryNotFoundError, HfHubHTTPError):
-        hub_api.delete_folder(repo_id=repo_id, path_in_repo="meta", repo_type="dataset", revision=branch)
-
-    hub_api.upload_folder(
-        repo_id=repo_id,
-        path_in_repo="data",
-        folder_path=v20_dir / "data",
-        repo_type="dataset",
-        revision=branch,
-    )
-    hub_api.upload_folder(
-        repo_id=repo_id,
-        path_in_repo="meta",
-        folder_path=v20_dir / "meta",
-        repo_type="dataset",
-        revision=branch,
-    )
-
-    card.push_to_hub(repo_id=repo_id, repo_type="dataset", revision=branch)
-
-    if not test_branch:
-        create_branch(repo_id=repo_id, branch=V20, repo_type="dataset")
-
-
-def make_robot_config(robot_type: str, **kwargs) -> RobotConfig:
-    if robot_type == "aloha":
-        raise NotImplementedError  # TODO
-
-    elif robot_type == "koch_follower":
-        from lerobot.robots.koch_follower import KochFollowerConfig
-
-        return KochFollowerConfig(**kwargs)
-    elif robot_type == "so100_follower":
-        from lerobot.robots.so100_follower import SO100FollowerConfig
-
-        return SO100FollowerConfig(**kwargs)
-    elif robot_type == "stretch":
-        from lerobot.robots.stretch3 import Stretch3RobotConfig
-
-        return Stretch3RobotConfig(**kwargs)
-    elif robot_type == "lekiwi":
-        from lerobot.robots.lekiwi import LeKiwiConfig
-
-        return LeKiwiConfig(**kwargs)
-    else:
-        raise ValueError(f"Robot type '{robot_type}' is not available.")
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    task_args = parser.add_mutually_exclusive_group(required=True)
-
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        required=True,
-        help="Repository identifier on Hugging Face: a community or a user name `/` the name of the dataset (e.g. `lerobot/pusht`, `cadene/aloha_sim_insertion_human`).",
-    )
-    task_args.add_argument(
-        "--single-task",
-        type=str,
-        help="A short but accurate description of the single task performed in the dataset.",
-    )
-    task_args.add_argument(
-        "--tasks-col",
-        type=str,
-        help="The name of the column containing language instructions",
-    )
-    task_args.add_argument(
-        "--tasks-path",
-        type=Path,
-        help="The path to a .json file containing one language instruction for each episode_index",
-    )
-    parser.add_argument(
-        "--robot",
-        type=str,
-        default=None,
-        help="Robot config used for the dataset during conversion (e.g. 'koch', 'aloha', 'so100', etc.)",
-    )
-    parser.add_argument(
-        "--local-dir",
-        type=Path,
-        default=None,
-        help="Local directory to store the dataset during conversion. Defaults to /tmp/lerobot_dataset_v2",
-    )
-    parser.add_argument(
-        "--license",
-        type=str,
-        default="apache-2.0",
-        help="Repo license. Must be one of https://huggingface.co/docs/hub/repositories-licenses. Defaults to mit.",
-    )
-    parser.add_argument(
-        "--test-branch",
-        type=str,
-        default=None,
-        help="Repo branch to test your conversion first (e.g. 'v2.0.test')",
-    )
-
-    args = parser.parse_args()
-    if not args.local_dir:
-        args.local_dir = Path("/tmp/lerobot_dataset_v2")
-
-    if args.robot is not None:
-        robot_config = make_robot_config(args.robot)
-
-    del args.robot
-
-    convert_dataset(**vars(args), robot_config=robot_config)
-
-
-if __name__ == "__main__":
-    main()
@@ -1,87 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import logging
-import traceback
-from pathlib import Path
-
-from datasets import get_dataset_config_info
-from huggingface_hub import HfApi
-
-from lerobot import available_datasets
-from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.datasets.utils import INFO_PATH, write_info
-from lerobot.datasets.v21.convert_dataset_v20_to_v21 import V20, SuppressWarnings
-
-LOCAL_DIR = Path("data/")
-
-hub_api = HfApi()
-
-
-def fix_dataset(repo_id: str) -> str:
-    if not hub_api.revision_exists(repo_id, V20, repo_type="dataset"):
-        return f"{repo_id}: skipped (not in {V20})."
-
-    dataset_info = get_dataset_config_info(repo_id, "default")
-    with SuppressWarnings():
-        lerobot_metadata = LeRobotDatasetMetadata(repo_id, revision=V20, force_cache_sync=True)
-
-    meta_features = {key for key, ft in lerobot_metadata.features.items() if ft["dtype"] != "video"}
-    parquet_features = set(dataset_info.features)
-
-    diff_parquet_meta = parquet_features - meta_features
-    diff_meta_parquet = meta_features - parquet_features
-
-    if diff_parquet_meta:
-        raise ValueError(f"In parquet not in info.json: {parquet_features - meta_features}")
-
-    if not diff_meta_parquet:
-        return f"{repo_id}: skipped (no diff)"
-
-    if diff_meta_parquet:
-        logging.warning(f"In info.json not in parquet: {meta_features - parquet_features}")
-        assert diff_meta_parquet == {"language_instruction"}
-        lerobot_metadata.features.pop("language_instruction")
-        write_info(lerobot_metadata.info, lerobot_metadata.root)
-        commit_info = hub_api.upload_file(
-            path_or_fileobj=lerobot_metadata.root / INFO_PATH,
-            path_in_repo=INFO_PATH,
-            repo_id=repo_id,
-            repo_type="dataset",
-            revision=V20,
-            commit_message="Remove 'language_instruction'",
-            create_pr=True,
-        )
-        return f"{repo_id}: success - PR: {commit_info.pr_url}"
-
-
-def batch_fix():
-    status = {}
-    LOCAL_DIR.mkdir(parents=True, exist_ok=True)
-    logfile = LOCAL_DIR / "fix_features_v20.txt"
-    for num, repo_id in enumerate(available_datasets):
-        print(f"\nConverting {repo_id} ({num}/{len(available_datasets)})")
-        print("---------------------------------------------------------")
-        try:
-            status = fix_dataset(repo_id)
-        except Exception:
-            status = f"{repo_id}: failed\n    {traceback.format_exc()}"
-
-        logging.info(status)
-        with open(logfile, "a") as file:
-            file.write(status + "\n")
-
-
-if __name__ == "__main__":
-    batch_fix()
@@ -1,54 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-This script is for internal use to convert all datasets under the 'lerobot' hub user account to v2.1.
-"""
-
-import traceback
-from pathlib import Path
-
-from huggingface_hub import HfApi
-
-from lerobot import available_datasets
-from lerobot.datasets.v21.convert_dataset_v20_to_v21 import V21, convert_dataset
-
-LOCAL_DIR = Path("data/")
-
-
-def batch_convert():
-    status = {}
-    LOCAL_DIR.mkdir(parents=True, exist_ok=True)
-    logfile = LOCAL_DIR / "conversion_log_v21.txt"
-    hub_api = HfApi()
-    for num, repo_id in enumerate(available_datasets):
-        print(f"\nConverting {repo_id} ({num}/{len(available_datasets)})")
-        print("---------------------------------------------------------")
-        try:
-            if hub_api.revision_exists(repo_id, V21, repo_type="dataset"):
-                status = f"{repo_id}: success (already in {V21})."
-            else:
-                convert_dataset(repo_id)
-                status = f"{repo_id}: success."
-        except Exception:
-            status = f"{repo_id}: failed\n    {traceback.format_exc()}"
-
-        with open(logfile, "a") as file:
-            file.write(status + "\n")
-
-
-if __name__ == "__main__":
-    batch_convert()
@@ -1,111 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 2.0 to
-2.1. It will:
-
- Generate per-episodes stats and writes them in `episodes_stats.jsonl`
- Check consistency between these new stats and the old ones.
- Remove the deprecated `stats.json`.
- Update codebase_version in `info.json`.
- Push this new version to the hub on the 'main' branch and tags it with "v2.1".
-
-Usage:
-
-```bash
-python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 \
-    --repo-id=aliberts/koch_tutorial
-```
-
-"""
-
-import argparse
-import logging
-
-from huggingface_hub import HfApi
-
-from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
-from lerobot.datasets.utils import STATS_PATH, load_stats, write_info
-from lerobot.datasets.v21.convert_stats import check_aggregate_stats, convert_stats
-
-V20 = "v2.0"
-V21 = "v2.1"
-
-
-class SuppressWarnings:
-    def __enter__(self):
-        self.previous_level = logging.getLogger().getEffectiveLevel()
-        logging.getLogger().setLevel(logging.ERROR)
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        logging.getLogger().setLevel(self.previous_level)
-
-
-def convert_dataset(
-    repo_id: str,
-    branch: str | None = None,
-    num_workers: int = 4,
-):
-    with SuppressWarnings():
-        dataset = LeRobotDataset(repo_id, revision=V20, force_cache_sync=True)
-
-    convert_stats(dataset, num_workers=num_workers)
-    ref_stats = load_stats(dataset.root)
-    check_aggregate_stats(dataset, ref_stats)
-
-    dataset.meta.info["codebase_version"] = CODEBASE_VERSION
-    write_info(dataset.meta.info, dataset.root)
-
-    dataset.push_to_hub(branch=branch, tag_version=False, allow_patterns="meta/")
-
-    # delete old stats.json file
-    if (dataset.root / STATS_PATH).is_file:
-        (dataset.root / STATS_PATH).unlink()
-
-    hub_api = HfApi()
-    if hub_api.file_exists(
-        repo_id=dataset.repo_id, filename=STATS_PATH, revision=branch, repo_type="dataset"
-    ):
-        hub_api.delete_file(
-            path_in_repo=STATS_PATH, repo_id=dataset.repo_id, revision=branch, repo_type="dataset"
-        )
-
-    hub_api.create_tag(repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        required=True,
-        help="Repository identifier on Hugging Face: a community or a user name `/` the name of the dataset "
-        "(e.g. `lerobot/pusht`, `cadene/aloha_sim_insertion_human`).",
-    )
-    parser.add_argument(
-        "--branch",
-        type=str,
-        default=None,
-        help="Repo branch to push your dataset. Defaults to the main branch.",
-    )
-    parser.add_argument(
-        "--num-workers",
-        type=int,
-        default=4,
-        help="Number of workers for parallelizing stats compute. Defaults to 4.",
-    )
-
-    args = parser.parse_args()
-    convert_dataset(**vars(args))
@@ -1,114 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from pathlib import Path
-
-import jsonlines
-import numpy as np
-from tqdm import tqdm
-
-from lerobot.datasets.compute_stats import aggregate_stats, get_feature_stats, sample_indices
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.utils import LEGACY_EPISODES_STATS_PATH, serialize_dict
-
-
-def append_jsonlines(data: dict, fpath: Path) -> None:
-    fpath.parent.mkdir(exist_ok=True, parents=True)
-    with jsonlines.open(fpath, "a") as writer:
-        writer.write(data)
-
-
-def legacy_write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path):
-    # We wrap episode_stats in a dictionary since `episode_stats["episode_index"]`
-    # is a dictionary of stats and not an integer.
-    episode_stats = {"episode_index": episode_index, "stats": serialize_dict(episode_stats)}
-    append_jsonlines(episode_stats, local_dir / LEGACY_EPISODES_STATS_PATH)
-
-
-def sample_episode_video_frames(dataset: LeRobotDataset, episode_index: int, ft_key: str) -> np.ndarray:
-    ep_len = dataset.meta.episodes[episode_index]["length"]
-    sampled_indices = sample_indices(ep_len)
-    query_timestamps = dataset._get_query_timestamps(0.0, {ft_key: sampled_indices})
-    video_frames = dataset._query_videos(query_timestamps, episode_index)
-    return video_frames[ft_key].numpy()
-
-
-def convert_episode_stats(dataset: LeRobotDataset, ep_idx: int):
-    ep_start_idx = dataset.episode_data_index["from"][ep_idx]
-    ep_end_idx = dataset.episode_data_index["to"][ep_idx]
-    ep_data = dataset.hf_dataset.select(range(ep_start_idx, ep_end_idx))
-
-    ep_stats = {}
-    for key, ft in dataset.features.items():
-        if ft["dtype"] == "video":
-            # We sample only for videos
-            ep_ft_data = sample_episode_video_frames(dataset, ep_idx, key)
-        else:
-            ep_ft_data = np.array(ep_data[key])
-
-        axes_to_reduce = (0, 2, 3) if ft["dtype"] in ["image", "video"] else 0
-        keepdims = True if ft["dtype"] in ["image", "video"] else ep_ft_data.ndim == 1
-        ep_stats[key] = get_feature_stats(ep_ft_data, axis=axes_to_reduce, keepdims=keepdims)
-
-        if ft["dtype"] in ["image", "video"]:  # remove batch dim
-            ep_stats[key] = {
-                k: v if k == "count" else np.squeeze(v, axis=0) for k, v in ep_stats[key].items()
-            }
-
-    dataset.meta.episodes_stats[ep_idx] = ep_stats
-
-
-def convert_stats(dataset: LeRobotDataset, num_workers: int = 0):
-    assert dataset.episodes is None
-    print("Computing episodes stats")
-    total_episodes = dataset.meta.total_episodes
-    if num_workers > 0:
-        with ThreadPoolExecutor(max_workers=num_workers) as executor:
-            futures = {
-                executor.submit(convert_episode_stats, dataset, ep_idx): ep_idx
-                for ep_idx in range(total_episodes)
-            }
-            for future in tqdm(as_completed(futures), total=total_episodes):
-                future.result()
-    else:
-        for ep_idx in tqdm(range(total_episodes)):
-            convert_episode_stats(dataset, ep_idx)
-
-    for ep_idx in tqdm(range(total_episodes)):
-        legacy_write_episode_stats(ep_idx, dataset.meta.episodes_stats[ep_idx], dataset.root)
-
-
-def check_aggregate_stats(
-    dataset: LeRobotDataset,
-    reference_stats: dict[str, dict[str, np.ndarray]],
-    video_rtol_atol: tuple[float] = (1e-2, 1e-2),
-    default_rtol_atol: tuple[float] = (5e-6, 6e-5),
-):
-    """Verifies that the aggregated stats from episodes_stats are close to reference stats."""
-    agg_stats = aggregate_stats(list(dataset.meta.episodes_stats.values()))
-    for key, ft in dataset.features.items():
-        # These values might need some fine-tuning
-        if ft["dtype"] == "video":
-            # to account for image sub-sampling
-            rtol, atol = video_rtol_atol
-        else:
-            rtol, atol = default_rtol_atol
-
-        for stat, val in agg_stats[key].items():
-            if key in reference_stats and stat in reference_stats[key]:
-                err_msg = f"feature='{key}' stats='{stat}'"
-                np.testing.assert_allclose(
-                    val, reference_stats[key][stat], rtol=rtol, atol=atol, err_msg=err_msg
-                )
@@ -62,7 +62,6 @@ from lerobot.datasets.utils import (
    flatten_dict,
    get_parquet_file_size_in_mb,
    get_parquet_num_frames,
-    get_video_duration_in_s,
    get_video_size_in_mb,
    load_info,
    update_chunk_file_indices,
@@ -71,7 +70,7 @@ from lerobot.datasets.utils import (
    write_stats,
    write_tasks,
 )
-from lerobot.datasets.video_utils import concat_video_files
+from lerobot.datasets.video_utils import concat_video_files, get_video_duration_in_s

 V21 = "v2.1"

@@ -173,7 +172,7 @@ def concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys):
    concatenated_df.to_parquet(path, index=False, schema=schema)


-def convert_data(root, new_root):
+def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
    data_dir = root / "data"
    ep_paths = sorted(data_dir.glob("*/*.parquet"))

@@ -201,7 +200,7 @@ def convert_data(root, new_root):
        episodes_metadata.append(ep_metadata)
        ep_idx += 1

-        if size_in_mb < DEFAULT_DATA_FILE_SIZE_IN_MB:
+        if size_in_mb < data_file_size_in_mb:
            paths_to_cat.append(ep_path)
            continue

@@ -235,7 +234,7 @@ def get_image_keys(root):
    return image_keys


-def convert_videos(root: Path, new_root: Path):
+def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
    video_keys = get_video_keys(root)
    if len(video_keys) == 0:
        return None
@@ -244,7 +243,7 @@ def convert_videos(root: Path, new_root: Path):

    eps_metadata_per_cam = []
    for camera in video_keys:
-        eps_metadata = convert_videos_of_camera(root, new_root, camera)
+        eps_metadata = convert_videos_of_camera(root, new_root, camera, video_file_size_in_mb)
        eps_metadata_per_cam.append(eps_metadata)

    num_eps_per_cam = [len(eps_cam_map) for eps_cam_map in eps_metadata_per_cam]
@@ -269,7 +268,7 @@ def convert_videos(root: Path, new_root: Path):
    return episods_metadata


-def convert_videos_of_camera(root: Path, new_root: Path, video_key):
+def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_file_size_in_mb: int):
    # Access old paths to mp4
    videos_dir = root / "videos"
    ep_paths = sorted(videos_dir.glob(f"*/{video_key}/*.mp4"))
@@ -286,7 +285,7 @@ def convert_videos_of_camera(root: Path, new_root: Path, video_key):
        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 >= DEFAULT_VIDEO_FILE_SIZE_IN_MB and len(paths_to_cat) > 0:
+        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
            concat_video_files(paths_to_cat, new_root, video_key, chunk_idx, file_idx)

@@ -387,13 +386,13 @@ def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_
    write_stats(stats, new_root)


-def convert_info(root, new_root):
+def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
    info = load_info(root)
    info["codebase_version"] = "v3.0"
    del info["total_chunks"]
    del info["total_videos"]
-    info["data_files_size_in_mb"] = DEFAULT_DATA_FILE_SIZE_IN_MB
-    info["video_files_size_in_mb"] = DEFAULT_VIDEO_FILE_SIZE_IN_MB
+    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
    info["fps"] = float(info["fps"])
@@ -408,12 +407,18 @@ def convert_info(root, new_root):
 def convert_dataset(
    repo_id: str,
    branch: str | None = None,
-    num_workers: int = 4,
+    data_file_size_in_mb: int | None = None,
+    video_file_size_in_mb: int | None = None,
 ):
    root = HF_LEROBOT_HOME / repo_id
    old_root = HF_LEROBOT_HOME / f"{repo_id}_old"
    new_root = HF_LEROBOT_HOME / f"{repo_id}_v30"

+    if data_file_size_in_mb is None:
+        data_file_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
+    if video_file_size_in_mb is None:
+        video_file_size_in_mb = DEFAULT_VIDEO_FILE_SIZE_IN_MB
+
    if old_root.is_dir() and root.is_dir():
        shutil.rmtree(str(root))
        shutil.move(str(old_root), str(root))
@@ -428,10 +433,10 @@ def convert_dataset(
        local_dir=root,
    )

-    convert_info(root, new_root)
+    convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb)
    convert_tasks(root, new_root)
-    episodes_metadata = convert_data(root, new_root)
-    episodes_videos_metadata = convert_videos(root, new_root)
+    episodes_metadata = convert_data(root, new_root, data_file_size_in_mb)
+    episodes_videos_metadata = convert_videos(root, new_root, video_file_size_in_mb)
    convert_episodes_metadata(root, new_root, episodes_metadata, episodes_videos_metadata)

    shutil.move(str(root), str(old_root))
@@ -470,10 +475,16 @@ if __name__ == "__main__":
        help="Repo branch to push your dataset. Defaults to the main branch.",
    )
    parser.add_argument(
-        "--num-workers",
+        "--data-file-size-in-mb",
        type=int,
-        default=4,
-        help="Number of workers for parallelizing stats compute. Defaults to 4.",
+        default=None,
+        help="File size in MB. Defaults to 100 for data and 500 for videos.",
+    )
+    parser.add_argument(
+        "--video-file-size-in-mb",
+        type=int,
+        default=None,
+        help="File size in MB. Defaults to 100 for data and 500 for videos.",
    )

    args = parser.parse_args()
@@ -13,18 +13,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import glob
 import importlib
-import json
 import logging
 import shutil
 import subprocess
 import tempfile
 import warnings
-from collections import OrderedDict
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, ClassVar

+import av
 import pyarrow as pa
 import torch
 import torchvision
@@ -106,7 +106,7 @@ def decode_video_frames_torchvision(
    keyframes_only = False
    torchvision.set_video_backend(backend)
    if backend == "pyav":
-        keyframes_only = True  # pyav doesnt support accuracte seek
+        keyframes_only = True  # pyav doesn't support accurate seek

    # set a video stream reader
    # TODO(rcadene): also load audio stream at the same time
@@ -159,7 +159,6 @@ def decode_video_frames_torchvision(
    )

    # get closest frames to the query timestamps
-    # TODO(rcadene): remove torch.stack
    closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_])
    closest_ts = loaded_ts[argmin_]

@@ -257,51 +256,83 @@ def encode_video_frames(
    g: int | None = 2,
    crf: int | None = 30,
    fast_decode: int = 0,
-    log_level: str | None = "quiet",
+    log_level: int | None = av.logging.ERROR,
    overwrite: bool = False,
 ) -> None:
    """More info on ffmpeg arguments tuning on `benchmark/video/README.md`"""
+    # Check encoder availability
+    if vcodec not in ["h264", "hevc", "libsvtav1"]:
+        raise ValueError(f"Unsupported video codec: {vcodec}. Supported codecs are: h264, hevc, libsvtav1.")
+
    video_path = Path(video_path)
    imgs_dir = Path(imgs_dir)
+
    video_path.parent.mkdir(parents=True, exist_ok=True)

-    ffmpeg_args = OrderedDict(
-        [
-            ("-f", "image2"),
-            ("-r", str(fps)),
-            ("-i", str(imgs_dir / "frame-%06d.png")),
-            ("-vcodec", vcodec),
-            ("-pix_fmt", pix_fmt),
-        ]
+    # Encoders/pixel formats incompatibility check
+    if (vcodec == "libsvtav1" or vcodec == "hevc") and pix_fmt == "yuv444p":
+        logging.warning(
+            f"Incompatible pixel format 'yuv444p' for codec {vcodec}, auto-selecting format 'yuv420p'"
+        )
+        pix_fmt = "yuv420p"
+
+    # Get input frames
+    template = "frame-" + ("[0-9]" * 6) + ".png"
+    input_list = sorted(
+        glob.glob(str(imgs_dir / template)), key=lambda x: int(x.split("-")[-1].split(".")[0])
    )

+    # Define video output frame size (assuming all input frames are the same size)
+    if len(input_list) == 0:
+        raise FileNotFoundError(f"No images found in {imgs_dir}.")
+    dummy_image = Image.open(input_list[0])
+    width, height = dummy_image.size
+
+    # Define video codec options
+    video_options = {}
+
    if g is not None:
-        ffmpeg_args["-g"] = str(g)
+        video_options["g"] = str(g)

    if crf is not None:
-        ffmpeg_args["-crf"] = str(crf)
+        video_options["crf"] = str(crf)

    if fast_decode:
-        key = "-svtav1-params" if vcodec == "libsvtav1" else "-tune"
+        key = "svtav1-params" if vcodec == "libsvtav1" else "tune"
        value = f"fast-decode={fast_decode}" if vcodec == "libsvtav1" else "fastdecode"
-        ffmpeg_args[key] = value
+        video_options[key] = value

+    # Set logging level
    if log_level is not None:
-        ffmpeg_args["-loglevel"] = str(log_level)
+        # "While less efficient, it is generally preferable to modify logging with Python's logging"
+        logging.getLogger("libav").setLevel(log_level)

-    ffmpeg_args = [item for pair in ffmpeg_args.items() for item in pair]
-    if overwrite:
-        ffmpeg_args.append("-y")
+    # Create and open output file (overwrite by default)
+    with av.open(str(video_path), "w") as output:
+        output_stream = output.add_stream(vcodec, fps, options=video_options)
+        output_stream.pix_fmt = pix_fmt
+        output_stream.width = width
+        output_stream.height = height

-    ffmpeg_cmd = ["ffmpeg"] + ffmpeg_args + [str(video_path)]
-    # redirect stdin to subprocess.DEVNULL to prevent reading random keyboard inputs from terminal
-    subprocess.run(ffmpeg_cmd, check=True, stdin=subprocess.DEVNULL)
+        # Loop through input frames and encode them
+        for input_data in input_list:
+            input_image = Image.open(input_data).convert("RGB")
+            input_frame = av.VideoFrame.from_image(input_image)
+            packet = output_stream.encode(input_frame)
+            if packet:
+                output.mux(packet)
+
+        # Flush the encoder
+        packet = output_stream.encode()
+        if packet:
+            output.mux(packet)
+
+    # Reset logging level
+    if log_level is not None:
+        av.logging.restore_default_callback()

    if not video_path.exists():
-        raise OSError(
-            f"Video encoding did not work. File not found: {video_path}. "
-            f"Try running the command manually to debug: `{''.join(ffmpeg_cmd)}`"
-        )
+        raise OSError(f"Video encoding did not work. File not found: {video_path}.")


 def concat_video_files(paths_to_cat: list[Path], root: Path, video_key: str, chunk_idx: int, file_idx: int):
@@ -325,6 +356,9 @@ def concat_video_files(paths_to_cat: list[Path], root: Path, video_key: str, chu
          codec, resolution, and frame rate for proper concatenation.
        - Output path follows the DEFAULT_VIDEO_PATH pattern with video_key, chunk_idx,
          and file_idx parameters.
+        - This function uses subprocess to call ffmpeg directly because PyAV doesn't have
+          built-in support for video concatenation. The concat demuxer in ffmpeg handles
+          all the complex timestamp adjustments automatically.
    """

    tmp_dir = Path(tempfile.mkdtemp(dir=root))
@@ -390,78 +424,68 @@ with warnings.catch_warnings():


 def get_audio_info(video_path: Path | str) -> dict:
-    ffprobe_audio_cmd = [
-        "ffprobe",
-        "-v",
-        "error",
-        "-select_streams",
-        "a:0",
-        "-show_entries",
-        "stream=channels,codec_name,bit_rate,sample_rate,bit_depth,channel_layout,duration",
-        "-of",
-        "json",
-        str(video_path),
-    ]
-    result = subprocess.run(ffprobe_audio_cmd, capture_output=True, text=True)
-    if result.returncode != 0:
-        raise RuntimeError(f"Error running ffprobe: {result.stderr}")
+    # Set logging level
+    logging.getLogger("libav").setLevel(av.logging.ERROR)

-    info = json.loads(result.stdout)
-    audio_stream_info = info["streams"][0] if info.get("streams") else None
-    if audio_stream_info is None:
-        return {"has_audio": False}
+    # Getting audio stream information
+    audio_info = {}
+    with av.open(str(video_path), "r") as audio_file:
+        try:
+            audio_stream = audio_file.streams.audio[0]
+        except IndexError:
+            # Reset logging level
+            av.logging.restore_default_callback()
+            return {"has_audio": False}

-    # Return the information, defaulting to None if no audio stream is present
-    return {
-        "has_audio": True,
-        "audio.channels": audio_stream_info.get("channels", None),
-        "audio.codec": audio_stream_info.get("codec_name", None),
-        "audio.bit_rate": int(audio_stream_info["bit_rate"]) if audio_stream_info.get("bit_rate") else None,
-        "audio.sample_rate": int(audio_stream_info["sample_rate"])
-        if audio_stream_info.get("sample_rate")
-        else None,
-        "audio.bit_depth": audio_stream_info.get("bit_depth", None),
-        "audio.channel_layout": audio_stream_info.get("channel_layout", None),
-    }
+        audio_info["audio.channels"] = audio_stream.channels
+        audio_info["audio.codec"] = audio_stream.codec.canonical_name
+        # In an ideal loseless case : bit depth x sample rate x channels = bit rate.
+        # In an actual compressed case, the bit rate is set according to the compression level : the lower the bit rate, the more compression is applied.
+        audio_info["audio.bit_rate"] = audio_stream.bit_rate
+        audio_info["audio.sample_rate"] = audio_stream.sample_rate  # Number of samples per second
+        # In an ideal loseless case : fixed number of bits per sample.
+        # In an actual compressed case : variable number of bits per sample (often reduced to match a given depth rate).
+        audio_info["audio.bit_depth"] = audio_stream.format.bits
+        audio_info["audio.channel_layout"] = audio_stream.layout.name
+        audio_info["has_audio"] = True
+
+    # Reset logging level
+    av.logging.restore_default_callback()
+
+    return audio_info


 def get_video_info(video_path: Path | str) -> dict:
-    ffprobe_video_cmd = [
-        "ffprobe",
-        "-v",
-        "error",
-        "-select_streams",
-        "v:0",
-        "-show_entries",
-        "stream=r_frame_rate,width,height,codec_name,nb_frames,duration,pix_fmt",
-        "-of",
-        "json",
-        str(video_path),
-    ]
-    result = subprocess.run(ffprobe_video_cmd, capture_output=True, text=True)
-    if result.returncode != 0:
-        raise RuntimeError(f"Error running ffprobe: {result.stderr}")
+    # Set logging level
+    logging.getLogger("libav").setLevel(av.logging.ERROR)

-    info = json.loads(result.stdout)
-    video_stream_info = info["streams"][0]
+    # Getting video stream information
+    video_info = {}
+    with av.open(str(video_path), "r") as video_file:
+        try:
+            video_stream = video_file.streams.video[0]
+        except IndexError:
+            # Reset logging level
+            av.logging.restore_default_callback()
+            return {}

-    # Calculate fps from r_frame_rate
-    r_frame_rate = video_stream_info["r_frame_rate"]
-    num, denom = map(int, r_frame_rate.split("/"))
-    fps = num / denom
+        video_info["video.height"] = video_stream.height
+        video_info["video.width"] = video_stream.width
+        video_info["video.codec"] = video_stream.codec.canonical_name
+        video_info["video.pix_fmt"] = video_stream.pix_fmt
+        video_info["video.is_depth_map"] = False

-    pixel_channels = get_video_pixel_channels(video_stream_info["pix_fmt"])
+        # Calculate fps from r_frame_rate
+        video_info["video.fps"] = int(video_stream.base_rate)

-    video_info = {
-        "video.fps": fps,
-        "video.height": video_stream_info["height"],
-        "video.width": video_stream_info["width"],
-        "video.channels": pixel_channels,
-        "video.codec": video_stream_info["codec_name"],
-        "video.pix_fmt": video_stream_info["pix_fmt"],
-        "video.is_depth_map": False,
-        **get_audio_info(video_path),
-    }
+        pixel_channels = get_video_pixel_channels(video_stream.pix_fmt)
+        video_info["video.channels"] = pixel_channels
+
+    # Reset logging level
+    av.logging.restore_default_callback()
+
+    # Adding audio stream information
+    video_info.update(**get_audio_info(video_path))

    return video_info

@@ -488,3 +512,25 @@ def get_image_pixel_channels(image: Image):
        return 4  # RGBA
    else:
        raise ValueError("Unknown format")
+
+
+def get_video_duration_in_s(video_path: Path | str) -> float:
+    """
+    Get the duration of a video file in seconds using PyAV.
+
+    Args:
+        video_path: Path to the video file.
+
+    Returns:
+        Duration of the video in seconds.
+    """
+    with av.open(str(video_path)) as container:
+        # Get the first video stream
+        video_stream = container.streams.video[0]
+        # Calculate duration: stream.duration * stream.time_base gives duration in seconds
+        if video_stream.duration is not None:
+            duration = float(video_stream.duration * video_stream.time_base)
+        else:
+            # Fallback to container duration if stream duration is not available
+            duration = float(container.duration / av.time_base)
+    return duration
@@ -107,6 +107,8 @@ X_SERIES_ENCODINGS_TABLE = {
    "Goal_PWM": X_SERIES_CONTROL_TABLE["Goal_PWM"][1],
    "Goal_Current": X_SERIES_CONTROL_TABLE["Goal_Current"][1],
    "Goal_Velocity": X_SERIES_CONTROL_TABLE["Goal_Velocity"][1],
+    "Goal_Position": X_SERIES_CONTROL_TABLE["Goal_Position"][1],
+    "Present_Position": X_SERIES_CONTROL_TABLE["Present_Position"][1],
    "Present_PWM": X_SERIES_CONTROL_TABLE["Present_PWM"][1],
    "Present_Current": X_SERIES_CONTROL_TABLE["Present_Current"][1],
    "Present_Velocity": X_SERIES_CONTROL_TABLE["Present_Velocity"][1],
@@ -115,11 +115,11 @@ If you uploaded your dataset to the hub with `--control.push_to_hub=true`, you c
 echo ${HF_USER}/aloha_test
 ```

-If you didn't upload with `--control.push_to_hub=false`, you can also visualize it locally with:
+If you didn't upload with `--control.push_to_hub=false`, you can also visualize it locally with [Rerun](https://github.com/rerun-io/rerun):

 ```bash
-python -m lerobot.scripts.visualize_dataset_html \
-  --repo-id ${HF_USER}/aloha_test
+python -m lerobot.scripts.visualize_dataset \
+  --repo-id ${HF_USER}/aloha_test --episode 0
 ```

 ## Replay an episode
@@ -302,11 +302,6 @@ class RobotClient:

                    self.logger.debug(f"Current latest action: {latest_action}")

-                    # Get queue state before changes
-                    old_size, old_timesteps = self._inspect_action_queue()
-                    if not old_timesteps:
-                        old_timesteps = [latest_action]  # queue was empty
-
                    # Get queue state before changes
                    old_size, old_timesteps = self._inspect_action_queue()
                    if not old_timesteps:
@@ -1,517 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Visualize data of **all** frames of any episode of a dataset of type LeRobotDataset.
-
-Note: The last frame of the episode doesnt always correspond to a final state.
-That's because our datasets are composed of transition from state to state up to
-the antepenultimate state associated to the ultimate action to arrive in the final state.
-However, there might not be a transition from a final state to another state.
-
-Note: This script aims to visualize the data used to train the neural networks.
-~What you see is what you get~. When visualizing image modality, it is often expected to observe
-lossly compression artifacts since these images have been decoded from compressed mp4 videos to
-save disk space. The compression factor applied has been tuned to not affect success rate.
-
-Example of usage:
-
- Visualize data stored on a local machine:
-```bash
-local$ python -m lerobot.scripts.visualize_dataset_html \
-    --repo-id lerobot/pusht
-
-local$ open http://localhost:9090
-```
-
- Visualize data stored on a distant machine with a local viewer:
-```bash
-distant$ python -m lerobot.scripts.visualize_dataset_html \
-    --repo-id lerobot/pusht
-
-local$ ssh -L 9090:localhost:9090 distant  # create a ssh tunnel
-local$ open http://localhost:9090
-```
-
- Select episodes to visualize:
-```bash
-python -m lerobot.scripts.visualize_dataset_html \
-    --repo-id lerobot/pusht \
-    --episodes 7 3 5 1 4
-```
-"""
-
-import argparse
-import csv
-import json
-import logging
-import re
-import shutil
-import tempfile
-from io import StringIO
-from pathlib import Path
-
-import numpy as np
-import pandas as pd
-import requests
-from flask import Flask, redirect, render_template, request, url_for
-
-from lerobot import available_datasets
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.utils import IterableNamespace
-from lerobot.utils.utils import init_logging
-
-
-def run_server(
-    dataset: LeRobotDataset | IterableNamespace | None,
-    episodes: list[int] | None,
-    host: str,
-    port: str,
-    static_folder: Path,
-    template_folder: Path,
-):
-    app = Flask(__name__, static_folder=static_folder.resolve(), template_folder=template_folder.resolve())
-    app.config["SEND_FILE_MAX_AGE_DEFAULT"] = 0  # specifying not to cache
-
-    @app.route("/")
-    def hommepage(dataset=dataset):
-        if dataset:
-            dataset_namespace, dataset_name = dataset.repo_id.split("/")
-            return redirect(
-                url_for(
-                    "show_episode",
-                    dataset_namespace=dataset_namespace,
-                    dataset_name=dataset_name,
-                    episode_id=0,
-                )
-            )
-
-        dataset_param, episode_param = None, None
-        all_params = request.args
-        if "dataset" in all_params:
-            dataset_param = all_params["dataset"]
-        if "episode" in all_params:
-            episode_param = int(all_params["episode"])
-
-        if dataset_param:
-            dataset_namespace, dataset_name = dataset_param.split("/")
-            return redirect(
-                url_for(
-                    "show_episode",
-                    dataset_namespace=dataset_namespace,
-                    dataset_name=dataset_name,
-                    episode_id=episode_param if episode_param is not None else 0,
-                )
-            )
-
-        featured_datasets = [
-            "lerobot/aloha_static_cups_open",
-            "lerobot/columbia_cairlab_pusht_real",
-            "lerobot/taco_play",
-        ]
-        return render_template(
-            "visualize_dataset_homepage.html",
-            featured_datasets=featured_datasets,
-            lerobot_datasets=available_datasets,
-        )
-
-    @app.route("/<string:dataset_namespace>/<string:dataset_name>")
-    def show_first_episode(dataset_namespace, dataset_name):
-        first_episode_id = 0
-        return redirect(
-            url_for(
-                "show_episode",
-                dataset_namespace=dataset_namespace,
-                dataset_name=dataset_name,
-                episode_id=first_episode_id,
-            )
-        )
-
-    @app.route("/<string:dataset_namespace>/<string:dataset_name>/episode_<int:episode_id>")
-    def show_episode(dataset_namespace, dataset_name, episode_id, dataset=dataset, episodes=episodes):
-        repo_id = f"{dataset_namespace}/{dataset_name}"
-        try:
-            if dataset is None:
-                dataset = get_dataset_info(repo_id)
-        except FileNotFoundError:
-            return (
-                "Make sure to convert your LeRobotDataset to v2 & above. See how to convert your dataset at https://github.com/huggingface/lerobot/pull/461",
-                400,
-            )
-        dataset_version = (
-            str(dataset.meta._version) if isinstance(dataset, LeRobotDataset) else dataset.codebase_version
-        )
-
-        # Check minimum version requirement
-        match = re.search(r"v(\d+)\.", dataset_version)
-        if match:
-            major_version = int(match.group(1))
-            if major_version < 2:
-                return "Make sure to convert your LeRobotDataset to v2 & above."
-
-        # Get episode data once
-        episode_data_csv_str, columns, ignored_columns = get_episode_data(dataset, episode_id)
-
-        dataset_info = {
-            "repo_id": f"{dataset_namespace}/{dataset_name}",
-            "num_samples": dataset.num_frames
-            if isinstance(dataset, LeRobotDataset)
-            else dataset.total_frames,
-            "num_episodes": dataset.num_episodes
-            if isinstance(dataset, LeRobotDataset)
-            else dataset.total_episodes,
-            "fps": dataset.fps,
-        }
-
-        if isinstance(dataset, LeRobotDataset):
-            # Handle local datasets
-            # Determine if this is a chunked video dataset (v3.0+)
-            is_v3_or_later = False
-            match = re.search(r"v(\d+)\.(\d+)", dataset_version)
-            if match:
-                major_version = int(match.group(1))
-                is_v3_or_later = major_version >= 3
-
-            # Create videos_info with unified structure
-            videos_info = []
-
-            for key in dataset.meta.video_keys:
-                video_path = dataset.meta.get_video_file_path(episode_id, key)
-
-                if is_v3_or_later:
-                    # For v3.0+ datasets, get episode timestamps from chunked videos
-                    episode = dataset.meta.episodes[episode_id]
-                    from_timestamp = episode.get(f"videos/{key}/from_timestamp", 0)
-                    to_timestamp = episode.get(f"videos/{key}/to_timestamp", None)
-                    filename = key
-                else:
-                    # For v2.1 and earlier, videos are already per-episode
-                    from_timestamp = None
-                    to_timestamp = None
-                    filename = video_path.parent.name
-
-                videos_info.append(
-                    {
-                        "url": url_for("static", filename=str(video_path).replace("\\", "/")),
-                        "filename": filename,
-                        "start_time": from_timestamp,
-                        "end_time": to_timestamp,
-                        "is_chunked": is_v3_or_later,
-                    }
-                )
-
-            tasks = dataset.meta.episodes[episode_id]["tasks"]
-        else:
-            # Handle remote datasets from HF Hub
-            video_keys = [key for key, ft in dataset.features.items() if ft["dtype"] == "video"]
-            videos_info = [
-                {
-                    "url": f"https://huggingface.co/datasets/{repo_id}/resolve/main/"
-                    + dataset.video_path.format(
-                        episode_chunk=int(episode_id) // dataset.chunks_size,
-                        video_key=video_key,
-                        episode_index=episode_id,
-                    ),
-                    "filename": video_key,
-                    "start_time": None,
-                    "end_time": None,
-                    "is_chunked": False,
-                }
-                for video_key in video_keys
-            ]
-
-            response = requests.get(
-                f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/episodes.jsonl", timeout=5
-            )
-            response.raise_for_status()
-            # Split into lines and parse each line as JSON
-            tasks_jsonl = [json.loads(line) for line in response.text.splitlines() if line.strip()]
-
-            filtered_tasks_jsonl = [row for row in tasks_jsonl if row["episode_index"] == episode_id]
-            tasks = filtered_tasks_jsonl[0]["tasks"]
-
-        videos_info[0]["language_instruction"] = tasks
-
-        if episodes is None:
-            episodes = list(
-                range(dataset.num_episodes if isinstance(dataset, LeRobotDataset) else dataset.total_episodes)
-            )
-
-        return render_template(
-            "visualize_dataset_template.html",
-            episode_id=episode_id,
-            episodes=episodes,
-            dataset_info=dataset_info,
-            videos_info=videos_info,
-            episode_data_csv_str=episode_data_csv_str,
-            columns=columns,
-            ignored_columns=ignored_columns,
-        )
-
-    app.run(host=host, port=port)
-
-
-def get_ep_csv_fname(episode_id: int):
-    ep_csv_fname = f"episode_{episode_id}.csv"
-    return ep_csv_fname
-
-
-def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index):
-    """Get a csv str containing timeseries data of an episode (e.g. state and action).
-    This file will be loaded by Dygraph javascript to plot data in real time."""
-    columns = []
-
-    selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] in ["float32", "int32"]]
-    selected_columns.remove("timestamp")
-
-    ignored_columns = []
-    for column_name in selected_columns:
-        shape = dataset.features[column_name]["shape"]
-        shape_dim = len(shape)
-        if shape_dim > 1:
-            selected_columns.remove(column_name)
-            ignored_columns.append(column_name)
-
-    # init header of csv with state and action names
-    header = ["timestamp"]
-
-    for column_name in selected_columns:
-        dim_state = (
-            dataset.meta.shapes[column_name][0]
-            if isinstance(dataset, LeRobotDataset)
-            else dataset.features[column_name].shape[0]
-        )
-
-        if "names" in dataset.features[column_name] and dataset.features[column_name]["names"]:
-            column_names = dataset.features[column_name]["names"]
-            while not isinstance(column_names, list):
-                column_names = list(column_names.values())[0]
-        else:
-            column_names = [f"{column_name}_{i}" for i in range(dim_state)]
-        columns.append({"key": column_name, "value": column_names})
-
-        header += column_names
-
-    selected_columns.insert(0, "timestamp")
-
-    if isinstance(dataset, LeRobotDataset):
-        from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
-        to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
-        data = (
-            dataset.hf_dataset.select(range(from_idx, to_idx))
-            .select_columns(selected_columns)
-            .with_format("pandas")
-        )
-    else:
-        repo_id = dataset.repo_id
-
-        url = f"https://huggingface.co/datasets/{repo_id}/resolve/main/" + dataset.data_path.format(
-            episode_chunk=int(episode_index) // dataset.chunks_size, episode_index=episode_index
-        )
-        df = pd.read_parquet(url)
-        data = df[selected_columns]  # Select specific columns
-
-    rows = np.hstack(
-        (
-            np.expand_dims(data["timestamp"], axis=1),
-            *[np.vstack(data[col]) for col in selected_columns[1:]],
-        )
-    ).tolist()
-
-    # Convert data to CSV string
-    csv_buffer = StringIO()
-    csv_writer = csv.writer(csv_buffer)
-    # Write header
-    csv_writer.writerow(header)
-    # Write data rows
-    csv_writer.writerows(rows)
-    csv_string = csv_buffer.getvalue()
-
-    return csv_string, columns, ignored_columns
-
-
-def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]:
-    # get first frame of episode (hack to get video_path of the episode)
-    first_frame_idx = dataset.meta.episodes["dataset_from_index"][ep_index]
-    return [
-        dataset.hf_dataset.select_columns(key)[first_frame_idx][key]["path"]
-        for key in dataset.meta.video_keys
-    ]
-
-
-def get_episode_language_instruction(dataset: LeRobotDataset, ep_index: int) -> list[str]:
-    # check if the dataset has language instructions
-    if "language_instruction" not in dataset.features:
-        return None
-
-    # get first frame index
-    first_frame_idx = dataset.meta.episodes["dataset_from_index"][ep_index]
-
-    language_instruction = dataset.hf_dataset[first_frame_idx]["language_instruction"]
-    # TODO (michel-aractingi) hack to get the sentence, some strings in openx are badly stored
-    # with the tf.tensor appearing in the string
-    return language_instruction.removeprefix("tf.Tensor(b'").removesuffix("', shape=(), dtype=string)")
-
-
-def get_dataset_info(repo_id: str) -> IterableNamespace:
-    response = requests.get(
-        f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json", timeout=5
-    )
-    response.raise_for_status()  # Raises an HTTPError for bad responses
-    dataset_info = response.json()
-    dataset_info["repo_id"] = repo_id
-    return IterableNamespace(dataset_info)
-
-
-def visualize_dataset_html(
-    dataset: LeRobotDataset | None,
-    episodes: list[int] | None = None,
-    output_dir: Path | None = None,
-    serve: bool = True,
-    host: str = "127.0.0.1",
-    port: int = 9090,
-    force_override: bool = False,
-) -> Path | None:
-    init_logging()
-
-    template_dir = Path(__file__).resolve().parent.parent / "templates"
-
-    if output_dir is None:
-        # Create a temporary directory that will be automatically cleaned up
-        output_dir = tempfile.mkdtemp(prefix="lerobot_visualize_dataset_")
-
-    output_dir = Path(output_dir)
-    if output_dir.exists():
-        if force_override:
-            shutil.rmtree(output_dir)
-        else:
-            logging.info(f"Output directory already exists. Loading from it: '{output_dir}'")
-
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    static_dir = output_dir / "static"
-    static_dir.mkdir(parents=True, exist_ok=True)
-
-    if dataset is None:
-        if serve:
-            run_server(
-                dataset=None,
-                episodes=None,
-                host=host,
-                port=port,
-                static_folder=static_dir,
-                template_folder=template_dir,
-            )
-    else:
-        # Create a simlink from the dataset video folder containing mp4 files to the output directory
-        # so that the http server can get access to the mp4 files.
-        if isinstance(dataset, LeRobotDataset):
-            ln_videos_dir = static_dir / "videos"
-            if not ln_videos_dir.exists():
-                ln_videos_dir.symlink_to((dataset.root / "videos").resolve().as_posix())
-
-        if serve:
-            run_server(dataset, episodes, host, port, static_dir, template_dir)
-
-
-def main():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        default=None,
-        help="Name of hugging face repositery containing a LeRobotDataset dataset (e.g. `lerobot/pusht` for https://huggingface.co/datasets/lerobot/pusht).",
-    )
-    parser.add_argument(
-        "--root",
-        type=Path,
-        default=None,
-        help="Root directory for a dataset stored locally (e.g. `--root data`). By default, the dataset will be loaded from hugging face cache folder, or downloaded from the hub if available.",
-    )
-    parser.add_argument(
-        "--load-from-hf-hub",
-        type=int,
-        default=0,
-        help="Load videos and parquet files from HF Hub rather than local system.",
-    )
-    parser.add_argument(
-        "--episodes",
-        type=int,
-        nargs="*",
-        default=None,
-        help="Episode indices to visualize (e.g. `0 1 5 6` to load episodes of index 0, 1, 5 and 6). By default loads all episodes.",
-    )
-    parser.add_argument(
-        "--output-dir",
-        type=Path,
-        default=None,
-        help="Directory path to write html files and kickoff a web server. By default write them to 'outputs/visualize_dataset/REPO_ID'.",
-    )
-    parser.add_argument(
-        "--serve",
-        type=int,
-        default=1,
-        help="Launch web server.",
-    )
-    parser.add_argument(
-        "--host",
-        type=str,
-        default="127.0.0.1",
-        help="Web host used by the http server.",
-    )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=9090,
-        help="Web port used by the http server.",
-    )
-    parser.add_argument(
-        "--force-override",
-        type=int,
-        default=0,
-        help="Delete the output directory if it exists already.",
-    )
-
-    parser.add_argument(
-        "--tolerance-s",
-        type=float,
-        default=1e-4,
-        help=(
-            "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
-            "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
-            "If not given, defaults to 1e-4."
-        ),
-    )
-
-    args = parser.parse_args()
-    kwargs = vars(args)
-    repo_id = kwargs.pop("repo_id")
-    load_from_hf_hub = kwargs.pop("load_from_hf_hub")
-    root = kwargs.pop("root")
-    tolerance_s = kwargs.pop("tolerance_s")
-
-    dataset = None
-    if repo_id:
-        dataset = (
-            LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
-            if not load_from_hf_hub
-            else get_dataset_info(repo_id)
-        )
-
-    visualize_dataset_html(dataset, **vars(args))
-
-
-if __name__ == "__main__":
-    main()
@@ -1,68 +0,0 @@
-<!DOCTYPE html>
-<html lang="en">
-<head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>Interactive Video Background Page</title>
-    <script src="https://cdn.tailwindcss.com"></script>
-    <script defer src="https://cdn.jsdelivr.net/npm/alpinejs@3.x.x/dist/cdn.min.js"></script>
-</head>
-<body class="h-screen overflow-hidden font-mono text-white" x-data="{
-    inputValue: '',
-    navigateToDataset() {
-        const trimmedValue = this.inputValue.trim();
-        if (trimmedValue) {
-            window.location.href = `/${trimmedValue}`;
-        }
-    }
-}">
-    <div class="fixed inset-0 w-full h-full overflow-hidden">
-        <video class="absolute min-w-full min-h-full w-auto h-auto top-1/2 left-1/2 transform -translate-x-1/2 -translate-y-1/2" autoplay muted loop>
-            <source src="https://huggingface.co/datasets/cadene/koch_bimanual_folding/resolve/v1.6/videos/observation.images.phone_episode_000037.mp4" type="video/mp4">
-            Your browser does not support HTML5 video.
-        </video>
-    </div>
-    <div class="fixed inset-0 bg-black bg-opacity-80"></div>
-    <div class="relative z-10 flex flex-col items-center justify-center h-screen">
-        <div class="text-center mb-8">
-            <h1 class="text-4xl font-bold mb-4">LeRobot Dataset Visualizer</h1>
-
-            <a href="https://x.com/RemiCadene/status/1825455895561859185" target="_blank" rel="noopener noreferrer" class="underline">create & train your own robots</a>
-
-            <p class="text-xl mb-4"></p>
-            <div class="text-left inline-block">
-                <h3 class="font-semibold mb-2 mt-4">Example Datasets:</h3>
-                <ul class="list-disc list-inside">
-                    {% for dataset in featured_datasets %}
-                        <li><a href="/{{ dataset }}" class="text-blue-300 hover:text-blue-100 hover:underline">{{ dataset }}</a></li>
-                    {% endfor %}
-                </ul>
-            </div>
-        </div>
-        <div class="flex w-full max-w-lg px-4 mb-4">
-            <input
-                type="text"
-                x-model="inputValue"
-                @keyup.enter="navigateToDataset"
-                placeholder="enter dataset id (ex: lerobot/droid_100)"
-                class="flex-grow px-4 py-2 rounded-l bg-white bg-opacity-20 text-white placeholder-gray-300 focus:outline-none focus:ring-2 focus:ring-blue-300"
-            >
-            <button
-                @click="navigateToDataset"
-                class="px-4 py-2 bg-blue-500 text-white rounded-r hover:bg-blue-600 focus:outline-none focus:ring-2 focus:ring-blue-300"
-            >
-                Go
-            </button>
-        </div>
-
-        <details class="mt-4 max-w-full px-4">
-            <summary>More example datasets</summary>
-            <ul class="list-disc list-inside max-h-28 overflow-y-auto break-all">
-                {% for dataset in lerobot_datasets %}
-                    <li><a href="/{{ dataset }}" class="text-blue-300 hover:text-blue-100 hover:underline">{{ dataset }}</a></li>
-                {% endfor %}
-            </ul>
-        </details>
-    </div>
-</body>
-</html>
@@ -1,546 +0,0 @@
-<!DOCTYPE html>
-<html lang="en">
-
-<head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <!-- # TODO(rcadene, mishig25): store the js files locally -->
-    <script src="https://cdnjs.cloudflare.com/ajax/libs/alpinejs/3.13.5/cdn.min.js" defer></script>
-    <script src="https://cdn.jsdelivr.net/npm/dygraphs@2.2.1/dist/dygraph.min.js" type="text/javascript"></script>
-    <script src="https://cdn.tailwindcss.com"></script>
-    <title>{{ dataset_info.repo_id }} episode {{ episode_id }}</title>
-</head>
-
-<!-- Use [Alpin.js](https://alpinejs.dev), a lightweight and easy to learn JS framework -->
-<!-- Use [tailwindcss](https://tailwindcss.com/), CSS classes for styling html -->
-<!-- Use [dygraphs](https://dygraphs.com/), a lightweight JS charting library -->
-<body class="flex flex-col md:flex-row h-screen max-h-screen bg-slate-950 text-gray-200" x-data="createAlpineData()">
-    <!-- Sidebar -->
-    <div x-ref="sidebar" class="bg-slate-900 p-5 break-words overflow-y-auto shrink-0 md:shrink md:w-60 md:max-h-screen">
-        <a href="https://github.com/huggingface/lerobot" target="_blank" class="hidden md:block">
-            <img src="https://github.com/huggingface/lerobot/raw/main/media/lerobot-logo-thumbnail.png">
-        </a>
-        <a href="https://huggingface.co/datasets/{{ dataset_info.repo_id }}" target="_blank">
-            <h1 class="mb-4 text-xl font-semibold">{{ dataset_info.repo_id }}</h1>
-        </a>
-
-        <ul>
-            <li>
-                Number of samples/frames: {{ dataset_info.num_samples }}
-            </li>
-            <li>
-                Number of episodes: {{ dataset_info.num_episodes }}
-            </li>
-            <li>
-                Frames per second: {{ dataset_info.fps }}
-            </li>
-        </ul>
-
-        <p>Episodes:</p>
-        <!-- episodes menu for medium & large screens -->
-        <div class="ml-2 hidden md:block" x-data="episodePagination">
-            <ul>
-                <template x-for="episode in paginatedEpisodes" :key="episode">
-                    <li class="font-mono text-sm mt-0.5">
-                        <a :href="'episode_' + episode"
-                           :class="{'underline': true, 'font-bold -ml-1': episode == {{ episode_id }}}"
-                           x-text="'Episode ' + episode"></a>
-                    </li>
-                </template>
-            </ul>
-
-            <div class="flex items-center mt-3 text-xs" x-show="totalPages > 1">
-                <button @click="prevPage()"
-                        class="px-2 py-1 bg-slate-800 rounded mr-2"
-                        :class="{'opacity-50 cursor-not-allowed': page === 1}"
-                        :disabled="page === 1">
-                    &laquo; Prev
-                </button>
-                <span class="font-mono mr-2" x-text="` ${page} / ${totalPages}`"></span>
-                <button @click="nextPage()"
-                        class="px-2 py-1 bg-slate-800 rounded"
-                        :class="{'opacity-50 cursor-not-allowed': page === totalPages}"
-                        :disabled="page === totalPages">
-                    Next &raquo;
-                </button>
-            </div>
-        </div>
-
-        <!-- episodes menu for small screens -->
-        <div class="flex overflow-x-auto md:hidden" x-data="episodePagination">
-            <button @click="prevPage()"
-                    class="px-2 bg-slate-800 rounded mr-2"
-                    :class="{'opacity-50 cursor-not-allowed': page === 1}"
-                    :disabled="page === 1">&laquo;</button>
-            <div class="flex">
-                <template x-for="(episode, index) in paginatedEpisodes" :key="episode">
-                    <p class="font-mono text-sm mt-0.5 px-2"
-                       :class="{
-                           'font-bold': episode == {{ episode_id }},
-                           'border-r': index !== paginatedEpisodes.length - 1
-                       }">
-                        <a :href="'episode_' + episode" x-text="episode"></a>
-                    </p>
-                </template>
-            </div>
-            <button @click="nextPage()"
-                    class="px-2 bg-slate-800 rounded ml-2"
-                    :class="{'opacity-50 cursor-not-allowed': page === totalPages}"
-                    :disabled="page === totalPages">&raquo; </button>
-        </div>
-
-    </div>
-
-    <!-- Toggle sidebar button -->
-    <button class="flex items-center opacity-50 hover:opacity-100 mx-1 hidden md:block"
-        @click="() => ($refs.sidebar.classList.toggle('hidden'))" title="Toggle sidebar">
-        <div class="bg-slate-500 w-2 h-10 rounded-full"></div>
-    </button>
-
-    <!-- Content -->
-    <div class="max-h-screen flex flex-col gap-4 overflow-y-auto md:flex-1">
-        <h1 class="text-xl font-bold mt-4 font-mono">
-            Episode {{ episode_id }}
-        </h1>
-
-        <!-- Error message -->
-        <div class="font-medium text-orange-700 hidden" :class="{ 'hidden': !videoCodecError }">
-            <p>Videos could NOT play because <a href="https://en.wikipedia.org/wiki/AV1" target="_blank" class="underline">AV1</a> decoding is not available on your browser.</p>
-            <ul class="list-decimal list-inside">
-                <li>If iPhone: <span class="italic">It is supported with A17 chip or higher.</span></li>
-                <li>If Mac with Safari: <span class="italic">It is supported on most browsers except Safari with M1 chip or higher and on Safari with M3 chip or higher.</span></li>
-                <li>Other: <span class="italic">Contact the maintainers on LeRobot discord channel:</span> <a href="https://discord.com/invite/s3KuuzsPFb" target="_blank" class="underline">https://discord.com/invite/s3KuuzsPFb</a></li>
-            </ul>
-        </div>
-
-        <!-- Videos -->
-        <div  class="max-w-32 relative text-sm mb-4 select-none"
-            @click.outside="isVideosDropdownOpen = false">
-            <div
-                @click="isVideosDropdownOpen = !isVideosDropdownOpen"
-                class="p-2 border border-slate-500 rounded flex justify-between items-center cursor-pointer"
-            >
-            <span class="truncate">filter videos</span>
-            <div class="transition-transform" :class="{ 'rotate-180': isVideosDropdownOpen }">🔽</div>
-            </div>
-
-            <div x-show="isVideosDropdownOpen"
-                class="absolute mt-1 border border-slate-500 rounded shadow-lg z-10">
-            <div>
-                <template x-for="option in videosKeys" :key="option">
-                <div
-                    @click="videosKeysSelected = videosKeysSelected.includes(option) ? videosKeysSelected.filter(v => v !== option) : [...videosKeysSelected, option]"
-                    class="p-2 cursor-pointer bg-slate-900"
-                    :class="{ 'bg-slate-700': videosKeysSelected.includes(option) }"
-                    x-text="option"
-                ></div>
-                </template>
-            </div>
-            </div>
-        </div>
-
-        <div class="flex flex-wrap gap-x-2 gap-y-6">
-            {% for video_info in videos_info %}
-            <div x-show="!videoCodecError && videosKeysSelected.includes('{{ video_info.filename }}')" class="max-w-96 relative">
-                <p class="absolute inset-x-0 -top-4 text-sm text-gray-300 bg-gray-800 px-2 rounded-t-xl truncate">{{ video_info.filename }}</p>
-                <video muted loop type="video/mp4" class="object-contain w-full h-full" @canplaythrough="videoCanPlay" @timeupdate="() => {
-                    if (video.duration) {
-                      const time = video.currentTime;
-                      const pc = (100 / video.duration) * time;
-                      $refs.slider.value = pc;
-                      dygraphTime = time;
-                      dygraphIndex = Math.floor(pc * dygraph.numRows() / 100);
-                      dygraph.setSelection(dygraphIndex, undefined, true, true);
-
-                      $refs.timer.textContent = formatTime(time) + ' / ' + formatTime(video.duration);
-
-                      updateTimeQuery(time.toFixed(2));
-                    }
-                }" @ended="() => {
-                    $refs.btnPlay.classList.remove('hidden');
-                    $refs.btnPause.classList.add('hidden');
-                }"
-                    @loadedmetadata="() => ($refs.timer.textContent = formatTime(0) + ' / ' + formatTime(video.duration))">
-                    <source src="{{ video_info.url }}">
-                    Your browser does not support the video tag.
-                </video>
-            </div>
-            {% endfor %}
-        </div>
-
-        <!-- Language instruction -->
-        {% if videos_info[0].language_instruction %}
-        <p class="font-medium mt-2">
-            Language Instruction: <span class="italic">{{ videos_info[0].language_instruction }}</span>
-        </p>
-        {% endif %}
-
-        <!-- Shortcuts info -->
-        <div class="text-sm hidden md:block">
-            Hotkeys: <span class="font-mono">Space</span> to pause/unpause, <span class="font-mono">Arrow Down</span> to go to next episode, <span class="font-mono">Arrow Up</span> to go to previous episode.
-        </div>
-
-        <!-- Controllers -->
-        <div class="flex gap-1 text-3xl items-center">
-            <button x-ref="btnPlay" class="-rotate-90" class="-rotate-90" title="Play. Toggle with Space" @click="() => {
-                videos.forEach(video => video.play());
-                $refs.btnPlay.classList.toggle('hidden');
-                $refs.btnPause.classList.toggle('hidden');
-            }">🔽</button>
-            <button x-ref="btnPause" class="hidden" title="Pause. Toggle with Space" @click="() => {
-                videos.forEach(video => video.pause());
-                $refs.btnPlay.classList.toggle('hidden');
-                $refs.btnPause.classList.toggle('hidden');
-            }">⏸️</button>
-            <button title="Jump backward 5 seconds"
-                @click="() => (videos.forEach(video => (video.currentTime -= 5)))">⏪</button>
-            <button title="Jump forward 5 seconds"
-                @click="() => (videos.forEach(video => (video.currentTime += 5)))">⏩</button>
-            <button title="Rewind from start"
-                @click="() => (videos.forEach(video => (video.currentTime = 0.0)))">↩️</button>
-            <input x-ref="slider" max="100" min="0" step="1" type="range" value="0" class="w-80 mx-2" @input="() => {
-                const sliderValue = $refs.slider.value;
-                videos.forEach(video => {
-                    const time = (video.duration * sliderValue) / 100;
-                    video.currentTime = time;
-                });
-            }" />
-            <div x-ref="timer" class="font-mono text-sm border border-slate-500 rounded-lg px-1 py-0.5 shrink-0">0:00 /
-                0:00
-            </div>
-        </div>
-
-        <!-- Graph -->
-        <div class="flex gap-2 mb-4 flex-wrap">
-            <div>
-                <div id="graph" @mouseleave="() => {
-                    dygraph.setSelection(dygraphIndex, undefined, true, true);
-                    dygraphTime = video.currentTime;
-                }">
-                </div>
-                <p x-ref="graphTimer" class="font-mono ml-14 mt-4"
-                    x-init="$watch('dygraphTime', value => ($refs.graphTimer.innerText = `Time: ${dygraphTime.toFixed(2)}s`))">
-                    Time: 0.00s
-                </p>
-            </div>
-
-            <div>
-                <table class="text-sm border-collapse border border-slate-700" x-show="currentFrameData">
-                    <thead>
-                        <tr>
-                            <th></th>
-                            <template x-for="(_, colIndex) in Array.from({length: columns.length}, (_, index) => index)">
-                                <th class="border border-slate-700">
-                                    <div class="flex gap-x-2 justify-between px-2">
-                                        <input type="checkbox" :checked="isColumnChecked(colIndex)"
-                                            @change="toggleColumn(colIndex)">
-                                        <p x-text="`${columns[colIndex].key}`"></p>
-                                    </div>
-                                </th>
-                            </template>
-                        </tr>
-                    </thead>
-                    <tbody>
-                        <template x-for="(row, rowIndex) in rows">
-                            <tr class="odd:bg-gray-800 even:bg-gray-900">
-                                <td class="border border-slate-700">
-                                    <div class="flex gap-x-2 max-w-64 font-semibold px-1 break-all">
-                                        <input type="checkbox" :checked="isRowChecked(rowIndex)"
-                                            @change="toggleRow(rowIndex)">
-                                    </div>
-                                </td>
-                                <template x-for="(cell, colIndex) in row">
-                                    <td x-show="cell" class="border border-slate-700">
-                                        <div class="flex gap-x-2 justify-between px-2" :class="{ 'hidden': cell.isNull }">
-                                            <div class="flex gap-x-2">
-                                                <input type="checkbox" x-model="cell.checked" @change="updateTableValues()">
-                                                <span x-text="`${!cell.isNull ? cell.label : null}`"></span>
-                                            </div>
-                                            <span class="w-14 text-right" x-text="`${!cell.isNull ? (typeof cell.value === 'number' ? cell.value.toFixed(2) : cell.value) : null}`"
-                                                :style="`color: ${cell.color}`"></span>
-                                        </div>
-                                    </td>
-                                </template>
-                            </tr>
-                        </template>
-                    </tbody>
-                </table>
-
-                <div id="labels" class="hidden">
-                </div>
-
-                {% if ignored_columns|length > 0 %}
-                <div class="m-2 text-orange-700 max-w-96">
-                    Columns {{ ignored_columns }} are NOT shown since the visualizer currently does not support 2D or 3D data.
-                </div>
-                {% endif %}
-            </div>
-
-        </div>
-    </div>
-
-    <script>
-        const parentOrigin = "https://huggingface.co";
-        const searchParams = new URLSearchParams();
-        searchParams.set("dataset", "{{ dataset_info.repo_id }}");
-        searchParams.set("episode", "{{ episode_id }}");
-		window.parent.postMessage({ queryString: searchParams.toString() }, parentOrigin);
-    </script>
-
-    <script>
-        function createAlpineData() {
-            return {
-                // state
-                dygraph: null,
-                currentFrameData: null,
-                checked: [],
-                dygraphTime: 0.0,
-                dygraphIndex: 0,
-                videos: null,
-                video: null,
-                colors: null,
-                nVideos: {{ videos_info | length }},
-                nVideoReadyToPlay: 0,
-                videoCodecError: false,
-                isVideosDropdownOpen: false,
-                videosKeys: {{ videos_info | map(attribute='filename') | list | tojson }},
-                videosKeysSelected: [],
-                columns: {{ columns | tojson }},
-
-                // alpine initialization
-                init() {
-                    // check if videos can play
-                    const dummyVideo = document.createElement('video');
-                    const canPlayVideos = dummyVideo.canPlayType('video/mp4; codecs="av01.0.05M.08"'); // codec source: https://huggingface.co/blog/video-encoding#results
-                    if(!canPlayVideos){
-                        this.videoCodecError = true;
-                    }
-                    this.videosKeysSelected = this.videosKeys.map(opt => opt)
-
-                    // process CSV data
-                    const csvDataStr = {{ episode_data_csv_str|tojson|safe }};
-                    // Create a Blob with the CSV data
-                    const blob = new Blob([csvDataStr], { type: 'text/csv;charset=utf-8;' });
-                    // Create a URL for the Blob
-                    const csvUrl = URL.createObjectURL(blob);
-
-                    // process CSV data
-                    this.videos = document.querySelectorAll('video');
-                    this.video = this.videos[0];
-                    this.dygraph = new Dygraph(document.getElementById("graph"), csvUrl, {
-                        pixelsPerPoint: 0.01,
-                        legend: 'always',
-                        labelsDiv: document.getElementById('labels'),
-                        labelsKMB: true,
-                        strokeWidth: 1.5,
-                        pointClickCallback: (event, point) => {
-                            this.dygraphTime = point.xval;
-                            this.updateTableValues(this.dygraphTime);
-                        },
-                        highlightCallback: (event, x, points, row, seriesName) => {
-                            this.dygraphTime = x;
-                            this.updateTableValues(this.dygraphTime);
-                        },
-                        drawCallback: (dygraph, is_initial) => {
-                            if (is_initial) {
-                                // dygraph initialization
-                                this.dygraph.setSelection(this.dygraphIndex, undefined, true, true);
-                                this.colors = this.dygraph.getColors();
-                                this.checked = Array(this.colors.length).fill(true);
-
-                                const colors = [];
-                                let lightness = 30; // const LIGHTNESS = [30, 65, 85]; // state_lightness, action_lightness, pred_action_lightness
-                                for(const column of this.columns){
-                                    const nValues = column.value.length;
-                                    for (let hue = 0; hue < 360; hue += parseInt(360/nValues)) {
-                                        const color = `hsl(${hue}, 100%, ${lightness}%)`;
-                                        colors.push(color);
-                                    }
-                                    lightness += 35;
-                                }
-
-                                this.dygraph.updateOptions({ colors });
-                                this.colors = colors;
-
-                                this.updateTableValues();
-
-                                let url = new URL(window.location.href);
-                                let params = new URLSearchParams(url.search);
-                                let time = params.get("t");
-                                if(time){
-                                    time = parseFloat(time);
-                                    this.videos.forEach(video => (video.currentTime = time));
-                                }
-                            }
-                        },
-                    });
-                },
-
-                //#region Table Data
-
-                // turn dygraph's 1D data (at a given time t) to 2D data that whose columns names are defined in this.columnNames.
-                // 2d data view is used to create html table element.
-                get rows() {
-                    if (!this.currentFrameData) {
-                        return [];
-                    }
-                    const rows = [];
-                    const nRows = Math.max(...this.columns.map(column => column.value.length));
-                    let rowIndex = 0;
-                    while(rowIndex < nRows){
-                        const row = [];
-                        // number of states may NOT match number of actions. In this case, we null-pad the 2D array to make a fully rectangular 2d array
-                        const nullCell = { isNull: true };
-                        // row consists of [state value, action value]
-                        let idx = rowIndex;
-                        for(const column of this.columns){
-                            const nColumn = column.value.length;
-                            row.push(rowIndex < nColumn ? this.currentFrameData[idx] : nullCell);
-                            idx += nColumn; // because this.currentFrameData = [state0, state1, ..., stateN, action0, action1, ..., actionN]
-                        }
-                        rowIndex += 1;
-                        rows.push(row);
-                    }
-                    return rows;
-                },
-                isRowChecked(rowIndex) {
-                    return this.rows[rowIndex].every(cell => cell && (cell.isNull || cell.checked));
-                },
-                isColumnChecked(colIndex) {
-                    return this.rows.every(row => row[colIndex] && (row[colIndex].isNull || row[colIndex].checked));
-                },
-                toggleRow(rowIndex) {
-                    const newState = !this.isRowChecked(rowIndex);
-                    this.rows[rowIndex].forEach(cell => {
-                        if (cell && !cell.isNull) cell.checked = newState;
-                    });
-                    this.updateTableValues();
-                },
-                toggleColumn(colIndex) {
-                    const newState = !this.isColumnChecked(colIndex);
-                    this.rows.forEach(row => {
-                        if (row[colIndex] && !row[colIndex].isNull) row[colIndex].checked = newState;
-                    });
-                    this.updateTableValues();
-                },
-
-                // given time t, update the values in the html table with "data[t]"
-                updateTableValues(time) {
-                    if (!this.colors) {
-                        return;
-                    }
-                    let pc = (100 / this.video.duration) * (time === undefined ? this.video.currentTime : time);
-                    if (isNaN(pc)) pc = 0;
-                    const index = Math.floor(pc * this.dygraph.numRows() / 100);
-                    // slice(1) to remove the timestamp point that we do not need
-                    const labels = this.dygraph.getLabels().slice(1);
-                    const values = this.dygraph.rawData_[index].slice(1);
-                    const checkedNew = this.currentFrameData ? this.currentFrameData.map(cell => cell.checked) : Array(
-                        this.colors.length).fill(true);
-                    this.currentFrameData = labels.map((label, idx) => ({
-                        label,
-                        value: values[idx],
-                        color: this.colors[idx],
-                        checked: checkedNew[idx],
-                    }));
-                    const shouldUpdateVisibility = !this.checked.every((value, index) => value === checkedNew[index]);
-                    if (shouldUpdateVisibility) {
-                        this.checked = checkedNew;
-                        this.dygraph.setVisibility(this.checked);
-                    }
-                },
-
-                //#endregion
-
-                updateTimeQuery(time) {
-                    let url = new URL(window.location.href);
-                    let params = new URLSearchParams(url.search);
-                    params.set("t", time);
-                    url.search = params.toString();
-                    window.history.replaceState({}, '', url.toString());
-                },
-
-                formatTime(time) {
-                    var hours = Math.floor(time / 3600);
-                    var minutes = Math.floor((time % 3600) / 60);
-                    var seconds = Math.floor(time % 60);
-                    return (hours > 0 ? hours + ':' : '') + (minutes < 10 ? '0' + minutes : minutes) + ':' + (seconds <
-                        10 ?
-                        '0' + seconds : seconds);
-                },
-
-                videoCanPlay() {
-                    this.nVideoReadyToPlay += 1;
-                    if(this.nVideoReadyToPlay == this.nVideos) {
-                        // start autoplay all videos in sync
-                        this.$refs.btnPlay.click();
-                    }
-                }
-            };
-        }
-
-        document.addEventListener('alpine:init', () => {
-            // Episode pagination component
-            Alpine.data('episodePagination', () => ({
-                episodes: {{ episodes }},
-                pageSize: 100,
-                page: 1,
-
-                init() {
-                    // Find which page contains the current episode_id
-                    const currentEpisodeId = {{ episode_id }};
-                    const episodeIndex = this.episodes.indexOf(currentEpisodeId);
-                    if (episodeIndex !== -1) {
-                        this.page = Math.floor(episodeIndex / this.pageSize) + 1;
-                    }
-                },
-
-                get totalPages() {
-                    return Math.ceil(this.episodes.length / this.pageSize);
-                },
-
-                get paginatedEpisodes() {
-                    const start = (this.page - 1) * this.pageSize;
-                    const end = start + this.pageSize;
-                    return this.episodes.slice(start, end);
-                },
-
-                nextPage() {
-                    if (this.page < this.totalPages) {
-                        this.page++;
-                    }
-                },
-
-                prevPage() {
-                    if (this.page > 1) {
-                        this.page--;
-                    }
-                }
-            }));
-        });
-    </script>
-
-    <script>
-        window.addEventListener('keydown', (e) => {
-            // Use the space bar to play and pause, instead of default action (e.g. scrolling)
-            const { keyCode, key } = e;
-
-            if (keyCode === 32 || key === ' ') {
-                e.preventDefault();
-                const btnPause = document.querySelector('[x-ref="btnPause"]');
-                const btnPlay = document.querySelector('[x-ref="btnPlay"]');
-                btnPause.classList.contains('hidden') ? btnPlay.click() : btnPause.click();
-            } else if (key === 'ArrowDown' || key === 'ArrowUp') {
-                const episodes = {{ episodes }};  // Access episodes directly from the Jinja template
-                const nextEpisodeId = key === 'ArrowDown' ? {{ episode_id }} + 1 : {{ episode_id }} - 1;
-                const lowestEpisodeId = episodes.at(0);
-                const highestEpisodeId = episodes.at(-1);
-                if (nextEpisodeId >= lowestEpisodeId && nextEpisodeId <= highestEpisodeId) {
-                    window.location.href = `./episode_${nextEpisodeId}`;
-                }
-            }
-        });
-    </script>
-</body>
-
-</html>
@@ -47,23 +47,19 @@ def save_dataset_to_safetensors(output_dir, repo_id="lerobot/pusht"):
    )

    # save 2 first frames of first episode
-    i = dataset.meta.episodes["dataset_from_index"][0].item()
+    i = dataset.meta.episodes["dataset_from_index"][0]
    save_file(dataset[i], repo_dir / f"frame_{i}.safetensors")
    save_file(dataset[i + 1], repo_dir / f"frame_{i + 1}.safetensors")

    # save 2 frames at the middle of first episode
    i = int(
-        (
-            dataset.meta.episodes["dataset_to_index"][0].item()
-            - dataset.meta.episodes["dataset_from_index"][0].item()
-        )
-        / 2
+        (dataset.meta.episodes["dataset_to_index"][0] - dataset.meta.episodes["dataset_from_index"][0]) / 2
    )
    save_file(dataset[i], repo_dir / f"frame_{i}.safetensors")
    save_file(dataset[i + 1], repo_dir / f"frame_{i + 1}.safetensors")

    # save 2 last frames of first episode
-    i = dataset.meta.episodes["dataset_to_index"][0].item()
+    i = dataset.meta.episodes["dataset_to_index"][0]
    save_file(dataset[i - 2], repo_dir / f"frame_{i - 2}.safetensors")
    save_file(dataset[i - 1], repo_dir / f"frame_{i - 1}.safetensors")

@@ -1,584 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tests for dataset tools utilities."""
-
-from unittest.mock import patch
-
-import numpy as np
-import pytest
-import torch
-
-from lerobot.datasets.dataset_tools import (
-    add_feature,
-    delete_episodes,
-    merge_datasets,
-    remove_feature,
-    split_dataset,
-)
-
-
-@pytest.fixture
-def sample_dataset(tmp_path, empty_lerobot_dataset_factory):
-    """Create a sample dataset for testing."""
-    # Create an empty dataset and add data manually
-    features = {
-        "action": {"dtype": "float32", "shape": (6,), "names": None},
-        "observation.state": {"dtype": "float32", "shape": (4,), "names": None},
-        "observation.images.top": {"dtype": "image", "shape": (224, 224, 3), "names": None},
-    }
-
-    dataset = empty_lerobot_dataset_factory(
-        root=tmp_path / "test_dataset",
-        features=features,
-    )
-
-    # Add episodes manually
-    for ep_idx in range(5):
-        for _ in range(10):
-            frame = {
-                "action": np.random.randn(6).astype(np.float32),
-                "observation.state": np.random.randn(4).astype(np.float32),
-                "observation.images.top": np.random.randint(0, 255, size=(224, 224, 3), dtype=np.uint8),
-                "task": f"task_{ep_idx % 2}",
-            }
-            dataset.add_frame(frame)
-        dataset.save_episode()
-
-    return dataset
-
-
-class TestDeleteEpisodes:
-    def test_delete_single_episode(self, sample_dataset, tmp_path):
-        """Test deleting a single episode."""
-        output_dir = tmp_path / "filtered"
-
-        # Delete episode 2
-        # Mock the revision check and snapshot_download to prevent Hub calls
-        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)
-
-            new_dataset = delete_episodes(
-                sample_dataset,
-                episode_indices=[2],
-                output_dir=output_dir,
-            )
-
-        # Check results
-        assert new_dataset.meta.total_episodes == 4
-        assert new_dataset.meta.total_frames == 40
-
-        # Check episode indices are renumbered
-        episode_indices = {int(idx.item()) for idx in new_dataset.hf_dataset["episode_index"]}
-        assert episode_indices == {0, 1, 2, 3}
-
-        # Check data integrity
-        assert len(new_dataset) == 40
-
-    def test_delete_multiple_episodes(self, sample_dataset, tmp_path):
-        """Test deleting multiple episodes."""
-        output_dir = tmp_path / "filtered"
-
-        # Delete episodes 1 and 3
-        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)
-
-            new_dataset = delete_episodes(
-                sample_dataset,
-                episode_indices=[1, 3],
-                output_dir=output_dir,
-            )
-
-        # Check results
-        assert new_dataset.meta.total_episodes == 3
-        assert new_dataset.meta.total_frames == 30
-
-        # Check episode indices
-        episode_indices = {int(idx.item()) for idx in new_dataset.hf_dataset["episode_index"]}
-        assert episode_indices == {0, 1, 2}
-
-    def test_delete_invalid_episodes(self, sample_dataset, tmp_path):
-        """Test error handling for invalid episode indices."""
-        with pytest.raises(ValueError, match="Invalid episode indices"):
-            delete_episodes(
-                sample_dataset,
-                episode_indices=[10, 20],  # Out of range
-                output_dir=tmp_path / "filtered",
-            )
-
-    def test_delete_all_episodes(self, sample_dataset, tmp_path):
-        """Test error when trying to delete all episodes."""
-        with pytest.raises(ValueError, match="Cannot delete all episodes"):
-            delete_episodes(
-                sample_dataset,
-                episode_indices=list(range(5)),  # All episodes
-                output_dir=tmp_path / "filtered",
-            )
-
-    def test_delete_empty_list(self, sample_dataset, tmp_path):
-        """Test error when no episodes specified."""
-        with pytest.raises(ValueError, match="No episodes to delete"):
-            delete_episodes(
-                sample_dataset,
-                episode_indices=[],
-                output_dir=tmp_path / "filtered",
-            )
-
-
-class TestSplitDataset:
-    def test_split_by_episodes(self, sample_dataset, tmp_path):
-        """Test splitting dataset by specific episode indices."""
-        splits = {
-            "train": [0, 1, 2],
-            "val": [3, 4],
-        }
-
-        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 to return the appropriate directory for each split
-            def mock_snapshot(repo_id, **kwargs):
-                if "train" in repo_id:
-                    return str(tmp_path / f"{sample_dataset.repo_id}_train")
-                elif "val" in repo_id:
-                    return str(tmp_path / f"{sample_dataset.repo_id}_val")
-                return str(kwargs.get("local_dir", tmp_path))
-
-            mock_snapshot_download.side_effect = mock_snapshot
-
-            result = split_dataset(
-                sample_dataset,
-                splits=splits,
-                output_dir=tmp_path,
-            )
-
-        # Check we got both splits
-        assert set(result.keys()) == {"train", "val"}
-
-        # Check train split
-        assert result["train"].meta.total_episodes == 3
-        assert result["train"].meta.total_frames == 30
-
-        # Check val split
-        assert result["val"].meta.total_episodes == 2
-        assert result["val"].meta.total_frames == 20
-
-        # Check episode renumbering
-        train_episodes = {int(idx.item()) for idx in result["train"].hf_dataset["episode_index"]}
-        assert train_episodes == {0, 1, 2}
-
-        val_episodes = {int(idx.item()) for idx in result["val"].hf_dataset["episode_index"]}
-        assert val_episodes == {0, 1}
-
-    def test_split_by_fractions(self, sample_dataset, tmp_path):
-        """Test splitting dataset by fractions."""
-        splits = {
-            "train": 0.6,  # 3 episodes
-            "val": 0.4,  # 2 episodes
-        }
-
-        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"
-
-            def mock_snapshot(repo_id, **kwargs):
-                for split_name in splits:
-                    if split_name in repo_id:
-                        return str(tmp_path / f"{sample_dataset.repo_id}_{split_name}")
-                return str(kwargs.get("local_dir", tmp_path))
-
-            mock_snapshot_download.side_effect = mock_snapshot
-
-            result = split_dataset(
-                sample_dataset,
-                splits=splits,
-                output_dir=tmp_path,
-            )
-
-        # Check splits
-        assert result["train"].meta.total_episodes == 3
-        assert result["val"].meta.total_episodes == 2
-
-    def test_split_overlapping_episodes(self, sample_dataset, tmp_path):
-        """Test error when episodes appear in multiple splits."""
-        splits = {
-            "train": [0, 1, 2],
-            "val": [2, 3, 4],  # Episode 2 appears in both
-        }
-
-        with pytest.raises(ValueError, match="Episodes cannot appear in multiple splits"):
-            split_dataset(sample_dataset, splits=splits, output_dir=tmp_path)
-
-    def test_split_invalid_fractions(self, sample_dataset, tmp_path):
-        """Test error when fractions sum to more than 1."""
-        splits = {
-            "train": 0.7,
-            "val": 0.5,  # Sum = 1.2
-        }
-
-        with pytest.raises(ValueError, match="Split fractions must sum to <= 1.0"):
-            split_dataset(sample_dataset, splits=splits, output_dir=tmp_path)
-
-    def test_split_empty(self, sample_dataset, tmp_path):
-        """Test error with empty splits."""
-        with pytest.raises(ValueError, match="No splits provided"):
-            split_dataset(sample_dataset, splits={}, output_dir=tmp_path)
-
-
-class TestMergeDatasets:
-    def test_merge_two_datasets(self, sample_dataset, tmp_path, empty_lerobot_dataset_factory):
-        """Test merging two datasets."""
-        # Create a second dataset manually
-        features = {
-            "action": {"dtype": "float32", "shape": (6,), "names": None},
-            "observation.state": {"dtype": "float32", "shape": (4,), "names": None},
-            "observation.images.top": {"dtype": "image", "shape": (224, 224, 3), "names": None},
-        }
-
-        dataset2 = empty_lerobot_dataset_factory(
-            root=tmp_path / "test_dataset2",
-            features=features,
-        )
-
-        # Add 3 episodes
-        for ep_idx in range(3):
-            for _ in range(10):
-                frame = {
-                    "action": np.random.randn(6).astype(np.float32),
-                    "observation.state": np.random.randn(4).astype(np.float32),
-                    "observation.images.top": np.random.randint(0, 255, size=(224, 224, 3), dtype=np.uint8),
-                    "task": f"task_{ep_idx % 2}",
-                }
-                dataset2.add_frame(frame)
-            dataset2.save_episode()
-
-        # Merge datasets
-        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(tmp_path / "merged_dataset")
-
-            merged = merge_datasets(
-                [sample_dataset, dataset2],
-                output_repo_id="merged_dataset",
-                output_dir=tmp_path / "merged_dataset",
-            )
-
-        # Check results
-        assert merged.meta.total_episodes == 8  # 5 + 3
-        assert merged.meta.total_frames == 80  # 50 + 30
-
-        # Check episode indices are sequential
-        episode_indices = sorted({int(idx.item()) for idx in merged.hf_dataset["episode_index"]})
-        assert episode_indices == list(range(8))
-
-    def test_merge_empty_list(self, tmp_path):
-        """Test error when merging empty list."""
-        with pytest.raises(ValueError, match="No datasets to merge"):
-            merge_datasets([], output_repo_id="merged", output_dir=tmp_path)
-
-
-class TestAddFeature:
-    def test_add_feature_with_values(self, sample_dataset, tmp_path):
-        """Test adding a feature with pre-computed values."""
-        # Create reward values for all frames
-        num_frames = sample_dataset.meta.total_frames
-        reward_values = np.random.randn(num_frames, 1).astype(np.float32)
-
-        feature_info = {
-            "dtype": "float32",
-            "shape": (1,),
-            "names": None,
-        }
-
-        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(tmp_path / "with_reward")
-
-            new_dataset = add_feature(
-                sample_dataset,
-                feature_name="reward",
-                feature_values=reward_values,
-                feature_info=feature_info,
-                output_dir=tmp_path / "with_reward",
-            )
-
-        # Check feature was added
-        assert "reward" in new_dataset.meta.features
-        assert new_dataset.meta.features["reward"] == feature_info
-
-        # Check values
-        assert len(new_dataset) == num_frames
-        sample_item = new_dataset[0]
-        assert "reward" in sample_item
-        # Scalar features don't have shape, just check it's a tensor
-        assert isinstance(sample_item["reward"], torch.Tensor)
-
-    def test_add_feature_with_callable(self, sample_dataset, tmp_path):
-        """Test adding a feature with a callable."""
-
-        def compute_reward(frame_dict, episode_idx, frame_idx):
-            # Simple reward based on episode and frame indices
-            return float(episode_idx * 10 + frame_idx)
-
-        feature_info = {
-            "dtype": "float32",
-            "shape": (1,),
-            "names": None,
-        }
-
-        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(tmp_path / "with_reward")
-
-            new_dataset = add_feature(
-                sample_dataset,
-                feature_name="reward",
-                feature_values=compute_reward,
-                feature_info=feature_info,
-                output_dir=tmp_path / "with_reward",
-            )
-
-        # Check feature was added
-        assert "reward" in new_dataset.meta.features
-
-        # Check computed values
-        # Episode 0, frame 0 should have reward 0
-        items = [new_dataset[i] for i in range(10)]
-        first_episode_items = [item for item in items if item["episode_index"] == 0]
-        assert len(first_episode_items) == 10
-
-        # Check first frame of first episode
-        first_frame = first_episode_items[0]
-        assert first_frame["frame_index"] == 0
-        assert float(first_frame["reward"]) == 0.0
-
-    def test_add_existing_feature(self, sample_dataset, tmp_path):
-        """Test error when adding an existing feature."""
-        feature_info = {"dtype": "float32", "shape": (1,)}
-
-        with pytest.raises(ValueError, match="Feature 'action' already exists"):
-            add_feature(
-                sample_dataset,
-                feature_name="action",  # Already exists
-                feature_values=np.zeros(50),
-                feature_info=feature_info,
-                output_dir=tmp_path / "modified",
-            )
-
-    def test_add_feature_invalid_info(self, sample_dataset, tmp_path):
-        """Test error with invalid feature info."""
-        with pytest.raises(ValueError, match="feature_info must contain keys"):
-            add_feature(
-                sample_dataset,
-                feature_name="reward",
-                feature_values=np.zeros(50),
-                feature_info={"dtype": "float32"},  # Missing 'shape'
-                output_dir=tmp_path / "modified",
-            )
-
-
-class TestRemoveFeature:
-    def test_remove_single_feature(self, sample_dataset, tmp_path):
-        """Test removing a single feature."""
-        # First add a feature to remove
-        feature_info = {"dtype": "float32", "shape": (1,), "names": None}
-
-        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.side_effect = lambda repo_id, **kwargs: str(
-                kwargs.get("local_dir", tmp_path)
-            )
-
-            dataset_with_reward = add_feature(
-                sample_dataset,
-                feature_name="reward",
-                feature_values=np.random.randn(50, 1).astype(np.float32),
-                feature_info=feature_info,
-                output_dir=tmp_path / "with_reward",
-            )
-
-            # Now remove it
-            dataset_without_reward = remove_feature(
-                dataset_with_reward,
-                feature_names="reward",
-                output_dir=tmp_path / "without_reward",
-            )
-
-        # Check feature was removed
-        assert "reward" not in dataset_without_reward.meta.features
-
-        # Check data
-        sample_item = dataset_without_reward[0]
-        assert "reward" not in sample_item
-
-    def test_remove_multiple_features(self, sample_dataset, tmp_path):
-        """Test removing multiple features at once."""
-        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.side_effect = lambda repo_id, **kwargs: str(
-                kwargs.get("local_dir", tmp_path)
-            )
-
-            # Add two features
-            dataset = sample_dataset
-            for feature_name in ["reward", "success"]:
-                feature_info = {"dtype": "float32", "shape": (1,), "names": None}
-                dataset = add_feature(
-                    dataset,
-                    feature_name=feature_name,
-                    feature_values=np.random.randn(dataset.meta.total_frames, 1).astype(np.float32),
-                    feature_info=feature_info,
-                    output_dir=tmp_path / f"with_{feature_name}",
-                )
-
-            # Remove both
-            dataset_clean = remove_feature(
-                dataset,
-                feature_names=["reward", "success"],
-                output_dir=tmp_path / "clean",
-            )
-
-        # Check both were removed
-        assert "reward" not in dataset_clean.meta.features
-        assert "success" not in dataset_clean.meta.features
-
-    def test_remove_nonexistent_feature(self, sample_dataset, tmp_path):
-        """Test error when removing non-existent feature."""
-        with pytest.raises(ValueError, match="Feature 'nonexistent' not found"):
-            remove_feature(
-                sample_dataset,
-                feature_names="nonexistent",
-                output_dir=tmp_path / "modified",
-            )
-
-    def test_remove_required_feature(self, sample_dataset, tmp_path):
-        """Test error when trying to remove required features."""
-        with pytest.raises(ValueError, match="Cannot remove required features"):
-            remove_feature(
-                sample_dataset,
-                feature_names="timestamp",  # Required feature
-                output_dir=tmp_path / "modified",
-            )
-
-    def test_remove_camera_feature(self, sample_dataset, tmp_path):
-        """Test removing a camera feature."""
-        camera_keys = sample_dataset.meta.camera_keys
-        if not camera_keys:
-            pytest.skip("No camera keys in dataset")
-
-        # Remove first camera
-        camera_to_remove = camera_keys[0]
-
-        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(tmp_path / "without_camera")
-
-            dataset_without_camera = remove_feature(
-                sample_dataset,
-                feature_names=camera_to_remove,
-                output_dir=tmp_path / "without_camera",
-            )
-
-        # Check camera was removed
-        assert camera_to_remove not in dataset_without_camera.meta.features
-        assert camera_to_remove not in dataset_without_camera.meta.camera_keys
-
-        # Check data
-        sample_item = dataset_without_camera[0]
-        assert camera_to_remove not in sample_item
-
-
-class TestIntegration:
-    def test_complex_workflow(self, sample_dataset, tmp_path):
-        """Test a complex workflow combining multiple operations."""
-        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.side_effect = lambda repo_id, **kwargs: str(
-                kwargs.get("local_dir", tmp_path)
-            )
-
-            # 1. Add a reward feature
-            dataset = add_feature(
-                sample_dataset,
-                feature_name="reward",
-                feature_values=np.random.randn(50, 1).astype(np.float32),
-                feature_info={"dtype": "float32", "shape": (1,), "names": None},
-                output_dir=tmp_path / "step1",
-            )
-
-            # 2. Delete an episode
-            dataset = delete_episodes(
-                dataset,
-                episode_indices=[2],
-                output_dir=tmp_path / "step2",
-            )
-
-            # 3. Split into train/val
-            splits = split_dataset(
-                dataset,
-                splits={"train": 0.75, "val": 0.25},
-                output_dir=tmp_path / "step3",
-            )
-
-            # 4. Merge them back
-            merged = merge_datasets(
-                list(splits.values()),
-                output_repo_id="final_dataset",
-                output_dir=tmp_path / "step4",
-            )
-
-        # Check final dataset
-        assert merged.meta.total_episodes == 4  # Started with 5, deleted 1
-        assert merged.meta.total_frames == 40
-        assert "reward" in merged.meta.features  # Feature preserved
-
-        # Check data integrity
-        assert len(merged) == 40
-        sample_item = merged[0]
-        assert "reward" in sample_item
@@ -35,7 +35,12 @@ from lerobot.datasets.lerobot_dataset import (
    MultiLeRobotDataset,
 )
 from lerobot.datasets.utils import (
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_DATA_FILE_SIZE_IN_MB,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
    create_branch,
+    get_hf_features_from_features,
+    hf_transform_to_torch,
    hw_to_dataset_features,
 )
 from lerobot.envs.factory import make_env_config
@@ -552,3 +557,235 @@ def test_create_branch():

    # Clean
    api.delete_repo(repo_id, repo_type=repo_type)
+
+
+def test_check_cached_episodes_sufficient(tmp_path, lerobot_dataset_factory):
+    """Test the _check_cached_episodes_sufficient method of LeRobotDataset."""
+    # Create a dataset with 5 episodes (0-4)
+    dataset = lerobot_dataset_factory(
+        root=tmp_path / "test",
+        total_episodes=5,
+        total_frames=200,
+        use_videos=False,
+    )
+
+    # Test hf_dataset is None
+    dataset.hf_dataset = None
+    assert dataset._check_cached_episodes_sufficient() is False
+
+    # Test hf_dataset is empty
+    import datasets
+
+    empty_features = get_hf_features_from_features(dataset.features)
+    dataset.hf_dataset = datasets.Dataset.from_dict(
+        {key: [] for key in empty_features}, features=empty_features
+    )
+    dataset.hf_dataset.set_transform(hf_transform_to_torch)
+    assert dataset._check_cached_episodes_sufficient() is False
+
+    # Restore the original dataset for remaining tests
+    dataset.hf_dataset = dataset.load_hf_dataset()
+
+    # Test all episodes requested (self.episodes = None) and all are available
+    dataset.episodes = None
+    assert dataset._check_cached_episodes_sufficient() is True
+
+    # Test specific episodes requested that are all available
+    dataset.episodes = [0, 2, 4]
+    assert dataset._check_cached_episodes_sufficient() is True
+
+    # Test request episodes that don't exist in the cached dataset
+    # Create a dataset with only episodes 0, 1, 2
+    limited_dataset = lerobot_dataset_factory(
+        root=tmp_path / "limited",
+        total_episodes=3,
+        total_frames=120,
+        use_videos=False,
+    )
+
+    # Request episodes that include non-existent ones
+    limited_dataset.episodes = [0, 1, 2, 3, 4]
+    assert limited_dataset._check_cached_episodes_sufficient() is False
+
+    # Test create a dataset with sparse episodes (e.g., only episodes 0, 2, 4)
+    # First create the full dataset structure
+    sparse_dataset = lerobot_dataset_factory(
+        root=tmp_path / "sparse",
+        total_episodes=5,
+        total_frames=200,
+        use_videos=False,
+    )
+
+    # Manually filter hf_dataset to only include episodes 0, 2, 4
+    episode_indices = sparse_dataset.hf_dataset["episode_index"]
+    mask = torch.zeros(len(episode_indices), dtype=torch.bool)
+    for ep in [0, 2, 4]:
+        mask |= torch.tensor(episode_indices) == ep
+
+    # Create a filtered dataset
+    filtered_data = {}
+    # Find image keys by checking features
+    image_keys = [key for key, ft in sparse_dataset.features.items() if ft.get("dtype") == "image"]
+
+    for key in sparse_dataset.hf_dataset.column_names:
+        values = sparse_dataset.hf_dataset[key]
+        # Filter values based on mask
+        filtered_values = [val for i, val in enumerate(values) if mask[i]]
+
+        # Convert float32 image tensors back to uint8 numpy arrays for HuggingFace dataset
+        if key in image_keys and len(filtered_values) > 0:
+            # Convert torch tensors (float32, [0, 1], CHW) back to numpy arrays (uint8, [0, 255], HWC)
+            filtered_values = [
+                (val.permute(1, 2, 0).numpy() * 255).astype(np.uint8) for val in filtered_values
+            ]
+
+        filtered_data[key] = filtered_values
+
+    sparse_dataset.hf_dataset = datasets.Dataset.from_dict(
+        filtered_data, features=get_hf_features_from_features(sparse_dataset.features)
+    )
+    sparse_dataset.hf_dataset.set_transform(hf_transform_to_torch)
+
+    # Test requesting all episodes when only some are cached
+    sparse_dataset.episodes = None
+    assert sparse_dataset._check_cached_episodes_sufficient() is False
+
+    # Test requesting only the available episodes
+    sparse_dataset.episodes = [0, 2, 4]
+    assert sparse_dataset._check_cached_episodes_sufficient() is True
+
+    # Test requesting a mix of available and unavailable episodes
+    sparse_dataset.episodes = [0, 1, 2]
+    assert sparse_dataset._check_cached_episodes_sufficient() is False
+
+
+def test_update_chunk_settings(tmp_path, empty_lerobot_dataset_factory):
+    """Test the update_chunk_settings functionality for both LeRobotDataset and LeRobotDatasetMetadata."""
+    features = {
+        "observation.state": {
+            "dtype": "float32",
+            "shape": (6,),
+            "names": ["shoulder_pan", "shoulder_lift", "elbow", "wrist_1", "wrist_2", "wrist_3"],
+        },
+        "action": {
+            "dtype": "float32",
+            "shape": (6,),
+            "names": ["shoulder_pan", "shoulder_lift", "elbow", "wrist_1", "wrist_2", "wrist_3"],
+        },
+    }
+
+    # Create dataset with default chunk settings
+    dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
+
+    # Test initial default values
+    initial_settings = dataset.meta.get_chunk_settings()
+    assert initial_settings["chunks_size"] == DEFAULT_CHUNK_SIZE
+    assert initial_settings["data_files_size_in_mb"] == DEFAULT_DATA_FILE_SIZE_IN_MB
+    assert initial_settings["video_files_size_in_mb"] == DEFAULT_VIDEO_FILE_SIZE_IN_MB
+
+    # Test updating all settings at once
+    new_chunks_size = 2000
+    new_data_size = 200
+    new_video_size = 1000
+
+    dataset.meta.update_chunk_settings(
+        chunks_size=new_chunks_size,
+        data_files_size_in_mb=new_data_size,
+        video_files_size_in_mb=new_video_size,
+    )
+
+    # Verify settings were updated
+    updated_settings = dataset.meta.get_chunk_settings()
+    assert updated_settings["chunks_size"] == new_chunks_size
+    assert updated_settings["data_files_size_in_mb"] == new_data_size
+    assert updated_settings["video_files_size_in_mb"] == new_video_size
+
+    # Test updating individual settings
+    dataset.meta.update_chunk_settings(chunks_size=1500)
+    settings_after_partial = dataset.meta.get_chunk_settings()
+    assert settings_after_partial["chunks_size"] == 1500
+    assert settings_after_partial["data_files_size_in_mb"] == new_data_size
+    assert settings_after_partial["video_files_size_in_mb"] == new_video_size
+
+    # Test updating only data file size
+    dataset.meta.update_chunk_settings(data_files_size_in_mb=150)
+    settings_after_data = dataset.meta.get_chunk_settings()
+    assert settings_after_data["chunks_size"] == 1500
+    assert settings_after_data["data_files_size_in_mb"] == 150
+    assert settings_after_data["video_files_size_in_mb"] == new_video_size
+
+    # Test updating only video file size
+    dataset.meta.update_chunk_settings(video_files_size_in_mb=800)
+    settings_after_video = dataset.meta.get_chunk_settings()
+    assert settings_after_video["chunks_size"] == 1500
+    assert settings_after_video["data_files_size_in_mb"] == 150
+    assert settings_after_video["video_files_size_in_mb"] == 800
+
+    # Test that settings persist in the info file
+    info_path = dataset.root / "meta" / "info.json"
+    assert info_path.exists()
+
+    # Verify the underlying metadata properties
+    assert dataset.meta.chunks_size == 1500
+    assert dataset.meta.data_files_size_in_mb == 150
+    assert dataset.meta.video_files_size_in_mb == 800
+
+    # Test error handling for invalid values
+    with pytest.raises(ValueError, match="chunks_size must be positive"):
+        dataset.meta.update_chunk_settings(chunks_size=0)
+
+    with pytest.raises(ValueError, match="chunks_size must be positive"):
+        dataset.meta.update_chunk_settings(chunks_size=-100)
+
+    with pytest.raises(ValueError, match="data_files_size_in_mb must be positive"):
+        dataset.meta.update_chunk_settings(data_files_size_in_mb=0)
+
+    with pytest.raises(ValueError, match="data_files_size_in_mb must be positive"):
+        dataset.meta.update_chunk_settings(data_files_size_in_mb=-50)
+
+    with pytest.raises(ValueError, match="video_files_size_in_mb must be positive"):
+        dataset.meta.update_chunk_settings(video_files_size_in_mb=0)
+
+    with pytest.raises(ValueError, match="video_files_size_in_mb must be positive"):
+        dataset.meta.update_chunk_settings(video_files_size_in_mb=-200)
+
+    # Test calling with None values (should not change anything)
+    settings_before_none = dataset.meta.get_chunk_settings()
+    dataset.meta.update_chunk_settings(
+        chunks_size=None, data_files_size_in_mb=None, video_files_size_in_mb=None
+    )
+    settings_after_none = dataset.meta.get_chunk_settings()
+    assert settings_before_none == settings_after_none
+
+    # Test metadata direct access
+    meta_settings = dataset.meta.get_chunk_settings()
+    assert meta_settings == dataset.meta.get_chunk_settings()
+
+    # Test updating via metadata directly
+    dataset.meta.update_chunk_settings(chunks_size=3000)
+    assert dataset.meta.get_chunk_settings()["chunks_size"] == 3000
+
+
+def test_update_chunk_settings_video_dataset(tmp_path):
+    """Test update_chunk_settings with a video dataset to ensure video-specific logic works."""
+    features = {
+        "observation.images.cam": {
+            "dtype": "video",
+            "shape": (480, 640, 3),
+            "names": ["height", "width", "channels"],
+        },
+        "action": {"dtype": "float32", "shape": (6,), "names": ["j1", "j2", "j3", "j4", "j5", "j6"]},
+    }
+
+    # Create video dataset
+    dataset = LeRobotDataset.create(
+        repo_id=DUMMY_REPO_ID, fps=30, features=features, root=tmp_path / "video_test", use_videos=True
+    )
+
+    # Test that video-specific settings work
+    original_video_size = dataset.meta.get_chunk_settings()["video_files_size_in_mb"]
+    new_video_size = original_video_size * 2
+
+    dataset.meta.update_chunk_settings(video_files_size_in_mb=new_video_size)
+    assert dataset.meta.get_chunk_settings()["video_files_size_in_mb"] == new_video_size
+    assert dataset.meta.video_files_size_in_mb == new_video_size
@@ -54,7 +54,6 @@ class LeRobotDatasetFactory(Protocol):


 def get_task_index(tasks: datasets.Dataset, task: str) -> int:
-    # TODO(rcadene): a bit complicated no? ^^
    task_idx = tasks.loc[task].task_index.item()
    return task_idx

@@ -17,8 +17,6 @@ from pathlib import Path
 import datasets
 import numpy as np
 import pandas as pd
-import pyarrow.compute as pc
-import pyarrow.parquet as pq
 import pytest
 from datasets import Dataset

@@ -142,16 +140,6 @@ def create_stats(stats_factory):
    return _create_stats


-# @pytest.fixture(scope="session")
-# def create_episodes_stats(episodes_stats_factory):
-#     def _create_episodes_stats(dir: Path, episodes_stats: Dataset | None = None):
-#         if episodes_stats is None:
-#             episodes_stats = episodes_stats_factory()
-#         write_episodes_stats(episodes_stats, dir)
-
-#     return _create_episodes_stats
-
-
@pytest.fixture(scope="session")
 def create_tasks(tasks_factory):
    def _create_tasks(dir: Path, tasks: pd.DataFrame | None = None):
@@ -186,53 +174,3 @@ def create_hf_dataset(hf_dataset_factory):
        write_hf_dataset(hf_dataset, dir, data_file_size_in_mb, chunk_size)

    return _create_hf_dataset
-
-
-@pytest.fixture(scope="session")
-def single_episode_parquet_path(hf_dataset_factory, info_factory):
-    def _create_single_episode_parquet(
-        dir: Path, ep_idx: int = 0, hf_dataset: datasets.Dataset | None = None, info: dict | None = None
-    ) -> Path:
-        raise NotImplementedError()
-        if info is None:
-            info = info_factory()
-        if hf_dataset is None:
-            hf_dataset = hf_dataset_factory()
-
-        data_path = info["data_path"]
-        chunks_size = info["chunks_size"]
-        ep_chunk = ep_idx // chunks_size
-        fpath = dir / data_path.format(episode_chunk=ep_chunk, episode_index=ep_idx)
-        fpath.parent.mkdir(parents=True, exist_ok=True)
-        table = hf_dataset.data.table
-        ep_table = table.filter(pc.equal(table["episode_index"], ep_idx))
-        pq.write_table(ep_table, fpath)
-        return fpath
-
-    return _create_single_episode_parquet
-
-
-@pytest.fixture(scope="session")
-def multi_episode_parquet_path(hf_dataset_factory, info_factory):
-    def _create_multi_episode_parquet(
-        dir: Path, hf_dataset: datasets.Dataset | None = None, info: dict | None = None
-    ) -> Path:
-        raise NotImplementedError()
-        if info is None:
-            info = info_factory()
-        if hf_dataset is None:
-            hf_dataset = hf_dataset_factory()
-
-        data_path = info["data_path"]
-        chunks_size = info["chunks_size"]
-        total_episodes = info["total_episodes"]
-        for ep_idx in range(total_episodes):
-            ep_chunk = ep_idx // chunks_size
-            fpath = dir / data_path.format(episode_chunk=ep_chunk, episode_index=ep_idx)
-            fpath.parent.mkdir(parents=True, exist_ok=True)
-            table = hf_dataset.data.table
-            ep_table = table.filter(pc.equal(table["episode_index"], ep_idx))
-            pq.write_table(ep_table, fpath)
-        return dir / "data"
-
-    return _create_multi_episode_parquet
Author	SHA1	Message	Date
Pepijn	42d35e47b2	fix port	2025-09-08 18:58:42 +02:00
Pepijn	1b9330a25a	fix import	2025-09-08 18:49:26 +02:00
Pepijn	490ffa89a5	add download script	2025-09-08 17:57:04 +02:00
Pepijn	3d31f2ad53	add port rlds script	2025-09-08 13:40:47 +02:00
Michel Aractingi	af79dda8d9	fix(caching) remove cache dir when collecting a dataset with each call to load_episodes and load_hf_dataset	2025-09-08 12:44:43 +02:00
Michel Aractingi	952f455446	fix(bug) in save_episode_data	2025-09-06 00:17:46 +02:00
Michel Aractingi	0747afdba7	Optimize dataset updates by incrementally concatenating new data instead of reloading from disk, reducing memory usage and improving performance.	2025-09-05 18:37:48 +02:00
Michel Aractingi	992fb177c3	further memory optimizations needed due to calling `pd.concat`	2025-09-03 18:49:30 +02:00
Michel Aractingi	1db3401159	remove unused Iterable Namespace	2025-09-03 16:23:36 +02:00
pre-commit-ci[bot]	7868df27dc	[pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci	2025-09-03 14:18:02 +00:00
Michel Aractingi	0e04f5fbbe	remove html templates and flask dependency	2025-09-03 16:17:10 +02:00
Michel Aractingi	fdccf7774b	fix(memory explosion) added delete to episodes and hf_dataset everytime we reload while collecting a dataset ot avoid memroy explosion	2025-09-03 15:31:28 +02:00
Michel Aractingi	2a3d62259e	visualize_dataset_html deprecated	2025-09-02 15:51:11 +02:00
Michel Aractingi	2df4e25558	added the file and video max size as arguments	2025-09-02 15:41:42 +02:00
pre-commit-ci[bot]	4062d0564a	[pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci	2025-09-01 19:32:19 +00:00
CarolinePascal	0a30636fc6	chore(dataset v2.0): drop support for dataset v2.0 format	2025-09-01 21:31:46 +02:00
CarolinePascal	adad3698e1	chore(dataset v1): drop support for dataset v1 format	2025-09-01 19:37:20 +02:00
Michel Aractingi	84ffc28854	moved `get_video_duration_in_s` to video_utils and replaced subprocess and ffmpeg with pyAV	2025-08-29 01:31:53 +02:00
Michel Aractingi	47aee1fdbe	revert back `video_utils.py` to using pyav while keeping concat_video_files function	2025-08-29 01:06:46 +02:00
Michel Aractingi	bbd64b9ce5	fixes in `datasets/utils.py`	2025-08-29 00:03:13 +02:00
Michel Aractingi	000e88760d	removed unused functions from tests/fixtures	2025-08-28 11:18:36 +02:00
Michel Aractingi	35f36e8fba	removed outdated todos	2025-08-28 10:10:17 +02:00
pre-commit-ci[bot]	213ffe02cf	[pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci	2025-08-28 07:53:10 +00:00
Michel Aractingi	2b03dec01f	Removed .item from save_dataset_to_safetensors	2025-08-28 09:52:37 +02:00
Michel Aractingi	64a9dd3763	Removed agibot files and moved port_droid to port_datasets	2025-08-28 09:44:38 +02:00
Francesco Capuano	2ca6edc19e	Merge branch 'main' into user/michel-aractingi/2025_06_30_dataset_v3 Signed-off-by: Francesco Capuano <74058581+fracapuano@users.noreply.github.com>	2025-08-25 16:34:44 +02:00
mgiac-hexagon	577cd10974	Removed dupicate lines of code (#1709 )	2025-08-25 12:39:32 +02:00
lxk	b0923ab74b	fix(dataset): Use provided episode_data in save_episode (#1740 ) The 'episode_data' parameter was previously ignored, causing an error if provided. This change ensures it is correctly used, which allows for asynchronous episode saving by passing a copy of the episode buffer, preventing conflicts with the main data collection loop.	2025-08-22 15:24:02 +02:00
Jack Vial	7f70b78f32	Add missing encoding table entries for Koch arm (#1534 )	2025-08-20 17:24:05 +02:00
Michel Aractingi	db36f01e8b	add update_chunk_settings method for LeRobotDatasetMetadata. Introduce tests for chunk settings updates and validation of parameters.	2025-08-18 00:00:23 +02:00
Steven Palma	55198de096	fix(ci): rename libegl1-mesa in deb13 trixie (#1735 )	2025-08-14 11:12:06 +02:00
Michel Aractingi	c7a3b02625	fixed tensor indicies in `_check_cached_episode_sufficient` in lerobot_dataset.py, added test	2025-08-13 16:16:32 +02:00