bdfc77147d
Signed-off-by: Meng <meng_chen99@163.com> Signed-off-by: Yangshen Deng <yangshen.d@outlook.com> Signed-off-by: Suli Wang <50828270+wsl2000@users.noreply.github.com> Co-authored-by: Meng <meng_chen99@163.com> Co-authored-by: Suli Wang <50828270+wsl2000@users.noreply.github.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
382 lines
13 KiB
Python
382 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Copied from the DreamZero repository's `test_client_AR.py`:
|
|
https://github.com/dreamzero0/dreamzero/blob/main/test_client_AR.py
|
|
|
|
Kept here for end-to-end compatibility / parity testing against the vLLM
|
|
OpenPI server.
|
|
|
|
Test client for AR_droid policy server using roboarena interface.
|
|
|
|
Sends real video frames from debug_image/ directory instead of zero dummy images.
|
|
|
|
Frame schedule (matching debug_inference.py):
|
|
- Step 0 (initial): send frame [0] (1 frame, H W 3)
|
|
- Step 1: send frames [0, 7, 15, 23] (4 frames, 4 H W 3)
|
|
- Step 2: send frames [24, 31, 39, 47] (4 frames)
|
|
- Step 3: send frames [48, 55, 63, 71] (4 frames)
|
|
- ...
|
|
|
|
Expected server configuration:
|
|
- image_resolution: (180, 320)
|
|
- n_external_cameras: 2
|
|
- needs_wrist_camera: True
|
|
- action_space: "joint_position"
|
|
|
|
Usage:
|
|
# Start server with roboarena interface:
|
|
torchrun --nproc_per_node=8 socket_test_optimized_AR.py --port 8000
|
|
|
|
# Run against the original DreamZero websocket server (default path):
|
|
python test_client_AR.py --host <server_host> --port 8000
|
|
|
|
# Run against the vLLM OpenPI server:
|
|
python test_client_AR.py --host <server_host> --port 8000 --path /v1/realtime/robot/openpi
|
|
|
|
# Use zero images instead of real video (old behavior):
|
|
python test_client_AR.py --host <server_host> --port 8000 --use-zero-images
|
|
"""
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from openpi_client import msgpack_numpy
|
|
|
|
_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
|
|
DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None
|
|
|
|
|
|
def _import_upstream_policy_modules():
|
|
if DREAMZERO_REPO is None:
|
|
raise ImportError("Set DREAMZERO_REPO to an upstream DreamZero checkout before using this helper.")
|
|
if DREAMZERO_REPO.exists() and str(DREAMZERO_REPO) not in sys.path:
|
|
sys.path.insert(0, str(DREAMZERO_REPO))
|
|
|
|
import eval_utils.policy_server as policy_server
|
|
from eval_utils.policy_client import WebsocketClientPolicy
|
|
|
|
return policy_server, WebsocketClientPolicy
|
|
|
|
|
|
policy_server, WebsocketClientPolicy = _import_upstream_policy_modules()
|
|
|
|
VIDEO_DIR = os.environ.get(
|
|
"DREAMZERO_VIDEO_DIR",
|
|
str(DREAMZERO_REPO / "debug_image") if DREAMZERO_REPO is not None else "debug_image",
|
|
)
|
|
|
|
# roboarena key -> video filename
|
|
CAMERA_FILES = {
|
|
"observation/exterior_image_0_left": "exterior_image_1_left.mp4",
|
|
"observation/exterior_image_1_left": "exterior_image_2_left.mp4",
|
|
"observation/wrist_image_left": "wrist_image_left.mp4",
|
|
}
|
|
|
|
# Frame schedule constants (matching debug_inference.py)
|
|
RELATIVE_OFFSETS = [-23, -16, -8, 0]
|
|
ACTION_HORIZON = 24
|
|
DEFAULT_WEBSOCKET_PATH = ""
|
|
|
|
|
|
class OpenPIWebsocketClientPolicy(WebsocketClientPolicy):
|
|
"""DreamZero websocket client with a configurable path suffix.
|
|
|
|
The original DreamZero client connects to ``ws://host:port``.
|
|
vLLM serves the compatible robot policy endpoint at
|
|
``/v1/realtime/robot/openpi`` when ``path`` is set accordingly.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str = "0.0.0.0",
|
|
port: int = 8000,
|
|
path: str = DEFAULT_WEBSOCKET_PATH,
|
|
) -> None:
|
|
self._uri = f"ws://{host}:{port}{path}"
|
|
self._packer = msgpack_numpy.Packer()
|
|
self._ws, self._server_metadata = self._wait_for_server()
|
|
|
|
|
|
def load_all_frames(video_path: str) -> np.ndarray:
|
|
"""Load all frames from a video file. Returns (N, H, W, 3) uint8 array (RGB)."""
|
|
cap = cv2.VideoCapture(video_path)
|
|
frames = []
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
|
cap.release()
|
|
if not frames:
|
|
raise RuntimeError(f"No frames loaded from {video_path}")
|
|
return np.stack(frames, axis=0)
|
|
|
|
|
|
def load_camera_frames() -> dict[str, np.ndarray]:
|
|
"""Load all video frames for each camera from the debug_image/ directory.
|
|
|
|
Returns:
|
|
Dict mapping roboarena camera keys to (N, H, W, 3) uint8 arrays.
|
|
"""
|
|
camera_frames: dict[str, np.ndarray] = {}
|
|
for cam_key, fname in CAMERA_FILES.items():
|
|
path = os.path.join(VIDEO_DIR, fname)
|
|
camera_frames[cam_key] = load_all_frames(path)
|
|
logging.info(f"Loaded {cam_key}: {camera_frames[cam_key].shape}")
|
|
return camera_frames
|
|
|
|
|
|
def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
|
|
"""Build the frame index schedule for multi-frame chunks.
|
|
|
|
Returns a list of frame-index lists. Each inner list has 4 indices.
|
|
"""
|
|
chunks: list[list[int]] = []
|
|
current_frame = 23 # first anchor frame
|
|
for _ in range(num_chunks):
|
|
indices = [max(current_frame + off, 0) for off in RELATIVE_OFFSETS]
|
|
if indices[-1] >= total_frames:
|
|
logging.info(f"Frame {indices[-1]} >= {total_frames}, stopping at {len(chunks)} chunks")
|
|
break
|
|
chunks.append(indices)
|
|
current_frame += ACTION_HORIZON
|
|
return chunks
|
|
|
|
|
|
def _make_obs_from_video(
|
|
camera_frames: dict[str, np.ndarray],
|
|
frame_indices: list[int],
|
|
prompt: str,
|
|
session_id: str,
|
|
) -> dict:
|
|
"""Build an observation dict from real video frames.
|
|
|
|
For 1 frame: each image key is (H, W, 3).
|
|
For 4 frames: each image key is (4, H, W, 3).
|
|
"""
|
|
obs: dict = {}
|
|
for cam_key, all_frames in camera_frames.items():
|
|
selected = all_frames[frame_indices] # (T, H, W, 3)
|
|
if len(frame_indices) == 1:
|
|
selected = selected[0] # (H, W, 3)
|
|
obs[cam_key] = selected
|
|
|
|
obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
|
|
obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
|
|
obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
|
|
obs["prompt"] = prompt
|
|
obs["session_id"] = session_id
|
|
return obs
|
|
|
|
|
|
def _make_zero_observation(
|
|
server_config: policy_server.PolicyServerConfig,
|
|
prompt: str = "pick up the object",
|
|
session_id: str | None = None,
|
|
) -> dict:
|
|
"""Create a dummy observation matching AR_droid expectations.
|
|
|
|
AR_droid expects:
|
|
- 2 external cameras (exterior_image_0_left, exterior_image_1_left)
|
|
- 1 wrist camera (wrist_image_left)
|
|
- Image resolution: 180x320 (H x W)
|
|
- joint_position: 7 DoF
|
|
- gripper_position: 1 DoF
|
|
"""
|
|
obs = {}
|
|
|
|
# Determine image resolution
|
|
if server_config.image_resolution is not None:
|
|
h, w = server_config.image_resolution
|
|
else:
|
|
# Default for AR_droid
|
|
h, w = 180, 320
|
|
|
|
# External cameras (0-indexed in roboarena)
|
|
for i in range(server_config.n_external_cameras):
|
|
obs[f"observation/exterior_image_{i}_left"] = np.zeros((h, w, 3), dtype=np.uint8)
|
|
if server_config.needs_stereo_camera:
|
|
obs[f"observation/exterior_image_{i}_right"] = np.zeros((h, w, 3), dtype=np.uint8)
|
|
|
|
# Wrist camera
|
|
if server_config.needs_wrist_camera:
|
|
obs["observation/wrist_image_left"] = np.zeros((h, w, 3), dtype=np.uint8)
|
|
if server_config.needs_stereo_camera:
|
|
obs["observation/wrist_image_right"] = np.zeros((h, w, 3), dtype=np.uint8)
|
|
|
|
# Session ID - should be passed in to ensure consistency within a session
|
|
if server_config.needs_session_id:
|
|
import uuid
|
|
|
|
# Generate unique session ID if not provided
|
|
obs["session_id"] = session_id if session_id else str(uuid.uuid4())
|
|
|
|
# State observations (AR_droid: 7 DoF arm + 1 gripper)
|
|
obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
|
|
obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
|
|
obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
|
|
|
|
# Language prompt
|
|
obs["prompt"] = prompt
|
|
|
|
return obs
|
|
|
|
|
|
def test_ar_droid_policy_server(
|
|
host: str = "localhost",
|
|
port: int = 8000,
|
|
path: str = DEFAULT_WEBSOCKET_PATH,
|
|
num_chunks: int = 15,
|
|
prompt: str = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
|
|
use_zero_images: bool = False,
|
|
):
|
|
"""Test the AR_droid policy server with roboarena interface.
|
|
|
|
When use_zero_images is False (default), loads real video frames from
|
|
debug_image/ and follows the frame schedule from debug_inference.py.
|
|
"""
|
|
logging.info(f"Connecting to AR_droid server at ws://{host}:{port}{path} ...")
|
|
|
|
client = OpenPIWebsocketClientPolicy(host=host, port=port, path=path)
|
|
|
|
# Validate server metadata
|
|
metadata = client.get_server_metadata()
|
|
logging.info(f"Server metadata: {metadata}")
|
|
assert isinstance(metadata, dict), "Metadata should be a dict"
|
|
|
|
try:
|
|
server_config = policy_server.PolicyServerConfig(**metadata)
|
|
except Exception as e:
|
|
logging.error(f"Error parsing metadata: {e}")
|
|
raise e
|
|
|
|
# Validate expected AR_droid configuration
|
|
logging.info(f"Server config: {server_config}")
|
|
assert server_config.n_external_cameras == 2, f"Expected 2 external cameras, got {server_config.n_external_cameras}"
|
|
assert server_config.needs_wrist_camera, "Expected wrist camera to be enabled"
|
|
assert server_config.action_space == "joint_position", (
|
|
f"Expected joint_position action space, got {server_config.action_space}"
|
|
)
|
|
|
|
logging.info("Server configuration validated for AR_droid")
|
|
|
|
# Generate unique session ID for this test run
|
|
session_id = str(uuid.uuid4())
|
|
logging.info(f"Session ID: {session_id}")
|
|
|
|
# ── Zero-image fallback mode ──────────────────────────────────────
|
|
if use_zero_images:
|
|
logging.info("Using ZERO dummy images (legacy mode)")
|
|
for i in range(num_chunks):
|
|
obs = _make_zero_observation(server_config, prompt=prompt, session_id=session_id)
|
|
logging.info(f"Inference {i + 1}/{num_chunks}: prompt='{prompt}'")
|
|
t0 = time.time()
|
|
actions = client.infer(obs)
|
|
dt = time.time() - t0
|
|
_log_action(actions, dt)
|
|
|
|
logging.info("Sending reset...")
|
|
client.reset({})
|
|
logging.info("Done (zero-image mode).")
|
|
return
|
|
|
|
# ── Real video frame mode ─────────────────────────────────────────
|
|
logging.info("Loading real video frames from debug_image/ directory")
|
|
camera_frames = load_camera_frames()
|
|
|
|
total_frames = min(v.shape[0] for v in camera_frames.values())
|
|
logging.info(f"Total frames available: {total_frames}")
|
|
|
|
# Build frame schedule
|
|
chunks = build_frame_schedule(total_frames, num_chunks)
|
|
|
|
logging.info("Frame schedule:")
|
|
logging.info(" Initial: [0]")
|
|
for i, indices in enumerate(chunks):
|
|
logging.info(f" Chunk {i}: {indices}")
|
|
|
|
# Step 0: initial single frame
|
|
logging.info("=== Initial: frame [0] ===")
|
|
obs = _make_obs_from_video(camera_frames, [0], prompt, session_id)
|
|
t0 = time.time()
|
|
actions = client.infer(obs)
|
|
dt = time.time() - t0
|
|
_log_action(actions, dt)
|
|
|
|
# Subsequent chunks: send 4 frames at a time
|
|
for chunk_idx, frame_indices in enumerate(chunks):
|
|
logging.info(f"=== Chunk {chunk_idx}: frames {frame_indices} ===")
|
|
obs = _make_obs_from_video(camera_frames, frame_indices, prompt, session_id)
|
|
t0 = time.time()
|
|
actions = client.infer(obs)
|
|
dt = time.time() - t0
|
|
_log_action(actions, dt)
|
|
|
|
# Reset triggers video save on the server
|
|
logging.info("Sending reset to save video...")
|
|
client.reset({})
|
|
|
|
logging.info("Done.")
|
|
|
|
|
|
def _log_action(actions: np.ndarray, dt: float) -> None:
|
|
"""Pretty-print action shape, range, and timing."""
|
|
assert isinstance(actions, np.ndarray), f"Expected numpy array, got {type(actions)}"
|
|
assert actions.ndim == 2, f"Expected 2D array, got shape {actions.shape}"
|
|
assert actions.shape[-1] == 8, f"Expected 8 action dims (7 joints + 1 gripper), got {actions.shape[-1]}"
|
|
logging.info(f" Action shape: {actions.shape}, range: [{actions.min():.4f}, {actions.max():.4f}], time: {dt:.2f}s")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Test AR_droid policy server with real video frames from debug_image/")
|
|
parser.add_argument("--host", default="localhost", help="Server hostname")
|
|
parser.add_argument("--port", type=int, default=8000, help="Server port")
|
|
parser.add_argument(
|
|
"--path",
|
|
default=DEFAULT_WEBSOCKET_PATH,
|
|
help="WebSocket path suffix (default: empty string for the original DreamZero server)",
|
|
)
|
|
parser.add_argument(
|
|
"--num-chunks",
|
|
type=int,
|
|
default=15,
|
|
help="Number of 4-frame chunks to send after the initial frame (default: 15)",
|
|
)
|
|
parser.add_argument(
|
|
"--prompt",
|
|
default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
|
|
help="Language prompt for the policy",
|
|
)
|
|
parser.add_argument(
|
|
"--use-zero-images",
|
|
action="store_true",
|
|
help="Use zero dummy images instead of real video frames (legacy mode)",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s [%(levelname)s] %(message)s",
|
|
)
|
|
|
|
test_ar_droid_policy_server(
|
|
host=args.host,
|
|
port=args.port,
|
|
path=args.path,
|
|
num_chunks=args.num_chunks,
|
|
prompt=args.prompt,
|
|
use_zero_images=args.use_zero_images,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|