groot_deployment / Isaac-GR00T /tests /gr00t /data /test_stats_cache_fingerprint.py

Add files using upload-large-folder tool

e479c46 verified 2 months ago

11 kB

# 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.

"""CPU-only tests for the schema-fingerprint cache in ``generate_rel_stats``.
Without the fingerprint guard, an existing ``meta/relative_stats.json`` was
reused whenever a per-embodiment ``action_key`` name matched, regardless of
whether the inputs that drive the computation (``delta_indices``, ``format``,
``state_key``, ...) had since changed — silently corrupting normalization at
training time.
"""

from dataclasses import replace
import json
from unittest.mock import patch

from gr00t.configs.data.embodiment_configs import MODALITY_CONFIGS
from gr00t.data.embodiment_tags import EmbodimentTag
from gr00t.data.stats import (
    LE_ROBOT_REL_STATS_FILENAME,
    STATS_FINGERPRINTS_KEY,
    _compute_relative_action_fingerprint,
    generate_rel_stats,
)
from gr00t.data.types import ActionFormat
import numpy as np
import pytest


# Real registered embodiment with two RELATIVE action keys (eef_9d, joint_position)
# and one ABSOLUTE (gripper_position) — exercises the cache path without
# needing parquet data once calculate_stats_for_key is mocked.
EMBODIMENT = EmbodimentTag.OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT
RELATIVE_KEYS = ("eef_9d", "joint_position")


def _stub_stats():
    """Synthetic stats payload shaped like calculate_stats_for_key's return."""
    return {
        "max": np.ones(9, dtype=np.float32),
        "min": -np.ones(9, dtype=np.float32),
        "q01": -np.ones(9, dtype=np.float32) * 0.99,
        "q99": np.ones(9, dtype=np.float32) * 0.99,
        "mean": np.zeros(9, dtype=np.float32),
        "std": np.ones(9, dtype=np.float32),
    }


@pytest.fixture
def dataset_dir(tmp_path):
    (tmp_path / "meta").mkdir()
    return tmp_path


@pytest.fixture
def mock_calculate(monkeypatch):
    """Replace the heavy parquet-driven computation with a counter-stub."""
    calls = []

    def fake(dataset_path, embodiment_tag, group_key, max_episodes=-1):
        calls.append((str(dataset_path), embodiment_tag.value, group_key))
        return _stub_stats()

    monkeypatch.setattr("gr00t.data.stats.calculate_stats_for_key", fake)
    return calls


# ---------------------------------------------------------------------------
# _compute_relative_action_fingerprint — pure helper, no I/O
# ---------------------------------------------------------------------------


class TestFingerprintHelper:
    def test_deterministic(self):
        a = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
        b = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
        assert a == b

    def test_format_prefix(self):
        fp = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
        assert fp.startswith("sha256:") and len(fp) == len("sha256:") + 64

    def test_distinct_per_action_key(self):
        a = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
        b = _compute_relative_action_fingerprint(EMBODIMENT, "joint_position")
        assert a != b

    @pytest.mark.parametrize(
        "mutator",
        [
            pytest.param(
                lambda cfg: replace(cfg, format=ActionFormat.XYZ_ROTVEC),
                id="format-change",
            ),
            pytest.param(
                lambda cfg: replace(cfg, state_key="other_state_key"),
                id="state-key-change",
            ),
        ],
    )
    def test_changes_when_action_config_changes(self, mutator):
        """Mutating fields the loader actually reads must invalidate the cache."""
        cfgs = MODALITY_CONFIGS[EMBODIMENT.value]
        action_modality = cfgs["action"]
        idx = action_modality.modality_keys.index("eef_9d")
        original_action_config = action_modality.action_configs[idx]

        baseline = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")

        new_configs = list(action_modality.action_configs)
        new_configs[idx] = mutator(original_action_config)
        with patch.dict(
            MODALITY_CONFIGS[EMBODIMENT.value]["action"].__dict__,
            {"action_configs": new_configs},
        ):
            mutated = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")

        assert mutated != baseline

    def test_changes_when_delta_indices_change(self):
        baseline = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
        action_modality = MODALITY_CONFIGS[EMBODIMENT.value]["action"]
        with patch.dict(
            action_modality.__dict__,
            {"delta_indices": [0, 1, 2]},
        ):
            mutated = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
        assert mutated != baseline


# ---------------------------------------------------------------------------
# generate_rel_stats — cache hit / miss / mixed via mocked calculate_stats_for_key
# ---------------------------------------------------------------------------


class TestGenerateRelStatsCache:
    def test_first_run_computes_and_persists_fingerprints(self, dataset_dir, mock_calculate):
        generate_rel_stats(dataset_dir, EMBODIMENT)

        assert sorted(c[2] for c in mock_calculate) == sorted(RELATIVE_KEYS)

        with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME) as f:
            payload = json.load(f)
        assert STATS_FINGERPRINTS_KEY in payload
        for key in RELATIVE_KEYS:
            assert key in payload
            assert key in payload[STATS_FINGERPRINTS_KEY]
            assert payload[STATS_FINGERPRINTS_KEY][key].startswith("sha256:")

    def test_second_run_is_full_cache_hit(self, dataset_dir, mock_calculate):
        generate_rel_stats(dataset_dir, EMBODIMENT)
        mock_calculate.clear()

        generate_rel_stats(dataset_dir, EMBODIMENT)

        assert mock_calculate == [], "fresh fingerprints must produce zero recompute"

    def test_legacy_file_without_fingerprints_is_regenerated(self, dataset_dir, mock_calculate):
        """Pre-existing relative_stats.json from before this fix must be recomputed."""
        legacy_payload = {key: {"_legacy": True} for key in RELATIVE_KEYS}
        with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME, "w") as f:
            json.dump(legacy_payload, f)

        generate_rel_stats(dataset_dir, EMBODIMENT)

        assert sorted(c[2] for c in mock_calculate) == sorted(RELATIVE_KEYS)
        with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME) as f:
            payload = json.load(f)
        for key in RELATIVE_KEYS:
            assert "_legacy" not in payload[key]
            assert payload[STATS_FINGERPRINTS_KEY][key].startswith("sha256:")

    def test_stale_fingerprint_triggers_regenerate(self, dataset_dir, mock_calculate):
        generate_rel_stats(dataset_dir, EMBODIMENT)

        with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME) as f:
            payload = json.load(f)
        payload[STATS_FINGERPRINTS_KEY]["eef_9d"] = "sha256:" + "0" * 64
        with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME, "w") as f:
            json.dump(payload, f)

        mock_calculate.clear()
        generate_rel_stats(dataset_dir, EMBODIMENT)

        recomputed = [c[2] for c in mock_calculate]
        assert recomputed == ["eef_9d"], f"only the stale key should regenerate; got {recomputed}"

    def test_modality_config_change_invalidates_cache(self, dataset_dir, mock_calculate):
        """The whole point of the fingerprint: silently changed config must not get a stale cache hit."""
        generate_rel_stats(dataset_dir, EMBODIMENT)
        mock_calculate.clear()

        action_modality = MODALITY_CONFIGS[EMBODIMENT.value]["action"]
        idx = action_modality.modality_keys.index("eef_9d")
        new_configs = list(action_modality.action_configs)
        new_configs[idx] = replace(new_configs[idx], format=ActionFormat.XYZ_ROTVEC)
        with patch.dict(action_modality.__dict__, {"action_configs": new_configs}):
            generate_rel_stats(dataset_dir, EMBODIMENT)

        recomputed = [c[2] for c in mock_calculate]
        assert "eef_9d" in recomputed, "format change must trigger eef_9d regeneration"
        assert "joint_position" not in recomputed, (
            "joint_position config did not change; must not be regenerated"
        )

    def test_partial_cache_only_recomputes_missing(self, dataset_dir, mock_calculate):
        """Pre-fill cache for one key only; the other should be the only one computed."""
        eef_fp = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
        prefilled = {
            "eef_9d": {k: v.tolist() for k, v in _stub_stats().items()},
            STATS_FINGERPRINTS_KEY: {"eef_9d": eef_fp},
        }
        with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME, "w") as f:
            json.dump(prefilled, f)

        generate_rel_stats(dataset_dir, EMBODIMENT)

        assert [c[2] for c in mock_calculate] == ["joint_position"]
        with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME) as f:
            payload = json.load(f)
        for key in RELATIVE_KEYS:
            assert key in payload
            assert key in payload[STATS_FINGERPRINTS_KEY]


# ---------------------------------------------------------------------------
# Backward compat — downstream consumers must not see __fingerprints__ as data
# ---------------------------------------------------------------------------


def test_fingerprints_key_does_not_collide_with_action_keys():
    """Reserved sentinel must not match any registered action key string anywhere."""
    for tag, cfg in MODALITY_CONFIGS.items():
        action = cfg.get("action")
        if action is None or action.action_configs is None:
            continue
        assert STATS_FINGERPRINTS_KEY not in action.modality_keys, (
            f"reserved sentinel {STATS_FINGERPRINTS_KEY!r} collides with an "
            f"action key in embodiment {tag!r}"
        )


def test_per_action_key_payload_unchanged_shape(dataset_dir, mock_calculate):
    """The per-action_key dict written to disk must contain only the stat fields,
    not the fingerprint — downstream loader iterates ``stats[key].keys()``
    expecting ``{mean, std, min, max, q01, q99}``.
    """
    generate_rel_stats(dataset_dir, EMBODIMENT)
    with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME) as f:
        payload = json.load(f)
    for key in RELATIVE_KEYS:
        assert set(payload[key].keys()) == {"mean", "std", "min", "max", "q01", "q99"}

	# 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.

	"""CPU-only tests for the schema-fingerprint cache in ``generate_rel_stats``.

	Without the fingerprint guard, an existing ``meta/relative_stats.json`` was
	reused whenever a per-embodiment ``action_key`` name matched, regardless of
	whether the inputs that drive the computation (``delta_indices``, ``format``,
	``state_key``, ...) had since changed — silently corrupting normalization at
	training time.
	"""

	from dataclasses import replace
	import json
	from unittest.mock import patch

	from gr00t.configs.data.embodiment_configs import MODALITY_CONFIGS
	from gr00t.data.embodiment_tags import EmbodimentTag
	from gr00t.data.stats import (
	LE_ROBOT_REL_STATS_FILENAME,
	STATS_FINGERPRINTS_KEY,
	_compute_relative_action_fingerprint,
	generate_rel_stats,
	)
	from gr00t.data.types import ActionFormat
	import numpy as np
	import pytest


	# Real registered embodiment with two RELATIVE action keys (eef_9d, joint_position)
	# and one ABSOLUTE (gripper_position) — exercises the cache path without
	# needing parquet data once calculate_stats_for_key is mocked.
	EMBODIMENT = EmbodimentTag.OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT
	RELATIVE_KEYS = ("eef_9d", "joint_position")


	def _stub_stats():
	"""Synthetic stats payload shaped like calculate_stats_for_key's return."""
	return {
	"max": np.ones(9, dtype=np.float32),
	"min": -np.ones(9, dtype=np.float32),
	"q01": -np.ones(9, dtype=np.float32) * 0.99,
	"q99": np.ones(9, dtype=np.float32) * 0.99,
	"mean": np.zeros(9, dtype=np.float32),
	"std": np.ones(9, dtype=np.float32),
	}


	@pytest.fixture
	def dataset_dir(tmp_path):
	(tmp_path / "meta").mkdir()
	return tmp_path


	@pytest.fixture
	def mock_calculate(monkeypatch):
	"""Replace the heavy parquet-driven computation with a counter-stub."""
	calls = []

	def fake(dataset_path, embodiment_tag, group_key, max_episodes=-1):
	calls.append((str(dataset_path), embodiment_tag.value, group_key))
	return _stub_stats()

	monkeypatch.setattr("gr00t.data.stats.calculate_stats_for_key", fake)
	return calls


	# ---------------------------------------------------------------------------
	# _compute_relative_action_fingerprint — pure helper, no I/O
	# ---------------------------------------------------------------------------


	class TestFingerprintHelper:
	def test_deterministic(self):
	a = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
	b = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
	assert a == b

	def test_format_prefix(self):
	fp = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
	assert fp.startswith("sha256:") and len(fp) == len("sha256:") + 64

	def test_distinct_per_action_key(self):
	a = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
	b = _compute_relative_action_fingerprint(EMBODIMENT, "joint_position")
	assert a != b

	@pytest.mark.parametrize(
	"mutator",
	[
	pytest.param(
	lambda cfg: replace(cfg, format=ActionFormat.XYZ_ROTVEC),
	id="format-change",
	),
	pytest.param(
	lambda cfg: replace(cfg, state_key="other_state_key"),
	id="state-key-change",
	),
	],
	)
	def test_changes_when_action_config_changes(self, mutator):
	"""Mutating fields the loader actually reads must invalidate the cache."""
	cfgs = MODALITY_CONFIGS[EMBODIMENT.value]
	action_modality = cfgs["action"]
	idx = action_modality.modality_keys.index("eef_9d")
	original_action_config = action_modality.action_configs[idx]

	baseline = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")

	new_configs = list(action_modality.action_configs)
	new_configs[idx] = mutator(original_action_config)
	with patch.dict(
	MODALITY_CONFIGS[EMBODIMENT.value]["action"].__dict__,
	{"action_configs": new_configs},
	):
	mutated = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")

	assert mutated != baseline

	def test_changes_when_delta_indices_change(self):
	baseline = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
	action_modality = MODALITY_CONFIGS[EMBODIMENT.value]["action"]
	with patch.dict(
	action_modality.__dict__,
	{"delta_indices": [0, 1, 2]},
	):
	mutated = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
	assert mutated != baseline


	# ---------------------------------------------------------------------------
	# generate_rel_stats — cache hit / miss / mixed via mocked calculate_stats_for_key
	# ---------------------------------------------------------------------------


	class TestGenerateRelStatsCache:
	def test_first_run_computes_and_persists_fingerprints(self, dataset_dir, mock_calculate):
	generate_rel_stats(dataset_dir, EMBODIMENT)

	assert sorted(c[2] for c in mock_calculate) == sorted(RELATIVE_KEYS)

	with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME) as f:
	payload = json.load(f)
	assert STATS_FINGERPRINTS_KEY in payload
	for key in RELATIVE_KEYS:
	assert key in payload
	assert key in payload[STATS_FINGERPRINTS_KEY]
	assert payload[STATS_FINGERPRINTS_KEY][key].startswith("sha256:")

	def test_second_run_is_full_cache_hit(self, dataset_dir, mock_calculate):
	generate_rel_stats(dataset_dir, EMBODIMENT)
	mock_calculate.clear()

	generate_rel_stats(dataset_dir, EMBODIMENT)

	assert mock_calculate == [], "fresh fingerprints must produce zero recompute"

	def test_legacy_file_without_fingerprints_is_regenerated(self, dataset_dir, mock_calculate):
	"""Pre-existing relative_stats.json from before this fix must be recomputed."""
	legacy_payload = {key: {"_legacy": True} for key in RELATIVE_KEYS}
	with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME, "w") as f:
	json.dump(legacy_payload, f)

	generate_rel_stats(dataset_dir, EMBODIMENT)

	assert sorted(c[2] for c in mock_calculate) == sorted(RELATIVE_KEYS)
	with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME) as f:
	payload = json.load(f)
	for key in RELATIVE_KEYS:
	assert "_legacy" not in payload[key]
	assert payload[STATS_FINGERPRINTS_KEY][key].startswith("sha256:")

	def test_stale_fingerprint_triggers_regenerate(self, dataset_dir, mock_calculate):
	generate_rel_stats(dataset_dir, EMBODIMENT)

	with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME) as f:
	payload = json.load(f)
	payload[STATS_FINGERPRINTS_KEY]["eef_9d"] = "sha256:" + "0" * 64
	with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME, "w") as f:
	json.dump(payload, f)

	mock_calculate.clear()
	generate_rel_stats(dataset_dir, EMBODIMENT)

	recomputed = [c[2] for c in mock_calculate]
	assert recomputed == ["eef_9d"], f"only the stale key should regenerate; got {recomputed}"

	def test_modality_config_change_invalidates_cache(self, dataset_dir, mock_calculate):
	"""The whole point of the fingerprint: silently changed config must not get a stale cache hit."""
	generate_rel_stats(dataset_dir, EMBODIMENT)
	mock_calculate.clear()

	action_modality = MODALITY_CONFIGS[EMBODIMENT.value]["action"]
	idx = action_modality.modality_keys.index("eef_9d")
	new_configs = list(action_modality.action_configs)
	new_configs[idx] = replace(new_configs[idx], format=ActionFormat.XYZ_ROTVEC)
	with patch.dict(action_modality.__dict__, {"action_configs": new_configs}):
	generate_rel_stats(dataset_dir, EMBODIMENT)

	recomputed = [c[2] for c in mock_calculate]
	assert "eef_9d" in recomputed, "format change must trigger eef_9d regeneration"
	assert "joint_position" not in recomputed, (
	"joint_position config did not change; must not be regenerated"
	)

	def test_partial_cache_only_recomputes_missing(self, dataset_dir, mock_calculate):
	"""Pre-fill cache for one key only; the other should be the only one computed."""
	eef_fp = _compute_relative_action_fingerprint(EMBODIMENT, "eef_9d")
	prefilled = {
	"eef_9d": {k: v.tolist() for k, v in _stub_stats().items()},
	STATS_FINGERPRINTS_KEY: {"eef_9d": eef_fp},
	}
	with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME, "w") as f:
	json.dump(prefilled, f)

	generate_rel_stats(dataset_dir, EMBODIMENT)

	assert [c[2] for c in mock_calculate] == ["joint_position"]
	with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME) as f:
	payload = json.load(f)
	for key in RELATIVE_KEYS:
	assert key in payload
	assert key in payload[STATS_FINGERPRINTS_KEY]


	# ---------------------------------------------------------------------------
	# Backward compat — downstream consumers must not see __fingerprints__ as data
	# ---------------------------------------------------------------------------


	def test_fingerprints_key_does_not_collide_with_action_keys():
	"""Reserved sentinel must not match any registered action key string anywhere."""
	for tag, cfg in MODALITY_CONFIGS.items():
	action = cfg.get("action")
	if action is None or action.action_configs is None:
	continue
	assert STATS_FINGERPRINTS_KEY not in action.modality_keys, (
	f"reserved sentinel {STATS_FINGERPRINTS_KEY!r} collides with an "
	f"action key in embodiment {tag!r}"
	)


	def test_per_action_key_payload_unchanged_shape(dataset_dir, mock_calculate):
	"""The per-action_key dict written to disk must contain only the stat fields,
	not the fingerprint — downstream loader iterates ``stats[key].keys()``
	expecting ``{mean, std, min, max, q01, q99}``.
	"""
	generate_rel_stats(dataset_dir, EMBODIMENT)
	with open(dataset_dir / LE_ROBOT_REL_STATS_FILENAME) as f:
	payload = json.load(f)
	for key in RELATIVE_KEYS:
	assert set(payload[key].keys()) == {"mean", "std", "min", "max", "q01", "q99"}