Add files using upload-large-folder tool

e479c46 verified 2 months ago

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

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,
    )

	# 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,
	)