[Bugfix][Model] GR00T-N1.7: honor the per-request seed for flow-matching noise (#7253)

Signed-off-by: liangmengh <liangmengh@nvidia.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

tests/diffusion/models/gr00t/test_noise_seed.py

@@ -0,0 +1,177 @@
+# 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)

tests/diffusion/models/gr00t/test_pipeline.py

@@ -1,4 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
+from typing import Any
 
 import numpy as np
 import pytest
@@ -12,12 +16,12 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class FakeGr00tPolicy:
-    instances = []
+    instances: list["FakeGr00tPolicy"] = []
 
     def __init__(self, **kwargs):
         self.kwargs = kwargs
         self.reset_calls = 0
-        self.seen_obs = None
+        self.seen_obs: dict[str, Any] | None = None
         self.embodiment_tag = SimpleNamespace(value="fake_embodiment")
         self.language_key = "annotation.language.language_instruction"
         self.modality_configs = {
@@ -40,7 +44,7 @@ class FakeGr00tPolicy:
         )
         FakeGr00tPolicy.instances.append(self)
 
-    def get_action(self, obs):
+    def get_action(self, obs, options=None):
         self.seen_obs = obs
         return {
             "arm": np.array([[[1.0, 2.0]]], dtype=np.float64),
@@ -107,6 +111,7 @@ def test_forward_returns_dict_actions_in_output():
     assert actions["arm"].dtype == np.float32
     np.testing.assert_allclose(actions["arm"], np.array([[[1.0, 2.0]]], dtype=np.float32))
     policy = FakeGr00tPolicy.instances[0]
+    assert policy.seen_obs is not None
     assert "video" in policy.seen_obs
     assert policy.seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]}
     assert "images" not in policy.seen_obs