[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>

docs/serving/openpi_api.md

@@ -64,7 +64,7 @@ connect
   <- msgpack(policy_server_config)
 
 infer
-  -> msgpack({"endpoint": "infer", "session_id": "...", ...observation})
+  -> msgpack({"endpoint": "infer", "session_id": "...", "seed": <optional int>, ...observation})
   <- msgpack(ndarray | dict[str, ndarray])
 
 reset
@@ -92,6 +92,12 @@ frames are not observation messages.
 - The policy pipeline owns observation transforms and persistent model state,
   normally keyed by `session_id`; the API layer forwards the raw observation
   dictionary.
+- An optional integer `seed` in an inference message becomes the engine
+  request's `sampling_params.seed`. A policy whose sampling consumes the
+  request generator (GR00T-N1.7) returns the same action chunk for the same
+  observation and seed. If omitted, the engine assigns a random per-request
+  seed, so unseeded inferences are independent of each other. The API layer
+  consumes `seed`; it is not forwarded inside the observation.
 
 ## Limits and Errors
 

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 = (

vllm_omni/config/environment_variable_inventory.py

@@ -175,7 +175,6 @@ _MODEL_PROMOTE = (
 
 _MODEL_REQUEST_SCOPE = (
     "DOTS_TTS_DIT_NUM_STEPS",
-    "GR00T_NOISE_SEED",
     "MAGI2_NEGATIVE_PROMPT",
     "STEP_AUDIO2_DEFAULT_PROMPT_WAV",
     "VLLM_GEPARD_END_FADE_MS",

vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py

@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import logging
-import os
 from typing import Any
 
 import torch
@@ -138,26 +137,22 @@ class Gr00tN1d7ActionHead(nn.Module):
             state_features: [B, state_horizon, input_embedding_dim]
             embodiment_id: [B] (embodiment IDs)
             backbone_output: Output from the backbone model
+            options: ``generator`` (a ``torch.Generator`` on the model device) draws
+                the initial action noise; ``None`` leaves ``torch.randn`` on the
+                global RNG, as upstream Isaac-GR00T does. The ``rtc_*`` keys drive
+                real-time chunking when ``action_input`` carries a previous chunk.
         """
         vl_embeds = backbone_features
 
         batch_size = vl_embeds.shape[0]
         device = vl_embeds.device
-        _seed_env = os.environ.get("GR00T_NOISE_SEED")
-        if _seed_env is not None:
-            _gen = torch.Generator(device=device).manual_seed(int(_seed_env))
-            actions = torch.randn(
-                size=(batch_size, self.config.action_horizon, self.action_dim),
-                dtype=vl_embeds.dtype,
-                device=device,
-                generator=_gen,
-            )
-        else:
-            actions = torch.randn(
-                size=(batch_size, self.config.action_horizon, self.action_dim),
-                dtype=vl_embeds.dtype,
-                device=device,
-            )
+        generator = options.get("generator") if options else None
+        actions = torch.randn(
+            size=(batch_size, self.config.action_horizon, self.action_dim),
+            dtype=vl_embeds.dtype,
+            device=device,
+            generator=generator,
+        )
 
         dt = 1.0 / self.num_inference_timesteps
         vel_strength = torch.ones_like(actions)

vllm_omni/diffusion/models/gr00t/pipeline_gr00t.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
 
 from collections.abc import Iterable, Mapping
 from typing import Any
@@ -13,6 +13,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.models.gr00t.policy import Gr00tPolicy
 from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 logger = init_logger(__name__)
 
@@ -30,6 +31,10 @@ class Gr00tN1d7Pipeline(nn.Module):
     vLLM-Omni owns the serving integration: OpenPI observations arrive through
     `sampling_params.extra_args["robot_obs"]`, this pipeline runs GR00T policy
     inference, and actions are returned through `DiffusionOutput.output["actions"]`.
+
+    The request's `seed` (materialised by the runner into `sampling_params.generator`)
+    draws the flow-matching noise, so repeating a request with the same seed reproduces
+    its action chunk.
     """
 
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
@@ -123,7 +128,8 @@ class Gr00tN1d7Pipeline(nn.Module):
             self.reset()
 
         policy_obs = _normalize_observation(robot_obs, language_key=self.policy.language_key)
-        result = self.policy.get_action(policy_obs)
+        generator = _noise_generator(req.sampling_params, self.device)
+        result = self.policy.get_action(policy_obs, {"generator": generator} if generator is not None else None)
         actions = result[0] if isinstance(result, tuple) else result
         if not isinstance(actions, Mapping):
             return DiffusionOutput(error=f"GR00T policy returned {type(actions).__name__}; expected dict actions.")
@@ -132,6 +138,35 @@ class Gr00tN1d7Pipeline(nn.Module):
         return DiffusionOutput(output={"actions": _to_float32_action_dict(actions)})
 
 
+def _noise_generator(sampling_params: OmniDiffusionSamplingParams, device: str) -> torch.Generator | None:
+    """Pick the generator that draws the action head's initial noise.
+
+    The runner materialises ``sampling_params.seed`` into ``sampling_params.generator`` on
+    the model device before every forward, and ``OmniDiffusionRequest`` assigns a random
+    seed when the caller omits one, so in serving the generator is always present. The
+    seed fallback serves callers that bypass the runner (direct use, unit tests). ``None``
+    leaves the action head on the global RNG, as upstream Isaac-GR00T runs. GR00T draws one
+    noise tensor per request, so a generator in any other form (for example a list of several
+    generators) is ignored with a warning and the seed fallback applies.
+    """
+    generator = sampling_params.generator
+    if isinstance(generator, list) and len(generator) == 1:
+        generator = generator[0]
+    if isinstance(generator, torch.Generator):
+        return generator
+    if generator is not None:
+        received = f"a list of {len(generator)} generators" if isinstance(generator, list) else type(generator).__name__
+        logger.warning(
+            "GR00T expects a single torch.Generator per request but got %s; ignoring it. The initial noise "
+            "falls back to sampling_params.seed=%s, or to the global RNG when that is None.",
+            received,
+            sampling_params.seed,
+        )
+    if sampling_params.seed is not None:
+        return torch.Generator(device=device).manual_seed(int(sampling_params.seed))
+    return None
+
+
 def _normalize_observation(robot_obs: Mapping[str, Any], *, language_key: str) -> dict[str, Any]:
     obs: dict[str, Any] = {}
     if "video" in robot_obs:

vllm_omni/diffusion/models/gr00t/policy.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
 
 from pathlib import Path
 from typing import Any
@@ -328,7 +328,8 @@ class Gr00tPolicy:
 
         Args:
             observation: Batched observation dictionary
-            options: Optional parameters (currently unused)
+            options: Optional per-request model parameters, forwarded verbatim to
+                the action head (e.g. ``generator``)
 
         Returns:
             Tuple of (actions_dict, info_dict)
@@ -351,7 +352,7 @@ class Gr00tPolicy:
 
         # Step 4: Run model inference to predict actions
         with torch.inference_mode():
-            model_pred = self.model.get_action(**collated_inputs)
+            model_pred = self.model.get_action(**collated_inputs, options=options)
         normalized_action = model_pred["action_pred"].float()
 
         # Step 5: Decode actions from normalized space back to physical units

vllm_omni/entrypoints/openpi/serving.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
 
 """Serving layer for robot policy inference via `/v1/realtime/robot/openpi`.
 
@@ -12,7 +12,7 @@ from __future__ import annotations
 from collections.abc import Mapping
 from dataclasses import dataclass
 from itertools import count
-from typing import Any
+from typing import Any, TypeAlias
 
 import numpy as np
 from omegaconf import OmegaConf
@@ -20,7 +20,7 @@ from vllm.logger import init_logger
 
 logger = init_logger(__name__)
 
-ActionOutput = np.ndarray | dict[str, np.ndarray]
+ActionOutput: TypeAlias = np.ndarray | dict[str, np.ndarray]
 
 
 def _to_builtin_container(value: Any) -> Any:
@@ -150,6 +150,9 @@ class ServingRealtimeRobotOpenPI:
         from vllm_omni.diffusion.request import OmniDiffusionRequest
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
+        # An optional integer ``seed`` in the inference message becomes the engine request's
+        # ``sampling_params.seed``; omitted, ``OmniDiffusionRequest`` assigns a random one.
+        seed = obs.pop("seed", None)
         extra_args = {
             "reset": reset,
             "session_id": session_id,
@@ -157,7 +160,10 @@ class ServingRealtimeRobotOpenPI:
         }
 
         prompt = obs.get("prompt", "")
-        sampling_params = OmniDiffusionSamplingParams(extra_args=extra_args)
+        sampling_params = OmniDiffusionSamplingParams(
+            seed=int(seed) if seed is not None else None,
+            extra_args=extra_args,
+        )
         return OmniDiffusionRequest(
             prompt=prompt,
             sampling_params=sampling_params,