add mcap dataset support

README.md

@@ -5,6 +5,7 @@
 - [Installation](#installation)
 - [Dataset Generation Pipeline](#dataset-generation-pipeline)
 - [Data Collection](#data-collection)
+- [Dataset Formats](#dataset-formats)
 - [License](#license)
 - [Acknowledgements](#acknowledgements)
 
@@ -115,11 +116,11 @@ INFO: 99% of raw data are used.
 INFO: Dropped demos: 0
 INFO: Saved dataset plan (2 episodes) to example_gopro13_dataset/dataset_plan.pkl
 INFO:
-############### 07_generate_replay_buffer ###############
-INFO: 2 videos used in total!
-100%|█████████████████████████████████████████████████████████████████████████████| 2/2 [00:19<00:00,  9.98s/it]
-INFO: Saving ReplayBuffer to example_gopro13_dataset/dataset.zarr.zip
-INFO: Done! 2 videos used in total!
+############### 07_generate_dataset (mcap) ###############
+INFO: Collected 2 episodes, 1 grippers, 1 cameras.
+INFO: Writing 2 episode MCAP files to example_gopro13_dataset/dataset_mcap
+Episodes: 100%|████████████████████████████████████████████████████████████████████| 2/2 [00:19<00:00,  9.98s/it]
+INFO: Done! 2 episode MCAP files written to example_gopro13_dataset/dataset_mcap
 ```
 
 For this dataset, 99% of the data are useable (successful SLAM), with 0 demonstrations dropped. If your dataset has a low SLAM success rate, double check if you carefully followed our [data collection instructions](#data-collection).
@@ -135,8 +136,8 @@ For this dataset, 99% of the data are useable (successful SLAM), with 0 demonstr
   Video: 119.88 fps  |  SLAM: 59.9401 fps (skip=2)  |  IMU/frame: 3.33667
   ```
   Pass the `skip` value as `--slam_frame_stride` to `dataset_generation_pipeline.py` (and to
-  `04_detect_aruco.py`, `06_generate_dataset_plan.py`, and `07_generate_replay_buffer.py` when
-  running steps manually). The default is `2` (120 fps → 60 fps SLAM).
+  `04_detect_aruco.py`, `06_generate_dataset_plan.py`, `07_generate_mcap_dataset.py`, and
+  `07_generate_zarr_dataset.py` when running steps manually). The default is `2` (120 fps → 60 fps SLAM).
 
 - To inspect intermediate results, visualization scripts are provided:
 
@@ -227,6 +228,25 @@ Record a short video of opening and closing the gripper 5 times. This is used to
 Record *N* demonstration videos. The number of demonstrations needed depends on task complexity and environment variability. We recommend 200 demonstrations for a single task in a fixed environment.
 
 
+## Dataset Formats
+
+The pipeline supports two output formats, selectable via `--format` (`-f`):
+
+```bash
+# MCAP (default)
+uv run python scripts/dataset_generation_pipeline.py <session_dir>
+
+# Zarr
+uv run python scripts/dataset_generation_pipeline.py -f zarr <session_dir>
+```
+
+**MCAP** — One `.mcap` file per episode in a directory (`dataset_mcap/`). Each file contains time-aligned robot state, JPEG-compressed camera images, and IMU telemetry (accelerometer + gyroscope) as typed, self-describing messages. MCAP files can be inspected with [Foxglove Studio](https://foxglove.dev/).
+
+**Zarr** — A single `dataset.zarr.zip` archive containing all episodes in a flat NumPy-backed replay buffer with JpegXl-compressed images.
+
+Both formats store per-step end-effector pose (position + axis-angle rotation), gripper width, demo start/end poses, and camera images. MCAP additionally includes raw IMU samples.
+
+
 ## License
 
 This repository is released under the MIT license. See [LICENSE](LICENSE) for additional details.

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "click>=8.3.1",
     "imagecodecs>=2026.3.6",
     "matplotlib>=3.10.8",
+    "mcap>=1.3.1",
     "opencv-python>=4.13.0.92",
     "pandas>=3.0.1",
     "py-gpmf-parser>=0.1.1",

scripts/dataset_generation_pipeline.py

@@ -8,7 +8,7 @@
     4. Detect ArUco tags in all demo videos.
     5. Run SLAM tag and gripper range calibrations.
     6. Generate dataset_plan.pkl.
-    7. Generate the Zarr replay buffer (dataset.zarr.zip).
+    7. Generate the dataset (MCAP or Zarr replay buffer).
 
 :Usage:
     uv run python scripts/dataset_generation_pipeline.py <session_dir> [<session_dir> ...]
@@ -38,20 +38,30 @@ def _run(cmd: list, step_name: str) -> None:
 @click.command()
 @click.argument("session_dir", nargs=-1, required=True)
 @click.option("-c", "--calibration_dir", default=None)
+@click.option(
+    "-f",
+    "--format",
+    "dataset_format",
+    type=click.Choice(["mcap", "zarr"], case_sensitive=False),
+    default="mcap",
+    show_default=True,
+    help="Output dataset format.",
+)
 @click.option(
     "-sfs",
     "--slam_frame_stride",
     type=int,
     default=2,
     help="Frame stride used by SLAM and ArUco (raw_fps / slam_fps).",
 )
-def main(session_dir, calibration_dir, slam_frame_stride):
+def main(session_dir, calibration_dir, dataset_format, slam_frame_stride):
     """Run the full dataset generation pipeline for each session directory.
 
     :param session_dir: One or more session directories to process.
     :param calibration_dir: Path to the calibration directory containing
         gopro13_intrinsics_2_7k.json and aruco_config.yaml. Defaults to
         example/calibration/ relative to the repository root.
+    :param dataset_format: Output format — 'mcap' (default) or 'zarr'.
     :param slam_frame_stride: Frame stride matching SLAM/ArUco detection rate
         (raw_fps / slam_fps). Default 2 matches 120 fps -> 60 fps SLAM.
     """
@@ -213,22 +223,39 @@ def main(session_dir, calibration_dir, slam_frame_stride):
         ]
         _run(cmd, "06_generate_dataset_plan")
 
-        # 07 generate replay buffer
-        logger.info("\n%s 07_generate_replay_buffer %s", "#" * 15, "#" * 15)
-        script_path = script_dir.joinpath("07_generate_replay_buffer.py")
-        if not script_path.is_file():
-            raise click.ClickException(f"Could not find script at: {script_path}")
-        output_path = session.joinpath("dataset.zarr.zip")
-        cmd = [
-            sys.executable,
-            str(script_path),
-            "--output",
-            str(output_path),
-            "--slam_frame_stride",
-            str(slam_frame_stride),
-            str(session),
-        ]
-        _run(cmd, "07_generate_replay_buffer")
+        # 07 generate dataset
+        logger.info(
+            "\n%s 07_generate_dataset (%s) %s", "#" * 15, dataset_format, "#" * 15
+        )
+        if dataset_format == "mcap":
+            script_path = script_dir.joinpath("07_generate_mcap_dataset.py")
+            if not script_path.is_file():
+                raise click.ClickException(f"Could not find script at: {script_path}")
+            output_path = session.joinpath("dataset_mcap")
+            cmd = [
+                sys.executable,
+                str(script_path),
+                "--output",
+                str(output_path),
+                "--slam_frame_stride",
+                str(slam_frame_stride),
+                str(session),
+            ]
+        else:
+            script_path = script_dir.joinpath("07_generate_zarr_dataset.py")
+            if not script_path.is_file():
+                raise click.ClickException(f"Could not find script at: {script_path}")
+            output_path = session.joinpath("dataset.zarr.zip")
+            cmd = [
+                sys.executable,
+                str(script_path),
+                "--output",
+                str(output_path),
+                "--slam_frame_stride",
+                str(slam_frame_stride),
+                str(session),
+            ]
+        _run(cmd, "07_generate_dataset")
 
 
 if __name__ == "__main__":