Isaac-GR00T/tests/gr00t/model/test_model_forward.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.

"""
Test Gr00tN1d7 model forward pass and action generation with dummy data.

These tests construct a minimal Gr00tN1d7 model (with mocked backbone) and
verify that forward() computes a scalar loss and get_action() produces
action predictions of the expected shape.
"""

from unittest.mock import MagicMock, patch

from gr00t.configs.model.gr00t_n1d7 import Gr00tN1d7Config
import pytest
import torch
from transformers.feature_extraction_utils import BatchFeature


def _make_small_config(**overrides) -> Gr00tN1d7Config:
    """Return a minimal config for fast instantiation."""
    defaults = dict(
        model_name="nvidia/Cosmos-Reason2-2B",
        backbone_model_type="qwen",
        backbone_embedding_dim=64,
        hidden_size=64,
        input_embedding_dim=64,
        max_state_dim=7,
        max_action_dim=7,
        action_horizon=4,
        state_history_length=1,
        num_inference_timesteps=2,
        max_num_embodiments=4,
        add_pos_embed=True,
        use_vlln=True,
        max_seq_len=32,
        use_alternate_vl_dit=False,
        select_layer=1,
        reproject_vision=False,
        use_flash_attention=False,
        load_bf16=False,
        tune_top_llm_layers=0,
        backbone_trainable_params_fp32=False,
        tune_llm=False,
        tune_visual=False,
        tune_projector=True,
        tune_diffusion_model=True,
        tune_vlln=True,
        state_dropout_prob=0.0,
        diffusion_model_cfg={
            "positional_embeddings": None,
            "num_layers": 2,
            "num_attention_heads": 2,
            "attention_head_dim": 32,
            "norm_type": "ada_norm",
            "dropout": 0.0,
            "final_dropout": False,
            "output_dim": 64,
            "interleave_self_attention": True,
        },
    )
    defaults.update(overrides)
    return Gr00tN1d7Config(**defaults)


def _make_mock_backbone(config, seq_len=8):
    """Return a mock backbone that produces correctly-shaped outputs."""
    backbone = MagicMock()

    def fake_forward(vl_input):
        B = 1
        # Try to infer batch size from input
        for v in vl_input.values():
            if isinstance(v, torch.Tensor) and v.dim() >= 2:
                B = v.shape[0]
                break
        device = next(
            (v.device for v in vl_input.values() if isinstance(v, torch.Tensor)),
            torch.device("cpu"),
        )
        dtype = next(
            (
                v.dtype
                for v in vl_input.values()
                if isinstance(v, torch.Tensor) and v.is_floating_point()
            ),
            torch.float32,
        )
        return BatchFeature(
            data={
                "backbone_features": torch.randn(
                    B, seq_len, config.backbone_embedding_dim, device=device, dtype=dtype
                ),
                "backbone_attention_mask": torch.ones(B, seq_len, device=device, dtype=torch.long),
                "image_mask": torch.ones(B, seq_len, device=device, dtype=torch.bool),
            }
        )

    backbone.side_effect = fake_forward
    backbone.prepare_input = lambda x: BatchFeature(data=x)
    return backbone


@pytest.fixture
def small_model():
    """Build a Gr00tN1d7 with mocked backbone (no GPU/download required)."""
    config = _make_small_config()

    with patch("gr00t.model.gr00t_n1d7.gr00t_n1d7.get_backbone_cls") as mock_get_cls:
        mock_get_cls.return_value = lambda **kwargs: _make_mock_backbone(config)
        with patch("gr00t.model.gr00t_n1d7.processing_gr00t_n1d7.build_processor"):
            from gr00t.model.gr00t_n1d7.gr00t_n1d7 import Gr00tN1d7

            model = Gr00tN1d7(config)

    model.eval()
    return model, config


def _make_dummy_inputs(config, batch_size=2):
    """Create dummy input tensors matching the model's expected format."""
    return {
        "state": torch.randn(batch_size, config.state_history_length, config.max_state_dim),
        "action": torch.randn(batch_size, config.action_horizon, config.max_action_dim),
        "embodiment_id": torch.zeros(batch_size, dtype=torch.long),
        "action_mask": torch.ones(batch_size, config.action_horizon, config.max_action_dim),
    }


class TestGr00tN1d7Forward:
    """Test model forward pass produces valid loss."""

    def test_forward_returns_loss(self, small_model):
        model, config = small_model
        inputs = _make_dummy_inputs(config)
        output = model.forward(inputs)
        assert "loss" in output
        assert output["loss"].dim() == 0, "loss should be scalar"
        assert torch.isfinite(output["loss"]), "loss should be finite"

    def test_forward_returns_action_loss_and_mask(self, small_model):
        model, config = small_model
        inputs = _make_dummy_inputs(config)
        output = model.forward(inputs)
        assert "action_loss" in output
        assert "action_mask" in output
        assert output["action_loss"].shape == (2, config.action_horizon, config.max_action_dim)

    def test_forward_loss_requires_grad(self, small_model):
        model, config = small_model
        model.train()
        inputs = _make_dummy_inputs(config)
        output = model.forward(inputs)
        assert output["loss"].requires_grad

    def test_forward_different_batch_sizes(self, small_model):
        model, config = small_model
        for bs in [1, 4]:
            inputs = _make_dummy_inputs(config, batch_size=bs)
            output = model.forward(inputs)
            assert output["loss"].dim() == 0


class TestGr00tN1d7GetAction:
    """Test model action generation."""

    def test_get_action_shape(self, small_model):
        model, config = small_model
        inputs = _make_dummy_inputs(config, batch_size=1)
        del inputs["action"]  # get_action uses diffusion denoising, not ground-truth
        output = model.get_action(inputs)
        assert "action_pred" in output
        assert output["action_pred"].shape == (1, config.action_horizon, config.max_action_dim)

    def test_get_action_no_grad(self, small_model):
        model, config = small_model
        inputs = _make_dummy_inputs(config, batch_size=1)
        del inputs["action"]
        output = model.get_action(inputs)
        assert not output["action_pred"].requires_grad


class TestGr00tN1d7Config:
    """Test config creation and serialization."""

    def test_default_config(self):
        config = Gr00tN1d7Config()
        assert config.model_type == "Gr00tN1d7"
        assert config.max_state_dim == 132
        assert config.action_horizon == 40

    def test_custom_config(self):
        config = Gr00tN1d7Config(max_state_dim=10, action_horizon=8)
        assert config.max_state_dim == 10
        assert config.action_horizon == 8

    def test_to_filtered_dict(self):
        config = Gr00tN1d7Config()
        d = config.to_filtered_dict(exclude_augment=True)
        assert "random_rotation_angle" not in d
        assert "hidden_size" in d

    def test_to_filtered_json(self):
        config = Gr00tN1d7Config()
        j = config.to_filtered_json()
        assert isinstance(j, str)
        import json

        parsed = json.loads(j)
        assert parsed["model_type"] == "Gr00tN1d7"