Stream ALOHA HDF5 conversion frame-by-frame to cut RAM use

examples/aloha_real/convert_aloha_data_to_lerobot.py

@@ -1,7 +1,16 @@
 """
-Script to convert Aloha hdf5 data to the LeRobot dataset v2.0 format.
+Script to convert Aloha hdf5 data to the LeRobot dataset format used by openpi.
 
-Example usage: uv run examples/aloha_real/convert_aloha_data_to_lerobot.py --raw-dir /path/to/raw/data --repo-id <org>/<dataset-name>
+Uses frame-at-a-time HDF5 reads so large multi-task exports do not load every
+camera frame into RAM at once (see openpi#565 / openpi#986).
+
+Example usage:
+  uv run examples/aloha_real/convert_aloha_data_to_lerobot.py \\
+    --raw-dir /path/to/raw/data --repo-id <org>/<dataset-name>
+
+Note: openpi training currently expects LeRobot v2.1 datasets on the pinned
+lerobot revision. This converter targets that format; v3 output is tracked in
+openpi#706.
 """
 
 import dataclasses
@@ -10,7 +19,7 @@
 from typing import Literal
 
 import h5py
-from lerobot.common.datasets.lerobot_dataset import LEROBOT_HOME
+from lerobot.common.datasets.lerobot_dataset import HF_LEROBOT_HOME
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.common.datasets.push_dataset_to_hub._download_raw import download_raw
 import numpy as np
@@ -30,6 +39,13 @@ class DatasetConfig:
 
 DEFAULT_DATASET_CONFIG = DatasetConfig()
 
+DEFAULT_CAMERAS = (
+    "cam_high",
+    "cam_low",
+    "cam_left_wrist",
+    "cam_right_wrist",
+)
+
 
 def create_empty_dataset(
     repo_id: str,
@@ -56,12 +72,7 @@ def create_empty_dataset(
         "left_wrist_rotate",
         "left_gripper",
     ]
-    cameras = [
-        "cam_high",
-        "cam_low",
-        "cam_left_wrist",
-        "cam_right_wrist",
-    ]
+    cameras = list(DEFAULT_CAMERAS)
 
     features = {
         "observation.state": {
@@ -109,8 +120,9 @@ def create_empty_dataset(
             ],
         }
 
-    if Path(LEROBOT_HOME / repo_id).exists():
-        shutil.rmtree(LEROBOT_HOME / repo_id)
+    output_path = HF_LEROBOT_HOME / repo_id
+    if output_path.exists():
+        shutil.rmtree(output_path)
 
     return LeRobotDataset.create(
         repo_id=repo_id,
@@ -127,7 +139,6 @@ def create_empty_dataset(
 
 def get_cameras(hdf5_files: list[Path]) -> list[str]:
     with h5py.File(hdf5_files[0], "r") as ep:
-        # ignore depth channel, not currently handled
         return [key for key in ep["/observations/images"].keys() if "depth" not in key]  # noqa: SIM118
 
 
@@ -141,87 +152,57 @@ def has_effort(hdf5_files: list[Path]) -> bool:
         return "/observations/effort" in ep
 
 
-def load_raw_images_per_camera(ep: h5py.File, cameras: list[str]) -> dict[str, np.ndarray]:
-    imgs_per_cam = {}
-    for camera in cameras:
-        uncompressed = ep[f"/observations/images/{camera}"].ndim == 4
+def load_image_frame(ep: h5py.File, camera: str, frame_idx: int) -> np.ndarray:
+    """Load a single camera frame without materializing the full episode tensor."""
+    dataset = ep[f"/observations/images/{camera}"]
+    if dataset.ndim == 4:
+        return dataset[frame_idx]
 
-        if uncompressed:
-            # load all images in RAM
-            imgs_array = ep[f"/observations/images/{camera}"][:]
-        else:
-            import cv2
+    import cv2
 
-            # load one compressed image after the other in RAM and uncompress
-            imgs_array = []
-            for data in ep[f"/observations/images/{camera}"]:
-                imgs_array.append(cv2.cvtColor(cv2.imdecode(data, 1), cv2.COLOR_BGR2RGB))
-            imgs_array = np.array(imgs_array)
+    encoded = dataset[frame_idx]
+    return cv2.cvtColor(cv2.imdecode(encoded, 1), cv2.COLOR_BGR2RGB)
 
-        imgs_per_cam[camera] = imgs_array
-    return imgs_per_cam
 
-
-def load_raw_episode_data(
-    ep_path: Path,
-) -> tuple[dict[str, np.ndarray], torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+def populate_episode(dataset: LeRobotDataset, ep_path: Path, *, task: str, cameras: list[str]) -> None:
     with h5py.File(ep_path, "r") as ep:
         state = torch.from_numpy(ep["/observations/qpos"][:])
         action = torch.from_numpy(ep["/action"][:])
+        velocity = torch.from_numpy(ep["/observations/qvel"][:]) if "/observations/qvel" in ep else None
+        effort = torch.from_numpy(ep["/observations/effort"][:]) if "/observations/effort" in ep else None
+        num_frames = state.shape[0]
 
-        velocity = None
-        if "/observations/qvel" in ep:
-            velocity = torch.from_numpy(ep["/observations/qvel"][:])
-
-        effort = None
-        if "/observations/effort" in ep:
-            effort = torch.from_numpy(ep["/observations/effort"][:])
-
-        imgs_per_cam = load_raw_images_per_camera(
-            ep,
-            [
-                "cam_high",
-                "cam_low",
-                "cam_left_wrist",
-                "cam_right_wrist",
-            ],
-        )
+        for frame_idx in range(num_frames):
+            frame = {
+                "observation.state": state[frame_idx],
+                "action": action[frame_idx],
+            }
+            for camera in cameras:
+                frame[f"observation.images.{camera}"] = load_image_frame(ep, camera, frame_idx)
+            if velocity is not None:
+                frame["observation.velocity"] = velocity[frame_idx]
+            if effort is not None:
+                frame["observation.effort"] = effort[frame_idx]
+            dataset.add_frame(frame)
 
-    return imgs_per_cam, state, action, velocity, effort
+    dataset.save_episode(task=task)
 
 
 def populate_dataset(
     dataset: LeRobotDataset,
     hdf5_files: list[Path],
     task: str,
     episodes: list[int] | None = None,
+    *,
+    cameras: list[str] | None = None,
 ) -> LeRobotDataset:
     if episodes is None:
         episodes = range(len(hdf5_files))
+    if cameras is None:
+        cameras = get_cameras(hdf5_files)
 
     for ep_idx in tqdm.tqdm(episodes):
-        ep_path = hdf5_files[ep_idx]
-
-        imgs_per_cam, state, action, velocity, effort = load_raw_episode_data(ep_path)
-        num_frames = state.shape[0]
-
-        for i in range(num_frames):
-            frame = {
-                "observation.state": state[i],
-                "action": action[i],
-            }
-
-            for camera, img_array in imgs_per_cam.items():
-                frame[f"observation.images.{camera}"] = img_array[i]
-
-            if velocity is not None:
-                frame["observation.velocity"] = velocity[i]
-            if effort is not None:
-                frame["observation.effort"] = effort[i]
-
-            dataset.add_frame(frame)
-
-        dataset.save_episode(task=task)
+        populate_episode(dataset, hdf5_files[ep_idx], task=task, cameras=cameras)
 
     return dataset
 
@@ -238,15 +219,17 @@ def port_aloha(
     mode: Literal["video", "image"] = "image",
     dataset_config: DatasetConfig = DEFAULT_DATASET_CONFIG,
 ):
-    if (LEROBOT_HOME / repo_id).exists():
-        shutil.rmtree(LEROBOT_HOME / repo_id)
+    output_path = HF_LEROBOT_HOME / repo_id
+    if output_path.exists():
+        shutil.rmtree(output_path)
 
     if not raw_dir.exists():
         if raw_repo_id is None:
             raise ValueError("raw_repo_id must be provided if raw_dir does not exist")
         download_raw(raw_dir, repo_id=raw_repo_id)
 
     hdf5_files = sorted(raw_dir.glob("episode_*.hdf5"))
+    cameras = get_cameras(hdf5_files)
 
     dataset = create_empty_dataset(
         repo_id,
@@ -261,6 +244,7 @@ def port_aloha(
         hdf5_files,
         task=task,
         episodes=episodes,
+        cameras=cameras,
     )
     dataset.consolidate()
 

examples/aloha_real/convert_aloha_test.py

@@ -0,0 +1,67 @@
+"""Tests for low-memory ALOHA → LeRobot conversion helpers."""
+
+import importlib.util
+from pathlib import Path
+
+import h5py
+import numpy as np
+import pytest
+
+_CONVERTER_PATH = Path(__file__).with_name("convert_aloha_data_to_lerobot.py")
+_SPEC = importlib.util.spec_from_file_location("convert_aloha_data_to_lerobot", _CONVERTER_PATH)
+assert _SPEC and _SPEC.loader
+converter = importlib.util.module_from_spec(_SPEC)
+_SPEC.loader.exec_module(converter)
+
+
+def _write_synthetic_episode(path: Path, num_frames: int = 4) -> None:
+    with h5py.File(path, "w") as f:
+        f.create_dataset("/observations/qpos", data=np.ones((num_frames, 14), dtype=np.float32))
+        f.create_dataset("/action", data=np.zeros((num_frames, 14), dtype=np.float32))
+        images = f.create_group("/observations/images")
+        for cam in converter.DEFAULT_CAMERAS:
+            images.create_dataset(
+                cam,
+                data=np.zeros((num_frames, 3, 480, 640), dtype=np.uint8),
+            )
+
+
+def test_load_image_frame_reads_single_frame(tmp_path: Path):
+    ep_path = tmp_path / "episode_0.hdf5"
+    _write_synthetic_episode(ep_path, num_frames=3)
+
+    with h5py.File(ep_path, "r") as ep:
+        frame = converter.load_image_frame(ep, "cam_high", 1)
+
+    assert frame.shape == (3, 480, 640)
+    assert frame.dtype == np.uint8
+
+
+def test_populate_episode_does_not_load_all_images_at_once(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    ep_path = tmp_path / "episode_0.hdf5"
+    _write_synthetic_episode(ep_path, num_frames=2)
+    calls: list[tuple[str, int]] = []
+
+    def spy_load_image_frame(ep, camera, frame_idx):
+        calls.append((camera, frame_idx))
+        return np.zeros((3, 480, 640), dtype=np.uint8)
+
+    monkeypatch.setattr(converter, "load_image_frame", spy_load_image_frame)
+
+    class FakeDataset:
+        def __init__(self):
+            self.frames: list[dict] = []
+
+        def add_frame(self, frame):
+            self.frames.append(frame)
+
+        def save_episode(self, *, task: str):
+            assert task == "test-task"
+
+    dataset = FakeDataset()
+    converter.populate_episode(dataset, ep_path, task="test-task", cameras=["cam_high", "cam_low"])
+
+    assert len(dataset.frames) == 2
+    assert len(calls) == 4
+    assert ("cam_high", 0) in calls
+    assert ("cam_low", 1) in calls