vllm-omni/tests/diffusion/models/gr00t/test_noise_seed.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The request's seed must reach GR00T's flow-matching noise.

The runner builds ``sampling_params.generator`` from ``sampling_params.seed``, the pipeline hands
that generator to the policy, and the action head draws its initial noise from it, so the same
seed reproduces the same action chunk and different seeds give independent ones. These tests pin
both ends of that plumbing.
"""

from __future__ import annotations

from types import SimpleNamespace

import numpy as np
import pytest
import torch
from torch import nn
from transformers.feature_extraction_utils import BatchFeature

from vllm_omni.diffusion.models.gr00t import pipeline_gr00t
from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config
from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7ActionHead
from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

ROBOT_OBS = {
    "images": {"cam": np.zeros((1, 1, 8, 8, 3), dtype=np.uint8)},
    "state": {"joint": np.zeros((1, 1, 2), dtype=np.float32)},
    "prompt": "pick the cube",
}


class _OptionsRecordingPolicy:
    """Stand-in for ``Gr00tPolicy``; the pipeline only needs ``language_key`` and ``get_action``."""

    def __init__(self, *, model_path: str, embodiment_tag: str, device: str, strict: bool) -> None:
        del model_path, embodiment_tag, device, strict  # the pipeline's constructor contract; unused here
        self.language_key = "annotation.language.language_instruction"
        self.seen_options: dict[str, object] | None = None

    def get_action(self, obs: dict[str, object], options: dict[str, object] | None = None):
        self.seen_options = options
        return {"arm": np.zeros((1, 2, 2), dtype=np.float32)}, {}


@pytest.fixture
def pipeline(monkeypatch) -> Gr00tN1d7Pipeline:
    monkeypatch.setattr(pipeline_gr00t, "Gr00tPolicy", _OptionsRecordingPolicy)
    od_config = SimpleNamespace(
        model="nvidia/GR00T-N1.7-3B",
        model_config={"embodiment_tag": "LIBERO_PANDA", "strict": False},
        custom_pipeline_args={},
    )
    return Gr00tN1d7Pipeline(od_config=od_config)


def _request(**sampling_kwargs) -> OmniDiffusionRequest:
    return OmniDiffusionRequest(
        prompt="pick",
        request_id="req",
        sampling_params=OmniDiffusionSamplingParams(extra_args={"robot_obs": ROBOT_OBS}, **sampling_kwargs),
    )


def test_runner_generator_is_passed_through_untouched(pipeline):
    """In serving the runner has already built ``sampling_params.generator``; the head must get that object."""
    generator = torch.Generator().manual_seed(99)

    assert pipeline.forward(_request(generator=generator)).error is None

    assert pipeline.policy.seen_options["generator"] is generator


def test_seed_without_generator_builds_one_on_the_pipeline_device(pipeline):
    """Callers that bypass the runner (direct use, tests) still get the seed honoured."""
    pipeline.forward(_request(seed=1234))

    generator = pipeline.policy.seen_options["generator"]
    assert isinstance(generator, torch.Generator)
    assert generator.initial_seed() == 1234
    assert generator.device.type == pipeline.device


def test_auto_assigned_seed_reaches_the_head(pipeline):
    """``OmniDiffusionRequest`` seeds every request, so the head is never left on the global RNG."""
    req = _request()
    assert req.sampling_params.seed is not None

    pipeline.forward(req)

    assert pipeline.policy.seen_options["generator"].initial_seed() == req.sampling_params.seed


def _captured_warnings(monkeypatch) -> list[tuple]:
    """vLLM loggers do not propagate to the root logger, so ``caplog`` cannot see them; record calls instead."""
    captured: list[tuple] = []
    monkeypatch.setattr(pipeline_gr00t.logger, "warning", lambda *args, **kwargs: captured.append(args))
    return captured


def test_unexpected_generator_form_warns_and_falls_back_to_the_seed(pipeline, monkeypatch):
    """GR00T draws one noise tensor per request, so a list of several generators cannot be honoured."""
    warnings = _captured_warnings(monkeypatch)
    generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]

    pipeline.forward(_request(generator=generators, seed=7))

    assert len(warnings) == 1
    assert "torch.Generator" in warnings[0][0]
    assert pipeline.policy.seen_options["generator"].initial_seed() == 7


def test_unexpected_generator_form_without_a_seed_warns_and_leaves_the_global_rng(pipeline, monkeypatch):
    warnings = _captured_warnings(monkeypatch)
    generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]

    pipeline.forward(_request(generator=generators))

    assert len(warnings) == 1
    assert pipeline.policy.seen_options is None


def _noise_only_head(action_horizon: int, action_dim: int) -> Gr00tN1d7ActionHead:
    """An action head whose single denoising step is the identity, so ``action_pred`` is the initial noise."""
    config = Gr00tN1d7Config(
        max_action_dim=action_dim,
        action_horizon=action_horizon,
        num_inference_timesteps=1,
        add_pos_embed=False,
        use_alternate_vl_dit=False,
    )
    head = Gr00tN1d7ActionHead.__new__(Gr00tN1d7ActionHead)
    nn.Module.__init__(head)
    head.config = config
    head.action_dim = config.max_action_dim
    head.action_horizon = config.action_horizon
    head.num_inference_timesteps = config.num_inference_timesteps
    head.num_timestep_buckets = config.num_timestep_buckets
    head.action_encoder = lambda actions, timesteps, embodiment_id: actions
    head.model = lambda hidden_states, encoder_hidden_states, timestep: hidden_states
    head.action_decoder = lambda hidden_states, embodiment_id: torch.zeros_like(hidden_states)
    return head


def _predict(head: Gr00tN1d7ActionHead, options) -> torch.Tensor:
    batch, seq, dim = 2, 3, head.action_dim
    return head.get_action_with_features(
        backbone_features=torch.zeros(batch, seq, dim),
        state_features=torch.zeros(batch, 1, dim),
        embodiment_id=torch.zeros(batch, dtype=torch.long),
        backbone_output=BatchFeature(data={}),
        action_input=BatchFeature(data={}),
        options=options,
    )["action_pred"]


def test_head_draws_its_noise_from_the_request_generator():
    head = _noise_only_head(action_horizon=4, action_dim=3)
    expected = torch.randn(2, 4, 3, generator=torch.Generator().manual_seed(5))

    assert torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(5)}), expected)
    assert not torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(6)}), expected)


def test_head_without_generator_stays_on_the_global_rng():
    """Upstream Isaac-GR00T parity: no generator means a bare ``torch.randn`` on the global RNG."""
    head = _noise_only_head(action_horizon=4, action_dim=3)

    torch.manual_seed(5)
    expected = torch.randn(2, 4, 3)
    torch.manual_seed(5)

    assert torch.equal(_predict(head, None), expected)