[Diffusion] add GR00T-N1.7 pipeline with OpenPI serving (#3798)

Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

tests/diffusion/models/gr00t/test_pipeline.py

@@ -0,0 +1,141 @@
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from vllm_omni.diffusion.models.gr00t import pipeline_gr00t
+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]
+
+
+class FakeGr00tPolicy:
+    instances = []
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+        self.reset_calls = 0
+        self.seen_obs = None
+        self.embodiment_tag = SimpleNamespace(value="fake_embodiment")
+        self.language_key = "annotation.language.language_instruction"
+        self.modality_configs = {
+            "action": SimpleNamespace(
+                delta_indices=[0, 1],
+                modality_keys=["arm", "gripper"],
+            )
+        }
+        self.processor = SimpleNamespace(
+            state_action_processor=SimpleNamespace(
+                norm_params={
+                    "fake_embodiment": {
+                        "action": {
+                            "arm": {"dim": np.array(2)},
+                            "gripper": {"dim": np.array(1)},
+                        }
+                    }
+                }
+            )
+        )
+        FakeGr00tPolicy.instances.append(self)
+
+    def get_action(self, obs):
+        self.seen_obs = obs
+        return {
+            "arm": np.array([[[1.0, 2.0]]], dtype=np.float64),
+            "gripper": [[[3.0]]],
+        }, {"latency_ms": 1.0}
+
+    def reset(self):
+        self.reset_calls += 1
+        return {"reset": True}
+
+
+@pytest.fixture(autouse=True)
+def fake_gr00t_policy(monkeypatch):
+    FakeGr00tPolicy.instances.clear()
+    monkeypatch.setattr(pipeline_gr00t, "Gr00tPolicy", FakeGr00tPolicy)
+
+
+def _pipeline():
+    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 test_pipeline_initializes_local_policy():
+    pipeline = _pipeline()
+
+    policy = FakeGr00tPolicy.instances[0]
+    assert policy.kwargs["model_path"] == "nvidia/GR00T-N1.7-3B"
+    assert policy.kwargs["embodiment_tag"] == "LIBERO_PANDA"
+    assert policy.kwargs["strict"] is False
+    assert pipeline.weights_sources == ()
+    assert pipeline.load_weights(iter(())) == set()
+
+
+def test_forward_returns_dict_actions_in_output():
+    pipeline = _pipeline()
+    req = OmniDiffusionRequest(
+        prompts=["pick"],
+        request_id="req",
+        sampling_params=OmniDiffusionSamplingParams(
+            extra_args={
+                "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",
+                    "session_id": "session-a",
+                },
+                "reset": True,
+            }
+        ),
+    )
+
+    output = pipeline.forward(req)
+
+    assert output.error is None
+    actions = output.output["actions"]
+    assert set(actions) == {"arm", "gripper"}
+    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 "video" in policy.seen_obs
+    assert policy.seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]}
+    assert "images" not in policy.seen_obs
+    assert "prompt" not in policy.seen_obs
+    assert "session_id" not in policy.seen_obs
+    assert policy.reset_calls == 1
+
+
+def test_dummy_warmup_returns_shape_correct_zero_actions():
+    pipeline = _pipeline()
+    req = OmniDiffusionRequest(
+        prompts=["dummy run"],
+        request_id="dummy_req_id",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+    )
+
+    output = pipeline.forward(req)
+
+    assert output.error is None
+    actions = output.output["actions"]
+    assert set(actions) == {"arm", "gripper"}
+    assert actions["arm"].shape == (1, 2, 2)
+    assert actions["gripper"].shape == (1, 2, 1)
+    assert not actions["arm"].any()
+    assert FakeGr00tPolicy.instances[0].seen_obs is None
+
+
+def test_reset_delegates_to_policy():
+    pipeline = _pipeline()
+
+    assert pipeline.reset() == {"reset": True}
+    assert FakeGr00tPolicy.instances[0].reset_calls == 1

tests/e2e/online_serving/test_gr00t_openpi.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end online serving test for GR00T N1.7 through the OpenPI robot endpoint."""
+
+import os
+
+import numpy as np
+import pytest
+
+from tests.gr00t import openpi_client_helper as openpi_client
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = "nvidia/GR00T-N1.7-3B"
+
+pytest.importorskip("websockets")
+pytest.importorskip("openpi_client.msgpack_numpy")
+
+test_params = [
+    pytest.param(
+        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"},
+            init_timeout=1200,
+            stage_init_timeout=900,
+        ),
+        id="gr00t-n1d7-openpi",
+    )
+]
+
+# Reference values captured from the Isaac-GR00T ZMQ server (nvidia/GR00T-N1.7-3B,
+# OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT embodiment) using build_droid_observation()
+# inputs: zero images (256×256), identity eef_9d, zero gripper/joint, "pick up the object".
+# Outputs are bit-reproducible across resets and across runs (max_diff=0.0).
+_REF_EEF_9D = np.array(
+    [
+        [
+            0.014888550154864788,
+            -0.00039258378092199564,
+            -0.013574761338531971,
+            0.9999837875366211,
+            0.005678884219378233,
+            -0.00034708858584053814,
+            -0.005677700508385897,
+            0.9999783635139465,
+            0.0033210734836757183,
+        ],
+        [
+            0.020188070833683014,
+            -0.0003105341165792197,
+            -0.02232760563492775,
+            0.9997346997261047,
+            0.02301345206797123,
+            0.000983047066256404,
+            -0.02302766777575016,
+            0.9995648860931396,
+            0.018431292846798897,
+        ],
+        [
+            -0.007266733795404434,
+            -0.05537768080830574,
+            0.03667901083827019,
+            0.992686927318573,
+            0.1206541359424591,
+            0.003906540106981993,
+            -0.12024091929197311,
+            0.9853788614273071,
+            0.12070894986391068,
+        ],
+    ],
+    dtype=np.float32,
+)  # rows = step 0, step 4, step 39
+
+_REF_GRIPPER = np.array([[0.0], [0.0078125], [0.939453125]], dtype=np.float32)  # steps 0, 4, 39
+
+_REF_JOINT = np.array(
+    [
+        [
+            -0.0010484338272362947,
+            0.0014262489276006818,
+            -0.003565810853615403,
+            -3.846167237497866e-05,
+            -0.0002604846959002316,
+            0.008521700277924538,
+            -0.006872728932648897,
+        ],
+        [
+            -0.009435676969587803,
+            0.0021475711837410927,
+            -0.0031688229646533728,
+            -1.8328893929719925e-05,
+            0.0005945992306806147,
+            0.019159257411956787,
+            0.0009468861389905214,
+        ],
+        [
+            -0.02944088727235794,
+            -0.08419207483530045,
+            -0.0251418836414814,
+            0.00540524534881115,
+            0.04752273112535477,
+            -0.012884500436484814,
+            0.024298785254359245,
+        ],
+    ],
+    dtype=np.float32,
+)  # rows = step 0, step 4, step 39
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_gr00t_n1d7_openpi_online(omni_server) -> None:
+    result = openpi_client.run_policy_session(
+        host=omni_server.host,
+        port=omni_server.port,
+        session_id="gr00t-online-e2e",
+    )
+    openpi_client.validate_session_result(result)
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@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) -> None:
+    """Assert actions match Isaac-GR00T reference (GR00T_NOISE_SEED=42, zero inputs)."""
+    client = openpi_client.OpenPIWebsocketClient(host=omni_server.host, port=omni_server.port)
+    try:
+        obs = openpi_client.build_droid_observation(session_id="gr00t-precision-e2e")
+        client.reset({})
+        actions = client.infer(obs)
+    finally:
+        client.close()
+
+    steps = [0, 4, 39]
+    np.testing.assert_allclose(
+        actions["eef_9d"][0, steps, :],
+        _REF_EEF_9D,
+        atol=1e-2,
+        rtol=0.0,
+        err_msg="eef_9d action mismatch vs Isaac-GR00T reference",
+    )
+    np.testing.assert_allclose(
+        actions["gripper_position"][0, steps, :],
+        _REF_GRIPPER,
+        atol=1e-2,
+        rtol=0.0,
+        err_msg="gripper_position action mismatch vs Isaac-GR00T reference",
+    )
+    np.testing.assert_allclose(
+        actions["joint_position"][0, steps, :],
+        _REF_JOINT,
+        atol=1e-2,
+        rtol=0.0,
+        err_msg="joint_position action mismatch vs Isaac-GR00T reference",
+    )

tests/gr00t/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/gr00t/openpi_client_helper.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+try:
+    import websockets.sync.client as websockets_client
+except ImportError:  # pragma: no cover - optional e2e dependency
+    websockets_client = None
+
+try:
+    from openpi_client import msgpack_numpy
+except ImportError:  # pragma: no cover - optional e2e dependency
+    msgpack_numpy = None
+
+PING_INTERVAL_SECS = 300
+PING_TIMEOUT_SECS = 3600
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 8000
+DEFAULT_PATH = "/v1/realtime/robot/openpi"
+DEFAULT_SESSION_ID = "gr00t-smoke"
+ACTION_KEYS = {"eef_9d", "gripper_position", "joint_position"}
+LANGUAGE_KEY = "annotation.language.language_instruction"
+
+
+def _identity_eef_9d_state() -> np.ndarray:
+    state = np.zeros((1, 1, 9), dtype=np.float32)
+    state[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32)
+    return state
+
+
+def require_dependencies() -> None:
+    missing = []
+    if websockets_client is None:
+        missing.append("websockets")
+    if msgpack_numpy is None:
+        missing.append("openpi-client")
+    if missing:
+        raise ModuleNotFoundError(f"GR00T OpenPI test dependencies are missing: {', '.join(missing)}")
+
+
+@dataclass(frozen=True)
+class Gr00tServerMetadata:
+    action_horizon: int
+    action_keys: set[str]
+    embodiment_tag: str
+    needs_session_id: bool
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]):
+        required_keys = ("action_horizon", "action_keys", "embodiment_tag", "needs_session_id")
+        missing_keys = [key for key in required_keys if key not in payload]
+        if missing_keys:
+            raise ValueError(f"Missing GR00T metadata keys: {missing_keys}")
+
+        return cls(
+            action_horizon=int(payload["action_horizon"]),
+            action_keys={str(key) for key in payload["action_keys"]},
+            embodiment_tag=str(payload["embodiment_tag"]),
+            needs_session_id=bool(payload["needs_session_id"]),
+        )
+
+
+class OpenPIWebsocketClient:
+    def __init__(
+        self,
+        *,
+        host: str = DEFAULT_HOST,
+        port: int = DEFAULT_PORT,
+        path: str = DEFAULT_PATH,
+    ) -> None:
+        require_dependencies()
+        self._uri = f"ws://{host}:{port}{path}"
+        self._packer = msgpack_numpy.Packer()
+        self._ws, self._server_metadata = self._connect()
+
+    def _connect(self):
+        conn = websockets_client.connect(
+            self._uri,
+            compression=None,
+            max_size=None,
+            ping_interval=PING_INTERVAL_SECS,
+            ping_timeout=PING_TIMEOUT_SECS,
+        )
+        metadata = msgpack_numpy.unpackb(conn.recv())
+        if not isinstance(metadata, dict):
+            raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
+        return conn, metadata
+
+    def get_server_metadata(self) -> dict[str, Any]:
+        return dict(self._server_metadata)
+
+    def infer(self, obs: dict[str, Any]) -> dict[str, np.ndarray]:
+        payload = dict(obs)
+        payload["endpoint"] = "infer"
+        self._ws.send(self._packer.pack(payload))
+        response = msgpack_numpy.unpackb(self._ws.recv())
+        if isinstance(response, dict) and response.get("type") == "error":
+            raise RuntimeError(f"Inference failed: {response['message']}")
+        if not isinstance(response, dict):
+            raise RuntimeError(f"Expected dict actions from GR00T OpenPI endpoint, got {type(response)!r}")
+        return {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()}
+
+    def reset(self, reset_info: dict[str, Any] | None = None) -> str:
+        payload = dict(reset_info or {})
+        payload["endpoint"] = "reset"
+        self._ws.send(self._packer.pack(payload))
+        response = msgpack_numpy.unpackb(self._ws.recv())
+        if not isinstance(response, dict) or response.get("status") != "reset successful":
+            raise RuntimeError(f"Unexpected reset response: {response!r}")
+        return str(response["status"])
+
+    def close(self) -> None:
+        self._ws.close()
+
+
+def build_droid_observation(*, session_id: str = DEFAULT_SESSION_ID) -> dict[str, Any]:
+    return {
+        "session_id": session_id,
+        "video": {
+            "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
+            "wrist_image_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
+        },
+        "state": {
+            "eef_9d": _identity_eef_9d_state(),
+            "gripper_position": np.zeros((1, 1, 1), dtype=np.float32),
+            "joint_position": np.zeros((1, 1, 7), dtype=np.float32),
+        },
+        "language": {LANGUAGE_KEY: [["pick up the object"]]},
+    }
+
+
+def run_policy_session(
+    *,
+    host: str = DEFAULT_HOST,
+    port: int = DEFAULT_PORT,
+    path: str = DEFAULT_PATH,
+    session_id: str = DEFAULT_SESSION_ID,
+) -> dict[str, Any]:
+    client = OpenPIWebsocketClient(host=host, port=port, path=path)
+    try:
+        metadata = client.get_server_metadata()
+        actions = client.infer(build_droid_observation(session_id=session_id))
+        reset_status = client.reset({})
+        return {
+            "metadata": metadata,
+            "actions": actions,
+            "reset_status": reset_status,
+            "session_id": session_id,
+        }
+    finally:
+        client.close()
+
+
+def validate_session_result(result: dict[str, Any]) -> None:
+    metadata = Gr00tServerMetadata.from_dict(result["metadata"])
+    if metadata.embodiment_tag != "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT":
+        raise AssertionError(f"Unexpected embodiment_tag: {metadata.embodiment_tag}")
+    if not metadata.needs_session_id:
+        raise AssertionError("GR00T test expects needs_session_id metadata")
+    if metadata.action_keys != ACTION_KEYS:
+        raise AssertionError(f"Unexpected action keys: {metadata.action_keys}")
+    if result["reset_status"] != "reset successful":
+        raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}")
+
+    actions = result["actions"]
+    if set(actions) != ACTION_KEYS:
+        raise AssertionError(f"Unexpected action keys: {set(actions)}")
+    expected_shapes = {
+        "eef_9d": (1, metadata.action_horizon, 9),
+        "gripper_position": (1, metadata.action_horizon, 1),
+        "joint_position": (1, metadata.action_horizon, 7),
+    }
+    for key, expected_shape in expected_shapes.items():
+        if actions[key].shape != expected_shape:
+            raise AssertionError(f"Action {key} shape mismatch: expected {expected_shape}, got {actions[key].shape}")
+        if actions[key].dtype != np.float32:
+            raise AssertionError(f"Action {key} dtype mismatch: expected float32, got {actions[key].dtype}")
+        if not np.isfinite(actions[key]).all():
+            raise AssertionError(f"Action {key} contains non-finite values")