sgl-project--sglang/python/sglang/multimodal_gen/benchmarks/bench_pi05_openpi.py

# SPDX-License-Identifier: Apache-2.0

"""Manual Pi0.5 SGLang vs OpenPI benchmark.

This script is intentionally outside the unit-test path. It needs GPU memory,
Pi0.5 checkpoints, and an OpenPI install for the baseline.
"""

from __future__ import annotations

import argparse
import asyncio
import dataclasses
import json
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import requests

SCRIPT_DIR = Path(__file__).resolve().parent
if sys.path and Path(sys.path[0]).resolve() == SCRIPT_DIR:
    sys.path.pop(0)

from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (  # noqa: E402
    pack_msgpack,
    unpack_msgpack,
)


@dataclass(frozen=True)
class Pi05BenchProfile:
    name: str
    sglang_model: str
    openpi_config: str
    openpi_checkpoint: str
    prompt: str
    sglang_action_horizon: int
    openpi_action_horizon: int
    action_dim: int
    output_action_dim: int


PROFILES = {
    "libero": Pi05BenchProfile(
        name="libero",
        sglang_model="lerobot/pi05_libero_base",
        openpi_config="pi05_libero",
        openpi_checkpoint="gs://openpi-assets/checkpoints/pi05_libero",
        prompt="pick up the object",
        sglang_action_horizon=50,
        openpi_action_horizon=10,
        action_dim=32,
        output_action_dim=7,
    ),
    "aloha": Pi05BenchProfile(
        name="aloha",
        sglang_model="lerobot/pi05_base",
        openpi_config="pi05_aloha",
        openpi_checkpoint="gs://openpi-assets/checkpoints/pi05_base",
        prompt="pick up the block",
        sglang_action_horizon=50,
        openpi_action_horizon=50,
        action_dim=32,
        output_action_dim=14,
    ),
}


def _stats_ms(samples: list[float]) -> dict[str, float]:
    if not samples:
        return {}
    values = np.asarray(samples, dtype=np.float64)
    return {
        "count": float(len(samples)),
        "mean_ms": float(np.mean(values)),
        "std_ms": float(np.std(values)),
        "p50_ms": float(np.quantile(values, 0.50)),
        "p90_ms": float(np.quantile(values, 0.90)),
        "p95_ms": float(np.quantile(values, 0.95)),
        "min_ms": float(np.min(values)),
        "max_ms": float(np.max(values)),
    }


def _inc(mapping: dict[str, int], key: object, value: int) -> None:
    key_str = str(key)
    mapping[key_str] = mapping.get(key_str, 0) + value


def _summarize_torch_module(module) -> dict[str, Any]:
    param_dtypes: dict[str, int] = {}
    param_dtype_examples: dict[str, list[str]] = {}
    buffer_dtypes: dict[str, int] = {}
    buffer_dtype_examples: dict[str, list[str]] = {}
    devices: dict[str, int] = {}
    param_count = 0
    buffer_count = 0
    trainable_param_count = 0
    for name, param in module.named_parameters(recurse=True):
        numel = int(param.numel())
        param_count += numel
        if param.requires_grad:
            trainable_param_count += numel
        _inc(param_dtypes, param.dtype, numel)
        _inc(devices, param.device, numel)
        examples = param_dtype_examples.setdefault(str(param.dtype), [])
        if len(examples) < 16:
            examples.append(name)
    for name, buffer in module.named_buffers(recurse=True):
        numel = int(buffer.numel())
        buffer_count += numel
        _inc(buffer_dtypes, buffer.dtype, numel)
        _inc(devices, buffer.device, numel)
        examples = buffer_dtype_examples.setdefault(str(buffer.dtype), [])
        if len(examples) < 16:
            examples.append(name)
    return {
        "class": module.__class__.__name__,
        "param_dtypes": param_dtypes,
        "param_dtype_examples": param_dtype_examples,
        "buffer_dtypes": buffer_dtypes,
        "buffer_dtype_examples": buffer_dtype_examples,
        "devices": devices,
        "param_count": param_count,
        "trainable_param_count": trainable_param_count,
        "buffer_count": buffer_count,
    }


def _torch_autocast_dtype(torch_module, device: str) -> str:
    get_autocast_dtype = getattr(torch_module, "get_autocast_dtype", None)
    if get_autocast_dtype is not None:
        return str(get_autocast_dtype(device))
    if device == "cuda":
        return str(torch_module.get_autocast_gpu_dtype())
    return str(torch_module.get_autocast_cpu_dtype())


def openpi_precision_metadata(policy) -> dict[str, Any]:
    import torch

    metadata: dict[str, Any] = {
        "policy_class": policy.__class__.__name__,
        "is_pytorch_model": bool(getattr(policy, "_is_pytorch_model", False)),
        "pytorch_device": str(getattr(policy, "_pytorch_device", "")),
        "torch_default_dtype": str(torch.get_default_dtype()),
        "torch_autocast_cpu_dtype": _torch_autocast_dtype(torch, "cpu"),
    }
    if torch.cuda.is_available():
        metadata["torch_autocast_cuda_dtype"] = _torch_autocast_dtype(torch, "cuda")

    modules = {}
    for name, value in vars(policy).items():
        if isinstance(value, torch.nn.Module):
            modules[name] = _summarize_torch_module(value)
    metadata["torch_modules"] = modules
    return metadata


def sglang_precision_metadata(pipeline) -> dict[str, Any]:
    import torch

    metadata: dict[str, Any] = {
        "torch_default_dtype": str(torch.get_default_dtype()),
        "torch_autocast_cpu_dtype": _torch_autocast_dtype(torch, "cpu"),
    }
    if torch.cuda.is_available():
        metadata["torch_autocast_cuda_dtype"] = _torch_autocast_dtype(torch, "cuda")

    modules = {}
    policy_model = pipeline.get_module("policy_model")
    if isinstance(policy_model, torch.nn.Module):
        modules["policy_model"] = _summarize_torch_module(policy_model)
        core_model = getattr(policy_model, "core_model", None)
        if isinstance(core_model, torch.nn.Module):
            modules["core_model"] = _summarize_torch_module(core_model)
    metadata["torch_modules"] = modules
    return metadata


def _image(rng: np.random.Generator, *, chw: bool = False) -> np.ndarray:
    image = rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8)
    if chw:
        return np.transpose(image, (2, 0, 1))
    return image


def _make_libero_observation(
    rng: np.random.Generator,
    prompt: str,
) -> tuple[dict[str, Any], dict[str, Any]]:
    base_image = _image(rng)
    wrist_image = _image(rng)
    state = rng.random(8, dtype=np.float32)
    openpi_obs = {
        "observation/state": state,
        "observation/image": base_image,
        "observation/wrist_image": wrist_image,
        "prompt": prompt,
    }
    sglang_observation = {
        "images": {
            "image": base_image,
            "image2": wrist_image,
        },
        "state": state,
    }
    return openpi_obs, sglang_observation


def _make_aloha_observation(
    rng: np.random.Generator,
    prompt: str,
) -> tuple[dict[str, Any], dict[str, Any]]:
    cam_high = _image(rng, chw=True)
    cam_left = _image(rng, chw=True)
    cam_right = _image(rng, chw=True)
    state = np.ones((14,), dtype=np.float32)
    openpi_obs = {
        "state": state,
        "images": {
            "cam_high": cam_high,
            "cam_low": _image(rng, chw=True),
            "cam_left_wrist": cam_left,
            "cam_right_wrist": cam_right,
        },
        "prompt": prompt,
    }
    sglang_state = np.zeros((32,), dtype=np.float32)
    sglang_state[: state.shape[0]] = state
    sglang_observation = {
        "images": {
            "base_0_rgb": np.transpose(cam_high, (1, 2, 0)),
            "left_wrist_0_rgb": np.transpose(cam_left, (1, 2, 0)),
            "right_wrist_0_rgb": np.transpose(cam_right, (1, 2, 0)),
        },
        "state": sglang_state,
    }
    return openpi_obs, sglang_observation


def build_observations(
    profile: Pi05BenchProfile,
    count: int,
    seed: int,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    rng = np.random.default_rng(seed)
    openpi_observations = []
    sglang_observations = []
    for _ in range(count):
        if profile.name == "libero":
            openpi_obs, sglang_obs = _make_libero_observation(rng, profile.prompt)
        elif profile.name == "aloha":
            openpi_obs, sglang_obs = _make_aloha_observation(rng, profile.prompt)
        else:
            raise ValueError(f"Unsupported Pi0.5 benchmark profile: {profile.name}")
        openpi_observations.append(openpi_obs)
        sglang_observations.append(sglang_obs)
    return openpi_observations, sglang_observations


def _json_tensor(array: np.ndarray) -> dict[str, Any]:
    return {
        "dtype": str(array.dtype),
        "shape": list(array.shape),
        "values": array.tolist(),
    }


def build_sglang_payload(
    profile: Pi05BenchProfile,
    observation: dict[str, Any],
    *,
    num_inference_steps: int,
    prefix_cache: bool,
    cuda_graph: bool,
    noise: np.ndarray | None,
    response_format: str = "envelope",
) -> dict[str, Any]:
    encoded_images = {
        key: _json_tensor(np.asarray(value))
        for key, value in observation["images"].items()
    }
    encoded_observation = {
        "images": encoded_images,
        "state": _json_tensor(np.asarray(observation["state"], dtype=np.float32)),
    }
    if noise is not None:
        encoded_observation["noise"] = _json_tensor(noise.astype(np.float32))
    return {
        "model": profile.sglang_model,
        "input": {
            "task": profile.prompt,
            "observation": encoded_observation,
        },
        "parameters": {
            "num_inference_steps": num_inference_steps,
        },
        "runtime": {
            "return_timing": True,
            "prefix_cache": prefix_cache,
            "cuda_graph": cuda_graph,
            "response_format": response_format,
        },
    }


def build_sglang_python_payload(
    profile: Pi05BenchProfile,
    observation: dict[str, Any],
    *,
    num_inference_steps: int,
    prefix_cache: bool,
    cuda_graph: bool,
    noise: np.ndarray | None,
    response_format: str = "envelope",
) -> dict[str, Any]:
    encoded_observation = {
        "images": {
            key: np.asarray(value) for key, value in observation["images"].items()
        },
        "state": np.asarray(observation["state"], dtype=np.float32),
    }
    if noise is not None:
        encoded_observation["noise"] = noise.astype(np.float32)
    return {
        "model": profile.sglang_model,
        "input": {
            "task": profile.prompt,
            "observation": encoded_observation,
        },
        "parameters": {
            "num_inference_steps": num_inference_steps,
        },
        "runtime": {
            "return_timing": True,
            "prefix_cache": prefix_cache,
            "cuda_graph": cuda_graph,
            "output_format": "numpy",
            "response_format": response_format,
        },
    }


def build_sglang_openpi_ws_payload(
    profile: Pi05BenchProfile,
    observation: dict[str, Any],
    *,
    num_inference_steps: int,
    prefix_cache: bool,
    cuda_graph: bool,
    noise: np.ndarray | None,
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "task": profile.prompt,
        "observation.state": np.asarray(observation["state"], dtype=np.float32),
        "num_inference_steps": num_inference_steps,
        "enable_pi_prefix_cache": prefix_cache,
        "enable_pi_cuda_graph": cuda_graph,
        "output_format": "numpy",
    }
    for key, value in observation["images"].items():
        payload[f"observation.images.{key}"] = np.asarray(value)
    if noise is not None:
        payload["observation.noise"] = noise.astype(np.float32)
    return payload


def _post_action(
    session: requests.Session,
    url: str,
    payload: dict[str, Any],
    timeout_s: float,
) -> dict[str, Any]:
    response = session.post(url, json=payload, timeout=timeout_s)
    response.raise_for_status()
    return response.json()


def _post_action_msgpack(
    session: requests.Session,
    url: str,
    payload: dict[str, Any],
    timeout_s: float,
) -> dict[str, Any]:
    response = session.post(
        url,
        data=pack_msgpack(payload),
        headers={
            "Content-Type": "application/msgpack",
            "Accept": "application/msgpack",
        },
        timeout=timeout_s,
    )
    response.raise_for_status()
    return unpack_msgpack(response.content)


def _get_action_metadata(session: requests.Session, url: str, timeout_s: float):
    response = session.get(
        url.rstrip("/") + "/v1/actions/metadata",
        timeout=timeout_s,
    )
    response.raise_for_status()
    return response.json()


def run_sglang_http(
    url: str,
    payloads: list[dict[str, Any]],
    *,
    warmup: int,
    repeats: int,
    batch_size: int,
    timeout_s: float,
    msgpack: bool = False,
) -> dict[str, Any]:
    endpoint = url.rstrip("/") + "/v1/actions/generations"
    post_action = _post_action_msgpack if msgpack else _post_action

    single_latencies = []
    single_outputs = []
    batch_latencies = []
    with requests.Session() as session:
        metadata = _get_action_metadata(session, url, timeout_s)
        for idx in range(min(warmup, len(payloads))):
            post_action(session, endpoint, payloads[idx], timeout_s)

        for idx in range(repeats):
            payload = payloads[idx % len(payloads)]
            start = time.perf_counter()
            output = post_action(session, endpoint, payload, timeout_s)
            single_latencies.append((time.perf_counter() - start) * 1000)
            single_outputs.append(output)

    if batch_size > 1:
        sessions = [requests.Session() for _ in range(batch_size)]
        try:
            with ThreadPoolExecutor(max_workers=batch_size) as pool:

                def post_item(item):
                    session, payload = item
                    return post_action(session, endpoint, payload, timeout_s)

                for warmup_idx in range(warmup):
                    batch = [
                        payloads[(warmup_idx * batch_size + offset) % len(payloads)]
                        for offset in range(batch_size)
                    ]
                    list(pool.map(post_item, zip(sessions, batch)))

                for start_idx in range(repeats):
                    batch = [
                        payloads[(start_idx * batch_size + offset) % len(payloads)]
                        for offset in range(batch_size)
                    ]
                    start = time.perf_counter()
                    list(pool.map(post_item, zip(sessions, batch)))
                    batch_latencies.append((time.perf_counter() - start) * 1000)
        finally:
            for session in sessions:
                session.close()

    stage_timings = {}
    for output in single_outputs:
        for key, value in output.get("timings", {}).items():
            stage_timings.setdefault(key, []).append(float(value))

    return {
        "single": _stats_ms(single_latencies),
        "batch": _stats_ms(batch_latencies),
        "batch_size": batch_size,
        "stage_timings": {
            key: _stats_ms(values) for key, values in stage_timings.items()
        },
        "first_output": single_outputs[0] if single_outputs else None,
        "batch_mode": "concurrent_http_msgpack" if msgpack else "concurrent_http_json",
        "metadata": metadata,
    }


def _action_ws_url(url: str) -> str:
    if url.startswith("https://"):
        return "wss://" + url[len("https://") :].rstrip("/") + "/openpi/policy"
    if url.startswith("http://"):
        return "ws://" + url[len("http://") :].rstrip("/") + "/openpi/policy"
    return url.rstrip("/") + "/openpi/policy"


async def _ws_send_recv(websocket, payload: dict[str, Any]) -> dict[str, Any]:
    await websocket.send(pack_msgpack(payload))
    response = await websocket.recv()
    if isinstance(response, str):
        raise RuntimeError(response)
    return unpack_msgpack(response)


async def _run_sglang_openpi_ws_async(
    url: str,
    payloads: list[dict[str, Any]],
    *,
    warmup: int,
    repeats: int,
    batch_size: int,
) -> dict[str, Any]:
    import websockets

    endpoint = _action_ws_url(url)
    async with websockets.connect(endpoint, max_size=None) as websocket:
        metadata = unpack_msgpack(await websocket.recv())
        for idx in range(min(warmup, len(payloads))):
            await _ws_send_recv(websocket, payloads[idx])

        single_latencies = []
        single_outputs = []
        for idx in range(repeats):
            payload = payloads[idx % len(payloads)]
            start = time.perf_counter()
            output = await _ws_send_recv(websocket, payload)
            single_latencies.append((time.perf_counter() - start) * 1000)
            single_outputs.append(output)

    batch_latencies = []
    if batch_size > 1:
        websockets_list = []
        try:
            for _ in range(batch_size):
                websocket = await websockets.connect(endpoint, max_size=None)
                await websocket.recv()
                websockets_list.append(websocket)

            for warmup_idx in range(warmup):
                batch = [
                    payloads[(warmup_idx * batch_size + offset) % len(payloads)]
                    for offset in range(batch_size)
                ]
                await asyncio.gather(
                    *[
                        _ws_send_recv(websocket, payload)
                        for websocket, payload in zip(websockets_list, batch)
                    ]
                )

            for start_idx in range(repeats):
                batch = [
                    payloads[(start_idx * batch_size + offset) % len(payloads)]
                    for offset in range(batch_size)
                ]
                start = time.perf_counter()
                await asyncio.gather(
                    *[
                        _ws_send_recv(websocket, payload)
                        for websocket, payload in zip(websockets_list, batch)
                    ]
                )
                batch_latencies.append((time.perf_counter() - start) * 1000)
        finally:
            for websocket in websockets_list:
                await websocket.close()

    stage_timings = {}
    server_timings = {}
    for output in single_outputs:
        for key, value in output.get("timings", {}).items():
            stage_timings.setdefault(key, []).append(float(value))
        for key, value in output.get("server_timing", {}).items():
            server_timings.setdefault(key, []).append(float(value))

    return {
        "single": _stats_ms(single_latencies),
        "batch": _stats_ms(batch_latencies),
        "batch_size": batch_size,
        "stage_timings": {
            key: _stats_ms(values) for key, values in stage_timings.items()
        },
        "server_timings": {
            key: _stats_ms(values) for key, values in server_timings.items()
        },
        "first_output": single_outputs[0] if single_outputs else None,
        "batch_mode": "persistent_openpi_websocket",
        "metadata": metadata,
    }


def run_sglang_openpi_ws(
    url: str,
    payloads: list[dict[str, Any]],
    *,
    warmup: int,
    repeats: int,
    batch_size: int,
) -> dict[str, Any]:
    return asyncio.run(
        _run_sglang_openpi_ws_async(
            url,
            payloads,
            warmup=warmup,
            repeats=repeats,
            batch_size=batch_size,
        )
    )


def create_sglang_python_pipeline(
    model_path: str,
    *,
    pipeline_config_path: str | None,
):
    from sglang.multimodal_gen.runtime.pipelines.pi05 import Pi05Pipeline
    from sglang.multimodal_gen.runtime.pipelines_core.executors.sync_executor import (
        SyncExecutor,
    )
    from sglang.multimodal_gen.runtime.server_args import (
        ServerArgs,
        set_global_server_args,
    )

    kwargs: dict[str, Any] = {
        "model_path": model_path,
        "warmup_mode": "off",
        "num_gpus": 1,
    }
    if pipeline_config_path:
        kwargs["pipeline_config_path"] = pipeline_config_path
    server_args = ServerArgs.from_kwargs(**kwargs)
    set_global_server_args(server_args)
    pipeline = Pi05Pipeline(
        model_path,
        server_args,
        executor=SyncExecutor(server_args=server_args),
    )
    return pipeline, server_args


def _make_sglang_python_req(server_args, payload: dict[str, Any]):
    from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
    from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
        build_action_sampling_params,
    )

    sampling_params = build_action_sampling_params(payload, server_args)
    req = prepare_request(server_args, sampling_params)
    req.suppress_logs = True
    return req


def _run_sglang_python_once(pipeline, server_args, payload: dict[str, Any]):
    req = _make_sglang_python_req(server_args, payload)
    output_batch = pipeline.forward(req, server_args)
    if output_batch.error:
        raise RuntimeError(output_batch.error)
    if not output_batch.output:
        raise RuntimeError("SGLang Python policy returned no output")
    return output_batch.output[0]


def _run_sglang_python_group(pipeline, server_args, payloads: list[dict[str, Any]]):
    reqs = [_make_sglang_python_req(server_args, payload) for payload in payloads]
    output_batches = pipeline.forward_batch(reqs, server_args)
    outputs = []
    for output_batch in output_batches:
        if output_batch.error:
            raise RuntimeError(output_batch.error)
        if not output_batch.output:
            raise RuntimeError("SGLang Python grouped policy returned no output")
        outputs.append(output_batch.output[0])
    return outputs


def run_sglang_python(
    model_path: str,
    payloads: list[dict[str, Any]],
    *,
    pipeline_config_path: str | None,
    warmup: int,
    repeats: int,
    batch_size: int,
    batch_mode: str,
) -> dict[str, Any]:
    pipeline, server_args = create_sglang_python_pipeline(
        model_path,
        pipeline_config_path=pipeline_config_path,
    )
    from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import action_metadata

    metadata = action_metadata(server_args)
    metadata["precision"] = sglang_precision_metadata(pipeline)
    for idx in range(min(warmup, len(payloads))):
        _run_sglang_python_once(pipeline, server_args, payloads[idx])

    single_latencies = []
    single_outputs = []
    for idx in range(repeats):
        payload = payloads[idx % len(payloads)]
        start = time.perf_counter()
        output = _run_sglang_python_once(pipeline, server_args, payload)
        single_latencies.append((time.perf_counter() - start) * 1000)
        single_outputs.append(output)

    batch_latencies = []
    batch_outputs = []
    if batch_size > 1:
        for warmup_idx in range(warmup):
            batch = [
                payloads[(warmup_idx * batch_size + offset) % len(payloads)]
                for offset in range(batch_size)
            ]
            if batch_mode == "grouped":
                _run_sglang_python_group(pipeline, server_args, batch)
            else:
                for payload in batch:
                    _run_sglang_python_once(pipeline, server_args, payload)
        for start_idx in range(repeats):
            batch = [
                payloads[(start_idx * batch_size + offset) % len(payloads)]
                for offset in range(batch_size)
            ]
            start = time.perf_counter()
            if batch_mode == "grouped":
                outputs = _run_sglang_python_group(pipeline, server_args, batch)
            else:
                outputs = []
                for payload in batch:
                    outputs.append(
                        _run_sglang_python_once(pipeline, server_args, payload)
                    )
            batch_latencies.append((time.perf_counter() - start) * 1000)
            batch_outputs.extend(outputs)

    stage_timings = {}
    for output in single_outputs:
        for key, value in output.get("timings", {}).items():
            stage_timings.setdefault(key, []).append(float(value))
    batch_stage_timings = {}
    for output in batch_outputs:
        for key, value in output.get("timings", {}).items():
            batch_stage_timings.setdefault(key, []).append(float(value))

    return {
        "single": _stats_ms(single_latencies),
        "batch": _stats_ms(batch_latencies),
        "batch_size": batch_size,
        "stage_timings": {
            key: _stats_ms(values) for key, values in stage_timings.items()
        },
        "batch_stage_timings": {
            key: _stats_ms(values) for key, values in batch_stage_timings.items()
        },
        "first_output": single_outputs[0] if single_outputs else None,
        "batch_mode": f"python_policy_{batch_mode}",
        "metadata": metadata,
    }


def create_openpi_policy(
    config_name: str,
    checkpoint_dir: str,
    *,
    pytorch_device: str,
    num_inference_steps: int,
    pytorch_compile_mode: str | None,
):
    from openpi.policies import policy_config
    from openpi.training import config as openpi_config

    train_config = openpi_config.get_config(config_name)
    if pytorch_compile_mode != "keep":
        train_config = dataclasses.replace(
            train_config,
            model=dataclasses.replace(
                train_config.model,
                pytorch_compile_mode=pytorch_compile_mode,
            ),
        )
    return policy_config.create_trained_policy(
        train_config,
        checkpoint_dir,
        sample_kwargs={"num_steps": num_inference_steps},
        pytorch_device=pytorch_device,
    )


def _openpi_infer(policy, observation: dict[str, Any], noise: np.ndarray | None):
    if noise is None:
        return policy.infer(observation)
    return policy.infer(observation, noise=noise)


def _openpi_direct_batch(
    policy,
    observations: list[dict[str, Any]],
    noises: list[np.ndarray] | None,
):
    import jax
    import numpy as onp
    from openpi.models import model as openpi_model

    inputs_list = [
        policy._input_transform(jax.tree.map(lambda value: value, observation))
        for observation in observations
    ]
    batched_inputs = jax.tree.map(
        lambda *values: onp.stack(values, axis=0),
        *inputs_list,
    )
    sample_kwargs = dict(policy._sample_kwargs)
    if noises is not None:
        noise = onp.stack(noises, axis=0)
        if policy._is_pytorch_model:
            import torch

            sample_kwargs["noise"] = torch.from_numpy(noise).to(policy._pytorch_device)
        else:
            import jax.numpy as jnp

            sample_kwargs["noise"] = jnp.asarray(noise)

    if policy._is_pytorch_model:
        import torch

        inputs = jax.tree.map(
            lambda value: torch.from_numpy(onp.asarray(value)).to(
                policy._pytorch_device
            ),
            batched_inputs,
        )
        sample_key_or_device = policy._pytorch_device
    else:
        import jax.numpy as jnp

        inputs = jax.tree.map(lambda value: jnp.asarray(value), batched_inputs)
        policy._rng, sample_key_or_device = jax.random.split(policy._rng)

    observation = openpi_model.Observation.from_dict(inputs)
    actions = policy._sample_actions(sample_key_or_device, observation, **sample_kwargs)
    if policy._is_pytorch_model:
        actions_np = actions.detach().cpu().numpy()
        states_np = inputs["state"].detach().cpu().numpy()
    else:
        actions_np = onp.asarray(actions)
        states_np = onp.asarray(inputs["state"])

    outputs = []
    for idx in range(actions_np.shape[0]):
        outputs.append(
            policy._output_transform(
                {
                    "state": states_np[idx],
                    "actions": actions_np[idx],
                }
            )
        )
    return outputs


def run_openpi_policy(
    policy,
    observations: list[dict[str, Any]],
    *,
    warmup: int,
    repeats: int,
    batch_size: int,
    noise: np.ndarray | None,
    batch_mode: str,
) -> dict[str, Any]:
    for idx in range(min(warmup, len(observations))):
        _openpi_infer(policy, observations[idx], noise)

    single_latencies = []
    single_outputs = []
    for idx in range(repeats):
        observation = observations[idx % len(observations)]
        start = time.perf_counter()
        output = _openpi_infer(policy, observation, noise)
        single_latencies.append((time.perf_counter() - start) * 1000)
        single_outputs.append(output)

    batch_latencies = []
    if batch_size > 1:
        for warmup_idx in range(warmup):
            batch = [
                observations[(warmup_idx * batch_size + offset) % len(observations)]
                for offset in range(batch_size)
            ]
            noises = [noise] * len(batch) if noise is not None else None
            if batch_mode == "direct_model":
                _openpi_direct_batch(policy, batch, noises)
            elif batch_mode == "policy_loop":
                for obs in batch:
                    _openpi_infer(policy, obs, noise)
            else:
                raise ValueError(f"Unsupported OpenPI batch mode: {batch_mode}")
        for start_idx in range(repeats):
            batch = [
                observations[(start_idx * batch_size + offset) % len(observations)]
                for offset in range(batch_size)
            ]
            noises = [noise] * len(batch) if noise is not None else None
            start = time.perf_counter()
            if batch_mode == "direct_model":
                _openpi_direct_batch(policy, batch, noises)
            elif batch_mode == "policy_loop":
                for obs in batch:
                    _openpi_infer(policy, obs, noise)
            else:
                raise ValueError(f"Unsupported OpenPI batch mode: {batch_mode}")
            batch_latencies.append((time.perf_counter() - start) * 1000)

    policy_timings = {}
    for output in single_outputs:
        for key, value in output.get("policy_timing", {}).items():
            policy_timings.setdefault(key, []).append(float(value))

    precision = openpi_precision_metadata(policy)
    first_actions = _openpi_actions(single_outputs[0]) if single_outputs else None
    if first_actions is not None:
        precision["output_action_dtype"] = str(first_actions.dtype)
        precision["output_action_shape"] = list(first_actions.shape)

    return {
        "single": _stats_ms(single_latencies),
        "batch": _stats_ms(batch_latencies),
        "batch_size": batch_size,
        "policy_timings": {
            key: _stats_ms(values) for key, values in policy_timings.items()
        },
        "first_output": single_outputs[0] if single_outputs else None,
        "batch_mode": batch_mode,
        "precision": precision,
    }


def _sglang_actions(output: dict[str, Any]) -> np.ndarray | None:
    if output is None:
        return None
    if "actions" in output:
        return np.asarray(output["actions"], dtype=np.float32)
    return np.asarray(output["data"][0]["action"]["values"], dtype=np.float32)


def _openpi_actions(output: dict[str, Any]) -> np.ndarray | None:
    if output is None:
        return None
    return np.asarray(output["actions"], dtype=np.float32)


def compare_first_actions(
    sglang_output: dict[str, Any] | None,
    openpi_output: dict[str, Any] | None,
) -> dict[str, Any]:
    sglang_actions = _sglang_actions(sglang_output)
    openpi_actions = _openpi_actions(openpi_output)
    if sglang_actions is None or openpi_actions is None:
        return {"available": False}
    horizon = min(sglang_actions.shape[0], openpi_actions.shape[0])
    dim = min(sglang_actions.shape[1], openpi_actions.shape[1])
    if horizon == 0 or dim == 0:
        return {
            "available": False,
            "sglang_shape": list(sglang_actions.shape),
            "openpi_shape": list(openpi_actions.shape),
        }
    diff = np.abs(sglang_actions[:horizon, :dim] - openpi_actions[:horizon, :dim])
    return {
        "available": True,
        "common_shape": [int(horizon), int(dim)],
        "sglang_shape": list(sglang_actions.shape),
        "openpi_shape": list(openpi_actions.shape),
        "max_abs_diff": float(np.max(diff)),
        "mean_abs_diff": float(np.mean(diff)),
    }


def print_summary(result: dict[str, Any]) -> None:
    print(json.dumps(result, indent=2, sort_keys=True))
    sglang_result = result.get("sglang") or {}
    openpi_result = result.get("openpi") or {}
    sgl = sglang_result.get("single", {}).get("mean_ms")
    opi = openpi_result.get("single", {}).get("mean_ms")
    if sgl and opi:
        print(
            "\nSingle mean latency: "
            f"SGLang={sgl:.2f} ms, OpenPI={opi:.2f} ms, "
            f"speedup={opi / sgl:.2f}x"
        )
    sgl_batch = sglang_result.get("batch", {}).get("mean_ms")
    opi_batch = openpi_result.get("batch", {}).get("mean_ms")
    batch_size = sglang_result.get("batch_size") or openpi_result.get("batch_size", 0)
    if sgl_batch and opi_batch and batch_size:
        print(
            "Batch mean latency: "
            f"SGLang={sgl_batch:.2f} ms/{batch_size} req, "
            f"OpenPI={opi_batch:.2f} ms/{batch_size} req, "
            f"speedup={opi_batch / sgl_batch:.2f}x"
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--profile", choices=sorted(PROFILES), default="libero")
    parser.add_argument("--sglang-url", default="http://127.0.0.1:30000")
    parser.add_argument(
        "--sglang-api",
        choices=("http", "http_msgpack", "openpi_ws", "python"),
        default="http",
    )
    parser.add_argument("--sglang-model", default=None)
    parser.add_argument("--sglang-pipeline-config-path", default=None)
    parser.add_argument(
        "--sglang-http-response-format",
        choices=("envelope", "raw"),
        default="envelope",
    )
    parser.add_argument(
        "--sglang-python-batch-mode",
        choices=("loop", "grouped"),
        default="loop",
    )
    parser.add_argument("--openpi-config", default=None)
    parser.add_argument("--openpi-checkpoint", default=None)
    parser.add_argument("--openpi-device", default="cuda")
    parser.add_argument(
        "--openpi-pytorch-compile-mode",
        choices=(
            "keep",
            "none",
            "default",
            "reduce-overhead",
            "max-autotune",
            "max-autotune-no-cudagraphs",
        ),
        default="keep",
    )
    parser.add_argument(
        "--openpi-batch-mode",
        choices=("direct_model", "policy_loop"),
        default="direct_model",
    )
    parser.add_argument("--num-inference-steps", type=int, default=10)
    parser.add_argument("--num-samples", type=int, default=16)
    parser.add_argument("--repeats", type=int, default=20)
    parser.add_argument("--warmup", type=int, default=3)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--timeout-s", type=float, default=180.0)
    parser.add_argument("--disable-prefix-cache", action="store_true")
    parser.add_argument("--disable-cuda-graph", action="store_true")
    parser.add_argument("--deterministic-noise", action="store_true")
    parser.add_argument("--skip-sglang", action="store_true")
    parser.add_argument("--skip-openpi", action="store_true")
    parser.add_argument("--output-file", default="")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    profile = PROFILES[args.profile]
    if args.sglang_model:
        profile = Pi05BenchProfile(
            name=profile.name,
            sglang_model=args.sglang_model,
            openpi_config=profile.openpi_config,
            openpi_checkpoint=profile.openpi_checkpoint,
            prompt=profile.prompt,
            sglang_action_horizon=profile.sglang_action_horizon,
            openpi_action_horizon=profile.openpi_action_horizon,
            action_dim=profile.action_dim,
            output_action_dim=profile.output_action_dim,
        )
    openpi_config = args.openpi_config or profile.openpi_config
    openpi_checkpoint = args.openpi_checkpoint or profile.openpi_checkpoint

    openpi_observations, sglang_observations = build_observations(
        profile,
        max(args.num_samples, args.batch_size),
        args.seed,
    )
    noise = None
    if args.skip_sglang and args.skip_openpi:
        raise ValueError("At least one backend must be enabled")
    if args.deterministic_noise:
        if profile.sglang_action_horizon != profile.openpi_action_horizon:
            raise ValueError(
                "--deterministic-noise requires matching SGLang and OpenPI "
                f"action horizons; profile {profile.name!r} has "
                f"{profile.sglang_action_horizon} vs {profile.openpi_action_horizon}"
            )
        rng = np.random.default_rng(args.seed + 1)
        noise = rng.standard_normal(
            (profile.sglang_action_horizon, profile.action_dim),
            dtype=np.float32,
        )

    payloads = []
    if args.skip_sglang:
        pass
    elif args.sglang_api in ("python", "http_msgpack"):
        payloads = [
            build_sglang_python_payload(
                profile,
                observation,
                num_inference_steps=args.num_inference_steps,
                prefix_cache=not args.disable_prefix_cache,
                cuda_graph=not args.disable_cuda_graph,
                noise=noise,
                response_format=args.sglang_http_response_format,
            )
            for observation in sglang_observations
        ]
    elif args.sglang_api == "openpi_ws":
        payloads = [
            build_sglang_openpi_ws_payload(
                profile,
                observation,
                num_inference_steps=args.num_inference_steps,
                prefix_cache=not args.disable_prefix_cache,
                cuda_graph=not args.disable_cuda_graph,
                noise=noise,
            )
            for observation in sglang_observations
        ]
    else:
        payloads = [
            build_sglang_payload(
                profile,
                observation,
                num_inference_steps=args.num_inference_steps,
                prefix_cache=not args.disable_prefix_cache,
                cuda_graph=not args.disable_cuda_graph,
                noise=noise,
                response_format=args.sglang_http_response_format,
            )
            for observation in sglang_observations
        ]

    openpi_policy = None
    if not args.skip_openpi:
        openpi_policy = create_openpi_policy(
            openpi_config,
            openpi_checkpoint,
            pytorch_device=args.openpi_device,
            num_inference_steps=args.num_inference_steps,
            pytorch_compile_mode=(
                None
                if args.openpi_pytorch_compile_mode == "none"
                else args.openpi_pytorch_compile_mode
            ),
        )

    sglang_result = None
    if args.skip_sglang:
        pass
    elif args.sglang_api == "python":
        sglang_result = run_sglang_python(
            profile.sglang_model,
            payloads,
            pipeline_config_path=args.sglang_pipeline_config_path,
            warmup=args.warmup,
            repeats=args.repeats,
            batch_size=args.batch_size,
            batch_mode=args.sglang_python_batch_mode,
        )
    elif args.sglang_api == "openpi_ws":
        sglang_result = run_sglang_openpi_ws(
            args.sglang_url,
            payloads,
            warmup=args.warmup,
            repeats=args.repeats,
            batch_size=args.batch_size,
        )
    else:
        sglang_result = run_sglang_http(
            args.sglang_url,
            payloads,
            warmup=args.warmup,
            repeats=args.repeats,
            batch_size=args.batch_size,
            timeout_s=args.timeout_s,
            msgpack=args.sglang_api == "http_msgpack",
        )
    openpi_result = None
    if openpi_policy is not None:
        openpi_result = run_openpi_policy(
            openpi_policy,
            openpi_observations,
            warmup=args.warmup,
            repeats=args.repeats,
            batch_size=args.batch_size,
            noise=noise,
            batch_mode=args.openpi_batch_mode,
        )

    result = {
        "profile": profile.name,
        "sglang_model": profile.sglang_model,
        "sglang_api": args.sglang_api,
        "sglang_pipeline_config_path": args.sglang_pipeline_config_path,
        "sglang_http_response_format": args.sglang_http_response_format,
        "sglang_python_batch_mode": args.sglang_python_batch_mode,
        "openpi_config": openpi_config,
        "openpi_checkpoint": openpi_checkpoint,
        "openpi_pytorch_compile_mode": args.openpi_pytorch_compile_mode,
        "num_inference_steps": args.num_inference_steps,
        "num_samples": args.num_samples,
        "repeats": args.repeats,
        "warmup": args.warmup,
        "deterministic_noise": args.deterministic_noise,
        "action_diff": compare_first_actions(
            None if sglang_result is None else sglang_result.get("first_output"),
            None if openpi_result is None else openpi_result.get("first_output"),
        ),
        "sglang": (
            None
            if sglang_result is None
            else {
                key: value
                for key, value in sglang_result.items()
                if key not in ("first_output",)
            }
        ),
        "openpi": (
            None
            if openpi_result is None
            else {
                key: value
                for key, value in openpi_result.items()
                if key not in ("first_output",)
            }
        ),
    }
    if args.output_file:
        with open(args.output_file, "w", encoding="utf-8") as f:
            json.dump(result, f, indent=2, sort_keys=True)
    print_summary(result)


if __name__ == "__main__":
    main()