Update v21 to v30 converter (#113)

Co-authored-by: Codex <codex@openai.com>

ds_version_convert/v21_to_v30/convert_dataset_v21_to_v30.py

@@ -26,11 +26,23 @@ This script will help you convert any LeRobot dataset already pushed to the hub
 
 Usage:
 
+Convert a dataset from the hub:
 ```bash
-python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
+python src/lerobot/scripts/convert_dataset_v21_to_v30.py \
     --repo-id=lerobot/pusht
 ```
 
+Convert a local dataset (works in place):
+```bash
+python src/lerobot/scripts/convert_dataset_v21_to_v30.py \
+    --repo-id=lerobot/pusht \
+    --root=/path/to/local/dataset/directory \
+    --push-to-hub=false
+
+N.B. Path semantics (v2): --root is the exact dataset folder containing
+meta/, data/, videos/. When omitted, defaults to $HF_LEROBOT_HOME/{repo_id}.
+```
+
 """
 
 import argparse
@@ -39,39 +51,47 @@ import shutil
 from pathlib import Path
 from typing import Any
 
+from lerobot.utils.import_utils import require_package
+
+require_package("jsonlines", extra="dataset")
+
 import jsonlines
 import pandas as pd
 import pyarrow as pa
 import tqdm
 from datasets import Dataset, Features, Image
 from huggingface_hub import HfApi, snapshot_download
-from lerobot.datasets.compute_stats import aggregate_stats
+from requests import HTTPError
-from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
+
+from lerobot.datasets import CODEBASE_VERSION, LeRobotDataset, aggregate_stats
+from lerobot.datasets.io_utils import (
+    cast_stats_to_numpy,
+    get_file_size_in_mb,
+    get_parquet_file_size_in_mb,
+    get_parquet_num_frames,
+    load_info,
+    load_json,
+    write_episodes,
+    write_info,
+    write_stats,
+    write_tasks,
+)
 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,
+    INFO_PATH,
     LEGACY_EPISODES_PATH,
     LEGACY_EPISODES_STATS_PATH,
     LEGACY_TASKS_PATH,
-    cast_stats_to_numpy,
+    DatasetInfo,
-    flatten_dict,
-    get_file_size_in_mb,
-    get_parquet_file_size_in_mb,
-    get_parquet_num_frames,
-    load_info,
     update_chunk_file_indices,
-    write_episodes,
-    write_info,
-    write_stats,
-    write_tasks,
 )
 from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
 from lerobot.utils.constants import HF_LEROBOT_HOME
-from lerobot.utils.utils import init_logging
+from lerobot.utils.utils import flatten_dict, init_logging
-from requests import HTTPError
 
 V21 = "v2.1"
 V30 = "v3.0"
@@ -95,7 +115,7 @@ episodes.jsonl
 {"episode_index": 1, "tasks": ["Put the blue block in the green bowl"], "length": 266}
 
 NEW
-meta/episodes/chunk-000/episodes_000.parquet
+meta/episodes/chunk-000/file_000.parquet
 episode_index | video_chunk_index | video_file_index | data_chunk_index | data_file_index | tasks | length
 -------------------------
 OLD
@@ -103,15 +123,16 @@ tasks.jsonl
 {"task_index": 1, "task": "Put the blue block in the green bowl"}
 
 NEW
-meta/tasks/chunk-000/file_000.parquet
+meta/tasks.parquet
 task_index | task
 -------------------------
 OLD
 episodes_stats.jsonl
+{"episode_index": 1, "stats": {"feature_name": {"min": ..., "max": ..., "mean": ..., "std": ..., "count": ...}}}
 
 NEW
-meta/episodes_stats/chunk-000/file_000.parquet
+meta/episodes/chunk-000/file_000.parquet
-episode_index | mean | std | min | max
+episode_index | feature_name/min | feature_name/max | feature_name/mean | feature_name/std | feature_name/count
 -------------------------
 UPDATE
 meta/info.json
@@ -147,7 +168,7 @@ def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]:
 def validate_local_dataset_version(local_path: Path) -> None:
     """Validate that the local dataset has the expected v2.1 version."""
     info = load_info(local_path)
-    dataset_version = info.get("codebase_version", "unknown")
+    dataset_version = info.codebase_version or "unknown"
     if dataset_version != V21:
         raise ValueError(
             f"Local dataset has codebase version '{dataset_version}', expected '{V21}'. "
@@ -160,7 +181,7 @@ def convert_tasks(root, new_root):
     tasks, _ = legacy_load_tasks(root)
     task_indices = tasks.keys()
     task_strings = tasks.values()
-    df_tasks = pd.DataFrame({"task_index": task_indices}, index=task_strings)
+    df_tasks = pd.DataFrame({"task_index": task_indices}, index=pd.Index(task_strings, name="task"))
     write_tasks(df_tasks, new_root)
 
 
@@ -191,7 +212,6 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
 
     image_keys = get_image_keys(root)
 
-    ep_idx = 0
     chunk_idx = 0
     file_idx = 0
     size_in_mb = 0
@@ -201,9 +221,23 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
 
     logging.info(f"Converting data files from {len(ep_paths)} episodes")
 
-    for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
+    for ep_idx, ep_path in enumerate(tqdm.tqdm(ep_paths, desc="convert data files")):
         ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
         ep_num_frames = get_parquet_num_frames(ep_path)
+
+        # Check if we need to start a new file BEFORE creating metadata
+        if size_in_mb + ep_size_in_mb >= data_file_size_in_mb and len(paths_to_cat) > 0:
+            # Write the accumulated data files
+            concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
+
+            # Move to next file
+            chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
+
+            # Reset for the next file
+            size_in_mb = 0
+            paths_to_cat = []
+
+        # Now create metadata with correct chunk/file indices
         ep_metadata = {
             "episode_index": ep_idx,
             "data/chunk_index": chunk_idx,
@@ -214,20 +248,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
         size_in_mb += ep_size_in_mb
         num_frames += ep_num_frames
         episodes_metadata.append(ep_metadata)
-        ep_idx += 1
+        paths_to_cat.append(ep_path)
-
-        if size_in_mb < data_file_size_in_mb:
-            paths_to_cat.append(ep_path)
-            continue
-
-        if paths_to_cat:
-            concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
-
-        # Reset for the next file
-        size_in_mb = ep_size_in_mb
-        paths_to_cat = [ep_path]
-
-        chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
 
     # Write remaining data if any
     if paths_to_cat:
@@ -238,14 +259,14 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
 
 def get_video_keys(root):
     info = load_info(root)
-    features = info["features"]
+    features = info.features
     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"]
+    features = info.features
     image_keys = [key for key, ft in features.items() if ft["dtype"] == "image"]
     return image_keys
 
@@ -268,7 +289,7 @@ def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
     if len(set(num_eps_per_cam)) != 1:
         raise ValueError(f"All cams dont have same number of episodes ({num_eps_per_cam}).")
 
-    episods_metadata = []
+    episodes_metadata = []
     num_cameras = len(video_keys)
     num_episodes = num_eps_per_cam[0]
     for ep_idx in tqdm.tqdm(range(num_episodes), desc="convert videos"):
@@ -281,9 +302,9 @@ def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
         ep_dict = {}
         for cam_idx in range(num_cameras):
             ep_dict.update(eps_metadata_per_cam[cam_idx][ep_idx])
-        episods_metadata.append(ep_dict)
+        episodes_metadata.append(ep_dict)
 
-    return episods_metadata
+    return episodes_metadata
 
 
 def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_file_size_in_mb: int):
@@ -416,7 +437,8 @@ def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_
 
 
 def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
-    info = load_info(root)
+    # Load as raw dict to remove legacy v2.1 fields before constructing DatasetInfo.
+    info = load_json(root / INFO_PATH)
     info["codebase_version"] = V30
     del info["total_chunks"]
     del info["total_videos"]
@@ -431,7 +453,9 @@ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
             # already has fps in video_info
             continue
         info["features"][key]["fps"] = info["fps"]
-    write_info(info, new_root)
+    # Convert raw dict to typed DatasetInfo before writing
+    dataset_info = DatasetInfo.from_dict(info)
+    write_info(dataset_info, new_root)
 
 
 def convert_dataset(
@@ -459,7 +483,7 @@ def convert_dataset(
 
     # Set root based on whether local dataset path is provided
     use_local_dataset = False
-    root = HF_LEROBOT_HOME / repo_id if root is None else Path(root) / repo_id
+    root = HF_LEROBOT_HOME / repo_id if root is None else Path(root)
     if root.exists():
         validate_local_dataset_version(root)
         use_local_dataset = True
@@ -519,7 +543,7 @@ if __name__ == "__main__":
         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`).",
+        "(e.g. `lerobot/pusht`, `<USER>/aloha_sim_insertion_human`).",
     )
     parser.add_argument(
         "--branch",
@@ -543,7 +567,7 @@ if __name__ == "__main__":
         "--root",
         type=str,
         default=None,
-        help="Local directory to use for downloading/writing the dataset.",
+        help="Local directory to use for downloading/writing the dataset. Defaults to $HF_LEROBOT_HOME/repo_id.",
     )
     parser.add_argument(
         "--push-to-hub",
@@ -558,4 +582,4 @@ if __name__ == "__main__":
     )
 
     args = parser.parse_args()
-    convert_dataset(**vars(args))
+    convert_dataset(**vars(args))