fix: modularize tests to improve readability

switch back to set_transform instead of set_format

Add video_files_size_in_mb

pre-commit run --all-files

README.md

@@ -116,7 +116,7 @@ pip install -e .
 ```
 
 > **NOTE:** If you encounter build errors, you may need to install additional dependencies (`cmake`, `build-essential`, and `ffmpeg libs`). On Linux, run:
-`sudo apt-get install cmake build-essential python-dev pkg-config libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev pkg-config`. For other systems, see: [Compiling PyAV](https://pyav.org/docs/develop/overview/installation.html#bring-your-own-ffmpeg)
+`sudo apt-get install cmake build-essential python3-dev pkg-config libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev pkg-config`. For other systems, see: [Compiling PyAV](https://pyav.org/docs/develop/overview/installation.html#bring-your-own-ffmpeg)
 
 For simulations, 🤗 LeRobot comes with gymnasium environments that can be installed as extras:
 - [aloha](https://github.com/huggingface/gym-aloha)

docker/lerobot-gpu-dev/Dockerfile

@@ -14,7 +14,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     tcpdump sysstat screen tmux \
     libglib2.0-0 libgl1-mesa-glx libegl1-mesa \
     speech-dispatcher portaudio19-dev libgeos-dev \
-    python${PYTHON_VERSION} python${PYTHON_VERSION}-venv \
+    python${PYTHON_VERSION} python${PYTHON_VERSION}-venv python${PYTHON_VERSION}-dev \
     && apt-get clean && rm -rf /var/lib/apt/lists/*
 
 # Install ffmpeg build dependencies. See:

examples/4_train_policy_with_script.md

@@ -4,7 +4,7 @@ This tutorial will explain the training script, how to use it, and particularly
 
 ## The training script
 
-LeRobot offers a training script at [`lerobot/scripts/train.py`](../../lerobot/scripts/train.py). At a high level it does the following:
+LeRobot offers a training script at [`lerobot/scripts/train.py`](../lerobot/scripts/train.py). At a high level it does the following:
 
 - Initialize/load a configuration for the following steps using.
 - Instantiates a dataset.
@@ -21,7 +21,7 @@ In the training script, the main function `train` expects a `TrainPipelineConfig
 def train(cfg: TrainPipelineConfig):
 ```
 
-You can inspect the `TrainPipelineConfig` defined in [`lerobot/configs/train.py`](../../lerobot/configs/train.py) (which is heavily commented and meant to be a reference to understand any option)
+You can inspect the `TrainPipelineConfig` defined in [`lerobot/configs/train.py`](../lerobot/configs/train.py) (which is heavily commented and meant to be a reference to understand any option)
 
 When running the script, inputs for the command line are parsed thanks to the `@parser.wrap()` decorator and an instance of this class is automatically generated. Under the hood, this is done with [Draccus](https://github.com/dlwh/draccus) which is a tool dedicated for this purpose. If you're familiar with Hydra, Draccus can similarly load configurations from config files (.json, .yaml) and also override their values through command line inputs. Unlike Hydra, these configurations are pre-defined in the code through dataclasses rather than being defined entirely in config files. This allows for more rigorous serialization/deserialization, typing, and to manipulate configuration as objects directly in the code and not as dictionaries or namespaces (which enables nice features in an IDE such as autocomplete, jump-to-def, etc.)
 
@@ -50,7 +50,7 @@ By default, every field takes its default value specified in the dataclass. If a
 
 ## Specifying values from the CLI
 
-Let's say that we want to train [Diffusion Policy](../../lerobot/common/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
+Let's say that we want to train [Diffusion Policy](../lerobot/common/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
 ```bash
 python lerobot/scripts/train.py \
     --dataset.repo_id=lerobot/pusht \
@@ -60,10 +60,10 @@ python lerobot/scripts/train.py \
 
 Let's break this down:
 - To specify the dataset, we just need to specify its `repo_id` on the hub which is the only required argument in the `DatasetConfig`. The rest of the fields have default values and in this case we are fine with those so we can just add the option `--dataset.repo_id=lerobot/pusht`.
-- To specify the policy, we can just select diffusion policy using `--policy` appended with `.type`. Here, `.type` is a special argument which allows us to select config classes inheriting from `draccus.ChoiceRegistry` and that have been decorated with the `register_subclass()` method. To have a better explanation of this feature, have a look at this [Draccus demo](https://github.com/dlwh/draccus?tab=readme-ov-file#more-flexible-configuration-with-choice-types). In our code, we use this mechanism mainly to select policies, environments, robots, and some other components like optimizers. The policies available to select are located in [lerobot/common/policies](../../lerobot/common/policies)
-- Similarly, we select the environment with `--env.type=pusht`. The different environment configs are available in [`lerobot/common/envs/configs.py`](../../lerobot/common/envs/configs.py)
+- To specify the policy, we can just select diffusion policy using `--policy` appended with `.type`. Here, `.type` is a special argument which allows us to select config classes inheriting from `draccus.ChoiceRegistry` and that have been decorated with the `register_subclass()` method. To have a better explanation of this feature, have a look at this [Draccus demo](https://github.com/dlwh/draccus?tab=readme-ov-file#more-flexible-configuration-with-choice-types). In our code, we use this mechanism mainly to select policies, environments, robots, and some other components like optimizers. The policies available to select are located in [lerobot/common/policies](../lerobot/common/policies)
+- Similarly, we select the environment with `--env.type=pusht`. The different environment configs are available in [`lerobot/common/envs/configs.py`](../lerobot/common/envs/configs.py)
 
-Let's see another example. Let's say you've been training [ACT](../../lerobot/common/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
+Let's see another example. Let's say you've been training [ACT](../lerobot/common/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
 ```bash
 python lerobot/scripts/train.py \
     --policy.type=act \
@@ -74,7 +74,7 @@ python lerobot/scripts/train.py \
 > Notice we added `--output_dir` to explicitly tell where to write outputs from this run (checkpoints, training state, configs etc.). This is not mandatory and if you don't specify it, a default directory will be created from the current date and time, env.type and policy.type. This will typically look like `outputs/train/2025-01-24/16-10-05_aloha_act`.
 
 We now want to train a different policy for aloha on another task. We'll change the dataset and use [lerobot/aloha_sim_transfer_cube_human](https://huggingface.co/datasets/lerobot/aloha_sim_transfer_cube_human) instead. Of course, we also need to change the task of the environment as well to match this other task.
-Looking at the [`AlohaEnv`](../../lerobot/common/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
+Looking at the [`AlohaEnv`](../lerobot/common/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
 ```bash
 python lerobot/scripts/train.py \
     --policy.type=act \

examples/port_datasets/droid_rlds/port_droid.py

@@ -22,7 +22,7 @@ from pathlib import Path
 import numpy as np
 import tensorflow_datasets as tfds
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
 from lerobot.common.utils.utils import get_elapsed_time_in_days_hours_minutes_seconds
 
 DROID_SHARDS = 2048
@@ -370,6 +370,25 @@ def port_droid(
         )
 
 
+def validate_dataset(repo_id):
+    """Sanity check that ensure meta data 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()
 

examples/port_datasets/droid_rlds/slurm_port_shards.py

@@ -6,26 +6,6 @@ from datatrove.executor.slurm import SlurmPipelineExecutor
 from datatrove.pipeline.base import PipelineStep
 
 from examples.port_datasets.droid_rlds.port_droid import DROID_SHARDS
-from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
-
-
-def validate_shard(repo_id):
-    """Sanity check that ensure meta data 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}")
 
 
 class PortDroidShards(PipelineStep):
@@ -41,7 +21,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 examples.port_datasets.droid_rlds.port_droid import port_droid
+        from examples.port_datasets.droid_rlds.port_droid import port_droid, validate_dataset
         from lerobot.common.utils.utils import init_logging
 
         init_logging()
@@ -49,6 +29,12 @@ class PortDroidShards(PipelineStep):
 
         shard_repo_id = f"{self.repo_id}_world_{world_size}_rank_{rank}"
 
+        try:
+            validate_dataset(shard_repo_id)
+            return
+        except:
+            pass
+
         port_droid(
             self.raw_dir,
             shard_repo_id,
@@ -57,7 +43,7 @@ class PortDroidShards(PipelineStep):
             shard_index=rank,
         )
 
-        validate_shard(shard_repo_id)
+        validate_dataset(shard_repo_id)
 
 
 def make_port_executor(

lerobot/common/datasets/aggregate.py

@@ -6,17 +6,23 @@ import pandas as pd
 import tqdm
 
 from lerobot.common.datasets.compute_stats import aggregate_stats
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.common.datasets.utils import (
     DEFAULT_CHUNK_SIZE,
+    DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_DATA_PATH,
     DEFAULT_EPISODES_PATH,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
     DEFAULT_VIDEO_PATH,
+    concat_video_files,
+    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,
 )
-from lerobot.common.utils.utils import init_logging
 
 
 def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
@@ -41,208 +47,313 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
     return fps, robot_type, features
 
 
-def get_update_episode_and_task_func(episode_index_to_add, old_tasks, new_tasks):
+def update_data_df(df, src_meta, dst_meta):
     def _update(row):
-        row["episode_index"] = row["episode_index"] + episode_index_to_add
-        task = old_tasks.iloc[row["task_index"]].name
-        row["task_index"] = new_tasks.loc[task].task_index.item()
+        row["episode_index"] = row["episode_index"] + dst_meta.info["total_episodes"]
+        row["index"] = row["index"] + dst_meta.info["total_frames"]
+        task = src_meta.tasks.iloc[row["task_index"]].name
+        row["task_index"] = dst_meta.tasks.loc[task].task_index.item()
         return row
 
-    return _update
+    return df.apply(_update, axis=1)
 
 
-def get_update_meta_func(
-    meta_chunk_index_to_add,
-    meta_file_index_to_add,
-    data_chunk_index_to_add,
-    data_file_index_to_add,
-    videos_chunk_index_to_add,
-    videos_file_index_to_add,
-    frame_index_to_add,
+def update_meta_data(
+    df,
+    dst_meta,
+    meta_idx,
+    data_idx,
+    videos_idx,
 ):
     def _update(row):
-        row["meta/episodes/chunk_index"] = row["meta/episodes/chunk_index"] + meta_chunk_index_to_add
-        row["meta/episodes/file_index"] = row["meta/episodes/file_index"] + meta_file_index_to_add
-        row["data/chunk_index"] = row["data/chunk_index"] + data_chunk_index_to_add
-        row["data/file_index"] = row["data/file_index"] + data_file_index_to_add
-        for key in videos_chunk_index_to_add:
-            row[f"videos/{key}/chunk_index"] = (
-                row[f"videos/{key}/chunk_index"] + videos_chunk_index_to_add[key]
+        row["meta/episodes/chunk_index"] = row["meta/episodes/chunk_index"] + meta_idx["chunk"]
+        row["meta/episodes/file_index"] = row["meta/episodes/file_index"] + meta_idx["file"]
+        row["data/chunk_index"] = row["data/chunk_index"] + data_idx["chunk"]
+        row["data/file_index"] = row["data/file_index"] + data_idx["file"]
+        for key, video_idx in videos_idx.items():
+            row[f"videos/{key}/chunk_index"] = row[f"videos/{key}/chunk_index"] + video_idx["chunk"]
+            row[f"videos/{key}/file_index"] = row[f"videos/{key}/file_index"] + video_idx["file"]
+            row[f"videos/{key}/from_timestamp"] = (
+                row[f"videos/{key}/from_timestamp"] + video_idx["latest_duration"]
             )
-            row[f"videos/{key}/file_index"] = row[f"videos/{key}/file_index"] + videos_file_index_to_add[key]
-        row["dataset_from_index"] = row["dataset_from_index"] + frame_index_to_add
-        row["dataset_to_index"] = row["dataset_to_index"] + frame_index_to_add
+            row[f"videos/{key}/to_timestamp"] = (
+                row[f"videos/{key}/to_timestamp"] + video_idx["latest_duration"]
+            )
+        row["dataset_from_index"] = row["dataset_from_index"] + dst_meta.info["total_frames"]
+        row["dataset_to_index"] = row["dataset_to_index"] + dst_meta.info["total_frames"]
+        row["episode_index"] = row["episode_index"] + dst_meta.info["total_episodes"]
         return row
 
-    return _update
+    return df.apply(_update, axis=1)
 
 
 def aggregate_datasets(repo_ids: list[str], aggr_repo_id: str, roots: list[Path] = None, aggr_root=None):
     logging.info("Start aggregate_datasets")
 
-    if roots is None:
-        all_metadata = [LeRobotDatasetMetadata(repo_id) for repo_id in repo_ids]
-    else:
-        all_metadata = [
+    # Load metadata
+    all_metadata = (
+        [LeRobotDatasetMetadata(repo_id) for repo_id in repo_ids]
+        if roots is None
+        else [
             LeRobotDatasetMetadata(repo_id, root=root) for repo_id, root in zip(repo_ids, roots, strict=False)
         ]
-
+    )
     fps, robot_type, features = validate_all_metadata(all_metadata)
     video_keys = [key for key in features if features[key]["dtype"] == "video"]
 
-    # Create resulting dataset folder
-    aggr_meta = LeRobotDatasetMetadata.create(
+    # Initialize output dataset metadata
+    dst_meta = LeRobotDatasetMetadata.create(
         repo_id=aggr_repo_id,
         fps=fps,
         robot_type=robot_type,
         features=features,
         root=aggr_root,
     )
-    aggr_root = aggr_meta.root
 
+    # Aggregate task info
     logging.info("Find all tasks")
-    unique_tasks = pd.concat([meta.tasks for meta in all_metadata]).index.unique()
-    aggr_meta.tasks = pd.DataFrame({"task_index": range(len(unique_tasks))}, index=unique_tasks)
+    unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
+    dst_meta.tasks = pd.DataFrame({"task_index": range(len(unique_tasks))}, index=unique_tasks)
 
-    num_episodes = 0
-    num_frames = 0
+    # Track counters and indices
+    meta_idx = {"chunk": 0, "file": 0}
+    data_idx = {"chunk": 0, "file": 0}
+    videos_idx = {
+        key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
+    }
 
-    aggr_meta_chunk_idx = 0
-    aggr_meta_file_idx = 0
+    dst_meta.episodes = {}
 
-    aggr_data_chunk_idx = 0
-    aggr_data_file_idx = 0
+    # Process each dataset
+    for src_meta in tqdm.tqdm(all_metadata, desc="Copy data and videos"):
+        videos_idx = aggregate_videos(src_meta, dst_meta, videos_idx)
+        data_idx = aggregate_data(src_meta, dst_meta, data_idx)
 
-    aggr_videos_chunk_idx = dict.fromkeys(video_keys, 0)
-    aggr_videos_file_idx = dict.fromkeys(video_keys, 0)
+        meta_idx = aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx)
 
-    for meta in tqdm.tqdm(all_metadata, desc="Copy data and videos"):
-        meta_chunk_file_ids = {
-            (c, f)
-            for c, f in zip(
-                meta.episodes["meta/episodes/chunk_index"],
-                meta.episodes["meta/episodes/file_index"],
+        dst_meta.info["total_episodes"] += src_meta.total_episodes
+        dst_meta.info["total_frames"] += src_meta.total_frames
+
+    finalize_aggregation(dst_meta, all_metadata)
+    logging.info("Aggregation complete.")
+
+
+# -------------------------------
+# Helper Functions
+# -------------------------------
+
+
+def aggregate_videos(src_meta, dst_meta, videos_idx):
+    """
+    Aggregates video chunks from a dataset into the aggregated dataset folder.
+    """
+    for key, video_idx in videos_idx.items():
+        # Get unique (chunk, file) combinations
+        unique_chunk_file_pairs = {
+            (chunk, file)
+            for chunk, file in zip(
+                src_meta.episodes[f"videos/{key}/chunk_index"],
+                src_meta.episodes[f"videos/{key}/file_index"],
                 strict=False,
             )
         }
-        for chunk_idx, file_idx in meta_chunk_file_ids:
-            path = meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
-            df = pd.read_parquet(path)
-            update_meta_func = get_update_meta_func(
-                aggr_meta_chunk_idx,
-                aggr_meta_file_idx,
-                aggr_data_chunk_idx,
-                aggr_data_file_idx,
-                aggr_videos_chunk_idx,
-                aggr_videos_file_idx,
-                num_frames,
+
+        # Current target chunk/file index
+        chunk_idx = video_idx["chunk"]
+        file_idx = video_idx["file"]
+
+        for src_chunk_idx, src_file_idx in unique_chunk_file_pairs:
+            src_path = src_meta.root / DEFAULT_VIDEO_PATH.format(
+                video_key=key,
+                chunk_index=src_chunk_idx,
+                file_index=src_file_idx,
             )
-            df = df.apply(update_meta_func, axis=1)
 
-            aggr_path = aggr_root / DEFAULT_EPISODES_PATH.format(
-                chunk_index=aggr_meta_chunk_idx, file_index=aggr_meta_file_idx
+            dst_path = dst_meta.root / DEFAULT_VIDEO_PATH.format(
+                video_key=key,
+                chunk_index=chunk_idx,
+                file_index=file_idx,
             )
-            aggr_path.parent.mkdir(parents=True, exist_ok=True)
-            df.to_parquet(aggr_path)
 
-            aggr_meta_file_idx += 1
-            if aggr_meta_file_idx >= DEFAULT_CHUNK_SIZE:
-                aggr_meta_file_idx = 0
-                aggr_meta_chunk_idx += 1
+            if not dst_path.exists():
+                # First write to this destination file
+                dst_path.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copy(str(src_path), str(dst_path))
+                continue
 
-        # cp videos
-        for key in video_keys:
-            video_chunk_file_ids = {
-                (c, f)
-                for c, f in zip(
-                    meta.episodes[f"videos/{key}/chunk_index"],
-                    meta.episodes[f"videos/{key}/file_index"],
-                    strict=False,
-                )
-            }
-            for chunk_idx, file_idx in video_chunk_file_ids:
-                path = meta.root / DEFAULT_VIDEO_PATH.format(
-                    video_key=key, chunk_index=chunk_idx, file_index=file_idx
-                )
-                aggr_path = aggr_root / DEFAULT_VIDEO_PATH.format(
+            # Check file sizes before appending
+            src_size = get_video_size_in_mb(src_path)
+            dst_size = get_video_size_in_mb(dst_path)
+
+            if dst_size + src_size >= DEFAULT_VIDEO_FILE_SIZE_IN_MB:
+                # Rotate to a new chunk/file
+                chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
+                dst_path = dst_meta.root / DEFAULT_VIDEO_PATH.format(
                     video_key=key,
-                    chunk_index=aggr_videos_chunk_idx[key],
-                    file_index=aggr_videos_file_idx[key],
+                    chunk_index=chunk_idx,
+                    file_index=file_idx,
+                )
+                dst_path.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copy(str(src_path), str(dst_path))
+            else:
+                # Append to existing video file
+                concat_video_files(
+                    [dst_path, src_path],
+                    dst_meta.root,
+                    key,
+                    chunk_idx,
+                    file_idx,
                 )
-                aggr_path.parent.mkdir(parents=True, exist_ok=True)
-                shutil.copy(str(path), str(aggr_path))
 
-                # copy_command = f"cp {video_path} {aggr_video_path} &"
-                # subprocess.Popen(copy_command, shell=True)
+        # Update the videos_idx with the final chunk and file indices for this key
+        videos_idx[key]["chunk"] = chunk_idx
+        videos_idx[key]["file"] = file_idx
 
-                aggr_videos_file_idx[key] += 1
-                if aggr_videos_file_idx[key] >= DEFAULT_CHUNK_SIZE:
-                    aggr_videos_file_idx[key] = 0
-                    aggr_videos_chunk_idx[key] += 1
+    return videos_idx
 
-        data_chunk_file_ids = {
-            (c, f)
-            for c, f in zip(meta.episodes["data/chunk_index"], meta.episodes["data/file_index"], strict=False)
-        }
-        for chunk_idx, file_idx in data_chunk_file_ids:
-            path = meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
-            df = pd.read_parquet(path)
-            update_data_func = get_update_episode_and_task_func(num_episodes, meta.tasks, aggr_meta.tasks)
-            df = df.apply(update_data_func, axis=1)
 
-            aggr_path = aggr_root / DEFAULT_DATA_PATH.format(
-                chunk_index=aggr_data_chunk_idx, file_index=aggr_data_file_idx
-            )
-            aggr_path.parent.mkdir(parents=True, exist_ok=True)
-            df.to_parquet(aggr_path)
+def aggregate_data(src_meta, dst_meta, data_idx):
+    unique_chunk_file_ids = {
+        (c, f)
+        for c, f in zip(
+            src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False
+        )
+    }
+    for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
+        src_path = src_meta.root / DEFAULT_DATA_PATH.format(
+            chunk_index=src_chunk_idx, file_index=src_file_idx
+        )
+        df = pd.read_parquet(src_path)
+        df = update_data_df(df, src_meta, dst_meta)
 
-            aggr_data_file_idx += 1
-            if aggr_data_file_idx >= DEFAULT_CHUNK_SIZE:
-                aggr_data_file_idx = 0
-                aggr_data_chunk_idx += 1
+        data_idx = append_or_create_parquet_file(
+            df,
+            src_path,
+            data_idx,
+            DEFAULT_DATA_FILE_SIZE_IN_MB,
+            DEFAULT_CHUNK_SIZE,
+            DEFAULT_DATA_PATH,
+            contains_images=len(dst_meta.image_keys) > 0,
+            aggr_root=dst_meta.root,
+        )
 
-        num_episodes += meta.total_episodes
-        num_frames += meta.total_frames
+    return data_idx
 
+
+def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
+    chunk_file_ids = {
+        (c, f)
+        for c, f in zip(
+            src_meta.episodes["meta/episodes/chunk_index"],
+            src_meta.episodes["meta/episodes/file_index"],
+            strict=False,
+        )
+    }
+
+    for chunk_idx, file_idx in chunk_file_ids:
+        src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
+        df = pd.read_parquet(src_path)
+        df = update_meta_data(
+            df,
+            dst_meta,
+            meta_idx,
+            data_idx,
+            videos_idx,
+        )
+
+        for k in videos_idx:
+            videos_idx[k]["latest_duration"] += videos_idx[k]["episode_duration"]
+
+        meta_idx = append_or_create_parquet_file(
+            df,
+            src_path,
+            meta_idx,
+            DEFAULT_DATA_FILE_SIZE_IN_MB,
+            DEFAULT_CHUNK_SIZE,
+            DEFAULT_EPISODES_PATH,
+            contains_images=False,
+            aggr_root=dst_meta.root,
+        )
+
+    return meta_idx
+
+
+def append_or_create_parquet_file(
+    df: pd.DataFrame,
+    src_path: Path,
+    idx: dict[str, int],
+    max_mb: float,
+    chunk_size: int,
+    default_path: str,
+    contains_images: bool = False,
+    aggr_root: Path = None,
+):
+    """
+    Safely appends or creates a Parquet file at dst_path based on size constraints.
+
+    Parameters:
+        df (pd.DataFrame): Data to write.
+        src_path (Path): Path to source file (used to get size).
+        idx (dict): Dictionary containing 'chunk' and 'file' indices.
+        max_mb (float): Maximum allowed file size in MB.
+        chunk_size (int): Maximum number of files per chunk.
+        default_path (str): Format string for generating a new file path.
+
+    Returns:
+        dict: Updated index dictionary.
+    """
+    # Initial destination path - use the correct default_path parameter
+    dst_path = aggr_root / default_path.format(chunk_index=idx["chunk"], file_index=idx["file"])
+
+    # If destination file doesn't exist, just write the new one
+    if not dst_path.exists():
+        dst_path.parent.mkdir(parents=True, exist_ok=True)
+        if contains_images:
+            to_parquet_with_hf_images(df, dst_path)
+        else:
+            df.to_parquet(dst_path)
+        return idx
+
+    # Otherwise, check if we exceed the size limit
+    src_size = get_parquet_file_size_in_mb(src_path)
+    dst_size = get_parquet_file_size_in_mb(dst_path)
+
+    if dst_size + src_size >= max_mb:
+        # File is too large, move to a new one
+        idx["chunk"], idx["file"] = update_chunk_file_indices(idx["chunk"], idx["file"], chunk_size)
+        new_path = aggr_root / default_path.format(chunk_index=idx["chunk"], file_index=idx["file"])
+        new_path.parent.mkdir(parents=True, exist_ok=True)
+        final_df = df
+        target_path = new_path
+    else:
+        # Append to existing file
+        existing_df = pd.read_parquet(dst_path)
+        final_df = pd.concat([existing_df, df], ignore_index=True)
+        target_path = dst_path
+
+    if contains_images:
+        to_parquet_with_hf_images(final_df, target_path)
+    else:
+        final_df.to_parquet(target_path)
+
+    return idx
+
+
+def finalize_aggregation(aggr_meta, all_metadata):
     logging.info("write tasks")
     write_tasks(aggr_meta.tasks, aggr_meta.root)
 
     logging.info("write info")
-    aggr_meta.info["total_episodes"] = sum([meta.total_episodes for meta in all_metadata])
-    aggr_meta.info["total_frames"] = sum([meta.total_frames for meta in all_metadata])
-    aggr_meta.info["splits"] = {"train": f"0:{aggr_meta.total_episodes}"}
+    aggr_meta.info.update(
+        {
+            "total_tasks": len(aggr_meta.tasks),
+            "total_episodes": sum(m.total_episodes for m in all_metadata),
+            "total_frames": sum(m.total_frames for m in all_metadata),
+            "splits": {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"},
+        }
+    )
     write_info(aggr_meta.info, aggr_meta.root)
 
     logging.info("write stats")
-    aggr_meta.stats = aggregate_stats([meta.stats for meta in all_metadata])
+    aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
     write_stats(aggr_meta.stats, aggr_meta.root)
-
-
-if __name__ == "__main__":
-    init_logging()
-    repo_id = "cadene/droid"
-    aggr_repo_id = "cadene/droid"
-    datetime = "2025-02-22_11-23-54"
-
-    # root = Path(f"/tmp/{repo_id}")
-    # if root.exists():
-    #     shutil.rmtree(root)
-    root = None
-
-    # all_metadata = [LeRobotDatasetMetadata(f"{repo_id}_{datetime}_world_2048_rank_{rank}") for rank in range(2048)]
-
-    # aggregate_datasets(
-    #     all_metadata,
-    #     aggr_repo_id,
-    #     root=root,
-    # )
-
-    aggr_dataset = LeRobotDataset(
-        repo_id=aggr_repo_id,
-        root=root,
-    )
-    aggr_dataset.push_to_hub(tags=["openx"])
-
-    # for meta in all_metadata:
-    #     dataset = LeRobotDataset(repo_id=meta.repo_id, root=meta.root)
-    #     dataset.push_to_hub(tags=["openx"])

lerobot/common/datasets/lerobot_dataset.py

@@ -56,12 +56,14 @@ from lerobot.common.datasets.utils import (
     get_safe_version,
     get_video_duration_in_s,
     get_video_size_in_mb,
+    hf_transform_to_torch,
     is_valid_version,
     load_episodes,
     load_info,
     load_nested_dataset,
     load_stats,
     load_tasks,
+    to_parquet_with_hf_images,
     update_chunk_file_indices,
     validate_episode_buffer,
     validate_frame,
@@ -136,14 +138,16 @@ class LeRobotDatasetMetadata:
         return packaging.version.parse(self.info["codebase_version"])
 
     def get_data_file_path(self, ep_index: int) -> Path:
-        chunk_idx = self.episodes["data/chunk_index"][ep_index]
-        file_idx = self.episodes["data/file_index"][ep_index]
+        ep = self.episodes[ep_index]
+        chunk_idx = ep["data/chunk_index"]
+        file_idx = ep["data/file_index"]
         fpath = self.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
         return Path(fpath)
 
     def get_video_file_path(self, ep_index: int, vid_key: str) -> Path:
-        chunk_idx = self.episodes[f"videos/{vid_key}/chunk_index"][ep_index]
-        file_idx = self.episodes[f"videos/{vid_key}/file_index"][ep_index]
+        ep = self.episodes[ep_index]
+        chunk_idx = ep[f"videos/{vid_key}/chunk_index"]
+        file_idx = ep[f"videos/{vid_key}/file_index"]
         fpath = self.video_path.format(video_key=vid_key, chunk_index=chunk_idx, file_index=file_idx)
         return Path(fpath)
 
@@ -218,9 +222,14 @@ class LeRobotDatasetMetadata:
         return self.info["chunks_size"]
 
     @property
-    def files_size_in_mb(self) -> int:
-        """Max size of file in mega bytes."""
-        return self.info["files_size_in_mb"]
+    def data_files_size_in_mb(self) -> int:
+        """Max size of data file in mega bytes."""
+        return self.info["data_files_size_in_mb"]
+
+    @property
+    def video_files_size_in_mb(self) -> int:
+        """Max size of video file in mega bytes."""
+        return self.info["video_files_size_in_mb"]
 
     def get_task_index(self, task: str) -> int | None:
         """
@@ -268,6 +277,7 @@ class LeRobotDatasetMetadata:
         ep_dataset = 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]
 
         if self.episodes is None:
             # Initialize indices and frame count for a new dataset made of the first episode data
@@ -275,26 +285,17 @@ class LeRobotDatasetMetadata:
             df["meta/episodes/chunk_index"] = [chunk_idx]
             df["meta/episodes/file_index"] = [file_idx]
             df["dataset_from_index"] = [0]
-            df["dataset_to_index"] = [len(df)]
+            df["dataset_to_index"] = [num_frames]
         else:
             # Retrieve information from the latest parquet file
-            latest_ep = self.episodes.with_format(
-                columns=[
-                    "meta/episodes/chunk_index",
-                    "meta/episodes/file_index",
-                    "dataset_from_index",
-                    "dataset_to_index",
-                ]
-            )[-1]
-            chunk_idx, file_idx = (
-                latest_ep["meta/episodes/chunk_index"],
-                latest_ep["meta/episodes/file_index"],
-            )
+            latest_ep = self.episodes[-1]
+            chunk_idx = latest_ep["meta/episodes/chunk_index"]
+            file_idx = latest_ep["meta/episodes/file_index"]
 
             latest_path = self.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
             latest_size_in_mb = get_parquet_file_size_in_mb(latest_path)
 
-            if latest_size_in_mb + ep_size_in_mb >= self.files_size_in_mb:
+            if latest_size_in_mb + ep_size_in_mb >= self.data_files_size_in_mb:
                 # Size limit is reached, prepare new parquet file
                 chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.chunks_size)
 
@@ -302,9 +303,9 @@ class LeRobotDatasetMetadata:
             df["meta/episodes/chunk_index"] = [chunk_idx]
             df["meta/episodes/file_index"] = [file_idx]
             df["dataset_from_index"] = [latest_ep["dataset_to_index"]]
-            df["dataset_to_index"] = [latest_ep["dataset_to_index"] + len(df)]
+            df["dataset_to_index"] = [latest_ep["dataset_to_index"] + num_frames]
 
-            if latest_size_in_mb + ep_size_in_mb < self.files_size_in_mb:
+            if latest_size_in_mb + ep_size_in_mb < self.data_files_size_in_mb:
                 # Size limit wasnt reached, concatenate latest dataframe with new one
                 latest_df = pd.read_parquet(latest_path)
                 df = pd.concat([latest_df, df], ignore_index=True)
@@ -339,6 +340,7 @@ class LeRobotDatasetMetadata:
         # Update info
         self.info["total_episodes"] += 1
         self.info["total_frames"] += episode_length
+        self.info["total_tasks"] = len(self.tasks)
         self.info["splits"] = {"train": f"0:{self.info['total_episodes']}"}
         if len(self.video_keys) > 0:
             self.update_video_info()
@@ -674,14 +676,14 @@ class LeRobotDataset(torch.utils.data.Dataset):
     def load_hf_dataset(self) -> datasets.Dataset:
         """hf_dataset contains all the observations, states, actions, rewards, etc."""
         hf_dataset = load_nested_dataset(self.root / "data")
-        hf_dataset.set_format("torch")
+        hf_dataset.set_transform(hf_transform_to_torch)
         return hf_dataset
 
     def create_hf_dataset(self) -> datasets.Dataset:
         features = get_hf_features_from_features(self.features)
         ft_dict = {col: [] for col in features}
         hf_dataset = datasets.Dataset.from_dict(ft_dict, features=features, split="train")
-        hf_dataset.set_format("torch")
+        hf_dataset.set_transform(hf_transform_to_torch)
         return hf_dataset
 
     @property
@@ -712,8 +714,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
             return get_hf_features_from_features(self.features)
 
     def _get_query_indices(self, idx: int, ep_idx: int) -> tuple[dict[str, list[int | bool]]]:
-        ep_start = self.meta.episodes["dataset_from_index"][ep_idx]
-        ep_end = self.meta.episodes["dataset_to_index"][ep_idx]
+        ep = self.meta.episodes[ep_idx]
+        ep_start = ep["dataset_from_index"]
+        ep_end = ep["dataset_to_index"]
         query_indices = {
             key: [max(ep_start, min(ep_end - 1, idx + delta)) for delta in delta_idx]
             for key, delta_idx in self.delta_indices.items()
@@ -734,8 +737,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
         query_timestamps = {}
         for key in self.meta.video_keys:
             if query_indices is not None and key in query_indices:
-                timestamps = self.hf_dataset.select(query_indices[key])["timestamp"]
-                query_timestamps[key] = timestamps.tolist()
+                timestamps = self.hf_dataset[query_indices[key]]["timestamp"]
+                query_timestamps[key] = torch.stack(timestamps).tolist()
             else:
                 query_timestamps[key] = [current_ts]
 
@@ -743,7 +746,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
     def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict:
         return {
-            key: self.hf_dataset.select(q_idx)[key]
+            key: torch.stack(self.hf_dataset[q_idx][key])
             for key, q_idx in query_indices.items()
             if key not in self.meta.video_keys
         }
@@ -754,12 +757,13 @@ class LeRobotDataset(torch.utils.data.Dataset):
         Segmentation Fault. This probably happens because a memory reference to the video loader is created in
         the main process and a subprocess fails to access it.
         """
+        ep = self.meta.episodes[ep_idx]
         item = {}
         for vid_key, query_ts in query_timestamps.items():
             # Episodes are stored sequentially on a single mp4 to reduce the number of files.
-            # Thus we load the start timestamp of the episode on this mp4 and,
+            # Thus we load the start timestamp of the episode on this mp4 and
             # shift the query timestamp accordingly.
-            from_timestamp = self.meta.episodes[f"videos/{vid_key}/from_timestamp"][ep_idx]
+            from_timestamp = ep[f"videos/{vid_key}/from_timestamp"]
             shifted_query_ts = [from_timestamp + ts for ts in query_ts]
 
             video_path = self.root / self.meta.get_video_file_path(ep_idx, vid_key)
@@ -984,15 +988,16 @@ class LeRobotDataset(torch.utils.data.Dataset):
             latest_num_frames = 0
         else:
             # Retrieve information from the latest parquet file
-            latest_ep = self.meta.episodes.with_format(columns=["data/chunk_index", "data/file_index"])[-1]
-            chunk_idx, file_idx = latest_ep["data/chunk_index"], latest_ep["data/file_index"]
+            latest_ep = self.meta.episodes[-1]
+            chunk_idx = latest_ep["data/chunk_index"]
+            file_idx = latest_ep["data/file_index"]
 
             latest_path = self.root / self.meta.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
             latest_size_in_mb = get_parquet_file_size_in_mb(latest_path)
             latest_num_frames = get_parquet_num_frames(latest_path)
 
             # Determine if a new parquet file is needed
-            if latest_size_in_mb + ep_size_in_mb >= self.meta.files_size_in_mb:
+            if latest_size_in_mb + ep_size_in_mb >= self.meta.data_files_size_in_mb:
                 # Size limit is reached, prepare new parquet file
                 chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.meta.chunks_size)
                 latest_num_frames = 0
@@ -1005,7 +1010,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         path = self.root / self.meta.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
         path.parent.mkdir(parents=True, exist_ok=True)
         if len(self.meta.image_keys) > 0:
-            datasets.Dataset.from_dict(df.to_dict(orient="list")).to_parquet(path)
+            to_parquet_with_hf_images(df, path)
         else:
             df.to_parquet(path)
 
@@ -1039,13 +1044,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
             shutil.move(str(ep_path), str(new_path))
         else:
             # Retrieve information from the latest video file
-            latest_ep = self.meta.episodes.with_format(
-                columns=[f"videos/{video_key}/chunk_index", f"videos/{video_key}/file_index"]
-            )[-1]
-            chunk_idx, file_idx = (
-                latest_ep[f"videos/{video_key}/chunk_index"],
-                latest_ep[f"videos/{video_key}/file_index"],
-            )
+            latest_ep = self.meta.episodes[-1]
+            chunk_idx = latest_ep[f"videos/{video_key}/chunk_index"]
+            file_idx = latest_ep[f"videos/{video_key}/file_index"]
 
             latest_path = self.root / self.meta.video_path.format(
                 video_key=video_key, chunk_index=chunk_idx, file_index=file_idx
@@ -1053,7 +1054,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
             latest_size_in_mb = get_video_size_in_mb(latest_path)
             latest_duration_in_s = get_video_duration_in_s(latest_path)
 
-            if latest_size_in_mb + ep_size_in_mb >= self.meta.files_size_in_mb:
+            if latest_size_in_mb + ep_size_in_mb >= self.meta.video_files_size_in_mb:
                 # Move temporary episode video to a new video file in the dataset
                 chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, self.meta.chunks_size)
                 new_path = self.root / self.meta.video_path.format(
@@ -1115,16 +1116,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         if self.image_writer is not None:
             self.image_writer.wait_until_done()
 
-    # TODO(rcadene): this method is currently not used
-    # def encode_videos(self) -> None:
-    #     """
-    #     Use ffmpeg to convert frames stored as png into mp4 videos.
-    #     Note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding,
-    #     since video encoding with ffmpeg is already using multithreading.
-    #     """
-    #     for ep_idx in range(self.meta.total_episodes):
-    #         self.encode_episode_videos(ep_idx)
-
     def _encode_temporary_episode_video(self, video_key: str, episode_index: int) -> dict:
         """
         Use ffmpeg to convert frames stored as png into mp4 videos.

lerobot/common/datasets/utils.py

@@ -50,7 +50,8 @@ from lerobot.common.utils.utils import is_valid_numpy_dtype_string
 from lerobot.configs.types import FeatureType, PolicyFeature
 
 DEFAULT_CHUNK_SIZE = 1000  # Max number of files per chunk
-DEFAULT_FILE_SIZE_IN_MB = 100.0  # Max size per file
+DEFAULT_DATA_FILE_SIZE_IN_MB = 100  # Max size per file
+DEFAULT_VIDEO_FILE_SIZE_IN_MB = 500  # Max size per file
 
 INFO_PATH = "meta/info.json"
 STATS_PATH = "meta/stats.json"
@@ -125,9 +126,8 @@ def load_nested_dataset(pq_dir: Path) -> Dataset:
         raise FileNotFoundError(f"Provided directory does not contain any parquet file: {pq_dir}")
 
     # TODO(rcadene): set num_proc to accelerate conversion to pyarrow
-    return concatenate_datasets(
-        [Dataset.from_parquet(str(path)) for path in sorted(pq_dir.glob("*/*.parquet"))]
-    )
+    datasets = [Dataset.from_parquet(str(path)) for path in paths]
+    return concatenate_datasets(datasets)
 
 
 def get_parquet_num_frames(parquet_path):
@@ -142,6 +142,7 @@ def get_video_size_in_mb(mp4_path: Path):
 
 
 def concat_video_files(paths_to_cat: list[Path], root: Path, video_key: str, chunk_idx: int, file_idx: int):
+    # TODO(rcadene): move to video_utils.py
     # TODO(rcadene): add docstring
     tmp_dir = Path(tempfile.mkdtemp(dir=root))
     # Create a text file with the list of files to concatenate
@@ -175,6 +176,7 @@ def concat_video_files(paths_to_cat: list[Path], root: Path, video_key: str, chu
 
 
 def get_video_duration_in_s(mp4_file: Path):
+    # TODO(rcadene): move to video_utils.py
     command = [
         "ffprobe",
         "-v",
@@ -290,7 +292,7 @@ def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]]:
 
 
 def write_hf_dataset(hf_dataset: Dataset, local_dir: Path):
-    if get_hf_dataset_size_in_mb(hf_dataset) > DEFAULT_FILE_SIZE_IN_MB:
+    if get_hf_dataset_size_in_mb(hf_dataset) > DEFAULT_DATA_FILE_SIZE_IN_MB:
         raise NotImplementedError("Contact a maintainer.")
 
     path = local_dir / DEFAULT_DATA_PATH.format(chunk_index=0, file_index=0)
@@ -310,7 +312,7 @@ def load_tasks(local_dir: Path):
 
 
 def write_episodes(episodes: Dataset, local_dir: Path):
-    if get_hf_dataset_size_in_mb(episodes) > DEFAULT_FILE_SIZE_IN_MB:
+    if get_hf_dataset_size_in_mb(episodes) > DEFAULT_DATA_FILE_SIZE_IN_MB:
         raise NotImplementedError("Contact a maintainer.")
 
     fpath = local_dir / DEFAULT_EPISODES_PATH.format(chunk_index=0, file_index=0)
@@ -318,9 +320,13 @@ def write_episodes(episodes: Dataset, local_dir: Path):
     episodes.to_parquet(fpath)
 
 
-def load_episodes(local_dir: Path):
-    hf_dataset = load_nested_dataset(local_dir / EPISODES_DIR)
-    return hf_dataset
+def load_episodes(local_dir: Path) -> datasets.Dataset:
+    episodes = load_nested_dataset(local_dir / EPISODES_DIR)
+    # Select episode features/columns containing references to episode data and videos
+    # (e.g. tasks, dataset_from_index, dataset_to_index, data/chunk_index, data/file_index, etc.)
+    # This is to speedup access to these data, instead of having to load episode stats.
+    episodes = episodes.select_columns([key for key in episodes.features if not key.startswith("stats/")])
+    return episodes
 
 
 def backward_compatible_episodes_stats(
@@ -528,9 +534,9 @@ def create_empty_dataset_info(
         "total_episodes": 0,
         "total_frames": 0,
         "total_tasks": 0,
-        "total_videos": 0,
         "chunks_size": DEFAULT_CHUNK_SIZE,
-        "files_size_in_mb": DEFAULT_FILE_SIZE_IN_MB,
+        "data_files_size_in_mb": DEFAULT_DATA_FILE_SIZE_IN_MB,
+        "video_files_size_in_mb": DEFAULT_VIDEO_FILE_SIZE_IN_MB,
         "fps": fps,
         "splits": {},
         "data_path": DEFAULT_DATA_PATH,
@@ -884,3 +890,11 @@ def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features:
             f"In episode_buffer not in features: {buffer_keys - set(features)}"
             f"In features not in episode_buffer: {set(features) - buffer_keys}"
         )
+
+
+def to_parquet_with_hf_images(df: pandas.DataFrame, path: Path):
+    """ This function correctly writes to parquet a panda DataFrame that contains images encoded by HF dataset.
+        This way, it can be loaded by HF dataset and correctly formated images are returned.
+    """
+    # TODO(qlhoest): replace this weird synthax by `df.to_parquet(path)` only
+    datasets.Dataset.from_dict(df.to_dict(orient="list")).to_parquet(path)

lerobot/common/datasets/v30/convert_dataset_v21_to_v30.py

@@ -24,8 +24,9 @@ from typing import Any
 
 import jsonlines
 import pandas as pd
+import pyarrow as pa
 import tqdm
-from datasets import Dataset
+from datasets import Dataset, Features, Image
 from huggingface_hub import HfApi, snapshot_download
 from requests import HTTPError
 
@@ -34,8 +35,9 @@ from lerobot.common.datasets.compute_stats import aggregate_stats
 from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
 from lerobot.common.datasets.utils import (
     DEFAULT_CHUNK_SIZE,
+    DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_DATA_PATH,
-    DEFAULT_FILE_SIZE_IN_MB,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
     DEFAULT_VIDEO_PATH,
     cast_stats_to_numpy,
     concat_video_files,
@@ -137,7 +139,7 @@ def convert_tasks(root, new_root):
     write_tasks(df_tasks, new_root)
 
 
-def concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx):
+def concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys):
     # TODO(rcadene): to save RAM use Dataset.from_parquet(file) and concatenate_datasets
     dataframes = [pd.read_parquet(file) for file in paths_to_cat]
     # Concatenate all DataFrames along rows
@@ -145,13 +147,25 @@ def concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx):
 
     path = new_root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
     path.parent.mkdir(parents=True, exist_ok=True)
-    concatenated_df.to_parquet(path, index=False)
+
+    if len(image_keys) > 0:
+        schema = pa.Schema.from_pandas(concatenated_df)
+        features = Features.from_arrow_schema(schema)
+        for key in image_keys:
+            features[key] = Image()
+        schema = features.arrow_schema
+    else:
+        schema = None
+
+    concatenated_df.to_parquet(path, index=False, schema=schema)
 
 
 def convert_data(root, new_root):
     data_dir = root / "data"
     ep_paths = sorted(data_dir.glob("*/*.parquet"))
 
+    image_keys = get_image_keys(root)
+
     ep_idx = 0
     chunk_idx = 0
     file_idx = 0
@@ -174,11 +188,11 @@ def convert_data(root, new_root):
         episodes_metadata.append(ep_metadata)
         ep_idx += 1
 
-        if size_in_mb < DEFAULT_FILE_SIZE_IN_MB:
+        if size_in_mb < DEFAULT_DATA_FILE_SIZE_IN_MB:
             paths_to_cat.append(ep_path)
             continue
 
-        concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx)
+        concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
 
         # Reset for the next file
         size_in_mb = ep_size_in_mb
@@ -189,7 +203,7 @@ def convert_data(root, new_root):
 
     # Write remaining data if any
     if paths_to_cat:
-        concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx)
+        concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
 
     return episodes_metadata
 
@@ -197,16 +211,22 @@ def convert_data(root, new_root):
 def get_video_keys(root):
     info = load_info(root)
     features = info["features"]
-    image_keys = [key for key, ft in features.items() if ft["dtype"] == "image"]
-    if len(image_keys) != 0:
-        raise NotImplementedError()
-
     video_keys = [key for key, ft in features.items() if ft["dtype"] == "video"]
     return video_keys
 
 
+def get_image_keys(root):
+    info = load_info(root)
+    features = info["features"]
+    image_keys = [key for key, ft in features.items() if ft["dtype"] == "image"]
+    return image_keys
+
+
 def convert_videos(root: Path, new_root: Path):
     video_keys = get_video_keys(root)
+    if len(video_keys) == 0:
+        return None
+
     video_keys = sorted(video_keys)
 
     eps_metadata_per_cam = []
@@ -263,7 +283,7 @@ def convert_videos_of_camera(root: Path, new_root: Path, video_key):
         episodes_metadata.append(ep_metadata)
         ep_idx += 1
 
-        if size_in_mb < DEFAULT_FILE_SIZE_IN_MB:
+        if size_in_mb < DEFAULT_VIDEO_FILE_SIZE_IN_MB:
             paths_to_cat.append(ep_path)
             continue
 
@@ -284,24 +304,32 @@ def convert_videos_of_camera(root: Path, new_root: Path, video_key):
 
 
 def generate_episode_metadata_dict(
-    episodes_legacy_metadata, episodes_metadata, episodes_videos, episodes_stats
+    episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_videos=None
 ):
-    for ep_legacy_metadata, ep_metadata, ep_video, ep_stats, ep_idx_stats in zip(
-        episodes_legacy_metadata.values(),
-        episodes_metadata,
-        episodes_videos,
-        episodes_stats.values(),
-        episodes_stats.keys(),
-        strict=False,
-    ):
-        ep_idx = ep_legacy_metadata["episode_index"]
-        ep_idx_data = ep_metadata["episode_index"]
-        ep_idx_video = ep_video["episode_index"]
+    num_episodes = len(episodes_metadata)
+    episodes_legacy_metadata_vals = list(episodes_legacy_metadata.values())
+    episodes_stats_vals = list(episodes_stats.values())
+    episodes_stats_keys = list(episodes_stats.keys())
 
-        if len({ep_idx, ep_idx_data, ep_idx_video, ep_idx_stats}) != 1:
-            raise ValueError(
-                f"Number of episodes is not the same ({ep_idx=},{ep_idx_data=},{ep_idx_video=},{ep_idx_stats=})."
-            )
+    for i in range(num_episodes):
+        ep_legacy_metadata = episodes_legacy_metadata_vals[i]
+        ep_metadata = episodes_metadata[i]
+        ep_stats = episodes_stats_vals[i]
+
+        ep_ids_set = {
+            ep_legacy_metadata["episode_index"],
+            ep_metadata["episode_index"],
+            episodes_stats_keys[i],
+        }
+
+        if episodes_videos is None:
+            ep_video = {}
+        else:
+            ep_video = episodes_videos[i]
+            ep_ids_set.add(ep_video["episode_index"])
+
+        if len(ep_ids_set) != 1:
+            raise ValueError(f"Number of episodes is not the same ({ep_ids_set}).")
 
         ep_dict = {**ep_metadata, **ep_video, **ep_legacy_metadata, **flatten_dict({"stats": ep_stats})}
         ep_dict["meta/episodes/chunk_index"] = 0
@@ -309,21 +337,20 @@ def generate_episode_metadata_dict(
         yield ep_dict
 
 
-def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_metadata):
+def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_metadata=None):
     episodes_legacy_metadata = legacy_load_episodes(root)
     episodes_stats = legacy_load_episodes_stats(root)
 
-    num_eps = len(episodes_legacy_metadata)
-    num_eps_metadata = len(episodes_metadata)
-    num_eps_video_metadata = len(episodes_video_metadata)
-    if len({num_eps, num_eps_metadata, num_eps_video_metadata}) != 1:
-        raise ValueError(
-            f"Number of episodes is not the same ({num_eps=},{num_eps_metadata=},{num_eps_video_metadata=})."
-        )
+    num_eps_set = {len(episodes_legacy_metadata), len(episodes_metadata)}
+    if episodes_video_metadata is not None:
+        num_eps_set.add(len(episodes_video_metadata))
+
+    if len(num_eps_set) != 1:
+        raise ValueError(f"Number of episodes is not the same ({num_eps_set}).")
 
     ds_episodes = Dataset.from_generator(
         lambda: generate_episode_metadata_dict(
-            episodes_legacy_metadata, episodes_metadata, episodes_video_metadata, episodes_stats
+            episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_video_metadata
         )
     )
     write_episodes(ds_episodes, new_root)
@@ -337,8 +364,8 @@ def convert_info(root, new_root):
     info["codebase_version"] = "v3.0"
     del info["total_chunks"]
     del info["total_videos"]
-    info["files_size_in_mb"] = DEFAULT_FILE_SIZE_IN_MB
-    # TODO(rcadene): chunk- or chunk_ or file- or file_
+    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_path"] = DEFAULT_DATA_PATH
     info["video_path"] = DEFAULT_VIDEO_PATH
     info["fps"] = float(info["fps"])

lerobot/common/datasets/video_utils.py

@@ -155,6 +155,7 @@ 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_]
 

lerobot/common/robot_devices/cameras/intelrealsense.py

@@ -512,13 +512,13 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--width",
-        type=str,
+        type=int,
         default=640,
         help="Set the width for all cameras. If not provided, use the default width of each camera.",
     )
     parser.add_argument(
         "--height",
-        type=str,
+        type=int,
         default=480,
         help="Set the height for all cameras. If not provided, use the default height of each camera.",
     )

lerobot/common/robot_devices/cameras/opencv.py

@@ -492,13 +492,13 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--width",
-        type=str,
+        type=int,
         default=None,
         help="Set the width for all cameras. If not provided, use the default width of each camera.",
     )
     parser.add_argument(
         "--height",
-        type=str,
+        type=int,
         default=None,
         help="Set the height for all cameras. If not provided, use the default height of each camera.",
     )

lerobot/scripts/visualize_dataset.py

@@ -79,8 +79,8 @@ from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 
 class EpisodeSampler(torch.utils.data.Sampler):
     def __init__(self, dataset: LeRobotDataset, episode_index: int):
-        from_idx = dataset.meta.episodes["dataset_from_index"][episode_index].item()
-        to_idx = dataset.meta.episodes["dataset_to_index"][episode_index].item()
+        from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
+        to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
         self.frame_ids = range(from_idx, to_idx)
 
     def __iter__(self) -> Iterator:
@@ -283,7 +283,7 @@ def main():
     tolerance_s = kwargs.pop("tolerance_s")
 
     logging.info("Loading dataset")
-    dataset = LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
+    dataset = LeRobotDataset(repo_id, episodes=[args.episode_index], root=root, tolerance_s=tolerance_s)
 
     visualize_dataset(dataset, **vars(args))
 

lerobot/scripts/visualize_dataset_html.py

@@ -174,7 +174,10 @@ def run_server(
                 dataset.meta.get_video_file_path(episode_id, key) for key in dataset.meta.video_keys
             ]
             videos_info = [
-                {"url": url_for("static", filename=video_path), "filename": video_path.parent.name}
+                {
+                    "url": url_for("static", filename=str(video_path).replace("\\", "/")),
+                    "filename": video_path.parent.name,
+                }
                 for video_path in video_paths
             ]
             tasks = dataset.meta.episodes[episode_id]["tasks"]
@@ -381,7 +384,7 @@ def visualize_dataset_html(
         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())
+                ln_videos_dir.symlink_to((dataset.root / "videos").resolve().as_posix())
 
         if serve:
             run_server(dataset, episodes, host, port, static_dir, template_dir)

tests/datasets/test_aggregate.py

@@ -1,20 +1,185 @@
+import torch
+
 from lerobot.common.datasets.aggregate import aggregate_datasets
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 from tests.fixtures.constants import DUMMY_REPO_ID
 
 
+def assert_episode_and_frame_counts(aggr_ds, expected_episodes, expected_frames):
+    """Test that total number of episodes and frames are correctly aggregated."""
+    assert aggr_ds.num_episodes == expected_episodes, (
+        f"Expected {expected_episodes} episodes, got {aggr_ds.num_episodes}"
+    )
+    assert aggr_ds.num_frames == expected_frames, (
+        f"Expected {expected_frames} frames, got {aggr_ds.num_frames}"
+    )
+
+
+def assert_dataset_content_integrity(aggr_ds, ds_0, ds_1):
+    """Test that the content of both datasets is preserved correctly in the aggregated dataset."""
+    # Test first part of dataset corresponds to ds_0, check first item (index 0) matches ds_0[0]
+    aggr_first_item = aggr_ds[0]
+    ds_0_first_item = ds_0[0]
+
+    # Compare all keys except episode_index and index which should be updated
+    for key in ds_0_first_item:
+        if key not in ["episode_index", "index"]:
+            # Handle both tensor and non-tensor data
+            if torch.is_tensor(aggr_first_item[key]) and torch.is_tensor(ds_0_first_item[key]):
+                assert torch.allclose(aggr_first_item[key], ds_0_first_item[key], atol=1e-6), (
+                    f"First item key '{key}' doesn't match between aggregated and ds_0"
+                )
+            else:
+                assert aggr_first_item[key] == ds_0_first_item[key], (
+                    f"First item key '{key}' doesn't match between aggregated and ds_0"
+                )
+
+    # Check last item of ds_0 part (index len(ds_0)-1) matches ds_0[-1]
+    aggr_ds_0_last_item = aggr_ds[len(ds_0) - 1]
+    ds_0_last_item = ds_0[-1]
+
+    for key in ds_0_last_item:
+        if key not in ["episode_index", "index"]:
+            # Handle both tensor and non-tensor data
+            if torch.is_tensor(aggr_ds_0_last_item[key]) and torch.is_tensor(ds_0_last_item[key]):
+                assert torch.allclose(aggr_ds_0_last_item[key], ds_0_last_item[key], atol=1e-6), (
+                    f"Last ds_0 item key '{key}' doesn't match between aggregated and ds_0"
+                )
+            else:
+                assert aggr_ds_0_last_item[key] == ds_0_last_item[key], (
+                    f"Last ds_0 item key '{key}' doesn't match between aggregated and ds_0"
+                )
+
+    # Test second part of dataset corresponds to ds_1
+    # Check first item of ds_1 part (index len(ds_0)) matches ds_1[0]
+    aggr_ds_1_first_item = aggr_ds[len(ds_0)]
+    ds_1_first_item = ds_1[0]
+
+    for key in ds_1_first_item:
+        if key not in ["episode_index", "index"]:
+            # Handle both tensor and non-tensor data
+            if torch.is_tensor(aggr_ds_1_first_item[key]) and torch.is_tensor(ds_1_first_item[key]):
+                assert torch.allclose(aggr_ds_1_first_item[key], ds_1_first_item[key], atol=1e-6), (
+                    f"First ds_1 item key '{key}' doesn't match between aggregated and ds_1"
+                )
+            else:
+                assert aggr_ds_1_first_item[key] == ds_1_first_item[key], (
+                    f"First ds_1 item key '{key}' doesn't match between aggregated and ds_1"
+                )
+
+    # Check last item matches ds_1[-1]
+    aggr_last_item = aggr_ds[-1]
+    ds_1_last_item = ds_1[-1]
+
+    for key in ds_1_last_item:
+        if key not in ["episode_index", "index"]:
+            # Handle both tensor and non-tensor data
+            if torch.is_tensor(aggr_last_item[key]) and torch.is_tensor(ds_1_last_item[key]):
+                assert torch.allclose(aggr_last_item[key], ds_1_last_item[key], atol=1e-6), (
+                    f"Last item key '{key}' doesn't match between aggregated and ds_1"
+                )
+            else:
+                assert aggr_last_item[key] == ds_1_last_item[key], (
+                    f"Last item key '{key}' doesn't match between aggregated and ds_1"
+                )
+
+
+def assert_metadata_consistency(aggr_ds, ds_0, ds_1):
+    """Test that metadata is correctly aggregated."""
+    # Test basic info
+    assert aggr_ds.fps == ds_0.fps == ds_1.fps, "FPS should be the same across all datasets"
+    assert aggr_ds.meta.info["robot_type"] == ds_0.meta.info["robot_type"] == ds_1.meta.info["robot_type"], (
+        "Robot type should be the same"
+    )
+
+    # Test features are the same
+    assert aggr_ds.features == ds_0.features == ds_1.features, "Features should be the same"
+
+    # Test tasks aggregation
+    expected_tasks = set(ds_0.meta.tasks.index) | set(ds_1.meta.tasks.index)
+    actual_tasks = set(aggr_ds.meta.tasks.index)
+    assert actual_tasks == expected_tasks, f"Expected tasks {expected_tasks}, got {actual_tasks}"
+
+
+def assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1):
+    """Test that episode indices are correctly updated after aggregation."""
+    # ds_0 episodes should have episode_index 0 to ds_0.num_episodes-1
+    for i in range(len(ds_0)):
+        assert aggr_ds[i]["episode_index"] < ds_0.num_episodes, (
+            f"Episode index {aggr_ds[i]['episode_index']} at position {i} should be < {ds_0.num_episodes}"
+        )
+
+    def ds1_episodes_condition(ep_idx):
+        return (ep_idx >= ds_0.num_episodes) and (ep_idx < ds_0.num_episodes + ds_1.num_episodes)
+
+    # ds_1 episodes should have episode_index ds_0.num_episodes to total_episodes-1
+    for i in range(len(ds_0), len(ds_0) + len(ds_1)):
+        expected_min_episode_idx = ds_0.num_episodes
+        assert ds1_episodes_condition(aggr_ds[i]["episode_index"]), (
+            f"Episode index {aggr_ds[i]['episode_index']} at position {i} should be >= {expected_min_episode_idx}"
+        )
+
+
+def assert_video_frames_integrity(aggr_ds, ds_0, ds_1):
+    """Test that video frames are correctly preserved and frame indices are updated."""
+
+    def visual_frames_equal(frame1, frame2):
+        return torch.allclose(frame1, frame2)
+
+    video_keys = list(
+        filter(
+            lambda key: aggr_ds.meta.info["features"][key]["dtype"] == "video",
+            aggr_ds.meta.info["features"].keys(),
+        )
+    )
+
+    # Test the section corresponding to the first dataset (ds_0)
+    for i in range(len(ds_0)):
+        assert aggr_ds[i]["index"] == i, (
+            f"Frame index at position {i} should be {i}, but got {aggr_ds[i]['index']}"
+        )
+        for key in video_keys:
+            assert visual_frames_equal(aggr_ds[i][key], ds_0[i][key]), (
+                f"Visual frames at position {i} should be equal between aggregated and ds_0"
+            )
+
+    # Test the section corresponding to the second dataset (ds_1)
+    for i in range(len(ds_0), len(ds_0) + len(ds_1)):
+        # The frame index in the aggregated dataset should also match its position.
+        assert aggr_ds[i]["index"] == i, (
+            f"Frame index at position {i} should be {i}, but got {aggr_ds[i]['index']}"
+        )
+        for key in video_keys:
+            assert visual_frames_equal(aggr_ds[i][key], ds_1[i - len(ds_0)][key]), (
+                f"Visual frames at position {i} should be equal between aggregated and ds_1"
+            )
+
+
+def assert_dataset_iteration_works(aggr_ds):
+    """Test that we can iterate through the entire dataset without errors."""
+    for _ in aggr_ds:
+        pass
+
+
 def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
+    """Test basic aggregation functionality with standard parameters."""
+    ds_0_num_frames = 400
+    ds_1_num_frames = 400
+    ds_0_num_episodes = 10
+    ds_1_num_episodes = 10
+
+    # Create two datasets with different number of frames and episodes
     ds_0 = lerobot_dataset_factory(
         root=tmp_path / "test_0",
         repo_id=f"{DUMMY_REPO_ID}_0",
-        total_episodes=10,
-        total_frames=400,
+        total_episodes=ds_0_num_episodes,
+        total_frames=ds_0_num_frames,
     )
     ds_1 = lerobot_dataset_factory(
         root=tmp_path / "test_1",
         repo_id=f"{DUMMY_REPO_ID}_1",
-        total_episodes=10,
-        total_frames=400,
+        total_episodes=ds_1_num_episodes,
+        total_frames=ds_1_num_frames,
     )
 
     aggregate_datasets(
@@ -25,5 +190,14 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
     )
 
     aggr_ds = LeRobotDataset(f"{DUMMY_REPO_ID}_aggr", root=tmp_path / "test_aggr")
-    for _ in aggr_ds:
-        pass
+
+    # Run all assertion functions
+    expected_total_episodes = ds_0.num_episodes + ds_1.num_episodes
+    expected_total_frames = ds_0.num_frames + ds_1.num_frames
+
+    assert_episode_and_frame_counts(aggr_ds, expected_total_episodes, expected_total_frames)
+    assert_dataset_content_integrity(aggr_ds, ds_0, ds_1)
+    assert_metadata_consistency(aggr_ds, ds_0, ds_1)
+    assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1)
+    assert_video_frames_integrity(aggr_ds, ds_0, ds_1)
+    assert_dataset_iteration_works(aggr_ds)

tests/datasets/test_datasets.py

@@ -13,10 +13,8 @@
 # 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 json
 import logging
 import re
-from copy import deepcopy
 from itertools import chain
 from pathlib import Path
 
@@ -36,8 +34,6 @@ from lerobot.common.datasets.lerobot_dataset import (
 )
 from lerobot.common.datasets.utils import (
     create_branch,
-    flatten_dict,
-    unflatten_dict,
 )
 from lerobot.common.envs.factory import make_env_config
 from lerobot.common.policies.factory import make_policy_config
@@ -75,7 +71,7 @@ def test_same_attributes_defined(tmp_path, lerobot_dataset_factory):
     dataset_create = LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=30, robot=robot, root=root_create)
 
     root_init = tmp_path / "init"
-    dataset_init = lerobot_dataset_factory(root=root_init)
+    dataset_init = lerobot_dataset_factory(root=root_init, total_episodes=1, total_frames=1)
 
     init_attr = set(vars(dataset_init).keys())
     create_attr = set(vars(dataset_create).keys())
@@ -100,6 +96,25 @@ def test_dataset_initialization(tmp_path, lerobot_dataset_factory):
     assert dataset.num_frames == len(dataset)
 
 
+# TODO(rcadene, aliberts): do not run LeRobotDataset.create, instead refactor LeRobotDatasetMetadata.create
+# and test the small resulting function that validates the features
+def test_dataset_feature_with_forward_slash_raises_error():
+    # make sure dir does not exist
+    from lerobot.common.constants import HF_LEROBOT_HOME
+
+    dataset_dir = HF_LEROBOT_HOME / "lerobot/test/with/slash"
+    # make sure does not exist
+    if dataset_dir.exists():
+        dataset_dir.rmdir()
+
+    with pytest.raises(ValueError):
+        LeRobotDataset.create(
+            repo_id="lerobot/test/with/slash",
+            fps=30,
+            features={"a/b": {"dtype": "float32", "shape": 2, "names": None}},
+        )
+
+
 def test_add_frame_missing_task(tmp_path, empty_lerobot_dataset_factory):
     features = {"state": {"dtype": "float32", "shape": (1,), "names": None}}
     dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
@@ -329,6 +344,13 @@ def test_image_array_to_pil_image_wrong_range_float_0_255():
 # - [ ] test push_to_hub
 # - [ ] test smaller methods
 
+# TODO(rcadene):
+# - [ ] fix code so that old test_factory + backward pass
+# - [ ] write new unit tests to test save_episode + getitem
+#   - [ ] save_episode : case where new dataset, concatenate same file, write new file (meta/episodes, data, videos)
+#   - [ ]
+# - [ ] remove old tests
+
 
 @pytest.mark.parametrize(
     "env_name, repo_id, policy_name",
@@ -436,30 +458,6 @@ def test_multidataset_frames():
             assert torch.equal(sub_dataset_item[k], dataset_item[k])
 
 
-# TODO(aliberts): Move to more appropriate location
-def test_flatten_unflatten_dict():
-    d = {
-        "obs": {
-            "min": 0,
-            "max": 1,
-            "mean": 2,
-            "std": 3,
-        },
-        "action": {
-            "min": 4,
-            "max": 5,
-            "mean": 6,
-            "std": 7,
-        },
-    }
-
-    original_d = deepcopy(d)
-    d = unflatten_dict(flatten_dict(d))
-
-    # test equality between nested dicts
-    assert json.dumps(original_d, sort_keys=True) == json.dumps(d, sort_keys=True), f"{original_d} != {d}"
-
-
 @pytest.mark.parametrize(
     "repo_id",
     [
@@ -569,20 +567,3 @@ def test_create_branch():
 
     # Clean
     api.delete_repo(repo_id, repo_type=repo_type)
-
-
-def test_dataset_feature_with_forward_slash_raises_error():
-    # make sure dir does not exist
-    from lerobot.common.constants import HF_LEROBOT_HOME
-
-    dataset_dir = HF_LEROBOT_HOME / "lerobot/test/with/slash"
-    # make sure does not exist
-    if dataset_dir.exists():
-        dataset_dir.rmdir()
-
-    with pytest.raises(ValueError):
-        LeRobotDataset.create(
-            repo_id="lerobot/test/with/slash",
-            fps=30,
-            features={"a/b": {"dtype": "float32", "shape": 2, "names": None}},
-        )

tests/datasets/test_delta_timestamps.py

@@ -56,8 +56,8 @@ def calculate_episode_data_index(hf_dataset: datasets.Dataset) -> dict[str, np.n
 def synced_timestamps_factory(hf_dataset_factory):
     def _create_synced_timestamps(fps: int = 30) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         hf_dataset = hf_dataset_factory(fps=fps)
-        timestamps = hf_dataset["timestamp"].numpy()
-        episode_indices = hf_dataset["episode_index"].numpy()
+        timestamps = torch.stack(hf_dataset["timestamp"]).numpy()
+        episode_indices = torch.stack(hf_dataset["episode_index"]).numpy()
         episode_data_index = calculate_episode_data_index(hf_dataset)
         return timestamps, episode_indices, episode_data_index
 

tests/datasets/test_utils.py

@@ -14,12 +14,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
+from copy import deepcopy
+
 import torch
 from datasets import Dataset
 from huggingface_hub import DatasetCard
 
 from lerobot.common.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
-from lerobot.common.datasets.utils import create_lerobot_dataset_card, hf_transform_to_torch
+from lerobot.common.datasets.utils import (
+    create_lerobot_dataset_card,
+    flatten_dict,
+    hf_transform_to_torch,
+    unflatten_dict,
+)
 
 
 def test_default_parameters():
@@ -53,3 +61,26 @@ def test_calculate_episode_data_index():
     episode_data_index = calculate_episode_data_index(dataset)
     assert torch.equal(episode_data_index["from"], torch.tensor([0, 2, 3]))
     assert torch.equal(episode_data_index["to"], torch.tensor([2, 3, 6]))
+
+
+def test_flatten_unflatten_dict():
+    d = {
+        "obs": {
+            "min": 0,
+            "max": 1,
+            "mean": 2,
+            "std": 3,
+        },
+        "action": {
+            "min": 4,
+            "max": 5,
+            "mean": 6,
+            "std": 7,
+        },
+    }
+
+    original_d = deepcopy(d)
+    d = unflatten_dict(flatten_dict(d))
+
+    # test equality between nested dicts
+    assert json.dumps(original_d, sort_keys=True) == json.dumps(d, sort_keys=True), f"{original_d} != {d}"

tests/fixtures/constants.py

@@ -29,8 +29,8 @@ DUMMY_MOTOR_FEATURES = {
     },
 }
 DUMMY_CAMERA_FEATURES = {
-    "laptop": {"shape": (480, 640, 3), "names": ["height", "width", "channels"], "info": None},
-    "phone": {"shape": (480, 640, 3), "names": ["height", "width", "channels"], "info": None},
+    "laptop": {"shape": (64, 96, 3), "names": ["height", "width", "channels"], "info": None},
+    "phone": {"shape": (64, 96, 3), "names": ["height", "width", "channels"], "info": None},
 }
 DEFAULT_FPS = 30
 DUMMY_VIDEO_INFO = {

tests/fixtures/dataset_factories.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import random
+import shutil
 from functools import partial
 from pathlib import Path
 from typing import Protocol
@@ -28,13 +29,16 @@ from datasets import Dataset
 from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset, LeRobotDatasetMetadata
 from lerobot.common.datasets.utils import (
     DEFAULT_CHUNK_SIZE,
+    DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_DATA_PATH,
     DEFAULT_FEATURES,
-    DEFAULT_FILE_SIZE_IN_MB,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
     DEFAULT_VIDEO_PATH,
     flatten_dict,
     get_hf_features_from_features,
+    hf_transform_to_torch,
 )
+from lerobot.common.datasets.video_utils import encode_video_frames
 from tests.fixtures.constants import (
     DEFAULT_FPS,
     DUMMY_CAMERA_FEATURES,
@@ -65,15 +69,49 @@ def img_tensor_factory():
 
 @pytest.fixture(scope="session")
 def img_array_factory():
-    def _create_img_array(height=100, width=100, channels=3, dtype=np.uint8) -> np.ndarray:
-        if np.issubdtype(dtype, np.unsignedinteger):
-            # Int array in [0, 255] range
-            img_array = np.random.randint(0, 256, size=(height, width, channels), dtype=dtype)
-        elif np.issubdtype(dtype, np.floating):
-            # Float array in [0, 1] range
-            img_array = np.random.rand(height, width, channels).astype(dtype)
+    def _create_img_array(height=100, width=100, channels=3, dtype=np.uint8, content=None) -> np.ndarray:
+        if content is None:
+            # Original random noise behavior
+            if np.issubdtype(dtype, np.unsignedinteger):
+                # Int array in [0, 255] range
+                img_array = np.random.randint(0, 256, size=(height, width, channels), dtype=dtype)
+            elif np.issubdtype(dtype, np.floating):
+                # Float array in [0, 1] range
+                img_array = np.random.rand(height, width, channels).astype(dtype)
+            else:
+                raise ValueError(dtype)
         else:
-            raise ValueError(dtype)
+            # Create image with text content using OpenCV
+            import cv2
+
+            # Create white background
+            img_array = np.ones((height, width, channels), dtype=np.uint8) * 255
+
+            # Font settings
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            font_scale = max(0.5, height / 200)  # Scale font with image size
+            font_color = (0, 0, 0)  # Black text
+            thickness = max(1, int(height / 100))
+
+            # Get text size to center it
+            text_size = cv2.getTextSize(content, font, font_scale, thickness)[0]
+            text_x = (width - text_size[0]) // 2
+            text_y = (height + text_size[1]) // 2
+
+            # Put text on image
+            cv2.putText(img_array, content, (text_x, text_y), font, font_scale, font_color, thickness)
+
+            # Handle single channel case
+            if channels == 1:
+                img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)
+                img_array = img_array[:, :, np.newaxis]
+
+            # Convert to target dtype
+            if np.issubdtype(dtype, np.floating):
+                img_array = img_array.astype(dtype) / 255.0
+            else:
+                img_array = img_array.astype(dtype)
+
         return img_array
 
     return _create_img_array
@@ -121,7 +159,8 @@ def info_factory(features_factory):
         total_tasks: int = 0,
         total_videos: int = 0,
         chunks_size: int = DEFAULT_CHUNK_SIZE,
-        files_size_in_mb: float = DEFAULT_FILE_SIZE_IN_MB,
+        data_files_size_in_mb: float = DEFAULT_DATA_FILE_SIZE_IN_MB,
+        video_files_size_in_mb: float = DEFAULT_VIDEO_FILE_SIZE_IN_MB,
         data_path: str = DEFAULT_DATA_PATH,
         video_path: str = DEFAULT_VIDEO_PATH,
         motor_features: dict = DUMMY_MOTOR_FEATURES,
@@ -137,7 +176,8 @@ def info_factory(features_factory):
             "total_tasks": total_tasks,
             "total_videos": total_videos,
             "chunks_size": chunks_size,
-            "files_size_in_mb": files_size_in_mb,
+            "data_files_size_in_mb": data_files_size_in_mb,
+            "video_files_size_in_mb": video_files_size_in_mb,
             "fps": fps,
             "splits": {},
             "data_path": data_path,
@@ -298,6 +338,41 @@ def episodes_factory(tasks_factory, stats_factory):
     return _create_episodes
 
 
+@pytest.fixture(scope="session")
+def create_videos(info_factory, img_array_factory):
+    def _create_video_directory(
+        root: Path,
+        info: dict | None = None,
+        total_episodes: int = 3,
+        total_frames: int = 150,
+        total_tasks: int = 1,
+    ):
+        if info is None:
+            info = info_factory(
+                total_episodes=total_episodes, total_frames=total_frames, total_tasks=total_tasks
+            )
+
+        video_feats = {key: feats for key, feats in info["features"].items() if feats["dtype"] == "video"}
+        for key, ft in video_feats.items():
+            # create and save images with identifiable content
+            tmp_dir = root / "tmp_images"
+            tmp_dir.mkdir(parents=True, exist_ok=True)
+            for frame_index in range(info["total_frames"]):
+                content = f"{key}-{frame_index}"
+                img = img_array_factory(height=ft["shape"][0], width=ft["shape"][1], content=content)
+                pil_img = PIL.Image.fromarray(img)
+                path = tmp_dir / f"frame-{frame_index:06d}.png"
+                pil_img.save(path)
+
+            video_path = root / DEFAULT_VIDEO_PATH.format(video_key=key, chunk_index=0, file_index=0)
+            video_path.parent.mkdir(parents=True, exist_ok=True)
+            # Use the global fps from info, not video-specific fps which might not exist
+            encode_video_frames(tmp_dir, video_path, fps=info["fps"])
+            shutil.rmtree(tmp_dir)
+
+    return _create_video_directory
+
+
 @pytest.fixture(scope="session")
 def hf_dataset_factory(features_factory, tasks_factory, episodes_factory, img_array_factory):
     def _create_hf_dataset(
@@ -334,8 +409,8 @@ def hf_dataset_factory(features_factory, tasks_factory, episodes_factory, img_ar
         for key, ft in features.items():
             if ft["dtype"] == "image":
                 robot_cols[key] = [
-                    img_array_factory(height=ft["shapes"][1], width=ft["shapes"][0])
-                    for _ in range(len(index_col))
+                    img_array_factory(height=ft["shape"][1], width=ft["shape"][0], content=f"{key}-{i}")
+                    for i in range(len(index_col))
                 ]
             elif ft["shape"][0] > 1 and ft["dtype"] != "video":
                 robot_cols[key] = np.random.random((len(index_col), ft["shape"][0])).astype(ft["dtype"])
@@ -352,7 +427,7 @@ def hf_dataset_factory(features_factory, tasks_factory, episodes_factory, img_ar
             },
             features=hf_features,
         )
-        dataset.set_format("torch")
+        dataset.set_transform(hf_transform_to_torch)
         return dataset
 
     return _create_hf_dataset
@@ -428,6 +503,7 @@ def lerobot_dataset_factory(
         total_frames: int = 150,
         total_tasks: int = 1,
         multi_task: bool = False,
+        use_videos: bool = True,
         info: dict | None = None,
         stats: dict | None = None,
         tasks: pd.DataFrame | None = None,
@@ -435,9 +511,13 @@ def lerobot_dataset_factory(
         hf_dataset: datasets.Dataset | None = None,
         **kwargs,
     ) -> LeRobotDataset:
+        # Instantiate objects
         if info is None:
             info = info_factory(
-                total_episodes=total_episodes, total_frames=total_frames, total_tasks=total_tasks
+                total_episodes=total_episodes,
+                total_frames=total_frames,
+                total_tasks=total_tasks,
+                use_videos=use_videos,
             )
         if stats is None:
             stats = stats_factory(features=info["features"])
@@ -454,9 +534,12 @@ def lerobot_dataset_factory(
                 tasks=tasks,
                 multi_task=multi_task,
             )
-        if not hf_dataset:
-            hf_dataset = hf_dataset_factory(tasks=tasks, episodes=episodes_metadata, fps=info["fps"])
+        if hf_dataset is None:
+            hf_dataset = hf_dataset_factory(
+                features=info["features"], tasks=tasks, episodes=episodes_metadata, fps=info["fps"]
+            )
 
+        # Write data on disk
         mock_snapshot_download = mock_snapshot_download_factory(
             info=info,
             stats=stats,

tests/fixtures/files.py

@@ -48,16 +48,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):

tests/fixtures/hub.py

@@ -22,6 +22,7 @@ from lerobot.common.datasets.utils import (
     DEFAULT_DATA_PATH,
     DEFAULT_EPISODES_PATH,
     DEFAULT_TASKS_PATH,
+    DEFAULT_VIDEO_PATH,
     INFO_PATH,
     STATS_PATH,
 )
@@ -40,6 +41,7 @@ def mock_snapshot_download_factory(
     create_episodes,
     hf_dataset_factory,
     create_hf_dataset,
+    create_videos,
 ):
     """
     This factory allows to patch snapshot_download such that when called, it will create expected files rather
@@ -91,40 +93,48 @@ def mock_snapshot_download_factory(
                 DEFAULT_DATA_PATH.format(chunk_index=0, file_index=0),
             ]
 
+            video_keys = [key for key, feats in info["features"].items() if feats["dtype"] == "video"]
+            for key in video_keys:
+                all_files.append(DEFAULT_VIDEO_PATH.format(video_key=key, chunk_index=0, file_index=0))
+
             allowed_files = filter_repo_objects(
                 all_files, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns
             )
 
-            has_info = False
-            has_tasks = False
-            has_episodes = False
-            has_stats = False
-            has_data = False
+            request_info = False
+            request_tasks = False
+            request_episodes = False
+            request_stats = False
+            request_data = False
+            request_videos = False
             for rel_path in allowed_files:
                 if rel_path.startswith("meta/info.json"):
-                    has_info = True
+                    request_info = True
                 elif rel_path.startswith("meta/stats"):
-                    has_stats = True
+                    request_stats = True
                 elif rel_path.startswith("meta/tasks"):
-                    has_tasks = True
+                    request_tasks = True
                 elif rel_path.startswith("meta/episodes"):
-                    has_episodes = True
+                    request_episodes = True
                 elif rel_path.startswith("data/"):
-                    has_data = True
+                    request_data = True
+                elif rel_path.startswith("videos/"):
+                    request_videos = True
                 else:
                     raise ValueError(f"{rel_path} not supported.")
 
-            if has_info:
+            if request_info:
                 create_info(local_dir, info)
-            if has_stats:
+            if request_stats:
                 create_stats(local_dir, stats)
-            if has_tasks:
+            if request_tasks:
                 create_tasks(local_dir, tasks)
-            if has_episodes:
+            if request_episodes:
                 create_episodes(local_dir, episodes)
-            # TODO(rcadene): create_videos?
-            if has_data:
+            if request_data:
                 create_hf_dataset(local_dir, hf_dataset)
+            if request_videos:
+                create_videos(root=local_dir, info=info)
 
             return str(local_dir)
 

tests/policies/test_policies.py

@@ -141,6 +141,7 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs):
     Note: We test various combinations of policy and dataset. The combinations are by no means exhaustive,
           and for now we add tests as we see fit.
     """
+    policy_kwargs["device"] = DEVICE
 
     train_cfg = TrainPipelineConfig(
         # TODO(rcadene, aliberts): remove dataset download