File size: 6,362 Bytes
e479c46 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Validate that bundled demo datasets conform to the GR00T LeRobot format
described in getting_started/data_preparation.md.
Checks:
- Required directory structure (meta/, data/chunk-*, videos/chunk-*)
- Required meta files (info.json, episodes.jsonl, tasks.jsonl, modality.json)
- modality.json schema (state/action keys with start/end, video keys)
- Parquet files exist and contain expected columns
- Video files exist as .mp4
"""
from __future__ import annotations
import json
import pathlib
import pytest
ROOT = pathlib.Path(__file__).resolve().parents[3]
DEMO_DATA_DIR = ROOT / "demo_data"
# All bundled demo datasets to validate.
DEMO_DATASETS = [
p for p in DEMO_DATA_DIR.iterdir() if p.is_dir() and (p / "meta" / "modality.json").exists()
]
@pytest.fixture(params=[str(d.relative_to(ROOT)) for d in DEMO_DATASETS])
def dataset_path(request):
return ROOT / request.param
class TestDatasetStructure:
"""Validate the directory structure required by data_preparation.md."""
def test_meta_dir_exists(self, dataset_path):
assert (dataset_path / "meta").is_dir(), f"Missing meta/ in {dataset_path}"
def test_data_dir_exists(self, dataset_path):
assert (dataset_path / "data").is_dir(), f"Missing data/ in {dataset_path}"
def test_videos_dir_exists(self, dataset_path):
assert (dataset_path / "videos").is_dir(), f"Missing videos/ in {dataset_path}"
def test_required_meta_files(self, dataset_path):
meta = dataset_path / "meta"
for filename in ["info.json", "episodes.jsonl", "tasks.jsonl", "modality.json"]:
assert (meta / filename).is_file(), f"Missing {filename} in {meta}"
def test_data_has_parquet_files(self, dataset_path):
parquets = list((dataset_path / "data").rglob("*.parquet"))
assert parquets, f"No .parquet files found in {dataset_path / 'data'}"
def test_videos_has_mp4_files(self, dataset_path):
mp4s = list((dataset_path / "videos").rglob("*.mp4"))
assert mp4s, f"No .mp4 files found in {dataset_path / 'videos'}"
class TestModalityJson:
"""Validate modality.json schema per data_preparation.md."""
def test_modality_json_is_valid_json(self, dataset_path):
modality_path = dataset_path / "meta" / "modality.json"
with open(modality_path) as f:
data = json.load(f)
assert isinstance(data, dict)
def test_modality_has_state_and_action(self, dataset_path):
with open(dataset_path / "meta" / "modality.json") as f:
data = json.load(f)
assert "state" in data, "modality.json missing 'state' key"
assert "action" in data, "modality.json missing 'action' key"
def test_state_keys_have_start_end(self, dataset_path):
with open(dataset_path / "meta" / "modality.json") as f:
data = json.load(f)
for key, spec in data["state"].items():
assert "start" in spec, f"state.{key} missing 'start'"
assert "end" in spec, f"state.{key} missing 'end'"
assert isinstance(spec["start"], int), f"state.{key}.start must be int"
assert isinstance(spec["end"], int), f"state.{key}.end must be int"
assert spec["end"] > spec["start"], (
f"state.{key}: end ({spec['end']}) must be > start ({spec['start']})"
)
def test_action_keys_have_start_end(self, dataset_path):
with open(dataset_path / "meta" / "modality.json") as f:
data = json.load(f)
for key, spec in data["action"].items():
assert "start" in spec, f"action.{key} missing 'start'"
assert "end" in spec, f"action.{key} missing 'end'"
assert isinstance(spec["start"], int), f"action.{key}.start must be int"
assert isinstance(spec["end"], int), f"action.{key}.end must be int"
assert spec["end"] > spec["start"], (
f"action.{key}: end ({spec['end']}) must be > start ({spec['start']})"
)
def test_video_keys_have_original_key(self, dataset_path):
with open(dataset_path / "meta" / "modality.json") as f:
data = json.load(f)
assert "video" in data, (
f"{dataset_path}: modality.json must define a 'video' section per data_preparation.md"
)
for key, spec in data["video"].items():
assert "original_key" in spec, f"video.{key} missing 'original_key'"
class TestEpisodesJsonl:
"""Validate episodes.jsonl format."""
def test_episodes_are_valid_jsonl(self, dataset_path):
episodes_path = dataset_path / "meta" / "episodes.jsonl"
episodes = []
with open(episodes_path) as f:
for line in f:
line = line.strip()
if line:
episodes.append(json.loads(line))
assert len(episodes) > 0, "episodes.jsonl is empty"
for ep in episodes:
assert "episode_index" in ep, f"Episode missing 'episode_index': {ep}"
assert "length" in ep, f"Episode missing 'length': {ep}"
class TestTasksJsonl:
"""Validate tasks.jsonl format."""
def test_tasks_are_valid_jsonl(self, dataset_path):
tasks_path = dataset_path / "meta" / "tasks.jsonl"
tasks = []
with open(tasks_path) as f:
for line in f:
line = line.strip()
if line:
tasks.append(json.loads(line))
assert len(tasks) > 0, "tasks.jsonl is empty"
for task in tasks:
assert "task_index" in task, f"Task missing 'task_index': {task}"
assert "task" in task, f"Task missing 'task': {task}"
|