diff --git a/docs/serving/openpi_api.md b/docs/serving/openpi_api.md index 27f092d87..52d94cbbf 100644 --- a/docs/serving/openpi_api.md +++ b/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": , ...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 diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 52ff1e274..7603945a0 100644 --- a/tests/config/test_environment_variables.py +++ b/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, diff --git a/tests/diffusion/models/gr00t/test_noise_seed.py b/tests/diffusion/models/gr00t/test_noise_seed.py new file mode 100644 index 000000000..a865d9a05 --- /dev/null +++ b/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) diff --git a/tests/diffusion/models/gr00t/test_pipeline.py b/tests/diffusion/models/gr00t/test_pipeline.py index 322112107..3310d049e 100644 --- a/tests/diffusion/models/gr00t/test_pipeline.py +++ b/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 diff --git a/tests/e2e/online_serving/test_gr00t_openpi_expansion.py b/tests/e2e/online_serving/test_gr00t_openpi_expansion.py index 70e11088d..2f81ab3fa 100644 --- a/tests/e2e/online_serving/test_gr00t_openpi_expansion.py +++ b/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), }, ], } diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index cdd833db0..072e734e4 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/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) diff --git a/tests/helpers/client.py b/tests/helpers/client.py index 0f8346200..5339f7195 100644 --- a/tests/helpers/client.py +++ b/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 = ( diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 1a1d40141..c1e380a22 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/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", diff --git a/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py index 3fcbce203..6116572ee 100644 --- a/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py +++ b/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) diff --git a/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py b/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py index 1154dcf71..87108522b 100644 --- a/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py +++ b/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: diff --git a/vllm_omni/diffusion/models/gr00t/policy.py b/vllm_omni/diffusion/models/gr00t/policy.py index c16b34463..118f57545 100644 --- a/vllm_omni/diffusion/models/gr00t/policy.py +++ b/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 diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index 4f1d58081..f2448dfa1 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/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,