[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/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 24,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 62,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 61,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
         EnvironmentVariableCategory.INTERNAL: 2,
     }
@@ -179,7 +179,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
     )
     assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
         ModelEnvironmentVariableDisposition.PROMOTE: 37,
-        ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 7,
+        ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 13,
         ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5,

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

tests/e2e/online_serving/test_gr00t_openpi_expansion.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """End-to-end online serving test for GR00T N1.7 through the OpenPI robot endpoint."""
 
 import os
@@ -29,8 +29,7 @@ test_params = [
         OmniServerParams(
             model=MODEL,
             stage_config_path=get_deploy_config_path("Gr00tN1d7.yaml"),
-            server_args=["--disable-log-stats"],
-            env_dict={"VLLM_DISABLE_COMPILE_CACHE": "1", "GR00T_NOISE_SEED": "42"},
+            env_dict={"VLLM_DISABLE_COMPILE_CACHE": "1"},
             init_timeout=1200,
             stage_init_timeout=900,
         ),
@@ -162,14 +161,14 @@ def test_gr00t_n1d7_openpi_online(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_gr00t_n1d7_openpi_precision(omni_server, openai_client) -> None:
-    """Assert actions match Isaac-GR00T reference (GR00T_NOISE_SEED=42, zero inputs)."""
+    """Assert actions match Isaac-GR00T reference (request seed=42, zero inputs)."""
     response = openai_client.send_robot_openpi_ws_request(
         {
             "operations": [
                 {"endpoint": "reset", "payload": {}},
                 {
                     "endpoint": "infer",
-                    "payload": build_openpi_droid_observation(session_id="gr00t-precision-e2e"),
+                    "payload": build_openpi_droid_observation(session_id="gr00t-precision-e2e", seed=42),
                 },
             ],
         }

tests/entrypoints/openai_api/test_openpi_serving.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
 import json
 import threading
@@ -218,6 +221,18 @@ def test_build_request_uses_unique_engine_request_id_per_inference():
     assert request_a.request_id != request_b.request_id
 
 
+def test_build_request_forwards_seed_to_sampling_params():
+    """``seed`` in the inference message is the engine-level seed; omitted, the request auto-seeds."""
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config())
+
+    seeded = serving._build_request({"prompt": "pick up the object", "seed": 42}, session_id="s", reset=True)
+    unseeded = serving._build_request({"prompt": "pick up the object"}, session_id="s", reset=False)
+
+    assert seeded.sampling_params.seed == 42
+    assert "seed" not in seeded.sampling_params.extra_args["robot_obs"]
+    assert isinstance(unseeded.sampling_params.seed, int)
+
+
 def test_infer_keeps_session_state_but_uses_unique_engine_request_ids():
     engine = RecordingEngine()
     serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine)

tests/helpers/client.py

@@ -152,11 +152,14 @@ class OpenPIWebSocketResponse:
     action_tensors: list[np.ndarray] | None = None
 
 
-def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[str, Any]:
-    """Build a minimal DROID-style observation payload for the OpenPI robot endpoint."""
+def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke", seed: int | None = None) -> dict[str, Any]:
+    """Build a minimal DROID-style observation payload for the OpenPI robot endpoint.
+
+    ``seed`` pins the policy's sampling noise for this request (see ``docs/serving/openpi_api.md``).
+    """
     identity_eef_9d = np.zeros((1, 1, 9), dtype=np.float32)
     identity_eef_9d[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32)
-    return {
+    observation: dict[str, Any] = {
         "session_id": session_id,
         "video": {
             "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
@@ -169,6 +172,9 @@ def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[s
         },
         "language": {"annotation.language.language_instruction": [["pick up the object"]]},
     }
+    if seed is not None:
+        observation["seed"] = seed
+    return observation
 
 
 DREAMZERO_DEFAULT_PROMPT = (