diff --git a/README.md b/README.md index ee74cee..2f9de10 100644 --- a/README.md +++ b/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 diff --git a/lerobot_validator/cli.py b/lerobot_validator/cli.py index 185ed7b..cbe861a 100644 --- a/lerobot_validator/cli.py +++ b/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() - diff --git a/lerobot_validator/metadata_validator.py b/lerobot_validator/metadata_validator.py index c23f467..42b3622 100644 --- a/lerobot_validator/metadata_validator.py +++ b/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,7 +80,7 @@ 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)}" @@ -72,11 +88,14 @@ def _check_required_columns(self) -> None: 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 - diff --git a/lerobot_validator/schemas.py b/lerobot_validator/schemas.py index 1532dfb..478599c 100644 --- a/lerobot_validator/schemas.py +++ b/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, } - diff --git a/lerobot_validator/v3_checks.py b/lerobot_validator/v3_checks.py index 155c5fb..3f603b4 100644 --- a/lerobot_validator/v3_checks.py +++ b/lerobot_validator/v3_checks.py @@ -13,22 +13,26 @@ V12: validate_start_timestamp -- start_timestamp must be plausible Unix epoch floats V13: validate_video_frame_count -- video frame counts must match data parquet row counts V14: validate_feature_dtypes -- warn about string-typed features that need special handling + V15: validate_umi_features -- canonical UMI cameras and tracked poses """ import inspect import json import logging +import re import subprocess from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, List, Optional, Union -import numpy as np import pandas as pd from cloudpathlib import AnyPath, CloudPath from lerobot_validator._episodes import load_episodes_df, video_indices -from lerobot_validator.schemas import REQUIRED_METADATA_COLUMNS +from lerobot_validator.schemas import ( + REQUIRED_METADATA_COLUMNS_BY_PROFILE, + DatasetProfile, +) logger = logging.getLogger(__name__) @@ -41,6 +45,25 @@ # Minimum columns required for the converter to function at all. _MIN_REQUIRED_COLUMNS = ["episode_index", "episode_id"] +_UMI_REQUIRED_IMAGE_FEATURE = "base_0_camera/rgb/image" +_UMI_CAMERA_IMAGE_PATTERN = re.compile( + r"^(?:base_\d+_camera|(?:left|right)_wrist_\d+_camera)/rgb/image$" +) +_UMI_REQUIRED_TRACKING_FEATURES = { + "left/position": [3], + "left/quaternion_xyzw": [4], + "left/gripper": [1], + "right/position": [3], + "right/quaternion_xyzw": [4], + "right/gripper": [1], +} +_UMI_CAMERA_TRACKING_PATTERN = re.compile( + r"^(?P.+_camera)/(?Pposition|quaternion_xyzw)$" +) +_UMI_TIMESTAMP_SHAPE = [1] +_UMI_TIMESTAMP_DTYPES = ("int64", "uint64") +_UMI_VIDEO_SHAPE = [480, 640, 3] + @dataclass class Issue: @@ -180,6 +203,141 @@ def validate_feature_shapes(dataset_path: Union[str, Path, CloudPath]) -> List[I return issues +def validate_umi_features(dataset_path: Union[str, Path, CloudPath]) -> List[Issue]: + """Check the public camera, gripper, and tracking contract for UMI datasets. + + Both the partner-facing names and their ``observation/``-prefixed forms + are accepted because the ingestion converter normalizes the former into + the latter. + """ + root = _to_path(dataset_path) + issues: List[Issue] = [] + info = _load_info(root) + + if info is None: + return issues + + raw_features = info.get("features", {}) + if not isinstance(raw_features, dict): + return [ + Issue.error( + "validate_umi_features", + "meta/info.json must contain a 'features' object for UMI validation.", + ) + ] + + features = { + _normalize_umi_feature_name(name): defn + for name, defn in raw_features.items() + if isinstance(name, str) and isinstance(defn, dict) + } + required_features = [_UMI_REQUIRED_IMAGE_FEATURE, *_UMI_REQUIRED_TRACKING_FEATURES] + missing = [name for name in required_features if name not in features] + if missing: + issues.append( + Issue.error( + "validate_umi_features", + f"UMI dataset is missing required features: {missing}", + ) + ) + + for name, defn in features.items(): + if not _UMI_CAMERA_IMAGE_PATTERN.fullmatch(name): + continue + if defn.get("dtype") != "video": + issues.append( + Issue.error( + "validate_umi_features", + f"UMI camera feature '{name}' must use dtype 'video' for MP4 encoding, " + f"got {defn.get('dtype')!r}.", + ) + ) + if defn.get("shape") != _UMI_VIDEO_SHAPE: + issues.append( + Issue.error( + "validate_umi_features", + f"UMI camera feature '{name}' must have shape {_UMI_VIDEO_SHAPE}, " + f"got {defn.get('shape')!r}.", + ) + ) + + tracked_features = dict(_UMI_REQUIRED_TRACKING_FEATURES) + camera_tracking_fields: dict[str, set[str]] = {} + for name in features: + match = _UMI_CAMERA_TRACKING_PATTERN.fullmatch(name) + if match is None: + continue + camera = match.group("camera") + field = match.group("field") + camera_tracking_fields.setdefault(camera, set()).add(field) + tracked_features[name] = [3] if field == "position" else [4] + + for camera, fields in camera_tracking_fields.items(): + missing_fields = {"position", "quaternion_xyzw"} - fields + if missing_fields: + issues.append( + Issue.error( + "validate_umi_features", + f"UMI camera tracking for '{camera}' must include both position and " + f"quaternion_xyzw; missing {sorted(missing_fields)}.", + ) + ) + + for name, expected_shape in tracked_features.items(): + defn = features.get(name) + if defn is not None and defn.get("shape") != expected_shape: + issues.append( + Issue.error( + "validate_umi_features", + f"UMI tracking feature '{name}' must have shape {expected_shape}, " + f"got {defn.get('shape')!r}.", + ) + ) + + for name, defn in features.items(): + if name.endswith("/intrinsics") and defn.get("shape") != [3, 3]: + issues.append( + Issue.error( + "validate_umi_features", + f"UMI camera intrinsics feature '{name}' must have shape [3, 3], " + f"got {defn.get('shape')!r}.", + ) + ) + + missing_timestamps = [] + for name in tracked_features: + if name not in features: + continue + timestamp_name = f"{name}/timestamp" + timestamp_defn = features.get(timestamp_name) + if timestamp_defn is None: + missing_timestamps.append(timestamp_name) + continue + if ( + timestamp_defn.get("shape") != _UMI_TIMESTAMP_SHAPE + or timestamp_defn.get("dtype") not in _UMI_TIMESTAMP_DTYPES + ): + issues.append( + Issue.error( + "validate_umi_features", + f"UMI timestamp feature '{timestamp_name}' must have shape " + f"{_UMI_TIMESTAMP_SHAPE} and dtype int64 or uint64 nanoseconds, " + f"got shape {timestamp_defn.get('shape')!r} and " + f"dtype {timestamp_defn.get('dtype')!r}.", + ) + ) + if missing_timestamps: + issues.append( + Issue.warning( + "validate_umi_features", + f"UMI tracking features should include nanosecond timestamp companions: " + f"{missing_timestamps}", + ) + ) + + return issues + + def validate_timestamps(dataset_path: Union[str, Path, CloudPath]) -> List[Issue]: """Check that data parquet timestamps are relative, not absolute Unix epoch. @@ -253,6 +411,7 @@ def validate_timestamps(dataset_path: Union[str, Path, CloudPath]) -> List[Issue def validate_custom_metadata_csv( dataset_path: Union[str, Path, CloudPath], _df_cache: Optional[Dict[str, pd.DataFrame]] = None, + dataset_profile: DatasetProfile = "robot", ) -> List[Issue]: """Check that meta/custom_metadata.csv exists and has required columns. @@ -305,7 +464,8 @@ def validate_custom_metadata_csv( ) # Warn about expected columns from the full schema that are missing. - missing_optional = [c for c in REQUIRED_METADATA_COLUMNS if c not in df.columns and c not in _MIN_REQUIRED_COLUMNS] + profile_columns = REQUIRED_METADATA_COLUMNS_BY_PROFILE[dataset_profile] + missing_optional = [c for c in profile_columns if c not in df.columns and c not in _MIN_REQUIRED_COLUMNS] if missing_optional: issues.append( Issue.warning( @@ -544,30 +704,39 @@ def validate_feature_dtypes(dataset_path: Union[str, Path, CloudPath]) -> List[I def validate_v3_dataset( dataset_path: Union[str, Path, CloudPath], + dataset_profile: DatasetProfile = "robot", ) -> List[Issue]: """Run all P0 validators and return a combined list of issues. Args: dataset_path: Path to the lerobot dataset directory. + dataset_profile: Dataset-specific validation contract. Returns: A list of Issue objects (errors and warnings). """ + if dataset_profile not in ("robot", "umi"): + raise ValueError(f"Unsupported dataset profile: {dataset_profile}") + all_issues: List[Issue] = [] # Shared cache so V12 reuses the CSV loaded by V11. df_cache: Dict[str, pd.DataFrame] = {} for validator_fn in _P0_VALIDATORS: try: sig = inspect.signature(validator_fn) + kwargs: Dict[str, Any] = {} if "_df_cache" in sig.parameters: - all_issues.extend(validator_fn(dataset_path, _df_cache=df_cache)) - else: - all_issues.extend(validator_fn(dataset_path)) + kwargs["_df_cache"] = df_cache + if "dataset_profile" in sig.parameters: + kwargs["dataset_profile"] = dataset_profile + all_issues.extend(validator_fn(dataset_path, **kwargs)) except Exception as exc: logger.warning("Validator %s raised: %s", validator_fn.__name__, exc) all_issues.append( Issue.error(validator_fn.__name__, f"Validator raised an unexpected exception: {exc}") ) + if dataset_profile == "umi": + all_issues.extend(validate_umi_features(dataset_path)) return all_issues @@ -583,6 +752,11 @@ def _to_path(dataset_path: Union[str, Path, CloudPath]) -> Any: return dataset_path +def _normalize_umi_feature_name(name: str) -> str: + prefix = "observation/" + return name[len(prefix):] if name.startswith(prefix) else name + + def _probe_frame_count(video_path: str) -> Optional[int]: """Use ffprobe to count frames in a video file. Returns None on failure.""" try: diff --git a/lerobot_validator/validator.py b/lerobot_validator/validator.py index 0574a1f..e524c02 100644 --- a/lerobot_validator/validator.py +++ b/lerobot_validator/validator.py @@ -10,6 +10,7 @@ from lerobot_validator.metadata_validator import MetadataValidator from lerobot_validator.annotation_validator import AnnotationValidator from lerobot_validator.lerobot_checks import LerobotDatasetChecker +from lerobot_validator.schemas import DatasetProfile from lerobot_validator.v3_metadata_checker import LerobotV3MetadataChecker from lerobot_validator.v3_checks import validate_v3_dataset @@ -23,6 +24,7 @@ def __init__( self, dataset_path: Union[str, Path, CloudPath], is_eval_data: Optional[bool] = None, + dataset_profile: DatasetProfile = "robot", ): """ Initialize the validator. @@ -31,6 +33,8 @@ def __init__( dataset_path: Path to the lerobot dataset directory (supports local Path or GCP CloudPath) is_eval_data: Optional flag indicating if this is eval data (True) or training data (False). If provided, validates that all episodes have matching is_eval_episode field. + dataset_profile: Dataset-specific validation contract. Use "umi" + for UMI demonstrations. The validator expects to find these files in the dataset's meta folder: - {dataset_path}/meta/custom_metadata.csv @@ -42,13 +46,17 @@ def __init__( else: self.dataset_path = dataset_path self.is_eval_data = is_eval_data + self.dataset_profile = dataset_profile # Construct expected paths in meta folder meta_dir = self.dataset_path / "meta" self.metadata_path = meta_dir / "custom_metadata.csv" self.annotation_path = meta_dir / "custom_annotation.json" - self.metadata_validator = MetadataValidator(self.metadata_path) + self.metadata_validator = MetadataValidator( + self.metadata_path, + dataset_profile=self.dataset_profile, + ) self.annotation_validator = AnnotationValidator(self.annotation_path) self.lerobot_checker = LerobotDatasetChecker(self.dataset_path) self.v3_checker = LerobotV3MetadataChecker(self.dataset_path) @@ -69,8 +77,8 @@ def validate(self) -> bool: # Run individual validators metadata_valid = self.metadata_validator.validate() annotation_valid = self.annotation_validator.validate() - lerobot_valid = self.lerobot_checker.validate() - v3_valid = self.v3_checker.validate() + self.lerobot_checker.validate() + self.v3_checker.validate() # Collect errors self.errors.extend(self.metadata_validator.get_errors()) @@ -79,7 +87,10 @@ def validate(self) -> bool: self.errors.extend(self.v3_checker.get_errors()) # Run P0 v3 validators - v3_issues = validate_v3_dataset(self.dataset_path) + v3_issues = validate_v3_dataset( + self.dataset_path, + dataset_profile=self.dataset_profile, + ) for issue in v3_issues: if issue.level == "error": self.errors.append(f"[{issue.validator}] {issue.message}") @@ -262,4 +273,3 @@ def print_results(self) -> None: print(f"✗ Validation failed with {len(self.errors)} error(s):\n") for i, error in enumerate(self.errors, 1): print(f"{i}. {error}") - diff --git a/tests/test_integration.py b/tests/test_integration.py index cc506f7..25bab88 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -31,7 +31,7 @@ def create_test_dataset(tmpdir): info = { "fps": 30, "codebase_version": "v3.0", - "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + "data_path": "data/chunk-{chunk_index:03d}/episode_{file_index:06d}.parquet", "features": { "action": {"dtype": "float32", "shape": [7]}, }, @@ -114,6 +114,63 @@ def test_full_validation_success(): assert len(validator.get_errors()) == 0 +def test_full_umi_validation_success() -> None: + """Test full validation for an unlabeled UMI dataset.""" + with tempfile.TemporaryDirectory() as tmpdir: + dataset_path = create_test_dataset(tmpdir) + meta_dir = dataset_path / "meta" + + with open(meta_dir / "info.json") as f: + info = json.load(f) + image = {"dtype": "video", "shape": [480, 640, 3]} + image_features = { + "observation/base_0_camera/rgb/image": dict(image), + } + tracking_features = { + "observation/left/position": {"dtype": "float32", "shape": [3]}, + "observation/left/quaternion_xyzw": {"dtype": "float32", "shape": [4]}, + "observation/left/gripper": {"dtype": "float32", "shape": [1]}, + "observation/right/position": {"dtype": "float32", "shape": [3]}, + "observation/right/quaternion_xyzw": {"dtype": "float32", "shape": [4]}, + "observation/right/gripper": {"dtype": "float32", "shape": [1]}, + } + timestamp_features = { + f"{name}/timestamp": {"dtype": "int64", "shape": [1]} + for name in tracking_features + } + info["features"] = image_features | tracking_features | timestamp_features + with open(meta_dir / "info.json", "w") as f: + json.dump(info, f) + + episodes = pd.read_parquet(meta_dir / "episodes.parquet") + for feature_name in image_features: + episodes[f"videos/{feature_name}/chunk_index"] = [0, 0] + episodes[f"videos/{feature_name}/from_timestamp"] = [0.0, 0.0] + episodes.to_parquet(meta_dir / "episodes.parquet", index=False) + + pd.DataFrame( + { + "episode_index": [0, 1], + "operator_id": ["op1", "op1"], + "is_eval_episode": [False, False], + "episode_id": ["umi_001", "umi_002"], + "start_timestamp": [1730455200, 1730458800], + "station_id": ["station_1", "station_1"], + "success": [None, None], + "provider_batch": ["pilot", "pilot"], + } + ).to_csv(meta_dir / "custom_metadata.csv", index=False) + + validator = LerobotDatasetValidator( + dataset_path, + is_eval_data=False, + dataset_profile="umi", + ) + + assert validator.validate() is True + assert validator.get_errors() == [] + + def test_intervention_non_eval_episode(): """Test that intervention on non-eval episode fails validation.""" with tempfile.TemporaryDirectory() as tmpdir: diff --git a/tests/test_is_eval_data_consistency.py b/tests/test_is_eval_data_consistency.py index 14350a3..00b353f 100644 --- a/tests/test_is_eval_data_consistency.py +++ b/tests/test_is_eval_data_consistency.py @@ -31,7 +31,7 @@ def create_test_dataset(tmpdir): info = { "fps": 30, "codebase_version": "v3.0", - "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + "data_path": "data/chunk-{chunk_index:03d}/episode_{file_index:06d}.parquet", "features": { "action": {"dtype": "float32", "shape": [7]}, }, diff --git a/tests/test_umi_profile.py b/tests/test_umi_profile.py new file mode 100644 index 0000000..4d83376 --- /dev/null +++ b/tests/test_umi_profile.py @@ -0,0 +1,148 @@ +"""Tests for UMI-specific metadata and feature validation.""" + +import json +from pathlib import Path +from typing import Any, Dict + +import pandas as pd + +from lerobot_validator.metadata_validator import MetadataValidator +from lerobot_validator.v3_checks import validate_umi_features + + +def test_umi_metadata_allows_unlabeled_non_robot_rows(tmp_path: Path) -> None: + metadata_path = tmp_path / "custom_metadata.csv" + pd.DataFrame( + { + "episode_index": [0], + "operator_id": ["collector@example.com"], + "is_eval_episode": [False], + "episode_id": ["umi_001"], + "start_timestamp": [1780591657.0], + "station_id": ["kitchen_1"], + "success": [None], + "provider_batch": ["pilot"], + } + ).to_csv(metadata_path, index=False) + + validator = MetadataValidator(metadata_path, dataset_profile="umi") + + assert validator.validate() is True + assert validator.get_errors() == [] + + +def test_robot_metadata_profile_stays_strict(tmp_path: Path) -> None: + metadata_path = tmp_path / "custom_metadata.csv" + pd.DataFrame( + { + "episode_index": [0], + "operator_id": ["collector@example.com"], + "is_eval_episode": [False], + "episode_id": ["umi_001"], + "start_timestamp": [1780591657.0], + "station_id": ["kitchen_1"], + } + ).to_csv(metadata_path, index=False) + + validator = MetadataValidator(metadata_path) + + assert validator.validate() is False + assert any("Missing required columns" in error for error in validator.get_errors()) + + +def test_umi_features_accept_public_contract_with_observation_prefix( + tmp_path: Path, +) -> None: + dataset_path = _write_info(tmp_path, _canonical_umi_features()) + + assert validate_umi_features(dataset_path) == [] + + +def test_umi_features_allow_optional_wrist_cameras_and_camera_tracking_to_be_omitted( + tmp_path: Path, +) -> None: + features = _canonical_umi_features() + for name in list(features): + if "wrist_0_camera" in name or "base_0_camera/position" in name or "base_0_camera/quaternion" in name: + del features[name] + dataset_path = _write_info(tmp_path, features) + + assert validate_umi_features(dataset_path) == [] + + +def test_umi_features_report_missing_fields_and_wrong_pose_shapes( + tmp_path: Path, +) -> None: + features = _canonical_umi_features() + del features["observation/right/gripper"] + features["observation/left/position"]["shape"] = [4] + features["observation/base_0_camera/intrinsics"]["shape"] = [4, 4] + dataset_path = _write_info(tmp_path, features) + + issues = validate_umi_features(dataset_path) + + assert any("right/gripper" in issue.message for issue in issues) + assert any( + "left/position" in issue.message and "shape [3]" in issue.message + for issue in issues + ) + assert any("base_0_camera/intrinsics" in issue.message and "shape [3, 3]" in issue.message for issue in issues) + + +def test_umi_features_warn_missing_tracking_timestamps(tmp_path: Path) -> None: + features = _canonical_umi_features() + del features["observation/left/gripper/timestamp"] + dataset_path = _write_info(tmp_path, features) + + issues = validate_umi_features(dataset_path) + + assert len(issues) == 1 + assert issues[0].level == "warning" + assert "left/gripper/timestamp" in issues[0].message + + +def _write_info(tmp_path: Path, features: Dict[str, Dict[str, Any]]) -> Path: + dataset_path = tmp_path / "dataset" + meta_path = dataset_path / "meta" + meta_path.mkdir(parents=True) + (meta_path / "info.json").write_text( + json.dumps( + { + "codebase_version": "v3.0", + "features": features, + } + ) + ) + return dataset_path + + +def _canonical_umi_features() -> Dict[str, Dict[str, Any]]: + image = {"dtype": "video", "shape": [480, 640, 3]} + tracked = { + "observation/left/position": {"dtype": "float32", "shape": [3]}, + "observation/left/quaternion_xyzw": {"dtype": "float32", "shape": [4]}, + "observation/left/gripper": {"dtype": "float32", "shape": [1]}, + "observation/right/position": {"dtype": "float32", "shape": [3]}, + "observation/right/quaternion_xyzw": {"dtype": "float32", "shape": [4]}, + "observation/right/gripper": {"dtype": "float32", "shape": [1]}, + "observation/base_0_camera/position": {"dtype": "float32", "shape": [3]}, + "observation/base_0_camera/quaternion_xyzw": { + "dtype": "float32", + "shape": [4], + }, + } + timestamps = { + f"{name}/timestamp": {"dtype": "int64", "shape": [1]} + for name in tracked + } + return { + "observation/base_0_camera/rgb/image": dict(image), + "observation/left_wrist_0_camera/rgb/image": dict(image), + "observation/right_wrist_0_camera/rgb/image": dict(image), + "observation/base_0_camera/intrinsics": { + "dtype": "float32", + "shape": [3, 3], + }, + **tracked, + **timestamps, + }