Isaac-GR00T/tests/getting_started/test_policy_md.py

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

from __future__ import annotations

import pathlib

import numpy as np
import pytest
from test_support.readme import extract_code_blocks, find_block, run_readme_python_blocks
from test_support.runtime import get_root
import torch
from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel


pytestmark = pytest.mark.serial


REPO_ROOT = get_root()
POLICY_README = REPO_ROOT / "getting_started" / "policy.md"


# ---------------------------------------------------------------------------
# TinyGr00t — minimal HuggingFace-compatible model registered once per session
# ---------------------------------------------------------------------------


class TinyGr00tConfig(PretrainedConfig):
    model_type = "TinyGr00t"

    def __init__(self, action_horizon: int = 16, max_action_dim: int = 128, **kwargs):
        super().__init__(**kwargs)
        self.action_horizon = action_horizon
        self.max_action_dim = max_action_dim


class TinyGr00tModel(PreTrainedModel):
    """Minimal stand-in for Gr00tN1d6: accepts any kwargs, returns zero action_pred."""

    config_class = TinyGr00tConfig

    def __init__(self, config: TinyGr00tConfig):
        super().__init__(config)
        self._dummy = torch.nn.Linear(1, 1)

    def get_action(self, **kwargs) -> dict:
        batch_size = 1
        for v in kwargs.values():
            if isinstance(v, torch.Tensor) and v.ndim >= 1:
                batch_size = v.shape[0]
                break
        return {
            "action_pred": torch.zeros(
                batch_size,
                self.config.action_horizon,
                self.config.max_action_dim,
                device=self._dummy.weight.device,
                dtype=self._dummy.weight.dtype,
            )
        }


AutoConfig.register("TinyGr00t", TinyGr00tConfig)
AutoModel.register(TinyGr00tConfig, TinyGr00tModel)


@pytest.fixture(scope="session")
def tiny_checkpoint(tmp_path_factory: pytest.TempPathFactory) -> pathlib.Path:
    """Save a TinyGr00t checkpoint to a temp directory and return its path."""
    ckpt_dir = tmp_path_factory.mktemp("tiny_gr00t")
    model = TinyGr00tModel(TinyGr00tConfig(action_horizon=16, max_action_dim=128))
    model.save_pretrained(ckpt_dir)
    return ckpt_dir


# ---------------------------------------------------------------------------
# MockProcessor — processor stand-in for policy.md inference tests
# ---------------------------------------------------------------------------

_POLICY_MD_ACTION_HORIZON = 16
_POLICY_MD_ACTION_DIM = 7


def _make_policy_md_modality_configs() -> dict:
    from gr00t.data.embodiment_tags import EmbodimentTag
    from gr00t.data.types import ModalityConfig

    tag = EmbodimentTag.NEW_EMBODIMENT.value
    return {
        tag: {
            "video": ModalityConfig(delta_indices=[0], modality_keys=["wrist_cam"]),
            "state": ModalityConfig(delta_indices=[0], modality_keys=["joints"]),
            "action": ModalityConfig(
                delta_indices=list(range(_POLICY_MD_ACTION_HORIZON)),
                modality_keys=["joints"],
            ),
            "language": ModalityConfig(delta_indices=[0], modality_keys=["task"]),
        }
    }


class MockProcessor:
    def eval(self) -> None:
        pass

    def get_modality_configs(self) -> dict:
        return _make_policy_md_modality_configs()

    def __call__(self, messages) -> dict:
        return {}

    @property
    def collator(self):
        def _collate(inputs: list) -> dict:
            return {}

        return _collate

    def decode_action(self, action_array, embodiment_tag, batched_states) -> dict:
        return {
            "joints": action_array[:, :_POLICY_MD_ACTION_HORIZON, :_POLICY_MD_ACTION_DIM].astype(
                np.float32
            )
        }


@pytest.fixture
def mock_processor() -> MockProcessor:
    """Return a MockProcessor instance for policy.md tests."""
    return MockProcessor()


# ---------------------------------------------------------------------------
# Tiny model smoke tests
# ---------------------------------------------------------------------------


def test_automodel_loads_tiny_checkpoint(tiny_checkpoint: pathlib.Path) -> None:
    """AutoModel.from_pretrained round-trips the tiny checkpoint."""
    model = AutoModel.from_pretrained(tiny_checkpoint)
    assert isinstance(model, TinyGr00tModel)


def test_get_action_returns_correct_shape(tiny_checkpoint: pathlib.Path) -> None:
    """get_action returns action_pred with shape (B, action_horizon, max_action_dim)."""
    model = AutoModel.from_pretrained(tiny_checkpoint)
    result = model.get_action()
    assert "action_pred" in result
    assert result["action_pred"].shape == (1, 16, 128)


def test_get_action_respects_batch_size(tiny_checkpoint: pathlib.Path) -> None:
    """get_action infers batch size from the first tensor kwarg."""
    model = AutoModel.from_pretrained(tiny_checkpoint)
    result = model.get_action(state=torch.zeros(4, 10))
    assert result["action_pred"].shape[0] == 4


# ---------------------------------------------------------------------------
# Policy.md integration test — extracts blocks from the README directly
# ---------------------------------------------------------------------------


def test_policy_md_steps(
    monkeypatch: pytest.MonkeyPatch,
    tiny_checkpoint: pathlib.Path,
    mock_processor: MockProcessor,
) -> None:
    """Run every named step from policy.md using extracted README code blocks."""
    import gr00t.policy.gr00t_policy as _policy_module

    loaded_model = AutoModel.from_pretrained(tiny_checkpoint)

    class _FakeAutoModel:
        @staticmethod
        def from_pretrained(path, **_):
            return loaded_model

    class _FakeAutoProcessor:
        @staticmethod
        def from_pretrained(path):
            return mock_processor

    monkeypatch.setattr(_policy_module, "AutoModel", _FakeAutoModel)
    monkeypatch.setattr(_policy_module, "AutoProcessor", _FakeAutoProcessor)

    blocks = extract_code_blocks(POLICY_README)

    # --- Loading the Policy ---
    loading_code = (
        find_block(blocks, "Gr00tPolicy(", language="python")
        .code.replace('"/path/to/your/checkpoint"', f'r"{tiny_checkpoint}"')
        .replace('"cuda:0"', '"cpu"')
    )

    # --- Querying Modality Configurations ---
    modality_code = find_block(blocks, "policy.get_modality_config()", language="python").code

    # Injected preamble: resolve undefined dimension variables used by the
    # batched-inference block.
    dims_preamble = (
        "import numpy as np\n"
        "T_video = video_horizon\n"
        "T_state = state_horizon\n"
        "H, W = 224, 224\n"
        "D_state = 7\n"
    )

    # --- Batched Inference ---
    batched_inference_code = find_block(blocks, "wrist_cam", language="python").code.replace(
        "batch_size = 4", "batch_size = 1"
    )

    # --- Running Inference — action access ---
    inference_code = find_block(blocks, "arm_action", language="python").code.replace(
        '"action_name"', "action_keys[0]"
    )

    # --- Resetting the Policy ---
    reset_code = find_block(blocks, "policy.reset()", language="python").code

    run_readme_python_blocks(
        [
            loading_code,
            modality_code,
            dims_preamble,
            batched_inference_code,
            inference_code,
            reset_code,
        ],
        readme_path=POLICY_README,
        repo_root=REPO_ROOT,
    )