lerobot/tests/datasets/test_dataset_metadata.py

#!/usr/bin/env python

# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# 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.
"""Contract tests for LeRobotDatasetMetadata."""

import json

import numpy as np
import pytest

pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")

from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.utils import INFO_PATH
from tests.fixtures.constants import DEFAULT_FPS, DUMMY_ROBOT_TYPE

# ── helpers ──────────────────────────────────────────────────────────

SIMPLE_FEATURES = {
    "state": {"dtype": "float32", "shape": (6,), "names": None},
    "action": {"dtype": "float32", "shape": (6,), "names": None},
}

VIDEO_FEATURES = {
    **SIMPLE_FEATURES,
    "observation.images.laptop": {
        "dtype": "video",
        "shape": (64, 96, 3),
        "names": ["height", "width", "channels"],
        "info": None,
    },
}

IMAGE_FEATURES = {
    **SIMPLE_FEATURES,
    "observation.images.laptop": {
        "dtype": "image",
        "shape": (64, 96, 3),
        "names": ["height", "width", "channels"],
        "info": None,
    },
}


def _make_dummy_stats(features: dict) -> dict:
    """Create minimal episode stats matching the given features."""
    stats = {}
    for key, ft in features.items():
        if ft["dtype"] in ("image", "video"):
            channels = ft["shape"][-1]
            stat_shape = (channels, 1, 1)
            stats[key] = {
                "max": np.ones(stat_shape, dtype=np.float32),
                "mean": np.full(stat_shape, 0.5, dtype=np.float32),
                "min": np.zeros(stat_shape, dtype=np.float32),
                "std": np.full(stat_shape, 0.25, dtype=np.float32),
                "count": np.array([5]),
            }
        elif ft["dtype"] in ("float32", "float64", "int64"):
            stats[key] = {
                "max": np.ones(ft["shape"], dtype=np.float32),
                "mean": np.full(ft["shape"], 0.5, dtype=np.float32),
                "min": np.zeros(ft["shape"], dtype=np.float32),
                "std": np.full(ft["shape"], 0.25, dtype=np.float32),
                "count": np.array([5]),
            }
    return stats


# ── Construction contracts ───────────────────────────────────────────


def test_create_produces_valid_info_on_disk(tmp_path):
    """create() writes info.json and the returned object reflects the provided settings."""
    root = tmp_path / "new_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/meta",
        fps=DEFAULT_FPS,
        features=SIMPLE_FEATURES,
        robot_type=DUMMY_ROBOT_TYPE,
        root=root,
        use_videos=False,
    )

    # info.json was written to disk
    assert (root / INFO_PATH).exists()
    with open(root / INFO_PATH) as f:
        info_on_disk = json.load(f)

    assert meta.fps == DEFAULT_FPS
    assert meta.robot_type == DUMMY_ROBOT_TYPE
    assert "state" in meta.features
    assert "action" in meta.features
    assert info_on_disk["fps"] == DEFAULT_FPS


def test_create_starts_with_zero_counts(tmp_path):
    """A freshly created metadata has zero episode/frame/task counts."""
    root = tmp_path / "empty_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/empty", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )

    assert meta.total_episodes == 0
    assert meta.total_frames == 0
    assert meta.total_tasks == 0
    assert meta.tasks is None
    assert meta.episodes is None
    assert meta.stats is None


def test_create_with_videos_sets_video_path(tmp_path):
    """When features include video-dtype keys, create() produces a non-None video_path."""
    root = tmp_path / "video_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/video", fps=DEFAULT_FPS, features=VIDEO_FEATURES, root=root, use_videos=True
    )

    assert meta.video_path is not None
    assert len(meta.video_keys) == 1
    assert "observation.images.laptop" in meta.video_keys


def test_create_without_videos_has_no_video_path(tmp_path):
    """When use_videos=False and no video features, video_path is None."""
    root = tmp_path / "no_video"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/novid", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )

    assert meta.video_path is None
    assert meta.video_keys == []


@pytest.mark.parametrize(
    ("marker_field", "marker_key"),
    [
        ("info", "is_depth_map"),
        ("info", "video.is_depth_map"),
        ("video_info", "video.is_depth_map"),
    ],
    ids=["info.is_depth_map", "info.video.is_depth_map_legacy", "video_info.video.is_depth_map_legacy"],
)
def test_depth_keys_property_filters_by_marker(tmp_path, marker_field, marker_key):
    """``depth_keys`` recognises the canonical and the two legacy marker variants."""
    depth_feature = {
        "dtype": "video",
        "shape": (64, 96, 1),
        "names": ["height", "width", "channels"],
        marker_field: {marker_key: True},
    }
    features = {
        **VIDEO_FEATURES,
        "observation.images.laptop_depth": depth_feature,
    }
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/depth_keys",
        fps=DEFAULT_FPS,
        features=features,
        root=tmp_path / f"depth_keys_{marker_field}_{marker_key.replace('.', '_')}",
    )

    assert set(meta.video_keys) == {"observation.images.laptop", "observation.images.laptop_depth"}
    assert meta.depth_keys == ["observation.images.laptop_depth"]


def test_depth_keys_empty_when_no_marker(tmp_path):
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/no_depth", fps=DEFAULT_FPS, features=VIDEO_FEATURES, root=tmp_path / "no_depth"
    )
    assert meta.depth_keys == []


def test_create_raises_on_existing_directory(tmp_path):
    """create() raises if root directory already exists."""
    root = tmp_path / "existing"
    root.mkdir()

    with pytest.raises(FileExistsError):
        LeRobotDatasetMetadata.create(
            repo_id="test/exists", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
        )


def test_init_loads_existing_metadata(tmp_path, lerobot_dataset_metadata_factory, info_factory):
    """When metadata files exist on disk, __init__ loads them correctly."""
    root = tmp_path / "load_test"
    info = info_factory(total_episodes=3, total_frames=150, total_tasks=1, use_videos=False)
    meta = lerobot_dataset_metadata_factory(root=root, info=info)

    assert meta.total_episodes == 3
    assert meta.total_frames == 150
    assert meta.fps == info.fps


# ── Property accessors ───────────────────────────────────────────────


def test_property_accessors_reflect_info(tmp_path):
    """Properties return values consistent with the info dict."""
    root = tmp_path / "props_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/props",
        fps=DEFAULT_FPS,
        features=IMAGE_FEATURES,
        robot_type=DUMMY_ROBOT_TYPE,
        root=root,
        use_videos=False,
    )

    assert meta.fps == DEFAULT_FPS
    assert meta.robot_type == DUMMY_ROBOT_TYPE
    # shapes should be tuples
    for _key, shape in meta.shapes.items():
        assert isinstance(shape, tuple)
    # image_keys should contain the image feature
    assert "observation.images.laptop" in meta.image_keys
    # camera_keys is a superset of image_keys and video_keys
    assert set(meta.image_keys + meta.video_keys) == set(meta.camera_keys)


def test_data_path_is_formattable(tmp_path):
    """data_path contains format placeholders that can be .format()-ed."""
    root = tmp_path / "fmt_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/fmt", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )

    formatted = meta.data_path.format(chunk_index=0, file_index=0)
    assert "chunk" in formatted.lower() or "0" in formatted


# ── Task management ──────────────────────────────────────────────────


def test_save_episode_tasks_creates_tasks_dataframe(tmp_path):
    """On a fresh metadata, save_episode_tasks() creates the tasks DataFrame."""
    root = tmp_path / "task_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/task", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )
    assert meta.tasks is None

    meta.save_episode_tasks(["Pick up the cube"])

    assert meta.tasks is not None
    assert len(meta.tasks) == 1
    assert "Pick up the cube" in meta.tasks.index


def test_save_episode_tasks_is_additive(tmp_path):
    """New tasks are added; existing tasks keep their original index."""
    root = tmp_path / "additive_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/add", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )

    meta.save_episode_tasks(["Task A"])
    idx_a = meta.get_task_index("Task A")

    meta.save_episode_tasks(["Task A", "Task B"])
    assert meta.get_task_index("Task A") == idx_a  # unchanged
    assert meta.get_task_index("Task B") is not None
    assert len(meta.tasks) == 2


def test_get_task_index_returns_none_for_unknown(tmp_path):
    """get_task_index() returns None for an unknown task."""
    root = tmp_path / "unknown_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/unknown", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )
    meta.save_episode_tasks(["Known task"])

    assert meta.get_task_index("Known task") == 0
    assert meta.get_task_index("Unknown task") is None


def test_save_episode_tasks_rejects_duplicates(tmp_path):
    """save_episode_tasks() raises ValueError on duplicate task strings."""
    root = tmp_path / "dup_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/dup", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )

    with pytest.raises(ValueError):
        meta.save_episode_tasks(["Same task", "Same task"])


# ── Episode saving ───────────────────────────────────────────────────


def test_save_episode_increments_counters(tmp_path):
    """After save_episode(), total_episodes and total_frames increase."""
    root = tmp_path / "ep_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/ep", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )
    meta.save_episode_tasks(["Task 1"])
    stats = _make_dummy_stats(meta.features)

    meta.save_episode(
        episode_index=0,
        episode_length=10,
        episode_tasks=["Task 1"],
        episode_stats=stats,
        episode_metadata={},
    )

    assert meta.total_episodes == 1
    assert meta.total_frames == 10


def test_save_episode_updates_stats(tmp_path):
    """After save_episode(), .stats is non-None and has feature keys."""
    root = tmp_path / "stats_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/stats", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )
    meta.save_episode_tasks(["Task 1"])
    stats = _make_dummy_stats(meta.features)

    meta.save_episode(
        episode_index=0,
        episode_length=5,
        episode_tasks=["Task 1"],
        episode_stats=stats,
        episode_metadata={},
    )

    assert meta.stats is not None
    # Stats should contain at least the user-defined feature keys
    for key in SIMPLE_FEATURES:
        assert key in meta.stats


# ── Chunk settings ───────────────────────────────────────────────────


def test_update_chunk_settings_persists(tmp_path):
    """update_chunk_settings() changes values and writes info.json."""
    root = tmp_path / "chunk_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/chunk", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )
    original = meta.get_chunk_settings()

    meta.update_chunk_settings(chunks_size=500)
    assert meta.chunks_size == 500
    assert meta.chunks_size != original["chunks_size"] or original["chunks_size"] == 500

    # Verify persisted
    with open(root / INFO_PATH) as f:
        info_on_disk = json.load(f)
    assert info_on_disk["chunks_size"] == 500


def test_update_chunk_settings_rejects_non_positive(tmp_path):
    """update_chunk_settings() raises ValueError for <= 0 values."""
    root = tmp_path / "bad_chunk"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/bad", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )

    with pytest.raises(ValueError):
        meta.update_chunk_settings(chunks_size=0)
    with pytest.raises(ValueError):
        meta.update_chunk_settings(data_files_size_in_mb=-1)


# ── Finalization ─────────────────────────────────────────────────────


def test_finalize_is_idempotent(tmp_path):
    """Calling finalize() multiple times does not raise."""
    root = tmp_path / "fin_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/fin", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
    )

    meta.finalize()
    meta.finalize()  # second call should not raise


def test_finalize_flushes_buffered_metadata(tmp_path):
    """Episodes saved before finalize() are written to parquet."""
    root = tmp_path / "flush_ds"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/flush",
        fps=DEFAULT_FPS,
        features=SIMPLE_FEATURES,
        root=root,
        use_videos=False,
        metadata_buffer_size=100,  # large buffer so nothing auto-flushes
    )
    meta.save_episode_tasks(["Task 1"])
    stats = _make_dummy_stats(meta.features)

    # Save a few episodes (won't auto-flush since buffer_size=100)
    for i in range(3):
        meta.save_episode(
            episode_index=i,
            episode_length=5,
            episode_tasks=["Task 1"],
            episode_stats=stats,
            episode_metadata={},
        )

    # Before finalize, the parquet might not exist yet
    meta.finalize()

    # After finalize, episodes parquet should exist
    episodes_dir = root / "meta" / "episodes"
    assert episodes_dir.exists()
    parquet_files = list(episodes_dir.rglob("*.parquet"))
    assert len(parquet_files) > 0


# ── Tools accessor ───────────────────────────────────────────────────


def test_tools_falls_back_to_default_when_info_has_no_tools_field(tmp_path):
    """meta.tools returns DEFAULT_TOOLS when info.json doesn't declare any."""
    from lerobot.datasets.language import DEFAULT_TOOLS

    root = tmp_path / "no_tools"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/no_tools",
        fps=DEFAULT_FPS,
        features=SIMPLE_FEATURES,
        root=root,
        use_videos=False,
    )

    assert meta.tools == DEFAULT_TOOLS
    # info.json on disk should NOT include a `tools` key for clean datasets
    with open(root / INFO_PATH) as f:
        info_on_disk = json.load(f)
    assert "tools" not in info_on_disk


def test_tools_reads_declared_tools_from_info_json(tmp_path):
    """A `tools` list written into info.json survives load → meta.tools.

    Regression test for the bug where ``DatasetInfo.from_dict`` silently
    dropped the ``tools`` key (no matching dataclass field), so
    ``meta.tools`` always returned ``DEFAULT_TOOLS`` regardless of
    what was on disk.
    """
    from lerobot.datasets.io_utils import load_info

    root = tmp_path / "with_tools"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/with_tools",
        fps=DEFAULT_FPS,
        features=SIMPLE_FEATURES,
        root=root,
        use_videos=False,
    )

    custom_tool = {
        "type": "function",
        "function": {
            "name": "record_observation",
            "description": "Capture a still image.",
            "parameters": {
                "type": "object",
                "properties": {"label": {"type": "string"}},
                "required": ["label"],
            },
        },
    }
    info_path = root / INFO_PATH
    with open(info_path) as f:
        raw = json.load(f)
    raw["tools"] = [custom_tool]
    with open(info_path, "w") as f:
        json.dump(raw, f)

    # Reload info from disk and rebind it on the metadata object
    meta.info = load_info(root)
    assert meta.tools == [custom_tool]


def test_tools_round_trip_through_dataset_info(tmp_path):
    """A `tools` list survives DatasetInfo.from_dict / to_dict."""
    from lerobot.datasets.utils import DatasetInfo

    raw = {
        "codebase_version": "v3.1",
        "fps": 30,
        "features": SIMPLE_FEATURES,
        "tools": [{"type": "function", "function": {"name": "say"}}],
    }
    info = DatasetInfo.from_dict(raw)
    assert info.tools == raw["tools"]
    assert info.to_dict()["tools"] == raw["tools"]


def test_tools_setter_persists_to_info_json_and_reloads(tmp_path):
    """Assigning meta.tools writes info.json and reloads meta.info."""
    from lerobot.datasets.io_utils import load_info

    root = tmp_path / "set_tools"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/set_tools",
        fps=DEFAULT_FPS,
        features=SIMPLE_FEATURES,
        root=root,
        use_videos=False,
    )

    custom_tool = {
        "type": "function",
        "function": {
            "name": "record_observation",
            "description": "Capture a still image.",
            "parameters": {
                "type": "object",
                "properties": {"label": {"type": "string"}},
                "required": ["label"],
            },
        },
    }
    meta.tools = [custom_tool]

    # In-memory metadata reflects the new catalog ...
    assert meta.tools == [custom_tool]
    assert meta.info.tools == [custom_tool]
    # ... and a fresh read from disk agrees.
    assert load_info(root).tools == [custom_tool]


def test_tools_setter_clears_key_when_set_to_none(tmp_path):
    """Setting meta.tools back to None drops the key and restores the default."""
    from lerobot.datasets.language import DEFAULT_TOOLS

    root = tmp_path / "clear_tools"
    meta = LeRobotDatasetMetadata.create(
        repo_id="test/clear_tools",
        fps=DEFAULT_FPS,
        features=SIMPLE_FEATURES,
        root=root,
        use_videos=False,
    )

    meta.tools = [{"type": "function", "function": {"name": "say"}}]
    meta.tools = None

    assert meta.tools == DEFAULT_TOOLS
    with open(root / INFO_PATH) as f:
        info_on_disk = json.load(f)
    assert "tools" not in info_on_disk