Yichen WangandDidan Deng
2a9f6aaa8e
Signed-off-by: Yichen Wang <yichen.wang@bytedance.com> Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
120 lines
4.2 KiB
Python
120 lines
4.2 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Minimal OpenPI client for the π0 (Pi-Zero) VLA policy server.
|
|
|
|
Connects to a running π0 server (see README.md), sends a single robot
|
|
observation (3 cameras + state + language instruction), and prints the returned
|
|
continuous action chunk ``[action_horizon, action_dim]``.
|
|
|
|
The server speaks the OpenPI websocket protocol:
|
|
connect -> server sends msgpack(policy_server_config metadata)
|
|
infer -> client sends msgpack(obs), server replies msgpack(actions ndarray)
|
|
|
|
Usage:
|
|
python openpi_client.py --host 127.0.0.1 --port 8000 \
|
|
--prompt "pick up the red block"
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import logging
|
|
import sys
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
try:
|
|
# Avoid shadowing the installed ``openpi_client`` package with this dir.
|
|
example_dir = str(Path(__file__).resolve().parent)
|
|
removed = sys.path and sys.path[0] == example_dir
|
|
if removed:
|
|
sys.path.pop(0)
|
|
try:
|
|
from openpi_client.websocket_client_policy import WebsocketClientPolicy
|
|
finally:
|
|
if removed:
|
|
sys.path.insert(0, example_dir)
|
|
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
raise ImportError("π0 OpenPI example requires `openpi-client`.") from exc
|
|
|
|
logger = logging.getLogger("pi0_openpi_client")
|
|
|
|
DEFAULT_HOST = "127.0.0.1"
|
|
DEFAULT_PORT = 8000
|
|
DEFAULT_PATH = "/v1/realtime/robot/openpi"
|
|
DEFAULT_PROMPT = "pick up the red block and place it in the bin"
|
|
|
|
# π0 / pi0_base camera identities (must match the server's image_feature_keys).
|
|
CAMERA_KEYS = (
|
|
"observation.images.base_0_rgb",
|
|
"observation.images.left_wrist_0_rgb",
|
|
"observation.images.right_wrist_0_rgb",
|
|
)
|
|
STATE_DIM = 32
|
|
|
|
|
|
def _make_dummy_obs(prompt: str, session_id: str, image_size: int = 224) -> dict[str, Any]:
|
|
"""Build a single π0 observation with blank cameras + zero state.
|
|
|
|
Replace the zeros with real camera frames (HWC uint8) and proprioceptive
|
|
state to drive an actual robot.
|
|
"""
|
|
obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in CAMERA_KEYS}
|
|
obs["state"] = np.zeros(STATE_DIM, dtype=np.float32)
|
|
obs["prompt"] = prompt
|
|
obs["session_id"] = session_id
|
|
return obs
|
|
|
|
|
|
def _ws_uri(host: str, port: int, path: str) -> str:
|
|
"""Build the full websocket URI. WebsocketClientPolicy takes the whole URI as
|
|
``host`` (it has no separate path arg)."""
|
|
if host.startswith(("ws://", "wss://")):
|
|
return host
|
|
return f"ws://{host}:{port}{path}"
|
|
|
|
|
|
def _as_action_array(response: Any) -> np.ndarray:
|
|
if isinstance(response, dict) and response.get("type") == "error":
|
|
raise RuntimeError(f"Inference failed: {response.get('message', response)}")
|
|
return np.asarray(response, dtype=np.float32)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="π0 OpenPI client")
|
|
parser.add_argument("--host", default=DEFAULT_HOST)
|
|
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
|
|
parser.add_argument("--path", default=DEFAULT_PATH)
|
|
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
|
|
parser.add_argument("--num-steps", type=int, default=2, help="Inference observations to send.")
|
|
args = parser.parse_args()
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
|
|
policy = WebsocketClientPolicy(host=_ws_uri(args.host, args.port, args.path))
|
|
metadata = dict(policy.get_server_metadata())
|
|
logger.info("Server metadata: %s", metadata)
|
|
|
|
session_id = str(uuid.uuid4())
|
|
for step in range(args.num_steps):
|
|
obs = _make_dummy_obs(args.prompt, session_id)
|
|
actions = _as_action_array(policy.infer(obs))
|
|
logger.info(
|
|
"[step %d] actions shape=%s mean=%.4f std=%.4f",
|
|
step,
|
|
actions.shape,
|
|
float(actions.mean()),
|
|
float(actions.std()),
|
|
)
|
|
if not np.isfinite(actions).all():
|
|
raise RuntimeError("Server returned non-finite actions.")
|
|
logger.info("Done — π0 OpenPI round-trip OK.")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|