Support UMI data validation

README.md

@@ -30,6 +30,12 @@ python validate.py validate \
 python validate.py validate \
   --dataset-path ./my-dataset \
   --data-type eval
+
+# Validate UMI demonstration data
+python validate.py validate \
+  --dataset-path ./my-umi-dataset \
+  --data-type teleop \
+  --dataset-profile umi
 ```
 
 ### 4. Get Upload Instructions
@@ -83,7 +89,7 @@ Example:
 
 ### custom_metadata.csv
 
-Must have exactly these columns:
+For the default `robot` profile, it must have exactly these columns:
 
 | Column | Type | Description |
 |--------|------|-------------|
@@ -111,6 +117,44 @@ episode_index,operator_id,is_eval_episode,episode_id,start_timestamp,checkpoint_
 - `checkpoint_path` should only be set for eval episodes (is_eval_episode=True)
 - `checkpoint_path` must be a valid GCS URI format: `gs://bucket/path/to/checkpoint`
 
+### UMI datasets
+
+Pass `--dataset-profile umi` for UMI demonstration data. The UMI profile:
+
+- Requires `episode_index`, `operator_id`, `is_eval_episode`, `episode_id`,
+  `start_timestamp`, and `station_id` in `custom_metadata.csv`
+- Allows `success`, `checkpoint_path`, and `robot_id` to be omitted
+- Allows missing values in an optional `success` column
+- Accepts provider-specific metadata columns
+- Requires a 640x480 MP4 base camera and tracked left/right gripper poses
+- Allows optional wrist cameras, camera tracking, and camera intrinsics
+- Accepts an optional `observation/` prefix on all UMI feature names
+
+```text
+# Required
+base_0_camera/rgb/image
+left/position
+left/quaternion_xyzw
+left/gripper
+right/position
+right/quaternion_xyzw
+right/gripper
+
+# Optional
+base_X_camera/rgb/image
+left_wrist_0_camera/rgb/image
+right_wrist_0_camera/rgb/image
+camera_name/intrinsics
+camera_name/position
+camera_name/quaternion_xyzw
+```
+
+Camera images must use `dtype: video` with shape `[480, 640, 3]`. Position,
+quaternion, gripper, and intrinsics features must have shapes `[3]`, `[4]`,
+`[1]`, and `[3, 3]`, respectively. Every provided tracking field should also
+include a `name/timestamp` companion with shape `[1]` and an `int64` or
+`uint64` nanosecond value.
+
 See `examples/example_dataset/meta/custom_metadata.csv` for a complete example.
 
 ### custom_annotation.json (Optional)
@@ -158,6 +202,7 @@ python validate.py validate \
 **Arguments:**
 - `--dataset-path`: Path to dataset directory (local or GCP URI like gs://bucket/path)
 - `--data-type`: Either `teleop` (training) or `eval` (evaluation)
+- `--dataset-profile`: Either `robot` (default) or `umi`
 
 **Examples:**
 ```bash
@@ -188,6 +233,7 @@ python validate.py compute-path \
 - `--dataset-name`: Dataset name for GCP path (required)
 - `--bucket-name`: GCS bucket name (required)
 - `--data-type`: Either `teleop` or `eval` (required)
+- `--dataset-profile`: Either `robot` (default) or `umi`
 - `--dataset-version`: Version string (optional, default: timestamp)
 - `--custom-folder-prefix`: Custom folder prefix (optional, e.g., "experiments/phase-1")
 - `--skip-validation`: Skip validation, only compute path (optional)
@@ -242,7 +288,7 @@ The validator uses two data types:
 
 ### Metadata CSV
 - All required columns must be present
-- No extra columns allowed
+- No extra columns allowed for the default `robot` profile
 - `episode_id` must be unique
 - `is_eval_episode` and `success` must be boolean
 - `start_timestamp` must be UTC seconds (Unix epoch time) in range 2000-2100

lerobot_validator/cli.py

@@ -15,6 +15,7 @@
 def validate(
     dataset_path: str,
     data_type: Literal["teleop", "eval"],
+    dataset_profile: Literal["robot", "umi"] = "robot",
 ):
     """
     Validate lerobot dataset metadata and annotations.
@@ -27,6 +28,8 @@ def validate(
         dataset_path: Path to the lerobot dataset directory (supports both local paths and GCP URIs like gs://bucket/path)
         data_type: Data type - must be "teleop" or "eval" (required). 
                    Teleop data (training) should have is_eval_episode=False, eval data should have is_eval_episode=True.
+        dataset_profile: Dataset-specific validation contract. Use "umi" for
+                         UMI demonstrations.
     """
     # Convert to AnyPath - supports both local Path and CloudPath (gs://)
     dataset_path_obj = AnyPath(dataset_path)
@@ -43,6 +46,7 @@ def validate(
     print(f"Metadata CSV:    {dataset_path_obj}/meta/custom_metadata.csv (required)")
     print(f"Annotation JSON: {dataset_path_obj}/meta/custom_annotation.json (optional)")
     print(f"Data type:       {data_type.capitalize()}")
+    print(f"Dataset profile: {dataset_profile.upper()}")
     print()
     print("Running validation...")
     print()
@@ -53,6 +57,7 @@ def validate(
     validator = LerobotDatasetValidator(
         dataset_path=dataset_path_obj,
         is_eval_data=is_eval_data,
+        dataset_profile=dataset_profile,
     )
 
     validation_passed = validator.validate()
@@ -78,6 +83,7 @@ def compute_upload_path(
     dataset_name: str,
     bucket_name: str,
     data_type: Literal["teleop", "eval"],
+    dataset_profile: Literal["robot", "umi"] = "robot",
     dataset_version: Optional[str] = None,
     custom_folder_prefix: Optional[str] = None,
     skip_validation: bool = False,
@@ -91,6 +97,8 @@ def compute_upload_path(
         bucket_name: GCS bucket name for upload destination (required)
         data_type: Data type - must be "teleop" or "eval" (required). 
                    Teleop data (training) should have is_eval_episode=False, eval data should have is_eval_episode=True.
+        dataset_profile: Dataset-specific validation contract. Use "umi" for
+                         UMI demonstrations.
         dataset_version: Dataset version (default: current timestamp)
         custom_folder_prefix: Custom folder prefix for GCP path (can include nested folders, e.g., 'foo/bar')
         skip_validation: Skip validation and only compute the path (default: False)
@@ -107,6 +115,7 @@ def compute_upload_path(
     print(f"Dataset name:    {dataset_name}")
     print(f"Bucket:          {bucket_name}")
     print(f"Data type:       {data_type.capitalize()}")
+    print(f"Dataset profile: {dataset_profile.upper()}")
     if dataset_version:
         print(f"Version:         {dataset_version}")
     if custom_folder_prefix:
@@ -123,6 +132,7 @@ def compute_upload_path(
         validator = LerobotDatasetValidator(
             dataset_path=dataset_path_obj,
             is_eval_data=is_eval_data,
+            dataset_profile=dataset_profile,
         )
 
         validation_passed = validator.validate()
@@ -172,4 +182,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-

lerobot_validator/metadata_validator.py

@@ -3,28 +3,44 @@
 """
 
 from pathlib import Path
-from typing import List, Dict, Any, Union
+from typing import List, Union
 
 import pandas as pd
 from cloudpathlib import CloudPath, AnyPath
 
-from lerobot_validator.schemas import REQUIRED_METADATA_COLUMNS
+from lerobot_validator.schemas import (
+    ALLOWED_METADATA_COLUMNS_BY_PROFILE,
+    REQUIRED_METADATA_COLUMNS_BY_PROFILE,
+    DatasetProfile,
+)
 
 
 class MetadataValidator:
     """Validates the custom_metadata.csv file."""
 
-    def __init__(self, metadata_path: Union[str, Path, CloudPath]):
+    def __init__(
+        self,
+        metadata_path: Union[str, Path, CloudPath],
+        dataset_profile: DatasetProfile = "robot",
+    ):
         """
         Initialize the metadata validator.
 
         Args:
             metadata_path: Path to the custom_metadata.csv file (supports local or cloud paths)
+            dataset_profile: Metadata contract to validate. UMI datasets may omit
+                robot/eval-only fields and include provider-specific columns.
         """
+        if dataset_profile not in REQUIRED_METADATA_COLUMNS_BY_PROFILE:
+            raise ValueError(f"Unsupported dataset profile: {dataset_profile}")
+
         if isinstance(metadata_path, str):
             self.metadata_path = AnyPath(metadata_path)
         else:
             self.metadata_path = metadata_path
+        self.dataset_profile = dataset_profile
+        self.required_columns = REQUIRED_METADATA_COLUMNS_BY_PROFILE[dataset_profile]
+        self.allowed_columns = ALLOWED_METADATA_COLUMNS_BY_PROFILE[dataset_profile]
         self.df = None
         self.errors: List[str] = []
 
@@ -64,19 +80,22 @@ def validate(self) -> bool:
 
     def _check_required_columns(self) -> None:
         """Check that all required columns are present."""
-        missing_columns = set(REQUIRED_METADATA_COLUMNS) - set(self.df.columns)
+        missing_columns = set(self.required_columns) - set(self.df.columns)
         if missing_columns:
             self.errors.append(
                 f"Missing required columns in metadata CSV: {sorted(missing_columns)}"
             )
 
     def _check_unexpected_columns(self) -> None:
         """Check for unexpected columns in the CSV."""
-        unexpected_columns = set(self.df.columns) - set(REQUIRED_METADATA_COLUMNS)
+        if self.allowed_columns is None:
+            return
+
+        unexpected_columns = set(self.df.columns) - set(self.allowed_columns)
         if unexpected_columns:
             self.errors.append(
                 f"Unexpected columns found in metadata CSV: {sorted(unexpected_columns)}. "
-                f"Only the following columns are allowed: {REQUIRED_METADATA_COLUMNS}"
+                f"Only the following columns are allowed: {self.allowed_columns}"
             )
 
     def _check_data_validity(self) -> None:
@@ -97,9 +116,12 @@ def _check_data_validity(self) -> None:
 
         # Check success is boolean
         if "success" in self.df.columns:
-            non_bool_values = self.df[
-                ~self.df["success"].isin([True, False, "True", "False", "true", "false", 0, 1])
-            ]
+            valid_success = self.df["success"].isin(
+                [True, False, "True", "False", "true", "false", 0, 1]
+            )
+            if self.dataset_profile == "umi":
+                valid_success |= self.df["success"].isna()
+            non_bool_values = self.df[~valid_success]
             if len(non_bool_values) > 0:
                 self.errors.append(
                     f"Column 'success' must contain boolean values. "
@@ -168,7 +190,7 @@ def _check_start_timestamp_format(self) -> None:
                 episode_id = row.get("episode_id", f"row_{idx}")
                 invalid_timestamps.append(
                     (idx, episode_id, timestamp, 
-                     f"not a valid numeric timestamp (expected UTC seconds since epoch)")
+                     "not a valid numeric timestamp (expected UTC seconds since epoch)")
                 )
 
         if invalid_timestamps:
@@ -177,7 +199,7 @@ def _check_start_timestamp_format(self) -> None:
                 error_details.append(f"  Row {idx} (episode '{episode_id}'): '{timestamp}' - {reason}")
 
             self.errors.append(
-                f"Column 'start_timestamp' must contain valid UTC timestamps in seconds (Unix epoch time).\n"
+                "Column 'start_timestamp' must contain valid UTC timestamps in seconds (Unix epoch time).\n"
                 + "\n".join(error_details)
             )
 
@@ -255,4 +277,3 @@ def get_errors(self) -> List[str]:
     def get_metadata_df(self) -> pd.DataFrame:
         """Get the loaded metadata DataFrame."""
         return self.df
-

lerobot_validator/schemas.py

@@ -2,7 +2,7 @@
 Schema definitions for metadata CSV and annotation JSON files.
 """
 
-from typing import Dict, Any
+from typing import Dict, Any, List, Literal, Optional
 
 # Expected columns in custom_metadata.csv (on top of what's already in lerobot dataset)
 REQUIRED_METADATA_COLUMNS = [
@@ -17,6 +17,32 @@
     "robot_id",    # the robot hardware
 ]
 
+DatasetProfile = Literal["robot", "umi"]
+
+# UMI datasets are human demonstrations, so they do not require robot/eval-only
+# metadata such as checkpoint_path, success, or robot_id.
+UMI_REQUIRED_METADATA_COLUMNS = [
+    "episode_index",
+    "operator_id",
+    "is_eval_episode",
+    "episode_id",
+    "start_timestamp",
+    "station_id",
+]
+
+REQUIRED_METADATA_COLUMNS_BY_PROFILE: Dict[DatasetProfile, List[str]] = {
+    "robot": REQUIRED_METADATA_COLUMNS,
+    "umi": UMI_REQUIRED_METADATA_COLUMNS,
+}
+
+# UMI partners may include provider-specific metadata fields. The ingestion
+# pipeline preserves or ignores these fields, so the validator should not
+# reject an otherwise valid UMI dataset for including them.
+ALLOWED_METADATA_COLUMNS_BY_PROFILE: Dict[DatasetProfile, Optional[List[str]]] = {
+    "robot": REQUIRED_METADATA_COLUMNS,
+    "umi": None,
+}
+
 # Required fields in the lerobot dataset itself
 REQUIRED_LEROBOT_FIELDS = [
     "fps",   # fps field in info.json (frequency of data collection)
@@ -58,4 +84,3 @@
     },
     "additionalProperties": False,
 }
-