From de390f27beac439196c85d29362d884c920bb9d9 Mon Sep 17 00:00:00 2001 From: zbzhu99 Date: Fri, 10 Apr 2026 21:48:59 +0800 Subject: [PATCH] Add zarr2lerobot converter Adds tools/zarr2lerobot/, an offline converter that reads the aggregated zarr produced by tools/process_demo_data.py and emits a LeRobot Dataset v3.0 directory via the official lerobot SDK. Features: - Both joint and cartesian control modes (cartesian enforced as 8-dim: tcp_position + tcp_orientation + gripper) - Task text from CLI --task with per-episode description fallback - Bulk per-episode zarr slicing (one chunk decompression per array per episode, not per frame) so large datasets are not O(frames) in read I/O - argparse CLI entry point (python -m tools.zarr2lerobot.convert_zarr_to_lerobot) - Depth frames are intentionally skipped in v1; the v3 ecosystem has no consensus storage format for depth yet Tests: - 22 pytest cases (unit helpers + end-to-end round-trip on synthetic fixture) passing on both joint and cartesian modes - Verified on a real 50-episode / 7075-frame cartesian recording; output loads back via LeRobotDataset with correct episode/frame boundaries The lerobot extra in setup.py is pinned to >=0.4.0,<0.5 from PyPI. 0.4.x is the last series that ships Dataset v3.0 while still supporting Python 3.10 (0.5.0+ requires >=3.12, which breaks the unicon env). Also adds a README section for the new export path and a .gitignore negation for tools/zarr2lerobot/tests/ (repo has a blanket tests/ ignore). --- .gitignore | 3 +- README.md | 13 + setup.py | 5 + tools/zarr2lerobot/__init__.py | 0 tools/zarr2lerobot/convert_zarr_to_lerobot.py | 306 ++++++++++++++++++ tools/zarr2lerobot/tests/__init__.py | 0 tools/zarr2lerobot/tests/conftest.py | 80 +++++ tools/zarr2lerobot/tests/test_convert.py | 218 +++++++++++++ 8 files changed, 624 insertions(+), 1 deletion(-) create mode 100644 tools/zarr2lerobot/__init__.py create mode 100644 tools/zarr2lerobot/convert_zarr_to_lerobot.py create mode 100644 tools/zarr2lerobot/tests/__init__.py create mode 100644 tools/zarr2lerobot/tests/conftest.py create mode 100644 tools/zarr2lerobot/tests/test_convert.py diff --git a/.gitignore b/.gitignore index 24c267f..f0a001c 100644 --- a/.gitignore +++ b/.gitignore @@ -93,4 +93,5 @@ references/ **/pretrained_weights **/outputs reference_repos/ -tests/ \ No newline at end of file +tests/ +!tools/zarr2lerobot/tests/ \ No newline at end of file diff --git a/README.md b/README.md index f77acf8..4af327b 100644 --- a/README.md +++ b/README.md @@ -151,6 +151,19 @@ python tools/process_demo_data.py ./data/ur5_recording ``` To support **vla training**, you can transfer the zarr dataset into **rlds dataset** easily, see [zarr2rlds](tools/zarr2rlds). +3. **Export to LeRobot Dataset v3.0**: +```bash +pip install -e '.[lerobot]' +python -m tools.zarr2lerobot.convert_zarr_to_lerobot \ + --zarr-path ./data/ur5_recording/run.joint.zarr \ + --output-dir ./data/ur5_recording_lerobot \ + --repo-id myuser/ur5-pick \ + --task "pick up the red cube" \ + --fps 30 \ + --robot-type ur5 +``` +Per-episode descriptions in the zarr's `meta/episode_descriptions` take precedence over `--task` when non-empty. Depth frames are not exported in v1. + ### Configuration Examples #### Teleoperate XArm6 with SpaceMouse Control diff --git a/setup.py b/setup.py index 59cd133..6daa485 100644 --- a/setup.py +++ b/setup.py @@ -50,6 +50,11 @@ "numba", ], # Optional dependency for RealSense cameras "franky_fr3_franky": ["franky-control==1.1.1"], + "lerobot": [ + # 0.4.x is the last series that still supports Python 3.10 + # (0.5.0+ requires 3.12) while shipping Dataset v3.0. + "lerobot>=0.4.0,<0.5", + ], }, python_requires=">=3.10", classifiers=[ diff --git a/tools/zarr2lerobot/__init__.py b/tools/zarr2lerobot/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tools/zarr2lerobot/convert_zarr_to_lerobot.py b/tools/zarr2lerobot/convert_zarr_to_lerobot.py new file mode 100644 index 0000000..e4254f4 --- /dev/null +++ b/tools/zarr2lerobot/convert_zarr_to_lerobot.py @@ -0,0 +1,306 @@ +"""Convert a ManiUniCon aggregated zarr to LeRobot Dataset v3.0. + +Input: a `{name}.{joint|cartesian}.zarr` directory produced by + ``tools/process_demo_data.py``. +Output: a local LeRobot v3.0 dataset directory. Push-to-hub is a manual + follow-up (``huggingface-cli upload``). + +Follow-ups not handled in v1: +- Depth frames (``data/obs/depths/*``) -- no consensus in the v3 ecosystem. +- Camera intrinsics/extrinsics -- not part of the LeRobot v3 core schema. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Iterator, Literal + +import numpy as np + +ControlMode = Literal["joint", "cartesian"] + + +def detect_control_mode(zarr_path: Path | str) -> ControlMode: + """Infer joint vs cartesian from the filename written by + ``tools/process_demo_data.py``: the name ends with + ``.joint.zarr`` or ``.cartesian.zarr``. + """ + name = Path(zarr_path).name + parts = name.split(".") + if len(parts) >= 3 and parts[-1] == "zarr" and parts[-2] in ("joint", "cartesian"): + return parts[-2] # type: ignore[return-value] + raise ValueError( + f"cannot infer control mode from {name!r}; expected '*.joint.zarr' or '*.cartesian.zarr'" + ) + + +def discover_camera_keys(root) -> list[str]: + """Return the sorted list of RGB camera array names under + ``data/obs/images``. Raises ValueError if the group is missing or empty. + """ + try: + images = root["data/obs/images"] + except KeyError: + raise ValueError("no RGB camera group at 'data/obs/images'") + keys = sorted(images.array_keys()) + if not keys: + raise ValueError("no RGB camera arrays found under 'data/obs/images'") + return keys + + +def build_features( + control_mode: ControlMode, + camera_keys: list[str], + state_dim: int, + action_dim: int, + image_hw: tuple[int, int], +) -> dict: + """Build the LeRobot v3 ``features`` dict for ``LeRobotDataset.create``. + + The LeRobot SDK auto-adds ``timestamp``, ``frame_index``, ``episode_index``, + ``index`` and ``task_index`` -- do not add them here. The dataset-level + fps is passed to ``LeRobotDataset.create`` directly and does not belong + in the per-feature spec. + """ + h, w = image_hw + if control_mode == "cartesian": + if state_dim != 8: + raise ValueError( + f"cartesian control mode requires state_dim=8 " + f"(tcp_position[3] + tcp_orientation[4] + gripper[1]); got {state_dim}" + ) + if control_mode == "joint": + state_names = [f"joint_{i}" for i in range(state_dim - 1)] + ["gripper"] + else: + state_names = ["x", "y", "z", "qx", "qy", "qz", "qw", "gripper"] + + feats: dict = { + "observation.state": { + "dtype": "float32", + "shape": [state_dim], + "names": state_names, + }, + "action": { + "dtype": "float32", + "shape": [action_dim], + "names": state_names if action_dim == state_dim else None, + }, + } + for cam in camera_keys: + feats[f"observation.images.{cam}"] = { + "dtype": "video", + "shape": [h, w, 3], + "names": ["height", "width", "channel"], + } + return feats + + +def resolve_task_text(cli_task: str, episode_desc: str) -> str: + """Per-episode description wins if non-empty; else CLI fallback. + + Raises ValueError if neither produces a non-empty string. + """ + if episode_desc and episode_desc.strip(): + return episode_desc.strip() + if cli_task and cli_task.strip(): + return cli_task.strip() + raise ValueError( + "no task text available: neither per-episode description nor --task was provided" + ) + + +def build_episode_arrays( + root, + start: int, + end: int, + control_mode: ControlMode, + camera_keys: list[str], +) -> dict[str, np.ndarray]: + """Bulk-read one episode's worth of state/action/image arrays. + + Returns a dict keyed by LeRobot feature name with first dimension + ``(end - start)``. Callers index into these arrays by local frame + offset to assemble ``add_frame`` payloads. + + Per-episode slicing matters because ``tools/process_demo_data.py`` writes + state and action arrays as single chunks spanning the full dataset -- + the zarr ``DirectoryStore`` has no chunk cache, so per-frame indexing + would decompress the whole chunk once per frame (thousands of times). + """ + obs = root["data/obs"] + gripper = np.asarray(obs["gripper_state"][start:end], dtype=np.float32) + if control_mode == "joint": + joint_positions = np.asarray( + obs["joint_positions"][start:end], dtype=np.float32 + ) + state = np.concatenate([joint_positions, gripper], axis=1) + else: + tcp_position = np.asarray(obs["tcp_position"][start:end], dtype=np.float32) + tcp_orientation = np.asarray( + obs["tcp_orientation"][start:end], dtype=np.float32 + ) + state = np.concatenate([tcp_position, tcp_orientation, gripper], axis=1) + + arrays: dict[str, np.ndarray] = { + "observation.state": state, + "action": np.asarray(root["data/action"][start:end], dtype=np.float32), + } + for cam in camera_keys: + arrays[f"observation.images.{cam}"] = np.asarray( + obs[f"images/{cam}"][start:end], dtype=np.uint8 + ) + return arrays + + +def iter_episode_bounds(root) -> Iterator[tuple[int, int, int]]: + """Yield ``(episode_index, start, end)`` half-open ranges into the global + time axis, derived from ``meta/episode_ends`` (cumulative convention used + by :class:`ReplayBuffer`). + """ + episode_ends = np.asarray(root["meta/episode_ends"][:], dtype=np.int64) + prev = 0 + for i, end in enumerate(episode_ends): + yield i, prev, int(end) + prev = int(end) + + +def convert_zarr_to_lerobot( + zarr_path: Path | str, + output_dir: Path | str, + repo_id: str, + cli_task: str, + fps: int, + robot_type: str, +) -> None: + """Convert a single aggregated zarr to a LeRobot v3.0 dataset on disk. + + Uses the official lerobot SDK's incremental writer. Always calls + ``finalize()`` -- skipping it would leave parquet files corrupt per the + LeRobot v3 docs "Common Issues" section. + """ + import zarr + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + zarr_path = Path(zarr_path) + output_dir = Path(output_dir) + + control_mode = detect_control_mode(zarr_path) + root = zarr.open(str(zarr_path), mode="r") + camera_keys = discover_camera_keys(root) + + gripper = root["data/obs/gripper_state"] + action_arr = root["data/action"] + if control_mode == "joint": + state_dim = int(root["data/obs/joint_positions"].shape[1]) + int( + gripper.shape[1] + ) + else: + state_dim = ( + int(root["data/obs/tcp_position"].shape[1]) + + int(root["data/obs/tcp_orientation"].shape[1]) + + int(gripper.shape[1]) + ) + action_dim = int(action_arr.shape[1]) + + first_cam = root[f"data/obs/images/{camera_keys[0]}"] + image_hw = (int(first_cam.shape[1]), int(first_cam.shape[2])) + + features = build_features( + control_mode=control_mode, + camera_keys=camera_keys, + state_dim=state_dim, + action_dim=action_dim, + image_hw=image_hw, + ) + # lerobot's validate_feature_numpy_array compares ``value.shape != + # feature["shape"]`` directly, but numpy shape tuples never equal JSON + # lists. Coerce shapes to tuples so the comparison works. + for spec in features.values(): + spec["shape"] = tuple(spec["shape"]) + + dataset = LeRobotDataset.create( + repo_id=repo_id, + fps=fps, + root=output_dir, + robot_type=robot_type, + features=features, + use_videos=True, + ) + + episode_descriptions = root["meta/episode_descriptions"][:] + try: + for ep_idx, start, end in iter_episode_bounds(root): + task_text = resolve_task_text(cli_task, str(episode_descriptions[ep_idx])) + ep_arrays = build_episode_arrays( + root, start, end, control_mode, camera_keys + ) + for local_t in range(end - start): + frame = {key: arr[local_t] for key, arr in ep_arrays.items()} + frame["task"] = task_text + dataset.add_frame(frame) + dataset.save_episode() + finally: + dataset.finalize() + + +def _build_parser(): + import argparse + + parser = argparse.ArgumentParser( + description="Convert a ManiUniCon aggregated zarr to LeRobot Dataset v3.0.", + ) + parser.add_argument( + "--zarr-path", + type=Path, + required=True, + help="Path to '*.joint.zarr' or '*.cartesian.zarr' from tools/process_demo_data.py.", + ) + parser.add_argument( + "--output-dir", + type=Path, + required=True, + help="Local directory to write the LeRobot dataset into (must not exist).", + ) + parser.add_argument( + "--repo-id", + type=str, + required=True, + help="LeRobot repo id, e.g. 'username/dataset-name'. Only used as metadata; push is manual.", + ) + parser.add_argument( + "--task", + type=str, + default="", + help="Fallback natural-language task description. Per-episode descriptions in " + "the zarr 'meta/episode_descriptions' take precedence when non-empty.", + ) + parser.add_argument( + "--fps", + type=int, + default=30, + help="Control frequency in Hz (matches configs/default.yaml policy_frequency=30).", + ) + parser.add_argument( + "--robot-type", + type=str, + default="unknown", + help="Robot type string stored in info.json (e.g. 'ur5', 'xarm6', 'fr3').", + ) + return parser + + +def main(): + args = _build_parser().parse_args() + convert_zarr_to_lerobot( + zarr_path=args.zarr_path, + output_dir=args.output_dir, + repo_id=args.repo_id, + cli_task=args.task, + fps=args.fps, + robot_type=args.robot_type, + ) + print(f"[zarr2lerobot] wrote LeRobot v3.0 dataset to {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/tools/zarr2lerobot/tests/__init__.py b/tools/zarr2lerobot/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tools/zarr2lerobot/tests/conftest.py b/tools/zarr2lerobot/tests/conftest.py new file mode 100644 index 0000000..def605e --- /dev/null +++ b/tools/zarr2lerobot/tests/conftest.py @@ -0,0 +1,80 @@ +from pathlib import Path + +import numpy as np +import pytest +import zarr + + +def _make_synthetic_zarr(tmp_path: Path, mode: str) -> Path: + """Build a 2-episode, 5-frame-each zarr matching process_demo_data output. + + Structure: + data/ + action float64 (10, A) + obs/ + joint_positions float64 (10, 6) + tcp_position float64 (10, 3) + tcp_orientation float64 (10, 4) + gripper_state float64 (10, 1) + images/ + camera_0 uint8 (10, 64, 64, 3) + camera_1 uint8 (10, 64, 64, 3) + meta/ + episode_ends int64 (2,) = [5, 10] + episode_descriptions Path: + return _make_synthetic_zarr(tmp_path, "joint") + + +@pytest.fixture +def cartesian_zarr(tmp_path: Path) -> Path: + return _make_synthetic_zarr(tmp_path, "cartesian") diff --git a/tools/zarr2lerobot/tests/test_convert.py b/tools/zarr2lerobot/tests/test_convert.py new file mode 100644 index 0000000..0dbc32e --- /dev/null +++ b/tools/zarr2lerobot/tests/test_convert.py @@ -0,0 +1,218 @@ +import numpy as np +import pytest +import zarr + +from tools.zarr2lerobot.convert_zarr_to_lerobot import ( + build_episode_arrays, + build_features, + detect_control_mode, + discover_camera_keys, + iter_episode_bounds, + resolve_task_text, +) + + +def test_joint_fixture_layout(joint_zarr): + root = zarr.open(str(joint_zarr), mode="r") + assert root["data/action"].shape == (10, 7) + assert root["data/obs/images/camera_0"].shape == (10, 64, 64, 3) + assert root["meta/episode_ends"][:].tolist() == [5, 10] + assert root["meta/episode_descriptions"][0] == "pick up the cube" + + +def test_cartesian_fixture_layout(cartesian_zarr): + root = zarr.open(str(cartesian_zarr), mode="r") + assert root["data/action"].shape == (10, 8) + + +def test_detect_control_mode_joint(joint_zarr): + assert detect_control_mode(joint_zarr) == "joint" + + +def test_detect_control_mode_cartesian(cartesian_zarr): + assert detect_control_mode(cartesian_zarr) == "cartesian" + + +def test_detect_control_mode_invalid(tmp_path): + bad = tmp_path / "bad.zarr" + bad.mkdir() + with pytest.raises(ValueError, match="cannot infer control mode"): + detect_control_mode(bad) + + +def test_discover_camera_keys_ordered(joint_zarr): + root = zarr.open(str(joint_zarr), mode="r") + assert discover_camera_keys(root) == ["camera_0", "camera_1"] + + +def test_discover_camera_keys_missing_group(tmp_path): + path = tmp_path / "empty.joint.zarr" + root = zarr.open(str(path), mode="w") + root.create_group("data").create_group("obs") + with pytest.raises(ValueError, match="no RGB camera"): + discover_camera_keys(root) + + +def test_discover_camera_keys_empty_group(tmp_path): + path = tmp_path / "empty2.joint.zarr" + root = zarr.open(str(path), mode="w") + root.create_group("data").create_group("obs").create_group("images") + with pytest.raises(ValueError, match="no RGB camera"): + discover_camera_keys(root) + + +def test_build_features_joint(): + feats = build_features( + control_mode="joint", + camera_keys=["camera_0"], + state_dim=7, + action_dim=7, + image_hw=(96, 96), + ) + assert feats["observation.state"]["shape"] == [7] + assert feats["observation.state"]["dtype"] == "float32" + assert feats["action"]["shape"] == [7] + assert feats["action"]["dtype"] == "float32" + assert feats["observation.images.camera_0"]["dtype"] == "video" + assert feats["observation.images.camera_0"]["shape"] == [96, 96, 3] + + +def test_build_features_cartesian_two_cams(): + feats = build_features( + control_mode="cartesian", + camera_keys=["front", "wrist"], + state_dim=8, + action_dim=8, + image_hw=(240, 320), + ) + assert feats["observation.state"]["shape"] == [8] + assert feats["action"]["shape"] == [8] + assert "observation.images.front" in feats + assert "observation.images.wrist" in feats + assert feats["observation.images.front"]["shape"] == [240, 320, 3] + + +def test_build_features_cartesian_rejects_wrong_state_dim(): + with pytest.raises(ValueError, match="cartesian control mode requires"): + build_features( + control_mode="cartesian", + camera_keys=["cam"], + state_dim=6, + action_dim=6, + image_hw=(64, 64), + ) + + +@pytest.mark.parametrize( + ("cli_task", "episode_desc", "expected"), + [ + ("fallback", "pick cube", "pick cube"), + ("fallback", "", "fallback"), + ("fallback", " ", "fallback"), + ("fallback", " hi ", "hi"), + ], +) +def test_resolve_task_text(cli_task, episode_desc, expected): + assert resolve_task_text(cli_task, episode_desc) == expected + + +def test_resolve_task_raises_if_both_empty(): + with pytest.raises(ValueError, match="no task text"): + resolve_task_text("", "") + + +def test_build_episode_arrays_joint(joint_zarr): + root = zarr.open(str(joint_zarr), mode="r") + arrays = build_episode_arrays( + root, start=0, end=5, control_mode="joint", camera_keys=["camera_0", "camera_1"] + ) + assert arrays["observation.state"].shape == (5, 7) + assert arrays["observation.state"].dtype == np.float32 + assert arrays["action"].shape == (5, 7) + assert arrays["action"].dtype == np.float32 + assert arrays["observation.images.camera_0"].shape == (5, 64, 64, 3) + assert arrays["observation.images.camera_0"].dtype == np.uint8 + assert arrays["observation.images.camera_1"].shape == (5, 64, 64, 3) + + +def test_build_episode_arrays_cartesian(cartesian_zarr): + root = zarr.open(str(cartesian_zarr), mode="r") + arrays = build_episode_arrays( + root, start=5, end=10, control_mode="cartesian", camera_keys=["camera_0"] + ) + assert arrays["observation.state"].shape == (5, 8) + assert arrays["observation.state"].dtype == np.float32 + assert arrays["action"].shape == (5, 8) + assert "observation.images.camera_0" in arrays + assert "observation.images.camera_1" not in arrays + + +def test_iter_episode_bounds_yields_half_open_ranges(joint_zarr): + root = zarr.open(str(joint_zarr), mode="r") + bounds = list(iter_episode_bounds(root)) + assert bounds == [(0, 0, 5), (1, 5, 10)] + + +def test_iter_episode_bounds_types(joint_zarr): + root = zarr.open(str(joint_zarr), mode="r") + for ep_idx, start, end in iter_episode_bounds(root): + assert isinstance(ep_idx, int) + assert isinstance(start, int) + assert isinstance(end, int) + + +def test_convert_joint_end_to_end(joint_zarr, tmp_path): + pytest.importorskip("lerobot") + from tools.zarr2lerobot.convert_zarr_to_lerobot import convert_zarr_to_lerobot + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + output_dir = tmp_path / "out_joint" + convert_zarr_to_lerobot( + zarr_path=joint_zarr, + output_dir=output_dir, + repo_id="test/fake-joint", + cli_task="pick up the red cube", + fps=10, + robot_type="ur5", + ) + + ds = LeRobotDataset("test/fake-joint", root=output_dir) + assert ds.meta.info["codebase_version"] == "v3.0" + assert ds.meta.info["total_episodes"] == 2 + assert ds.meta.info["total_frames"] == 10 + assert ds.meta.info["fps"] == 10 + + sample = ds[0] + assert "observation.state" in sample + assert "action" in sample + assert "observation.images.camera_0" in sample + assert "observation.images.camera_1" in sample + + # Different lerobot pins expose tasks differently (dict vs DataFrame vs + # ndarray), so stringify the whole metadata blob and check substring + # membership -- tolerant across API versions. + info_str = str(ds.meta.info) + str(getattr(ds.meta, "tasks", "")) + assert "pick up the cube" in info_str + assert "pick up the red cube" in info_str + + +def test_convert_cartesian_end_to_end(cartesian_zarr, tmp_path): + pytest.importorskip("lerobot") + from tools.zarr2lerobot.convert_zarr_to_lerobot import convert_zarr_to_lerobot + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + output_dir = tmp_path / "out_cart" + convert_zarr_to_lerobot( + zarr_path=cartesian_zarr, + output_dir=output_dir, + repo_id="test/fake-cartesian", + cli_task="default task", + fps=30, + robot_type="xarm6", + ) + ds = LeRobotDataset("test/fake-cartesian", root=output_dir) + sample = ds[0] + assert "observation.state" in sample + assert "action" in sample + assert ds.meta.info["total_episodes"] == 2 + assert ds.meta.info["fps"] == 30