From 02e416098cba7cd2389f1bc394d1141a43f06327 Mon Sep 17 00:00:00 2001 From: qyy Date: Wed, 24 Jun 2026 20:43:37 +0800 Subject: [PATCH 01/29] feat(phyai): add PI05 LIBERO policy adapter --- phyai/src/phyai/env.py | 4 + phyai/src/phyai/policies/__init__.py | 5 + phyai/src/phyai/policies/pi05_libero.py | 445 ++++++++++++++++++++++++ 3 files changed, 454 insertions(+) create mode 100644 phyai/src/phyai/policies/__init__.py create mode 100644 phyai/src/phyai/policies/pi05_libero.py diff --git a/phyai/src/phyai/env.py b/phyai/src/phyai/env.py index 545b12d..08acede 100644 --- a/phyai/src/phyai/env.py +++ b/phyai/src/phyai/env.py @@ -155,6 +155,10 @@ class envs: # ---------- runtime ---------- # PHYAI_USE_CUDA_GRAPH = EnvField("PHYAI_USE_CUDA_GRAPH", None, _parse_bool) + # ---------- policy adapters ---------- # + PHYAI_CAMERA_MODE = EnvField("PHYAI_CAMERA_MODE", None, str) + PHYAI_TOKENIZER_PATH = EnvField("PHYAI_TOKENIZER_PATH", None, str) + # ---------- parallel ---------- # PHYAI_WORLD_SIZE = EnvField("PHYAI_WORLD_SIZE", None, int) PHYAI_DP_SIZE = EnvField("PHYAI_DP_SIZE", None, int) diff --git a/phyai/src/phyai/policies/__init__.py b/phyai/src/phyai/policies/__init__.py new file mode 100644 index 0000000..9a1a0cf --- /dev/null +++ b/phyai/src/phyai/policies/__init__.py @@ -0,0 +1,5 @@ +"""High-level policy wrappers.""" + +from phyai.policies.pi05_libero import PI05LiberoPolicy + +__all__ = ["PI05LiberoPolicy"] diff --git a/phyai/src/phyai/policies/pi05_libero.py b/phyai/src/phyai/policies/pi05_libero.py new file mode 100644 index 0000000..bcaea76 --- /dev/null +++ b/phyai/src/phyai/policies/pi05_libero.py @@ -0,0 +1,445 @@ +"""Thin LIBERO adapter for pi0.5 PhyAI inference.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from safetensors.torch import load_file + +from phyai.engine import Engine, EngineArgs +from phyai.engine_config import BackendConfig, DeviceConfig, EngineConfig, RuntimeConfig +from phyai.env import envs +from phyai.models.pi05.configuration_pi05 import PI05Config +from phyai.models.pi05.main_pi05 import PI05Args +from phyai.models.pi05.scheduler_ws1_pi05 import PI05Request +from phyai_utils_tools.models.pi05 import PI05_DEFAULT_TOKENIZER_NAME, PI05Processor +from phyai_utils_tools.processing.transition import IMAGES, STATE, TASK + +LIBERO_AGENTVIEW_KEYS: tuple[str, ...] = ( + "agentview", + "agentview_image", + "image", + "observation.images.image", +) +LIBERO_WRIST_KEYS: tuple[str, ...] = ( + "wrist", + "robot0_eye_in_hand_image", + "wrist_image", + "image2", + "observation.images.image2", +) + + +def _lerobot_pi05_weight_remap(key: str) -> str | None: + """Strip LeRobot's outer model prefix and drop inference-unused keys.""" + if key.startswith("model."): + key = key[len("model.") :] + if key == "paligemma_with_expert.gemma_expert.lm_head.weight": + return None + return key + + +class PI05LiberoPolicy: + """Adapt vla-evaluation-harness LIBERO observations to ``PI05Processor``.""" + + def __init__( + self, + checkpoint_dir: str | Path, + *, + device: str = "cuda", + params_dtype: torch.dtype = torch.bfloat16, + max_batch_size: int = 1, + use_cuda_graph: bool = True, + attn_backend: str = "flashinfer", + norm_backend: str = "phyai-kernel", + linear_backend: str | None = "flashinfer", + flashinfer_workspace_bytes: int = 512 * 1024 * 1024, + tokenizer_name: str | None = None, + camera_mode: str | None = None, + ) -> None: + self.checkpoint_dir = Path(checkpoint_dir) + self.device = device + self.params_dtype = params_dtype + self.max_batch_size = int(max_batch_size) + self.config = self._read_config() + self.image_size = self._resolve_image_size(self.config) + self._action_dim = self._resolve_action_dim(self.config) + self.max_action_dim = int(self.config.get("max_action_dim", 32)) + self._chunk_size = int(self.config.get("chunk_size", PI05Config().chunk_size)) + self.camera_names = self._resolve_camera_names(camera_mode) + self.tokenizer_name = self._resolve_tokenizer_name(tokenizer_name) + self.prompt_mode = str( + self.config.get("phyai_prompt_mode", "lerobot_state_bins") + ) + self.normalization_mode = str( + self.config.get("phyai_normalization_mode", "mean_std") + ) + self._use_phyai_compat = ( + "phyai_prompt_mode" in self.config + or "phyai_normalization_mode" in self.config + ) + self._normalizer_stats = self._load_processor_state( + "policy_preprocessor.json", "normalizer_processor" + ) + self._unnormalizer_stats = self._load_processor_state( + "policy_postprocessor.json", "unnormalizer_processor" + ) + if self._use_phyai_compat: + self._validate_compat_stats() + self._tokenizer = None + self.processor = PI05Processor.from_pretrained( + self.checkpoint_dir, + tokenizer_name=self.tokenizer_name, + image_size=self.image_size, + num_channels=3, + num_images=len(self.camera_names), + action_dim=self._action_dim, + normalize_pixels=True, + device=device, + params_dtype=params_dtype, + ) + self.engine = Engine( + EngineArgs( + plugin="pi05", + plugin_args=PI05Args( + checkpoint_dir=self.checkpoint_dir, + max_batch_size=self.max_batch_size, + weight_remap=_lerobot_pi05_weight_remap, + inputs_image_shape=[ + [self.image_size, self.image_size, 3] for _ in self.camera_names + ], + ), + config=EngineConfig( + backends=BackendConfig( + attn=attn_backend, norm=norm_backend, linear=linear_backend + ), + device=DeviceConfig(target=device, params_dtype=params_dtype), + runtime=RuntimeConfig( + use_cuda_graph=use_cuda_graph, + flashinfer_workspace_bytes=flashinfer_workspace_bytes, + force_linear_kernel=linear_backend, + ), + ), + ) + ) + + @property + def chunk_size(self) -> int: + return self._chunk_size + + @property + def action_dim(self) -> int: + return int(self.processor.action_dim or self._action_dim) + + @staticmethod + def _resolve_image_size(config: dict[str, Any]) -> int: + resolution = config.get("image_resolution") + if isinstance(resolution, list) and resolution: + return int(resolution[0]) + return PI05Config().vision.image_size + + @staticmethod + def _resolve_action_dim(config: dict[str, Any]) -> int: + shape = config.get("output_features", {}).get("action", {}).get("shape") + if isinstance(shape, list) and shape: + return int(shape[-1]) + return 7 + + def _read_config(self) -> dict[str, Any]: + path = self.checkpoint_dir / "config.json" + if not path.exists(): + return {} + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + def _resolve_camera_names(self, camera_mode: str | None) -> list[str]: + mode = camera_mode or envs.PHYAI_CAMERA_MODE.get() or "three_camera" + if mode == "two_camera": + return ["agentview", "wrist"] + if mode == "three_camera": + return ["agentview", "wrist", "empty"] + raise ValueError(f"Unsupported PHYAI_CAMERA_MODE={mode!r}.") + + def _resolve_tokenizer_name(self, tokenizer_name: str | None) -> str: + if tokenizer_name: + return tokenizer_name + if env_tokenizer := envs.PHYAI_TOKENIZER_PATH.get(): + return env_tokenizer + if config_tokenizer := self.config.get("tokenizer_name"): + return str(config_tokenizer) + return PI05_DEFAULT_TOKENIZER_NAME + + def _load_processor_state( + self, config_name: str, registry_name: str + ) -> dict[str, torch.Tensor]: + path = self.checkpoint_dir / config_name + if not path.exists(): + return {} + with path.open("r", encoding="utf-8") as f: + config = json.load(f) + for step in config.get("steps", []): + if step.get("registry_name") != registry_name: + continue + state_file = step.get("state_file") + if not state_file: + return {} + return load_file(str(self.checkpoint_dir / state_file)) + return {} + + def _validate_compat_stats(self) -> None: + if self.normalization_mode == "openpi_quantile": + normalizer_keys = ("observation.state.min", "observation.state.max") + unnormalizer_keys = ("action.min", "action.max") + else: + normalizer_keys = ("observation.state.mean", "observation.state.std") + unnormalizer_keys = ("action.mean", "action.std") + missing = [ + f"normalizer:{key}" + for key in normalizer_keys + if key not in self._normalizer_stats + ] + missing.extend( + f"unnormalizer:{key}" + for key in unnormalizer_keys + if key not in self._unnormalizer_stats + ) + if missing: + raise ValueError( + f"{self.checkpoint_dir}: compat normalization requires missing stats " + f"{', '.join(missing)}" + ) + + @property + def tokenizer(self): + if self._tokenizer is None: + from transformers import AutoTokenizer + + self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name) + return self._tokenizer + + def observation_to_raw(self, obs: dict[str, Any]) -> dict[str, Any]: + return { + IMAGES: [ + self._extract_camera_tensor(obs, name) for name in self.camera_names + ], + STATE: self._extract_state(obs), + TASK: [self._extract_task(obs)], + } + + def observation_to_request_inputs( + self, obs: dict[str, Any] + ) -> dict[str, torch.Tensor]: + if not self._use_phyai_compat: + processed = self.processor.preprocess(self.observation_to_raw(obs)) + return { + "pixel_values": processed.pixel_values, + "input_ids": processed.input_ids, + "lang_lens": processed.lang_lens, + } + pixel_values = ( + torch.stack( + [ + self._extract_camera_model_tensor(obs, name).squeeze(0) + for name in self.camera_names + ], + dim=0, + ) + .unsqueeze(0) + .to(self.device) + ) + state = self._normalize_state(self._extract_state(obs)) + input_ids, lang_lens = self._tokenize_inputs([self._extract_task(obs)], state) + return { + "pixel_values": pixel_values, + "input_ids": input_ids.to(self.device), + "lang_lens": lang_lens.to(self.device), + } + + def _extract_camera_tensor( + self, obs: dict[str, Any], camera_name: str + ) -> torch.Tensor: + image = self._extract_camera_image(obs, camera_name) + return self._image_to_raw_tensor(image) + + def _extract_camera_model_tensor( + self, obs: dict[str, Any], camera_name: str + ) -> torch.Tensor: + image = self._extract_camera_image(obs, camera_name) + return self._image_to_model_tensor(image) + + def _extract_camera_image( + self, obs: dict[str, Any], camera_name: str + ) -> np.ndarray: + if camera_name == "agentview": + return self._extract_image(obs, LIBERO_AGENTVIEW_KEYS) + if camera_name == "wrist": + return self._extract_image(obs, LIBERO_WRIST_KEYS) + if camera_name == "empty": + return np.zeros((self.image_size, self.image_size, 3), dtype=np.uint8) + raise ValueError(f"Unsupported camera_name={camera_name!r}.") + + @staticmethod + def _extract_image(obs: dict[str, Any], keys: tuple[str, ...]) -> np.ndarray: + candidates: list[Any] = [] + images = obs.get("images") + if isinstance(images, dict): + candidates.extend(images.get(k) for k in keys) + candidates.extend(obs.get(k) for k in keys) + for candidate in candidates: + if candidate is None: + continue + array = np.asarray(candidate) + if array.ndim == 4: + array = array[0] + if array.ndim != 3: + continue + if array.shape[0] == 3 and array.shape[-1] != 3: + array = np.transpose(array, (1, 2, 0)) + if array.shape[-1] == 3: + return array + raise KeyError(f"LIBERO observation does not contain any image keys: {keys}.") + + @staticmethod + def _image_to_raw_tensor(image: np.ndarray) -> torch.Tensor: + array = np.asarray(image, dtype=np.float32) + if array.max(initial=0.0) > 1.0: + array = array / 255.0 + return ( + torch.from_numpy(np.ascontiguousarray(array.transpose(2, 0, 1))) + .unsqueeze(0) + .contiguous() + ) + + def _image_to_model_tensor(self, image: np.ndarray) -> torch.Tensor: + tensor = self._image_to_raw_tensor(image) + if tensor.shape[-2:] != (self.image_size, self.image_size): + tensor = self._resize_with_pad(tensor, self.image_size, self.image_size) + return (tensor * 2.0 - 1.0).contiguous() + + @staticmethod + def _resize_with_pad(images: torch.Tensor, height: int, width: int) -> torch.Tensor: + _, _, cur_height, cur_width = images.shape + ratio = max(cur_width / width, cur_height / height) + resized_height = int(cur_height / ratio) + resized_width = int(cur_width / ratio) + resized = F.interpolate( + images, + size=(resized_height, resized_width), + mode="bilinear", + align_corners=False, + ) + resized = resized.clamp(0.0, 1.0) + pad_h0, rem_h = divmod(height - resized_height, 2) + pad_w0, rem_w = divmod(width - resized_width, 2) + return F.pad( + resized, + (pad_w0, pad_w0 + rem_w, pad_h0, pad_h0 + rem_h), + mode="constant", + value=0.0, + ) + + @staticmethod + def _extract_state(obs: dict[str, Any]) -> torch.Tensor: + state = obs.get("states", obs.get("state")) + if state is None: + raise KeyError("LIBERO observation must contain 'states' or 'state'.") + array = np.asarray(state, dtype=np.float32) + if array.ndim == 1: + array = array[None, :] + return torch.from_numpy(np.ascontiguousarray(array)) + + @staticmethod + def _extract_task(obs: dict[str, Any]) -> str: + task = obs.get("task_description", obs.get("task", "")) + if isinstance(task, (list, tuple)): + task = task[0] if task else "" + return str(task) + + def _normalize_state(self, state: torch.Tensor) -> torch.Tensor: + if self.normalization_mode == "openpi_quantile": + min_v = self._normalizer_stats.get("observation.state.min") + max_v = self._normalizer_stats.get("observation.state.max") + if min_v is None or max_v is None: + return state + return (state - min_v.to(state)) / ( + max_v.to(state) - min_v.to(state) + 1e-6 + ) * 2.0 - 1.0 + mean = self._normalizer_stats.get("observation.state.mean") + std = self._normalizer_stats.get("observation.state.std") + if mean is None or std is None: + return state + return (state - mean.to(state)) / torch.clamp(std.to(state), min=1e-8) + + def _tokenize_inputs( + self, tasks: list[str], states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.prompt_mode == "openpi_task": + prompts = [ + task.strip().replace("_", " ").replace("\n", " ") + "\n" + for task in tasks + ] + else: + state_np = states.detach().cpu().numpy() + bins = np.linspace(-1.0, 1.0, 257)[:-1] + discretized = np.digitize(state_np, bins=bins) - 1 + discretized = np.clip(discretized, 0, 255) + prompts = [] + for task, state_bins in zip(tasks, discretized): + cleaned = task.strip().replace("_", " ").replace("\n", " ") + state_str = " ".join(map(str, state_bins)) + prompts.append(f"Task: {cleaned}, State: {state_str};\nAction: ") + encoded = self.tokenizer( + prompts, + max_length=int(self.config.get("tokenizer_max_length", 200)), + padding="max_length", + padding_side="right", + truncation=True, + return_tensors="pt", + ) + return encoded["input_ids"].to(torch.int64), encoded["attention_mask"].sum( + dim=-1 + ).to(torch.int64) + + def _postprocess_actions(self, raw_actions: torch.Tensor) -> np.ndarray: + action = raw_actions[..., : self.action_dim].detach().float() + if not self._use_phyai_compat: + actions = self.processor.postprocess(action) + if isinstance(actions, torch.Tensor): + actions = actions.detach().cpu().numpy() + return np.asarray(actions, dtype=np.float32) + action = action.cpu() + if self.normalization_mode == "openpi_quantile": + min_v = self._unnormalizer_stats.get("action.min") + max_v = self._unnormalizer_stats.get("action.max") + if min_v is not None and max_v is not None: + action = (action + 1.0) / 2.0 * ( + max_v.to(action) - min_v.to(action) + 1e-6 + ) + min_v.to(action) + else: + mean = self._unnormalizer_stats.get("action.mean") + std = self._unnormalizer_stats.get("action.std") + if mean is not None and std is not None: + action = action * torch.clamp(std.to(action), min=1e-8) + mean.to( + action + ) + return action.numpy().astype(np.float32) + + def infer( + self, obs: dict[str, Any], *, noise: torch.Tensor | np.ndarray | None = None + ) -> dict[str, np.ndarray]: + request_kwargs = self.observation_to_request_inputs(obs) + if noise is not None: + request_kwargs["noise"] = torch.as_tensor(noise, device=self.device) + request = PI05Request(**request_kwargs) + with torch.inference_mode(): + raw_actions = self.engine.step(request) + actions = self._postprocess_actions(raw_actions) + return {"actions": actions} + + def close(self) -> None: + self.engine.close() From c1c08d0d1d60efd0063b1287747ad6d076405e57 Mon Sep 17 00:00:00 2001 From: rebecca26358 <153866717+rebecca26358@users.noreply.github.com> Date: Mon, 6 Jul 2026 06:54:40 +0000 Subject: [PATCH 02/29] benchmark: add external pi05 latency wrappers --- .../pi05/README_external_pi05_benchmarks.md | 129 ++++++++++ benchmark/pi05/bench_flashrt_pi05.py | 210 +++++++++++++++ benchmark/pi05/bench_realtime_vla_pi05.py | 209 +++++++++++++++ benchmark/pi05/bench_vlacpp_pi05_client.py | 241 ++++++++++++++++++ 4 files changed, 789 insertions(+) create mode 100644 benchmark/pi05/README_external_pi05_benchmarks.md create mode 100755 benchmark/pi05/bench_flashrt_pi05.py create mode 100755 benchmark/pi05/bench_realtime_vla_pi05.py create mode 100755 benchmark/pi05/bench_vlacpp_pi05_client.py diff --git a/benchmark/pi05/README_external_pi05_benchmarks.md b/benchmark/pi05/README_external_pi05_benchmarks.md new file mode 100644 index 0000000..c395838 --- /dev/null +++ b/benchmark/pi05/README_external_pi05_benchmarks.md @@ -0,0 +1,129 @@ +# External PI0.5 Benchmark Wrappers + +This directory contains wrappers for benchmarking external PI0.5 inference runtimes with the same high-level settings and common `bench_n_batch.py` runner used by the PhyAI PI0.5 benchmarks. + +These wrappers do **not** call PhyAI's engine. Each wrapper calls the target runtime directly, and all machine-local paths must be passed by command-line flags or environment variables. + +| File | Runtime measured | Main timing scope | +| --- | --- | --- | +| `bench_flashrt_pi05.py` | FlashRT | `flash_rt.load_model(...).predict(...)` public hot path after the first graph-building call | +| `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward()` hot path | +| `bench_vlacpp_pi05_client.py` | vla.cpp | ZMQ client request wall time; server phase timings are stored in JSONL `extras` | + +## Common Settings + +Use the same settings across machines when possible: + +| Setting | Recommended value | +| --- | --- | +| batch size | 1 via `--batch-sizes 1` | +| views | 2 real views | +| chunk size | 50 for strict comparison rows | +| warmup | 100 via `--n-warmup` | +| timed iterations | 100 via `--n-timed` | +| result format | JSONL via `--result-file` | +| prompt | fixed prompt or fixed prompt length, recorded in JSONL `extras` | +| precision | BF16 fair row; optimized rows must be labeled separately | + +Before running, verify the GPU is idle: + +```bash +nvidia-smi +uptime +``` + +The examples below use placeholders rather than host-specific paths: + +| Placeholder | Meaning | +| --- | --- | +| `` | This PhyAI checkout containing `benchmark/pi05/` | +| `` | FlashRT repository clone | +| `` | realtime-vla repository clone | +| `` | vla.cpp repository clone | +| `` | PI0.5 checkpoint directory or file for the selected runtime | +| `` | Local tokenizer directory or HF id used by vla.cpp client | +| `` | Local LIBERO `meta/stats.json` for PI0.5 state tokenization | + +## FlashRT + +`--flashrt-root` can be omitted if `FLASHRT_ROOT=` is set. + +```bash +cd +python benchmark/pi05/bench_flashrt_pi05.py \ + --flashrt-root \ + --checkpoint \ + --precision bf16 \ + --num-views 2 \ + --chunk-size 50 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/flashrt_pi05_external.jsonl +``` + +For FlashRT optimized precision, use `--precision fp8_bf16` and keep it in a separate table row. + +## realtime-vla + +`--realtime-vla-root` can be omitted if `REALTIME_VLA_ROOT=` is set. If `--checkpoint` points to a PI0.5 `model.safetensors` file or a directory containing `model.safetensors`, pass `--flashrt-root ` as well because the wrapper reuses FlashRT's PI0.5 safetensors conversion helper. + +```bash +cd +python benchmark/pi05/bench_realtime_vla_pi05.py \ + --realtime-vla-root \ + --flashrt-root \ + --checkpoint \ + --num-views 2 \ + --chunk-size 50 \ + --prompt-len 16 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/realtime_vla_pi05_external.jsonl +``` + +If you already have a realtime-vla `.pkl`, `.pt`, or `.pth` checkpoint, `--flashrt-root` is not needed. + +## vla.cpp + +Start `vla-server` first. The PI0.5 GGUF server reports server phase timing when started with `--timing-detail phase`. + +Example server: + +```bash +/build_sm120/vla-server \ + --bind tcp://127.0.0.1:5555 \ + --timing-detail phase \ + \ + +``` + +Then run the client benchmark. `--vlacpp-root` can be omitted if `VLACPP_ROOT=` is set. + +```bash +cd +python benchmark/pi05/bench_vlacpp_pi05_client.py \ + --vlacpp-root \ + --addr tcp://127.0.0.1:5555 \ + --arch pi05 \ + --tokenizer \ + --stats-json \ + --num-views 2 \ + --chunk-size 50 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/vlacpp_pi05_external.jsonl +``` + +For reproducibility, prefer a local tokenizer path for `--tokenizer`. The default PI0.5 tokenizer id, `google/paligemma-3b-pt-224`, may require HuggingFace access and authentication. Passing `--stats-json` avoids implicit network fetches for LIBERO state quantile statistics. + +## Notes + +- `bench_flashrt_pi05.py` and `bench_realtime_vla_pi05.py` reuse the PhyAI `NBatchBenchRunner` JSONL schema. +- `bench_flashrt_pi05.py` records wall latency with the runner's perf-counter path because FlashRT may run work on an internal CUDA stream; the step synchronizes CUDA before returning. +- `bench_realtime_vla_pi05.py` uses the runner's CUDA event timing around `Pi05Inference.forward()`. +- `bench_vlacpp_pi05_client.py` requires a running `vla-server`; the runner records client wall latency, and server phase latency is recorded in JSONL `extras.server_phase_latency_ms`. +- vla.cpp phase timing may expose `vision` and combined `inference`, not necessarily separate prefix/expert timing. +- These wrappers are intended for latency reproduction and support only `--batch-sizes 1` for now. Component MFU tables still require the shared PI0.5 FLOP model and component timings, as described in `thor_pi05_benchmark_plan.md`. diff --git a/benchmark/pi05/bench_flashrt_pi05.py b/benchmark/pi05/bench_flashrt_pi05.py new file mode 100755 index 0000000..d96b744 --- /dev/null +++ b/benchmark/pi05/bench_flashrt_pi05.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 +"""FlashRT PI0.5 latency benchmark using the PhyAI bench runner. + +This is a thin adapter around FlashRT's public ``flash_rt.load_model(...).predict`` +path. It mirrors ``benchmark/bench_n_batch_ws1_pi05.py``: framework-specific +setup lives here, while warmup, timed iterations, JSONL output, and optional +profiling are handled by ``benchmark/bench_n_batch.py``. + +Run:: + + python benchmark/pi05/bench_flashrt_pi05.py \ + --flashrt-root \ + --checkpoint \ + --batch-sizes 1 --n-warmup 100 --n-timed 100 \ + --result-file results/flashrt_pi05.jsonl + +Only batch size 1 is supported because the FlashRT PI0.5 public predict path +used here takes one robot request at a time. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import statistics +import sys +from typing import Any + +import numpy as np +import torch + +_BENCHMARK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_BENCHMARK_DIR)) +import bench_n_batch as bnb # noqa: E402 +from phyai.utils.profile import ( # noqa: E402 + add_profile_cli_args, + install_profiler, + profile_config_from_args, +) + + +def env_path(name: str) -> Path | None: + value = os.environ.get(name) + return Path(value) if value else None + + +def summarize(values: list[float]) -> dict[str, float] | None: + if not values: + return None + xs = sorted(float(x) for x in values) + return { + "count": len(xs), + "mean_ms": float(statistics.fmean(xs)), + "median_ms": float(statistics.median(xs)), + "p50_ms": float(np.percentile(xs, 50)), + "p90_ms": float(np.percentile(xs, 90)), + "p99_ms": float(np.percentile(xs, 99)), + "min_ms": xs[0], + "max_ms": xs[-1], + } + + +def make_images(num_views: int, seed: int) -> list[np.ndarray]: + """Deterministic synthetic HWC uint8 images for latency-only runs.""" + rng = np.random.default_rng(seed) + return [ + rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8) + for _ in range(num_views) + ] + + +def import_flashrt(repo: Path): + # Use the checked-out FlashRT repository directly; installation is optional. + sys.path.insert(0, str(repo)) + import flash_rt + + return flash_rt + + +def make_setup_fn(args: argparse.Namespace): + flash_rt = import_flashrt(args.flashrt_root) + + def setup_fn(batch_size: int) -> bnb.BenchSpec: + if batch_size != 1: + raise ValueError("FlashRT PI0.5 wrapper supports only batch_size=1") + + if args.precision == "bf16": + # FlashRT uses this environment switch to force the BF16 PI0.5 RTX path. + os.environ["FVK_PI05_RTX_FORCE_BF16"] = "1" + else: + os.environ.pop("FVK_PI05_RTX_FORCE_BF16", None) + + model = flash_rt.load_model( + args.checkpoint, + framework="torch", + config="pi05", + hardware=args.hardware, + num_views=args.num_views, + num_steps=args.chunk_size, + cache_frames=1, + use_fp8=(args.precision == "fp8_bf16"), + use_fp16=(args.precision == "bf16"), + ) + images = make_images(args.num_views, args.seed) + + # First call sets the prompt and lazily builds/calibrates the runtime graph. + model.predict(images, prompt=args.prompt) + torch.cuda.synchronize() + + internal_summary: dict[str, Any] = {} + call_count = 0 + + def step() -> None: + nonlocal call_count + if call_count == args.n_warmup and hasattr(model._pipe, "latency_records"): + model._pipe.latency_records.clear() + model.predict(images) + # FlashRT PI0.5 may use an internal CUDA stream, so synchronize inside the + # step to make the common runner's timing boundary conservative. + torch.cuda.synchronize() + call_count += 1 + if call_count > args.n_warmup and hasattr(model._pipe, "latency_records"): + stats = summarize([float(x) for x in model._pipe.latency_records]) + internal_summary.clear() + if stats is not None: + internal_summary.update(stats) + + spec = bnb.BenchSpec( + name="flashrt_pi05", + step_callable=step, + teardown_callable=lambda: None, + ) + # extras_fn receives this object before timed steps run; keeping a + # mutable dict lets the final BenchResult include FlashRT's own + # internal latency_records after the timed loop completes. + spec.flashrt_internal_latency_ms = internal_summary # type: ignore[attr-defined] + return spec + + return setup_fn + + +def make_extras_fn(args: argparse.Namespace): + def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: + return { + "runtime": "FlashRT", + "checkpoint": str(args.checkpoint), + "flashrt_root": str(args.flashrt_root), + "precision": args.precision, + "hardware": args.hardware, + "batch_size_contract": 1, + "num_views": args.num_views, + "chunk_size": args.chunk_size, + "prompt": args.prompt, + "seed": args.seed, + "internal_latency_ms": getattr(spec, "flashrt_internal_latency_ms", {}), + "timing_scope": "FlashRT predict hot path after first graph-building call; common runner uses perf-counter wall time because FlashRT may run work on a non-default CUDA stream", + } + + return extras_fn + + +def main() -> None: + flashrt_default = env_path("FLASHRT_ROOT") + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--flashrt-root", + type=Path, + default=flashrt_default, + required=flashrt_default is None, + help="Path to the FlashRT repository clone. Can also be set with FLASHRT_ROOT.", + ) + parser.add_argument( + "--checkpoint", + type=Path, + required=True, + help="PI0.5 checkpoint directory readable by FlashRT.", + ) + parser.add_argument("--num-views", type=int, default=2) + parser.add_argument("--chunk-size", type=int, default=50) + parser.add_argument("--prompt", default="do something") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--precision", choices=("bf16", "fp8_bf16"), default="bf16") + parser.add_argument("--hardware", default="auto") + + bnb.add_bench_cli_args(parser) + parser.set_defaults(bench_name="flashrt_pi05", batch_sizes=[1]) + add_profile_cli_args(parser) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for FlashRT PI0.5 benchmarking") + + install_profiler(profile_config_from_args(args)) + runner = bnb.NBatchBenchRunner( + setup_fn=make_setup_fn(args), + extras_fn=make_extras_fn(args), + # FlashRT may execute on an internal CUDA stream. Force the common + # runner's perf-counter path; step() synchronizes CUDA before returning. + device=torch.device("cpu"), + **bnb.bench_runner_kwargs_from_args(args), + ) + runner.run() + + +if __name__ == "__main__": + main() diff --git a/benchmark/pi05/bench_realtime_vla_pi05.py b/benchmark/pi05/bench_realtime_vla_pi05.py new file mode 100755 index 0000000..30561ac --- /dev/null +++ b/benchmark/pi05/bench_realtime_vla_pi05.py @@ -0,0 +1,209 @@ +#!/usr/bin/env python3 +"""realtime-vla PI0.5 latency benchmark using the PhyAI bench runner. + +This script is a thin adapter around realtime-vla's ``Pi05Inference.forward``. +It mirrors ``benchmark/bench_n_batch_ws1_pi05.py`` by delegating warmup, timed +iterations, JSONL output, and optional profiling to ``benchmark/bench_n_batch.py``. + +Run:: + + python benchmark/pi05/bench_realtime_vla_pi05.py \ + --realtime-vla-root \ + --flashrt-root \ + --checkpoint \ + --batch-sizes 1 --n-warmup 100 --n-timed 100 \ + --result-file results/realtime_vla_pi05.jsonl + +Only batch size 1 is supported because realtime-vla's PI0.5 inference object +used here is allocated for one synthetic request. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import pickle +import sys +from typing import Any + +import torch + +_BENCHMARK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_BENCHMARK_DIR)) +import bench_n_batch as bnb # noqa: E402 +from phyai.utils.profile import ( # noqa: E402 + add_profile_cli_args, + install_profiler, + profile_config_from_args, +) + + +def env_path(name: str) -> Path | None: + value = os.environ.get(name) + return Path(value) if value else None + + +def load_checkpoint(path: Path, flashrt_root: Path | None): + if path.is_dir(): + path = path / "model.safetensors" + if path.suffix == ".safetensors": + if flashrt_root is None: + raise ValueError( + "--flashrt-root or FLASHRT_ROOT is required when loading a safetensors checkpoint" + ) + # Reuse FlashRT's tested PI0.5 key conversion instead of duplicating it here. + sys.path.insert(0, str(flashrt_root)) + from flash_rt.frontends.torch.pi05_rtx import convert_pi05_safetensors + + return convert_pi05_safetensors(path) + if path.suffix in {".pt", ".pth"}: + return torch.load(path, map_location="cpu", weights_only=True) + with path.open("rb") as f: + return pickle.load(f) + + +def make_setup_fn(args: argparse.Namespace): + # Use the checked-out realtime-vla repository directly; installation is optional. + sys.path.insert(0, str(args.realtime_vla_root)) + from pi05_infer import Pi05Inference + + def setup_fn(batch_size: int) -> bnb.BenchSpec: + if batch_size != 1: + raise ValueError("realtime-vla PI0.5 wrapper supports only batch_size=1") + + checkpoint = load_checkpoint(args.checkpoint, args.flashrt_root) + if "language_embeds" not in checkpoint: + # Latency-only fallback for checkpoints that do not include prompt embeds. + checkpoint["language_embeds"] = torch.randn( + args.prompt_len, 2048, dtype=torch.bfloat16 + ) + + infer = Pi05Inference( + checkpoint=checkpoint, + num_views=args.num_views, + chunk_size=args.chunk_size, + tokenizer_path=str(args.tokenizer) if args.tokenizer else None, + discrete_state_input=args.discrete_state_input, + ) + torch.manual_seed(args.seed) + input_image = torch.randn( + args.num_views, 224, 224, 3, dtype=torch.bfloat16, device="cuda" + ) + input_noise = torch.randn( + args.chunk_size, 32, dtype=torch.bfloat16, device="cuda" + ) + + state_tokens = None + if args.discrete_state_input: + import numpy as np + + state_tokens = np.zeros(args.state_dim, dtype=np.int64) + + def step() -> None: + infer.forward( + input_image, + input_noise, + task_prompt=args.prompt, + state_tokens=state_tokens, + ) + + return bnb.BenchSpec( + name="realtime_vla_pi05", + step_callable=step, + teardown_callable=lambda: None, + ) + + return setup_fn + + +def make_extras_fn(args: argparse.Namespace): + def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: + return { + "runtime": "realtime-vla", + "checkpoint": str(args.checkpoint), + "realtime_vla_root": str(args.realtime_vla_root), + "flashrt_root": str(args.flashrt_root) if args.flashrt_root else None, + "precision": "bf16", + "batch_size_contract": 1, + "num_views": args.num_views, + "chunk_size": args.chunk_size, + "prompt": args.prompt, + "prompt_len": args.prompt_len, + "seed": args.seed, + "discrete_state_input": args.discrete_state_input, + "timing_scope": "Pi05Inference.forward hot path; common runner CUDA event wraps one forward call", + } + + return extras_fn + + +def main() -> None: + realtime_default = env_path("REALTIME_VLA_ROOT") + flashrt_default = env_path("FLASHRT_ROOT") + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--realtime-vla-root", + type=Path, + default=realtime_default, + required=realtime_default is None, + help="Path to the realtime-vla repository clone. Can also be set with REALTIME_VLA_ROOT.", + ) + parser.add_argument( + "--flashrt-root", + type=Path, + default=flashrt_default, + help="Path to FlashRT. Required only when --checkpoint is a PI0.5 safetensors checkpoint.", + ) + parser.add_argument( + "--checkpoint", + type=Path, + required=True, + help="realtime-vla .pkl/.pt checkpoint, PI0.5 model.safetensors, or a directory containing model.safetensors.", + ) + parser.add_argument("--num-views", type=int, default=2) + parser.add_argument("--chunk-size", type=int, default=50) + parser.add_argument("--prompt", default="do something") + parser.add_argument("--prompt-len", type=int, default=16) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--discrete-state-input", + action="store_true", + help="Use realtime-vla tokenizer/state-token prompt path instead of precomputed language_embeds.", + ) + parser.add_argument( + "--tokenizer", + type=Path, + default=None, + help="Tokenizer path for --discrete-state-input.", + ) + parser.add_argument( + "--state-dim", + type=int, + default=8, + help="Synthetic state token count for --discrete-state-input.", + ) + + bnb.add_bench_cli_args(parser) + parser.set_defaults(bench_name="realtime_vla_pi05", batch_sizes=[1]) + add_profile_cli_args(parser) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for realtime-vla PI0.5 benchmarking") + + install_profiler(profile_config_from_args(args)) + runner = bnb.NBatchBenchRunner( + setup_fn=make_setup_fn(args), + extras_fn=make_extras_fn(args), + device=torch.device("cuda", torch.cuda.current_device()), + **bnb.bench_runner_kwargs_from_args(args), + ) + runner.run() + + +if __name__ == "__main__": + main() diff --git a/benchmark/pi05/bench_vlacpp_pi05_client.py b/benchmark/pi05/bench_vlacpp_pi05_client.py new file mode 100755 index 0000000..12a537b --- /dev/null +++ b/benchmark/pi05/bench_vlacpp_pi05_client.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +"""vla.cpp PI0.5 ZMQ client benchmark using the PhyAI bench runner. + +Start ``vla-server`` separately, then run this client. The common PhyAI runner +handles warmup, timed iterations, JSONL output, and optional profiling. This +script forces CPU timing in the runner because the measured operation is a ZMQ +request to an external server process; CUDA events in the client process would +not cover server-side GPU work. + +Run:: + + python benchmark/pi05/bench_vlacpp_pi05_client.py \ + --vlacpp-root \ + --addr tcp://127.0.0.1:5555 \ + --tokenizer \ + --stats-json \ + --batch-sizes 1 --n-warmup 100 --n-timed 100 \ + --result-file results/vlacpp_pi05.jsonl + +Only batch size 1 is supported because the vla.cpp Python client sends one +request at a time. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import statistics +import sys +from typing import Any + +import numpy as np +import torch + +_BENCHMARK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_BENCHMARK_DIR)) +import bench_n_batch as bnb # noqa: E402 +from phyai.utils.profile import ( # noqa: E402 + add_profile_cli_args, + install_profiler, + profile_config_from_args, +) + + +def env_path(name: str) -> Path | None: + value = os.environ.get(name) + return Path(value) if value else None + + +def summarize(values: list[float]) -> dict[str, float] | None: + if not values: + return None + xs = sorted(float(x) for x in values) + return { + "count": len(xs), + "mean_ms": float(statistics.fmean(xs)), + "median_ms": float(statistics.median(xs)), + "p50_ms": float(np.percentile(xs, 50)), + "p90_ms": float(np.percentile(xs, 90)), + "p99_ms": float(np.percentile(xs, 99)), + "min_ms": xs[0], + "max_ms": xs[-1], + } + + +def make_obs(seed: int, num_views: int, prompt: str) -> dict[str, Any]: + """Create deterministic synthetic observations for latency-only requests. + + vla.cpp's official PI0.5 client expects CHW float32 images in [0, 1] and an + unnormalized robot state vector; it converts these into server protobufs. + """ + rng = np.random.default_rng(seed) + obs: dict[str, Any] = { + "observation.images.image": rng.random((3, 224, 224), dtype=np.float32), + "observation.state": rng.standard_normal(8).astype(np.float32), + "task": prompt, + "prompt": prompt, + } + if num_views >= 2: + obs["observation.images.image2"] = rng.random((3, 224, 224), dtype=np.float32) + if num_views >= 3: + obs["observation.images.image3"] = rng.random((3, 224, 224), dtype=np.float32) + return obs + + +def make_setup_fn(args: argparse.Namespace): + eval_root = args.vlacpp_root / "eval" + sys.path.insert(0, str(eval_root)) + sys.path.insert(0, str(eval_root / "client")) + from client.vla_cpp_client import VlaCppClient + + def setup_fn(batch_size: int) -> bnb.BenchSpec: + if batch_size != 1: + raise ValueError("vla.cpp PI0.5 wrapper supports only batch_size=1") + + image_keys = ["observation.images.image"] + if args.num_views >= 2: + image_keys.append("observation.images.image2") + if args.num_views >= 3: + image_keys.append("observation.images.image3") + + client = VlaCppClient( + vla_addr=args.addr, + arch=args.arch, + tokenizer_name=args.tokenizer, + image_keys=image_keys, + max_length=args.max_length, + real_action_dim=args.real_action_dim, + n_action_steps=1, + stats_json=args.stats_json, + ) + obs = make_obs(args.seed, args.num_views, args.prompt) + phase_samples: dict[str, list[float]] = { + "server_total_latency_ms": [], + "server_vision_latency_ms": [], + "server_inference_latency_ms": [], + "server_prefill_latency_ms": [], + "server_denoise_latency_ms": [], + } + phase_summary: dict[str, Any] = {} + call_count = 0 + + def update_phase_summary() -> None: + phase_summary.clear() + for key, values in phase_samples.items(): + phase_summary[key] = summarize(values) + + def step() -> None: + nonlocal call_count + client.get_action(obs) + call_count += 1 + if call_count <= args.n_warmup: + return + r = getattr(client, "_last_response", None) + if r is None: + return + phase_samples["server_total_latency_ms"].append(float(r.latency_ms_total)) + phase_samples["server_vision_latency_ms"].append(float(r.latency_ms_vision)) + phase_samples["server_inference_latency_ms"].append( + float(r.latency_ms_inference) + ) + phase_samples["server_prefill_latency_ms"].append( + float(r.latency_ms_prefill) + ) + phase_samples["server_denoise_latency_ms"].append( + float(r.latency_ms_denoise) + ) + update_phase_summary() + + def teardown() -> None: + sock = getattr(client, "sock", None) + if sock is not None: + sock.close(linger=0) + + spec = bnb.BenchSpec( + name="vlacpp_pi05_zmq_client", + step_callable=step, + teardown_callable=teardown, + ) + # Attach dynamic summary for extras_fn. The runner copies the dict after + # timed steps finish, so it records the final server phase statistics. + spec.vlacpp_phase_summary = phase_summary # type: ignore[attr-defined] + return spec + + return setup_fn + + +def make_extras_fn(args: argparse.Namespace): + def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: + return { + "runtime": "vla.cpp", + "vlacpp_root": str(args.vlacpp_root), + "addr": args.addr, + "arch": args.arch, + "tokenizer": args.tokenizer, + "stats_json": str(args.stats_json) if args.stats_json else None, + "batch_size_contract": 1, + "num_views": args.num_views, + "chunk_size_metadata": args.chunk_size, + "prompt": args.prompt, + "seed": args.seed, + "server_phase_latency_ms": getattr(spec, "vlacpp_phase_summary", {}), + "timing_scope": "client ZMQ request wall time; server phase timings are copied from PredictResponse extras", + "notes": "vla.cpp server enforces the GGUF chunk size; prefix/expert may be combined in server phase timing.", + } + + return extras_fn + + +def main() -> None: + vlacpp_default = env_path("VLACPP_ROOT") + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--vlacpp-root", + type=Path, + default=vlacpp_default, + required=vlacpp_default is None, + help="Path to the vla.cpp repository clone. Can also be set with VLACPP_ROOT.", + ) + parser.add_argument("--addr", default="tcp://127.0.0.1:5555") + parser.add_argument("--arch", default="pi05") + parser.add_argument( + "--tokenizer", + default="google/paligemma-3b-pt-224", + help="Tokenizer name or local tokenizer path. Prefer a local path if the HF repo is gated.", + ) + parser.add_argument( + "--stats-json", + type=Path, + default=None, + help="Local LIBERO meta/stats.json for arch=pi05; avoids network fetch.", + ) + parser.add_argument("--num-views", type=int, default=2) + parser.add_argument("--chunk-size", type=int, default=50) + parser.add_argument("--prompt", default="do something") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--max-length", type=int, default=200) + parser.add_argument("--real-action-dim", type=int, default=7) + + bnb.add_bench_cli_args(parser) + parser.set_defaults(bench_name="vlacpp_pi05", batch_sizes=[1]) + add_profile_cli_args(parser) + args = parser.parse_args() + + install_profiler(profile_config_from_args(args)) + runner = bnb.NBatchBenchRunner( + setup_fn=make_setup_fn(args), + extras_fn=make_extras_fn(args), + # Force perf-counter timing: the GPU work happens in the vla-server process. + device=torch.device("cpu"), + **bnb.bench_runner_kwargs_from_args(args), + ) + runner.run() + + +if __name__ == "__main__": + main() From 999fd1d8eb551e200d868fdb0e58c07bd2d9f4aa Mon Sep 17 00:00:00 2001 From: rebecca26358 <153866717+rebecca26358@users.noreply.github.com> Date: Mon, 6 Jul 2026 07:28:46 +0000 Subject: [PATCH 03/29] benchmark: address external pi05 wrapper review --- benchmark/pi05/bench_flashrt_pi05.py | 33 +++++++++++++++------- benchmark/pi05/bench_realtime_vla_pi05.py | 24 +++++++++++++--- benchmark/pi05/bench_vlacpp_pi05_client.py | 28 +++++++++++------- 3 files changed, 61 insertions(+), 24 deletions(-) diff --git a/benchmark/pi05/bench_flashrt_pi05.py b/benchmark/pi05/bench_flashrt_pi05.py index d96b744..1eb5f42 100755 --- a/benchmark/pi05/bench_flashrt_pi05.py +++ b/benchmark/pi05/bench_flashrt_pi05.py @@ -45,6 +45,22 @@ def env_path(name: str) -> Path | None: return Path(value) if value else None +class LazySummary(dict): + def __init__(self, values_fn_or_items): + if callable(values_fn_or_items): + super().__init__() + self._values_fn = values_fn_or_items + else: + super().__init__(values_fn_or_items) + self._values_fn = None + + def items(self): + if self._values_fn is None: + return super().items() + summary = summarize(self._values_fn()) + return (summary or {}).items() + + def summarize(values: list[float]) -> dict[str, float] | None: if not values: return None @@ -108,7 +124,6 @@ def setup_fn(batch_size: int) -> bnb.BenchSpec: model.predict(images, prompt=args.prompt) torch.cuda.synchronize() - internal_summary: dict[str, Any] = {} call_count = 0 def step() -> None: @@ -120,21 +135,19 @@ def step() -> None: # step to make the common runner's timing boundary conservative. torch.cuda.synchronize() call_count += 1 - if call_count > args.n_warmup and hasattr(model._pipe, "latency_records"): - stats = summarize([float(x) for x in model._pipe.latency_records]) - internal_summary.clear() - if stats is not None: - internal_summary.update(stats) spec = bnb.BenchSpec( name="flashrt_pi05", step_callable=step, teardown_callable=lambda: None, ) - # extras_fn receives this object before timed steps run; keeping a - # mutable dict lets the final BenchResult include FlashRT's own - # internal latency_records after the timed loop completes. - spec.flashrt_internal_latency_ms = internal_summary # type: ignore[attr-defined] + spec.flashrt_internal_latency_ms = LazySummary( + lambda: ( + [float(x) for x in model._pipe.latency_records] + if hasattr(model, "_pipe") and hasattr(model._pipe, "latency_records") + else [] + ) + ) # type: ignore[attr-defined] return spec return setup_fn diff --git a/benchmark/pi05/bench_realtime_vla_pi05.py b/benchmark/pi05/bench_realtime_vla_pi05.py index 30561ac..f91a6e4 100755 --- a/benchmark/pi05/bench_realtime_vla_pi05.py +++ b/benchmark/pi05/bench_realtime_vla_pi05.py @@ -44,7 +44,7 @@ def env_path(name: str) -> Path | None: return Path(value) if value else None -def load_checkpoint(path: Path, flashrt_root: Path | None): +def load_checkpoint(path: Path, flashrt_root: Path | None, trust_pickle: bool): if path.is_dir(): path = path / "model.safetensors" if path.suffix == ".safetensors": @@ -59,8 +59,17 @@ def load_checkpoint(path: Path, flashrt_root: Path | None): return convert_pi05_safetensors(path) if path.suffix in {".pt", ".pth"}: return torch.load(path, map_location="cpu", weights_only=True) - with path.open("rb") as f: - return pickle.load(f) + if path.suffix in {".pkl", ".pickle"}: + if not trust_pickle: + raise ValueError( + "Refusing to load pickle checkpoint without --trust-pickle-checkpoint. " + "Only use that flag for checkpoints from a trusted source." + ) + with path.open("rb") as f: + return pickle.load(f) # nosec B301: guarded by --trust-pickle-checkpoint. + raise ValueError( + f"Unsupported checkpoint suffix {path.suffix!r}; expected .safetensors, .pt, .pth, .pkl, or .pickle" + ) def make_setup_fn(args: argparse.Namespace): @@ -72,7 +81,9 @@ def setup_fn(batch_size: int) -> bnb.BenchSpec: if batch_size != 1: raise ValueError("realtime-vla PI0.5 wrapper supports only batch_size=1") - checkpoint = load_checkpoint(args.checkpoint, args.flashrt_root) + checkpoint = load_checkpoint( + args.checkpoint, args.flashrt_root, args.trust_pickle_checkpoint + ) if "language_embeds" not in checkpoint: # Latency-only fallback for checkpoints that do not include prompt embeds. checkpoint["language_embeds"] = torch.randn( @@ -164,6 +175,11 @@ def main() -> None: required=True, help="realtime-vla .pkl/.pt checkpoint, PI0.5 model.safetensors, or a directory containing model.safetensors.", ) + parser.add_argument( + "--trust-pickle-checkpoint", + action="store_true", + help="Allow loading .pkl/.pickle checkpoints. Only use with trusted checkpoint files.", + ) parser.add_argument("--num-views", type=int, default=2) parser.add_argument("--chunk-size", type=int, default=50) parser.add_argument("--prompt", default="do something") diff --git a/benchmark/pi05/bench_vlacpp_pi05_client.py b/benchmark/pi05/bench_vlacpp_pi05_client.py index 12a537b..6e07a67 100755 --- a/benchmark/pi05/bench_vlacpp_pi05_client.py +++ b/benchmark/pi05/bench_vlacpp_pi05_client.py @@ -48,6 +48,23 @@ def env_path(name: str) -> Path | None: return Path(value) if value else None +class LazySummaryMap(dict): + def __init__(self, samples_or_items): + if isinstance(samples_or_items, dict) and all( + isinstance(v, list) for v in samples_or_items.values() + ): + super().__init__() + self._samples = samples_or_items + else: + super().__init__(samples_or_items) + self._samples = None + + def items(self): + if self._samples is None: + return super().items() + return {key: summarize(values) for key, values in self._samples.items()}.items() + + def summarize(values: list[float]) -> dict[str, float] | None: if not values: return None @@ -118,14 +135,8 @@ def setup_fn(batch_size: int) -> bnb.BenchSpec: "server_prefill_latency_ms": [], "server_denoise_latency_ms": [], } - phase_summary: dict[str, Any] = {} call_count = 0 - def update_phase_summary() -> None: - phase_summary.clear() - for key, values in phase_samples.items(): - phase_summary[key] = summarize(values) - def step() -> None: nonlocal call_count client.get_action(obs) @@ -146,7 +157,6 @@ def step() -> None: phase_samples["server_denoise_latency_ms"].append( float(r.latency_ms_denoise) ) - update_phase_summary() def teardown() -> None: sock = getattr(client, "sock", None) @@ -158,9 +168,7 @@ def teardown() -> None: step_callable=step, teardown_callable=teardown, ) - # Attach dynamic summary for extras_fn. The runner copies the dict after - # timed steps finish, so it records the final server phase statistics. - spec.vlacpp_phase_summary = phase_summary # type: ignore[attr-defined] + spec.vlacpp_phase_summary = LazySummaryMap(phase_samples) # type: ignore[attr-defined] return spec return setup_fn From e45def609cb7c02f39ecd25d39b510fb2c8acbe7 Mon Sep 17 00:00:00 2001 From: rebecca26358 <153866717+rebecca26358@users.noreply.github.com> Date: Mon, 6 Jul 2026 07:44:24 +0000 Subject: [PATCH 04/29] benchmark: use direct FlashRT pi05 frontend --- .../pi05/README_external_pi05_benchmarks.md | 2 +- benchmark/pi05/bench_flashrt_pi05.py | 72 ++++++++++--------- 2 files changed, 40 insertions(+), 34 deletions(-) diff --git a/benchmark/pi05/README_external_pi05_benchmarks.md b/benchmark/pi05/README_external_pi05_benchmarks.md index c395838..76e0dde 100644 --- a/benchmark/pi05/README_external_pi05_benchmarks.md +++ b/benchmark/pi05/README_external_pi05_benchmarks.md @@ -6,7 +6,7 @@ These wrappers do **not** call PhyAI's engine. Each wrapper calls the target run | File | Runtime measured | Main timing scope | | --- | --- | --- | -| `bench_flashrt_pi05.py` | FlashRT | `flash_rt.load_model(...).predict(...)` public hot path after the first graph-building call | +| `bench_flashrt_pi05.py` | FlashRT | direct `Pi05TorchFrontendRtx.set_prompt(...); infer(...)` hot path after the first graph-building infer | | `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward()` hot path | | `bench_vlacpp_pi05_client.py` | vla.cpp | ZMQ client request wall time; server phase timings are stored in JSONL `extras` | diff --git a/benchmark/pi05/bench_flashrt_pi05.py b/benchmark/pi05/bench_flashrt_pi05.py index 1eb5f42..9d390a2 100755 --- a/benchmark/pi05/bench_flashrt_pi05.py +++ b/benchmark/pi05/bench_flashrt_pi05.py @@ -1,10 +1,12 @@ #!/usr/bin/env python3 """FlashRT PI0.5 latency benchmark using the PhyAI bench runner. -This is a thin adapter around FlashRT's public ``flash_rt.load_model(...).predict`` -path. It mirrors ``benchmark/bench_n_batch_ws1_pi05.py``: framework-specific -setup lives here, while warmup, timed iterations, JSONL output, and optional -profiling are handled by ``benchmark/bench_n_batch.py``. +This is a thin adapter around FlashRT's direct ``Pi05TorchFrontendRtx`` path. +The direct frontend is used because action chunk size is a frontend constructor +argument; ``flash_rt.load_model(..., num_steps=...)`` controls denoise steps, +not action chunk size. It mirrors ``benchmark/bench_n_batch_ws1_pi05.py``: +framework-specific setup lives here, while warmup, timed iterations, JSONL +output, and optional profiling are handled by ``benchmark/bench_n_batch.py``. Run:: @@ -14,7 +16,7 @@ --batch-sizes 1 --n-warmup 100 --n-timed 100 \ --result-file results/flashrt_pi05.jsonl -Only batch size 1 is supported because the FlashRT PI0.5 public predict path +Only batch size 1 is supported because the FlashRT PI0.5 direct frontend path used here takes one robot request at a time. """ @@ -77,25 +79,34 @@ def summarize(values: list[float]) -> dict[str, float] | None: } -def make_images(num_views: int, seed: int) -> list[np.ndarray]: - """Deterministic synthetic HWC uint8 images for latency-only runs.""" +def make_observation(num_views: int, seed: int, prompt: str) -> dict[str, Any]: + """Deterministic synthetic observation for latency-only runs.""" rng = np.random.default_rng(seed) - return [ - rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8) - for _ in range(num_views) - ] + obs: dict[str, Any] = { + "image": rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8), + "state": rng.standard_normal(8).astype(np.float32), + "task": prompt, + "prompt": prompt, + } + if num_views >= 2: + obs["wrist_image"] = rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8) + if num_views >= 3: + obs["wrist_image_right"] = rng.integers( + 0, 256, size=(224, 224, 3), dtype=np.uint8 + ) + return obs -def import_flashrt(repo: Path): +def import_flashrt_frontend(repo: Path): # Use the checked-out FlashRT repository directly; installation is optional. sys.path.insert(0, str(repo)) - import flash_rt + from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx - return flash_rt + return Pi05TorchFrontendRtx def make_setup_fn(args: argparse.Namespace): - flash_rt = import_flashrt(args.flashrt_root) + frontend_cls = import_flashrt_frontend(args.flashrt_root) def setup_fn(batch_size: int) -> bnb.BenchSpec: if batch_size != 1: @@ -107,30 +118,29 @@ def setup_fn(batch_size: int) -> bnb.BenchSpec: else: os.environ.pop("FVK_PI05_RTX_FORCE_BF16", None) - model = flash_rt.load_model( + model = frontend_cls( args.checkpoint, - framework="torch", - config="pi05", - hardware=args.hardware, num_views=args.num_views, - num_steps=args.chunk_size, + chunk_size=args.chunk_size, cache_frames=1, use_fp8=(args.precision == "fp8_bf16"), - use_fp16=(args.precision == "bf16"), + hardware=args.hardware, ) - images = make_images(args.num_views, args.seed) + obs = make_observation(args.num_views, args.seed, args.prompt) - # First call sets the prompt and lazily builds/calibrates the runtime graph. - model.predict(images, prompt=args.prompt) + # set_prompt builds the prompt-specific pipeline; calibration captures the graph. + model.set_prompt(args.prompt) + model.calibrate_with_real_data([obs]) + model.infer(obs) torch.cuda.synchronize() call_count = 0 def step() -> None: nonlocal call_count - if call_count == args.n_warmup and hasattr(model._pipe, "latency_records"): - model._pipe.latency_records.clear() - model.predict(images) + if call_count == args.n_warmup: + model.latency_records.clear() + model.infer(obs) # FlashRT PI0.5 may use an internal CUDA stream, so synchronize inside the # step to make the common runner's timing boundary conservative. torch.cuda.synchronize() @@ -142,11 +152,7 @@ def step() -> None: teardown_callable=lambda: None, ) spec.flashrt_internal_latency_ms = LazySummary( - lambda: ( - [float(x) for x in model._pipe.latency_records] - if hasattr(model, "_pipe") and hasattr(model._pipe, "latency_records") - else [] - ) + lambda: [float(x) for x in getattr(model, "latency_records", [])] ) # type: ignore[attr-defined] return spec @@ -167,7 +173,7 @@ def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: "prompt": args.prompt, "seed": args.seed, "internal_latency_ms": getattr(spec, "flashrt_internal_latency_ms", {}), - "timing_scope": "FlashRT predict hot path after first graph-building call; common runner uses perf-counter wall time because FlashRT may run work on a non-default CUDA stream", + "timing_scope": "FlashRT direct Pi05TorchFrontendRtx.infer hot path after set_prompt and first graph-building infer; common runner uses perf-counter wall time because FlashRT runs work on an internal CUDA stream", } return extras_fn From 4ef9efdb969eb9450aa87d6ee68d04197fec7925 Mon Sep 17 00:00:00 2001 From: rebecca26358 <153866717+rebecca26358@users.noreply.github.com> Date: Mon, 6 Jul 2026 09:25:27 +0000 Subject: [PATCH 05/29] docs: add pi05 external runtime setup notes --- .../pi05/README_external_pi05_benchmarks.md | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/benchmark/pi05/README_external_pi05_benchmarks.md b/benchmark/pi05/README_external_pi05_benchmarks.md index 76e0dde..88a9d32 100644 --- a/benchmark/pi05/README_external_pi05_benchmarks.md +++ b/benchmark/pi05/README_external_pi05_benchmarks.md @@ -44,6 +44,51 @@ The examples below use placeholders rather than host-specific paths: | `` | Local tokenizer directory or HF id used by vla.cpp client | | `` | Local LIBERO `meta/stats.json` for PI0.5 state tokenization | +## Environment Setup + +These scripts assume the target runtime can already be imported or executed. Keep each official repo at a known commit and record it with your results. + +FlashRT: + +```bash +git clone https://github.com/flashrt-project/FlashRT +cd +# Install/build FlashRT following its official README for your GPU/CUDA stack. +export FLASHRT_ROOT= +``` + +Notes: on RTX 5090/SM120, make sure the CUDA toolkit used to build extensions supports the GPU. For the BF16 fair row, the wrapper sets `FVK_PI05_RTX_FORCE_BF16=1`. Do not use `load_model(..., num_steps=50)` for chunk size; FlashRT `num_steps` means denoise steps. + +realtime-vla: + +```bash +git clone https://github.com/Dexmal/realtime-vla +cd +# Install realtime-vla following its official README. +export REALTIME_VLA_ROOT= +``` + +If you pass a PI0.5 `model.safetensors`, also set `FLASHRT_ROOT` because the wrapper reuses FlashRT's checkpoint conversion helper. Loading `.pkl/.pickle` checkpoints requires `--trust-pickle-checkpoint` and should only be used for trusted files. + +vla.cpp: + +```bash +git clone https://github.com/VinRobotics/vla.cpp +cd +# Build vla-server following its official README, for example into build_sm120/. +export VLACPP_ROOT= +``` + +Prepare the PI0.5 GGUF model, `mmproj` GGUF, tokenizer, and LIBERO `stats.json` before running. Prefer local tokenizer and stats paths to avoid network/auth issues. + +Quick smoke test after setup: + +```bash +python benchmark/pi05/bench_flashrt_pi05.py ... --n-warmup 1 --n-timed 1 +python benchmark/pi05/bench_realtime_vla_pi05.py ... --n-warmup 1 --n-timed 1 +python benchmark/pi05/bench_vlacpp_pi05_client.py ... --n-warmup 1 --n-timed 1 +``` + ## FlashRT `--flashrt-root` can be omitted if `FLASHRT_ROOT=` is set. From 96355cbd28845ab83c0a9ba7b0bcc0845c0b662e Mon Sep 17 00:00:00 2001 From: rebecca26358 <153866717+rebecca26358@users.noreply.github.com> Date: Mon, 6 Jul 2026 15:28:50 +0000 Subject: [PATCH 06/29] docs: align external pi05 benchmark docs --- .../pi05/README_external_pi05_benchmarks.md | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/benchmark/pi05/README_external_pi05_benchmarks.md b/benchmark/pi05/README_external_pi05_benchmarks.md index 88a9d32..fa08f53 100644 --- a/benchmark/pi05/README_external_pi05_benchmarks.md +++ b/benchmark/pi05/README_external_pi05_benchmarks.md @@ -1,4 +1,4 @@ -# External PI0.5 Benchmark Wrappers +# External PI0.5 benchmark wrappers This directory contains wrappers for benchmarking external PI0.5 inference runtimes with the same high-level settings and common `bench_n_batch.py` runner used by the PhyAI PI0.5 benchmarks. @@ -6,18 +6,18 @@ These wrappers do **not** call PhyAI's engine. Each wrapper calls the target run | File | Runtime measured | Main timing scope | | --- | --- | --- | -| `bench_flashrt_pi05.py` | FlashRT | direct `Pi05TorchFrontendRtx.set_prompt(...); infer(...)` hot path after the first graph-building infer | +| `bench_flashrt_pi05.py` | FlashRT | direct `Pi05TorchFrontendRtx.infer(...)` hot path after `set_prompt`, calibration, and the first graph-building infer | | `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward()` hot path | | `bench_vlacpp_pi05_client.py` | vla.cpp | ZMQ client request wall time; server phase timings are stored in JSONL `extras` | -## Common Settings +## Common settings Use the same settings across machines when possible: | Setting | Recommended value | | --- | --- | | batch size | 1 via `--batch-sizes 1` | -| views | 2 real views | +| views | 2 synthetic views / camera streams for latency-only runs | | chunk size | 50 for strict comparison rows | | warmup | 100 via `--n-warmup` | | timed iterations | 100 via `--n-timed` | @@ -44,7 +44,7 @@ The examples below use placeholders rather than host-specific paths: | `` | Local tokenizer directory or HF id used by vla.cpp client | | `` | Local LIBERO `meta/stats.json` for PI0.5 state tokenization | -## Environment Setup +## Environment setup These scripts assume the target runtime can already be imported or executed. Keep each official repo at a known commit and record it with your results. @@ -81,13 +81,7 @@ export VLACPP_ROOT= Prepare the PI0.5 GGUF model, `mmproj` GGUF, tokenizer, and LIBERO `stats.json` before running. Prefer local tokenizer and stats paths to avoid network/auth issues. -Quick smoke test after setup: - -```bash -python benchmark/pi05/bench_flashrt_pi05.py ... --n-warmup 1 --n-timed 1 -python benchmark/pi05/bench_realtime_vla_pi05.py ... --n-warmup 1 --n-timed 1 -python benchmark/pi05/bench_vlacpp_pi05_client.py ... --n-warmup 1 --n-timed 1 -``` +Quick smoke test after setup: use the runtime-specific command below, replace the real output path with a scratch file, and set `--n-warmup 1 --n-timed 1`. ## FlashRT From 6337ea020f275461848194c3ba23d1f5f728feb5 Mon Sep 17 00:00:00 2001 From: qyy Date: Sun, 19 Jul 2026 20:29:18 +0800 Subject: [PATCH 07/29] docs: add pi05 libero reproduction guide --- .../pi05/README_external_pi05_benchmarks.md | 168 ----- benchmark/pi05/bench_flashrt_pi05.py | 229 ------- benchmark/pi05/bench_realtime_vla_pi05.py | 225 ------- benchmark/pi05/bench_vlacpp_pi05_client.py | 249 -------- ...ai_pi05_libero_four_suites_reproduction.md | 574 ++++++++++++++++++ 5 files changed, 574 insertions(+), 871 deletions(-) delete mode 100644 benchmark/pi05/README_external_pi05_benchmarks.md delete mode 100755 benchmark/pi05/bench_flashrt_pi05.py delete mode 100755 benchmark/pi05/bench_realtime_vla_pi05.py delete mode 100755 benchmark/pi05/bench_vlacpp_pi05_client.py create mode 100644 docs/phyai_pi05_libero_four_suites_reproduction.md diff --git a/benchmark/pi05/README_external_pi05_benchmarks.md b/benchmark/pi05/README_external_pi05_benchmarks.md deleted file mode 100644 index fa08f53..0000000 --- a/benchmark/pi05/README_external_pi05_benchmarks.md +++ /dev/null @@ -1,168 +0,0 @@ -# External PI0.5 benchmark wrappers - -This directory contains wrappers for benchmarking external PI0.5 inference runtimes with the same high-level settings and common `bench_n_batch.py` runner used by the PhyAI PI0.5 benchmarks. - -These wrappers do **not** call PhyAI's engine. Each wrapper calls the target runtime directly, and all machine-local paths must be passed by command-line flags or environment variables. - -| File | Runtime measured | Main timing scope | -| --- | --- | --- | -| `bench_flashrt_pi05.py` | FlashRT | direct `Pi05TorchFrontendRtx.infer(...)` hot path after `set_prompt`, calibration, and the first graph-building infer | -| `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward()` hot path | -| `bench_vlacpp_pi05_client.py` | vla.cpp | ZMQ client request wall time; server phase timings are stored in JSONL `extras` | - -## Common settings - -Use the same settings across machines when possible: - -| Setting | Recommended value | -| --- | --- | -| batch size | 1 via `--batch-sizes 1` | -| views | 2 synthetic views / camera streams for latency-only runs | -| chunk size | 50 for strict comparison rows | -| warmup | 100 via `--n-warmup` | -| timed iterations | 100 via `--n-timed` | -| result format | JSONL via `--result-file` | -| prompt | fixed prompt or fixed prompt length, recorded in JSONL `extras` | -| precision | BF16 fair row; optimized rows must be labeled separately | - -Before running, verify the GPU is idle: - -```bash -nvidia-smi -uptime -``` - -The examples below use placeholders rather than host-specific paths: - -| Placeholder | Meaning | -| --- | --- | -| `` | This PhyAI checkout containing `benchmark/pi05/` | -| `` | FlashRT repository clone | -| `` | realtime-vla repository clone | -| `` | vla.cpp repository clone | -| `` | PI0.5 checkpoint directory or file for the selected runtime | -| `` | Local tokenizer directory or HF id used by vla.cpp client | -| `` | Local LIBERO `meta/stats.json` for PI0.5 state tokenization | - -## Environment setup - -These scripts assume the target runtime can already be imported or executed. Keep each official repo at a known commit and record it with your results. - -FlashRT: - -```bash -git clone https://github.com/flashrt-project/FlashRT -cd -# Install/build FlashRT following its official README for your GPU/CUDA stack. -export FLASHRT_ROOT= -``` - -Notes: on RTX 5090/SM120, make sure the CUDA toolkit used to build extensions supports the GPU. For the BF16 fair row, the wrapper sets `FVK_PI05_RTX_FORCE_BF16=1`. Do not use `load_model(..., num_steps=50)` for chunk size; FlashRT `num_steps` means denoise steps. - -realtime-vla: - -```bash -git clone https://github.com/Dexmal/realtime-vla -cd -# Install realtime-vla following its official README. -export REALTIME_VLA_ROOT= -``` - -If you pass a PI0.5 `model.safetensors`, also set `FLASHRT_ROOT` because the wrapper reuses FlashRT's checkpoint conversion helper. Loading `.pkl/.pickle` checkpoints requires `--trust-pickle-checkpoint` and should only be used for trusted files. - -vla.cpp: - -```bash -git clone https://github.com/VinRobotics/vla.cpp -cd -# Build vla-server following its official README, for example into build_sm120/. -export VLACPP_ROOT= -``` - -Prepare the PI0.5 GGUF model, `mmproj` GGUF, tokenizer, and LIBERO `stats.json` before running. Prefer local tokenizer and stats paths to avoid network/auth issues. - -Quick smoke test after setup: use the runtime-specific command below, replace the real output path with a scratch file, and set `--n-warmup 1 --n-timed 1`. - -## FlashRT - -`--flashrt-root` can be omitted if `FLASHRT_ROOT=` is set. - -```bash -cd -python benchmark/pi05/bench_flashrt_pi05.py \ - --flashrt-root \ - --checkpoint \ - --precision bf16 \ - --num-views 2 \ - --chunk-size 50 \ - --batch-sizes 1 \ - --n-warmup 100 \ - --n-timed 100 \ - --result-file results/flashrt_pi05_external.jsonl -``` - -For FlashRT optimized precision, use `--precision fp8_bf16` and keep it in a separate table row. - -## realtime-vla - -`--realtime-vla-root` can be omitted if `REALTIME_VLA_ROOT=` is set. If `--checkpoint` points to a PI0.5 `model.safetensors` file or a directory containing `model.safetensors`, pass `--flashrt-root ` as well because the wrapper reuses FlashRT's PI0.5 safetensors conversion helper. - -```bash -cd -python benchmark/pi05/bench_realtime_vla_pi05.py \ - --realtime-vla-root \ - --flashrt-root \ - --checkpoint \ - --num-views 2 \ - --chunk-size 50 \ - --prompt-len 16 \ - --batch-sizes 1 \ - --n-warmup 100 \ - --n-timed 100 \ - --result-file results/realtime_vla_pi05_external.jsonl -``` - -If you already have a realtime-vla `.pkl`, `.pt`, or `.pth` checkpoint, `--flashrt-root` is not needed. - -## vla.cpp - -Start `vla-server` first. The PI0.5 GGUF server reports server phase timing when started with `--timing-detail phase`. - -Example server: - -```bash -/build_sm120/vla-server \ - --bind tcp://127.0.0.1:5555 \ - --timing-detail phase \ - \ - -``` - -Then run the client benchmark. `--vlacpp-root` can be omitted if `VLACPP_ROOT=` is set. - -```bash -cd -python benchmark/pi05/bench_vlacpp_pi05_client.py \ - --vlacpp-root \ - --addr tcp://127.0.0.1:5555 \ - --arch pi05 \ - --tokenizer \ - --stats-json \ - --num-views 2 \ - --chunk-size 50 \ - --batch-sizes 1 \ - --n-warmup 100 \ - --n-timed 100 \ - --result-file results/vlacpp_pi05_external.jsonl -``` - -For reproducibility, prefer a local tokenizer path for `--tokenizer`. The default PI0.5 tokenizer id, `google/paligemma-3b-pt-224`, may require HuggingFace access and authentication. Passing `--stats-json` avoids implicit network fetches for LIBERO state quantile statistics. - -## Notes - -- `bench_flashrt_pi05.py` and `bench_realtime_vla_pi05.py` reuse the PhyAI `NBatchBenchRunner` JSONL schema. -- `bench_flashrt_pi05.py` records wall latency with the runner's perf-counter path because FlashRT may run work on an internal CUDA stream; the step synchronizes CUDA before returning. -- `bench_realtime_vla_pi05.py` uses the runner's CUDA event timing around `Pi05Inference.forward()`. -- `bench_vlacpp_pi05_client.py` requires a running `vla-server`; the runner records client wall latency, and server phase latency is recorded in JSONL `extras.server_phase_latency_ms`. -- vla.cpp phase timing may expose `vision` and combined `inference`, not necessarily separate prefix/expert timing. -- These wrappers are intended for latency reproduction and support only `--batch-sizes 1` for now. Component MFU tables still require the shared PI0.5 FLOP model and component timings, as described in `thor_pi05_benchmark_plan.md`. diff --git a/benchmark/pi05/bench_flashrt_pi05.py b/benchmark/pi05/bench_flashrt_pi05.py deleted file mode 100755 index 9d390a2..0000000 --- a/benchmark/pi05/bench_flashrt_pi05.py +++ /dev/null @@ -1,229 +0,0 @@ -#!/usr/bin/env python3 -"""FlashRT PI0.5 latency benchmark using the PhyAI bench runner. - -This is a thin adapter around FlashRT's direct ``Pi05TorchFrontendRtx`` path. -The direct frontend is used because action chunk size is a frontend constructor -argument; ``flash_rt.load_model(..., num_steps=...)`` controls denoise steps, -not action chunk size. It mirrors ``benchmark/bench_n_batch_ws1_pi05.py``: -framework-specific setup lives here, while warmup, timed iterations, JSONL -output, and optional profiling are handled by ``benchmark/bench_n_batch.py``. - -Run:: - - python benchmark/pi05/bench_flashrt_pi05.py \ - --flashrt-root \ - --checkpoint \ - --batch-sizes 1 --n-warmup 100 --n-timed 100 \ - --result-file results/flashrt_pi05.jsonl - -Only batch size 1 is supported because the FlashRT PI0.5 direct frontend path -used here takes one robot request at a time. -""" - -from __future__ import annotations - -import argparse -import os -from pathlib import Path -import statistics -import sys -from typing import Any - -import numpy as np -import torch - -_BENCHMARK_DIR = Path(__file__).resolve().parents[1] -sys.path.insert(0, str(_BENCHMARK_DIR)) -import bench_n_batch as bnb # noqa: E402 -from phyai.utils.profile import ( # noqa: E402 - add_profile_cli_args, - install_profiler, - profile_config_from_args, -) - - -def env_path(name: str) -> Path | None: - value = os.environ.get(name) - return Path(value) if value else None - - -class LazySummary(dict): - def __init__(self, values_fn_or_items): - if callable(values_fn_or_items): - super().__init__() - self._values_fn = values_fn_or_items - else: - super().__init__(values_fn_or_items) - self._values_fn = None - - def items(self): - if self._values_fn is None: - return super().items() - summary = summarize(self._values_fn()) - return (summary or {}).items() - - -def summarize(values: list[float]) -> dict[str, float] | None: - if not values: - return None - xs = sorted(float(x) for x in values) - return { - "count": len(xs), - "mean_ms": float(statistics.fmean(xs)), - "median_ms": float(statistics.median(xs)), - "p50_ms": float(np.percentile(xs, 50)), - "p90_ms": float(np.percentile(xs, 90)), - "p99_ms": float(np.percentile(xs, 99)), - "min_ms": xs[0], - "max_ms": xs[-1], - } - - -def make_observation(num_views: int, seed: int, prompt: str) -> dict[str, Any]: - """Deterministic synthetic observation for latency-only runs.""" - rng = np.random.default_rng(seed) - obs: dict[str, Any] = { - "image": rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8), - "state": rng.standard_normal(8).astype(np.float32), - "task": prompt, - "prompt": prompt, - } - if num_views >= 2: - obs["wrist_image"] = rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8) - if num_views >= 3: - obs["wrist_image_right"] = rng.integers( - 0, 256, size=(224, 224, 3), dtype=np.uint8 - ) - return obs - - -def import_flashrt_frontend(repo: Path): - # Use the checked-out FlashRT repository directly; installation is optional. - sys.path.insert(0, str(repo)) - from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx - - return Pi05TorchFrontendRtx - - -def make_setup_fn(args: argparse.Namespace): - frontend_cls = import_flashrt_frontend(args.flashrt_root) - - def setup_fn(batch_size: int) -> bnb.BenchSpec: - if batch_size != 1: - raise ValueError("FlashRT PI0.5 wrapper supports only batch_size=1") - - if args.precision == "bf16": - # FlashRT uses this environment switch to force the BF16 PI0.5 RTX path. - os.environ["FVK_PI05_RTX_FORCE_BF16"] = "1" - else: - os.environ.pop("FVK_PI05_RTX_FORCE_BF16", None) - - model = frontend_cls( - args.checkpoint, - num_views=args.num_views, - chunk_size=args.chunk_size, - cache_frames=1, - use_fp8=(args.precision == "fp8_bf16"), - hardware=args.hardware, - ) - obs = make_observation(args.num_views, args.seed, args.prompt) - - # set_prompt builds the prompt-specific pipeline; calibration captures the graph. - model.set_prompt(args.prompt) - model.calibrate_with_real_data([obs]) - model.infer(obs) - torch.cuda.synchronize() - - call_count = 0 - - def step() -> None: - nonlocal call_count - if call_count == args.n_warmup: - model.latency_records.clear() - model.infer(obs) - # FlashRT PI0.5 may use an internal CUDA stream, so synchronize inside the - # step to make the common runner's timing boundary conservative. - torch.cuda.synchronize() - call_count += 1 - - spec = bnb.BenchSpec( - name="flashrt_pi05", - step_callable=step, - teardown_callable=lambda: None, - ) - spec.flashrt_internal_latency_ms = LazySummary( - lambda: [float(x) for x in getattr(model, "latency_records", [])] - ) # type: ignore[attr-defined] - return spec - - return setup_fn - - -def make_extras_fn(args: argparse.Namespace): - def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: - return { - "runtime": "FlashRT", - "checkpoint": str(args.checkpoint), - "flashrt_root": str(args.flashrt_root), - "precision": args.precision, - "hardware": args.hardware, - "batch_size_contract": 1, - "num_views": args.num_views, - "chunk_size": args.chunk_size, - "prompt": args.prompt, - "seed": args.seed, - "internal_latency_ms": getattr(spec, "flashrt_internal_latency_ms", {}), - "timing_scope": "FlashRT direct Pi05TorchFrontendRtx.infer hot path after set_prompt and first graph-building infer; common runner uses perf-counter wall time because FlashRT runs work on an internal CUDA stream", - } - - return extras_fn - - -def main() -> None: - flashrt_default = env_path("FLASHRT_ROOT") - parser = argparse.ArgumentParser( - description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - parser.add_argument( - "--flashrt-root", - type=Path, - default=flashrt_default, - required=flashrt_default is None, - help="Path to the FlashRT repository clone. Can also be set with FLASHRT_ROOT.", - ) - parser.add_argument( - "--checkpoint", - type=Path, - required=True, - help="PI0.5 checkpoint directory readable by FlashRT.", - ) - parser.add_argument("--num-views", type=int, default=2) - parser.add_argument("--chunk-size", type=int, default=50) - parser.add_argument("--prompt", default="do something") - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--precision", choices=("bf16", "fp8_bf16"), default="bf16") - parser.add_argument("--hardware", default="auto") - - bnb.add_bench_cli_args(parser) - parser.set_defaults(bench_name="flashrt_pi05", batch_sizes=[1]) - add_profile_cli_args(parser) - args = parser.parse_args() - - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required for FlashRT PI0.5 benchmarking") - - install_profiler(profile_config_from_args(args)) - runner = bnb.NBatchBenchRunner( - setup_fn=make_setup_fn(args), - extras_fn=make_extras_fn(args), - # FlashRT may execute on an internal CUDA stream. Force the common - # runner's perf-counter path; step() synchronizes CUDA before returning. - device=torch.device("cpu"), - **bnb.bench_runner_kwargs_from_args(args), - ) - runner.run() - - -if __name__ == "__main__": - main() diff --git a/benchmark/pi05/bench_realtime_vla_pi05.py b/benchmark/pi05/bench_realtime_vla_pi05.py deleted file mode 100755 index f91a6e4..0000000 --- a/benchmark/pi05/bench_realtime_vla_pi05.py +++ /dev/null @@ -1,225 +0,0 @@ -#!/usr/bin/env python3 -"""realtime-vla PI0.5 latency benchmark using the PhyAI bench runner. - -This script is a thin adapter around realtime-vla's ``Pi05Inference.forward``. -It mirrors ``benchmark/bench_n_batch_ws1_pi05.py`` by delegating warmup, timed -iterations, JSONL output, and optional profiling to ``benchmark/bench_n_batch.py``. - -Run:: - - python benchmark/pi05/bench_realtime_vla_pi05.py \ - --realtime-vla-root \ - --flashrt-root \ - --checkpoint \ - --batch-sizes 1 --n-warmup 100 --n-timed 100 \ - --result-file results/realtime_vla_pi05.jsonl - -Only batch size 1 is supported because realtime-vla's PI0.5 inference object -used here is allocated for one synthetic request. -""" - -from __future__ import annotations - -import argparse -import os -from pathlib import Path -import pickle -import sys -from typing import Any - -import torch - -_BENCHMARK_DIR = Path(__file__).resolve().parents[1] -sys.path.insert(0, str(_BENCHMARK_DIR)) -import bench_n_batch as bnb # noqa: E402 -from phyai.utils.profile import ( # noqa: E402 - add_profile_cli_args, - install_profiler, - profile_config_from_args, -) - - -def env_path(name: str) -> Path | None: - value = os.environ.get(name) - return Path(value) if value else None - - -def load_checkpoint(path: Path, flashrt_root: Path | None, trust_pickle: bool): - if path.is_dir(): - path = path / "model.safetensors" - if path.suffix == ".safetensors": - if flashrt_root is None: - raise ValueError( - "--flashrt-root or FLASHRT_ROOT is required when loading a safetensors checkpoint" - ) - # Reuse FlashRT's tested PI0.5 key conversion instead of duplicating it here. - sys.path.insert(0, str(flashrt_root)) - from flash_rt.frontends.torch.pi05_rtx import convert_pi05_safetensors - - return convert_pi05_safetensors(path) - if path.suffix in {".pt", ".pth"}: - return torch.load(path, map_location="cpu", weights_only=True) - if path.suffix in {".pkl", ".pickle"}: - if not trust_pickle: - raise ValueError( - "Refusing to load pickle checkpoint without --trust-pickle-checkpoint. " - "Only use that flag for checkpoints from a trusted source." - ) - with path.open("rb") as f: - return pickle.load(f) # nosec B301: guarded by --trust-pickle-checkpoint. - raise ValueError( - f"Unsupported checkpoint suffix {path.suffix!r}; expected .safetensors, .pt, .pth, .pkl, or .pickle" - ) - - -def make_setup_fn(args: argparse.Namespace): - # Use the checked-out realtime-vla repository directly; installation is optional. - sys.path.insert(0, str(args.realtime_vla_root)) - from pi05_infer import Pi05Inference - - def setup_fn(batch_size: int) -> bnb.BenchSpec: - if batch_size != 1: - raise ValueError("realtime-vla PI0.5 wrapper supports only batch_size=1") - - checkpoint = load_checkpoint( - args.checkpoint, args.flashrt_root, args.trust_pickle_checkpoint - ) - if "language_embeds" not in checkpoint: - # Latency-only fallback for checkpoints that do not include prompt embeds. - checkpoint["language_embeds"] = torch.randn( - args.prompt_len, 2048, dtype=torch.bfloat16 - ) - - infer = Pi05Inference( - checkpoint=checkpoint, - num_views=args.num_views, - chunk_size=args.chunk_size, - tokenizer_path=str(args.tokenizer) if args.tokenizer else None, - discrete_state_input=args.discrete_state_input, - ) - torch.manual_seed(args.seed) - input_image = torch.randn( - args.num_views, 224, 224, 3, dtype=torch.bfloat16, device="cuda" - ) - input_noise = torch.randn( - args.chunk_size, 32, dtype=torch.bfloat16, device="cuda" - ) - - state_tokens = None - if args.discrete_state_input: - import numpy as np - - state_tokens = np.zeros(args.state_dim, dtype=np.int64) - - def step() -> None: - infer.forward( - input_image, - input_noise, - task_prompt=args.prompt, - state_tokens=state_tokens, - ) - - return bnb.BenchSpec( - name="realtime_vla_pi05", - step_callable=step, - teardown_callable=lambda: None, - ) - - return setup_fn - - -def make_extras_fn(args: argparse.Namespace): - def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: - return { - "runtime": "realtime-vla", - "checkpoint": str(args.checkpoint), - "realtime_vla_root": str(args.realtime_vla_root), - "flashrt_root": str(args.flashrt_root) if args.flashrt_root else None, - "precision": "bf16", - "batch_size_contract": 1, - "num_views": args.num_views, - "chunk_size": args.chunk_size, - "prompt": args.prompt, - "prompt_len": args.prompt_len, - "seed": args.seed, - "discrete_state_input": args.discrete_state_input, - "timing_scope": "Pi05Inference.forward hot path; common runner CUDA event wraps one forward call", - } - - return extras_fn - - -def main() -> None: - realtime_default = env_path("REALTIME_VLA_ROOT") - flashrt_default = env_path("FLASHRT_ROOT") - parser = argparse.ArgumentParser( - description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - parser.add_argument( - "--realtime-vla-root", - type=Path, - default=realtime_default, - required=realtime_default is None, - help="Path to the realtime-vla repository clone. Can also be set with REALTIME_VLA_ROOT.", - ) - parser.add_argument( - "--flashrt-root", - type=Path, - default=flashrt_default, - help="Path to FlashRT. Required only when --checkpoint is a PI0.5 safetensors checkpoint.", - ) - parser.add_argument( - "--checkpoint", - type=Path, - required=True, - help="realtime-vla .pkl/.pt checkpoint, PI0.5 model.safetensors, or a directory containing model.safetensors.", - ) - parser.add_argument( - "--trust-pickle-checkpoint", - action="store_true", - help="Allow loading .pkl/.pickle checkpoints. Only use with trusted checkpoint files.", - ) - parser.add_argument("--num-views", type=int, default=2) - parser.add_argument("--chunk-size", type=int, default=50) - parser.add_argument("--prompt", default="do something") - parser.add_argument("--prompt-len", type=int, default=16) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument( - "--discrete-state-input", - action="store_true", - help="Use realtime-vla tokenizer/state-token prompt path instead of precomputed language_embeds.", - ) - parser.add_argument( - "--tokenizer", - type=Path, - default=None, - help="Tokenizer path for --discrete-state-input.", - ) - parser.add_argument( - "--state-dim", - type=int, - default=8, - help="Synthetic state token count for --discrete-state-input.", - ) - - bnb.add_bench_cli_args(parser) - parser.set_defaults(bench_name="realtime_vla_pi05", batch_sizes=[1]) - add_profile_cli_args(parser) - args = parser.parse_args() - - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required for realtime-vla PI0.5 benchmarking") - - install_profiler(profile_config_from_args(args)) - runner = bnb.NBatchBenchRunner( - setup_fn=make_setup_fn(args), - extras_fn=make_extras_fn(args), - device=torch.device("cuda", torch.cuda.current_device()), - **bnb.bench_runner_kwargs_from_args(args), - ) - runner.run() - - -if __name__ == "__main__": - main() diff --git a/benchmark/pi05/bench_vlacpp_pi05_client.py b/benchmark/pi05/bench_vlacpp_pi05_client.py deleted file mode 100755 index 6e07a67..0000000 --- a/benchmark/pi05/bench_vlacpp_pi05_client.py +++ /dev/null @@ -1,249 +0,0 @@ -#!/usr/bin/env python3 -"""vla.cpp PI0.5 ZMQ client benchmark using the PhyAI bench runner. - -Start ``vla-server`` separately, then run this client. The common PhyAI runner -handles warmup, timed iterations, JSONL output, and optional profiling. This -script forces CPU timing in the runner because the measured operation is a ZMQ -request to an external server process; CUDA events in the client process would -not cover server-side GPU work. - -Run:: - - python benchmark/pi05/bench_vlacpp_pi05_client.py \ - --vlacpp-root \ - --addr tcp://127.0.0.1:5555 \ - --tokenizer \ - --stats-json \ - --batch-sizes 1 --n-warmup 100 --n-timed 100 \ - --result-file results/vlacpp_pi05.jsonl - -Only batch size 1 is supported because the vla.cpp Python client sends one -request at a time. -""" - -from __future__ import annotations - -import argparse -import os -from pathlib import Path -import statistics -import sys -from typing import Any - -import numpy as np -import torch - -_BENCHMARK_DIR = Path(__file__).resolve().parents[1] -sys.path.insert(0, str(_BENCHMARK_DIR)) -import bench_n_batch as bnb # noqa: E402 -from phyai.utils.profile import ( # noqa: E402 - add_profile_cli_args, - install_profiler, - profile_config_from_args, -) - - -def env_path(name: str) -> Path | None: - value = os.environ.get(name) - return Path(value) if value else None - - -class LazySummaryMap(dict): - def __init__(self, samples_or_items): - if isinstance(samples_or_items, dict) and all( - isinstance(v, list) for v in samples_or_items.values() - ): - super().__init__() - self._samples = samples_or_items - else: - super().__init__(samples_or_items) - self._samples = None - - def items(self): - if self._samples is None: - return super().items() - return {key: summarize(values) for key, values in self._samples.items()}.items() - - -def summarize(values: list[float]) -> dict[str, float] | None: - if not values: - return None - xs = sorted(float(x) for x in values) - return { - "count": len(xs), - "mean_ms": float(statistics.fmean(xs)), - "median_ms": float(statistics.median(xs)), - "p50_ms": float(np.percentile(xs, 50)), - "p90_ms": float(np.percentile(xs, 90)), - "p99_ms": float(np.percentile(xs, 99)), - "min_ms": xs[0], - "max_ms": xs[-1], - } - - -def make_obs(seed: int, num_views: int, prompt: str) -> dict[str, Any]: - """Create deterministic synthetic observations for latency-only requests. - - vla.cpp's official PI0.5 client expects CHW float32 images in [0, 1] and an - unnormalized robot state vector; it converts these into server protobufs. - """ - rng = np.random.default_rng(seed) - obs: dict[str, Any] = { - "observation.images.image": rng.random((3, 224, 224), dtype=np.float32), - "observation.state": rng.standard_normal(8).astype(np.float32), - "task": prompt, - "prompt": prompt, - } - if num_views >= 2: - obs["observation.images.image2"] = rng.random((3, 224, 224), dtype=np.float32) - if num_views >= 3: - obs["observation.images.image3"] = rng.random((3, 224, 224), dtype=np.float32) - return obs - - -def make_setup_fn(args: argparse.Namespace): - eval_root = args.vlacpp_root / "eval" - sys.path.insert(0, str(eval_root)) - sys.path.insert(0, str(eval_root / "client")) - from client.vla_cpp_client import VlaCppClient - - def setup_fn(batch_size: int) -> bnb.BenchSpec: - if batch_size != 1: - raise ValueError("vla.cpp PI0.5 wrapper supports only batch_size=1") - - image_keys = ["observation.images.image"] - if args.num_views >= 2: - image_keys.append("observation.images.image2") - if args.num_views >= 3: - image_keys.append("observation.images.image3") - - client = VlaCppClient( - vla_addr=args.addr, - arch=args.arch, - tokenizer_name=args.tokenizer, - image_keys=image_keys, - max_length=args.max_length, - real_action_dim=args.real_action_dim, - n_action_steps=1, - stats_json=args.stats_json, - ) - obs = make_obs(args.seed, args.num_views, args.prompt) - phase_samples: dict[str, list[float]] = { - "server_total_latency_ms": [], - "server_vision_latency_ms": [], - "server_inference_latency_ms": [], - "server_prefill_latency_ms": [], - "server_denoise_latency_ms": [], - } - call_count = 0 - - def step() -> None: - nonlocal call_count - client.get_action(obs) - call_count += 1 - if call_count <= args.n_warmup: - return - r = getattr(client, "_last_response", None) - if r is None: - return - phase_samples["server_total_latency_ms"].append(float(r.latency_ms_total)) - phase_samples["server_vision_latency_ms"].append(float(r.latency_ms_vision)) - phase_samples["server_inference_latency_ms"].append( - float(r.latency_ms_inference) - ) - phase_samples["server_prefill_latency_ms"].append( - float(r.latency_ms_prefill) - ) - phase_samples["server_denoise_latency_ms"].append( - float(r.latency_ms_denoise) - ) - - def teardown() -> None: - sock = getattr(client, "sock", None) - if sock is not None: - sock.close(linger=0) - - spec = bnb.BenchSpec( - name="vlacpp_pi05_zmq_client", - step_callable=step, - teardown_callable=teardown, - ) - spec.vlacpp_phase_summary = LazySummaryMap(phase_samples) # type: ignore[attr-defined] - return spec - - return setup_fn - - -def make_extras_fn(args: argparse.Namespace): - def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: - return { - "runtime": "vla.cpp", - "vlacpp_root": str(args.vlacpp_root), - "addr": args.addr, - "arch": args.arch, - "tokenizer": args.tokenizer, - "stats_json": str(args.stats_json) if args.stats_json else None, - "batch_size_contract": 1, - "num_views": args.num_views, - "chunk_size_metadata": args.chunk_size, - "prompt": args.prompt, - "seed": args.seed, - "server_phase_latency_ms": getattr(spec, "vlacpp_phase_summary", {}), - "timing_scope": "client ZMQ request wall time; server phase timings are copied from PredictResponse extras", - "notes": "vla.cpp server enforces the GGUF chunk size; prefix/expert may be combined in server phase timing.", - } - - return extras_fn - - -def main() -> None: - vlacpp_default = env_path("VLACPP_ROOT") - parser = argparse.ArgumentParser( - description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - parser.add_argument( - "--vlacpp-root", - type=Path, - default=vlacpp_default, - required=vlacpp_default is None, - help="Path to the vla.cpp repository clone. Can also be set with VLACPP_ROOT.", - ) - parser.add_argument("--addr", default="tcp://127.0.0.1:5555") - parser.add_argument("--arch", default="pi05") - parser.add_argument( - "--tokenizer", - default="google/paligemma-3b-pt-224", - help="Tokenizer name or local tokenizer path. Prefer a local path if the HF repo is gated.", - ) - parser.add_argument( - "--stats-json", - type=Path, - default=None, - help="Local LIBERO meta/stats.json for arch=pi05; avoids network fetch.", - ) - parser.add_argument("--num-views", type=int, default=2) - parser.add_argument("--chunk-size", type=int, default=50) - parser.add_argument("--prompt", default="do something") - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--max-length", type=int, default=200) - parser.add_argument("--real-action-dim", type=int, default=7) - - bnb.add_bench_cli_args(parser) - parser.set_defaults(bench_name="vlacpp_pi05", batch_sizes=[1]) - add_profile_cli_args(parser) - args = parser.parse_args() - - install_profiler(profile_config_from_args(args)) - runner = bnb.NBatchBenchRunner( - setup_fn=make_setup_fn(args), - extras_fn=make_extras_fn(args), - # Force perf-counter timing: the GPU work happens in the vla-server process. - device=torch.device("cpu"), - **bnb.bench_runner_kwargs_from_args(args), - ) - runner.run() - - -if __name__ == "__main__": - main() diff --git a/docs/phyai_pi05_libero_four_suites_reproduction.md b/docs/phyai_pi05_libero_four_suites_reproduction.md new file mode 100644 index 0000000..e25b92d --- /dev/null +++ b/docs/phyai_pi05_libero_four_suites_reproduction.md @@ -0,0 +1,574 @@ +# PhyAI pi0.5 跑 LIBERO 四组件完整复现文档 + +本文档说明如何在一台新机器上使用 PhyAI 推理 pi0.5，并通过 vla-evaluation-harness 跑完整 LIBERO 四组件 benchmark。文档不依赖当前机器的本地路径，所有路径都用环境变量表示。 + +四组件指： + +```text +libero_spatial -> configs/benchmarks/libero/spatial.yaml +libero_object -> configs/benchmarks/libero/object.yaml +libero_goal -> configs/benchmarks/libero/goal.yaml +libero_10 -> configs/benchmarks/libero/10.yaml +``` + +完整口径： + +```text +模式：sync +chunk_size：10 +每个组件：10 tasks x 50 episodes = 500 episodes +总量：4 suites x 500 episodes = 2000 episodes +模型：PhyAI pi0.5 LIBERO converted checkpoint +仿真：vla-evaluation-harness LIBERO Docker 容器 +结果：每个 suite 一个 JSON，记录 success、steps、timing、chunk size +``` + +## 1. 机器与资源要求 + +建议机器： + +```text +GPU：至少 1 张 CUDA GPU，显存建议 >= 48GB +系统：Linux +容器：Docker + NVIDIA Container Toolkit +Python 环境管理：uv +辅助工具：tmux, nvidia-smi, ss +``` + +必须准备的模型资源： + +```text +PhyAI converted checkpoint：pi05_libero_phyai_converted +PaLI-Gemma tokenizer / processor：paligemma-3b-pt-224 +``` + +`paligemma-3b-pt-224` 是 gated 资源，推荐从已有机器同步，不建议在复现时现场下载。 + +## 2. 推荐环境变量 + +根据目标机器实际路径设置： + +```bash +export PHYAI_ROOT=$HOME/phyai +export VLA_ROOT=$HOME/vla-evaluation-harness +export MODEL_ROOT=$HOME/phyai_models +export PHYAI_CONTAINER=phyai_libero_eval + +export PHYAI_CKPT_HOST=$MODEL_ROOT/pi05_libero_phyai_converted +export TOKENIZER_HOST=$MODEL_ROOT/paligemma-3b-pt-224 + +export PHYAI_CKPT_IN_CONTAINER=/data/share/pi05_libero_phyai_converted +export TOKENIZER_IN_CONTAINER=/data/share/paligemma-3b-pt-224 + +export LIBERO_IMAGE=ghcr.io/allenai/vla-evaluation-harness/libero:latest +``` + +确认资源存在： + +```bash +test -d "$PHYAI_CKPT_HOST" +test -d "$TOKENIZER_HOST" +``` + +如果模型在另一台机器上，示例同步命令如下： + +```bash +rsync -azP \ + /path/to/pi05_libero_phyai_converted \ + /path/to/paligemma-3b-pt-224 \ + user@target-host:$MODEL_ROOT/ +``` + +## 3. 获取代码并创建环境 + +clone PhyAI： + +```bash +git clone https://github.com/MEmbodied/phyai.git "$PHYAI_ROOT" +cd "$PHYAI_ROOT" +uv sync +``` + +clone vla-evaluation-harness： + +```bash +git clone https://github.com/allenai/vla-evaluation-harness.git "$VLA_ROOT" +cd "$VLA_ROOT" +uv sync +./.venv/bin/vla-eval --help >/tmp/vla_eval_help.log +``` + +如果目标机器不能联网，可以在能联网机器上 clone 后用 `rsync` 同步两个仓库；同步后仍建议在目标机器上执行 `uv sync`，让本机 Python、CUDA、依赖和 editable path 正确。 + +## 4. 准备 LIBERO Docker 镜像 + +vla-harness 跑 LIBERO 时会启动 LIBERO benchmark 容器。x86_64 机器可以直接使用官方镜像： + +```bash +docker pull "$LIBERO_IMAGE" +``` + +如果目标机器是 ARM64，例如 Thor，而官方镜像只有 amd64，需要在目标机器本地构建 ARM64 镜像： + +```bash +cd "$VLA_ROOT" +export DOCKER_DEFAULT_PLATFORM=linux/arm64 +docker/build.sh libero + +docker image inspect "$LIBERO_IMAGE" \ + --format '{{.Architecture}} {{.Os}}' +``` + +期望输出： + +```text +arm64 linux +``` + +如果是 x86_64，期望输出一般是： + +```text +amd64 linux +``` + +## 5. 创建 PhyAI Docker 容器 + +建议在 Docker 中运行 PhyAI server，容器挂载 PhyAI 代码、vla-harness 代码和模型目录： + +```bash +docker run -dit --gpus all \ + -v "$PHYAI_ROOT":/phyai_workspace \ + -v "$VLA_ROOT":/vla-evaluation-harness \ + -v "$MODEL_ROOT":/data/share \ + -w /phyai_workspace \ + --cap-add=SYS_ADMIN \ + --ipc=host \ + --cap-add=SYS_PTRACE \ + --shm-size=4G \ + --security-opt seccomp=unconfined \ + --security-opt apparmor=unconfined \ + --name "$PHYAI_CONTAINER" \ + nvcr.io/nvidia/pytorch:25.12-py3 bash +``` + +容器内同步 PhyAI 环境： + +```bash +docker exec "$PHYAI_CONTAINER" bash -lc ' +cd /phyai_workspace +python3 -m pip install -U uv +uv sync +' +``` + +如果 `uv sync` 生成的 editable path 与容器路径不一致，可加兼容 symlink。只有在 import 报错指向旧宿主路径时才需要执行；`HOST_USER` 填目标机器上的用户名： + +```bash +export HOST_USER=$(id -un) +export COMPAT_PARENT=/compat_mount + +docker exec "$PHYAI_CONTAINER" bash -lc " +mkdir -p $COMPAT_PARENT/$HOST_USER +ln -sfn /phyai_workspace $COMPAT_PARENT/$HOST_USER/phyai +" +``` + +确认容器内 import： + +```bash +docker exec "$PHYAI_CONTAINER" bash -lc ' +cd /phyai_workspace +export PYTHONPATH=/phyai_workspace/phyai/src:/phyai_workspace/phyai-kernel:/phyai_workspace/phyai-utils-tools/src:/vla-evaluation-harness/src +/phyai_workspace/.venv/bin/python - <&1 | tee $VLA_ROOT/results/phyai_pi05_libero_server.log +" +``` + +关键配置含义： + +| 配置 | 值 | 说明 | +| --- | --- | --- | +| `--checkpoint_path` | `/data/share/pi05_libero_phyai_converted` | PhyAI converted pi0.5 LIBERO 权重 | +| `PHYAI_TOKENIZER_PATH` | `/data/share/paligemma-3b-pt-224` | tokenizer / processor 文件 | +| `PHYAI_CAMERA_MODE` | `two_camera` | LIBERO 发送 agentview 与 wrist 两路图像 | +| `--params_dtype` | `bfloat16` | 参数 dtype | +| `--attn_backend` | `flashinfer` | attention 后端 | +| `--norm_backend` | `phyai-kernel` | norm 后端 | +| `--linear_backend` | `flashinfer` | linear 后端 | +| `--flashinfer_workspace_bytes` | `536870912` | 512MiB workspace | +| `--chunk_size` | `10` | 每次推理产出 10 个 action | +| CUDA graph | 默认开启 | 不要传 `--no-use_cuda_graph` | + +等待 ready： + +```bash +tail -f "$VLA_ROOT/results/phyai_pi05_libero_server.log" +``` + +必须看到： + +```text +capturing vision-tower CUDA graph +capturing 4 prefix-forward CUDA graph(s) +capturing the full 10-step Euler loop as one CUDA graph +Starting server on ws://0.0.0.0:8000 +``` + +## 8. 创建 smoke 配置并验证 + +先跑一个最小 smoke，确认模型、LIBERO Docker、WebSocket、timing 字段都可用。 + +```bash +cat > "$VLA_ROOT/configs/benchmarks/libero/smoke_test_phyai_local.yaml" < "$VLA_ROOT/run_phyai_pi05_libero_four_suites.sh" <<'SH' +#!/usr/bin/env bash +set -euo pipefail + +: "${PHYAI_SERVER_URL:?must set PHYAI_SERVER_URL}" +: "${PHYAI_CKPT_IN_CONTAINER:?must set PHYAI_CKPT_IN_CONTAINER}" +: "${LIBERO_IMAGE:=ghcr.io/allenai/vla-evaluation-harness/libero:latest}" + +cd "$(dirname "$0")" + +RUN_ID="phyai_pi05_libero_four_$(date +%Y%m%d_%H%M%S)" +OUT="results/${RUN_ID}" +mkdir -p "$OUT/configs" + +{ + echo "RUN_ID=${RUN_ID}" + echo "START=$(date -Is)" + echo "MODEL=phyai_pi05" + echo "SERVER_URL=${PHYAI_SERVER_URL}" + echo "CHECKPOINT=${PHYAI_CKPT_IN_CONTAINER}" + echo "PHYAI_CAMERA_MODE=two_camera" + echo "MODE=sync" + echo "CHUNK_SIZE=10" + echo "SERVER_CONFIG=use_cuda_graph=True attn=flashinfer norm=phyai-kernel linear=flashinfer workspace=536870912 params_dtype=bfloat16" +} | tee "$OUT/run_summary.log" + +make_cfg() { + local suite_name="$1" + local suite="$2" + local cfg="$3" + cat > "$cfg" <&1 | tee "$log" + status=${PIPESTATUS[0]} + + echo "SUITE_END model=phyai suite=${name} status=${status} log=${log} time=$(date -Is)" | tee -a "$OUT/run_summary.log" + + result_json=$(ls -t "$OUT/${suite}_sync_"*.json "$OUT"/*"${suite}"*_sync_*.json 2>/dev/null | head -1 || true) + if [ -n "$result_json" ]; then + ./.venv/bin/python scripts/summarize_timing.py "$result_json" | sed "s/^/TIMING phyai_${name} /" | tee -a "$OUT/run_summary.log" + else + echo "WARN no result json found for ${name}" | tee -a "$OUT/run_summary.log" + fi + + if [ "$status" -ne 0 ]; then + exit "$status" + fi +done + +echo "ALL_DONE $(date -Is)" | tee -a "$OUT/run_summary.log" +SH +chmod +x "$VLA_ROOT/run_phyai_pi05_libero_four_suites.sh" +``` + +启动长跑： + +```bash +cd "$VLA_ROOT" +tmux new-session -d -s phyai_pi05_libero_four \ + "PHYAI_SERVER_URL=$PHYAI_SERVER_URL PHYAI_CKPT_IN_CONTAINER=$PHYAI_CKPT_IN_CONTAINER LIBERO_IMAGE=$LIBERO_IMAGE ./run_phyai_pi05_libero_four_suites.sh" +``` + +查看进度： + +```bash +tmux ls +cd "$VLA_ROOT" +latest=$(ls -td results/phyai_pi05_libero_four_* | head -1) +sed -n '1,220p' "$latest/run_summary.log" +tail -80 "$latest"/phyai_spatial.log +``` + +## 10. 解析成功率与 timing + +长跑完成后： + +```bash +cd "$VLA_ROOT" +latest=$(ls -td results/phyai_pi05_libero_four_* | head -1) +cat "$latest/run_summary.log" +``` + +用 `summarize_timing.py` 重新汇总全部 JSON： + +```bash +cd "$VLA_ROOT" +latest=$(ls -td results/phyai_pi05_libero_four_* | head -1) +./.venv/bin/python scripts/summarize_timing.py "$latest"/*.json +``` + +统计成功率： + +```bash +cd "$VLA_ROOT" +latest=$(ls -td results/phyai_pi05_libero_four_* | head -1) +./.venv/bin/python - <<'PY' "$latest"/*.json +import json +import sys +from pathlib import Path + +for arg in sys.argv[1:]: + p = Path(arg) + data = json.loads(p.read_text()) + eps = [ep for task in data.get("tasks", []) for ep in task.get("episodes", [])] + succ = sum(1 for ep in eps if ep.get("metrics", {}).get("success")) + total = len(eps) + rate = succ / total * 100.0 if total else 0.0 + steps = sum(int(ep.get("steps", 0)) for ep in eps) + print(f"{p.name}: success={succ}/{total} rate={rate:.1f}% steps={steps} mean_success={data.get('mean_success')}") +PY +``` + +必须记录的字段： + +```text +RUN_ID +结果目录 +checkpoint +server_url +suite +success / total +success rate +steps +/usr/bin/time -p real +model_wait_sec +model_inference_sec +env_step_sec +obs_sec +avg_model_inference_ms +model_inference_calls +model_buffer_hits +raw_chunk_size_max +served_chunk_size_max +``` + +`raw_chunk_size_max=10` 且 `served_chunk_size_max=10` 是四组件复现的关键校验项。 + +## 11. 预期参考结果 + +不同机器、GPU、驱动、负载会影响耗时；成功率也可能有少量随机波动。此前同口径参考结果如下： + +| 组件 | 成功率 | 成功数 | steps | 总耗时 | 模型纯推理时间 | env step 时间 | benchmark 等待 action 时间 | 平均单次模型推理 | 推理调用次数 | buffer hit | chunk 验证 | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | +| `libero_spatial` | 97.8% | 489/500 | 52,926 | 4,198.73s | 200.56s | 1,239.61s | 2,269.16s | 36.33ms | 5,520 | 47,406 | raw=10, served=10 | +| `libero_object` | 99.8% | 499/500 | 68,712 | 5,199.23s | 256.92s | 1,169.84s | 3,453.50s | 36.18ms | 7,102 | 61,610 | raw=10, served=10 | +| `libero_goal` | 98.0% | 490/500 | 56,292 | 3,981.09s | 210.86s | 1,045.35s | 2,365.71s | 36.10ms | 5,841 | 50,451 | raw=10, served=10 | +| `libero_10` | 94.2% | 471/500 | 134,962 | 9,150.13s | 496.56s | 2,227.75s | 6,237.24s | 36.21ms | 13,713 | 121,249 | raw=10, served=10 | + +参考值不是验收硬阈值。复现时优先确认： + +```text +四组件都完成 500 episodes +chunk 验证 raw=10 served=10 +server 日志确认 CUDA graph capture +结果 JSON 包含 timing 字段 +``` + +## 12. 常见问题 + +### 12.1 LIBERO Docker 镜像架构不匹配 + +如果目标机器是 ARM64，而官方镜像只有 amd64，需要本地构建： + +```bash +cd "$VLA_ROOT" +export DOCKER_DEFAULT_PLATFORM=linux/arm64 +docker/build.sh libero +``` + +### 12.2 容器连不上 PhyAI server + +PhyAI server 如果跑在 bridge Docker 容器里，宿主侧 `vla-eval` 需要连容器 IP： + +```bash +export PHYAI_CONTAINER_IP=$(docker inspect -f '{{range.NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$PHYAI_CONTAINER") +export PHYAI_SERVER_URL=ws://$PHYAI_CONTAINER_IP:8000 +``` + +### 12.3 JSON 没有 timing 字段 + +运行 benchmark 必须加 `--dev`。这会把宿主 `$VLA_ROOT/src` 挂载进 LIBERO 容器，确保使用带 timing 的 runner。 + +### 12.4 Paligemma tokenizer 缺失 + +`paligemma-3b-pt-224` 是 gated 资源。建议从已有机器同步到 `$MODEL_ROOT/paligemma-3b-pt-224`，不要依赖复现机器现场下载。 + +### 12.5 PhyAI server 没有 CUDA graph capture 日志 + +检查启动命令是否误传了 `--no-use_cuda_graph`，或是否走了错误 server adapter。正确日志必须包含： + +```text +capturing vision-tower CUDA graph +capturing 4 prefix-forward CUDA graph(s) +capturing the full 10-step Euler loop as one CUDA graph +``` + +### 12.6 释放资源 + +```bash +tmux kill-session -t phyai_pi05_libero_four || true +tmux kill-session -t phyai_pi05_libero_server || true +docker stop "$PHYAI_CONTAINER" || true +ss -ltnp | grep ':8000' || true +nvidia-smi --query-gpu=index,memory.used,memory.total,utilization.gpu --format=csv,noheader,nounits +``` + From 6558f0e92c6d398e1c85e8485619a498268c6478 Mon Sep 17 00:00:00 2001 From: qyy Date: Sun, 19 Jul 2026 20:40:47 +0800 Subject: [PATCH 08/29] docs: translate pi05 libero guide to english --- ...ai_pi05_libero_four_suites_reproduction.md | 226 ++++++++++-------- 1 file changed, 122 insertions(+), 104 deletions(-) diff --git a/docs/phyai_pi05_libero_four_suites_reproduction.md b/docs/phyai_pi05_libero_four_suites_reproduction.md index e25b92d..e405e8d 100644 --- a/docs/phyai_pi05_libero_four_suites_reproduction.md +++ b/docs/phyai_pi05_libero_four_suites_reproduction.md @@ -1,8 +1,14 @@ -# PhyAI pi0.5 跑 LIBERO 四组件完整复现文档 +--- +title: PhyAI pi0.5 LIBERO four-suite reproduction +description: Run the pi0.5 LIBERO policy with PhyAI on all four LIBERO benchmark suites. +--- -本文档说明如何在一台新机器上使用 PhyAI 推理 pi0.5，并通过 vla-evaluation-harness 跑完整 LIBERO 四组件 benchmark。文档不依赖当前机器的本地路径，所有路径都用环境变量表示。 +# PhyAI pi0.5 LIBERO four-suite reproduction -四组件指： +This guide shows how to run the pi0.5 LIBERO policy with PhyAI and evaluate it with `vla-evaluation-harness` on all four LIBERO suites. +It is written for a fresh machine and avoids local machine-specific paths by using environment variables. + +The four suites are: ```text libero_spatial -> configs/benchmarks/libero/spatial.yaml @@ -11,42 +17,41 @@ libero_goal -> configs/benchmarks/libero/goal.yaml libero_10 -> configs/benchmarks/libero/10.yaml ``` -完整口径： +The benchmark setup in this guide uses: ```text -模式：sync -chunk_size：10 -每个组件：10 tasks x 50 episodes = 500 episodes -总量：4 suites x 500 episodes = 2000 episodes -模型：PhyAI pi0.5 LIBERO converted checkpoint -仿真：vla-evaluation-harness LIBERO Docker 容器 -结果：每个 suite 一个 JSON，记录 success、steps、timing、chunk size +Mode: sync +Chunk size: 10 +Episodes per suite: 10 tasks x 50 episodes = 500 episodes +Total episodes: 4 suites x 500 episodes = 2000 episodes +Model: PhyAI pi0.5 LIBERO converted checkpoint +Simulator: vla-evaluation-harness LIBERO Docker container +Output: one JSON result file per suite with success, steps, timing, and chunk-size fields ``` -## 1. 机器与资源要求 +## 1. Prerequisites -建议机器： +Use a Linux machine with: ```text -GPU：至少 1 张 CUDA GPU，显存建议 >= 48GB -系统：Linux -容器：Docker + NVIDIA Container Toolkit -Python 环境管理：uv -辅助工具：tmux, nvidia-smi, ss +GPU: at least 1 CUDA GPU, 48 GB or more GPU memory recommended +Container runtime: Docker and NVIDIA Container Toolkit +Python environment manager: uv +Utility tools: tmux, nvidia-smi, ss ``` -必须准备的模型资源： +Prepare these model resources before you start: ```text -PhyAI converted checkpoint：pi05_libero_phyai_converted -PaLI-Gemma tokenizer / processor：paligemma-3b-pt-224 +PhyAI converted checkpoint: pi05_libero_phyai_converted +PaLI-Gemma tokenizer / processor: paligemma-3b-pt-224 ``` -`paligemma-3b-pt-224` 是 gated 资源，推荐从已有机器同步，不建议在复现时现场下载。 +`paligemma-3b-pt-224` is a gated resource. Prefer syncing it from a machine that already has access instead of downloading it during reproduction. -## 2. 推荐环境变量 +## 2. Set environment variables -根据目标机器实际路径设置： +Set paths for the target machine: ```bash export PHYAI_ROOT=$HOME/phyai @@ -63,14 +68,14 @@ export TOKENIZER_IN_CONTAINER=/data/share/paligemma-3b-pt-224 export LIBERO_IMAGE=ghcr.io/allenai/vla-evaluation-harness/libero:latest ``` -确认资源存在： +Check that the model directories exist: ```bash test -d "$PHYAI_CKPT_HOST" test -d "$TOKENIZER_HOST" ``` -如果模型在另一台机器上，示例同步命令如下： +If the models are on another machine, sync them to the target machine: ```bash rsync -azP \ @@ -79,9 +84,9 @@ rsync -azP \ user@target-host:$MODEL_ROOT/ ``` -## 3. 获取代码并创建环境 +## 3. Clone source code and create environments -clone PhyAI： +Clone PhyAI: ```bash git clone https://github.com/MEmbodied/phyai.git "$PHYAI_ROOT" @@ -89,7 +94,7 @@ cd "$PHYAI_ROOT" uv sync ``` -clone vla-evaluation-harness： +Clone `vla-evaluation-harness`: ```bash git clone https://github.com/allenai/vla-evaluation-harness.git "$VLA_ROOT" @@ -98,17 +103,19 @@ uv sync ./.venv/bin/vla-eval --help >/tmp/vla_eval_help.log ``` -如果目标机器不能联网，可以在能联网机器上 clone 后用 `rsync` 同步两个仓库；同步后仍建议在目标机器上执行 `uv sync`，让本机 Python、CUDA、依赖和 editable path 正确。 +If the target machine has no network access, clone both repositories on a networked machine and sync them with `rsync`. +After syncing, still run `uv sync` on the target machine so editable paths, Python versions, CUDA libraries, and local dependencies are resolved correctly. -## 4. 准备 LIBERO Docker 镜像 +## 4. Prepare the LIBERO Docker image -vla-harness 跑 LIBERO 时会启动 LIBERO benchmark 容器。x86_64 机器可以直接使用官方镜像： +`vla-evaluation-harness` runs LIBERO inside a benchmark container. +On an `x86_64` machine, pull the official image: ```bash docker pull "$LIBERO_IMAGE" ``` -如果目标机器是 ARM64，例如 Thor，而官方镜像只有 amd64，需要在目标机器本地构建 ARM64 镜像： +If the target machine is ARM64 and the official image is only available for `amd64`, build an ARM64 LIBERO image locally: ```bash cd "$VLA_ROOT" @@ -119,21 +126,22 @@ docker image inspect "$LIBERO_IMAGE" \ --format '{{.Architecture}} {{.Os}}' ``` -期望输出： +Expected output on ARM64: ```text arm64 linux ``` -如果是 x86_64，期望输出一般是： +Expected output on `x86_64`: ```text amd64 linux ``` -## 5. 创建 PhyAI Docker 容器 +## 5. Create the PhyAI Docker container -建议在 Docker 中运行 PhyAI server，容器挂载 PhyAI 代码、vla-harness 代码和模型目录： +Run the PhyAI server inside a Docker container. +Mount the PhyAI source tree, the `vla-evaluation-harness` source tree, and the model directory: ```bash docker run -dit --gpus all \ @@ -151,7 +159,7 @@ docker run -dit --gpus all \ nvcr.io/nvidia/pytorch:25.12-py3 bash ``` -容器内同步 PhyAI 环境： +Install the PhyAI environment inside the container: ```bash docker exec "$PHYAI_CONTAINER" bash -lc ' @@ -161,7 +169,8 @@ uv sync ' ``` -如果 `uv sync` 生成的 editable path 与容器路径不一致，可加兼容 symlink。只有在 import 报错指向旧宿主路径时才需要执行；`HOST_USER` 填目标机器上的用户名： +If `uv sync` produces editable paths that point to the host path instead of the container path, create a compatibility symlink. +Only run this if imports fail because a stale host path is referenced: ```bash export HOST_USER=$(id -un) @@ -173,7 +182,7 @@ ln -sfn /phyai_workspace $COMPAT_PARENT/$HOST_USER/phyai " ``` -确认容器内 import： +Verify imports inside the container: ```bash docker exec "$PHYAI_CONTAINER" bash -lc ' @@ -188,9 +197,9 @@ PY ' ``` -## 6. 运行前检查 +## 6. Check the machine before running -检查 GPU、端口和 Docker： +Check GPU usage, ports, and containers before starting the benchmark: ```bash nvidia-smi --query-gpu=index,name,memory.used,memory.total,utilization.gpu --format=csv,noheader @@ -198,11 +207,12 @@ ss -ltnp | grep -E ':8000|:8001' || true docker ps -a --format 'table {{.Names}}\t{{.Status}}\t{{.Image}}' ``` -如 GPU 上已有重负载，完整 benchmark 的耗时会不稳定，建议换空闲 GPU 或等待资源释放。 +If the GPU is under heavy load, benchmark timing can become unstable. +Use an idle GPU or wait for other jobs to finish. -## 7. 启动 PhyAI pi0.5 server +## 7. Start the PhyAI pi0.5 server -取 PhyAI 容器 IP： +Get the PhyAI container IP and construct the WebSocket URL: ```bash export PHYAI_CONTAINER_IP=$(docker inspect -f '{{range.NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$PHYAI_CONTAINER") @@ -210,7 +220,7 @@ export PHYAI_SERVER_URL=ws://$PHYAI_CONTAINER_IP:8000 echo "$PHYAI_SERVER_URL" ``` -启动 server： +Start the server in `tmux`: ```bash mkdir -p "$VLA_ROOT/results" @@ -236,28 +246,28 @@ export PHYAI_CAMERA_MODE=two_camera " ``` -关键配置含义： +Key settings: -| 配置 | 值 | 说明 | +| Setting | Value | Purpose | | --- | --- | --- | -| `--checkpoint_path` | `/data/share/pi05_libero_phyai_converted` | PhyAI converted pi0.5 LIBERO 权重 | -| `PHYAI_TOKENIZER_PATH` | `/data/share/paligemma-3b-pt-224` | tokenizer / processor 文件 | -| `PHYAI_CAMERA_MODE` | `two_camera` | LIBERO 发送 agentview 与 wrist 两路图像 | -| `--params_dtype` | `bfloat16` | 参数 dtype | -| `--attn_backend` | `flashinfer` | attention 后端 | -| `--norm_backend` | `phyai-kernel` | norm 后端 | -| `--linear_backend` | `flashinfer` | linear 后端 | -| `--flashinfer_workspace_bytes` | `536870912` | 512MiB workspace | -| `--chunk_size` | `10` | 每次推理产出 10 个 action | -| CUDA graph | 默认开启 | 不要传 `--no-use_cuda_graph` | - -等待 ready： +| `--checkpoint_path` | `/data/share/pi05_libero_phyai_converted` | PhyAI converted pi0.5 LIBERO checkpoint | +| `PHYAI_TOKENIZER_PATH` | `/data/share/paligemma-3b-pt-224` | Tokenizer and processor directory | +| `PHYAI_CAMERA_MODE` | `two_camera` | LIBERO sends both agent-view and wrist-camera images | +| `--params_dtype` | `bfloat16` | Parameter dtype | +| `--attn_backend` | `flashinfer` | Attention backend | +| `--norm_backend` | `phyai-kernel` | Normalization backend | +| `--linear_backend` | `flashinfer` | Linear backend | +| `--flashinfer_workspace_bytes` | `536870912` | 512 MiB FlashInfer workspace | +| `--chunk_size` | `10` | The policy returns 10 actions per inference call | +| CUDA graph | Enabled by default | Do not pass `--no-use_cuda_graph` | + +Follow the server log: ```bash tail -f "$VLA_ROOT/results/phyai_pi05_libero_server.log" ``` -必须看到： +Wait until the log contains: ```text capturing vision-tower CUDA graph @@ -266,9 +276,10 @@ capturing the full 10-step Euler loop as one CUDA graph Starting server on ws://0.0.0.0:8000 ``` -## 8. 创建 smoke 配置并验证 +## 8. Run a smoke test -先跑一个最小 smoke，确认模型、LIBERO Docker、WebSocket、timing 字段都可用。 +Run a minimal smoke test before launching the full benchmark. +This checks the model, LIBERO Docker image, WebSocket connection, timing fields, and chunk-size fields. ```bash cat > "$VLA_ROOT/configs/benchmarks/libero/smoke_test_phyai_local.yaml" < "$VLA_ROOT/run_phyai_pi05_libero_four_suites.sh" <<'SH' @@ -420,7 +433,7 @@ SH chmod +x "$VLA_ROOT/run_phyai_pi05_libero_four_suites.sh" ``` -启动长跑： +Start the full run: ```bash cd "$VLA_ROOT" @@ -428,7 +441,7 @@ tmux new-session -d -s phyai_pi05_libero_four \ "PHYAI_SERVER_URL=$PHYAI_SERVER_URL PHYAI_CKPT_IN_CONTAINER=$PHYAI_CKPT_IN_CONTAINER LIBERO_IMAGE=$LIBERO_IMAGE ./run_phyai_pi05_libero_four_suites.sh" ``` -查看进度： +Check progress: ```bash tmux ls @@ -438,9 +451,9 @@ sed -n '1,220p' "$latest/run_summary.log" tail -80 "$latest"/phyai_spatial.log ``` -## 10. 解析成功率与 timing +## 10. Summarize success rate and timing -长跑完成后： +After the run finishes, print the run summary: ```bash cd "$VLA_ROOT" @@ -448,7 +461,7 @@ latest=$(ls -td results/phyai_pi05_libero_four_* | head -1) cat "$latest/run_summary.log" ``` -用 `summarize_timing.py` 重新汇总全部 JSON： +Summarize timing from all result JSON files: ```bash cd "$VLA_ROOT" @@ -456,7 +469,7 @@ latest=$(ls -td results/phyai_pi05_libero_four_* | head -1) ./.venv/bin/python scripts/summarize_timing.py "$latest"/*.json ``` -统计成功率： +Summarize success rate: ```bash cd "$VLA_ROOT" @@ -478,17 +491,17 @@ for arg in sys.argv[1:]: PY ``` -必须记录的字段： +Record these fields for each suite: ```text RUN_ID -结果目录 -checkpoint -server_url -suite -success / total -success rate -steps +Result directory +Checkpoint +Server URL +Suite +Success / total +Success rate +Steps /usr/bin/time -p real model_wait_sec model_inference_sec @@ -501,33 +514,36 @@ raw_chunk_size_max served_chunk_size_max ``` -`raw_chunk_size_max=10` 且 `served_chunk_size_max=10` 是四组件复现的关键校验项。 +`raw_chunk_size_max=10` and `served_chunk_size_max=10` are key checks for this four-suite reproduction. -## 11. 预期参考结果 +## 11. Reference results -不同机器、GPU、驱动、负载会影响耗时；成功率也可能有少量随机波动。此前同口径参考结果如下： +Timing depends on the GPU, driver, machine load, and container environment. +Success rate can also vary slightly across runs. +The following results are from a previous run with the same evaluation setup: -| 组件 | 成功率 | 成功数 | steps | 总耗时 | 模型纯推理时间 | env step 时间 | benchmark 等待 action 时间 | 平均单次模型推理 | 推理调用次数 | buffer hit | chunk 验证 | +| Suite | Success rate | Success | Steps | Wall time | Model inference time | Env step time | Benchmark action-wait time | Average model inference | Inference calls | Buffer hits | Chunk check | | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | | `libero_spatial` | 97.8% | 489/500 | 52,926 | 4,198.73s | 200.56s | 1,239.61s | 2,269.16s | 36.33ms | 5,520 | 47,406 | raw=10, served=10 | | `libero_object` | 99.8% | 499/500 | 68,712 | 5,199.23s | 256.92s | 1,169.84s | 3,453.50s | 36.18ms | 7,102 | 61,610 | raw=10, served=10 | | `libero_goal` | 98.0% | 490/500 | 56,292 | 3,981.09s | 210.86s | 1,045.35s | 2,365.71s | 36.10ms | 5,841 | 50,451 | raw=10, served=10 | | `libero_10` | 94.2% | 471/500 | 134,962 | 9,150.13s | 496.56s | 2,227.75s | 6,237.24s | 36.21ms | 13,713 | 121,249 | raw=10, served=10 | -参考值不是验收硬阈值。复现时优先确认： +Use these numbers as references, not strict pass/fail thresholds. +For reproduction, first confirm: ```text -四组件都完成 500 episodes -chunk 验证 raw=10 served=10 -server 日志确认 CUDA graph capture -结果 JSON 包含 timing 字段 +All four suites finish 500 episodes each +Chunk check is raw=10 served=10 +The server log confirms CUDA graph capture +The result JSON files contain timing fields ``` -## 12. 常见问题 +## 12. Troubleshooting -### 12.1 LIBERO Docker 镜像架构不匹配 +### 12.1 LIBERO Docker image architecture mismatch -如果目标机器是 ARM64，而官方镜像只有 amd64，需要本地构建： +If the target machine is ARM64 and the official image is only available for `amd64`, build the image locally: ```bash cd "$VLA_ROOT" @@ -535,26 +551,29 @@ export DOCKER_DEFAULT_PLATFORM=linux/arm64 docker/build.sh libero ``` -### 12.2 容器连不上 PhyAI server +### 12.2 The benchmark cannot connect to the PhyAI server -PhyAI server 如果跑在 bridge Docker 容器里，宿主侧 `vla-eval` 需要连容器 IP： +If the PhyAI server runs in a bridge Docker container, the host-side `vla-eval` process should connect to the container IP: ```bash export PHYAI_CONTAINER_IP=$(docker inspect -f '{{range.NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$PHYAI_CONTAINER") export PHYAI_SERVER_URL=ws://$PHYAI_CONTAINER_IP:8000 ``` -### 12.3 JSON 没有 timing 字段 +### 12.3 The result JSON has no timing fields -运行 benchmark 必须加 `--dev`。这会把宿主 `$VLA_ROOT/src` 挂载进 LIBERO 容器，确保使用带 timing 的 runner。 +Run the benchmark with `--dev`. +This mounts the host `$VLA_ROOT/src` tree into the LIBERO container and ensures the benchmark uses the runner that records timing fields. -### 12.4 Paligemma tokenizer 缺失 +### 12.4 The PaLI-Gemma tokenizer is missing -`paligemma-3b-pt-224` 是 gated 资源。建议从已有机器同步到 `$MODEL_ROOT/paligemma-3b-pt-224`，不要依赖复现机器现场下载。 +`paligemma-3b-pt-224` is a gated resource. +Sync it from an existing machine to `$MODEL_ROOT/paligemma-3b-pt-224` instead of relying on an on-the-fly download. -### 12.5 PhyAI server 没有 CUDA graph capture 日志 +### 12.5 The PhyAI server log has no CUDA graph capture messages -检查启动命令是否误传了 `--no-use_cuda_graph`，或是否走了错误 server adapter。正确日志必须包含： +Check whether the server command accidentally passed `--no-use_cuda_graph` or whether it started the wrong server adapter. +The expected log must include: ```text capturing vision-tower CUDA graph @@ -562,7 +581,7 @@ capturing 4 prefix-forward CUDA graph(s) capturing the full 10-step Euler loop as one CUDA graph ``` -### 12.6 释放资源 +### 12.6 Release resources ```bash tmux kill-session -t phyai_pi05_libero_four || true @@ -571,4 +590,3 @@ docker stop "$PHYAI_CONTAINER" || true ss -ltnp | grep ':8000' || true nvidia-smi --query-gpu=index,memory.used,memory.total,utilization.gpu --format=csv,noheader,nounits ``` - From a40f5d79d4484f3b642b817acc1a39c0b9f414bf Mon Sep 17 00:00:00 2001 From: rebecca26358 <153866717+rebecca26358@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:20:38 +0000 Subject: [PATCH 09/29] benchmark: add external pi05 runtime latency wrappers --- .../pi05/README_external_runtime_latency.md | 214 ++++++++++++++++ benchmark/pi05/bench_flashrt_pi05.py | 230 +++++++++++++++++ benchmark/pi05/bench_realtime_vla_pi05.py | 225 +++++++++++++++++ benchmark/pi05/bench_vlacpp_pi05_client.py | 235 ++++++++++++++++++ 4 files changed, 904 insertions(+) create mode 100644 benchmark/pi05/README_external_runtime_latency.md create mode 100755 benchmark/pi05/bench_flashrt_pi05.py create mode 100755 benchmark/pi05/bench_realtime_vla_pi05.py create mode 100755 benchmark/pi05/bench_vlacpp_pi05_client.py diff --git a/benchmark/pi05/README_external_runtime_latency.md b/benchmark/pi05/README_external_runtime_latency.md new file mode 100644 index 0000000..a16a7af --- /dev/null +++ b/benchmark/pi05/README_external_runtime_latency.md @@ -0,0 +1,214 @@ +# External PI0.5 runtime latency wrappers + +This document explains how to set up and run the three external PI0.5 latency +wrappers under `benchmark/pi05/`. + +These scripts do not use the PhyAI engine for inference. Each script calls the +target runtime directly, while reusing PhyAI's common benchmark runner for +warmup, timing, and JSONL output. + +| Script | Runtime measured | Timed call | +| --- | --- | --- | +| `bench_flashrt_pi05.py` | FlashRT | `Pi05TorchFrontendRtx.infer(obs)` | +| `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward(...)` | +| `bench_vlacpp_pi05_client.py` | vla.cpp | one ZMQ request to a running `vla-server` | + +## Common setup + +Use one Python environment that can import PhyAI, PyTorch, and +`benchmark/bench_n_batch.py`. + +```bash +cd +python -c "import torch; import phyai; import benchmark.bench_n_batch" +nvidia-smi +``` + +Use the same benchmark settings when comparing runtimes: + +```text +batch size: 1 +views / camera streams: 2 +chunk size: 50 +prompt: keep the same text across runs +warmup / timed iterations: use the same values across runs +precision: label each row by the runtime's real precision path +``` + +The wrappers generate synthetic image/state inputs. They are for latency-only +measurements, not LIBERO accuracy evaluation. + +## Placeholders + +Use your own paths for these placeholders: + +| Placeholder | Meaning | +| --- | --- | +| `` | PhyAI checkout containing `benchmark/pi05/` | +| `` | FlashRT checkout | +| `` | realtime-vla checkout | +| `` | vla.cpp checkout | +| `` | compiled vla.cpp `vla-server` binary | +| `` | PI0.5 safetensors checkpoint directory or file | +| `` | PI0.5 GGUF file for vla.cpp | +| `` | vla.cpp multimodal projector GGUF | +| `` | local tokenizer directory | +| `` | local LIBERO `meta/stats.json` with `observation.state.q01/q99` | + +## FlashRT + +Install FlashRT following its official README, then make the checkout visible: + +```bash +export FLASHRT_ROOT= +cd +python -c "import sys; sys.path.insert(0, '$FLASHRT_ROOT'); from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx; print(Pi05TorchFrontendRtx)" +``` + +Run latency: + +```bash +cd +python benchmark/pi05/bench_flashrt_pi05.py \ + --flashrt-root \ + --checkpoint \ + --precision bf16 \ + --num-views 2 \ + --chunk-size 50 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/flashrt_pi05.jsonl +``` + +Notes: + +- The script uses FlashRT's direct `Pi05TorchFrontendRtx` API because `chunk_size` + is a frontend constructor argument. +- Do not use `load_model(..., num_steps=50)` to set action chunk size. In + FlashRT, `num_steps` means denoise steps. +- `--precision bf16` sets FlashRT's forced-BF16 PI0.5 RTX path. Use + `--precision fp8_bf16` for FlashRT's optimized FP8/BF16 path and label that + result separately. + +## realtime-vla + +Install realtime-vla following its official README. The wrapper can read a +converted `.pt` / `.pth` checkpoint directly. If you pass a PI0.5 safetensors +checkpoint, also provide FlashRT so the wrapper can reuse FlashRT's PI0.5 +conversion helper. + +```bash +export REALTIME_VLA_ROOT= +export FLASHRT_ROOT= +cd +python -c "import sys; sys.path.insert(0, '$REALTIME_VLA_ROOT'); from pi05_infer import Pi05Inference; print(Pi05Inference)" +``` + +Run latency: + +```bash +cd +python benchmark/pi05/bench_realtime_vla_pi05.py \ + --realtime-vla-root \ + --flashrt-root \ + --checkpoint \ + --num-views 2 \ + --chunk-size 50 \ + --prompt-len 16 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/realtime_vla_pi05.jsonl +``` + +Notes: + +- The wrapper uses BF16 synthetic inputs. +- For `.pkl` / `.pickle` checkpoints, add `--trust-pickle-checkpoint` only when + the file is trusted. +- If the checkpoint does not contain `language_embeds`, the wrapper creates a + synthetic prompt embedding for latency-only runs. + +## vla.cpp + +vla.cpp uses a server/client flow. Build `vla-server` with CUDA enabled, then +start the server in one shell and run the Python benchmark client in another. + +Basic checks: + +```bash +test -x +test -f +test -f +test -f /tokenizer.json +test -f +``` + +Start server: + +```bash + \ + --bind tcp://127.0.0.1:5555 \ + --timing-detail phase \ + \ + +``` + +Run latency client: + +```bash +cd +python benchmark/pi05/bench_vlacpp_pi05_client.py \ + --vlacpp-root \ + --addr tcp://127.0.0.1:5555 \ + --arch pi05 \ + --tokenizer \ + --stats-json \ + --num-views 2 \ + --chunk-size 50 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/vlacpp_pi05.jsonl +``` + +Notes: + +- vla.cpp needs GGUF files; a safetensors checkpoint is not enough. +- Prefer local tokenizer and stats files to avoid network or HuggingFace auth + issues during benchmarking. For PI0.5, vla.cpp expects a lerobot-style + `meta/stats.json`; OpenPI-style `norm_stats.json` is not the same format. +- The wrapper records client wall latency. If the server returns phase timing, + it is written under `extras.server_phase_latency_ms`. + +## Quick validation + +After each runtime is installed, reduce iterations to check that the wrapper can +start and write JSONL: + +```bash +--n-warmup 1 --n-timed 1 --result-file results/pi05_smoke.jsonl +``` + +For vla.cpp, keep `vla-server` running before starting the client smoke test. + +## Troubleshooting + +| Symptom | Check | +| --- | --- | +| `No module named flash_rt` | Pass `--flashrt-root` or set `FLASHRT_ROOT`. | +| `No module named pi05_infer` | Pass `--realtime-vla-root` or set `REALTIME_VLA_ROOT`. | +| realtime-vla safetensors conversion fails | Also pass `--flashrt-root`; verify FlashRT import works. | +| vla.cpp client cannot connect | Confirm the server printed that it is ready and the `--addr` matches `--bind`. | +| vla.cpp tokenizer downloads or asks for auth | Use a local tokenizer directory. | +| GPU architecture build error | Check CUDA, PyTorch CUDA, driver, and build flags for the target GPU. | +| Latency is much slower than expected | Check `nvidia-smi`, rerun after warmup/JIT, and make sure no other process is using the GPU. | + +## Timing scope + +- FlashRT: wall time around steady-state `Pi05TorchFrontendRtx.infer(obs)`, + after prompt setup, calibration, and first graph-building call. +- realtime-vla: CUDA-event time around one `Pi05Inference.forward(...)` call. +- vla.cpp: client wall time for one ZMQ request; server phase timing is copied + from the response when available. diff --git a/benchmark/pi05/bench_flashrt_pi05.py b/benchmark/pi05/bench_flashrt_pi05.py new file mode 100755 index 0000000..c045949 --- /dev/null +++ b/benchmark/pi05/bench_flashrt_pi05.py @@ -0,0 +1,230 @@ +#!/usr/bin/env python3 +"""FlashRT PI0.5 latency benchmark using the PhyAI bench runner. + +This is a thin adapter around FlashRT's direct ``Pi05TorchFrontendRtx`` path. +The direct frontend is used because action chunk size is a frontend constructor +argument; ``flash_rt.load_model(..., num_steps=...)`` controls denoise steps, +not action chunk size. It mirrors ``benchmark/bench_n_batch_ws1_pi05.py``: +framework-specific setup lives here, while warmup, timed iterations, JSONL +output, and optional profiling are handled by ``benchmark/bench_n_batch.py``. + +Run:: + + python benchmark/pi05/bench_flashrt_pi05.py \ + --flashrt-root \ + --checkpoint \ + --batch-sizes 1 --n-warmup 100 --n-timed 100 \ + --result-file results/flashrt_pi05.jsonl + +Only batch size 1 is supported because the FlashRT PI0.5 direct frontend path +used here takes one robot request at a time. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import statistics +import sys +from typing import Any + +import numpy as np +import torch + +_BENCHMARK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_BENCHMARK_DIR)) +import bench_n_batch as bnb # noqa: E402 +from phyai.utils.profile import ( # noqa: E402 + add_profile_cli_args, + install_profiler, + profile_config_from_args, +) + + +def env_path(name: str) -> Path | None: + value = os.environ.get(name) + return Path(value) if value else None + + +class LazySummary(dict): + def __init__(self, values_fn_or_items): + if callable(values_fn_or_items): + super().__init__() + self._values_fn = values_fn_or_items + else: + super().__init__(values_fn_or_items) + self._values_fn = None + + def items(self): + if self._values_fn is None: + return super().items() + summary = summarize(self._values_fn()) + return (summary or {}).items() + + +def summarize(values: list[float]) -> dict[str, float] | None: + if not values: + return None + xs = sorted(float(x) for x in values) + return { + "count": len(xs), + "mean_ms": float(statistics.fmean(xs)), + "median_ms": float(statistics.median(xs)), + "p50_ms": float(np.percentile(xs, 50)), + "p90_ms": float(np.percentile(xs, 90)), + "p99_ms": float(np.percentile(xs, 99)), + "min_ms": xs[0], + "max_ms": xs[-1], + } + + +def make_observation(num_views: int, seed: int, prompt: str) -> dict[str, Any]: + """Deterministic synthetic observation for latency-only runs.""" + rng = np.random.default_rng(seed) + obs: dict[str, Any] = { + "image": rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8), + "state": rng.standard_normal(8).astype(np.float32), + "task": prompt, + "prompt": prompt, + } + if num_views >= 2: + obs["wrist_image"] = rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8) + if num_views >= 3: + obs["wrist_image_right"] = rng.integers( + 0, 256, size=(224, 224, 3), dtype=np.uint8 + ) + return obs + + +def import_flashrt_frontend(repo: Path): + # Use the checked-out FlashRT repository directly; installation is optional. + sys.path.insert(0, str(repo)) + from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx + + return Pi05TorchFrontendRtx + + +def make_setup_fn(args: argparse.Namespace): + frontend_cls = import_flashrt_frontend(args.flashrt_root) + + def setup_fn(batch_size: int) -> bnb.BenchSpec: + if batch_size != 1: + raise ValueError("FlashRT PI0.5 wrapper supports only batch_size=1") + + if args.precision == "bf16": + # FlashRT uses this environment switch to force the BF16 PI0.5 RTX path. + os.environ["FVK_PI05_RTX_FORCE_BF16"] = "1" + else: + os.environ.pop("FVK_PI05_RTX_FORCE_BF16", None) + + model = frontend_cls( + args.checkpoint, + num_views=args.num_views, + chunk_size=args.chunk_size, + cache_frames=1, + use_fp8=(args.precision == "fp8_bf16"), + hardware=args.hardware, + ) + obs = make_observation(args.num_views, args.seed, args.prompt) + + # set_prompt builds the prompt-specific pipeline; calibration captures the graph. + # These setup calls are outside the measured latency window. + model.set_prompt(args.prompt) + model.calibrate_with_real_data([obs]) + model.infer(obs) + torch.cuda.synchronize() + + call_count = 0 + + def step() -> None: + nonlocal call_count + if call_count == args.n_warmup: + model.latency_records.clear() + model.infer(obs) + # FlashRT PI0.5 may use an internal CUDA stream, so synchronize inside the + # step to make the common runner's timing boundary conservative. + torch.cuda.synchronize() + call_count += 1 + + spec = bnb.BenchSpec( + name="flashrt_pi05", + step_callable=step, + teardown_callable=lambda: None, + ) + spec.flashrt_internal_latency_ms = LazySummary( + lambda: [float(x) for x in getattr(model, "latency_records", [])] + ) # type: ignore[attr-defined] + return spec + + return setup_fn + + +def make_extras_fn(args: argparse.Namespace): + def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: + return { + "runtime": "FlashRT", + "checkpoint": str(args.checkpoint), + "flashrt_root": str(args.flashrt_root), + "precision": args.precision, + "hardware": args.hardware, + "batch_size_contract": 1, + "num_views": args.num_views, + "chunk_size": args.chunk_size, + "prompt": args.prompt, + "seed": args.seed, + "internal_latency_ms": getattr(spec, "flashrt_internal_latency_ms", {}), + "timing_scope": "FlashRT direct Pi05TorchFrontendRtx.infer hot path after set_prompt and first graph-building infer; common runner uses perf-counter wall time because FlashRT runs work on an internal CUDA stream", + } + + return extras_fn + + +def main() -> None: + flashrt_default = env_path("FLASHRT_ROOT") + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--flashrt-root", + type=Path, + default=flashrt_default, + required=flashrt_default is None, + help="Path to the FlashRT repository clone. Can also be set with FLASHRT_ROOT.", + ) + parser.add_argument( + "--checkpoint", + type=Path, + required=True, + help="PI0.5 checkpoint directory readable by FlashRT.", + ) + parser.add_argument("--num-views", type=int, default=2) + parser.add_argument("--chunk-size", type=int, default=50) + parser.add_argument("--prompt", default="do something") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--precision", choices=("bf16", "fp8_bf16"), default="bf16") + parser.add_argument("--hardware", default="auto") + + bnb.add_bench_cli_args(parser) + parser.set_defaults(bench_name="flashrt_pi05", batch_sizes=[1]) + add_profile_cli_args(parser) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for FlashRT PI0.5 benchmarking") + + install_profiler(profile_config_from_args(args)) + runner = bnb.NBatchBenchRunner( + setup_fn=make_setup_fn(args), + extras_fn=make_extras_fn(args), + # FlashRT may execute on an internal CUDA stream. Force the common + # runner's perf-counter path; step() synchronizes CUDA before returning. + device=torch.device("cpu"), + **bnb.bench_runner_kwargs_from_args(args), + ) + runner.run() + + +if __name__ == "__main__": + main() diff --git a/benchmark/pi05/bench_realtime_vla_pi05.py b/benchmark/pi05/bench_realtime_vla_pi05.py new file mode 100755 index 0000000..f91a6e4 --- /dev/null +++ b/benchmark/pi05/bench_realtime_vla_pi05.py @@ -0,0 +1,225 @@ +#!/usr/bin/env python3 +"""realtime-vla PI0.5 latency benchmark using the PhyAI bench runner. + +This script is a thin adapter around realtime-vla's ``Pi05Inference.forward``. +It mirrors ``benchmark/bench_n_batch_ws1_pi05.py`` by delegating warmup, timed +iterations, JSONL output, and optional profiling to ``benchmark/bench_n_batch.py``. + +Run:: + + python benchmark/pi05/bench_realtime_vla_pi05.py \ + --realtime-vla-root \ + --flashrt-root \ + --checkpoint \ + --batch-sizes 1 --n-warmup 100 --n-timed 100 \ + --result-file results/realtime_vla_pi05.jsonl + +Only batch size 1 is supported because realtime-vla's PI0.5 inference object +used here is allocated for one synthetic request. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import pickle +import sys +from typing import Any + +import torch + +_BENCHMARK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_BENCHMARK_DIR)) +import bench_n_batch as bnb # noqa: E402 +from phyai.utils.profile import ( # noqa: E402 + add_profile_cli_args, + install_profiler, + profile_config_from_args, +) + + +def env_path(name: str) -> Path | None: + value = os.environ.get(name) + return Path(value) if value else None + + +def load_checkpoint(path: Path, flashrt_root: Path | None, trust_pickle: bool): + if path.is_dir(): + path = path / "model.safetensors" + if path.suffix == ".safetensors": + if flashrt_root is None: + raise ValueError( + "--flashrt-root or FLASHRT_ROOT is required when loading a safetensors checkpoint" + ) + # Reuse FlashRT's tested PI0.5 key conversion instead of duplicating it here. + sys.path.insert(0, str(flashrt_root)) + from flash_rt.frontends.torch.pi05_rtx import convert_pi05_safetensors + + return convert_pi05_safetensors(path) + if path.suffix in {".pt", ".pth"}: + return torch.load(path, map_location="cpu", weights_only=True) + if path.suffix in {".pkl", ".pickle"}: + if not trust_pickle: + raise ValueError( + "Refusing to load pickle checkpoint without --trust-pickle-checkpoint. " + "Only use that flag for checkpoints from a trusted source." + ) + with path.open("rb") as f: + return pickle.load(f) # nosec B301: guarded by --trust-pickle-checkpoint. + raise ValueError( + f"Unsupported checkpoint suffix {path.suffix!r}; expected .safetensors, .pt, .pth, .pkl, or .pickle" + ) + + +def make_setup_fn(args: argparse.Namespace): + # Use the checked-out realtime-vla repository directly; installation is optional. + sys.path.insert(0, str(args.realtime_vla_root)) + from pi05_infer import Pi05Inference + + def setup_fn(batch_size: int) -> bnb.BenchSpec: + if batch_size != 1: + raise ValueError("realtime-vla PI0.5 wrapper supports only batch_size=1") + + checkpoint = load_checkpoint( + args.checkpoint, args.flashrt_root, args.trust_pickle_checkpoint + ) + if "language_embeds" not in checkpoint: + # Latency-only fallback for checkpoints that do not include prompt embeds. + checkpoint["language_embeds"] = torch.randn( + args.prompt_len, 2048, dtype=torch.bfloat16 + ) + + infer = Pi05Inference( + checkpoint=checkpoint, + num_views=args.num_views, + chunk_size=args.chunk_size, + tokenizer_path=str(args.tokenizer) if args.tokenizer else None, + discrete_state_input=args.discrete_state_input, + ) + torch.manual_seed(args.seed) + input_image = torch.randn( + args.num_views, 224, 224, 3, dtype=torch.bfloat16, device="cuda" + ) + input_noise = torch.randn( + args.chunk_size, 32, dtype=torch.bfloat16, device="cuda" + ) + + state_tokens = None + if args.discrete_state_input: + import numpy as np + + state_tokens = np.zeros(args.state_dim, dtype=np.int64) + + def step() -> None: + infer.forward( + input_image, + input_noise, + task_prompt=args.prompt, + state_tokens=state_tokens, + ) + + return bnb.BenchSpec( + name="realtime_vla_pi05", + step_callable=step, + teardown_callable=lambda: None, + ) + + return setup_fn + + +def make_extras_fn(args: argparse.Namespace): + def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: + return { + "runtime": "realtime-vla", + "checkpoint": str(args.checkpoint), + "realtime_vla_root": str(args.realtime_vla_root), + "flashrt_root": str(args.flashrt_root) if args.flashrt_root else None, + "precision": "bf16", + "batch_size_contract": 1, + "num_views": args.num_views, + "chunk_size": args.chunk_size, + "prompt": args.prompt, + "prompt_len": args.prompt_len, + "seed": args.seed, + "discrete_state_input": args.discrete_state_input, + "timing_scope": "Pi05Inference.forward hot path; common runner CUDA event wraps one forward call", + } + + return extras_fn + + +def main() -> None: + realtime_default = env_path("REALTIME_VLA_ROOT") + flashrt_default = env_path("FLASHRT_ROOT") + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--realtime-vla-root", + type=Path, + default=realtime_default, + required=realtime_default is None, + help="Path to the realtime-vla repository clone. Can also be set with REALTIME_VLA_ROOT.", + ) + parser.add_argument( + "--flashrt-root", + type=Path, + default=flashrt_default, + help="Path to FlashRT. Required only when --checkpoint is a PI0.5 safetensors checkpoint.", + ) + parser.add_argument( + "--checkpoint", + type=Path, + required=True, + help="realtime-vla .pkl/.pt checkpoint, PI0.5 model.safetensors, or a directory containing model.safetensors.", + ) + parser.add_argument( + "--trust-pickle-checkpoint", + action="store_true", + help="Allow loading .pkl/.pickle checkpoints. Only use with trusted checkpoint files.", + ) + parser.add_argument("--num-views", type=int, default=2) + parser.add_argument("--chunk-size", type=int, default=50) + parser.add_argument("--prompt", default="do something") + parser.add_argument("--prompt-len", type=int, default=16) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--discrete-state-input", + action="store_true", + help="Use realtime-vla tokenizer/state-token prompt path instead of precomputed language_embeds.", + ) + parser.add_argument( + "--tokenizer", + type=Path, + default=None, + help="Tokenizer path for --discrete-state-input.", + ) + parser.add_argument( + "--state-dim", + type=int, + default=8, + help="Synthetic state token count for --discrete-state-input.", + ) + + bnb.add_bench_cli_args(parser) + parser.set_defaults(bench_name="realtime_vla_pi05", batch_sizes=[1]) + add_profile_cli_args(parser) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for realtime-vla PI0.5 benchmarking") + + install_profiler(profile_config_from_args(args)) + runner = bnb.NBatchBenchRunner( + setup_fn=make_setup_fn(args), + extras_fn=make_extras_fn(args), + device=torch.device("cuda", torch.cuda.current_device()), + **bnb.bench_runner_kwargs_from_args(args), + ) + runner.run() + + +if __name__ == "__main__": + main() diff --git a/benchmark/pi05/bench_vlacpp_pi05_client.py b/benchmark/pi05/bench_vlacpp_pi05_client.py new file mode 100755 index 0000000..d3ff5aa --- /dev/null +++ b/benchmark/pi05/bench_vlacpp_pi05_client.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +"""vla.cpp PI0.5 ZMQ client benchmark using the PhyAI bench runner. + +Start ``vla-server`` separately, then run this client. The common PhyAI runner +handles warmup, timed iterations, JSONL output, and optional profiling. This +script forces CPU timing in the runner because the measured operation is a ZMQ +request to an external server process; CUDA events in the client process would +not cover server-side GPU work. + +Run:: + + python benchmark/pi05/bench_vlacpp_pi05_client.py \ + --vlacpp-root \ + --addr tcp://127.0.0.1:5555 \ + --tokenizer \ + --stats-json \ + --batch-sizes 1 --n-warmup 100 --n-timed 100 \ + --result-file results/vlacpp_pi05.jsonl + +Only batch size 1 is supported because the vla.cpp Python client sends one +request at a time. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import statistics +import sys +from typing import Any + +import numpy as np +import torch + +_BENCHMARK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_BENCHMARK_DIR)) +import bench_n_batch as bnb +from phyai.utils.profile import ( + add_profile_cli_args, + install_profiler, + profile_config_from_args, +) + + +def env_path(name: str) -> Path | None: + value = os.environ.get(name) + return Path(value) if value else None + + +def summarize(values: list[float]) -> dict[str, float] | None: + if not values: + return None + xs = sorted(float(x) for x in values) + return { + "count": len(xs), + "mean_ms": float(statistics.fmean(xs)), + "median_ms": float(statistics.median(xs)), + "p50_ms": float(np.percentile(xs, 50)), + "p90_ms": float(np.percentile(xs, 90)), + "p99_ms": float(np.percentile(xs, 99)), + "min_ms": xs[0], + "max_ms": xs[-1], + } + + +def make_obs(seed: int, num_views: int, prompt: str) -> dict[str, Any]: + """Create deterministic synthetic observations for latency-only requests. + + vla.cpp's official PI0.5 client expects CHW float32 images in [0, 1] and an + unnormalized robot state vector; it converts these into server protobufs. + """ + rng = np.random.default_rng(seed) + obs: dict[str, Any] = { + "observation.images.image": rng.random((3, 224, 224), dtype=np.float32), + "observation.state": rng.standard_normal(8).astype(np.float32), + "task": prompt, + "prompt": prompt, + } + if num_views >= 2: + obs["observation.images.image2"] = rng.random((3, 224, 224), dtype=np.float32) + if num_views >= 3: + obs["observation.images.image3"] = rng.random((3, 224, 224), dtype=np.float32) + return obs + + +def make_setup_fn(args: argparse.Namespace): + eval_root = args.vlacpp_root / "eval" + sys.path.insert(0, str(eval_root)) + sys.path.insert(0, str(eval_root / "client")) + from client.vla_cpp_client import VlaCppClient + + def setup_fn(batch_size: int) -> bnb.BenchSpec: + if batch_size != 1: + raise ValueError("vla.cpp PI0.5 wrapper supports only batch_size=1") + + image_keys = ["observation.images.image"] + if args.num_views >= 2: + image_keys.append("observation.images.image2") + if args.num_views >= 3: + image_keys.append("observation.images.image3") + + client = VlaCppClient( + vla_addr=args.addr, + arch=args.arch, + tokenizer_name=args.tokenizer, + image_keys=image_keys, + max_length=args.max_length, + real_action_dim=args.real_action_dim, + n_action_steps=1, + stats_json=args.stats_json, + ) + obs = make_obs(args.seed, args.num_views, args.prompt) + phase_samples: dict[str, list[float]] = { + "server_total_latency_ms": [], + "server_vision_latency_ms": [], + "server_inference_latency_ms": [], + "server_prefill_latency_ms": [], + "server_denoise_latency_ms": [], + } + phase_summary: dict[str, Any] = {} + call_count = 0 + + def update_phase_summary() -> None: + phase_summary.clear() + for key, values in phase_samples.items(): + phase_summary[key] = summarize(values) + + def step() -> None: + nonlocal call_count + client.get_action(obs) + call_count += 1 + if call_count <= args.n_warmup: + return + r = getattr(client, "_last_response", None) + if r is None: + return + phase_samples["server_total_latency_ms"].append(float(r.latency_ms_total)) + phase_samples["server_vision_latency_ms"].append(float(r.latency_ms_vision)) + phase_samples["server_inference_latency_ms"].append(float(r.latency_ms_inference)) + phase_samples["server_prefill_latency_ms"].append(float(r.latency_ms_prefill)) + phase_samples["server_denoise_latency_ms"].append(float(r.latency_ms_denoise)) + update_phase_summary() + + def teardown() -> None: + sock = getattr(client, "sock", None) + if sock is not None: + sock.close(linger=0) + + spec = bnb.BenchSpec( + name="vlacpp_pi05_zmq_client", + step_callable=step, + teardown_callable=teardown, + ) + # Attach dynamic summary for extras_fn. The runner copies the dict after + # timed steps finish, so it records the final server phase statistics. + spec.vlacpp_phase_summary = phase_summary # type: ignore[attr-defined] + return spec + + return setup_fn + + +def make_extras_fn(args: argparse.Namespace): + def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: + return { + "runtime": "vla.cpp", + "vlacpp_root": str(args.vlacpp_root), + "addr": args.addr, + "arch": args.arch, + "tokenizer": args.tokenizer, + "stats_json": str(args.stats_json) if args.stats_json else None, + "batch_size_contract": 1, + "num_views": args.num_views, + "chunk_size_metadata": args.chunk_size, + "prompt": args.prompt, + "seed": args.seed, + "server_phase_latency_ms": getattr(spec, "vlacpp_phase_summary", {}), + "timing_scope": "client ZMQ request wall time; server phase timings are copied from PredictResponse extras", + "notes": "vla.cpp server enforces the GGUF chunk size; prefix/expert may be combined in server phase timing.", + } + + return extras_fn + + +def main() -> None: + vlacpp_default = env_path("VLACPP_ROOT") + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--vlacpp-root", + type=Path, + default=vlacpp_default, + required=vlacpp_default is None, + help="Path to the vla.cpp repository clone. Can also be set with VLACPP_ROOT.", + ) + parser.add_argument("--addr", default="tcp://127.0.0.1:5555") + parser.add_argument("--arch", default="pi05") + parser.add_argument( + "--tokenizer", + default="google/paligemma-3b-pt-224", + help="Tokenizer name or local tokenizer path. Prefer a local path if the HF repo is gated.", + ) + parser.add_argument( + "--stats-json", + type=Path, + default=None, + help="Local LIBERO meta/stats.json for arch=pi05; avoids network fetch.", + ) + parser.add_argument("--num-views", type=int, default=2) + parser.add_argument("--chunk-size", type=int, default=50) + parser.add_argument("--prompt", default="do something") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--max-length", type=int, default=200) + parser.add_argument("--real-action-dim", type=int, default=7) + + bnb.add_bench_cli_args(parser) + parser.set_defaults(bench_name="vlacpp_pi05", batch_sizes=[1]) + add_profile_cli_args(parser) + args = parser.parse_args() + + install_profiler(profile_config_from_args(args)) + runner = bnb.NBatchBenchRunner( + setup_fn=make_setup_fn(args), + extras_fn=make_extras_fn(args), + # Force perf-counter timing: the GPU work happens in the vla-server process. + device=torch.device("cpu"), + **bnb.bench_runner_kwargs_from_args(args), + ) + runner.run() + + +if __name__ == "__main__": + main() From c3e0f94f8d0dbed7df0560a95ea45a2f29ff9f6d Mon Sep 17 00:00:00 2001 From: rebecca26358 <153866717+rebecca26358@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:25:12 +0000 Subject: [PATCH 10/29] style: format external pi05 benchmark wrapper --- benchmark/pi05/bench_vlacpp_pi05_client.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/benchmark/pi05/bench_vlacpp_pi05_client.py b/benchmark/pi05/bench_vlacpp_pi05_client.py index d3ff5aa..b29924b 100755 --- a/benchmark/pi05/bench_vlacpp_pi05_client.py +++ b/benchmark/pi05/bench_vlacpp_pi05_client.py @@ -137,9 +137,15 @@ def step() -> None: return phase_samples["server_total_latency_ms"].append(float(r.latency_ms_total)) phase_samples["server_vision_latency_ms"].append(float(r.latency_ms_vision)) - phase_samples["server_inference_latency_ms"].append(float(r.latency_ms_inference)) - phase_samples["server_prefill_latency_ms"].append(float(r.latency_ms_prefill)) - phase_samples["server_denoise_latency_ms"].append(float(r.latency_ms_denoise)) + phase_samples["server_inference_latency_ms"].append( + float(r.latency_ms_inference) + ) + phase_samples["server_prefill_latency_ms"].append( + float(r.latency_ms_prefill) + ) + phase_samples["server_denoise_latency_ms"].append( + float(r.latency_ms_denoise) + ) update_phase_summary() def teardown() -> None: From 493a9af70f4477680f75438dc2b71b2ca5ee64e8 Mon Sep 17 00:00:00 2001 From: rebecca26358 <153866717+rebecca26358@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:45:42 +0000 Subject: [PATCH 11/29] docs: add pi05 eight-gpu inference tutorial --- .../pi05/eight_gpu_inference_tutorial.en.md | 193 ++++++++++++++++++ 1 file changed, 193 insertions(+) create mode 100644 benchmark/pi05/eight_gpu_inference_tutorial.en.md diff --git a/benchmark/pi05/eight_gpu_inference_tutorial.en.md b/benchmark/pi05/eight_gpu_inference_tutorial.en.md new file mode 100644 index 0000000..ba6d6cd --- /dev/null +++ b/benchmark/pi05/eight_gpu_inference_tutorial.en.md @@ -0,0 +1,193 @@ +# PhyAI pi0.5 Eight-GPU Inference Tutorial + +Reproduce **PhyAI `pi05_wn` 8-GPU DP inference + concurrent LIBERO demos**. +Chinese: [`八卡推理从零开始教程.md`](./八卡推理从零开始教程.md) + +--- + +## 1. What you get + +```text +8-GPU PhyAI WebSocket server (port 8000) + ↑ +32 LIBERO shards (4 suites × 8 tasks) + ↓ +Per-shard JSON (success / timing) + optional wait videos +``` + +Recommended settings: + +| Item | Value | +| --- | --- | +| GPUs / batch | 8 GPUs, `MAX_BATCH_SIZE=32` (B=4 per GPU) | +| Chunk | `CHUNK_SIZE=10`, `SEND_ACTION_CHUNKS=1` (true chunk=10) | +| Batching wait | `MAX_WAIT_TIME=0.02` | +| Recording (optional) | `continuous` + 20fps | + +Architecture: all LIBERO clients talk to rank0; after the batch fills (or timeout), DP scatter → each GPU runs its slice → gather and reply. All 8 ranks sync in one step. CUDA graphs use a fixed padded shape, so a partial batch is not proportionally faster. + +Chunk modes: + +| Mode | Flag | Behavior | +| --- | --- | --- | +| True chunk=10 | `SEND_ACTION_CHUNKS=1` | Return 10 actions once; client runs them locally, then requests again | +| Pseudo chunk=1 | `SEND_ACTION_CHUNKS=0` | Model still produces 10; server returns 1 and buffers the rest; request every step | + +--- + +## 2. Setup + +Needs: 8 free GPUs, Docker, `tmux`. + +```bash +export WORKSPACE="$HOME/phyai_workplace" # change me: must contain phyai / vla-evaluation-harness / phyai_models +export PHYAI_ROOT="$WORKSPACE/phyai" +export VLA_ROOT="$WORKSPACE/vla-evaluation-harness" +export MODEL_ROOT="$WORKSPACE/phyai_models" +export DEMO_ROOT="$PHYAI_ROOT/libero_wn_demo" +``` + +Checks: + +```bash +ls "$MODEL_ROOT/pi05_libero_phyai_converted" "$MODEL_ROOT/paligemma-3b-pt-224" +ls "$PHYAI_ROOT/.venv/bin/torchrun" "$VLA_ROOT/.venv/bin/vla-eval" +``` + +Images: + +```bash +sg docker -c 'docker pull nvcr.io/nvidia/pytorch:25.12-py3' +sg docker -c 'docker pull ghcr.io/allenai/vla-evaluation-harness/libero:latest' +``` + +Tokenizer must be offline: place `$MODEL_ROOT/paligemma-3b-pt-224` first (`HF_HUB_OFFLINE=1` is set by the script). +Always pass `PHYAI_ROOT` / `VLA_ROOT` / `MODEL_ROOT` explicitly; do not copy another machine’s absolute paths. + +--- + +## 3. Shortest path (batch32 + true chunk10) + +### A. Confirm idle + +```bash +nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader,nounits +ss -ltnp | grep 8000 || echo '8000 free' +``` + +Cleanup leftovers if needed: + +```bash +tmux kill-session -t phyai_pi05_wn_demo 2>/dev/null || true +sg docker -c 'docker rm -f phyai_pi05_wn_demo' 2>/dev/null || true +``` + +### B. Start server + +```bash +sg docker -c " +SESSION=phyai_pi05_wn_demo CONTAINER=phyai_pi05_wn_demo \ +PHYAI_ROOT=$PHYAI_ROOT VLA_ROOT=$VLA_ROOT MODEL_ROOT=$MODEL_ROOT \ +MAX_BATCH_SIZE=32 MAX_WAIT_TIME=0.02 CHUNK_SIZE=10 SEND_ACTION_CHUNKS=1 \ +MASTER_PORT=29619 \ +bash $DEMO_ROOT/start_pi05_wn_server.sh +" +``` + +After ~1–3 minutes: + +```bash +curl -sS http://127.0.0.1:8000/config +# expect max_batch_size=32 +tmux attach -t phyai_pi05_wn_demo # detach: Ctrl-b d +``` + +### C. Run 32 shards + +```bash +sg docker -c "bash $DEMO_ROOT/run_libero_32_clientchunk10_maxwait002.sh" +``` + +Outputs under `$VLA_ROOT/results/clientchunk10_maxwait002/`: suite `*.json` and `wait_videos/*.mp4` (pixels are only in the videos). + +### D. Cleanup + +```bash +tmux kill-session -t phyai_pi05_wn_demo +sg docker -c 'docker rm -f phyai_pi05_wn_demo' 2>/dev/null || true +``` + +--- + +## 4. Variants + +**Pseudo chunk=1 (batch32)** + +```bash +# same server command with SEND_ACTION_CHUNKS=0 and new SESSION/CONTAINER/MASTER_PORT +sg docker -c "bash $DEMO_ROOT/run_libero_32_serverchunk1_maxwait002.sh" +``` + +**batch16 (B=2 per GPU, 16 shards)** + +```bash +# MAX_BATCH_SIZE=16, SEND_ACTION_CHUNKS=0 or 1 +sg docker -c "bash $DEMO_ROOT/run_libero_batch16_serverchunk1_maxwait002.sh" +``` + +Start batch16 and batch32 as separate servers; do not switch inside one `torchrun`. + +**Custom experiment**: copy `$DEMO_ROOT/clientchunk10_maxwait002/`, edit yaml `output_dir` / suite, then point the client script’s `DEMO_ROOT` and `RESULTS_ROOT` at your dirs. + +**Pure inference latency (no LIBERO)**: + +```bash +sg docker -c "bash $PHYAI_ROOT/benchmark/run_pi05_wn_latency_dp8_docker.sh" +``` + +Reference: batch16 ~60ms, batch32 ~97ms (pure `Engine.step`, not end-to-end). + +--- + +## 5. Recording and results + +For end-to-end stalls use `continuous`, not `step`: + +```yaml +docker: + env: + - VLA_EVAL_WAIT_VIDEO_DIR=/workspace/results/wait_videos + - VLA_EVAL_WAIT_VIDEO_MODE=continuous + - VLA_EVAL_WAIT_VIDEO_FPS=20 + - VLA_EVAL_WAIT_VIDEO_MODEL= +``` + +`continuous` freezes while waiting on the model; freeze length ≈ real wait. `step` removes waits. Changing export fps does not shorten real stalls. + +Useful JSON fields: `metrics.success`, `avg_model_wait_ms`, `model_buffer_hits`, `model_inference_calls`, `wait_video_mode`. +End-to-end wait ≈ queue + predict_batch + ws, often much larger than pure GPU bench. + +--- + +## 6. Troubleshooting + +| Symptom | Fix | +| --- | --- | +| Stuck in setup | Offline tokenizer; change `MASTER_PORT`; remove same-name container; free GPUs | +| `:8000/config` fails | Wait for graph capture; `tmux capture-pane -t -p -S -80` | +| Shard cannot connect | Host networking; URL=`ws://127.0.0.1:8000`; `NO_PROXY='*'` | +| Partial results | Check `$VLA_ROOT/results//*_logs/`; `taskunknown` ≈ task0 | +| FlashInfer/CUDA clash | Use Docker; avoid bare-metal host runs | + +--- + +## 7. File index (under `$DEMO_ROOT`) + +| File | Purpose | +| --- | --- | +| `start_pi05_wn_server.sh` | 8-GPU server | +| `run_libero_32_clientchunk10_maxwait002.sh` | batch32 true chunk10 | +| `run_libero_32_serverchunk1_maxwait002.sh` | batch32 pseudo chunk=1 | +| `run_libero_batch16_serverchunk1_maxwait002.sh` | batch16 pseudo chunk=1 | +| `clientchunk10_maxwait002/` etc. | continuous demo configs | +| `experiment_setup.md` | Historical experiment notes | From 48e2623761aa72d3efe1e07206a2a5fb8d2bf86a Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 04:17:28 +0000 Subject: [PATCH 12/29] benchmark: remove flashrt pi05 latency wrapper --- benchmark/pi05/bench_flashrt_pi05.py | 230 --------------------------- 1 file changed, 230 deletions(-) delete mode 100755 benchmark/pi05/bench_flashrt_pi05.py diff --git a/benchmark/pi05/bench_flashrt_pi05.py b/benchmark/pi05/bench_flashrt_pi05.py deleted file mode 100755 index c045949..0000000 --- a/benchmark/pi05/bench_flashrt_pi05.py +++ /dev/null @@ -1,230 +0,0 @@ -#!/usr/bin/env python3 -"""FlashRT PI0.5 latency benchmark using the PhyAI bench runner. - -This is a thin adapter around FlashRT's direct ``Pi05TorchFrontendRtx`` path. -The direct frontend is used because action chunk size is a frontend constructor -argument; ``flash_rt.load_model(..., num_steps=...)`` controls denoise steps, -not action chunk size. It mirrors ``benchmark/bench_n_batch_ws1_pi05.py``: -framework-specific setup lives here, while warmup, timed iterations, JSONL -output, and optional profiling are handled by ``benchmark/bench_n_batch.py``. - -Run:: - - python benchmark/pi05/bench_flashrt_pi05.py \ - --flashrt-root \ - --checkpoint \ - --batch-sizes 1 --n-warmup 100 --n-timed 100 \ - --result-file results/flashrt_pi05.jsonl - -Only batch size 1 is supported because the FlashRT PI0.5 direct frontend path -used here takes one robot request at a time. -""" - -from __future__ import annotations - -import argparse -import os -from pathlib import Path -import statistics -import sys -from typing import Any - -import numpy as np -import torch - -_BENCHMARK_DIR = Path(__file__).resolve().parents[1] -sys.path.insert(0, str(_BENCHMARK_DIR)) -import bench_n_batch as bnb # noqa: E402 -from phyai.utils.profile import ( # noqa: E402 - add_profile_cli_args, - install_profiler, - profile_config_from_args, -) - - -def env_path(name: str) -> Path | None: - value = os.environ.get(name) - return Path(value) if value else None - - -class LazySummary(dict): - def __init__(self, values_fn_or_items): - if callable(values_fn_or_items): - super().__init__() - self._values_fn = values_fn_or_items - else: - super().__init__(values_fn_or_items) - self._values_fn = None - - def items(self): - if self._values_fn is None: - return super().items() - summary = summarize(self._values_fn()) - return (summary or {}).items() - - -def summarize(values: list[float]) -> dict[str, float] | None: - if not values: - return None - xs = sorted(float(x) for x in values) - return { - "count": len(xs), - "mean_ms": float(statistics.fmean(xs)), - "median_ms": float(statistics.median(xs)), - "p50_ms": float(np.percentile(xs, 50)), - "p90_ms": float(np.percentile(xs, 90)), - "p99_ms": float(np.percentile(xs, 99)), - "min_ms": xs[0], - "max_ms": xs[-1], - } - - -def make_observation(num_views: int, seed: int, prompt: str) -> dict[str, Any]: - """Deterministic synthetic observation for latency-only runs.""" - rng = np.random.default_rng(seed) - obs: dict[str, Any] = { - "image": rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8), - "state": rng.standard_normal(8).astype(np.float32), - "task": prompt, - "prompt": prompt, - } - if num_views >= 2: - obs["wrist_image"] = rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8) - if num_views >= 3: - obs["wrist_image_right"] = rng.integers( - 0, 256, size=(224, 224, 3), dtype=np.uint8 - ) - return obs - - -def import_flashrt_frontend(repo: Path): - # Use the checked-out FlashRT repository directly; installation is optional. - sys.path.insert(0, str(repo)) - from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx - - return Pi05TorchFrontendRtx - - -def make_setup_fn(args: argparse.Namespace): - frontend_cls = import_flashrt_frontend(args.flashrt_root) - - def setup_fn(batch_size: int) -> bnb.BenchSpec: - if batch_size != 1: - raise ValueError("FlashRT PI0.5 wrapper supports only batch_size=1") - - if args.precision == "bf16": - # FlashRT uses this environment switch to force the BF16 PI0.5 RTX path. - os.environ["FVK_PI05_RTX_FORCE_BF16"] = "1" - else: - os.environ.pop("FVK_PI05_RTX_FORCE_BF16", None) - - model = frontend_cls( - args.checkpoint, - num_views=args.num_views, - chunk_size=args.chunk_size, - cache_frames=1, - use_fp8=(args.precision == "fp8_bf16"), - hardware=args.hardware, - ) - obs = make_observation(args.num_views, args.seed, args.prompt) - - # set_prompt builds the prompt-specific pipeline; calibration captures the graph. - # These setup calls are outside the measured latency window. - model.set_prompt(args.prompt) - model.calibrate_with_real_data([obs]) - model.infer(obs) - torch.cuda.synchronize() - - call_count = 0 - - def step() -> None: - nonlocal call_count - if call_count == args.n_warmup: - model.latency_records.clear() - model.infer(obs) - # FlashRT PI0.5 may use an internal CUDA stream, so synchronize inside the - # step to make the common runner's timing boundary conservative. - torch.cuda.synchronize() - call_count += 1 - - spec = bnb.BenchSpec( - name="flashrt_pi05", - step_callable=step, - teardown_callable=lambda: None, - ) - spec.flashrt_internal_latency_ms = LazySummary( - lambda: [float(x) for x in getattr(model, "latency_records", [])] - ) # type: ignore[attr-defined] - return spec - - return setup_fn - - -def make_extras_fn(args: argparse.Namespace): - def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: - return { - "runtime": "FlashRT", - "checkpoint": str(args.checkpoint), - "flashrt_root": str(args.flashrt_root), - "precision": args.precision, - "hardware": args.hardware, - "batch_size_contract": 1, - "num_views": args.num_views, - "chunk_size": args.chunk_size, - "prompt": args.prompt, - "seed": args.seed, - "internal_latency_ms": getattr(spec, "flashrt_internal_latency_ms", {}), - "timing_scope": "FlashRT direct Pi05TorchFrontendRtx.infer hot path after set_prompt and first graph-building infer; common runner uses perf-counter wall time because FlashRT runs work on an internal CUDA stream", - } - - return extras_fn - - -def main() -> None: - flashrt_default = env_path("FLASHRT_ROOT") - parser = argparse.ArgumentParser( - description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - parser.add_argument( - "--flashrt-root", - type=Path, - default=flashrt_default, - required=flashrt_default is None, - help="Path to the FlashRT repository clone. Can also be set with FLASHRT_ROOT.", - ) - parser.add_argument( - "--checkpoint", - type=Path, - required=True, - help="PI0.5 checkpoint directory readable by FlashRT.", - ) - parser.add_argument("--num-views", type=int, default=2) - parser.add_argument("--chunk-size", type=int, default=50) - parser.add_argument("--prompt", default="do something") - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--precision", choices=("bf16", "fp8_bf16"), default="bf16") - parser.add_argument("--hardware", default="auto") - - bnb.add_bench_cli_args(parser) - parser.set_defaults(bench_name="flashrt_pi05", batch_sizes=[1]) - add_profile_cli_args(parser) - args = parser.parse_args() - - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required for FlashRT PI0.5 benchmarking") - - install_profiler(profile_config_from_args(args)) - runner = bnb.NBatchBenchRunner( - setup_fn=make_setup_fn(args), - extras_fn=make_extras_fn(args), - # FlashRT may execute on an internal CUDA stream. Force the common - # runner's perf-counter path; step() synchronizes CUDA before returning. - device=torch.device("cpu"), - **bnb.bench_runner_kwargs_from_args(args), - ) - runner.run() - - -if __name__ == "__main__": - main() From 694cea90b039054a02e0433451c356a8699c4c39 Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 04:17:28 +0000 Subject: [PATCH 13/29] benchmark: add flashrt pi05 latency wrapper --- benchmark/pi05/bench_flashrt_pi05.py | 230 +++++++++++++++++++++++++++ 1 file changed, 230 insertions(+) create mode 100755 benchmark/pi05/bench_flashrt_pi05.py diff --git a/benchmark/pi05/bench_flashrt_pi05.py b/benchmark/pi05/bench_flashrt_pi05.py new file mode 100755 index 0000000..c045949 --- /dev/null +++ b/benchmark/pi05/bench_flashrt_pi05.py @@ -0,0 +1,230 @@ +#!/usr/bin/env python3 +"""FlashRT PI0.5 latency benchmark using the PhyAI bench runner. + +This is a thin adapter around FlashRT's direct ``Pi05TorchFrontendRtx`` path. +The direct frontend is used because action chunk size is a frontend constructor +argument; ``flash_rt.load_model(..., num_steps=...)`` controls denoise steps, +not action chunk size. It mirrors ``benchmark/bench_n_batch_ws1_pi05.py``: +framework-specific setup lives here, while warmup, timed iterations, JSONL +output, and optional profiling are handled by ``benchmark/bench_n_batch.py``. + +Run:: + + python benchmark/pi05/bench_flashrt_pi05.py \ + --flashrt-root \ + --checkpoint \ + --batch-sizes 1 --n-warmup 100 --n-timed 100 \ + --result-file results/flashrt_pi05.jsonl + +Only batch size 1 is supported because the FlashRT PI0.5 direct frontend path +used here takes one robot request at a time. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import statistics +import sys +from typing import Any + +import numpy as np +import torch + +_BENCHMARK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_BENCHMARK_DIR)) +import bench_n_batch as bnb # noqa: E402 +from phyai.utils.profile import ( # noqa: E402 + add_profile_cli_args, + install_profiler, + profile_config_from_args, +) + + +def env_path(name: str) -> Path | None: + value = os.environ.get(name) + return Path(value) if value else None + + +class LazySummary(dict): + def __init__(self, values_fn_or_items): + if callable(values_fn_or_items): + super().__init__() + self._values_fn = values_fn_or_items + else: + super().__init__(values_fn_or_items) + self._values_fn = None + + def items(self): + if self._values_fn is None: + return super().items() + summary = summarize(self._values_fn()) + return (summary or {}).items() + + +def summarize(values: list[float]) -> dict[str, float] | None: + if not values: + return None + xs = sorted(float(x) for x in values) + return { + "count": len(xs), + "mean_ms": float(statistics.fmean(xs)), + "median_ms": float(statistics.median(xs)), + "p50_ms": float(np.percentile(xs, 50)), + "p90_ms": float(np.percentile(xs, 90)), + "p99_ms": float(np.percentile(xs, 99)), + "min_ms": xs[0], + "max_ms": xs[-1], + } + + +def make_observation(num_views: int, seed: int, prompt: str) -> dict[str, Any]: + """Deterministic synthetic observation for latency-only runs.""" + rng = np.random.default_rng(seed) + obs: dict[str, Any] = { + "image": rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8), + "state": rng.standard_normal(8).astype(np.float32), + "task": prompt, + "prompt": prompt, + } + if num_views >= 2: + obs["wrist_image"] = rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8) + if num_views >= 3: + obs["wrist_image_right"] = rng.integers( + 0, 256, size=(224, 224, 3), dtype=np.uint8 + ) + return obs + + +def import_flashrt_frontend(repo: Path): + # Use the checked-out FlashRT repository directly; installation is optional. + sys.path.insert(0, str(repo)) + from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx + + return Pi05TorchFrontendRtx + + +def make_setup_fn(args: argparse.Namespace): + frontend_cls = import_flashrt_frontend(args.flashrt_root) + + def setup_fn(batch_size: int) -> bnb.BenchSpec: + if batch_size != 1: + raise ValueError("FlashRT PI0.5 wrapper supports only batch_size=1") + + if args.precision == "bf16": + # FlashRT uses this environment switch to force the BF16 PI0.5 RTX path. + os.environ["FVK_PI05_RTX_FORCE_BF16"] = "1" + else: + os.environ.pop("FVK_PI05_RTX_FORCE_BF16", None) + + model = frontend_cls( + args.checkpoint, + num_views=args.num_views, + chunk_size=args.chunk_size, + cache_frames=1, + use_fp8=(args.precision == "fp8_bf16"), + hardware=args.hardware, + ) + obs = make_observation(args.num_views, args.seed, args.prompt) + + # set_prompt builds the prompt-specific pipeline; calibration captures the graph. + # These setup calls are outside the measured latency window. + model.set_prompt(args.prompt) + model.calibrate_with_real_data([obs]) + model.infer(obs) + torch.cuda.synchronize() + + call_count = 0 + + def step() -> None: + nonlocal call_count + if call_count == args.n_warmup: + model.latency_records.clear() + model.infer(obs) + # FlashRT PI0.5 may use an internal CUDA stream, so synchronize inside the + # step to make the common runner's timing boundary conservative. + torch.cuda.synchronize() + call_count += 1 + + spec = bnb.BenchSpec( + name="flashrt_pi05", + step_callable=step, + teardown_callable=lambda: None, + ) + spec.flashrt_internal_latency_ms = LazySummary( + lambda: [float(x) for x in getattr(model, "latency_records", [])] + ) # type: ignore[attr-defined] + return spec + + return setup_fn + + +def make_extras_fn(args: argparse.Namespace): + def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: + return { + "runtime": "FlashRT", + "checkpoint": str(args.checkpoint), + "flashrt_root": str(args.flashrt_root), + "precision": args.precision, + "hardware": args.hardware, + "batch_size_contract": 1, + "num_views": args.num_views, + "chunk_size": args.chunk_size, + "prompt": args.prompt, + "seed": args.seed, + "internal_latency_ms": getattr(spec, "flashrt_internal_latency_ms", {}), + "timing_scope": "FlashRT direct Pi05TorchFrontendRtx.infer hot path after set_prompt and first graph-building infer; common runner uses perf-counter wall time because FlashRT runs work on an internal CUDA stream", + } + + return extras_fn + + +def main() -> None: + flashrt_default = env_path("FLASHRT_ROOT") + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--flashrt-root", + type=Path, + default=flashrt_default, + required=flashrt_default is None, + help="Path to the FlashRT repository clone. Can also be set with FLASHRT_ROOT.", + ) + parser.add_argument( + "--checkpoint", + type=Path, + required=True, + help="PI0.5 checkpoint directory readable by FlashRT.", + ) + parser.add_argument("--num-views", type=int, default=2) + parser.add_argument("--chunk-size", type=int, default=50) + parser.add_argument("--prompt", default="do something") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--precision", choices=("bf16", "fp8_bf16"), default="bf16") + parser.add_argument("--hardware", default="auto") + + bnb.add_bench_cli_args(parser) + parser.set_defaults(bench_name="flashrt_pi05", batch_sizes=[1]) + add_profile_cli_args(parser) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for FlashRT PI0.5 benchmarking") + + install_profiler(profile_config_from_args(args)) + runner = bnb.NBatchBenchRunner( + setup_fn=make_setup_fn(args), + extras_fn=make_extras_fn(args), + # FlashRT may execute on an internal CUDA stream. Force the common + # runner's perf-counter path; step() synchronizes CUDA before returning. + device=torch.device("cpu"), + **bnb.bench_runner_kwargs_from_args(args), + ) + runner.run() + + +if __name__ == "__main__": + main() From 472b5da07f960d40298dbf9f958848892095b90f Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 04:24:39 +0000 Subject: [PATCH 14/29] benchmark: remove remaining external pi05 runtime wrappers --- benchmark/pi05/bench_realtime_vla_pi05.py | 225 ------------------- benchmark/pi05/bench_vlacpp_pi05_client.py | 241 --------------------- 2 files changed, 466 deletions(-) delete mode 100755 benchmark/pi05/bench_realtime_vla_pi05.py delete mode 100755 benchmark/pi05/bench_vlacpp_pi05_client.py diff --git a/benchmark/pi05/bench_realtime_vla_pi05.py b/benchmark/pi05/bench_realtime_vla_pi05.py deleted file mode 100755 index f91a6e4..0000000 --- a/benchmark/pi05/bench_realtime_vla_pi05.py +++ /dev/null @@ -1,225 +0,0 @@ -#!/usr/bin/env python3 -"""realtime-vla PI0.5 latency benchmark using the PhyAI bench runner. - -This script is a thin adapter around realtime-vla's ``Pi05Inference.forward``. -It mirrors ``benchmark/bench_n_batch_ws1_pi05.py`` by delegating warmup, timed -iterations, JSONL output, and optional profiling to ``benchmark/bench_n_batch.py``. - -Run:: - - python benchmark/pi05/bench_realtime_vla_pi05.py \ - --realtime-vla-root \ - --flashrt-root \ - --checkpoint \ - --batch-sizes 1 --n-warmup 100 --n-timed 100 \ - --result-file results/realtime_vla_pi05.jsonl - -Only batch size 1 is supported because realtime-vla's PI0.5 inference object -used here is allocated for one synthetic request. -""" - -from __future__ import annotations - -import argparse -import os -from pathlib import Path -import pickle -import sys -from typing import Any - -import torch - -_BENCHMARK_DIR = Path(__file__).resolve().parents[1] -sys.path.insert(0, str(_BENCHMARK_DIR)) -import bench_n_batch as bnb # noqa: E402 -from phyai.utils.profile import ( # noqa: E402 - add_profile_cli_args, - install_profiler, - profile_config_from_args, -) - - -def env_path(name: str) -> Path | None: - value = os.environ.get(name) - return Path(value) if value else None - - -def load_checkpoint(path: Path, flashrt_root: Path | None, trust_pickle: bool): - if path.is_dir(): - path = path / "model.safetensors" - if path.suffix == ".safetensors": - if flashrt_root is None: - raise ValueError( - "--flashrt-root or FLASHRT_ROOT is required when loading a safetensors checkpoint" - ) - # Reuse FlashRT's tested PI0.5 key conversion instead of duplicating it here. - sys.path.insert(0, str(flashrt_root)) - from flash_rt.frontends.torch.pi05_rtx import convert_pi05_safetensors - - return convert_pi05_safetensors(path) - if path.suffix in {".pt", ".pth"}: - return torch.load(path, map_location="cpu", weights_only=True) - if path.suffix in {".pkl", ".pickle"}: - if not trust_pickle: - raise ValueError( - "Refusing to load pickle checkpoint without --trust-pickle-checkpoint. " - "Only use that flag for checkpoints from a trusted source." - ) - with path.open("rb") as f: - return pickle.load(f) # nosec B301: guarded by --trust-pickle-checkpoint. - raise ValueError( - f"Unsupported checkpoint suffix {path.suffix!r}; expected .safetensors, .pt, .pth, .pkl, or .pickle" - ) - - -def make_setup_fn(args: argparse.Namespace): - # Use the checked-out realtime-vla repository directly; installation is optional. - sys.path.insert(0, str(args.realtime_vla_root)) - from pi05_infer import Pi05Inference - - def setup_fn(batch_size: int) -> bnb.BenchSpec: - if batch_size != 1: - raise ValueError("realtime-vla PI0.5 wrapper supports only batch_size=1") - - checkpoint = load_checkpoint( - args.checkpoint, args.flashrt_root, args.trust_pickle_checkpoint - ) - if "language_embeds" not in checkpoint: - # Latency-only fallback for checkpoints that do not include prompt embeds. - checkpoint["language_embeds"] = torch.randn( - args.prompt_len, 2048, dtype=torch.bfloat16 - ) - - infer = Pi05Inference( - checkpoint=checkpoint, - num_views=args.num_views, - chunk_size=args.chunk_size, - tokenizer_path=str(args.tokenizer) if args.tokenizer else None, - discrete_state_input=args.discrete_state_input, - ) - torch.manual_seed(args.seed) - input_image = torch.randn( - args.num_views, 224, 224, 3, dtype=torch.bfloat16, device="cuda" - ) - input_noise = torch.randn( - args.chunk_size, 32, dtype=torch.bfloat16, device="cuda" - ) - - state_tokens = None - if args.discrete_state_input: - import numpy as np - - state_tokens = np.zeros(args.state_dim, dtype=np.int64) - - def step() -> None: - infer.forward( - input_image, - input_noise, - task_prompt=args.prompt, - state_tokens=state_tokens, - ) - - return bnb.BenchSpec( - name="realtime_vla_pi05", - step_callable=step, - teardown_callable=lambda: None, - ) - - return setup_fn - - -def make_extras_fn(args: argparse.Namespace): - def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: - return { - "runtime": "realtime-vla", - "checkpoint": str(args.checkpoint), - "realtime_vla_root": str(args.realtime_vla_root), - "flashrt_root": str(args.flashrt_root) if args.flashrt_root else None, - "precision": "bf16", - "batch_size_contract": 1, - "num_views": args.num_views, - "chunk_size": args.chunk_size, - "prompt": args.prompt, - "prompt_len": args.prompt_len, - "seed": args.seed, - "discrete_state_input": args.discrete_state_input, - "timing_scope": "Pi05Inference.forward hot path; common runner CUDA event wraps one forward call", - } - - return extras_fn - - -def main() -> None: - realtime_default = env_path("REALTIME_VLA_ROOT") - flashrt_default = env_path("FLASHRT_ROOT") - parser = argparse.ArgumentParser( - description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - parser.add_argument( - "--realtime-vla-root", - type=Path, - default=realtime_default, - required=realtime_default is None, - help="Path to the realtime-vla repository clone. Can also be set with REALTIME_VLA_ROOT.", - ) - parser.add_argument( - "--flashrt-root", - type=Path, - default=flashrt_default, - help="Path to FlashRT. Required only when --checkpoint is a PI0.5 safetensors checkpoint.", - ) - parser.add_argument( - "--checkpoint", - type=Path, - required=True, - help="realtime-vla .pkl/.pt checkpoint, PI0.5 model.safetensors, or a directory containing model.safetensors.", - ) - parser.add_argument( - "--trust-pickle-checkpoint", - action="store_true", - help="Allow loading .pkl/.pickle checkpoints. Only use with trusted checkpoint files.", - ) - parser.add_argument("--num-views", type=int, default=2) - parser.add_argument("--chunk-size", type=int, default=50) - parser.add_argument("--prompt", default="do something") - parser.add_argument("--prompt-len", type=int, default=16) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument( - "--discrete-state-input", - action="store_true", - help="Use realtime-vla tokenizer/state-token prompt path instead of precomputed language_embeds.", - ) - parser.add_argument( - "--tokenizer", - type=Path, - default=None, - help="Tokenizer path for --discrete-state-input.", - ) - parser.add_argument( - "--state-dim", - type=int, - default=8, - help="Synthetic state token count for --discrete-state-input.", - ) - - bnb.add_bench_cli_args(parser) - parser.set_defaults(bench_name="realtime_vla_pi05", batch_sizes=[1]) - add_profile_cli_args(parser) - args = parser.parse_args() - - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required for realtime-vla PI0.5 benchmarking") - - install_profiler(profile_config_from_args(args)) - runner = bnb.NBatchBenchRunner( - setup_fn=make_setup_fn(args), - extras_fn=make_extras_fn(args), - device=torch.device("cuda", torch.cuda.current_device()), - **bnb.bench_runner_kwargs_from_args(args), - ) - runner.run() - - -if __name__ == "__main__": - main() diff --git a/benchmark/pi05/bench_vlacpp_pi05_client.py b/benchmark/pi05/bench_vlacpp_pi05_client.py deleted file mode 100755 index b29924b..0000000 --- a/benchmark/pi05/bench_vlacpp_pi05_client.py +++ /dev/null @@ -1,241 +0,0 @@ -#!/usr/bin/env python3 -"""vla.cpp PI0.5 ZMQ client benchmark using the PhyAI bench runner. - -Start ``vla-server`` separately, then run this client. The common PhyAI runner -handles warmup, timed iterations, JSONL output, and optional profiling. This -script forces CPU timing in the runner because the measured operation is a ZMQ -request to an external server process; CUDA events in the client process would -not cover server-side GPU work. - -Run:: - - python benchmark/pi05/bench_vlacpp_pi05_client.py \ - --vlacpp-root \ - --addr tcp://127.0.0.1:5555 \ - --tokenizer \ - --stats-json \ - --batch-sizes 1 --n-warmup 100 --n-timed 100 \ - --result-file results/vlacpp_pi05.jsonl - -Only batch size 1 is supported because the vla.cpp Python client sends one -request at a time. -""" - -from __future__ import annotations - -import argparse -import os -from pathlib import Path -import statistics -import sys -from typing import Any - -import numpy as np -import torch - -_BENCHMARK_DIR = Path(__file__).resolve().parents[1] -sys.path.insert(0, str(_BENCHMARK_DIR)) -import bench_n_batch as bnb -from phyai.utils.profile import ( - add_profile_cli_args, - install_profiler, - profile_config_from_args, -) - - -def env_path(name: str) -> Path | None: - value = os.environ.get(name) - return Path(value) if value else None - - -def summarize(values: list[float]) -> dict[str, float] | None: - if not values: - return None - xs = sorted(float(x) for x in values) - return { - "count": len(xs), - "mean_ms": float(statistics.fmean(xs)), - "median_ms": float(statistics.median(xs)), - "p50_ms": float(np.percentile(xs, 50)), - "p90_ms": float(np.percentile(xs, 90)), - "p99_ms": float(np.percentile(xs, 99)), - "min_ms": xs[0], - "max_ms": xs[-1], - } - - -def make_obs(seed: int, num_views: int, prompt: str) -> dict[str, Any]: - """Create deterministic synthetic observations for latency-only requests. - - vla.cpp's official PI0.5 client expects CHW float32 images in [0, 1] and an - unnormalized robot state vector; it converts these into server protobufs. - """ - rng = np.random.default_rng(seed) - obs: dict[str, Any] = { - "observation.images.image": rng.random((3, 224, 224), dtype=np.float32), - "observation.state": rng.standard_normal(8).astype(np.float32), - "task": prompt, - "prompt": prompt, - } - if num_views >= 2: - obs["observation.images.image2"] = rng.random((3, 224, 224), dtype=np.float32) - if num_views >= 3: - obs["observation.images.image3"] = rng.random((3, 224, 224), dtype=np.float32) - return obs - - -def make_setup_fn(args: argparse.Namespace): - eval_root = args.vlacpp_root / "eval" - sys.path.insert(0, str(eval_root)) - sys.path.insert(0, str(eval_root / "client")) - from client.vla_cpp_client import VlaCppClient - - def setup_fn(batch_size: int) -> bnb.BenchSpec: - if batch_size != 1: - raise ValueError("vla.cpp PI0.5 wrapper supports only batch_size=1") - - image_keys = ["observation.images.image"] - if args.num_views >= 2: - image_keys.append("observation.images.image2") - if args.num_views >= 3: - image_keys.append("observation.images.image3") - - client = VlaCppClient( - vla_addr=args.addr, - arch=args.arch, - tokenizer_name=args.tokenizer, - image_keys=image_keys, - max_length=args.max_length, - real_action_dim=args.real_action_dim, - n_action_steps=1, - stats_json=args.stats_json, - ) - obs = make_obs(args.seed, args.num_views, args.prompt) - phase_samples: dict[str, list[float]] = { - "server_total_latency_ms": [], - "server_vision_latency_ms": [], - "server_inference_latency_ms": [], - "server_prefill_latency_ms": [], - "server_denoise_latency_ms": [], - } - phase_summary: dict[str, Any] = {} - call_count = 0 - - def update_phase_summary() -> None: - phase_summary.clear() - for key, values in phase_samples.items(): - phase_summary[key] = summarize(values) - - def step() -> None: - nonlocal call_count - client.get_action(obs) - call_count += 1 - if call_count <= args.n_warmup: - return - r = getattr(client, "_last_response", None) - if r is None: - return - phase_samples["server_total_latency_ms"].append(float(r.latency_ms_total)) - phase_samples["server_vision_latency_ms"].append(float(r.latency_ms_vision)) - phase_samples["server_inference_latency_ms"].append( - float(r.latency_ms_inference) - ) - phase_samples["server_prefill_latency_ms"].append( - float(r.latency_ms_prefill) - ) - phase_samples["server_denoise_latency_ms"].append( - float(r.latency_ms_denoise) - ) - update_phase_summary() - - def teardown() -> None: - sock = getattr(client, "sock", None) - if sock is not None: - sock.close(linger=0) - - spec = bnb.BenchSpec( - name="vlacpp_pi05_zmq_client", - step_callable=step, - teardown_callable=teardown, - ) - # Attach dynamic summary for extras_fn. The runner copies the dict after - # timed steps finish, so it records the final server phase statistics. - spec.vlacpp_phase_summary = phase_summary # type: ignore[attr-defined] - return spec - - return setup_fn - - -def make_extras_fn(args: argparse.Namespace): - def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: - return { - "runtime": "vla.cpp", - "vlacpp_root": str(args.vlacpp_root), - "addr": args.addr, - "arch": args.arch, - "tokenizer": args.tokenizer, - "stats_json": str(args.stats_json) if args.stats_json else None, - "batch_size_contract": 1, - "num_views": args.num_views, - "chunk_size_metadata": args.chunk_size, - "prompt": args.prompt, - "seed": args.seed, - "server_phase_latency_ms": getattr(spec, "vlacpp_phase_summary", {}), - "timing_scope": "client ZMQ request wall time; server phase timings are copied from PredictResponse extras", - "notes": "vla.cpp server enforces the GGUF chunk size; prefix/expert may be combined in server phase timing.", - } - - return extras_fn - - -def main() -> None: - vlacpp_default = env_path("VLACPP_ROOT") - parser = argparse.ArgumentParser( - description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - parser.add_argument( - "--vlacpp-root", - type=Path, - default=vlacpp_default, - required=vlacpp_default is None, - help="Path to the vla.cpp repository clone. Can also be set with VLACPP_ROOT.", - ) - parser.add_argument("--addr", default="tcp://127.0.0.1:5555") - parser.add_argument("--arch", default="pi05") - parser.add_argument( - "--tokenizer", - default="google/paligemma-3b-pt-224", - help="Tokenizer name or local tokenizer path. Prefer a local path if the HF repo is gated.", - ) - parser.add_argument( - "--stats-json", - type=Path, - default=None, - help="Local LIBERO meta/stats.json for arch=pi05; avoids network fetch.", - ) - parser.add_argument("--num-views", type=int, default=2) - parser.add_argument("--chunk-size", type=int, default=50) - parser.add_argument("--prompt", default="do something") - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--max-length", type=int, default=200) - parser.add_argument("--real-action-dim", type=int, default=7) - - bnb.add_bench_cli_args(parser) - parser.set_defaults(bench_name="vlacpp_pi05", batch_sizes=[1]) - add_profile_cli_args(parser) - args = parser.parse_args() - - install_profiler(profile_config_from_args(args)) - runner = bnb.NBatchBenchRunner( - setup_fn=make_setup_fn(args), - extras_fn=make_extras_fn(args), - # Force perf-counter timing: the GPU work happens in the vla-server process. - device=torch.device("cpu"), - **bnb.bench_runner_kwargs_from_args(args), - ) - runner.run() - - -if __name__ == "__main__": - main() From 26d8436a570b0d28aacb30b40835a801dcee51be Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 04:24:39 +0000 Subject: [PATCH 15/29] benchmark: add remaining external pi05 runtime wrappers --- benchmark/pi05/bench_realtime_vla_pi05.py | 225 +++++++++++++++++++ benchmark/pi05/bench_vlacpp_pi05_client.py | 241 +++++++++++++++++++++ 2 files changed, 466 insertions(+) create mode 100755 benchmark/pi05/bench_realtime_vla_pi05.py create mode 100755 benchmark/pi05/bench_vlacpp_pi05_client.py diff --git a/benchmark/pi05/bench_realtime_vla_pi05.py b/benchmark/pi05/bench_realtime_vla_pi05.py new file mode 100755 index 0000000..f91a6e4 --- /dev/null +++ b/benchmark/pi05/bench_realtime_vla_pi05.py @@ -0,0 +1,225 @@ +#!/usr/bin/env python3 +"""realtime-vla PI0.5 latency benchmark using the PhyAI bench runner. + +This script is a thin adapter around realtime-vla's ``Pi05Inference.forward``. +It mirrors ``benchmark/bench_n_batch_ws1_pi05.py`` by delegating warmup, timed +iterations, JSONL output, and optional profiling to ``benchmark/bench_n_batch.py``. + +Run:: + + python benchmark/pi05/bench_realtime_vla_pi05.py \ + --realtime-vla-root \ + --flashrt-root \ + --checkpoint \ + --batch-sizes 1 --n-warmup 100 --n-timed 100 \ + --result-file results/realtime_vla_pi05.jsonl + +Only batch size 1 is supported because realtime-vla's PI0.5 inference object +used here is allocated for one synthetic request. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import pickle +import sys +from typing import Any + +import torch + +_BENCHMARK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_BENCHMARK_DIR)) +import bench_n_batch as bnb # noqa: E402 +from phyai.utils.profile import ( # noqa: E402 + add_profile_cli_args, + install_profiler, + profile_config_from_args, +) + + +def env_path(name: str) -> Path | None: + value = os.environ.get(name) + return Path(value) if value else None + + +def load_checkpoint(path: Path, flashrt_root: Path | None, trust_pickle: bool): + if path.is_dir(): + path = path / "model.safetensors" + if path.suffix == ".safetensors": + if flashrt_root is None: + raise ValueError( + "--flashrt-root or FLASHRT_ROOT is required when loading a safetensors checkpoint" + ) + # Reuse FlashRT's tested PI0.5 key conversion instead of duplicating it here. + sys.path.insert(0, str(flashrt_root)) + from flash_rt.frontends.torch.pi05_rtx import convert_pi05_safetensors + + return convert_pi05_safetensors(path) + if path.suffix in {".pt", ".pth"}: + return torch.load(path, map_location="cpu", weights_only=True) + if path.suffix in {".pkl", ".pickle"}: + if not trust_pickle: + raise ValueError( + "Refusing to load pickle checkpoint without --trust-pickle-checkpoint. " + "Only use that flag for checkpoints from a trusted source." + ) + with path.open("rb") as f: + return pickle.load(f) # nosec B301: guarded by --trust-pickle-checkpoint. + raise ValueError( + f"Unsupported checkpoint suffix {path.suffix!r}; expected .safetensors, .pt, .pth, .pkl, or .pickle" + ) + + +def make_setup_fn(args: argparse.Namespace): + # Use the checked-out realtime-vla repository directly; installation is optional. + sys.path.insert(0, str(args.realtime_vla_root)) + from pi05_infer import Pi05Inference + + def setup_fn(batch_size: int) -> bnb.BenchSpec: + if batch_size != 1: + raise ValueError("realtime-vla PI0.5 wrapper supports only batch_size=1") + + checkpoint = load_checkpoint( + args.checkpoint, args.flashrt_root, args.trust_pickle_checkpoint + ) + if "language_embeds" not in checkpoint: + # Latency-only fallback for checkpoints that do not include prompt embeds. + checkpoint["language_embeds"] = torch.randn( + args.prompt_len, 2048, dtype=torch.bfloat16 + ) + + infer = Pi05Inference( + checkpoint=checkpoint, + num_views=args.num_views, + chunk_size=args.chunk_size, + tokenizer_path=str(args.tokenizer) if args.tokenizer else None, + discrete_state_input=args.discrete_state_input, + ) + torch.manual_seed(args.seed) + input_image = torch.randn( + args.num_views, 224, 224, 3, dtype=torch.bfloat16, device="cuda" + ) + input_noise = torch.randn( + args.chunk_size, 32, dtype=torch.bfloat16, device="cuda" + ) + + state_tokens = None + if args.discrete_state_input: + import numpy as np + + state_tokens = np.zeros(args.state_dim, dtype=np.int64) + + def step() -> None: + infer.forward( + input_image, + input_noise, + task_prompt=args.prompt, + state_tokens=state_tokens, + ) + + return bnb.BenchSpec( + name="realtime_vla_pi05", + step_callable=step, + teardown_callable=lambda: None, + ) + + return setup_fn + + +def make_extras_fn(args: argparse.Namespace): + def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: + return { + "runtime": "realtime-vla", + "checkpoint": str(args.checkpoint), + "realtime_vla_root": str(args.realtime_vla_root), + "flashrt_root": str(args.flashrt_root) if args.flashrt_root else None, + "precision": "bf16", + "batch_size_contract": 1, + "num_views": args.num_views, + "chunk_size": args.chunk_size, + "prompt": args.prompt, + "prompt_len": args.prompt_len, + "seed": args.seed, + "discrete_state_input": args.discrete_state_input, + "timing_scope": "Pi05Inference.forward hot path; common runner CUDA event wraps one forward call", + } + + return extras_fn + + +def main() -> None: + realtime_default = env_path("REALTIME_VLA_ROOT") + flashrt_default = env_path("FLASHRT_ROOT") + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--realtime-vla-root", + type=Path, + default=realtime_default, + required=realtime_default is None, + help="Path to the realtime-vla repository clone. Can also be set with REALTIME_VLA_ROOT.", + ) + parser.add_argument( + "--flashrt-root", + type=Path, + default=flashrt_default, + help="Path to FlashRT. Required only when --checkpoint is a PI0.5 safetensors checkpoint.", + ) + parser.add_argument( + "--checkpoint", + type=Path, + required=True, + help="realtime-vla .pkl/.pt checkpoint, PI0.5 model.safetensors, or a directory containing model.safetensors.", + ) + parser.add_argument( + "--trust-pickle-checkpoint", + action="store_true", + help="Allow loading .pkl/.pickle checkpoints. Only use with trusted checkpoint files.", + ) + parser.add_argument("--num-views", type=int, default=2) + parser.add_argument("--chunk-size", type=int, default=50) + parser.add_argument("--prompt", default="do something") + parser.add_argument("--prompt-len", type=int, default=16) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--discrete-state-input", + action="store_true", + help="Use realtime-vla tokenizer/state-token prompt path instead of precomputed language_embeds.", + ) + parser.add_argument( + "--tokenizer", + type=Path, + default=None, + help="Tokenizer path for --discrete-state-input.", + ) + parser.add_argument( + "--state-dim", + type=int, + default=8, + help="Synthetic state token count for --discrete-state-input.", + ) + + bnb.add_bench_cli_args(parser) + parser.set_defaults(bench_name="realtime_vla_pi05", batch_sizes=[1]) + add_profile_cli_args(parser) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for realtime-vla PI0.5 benchmarking") + + install_profiler(profile_config_from_args(args)) + runner = bnb.NBatchBenchRunner( + setup_fn=make_setup_fn(args), + extras_fn=make_extras_fn(args), + device=torch.device("cuda", torch.cuda.current_device()), + **bnb.bench_runner_kwargs_from_args(args), + ) + runner.run() + + +if __name__ == "__main__": + main() diff --git a/benchmark/pi05/bench_vlacpp_pi05_client.py b/benchmark/pi05/bench_vlacpp_pi05_client.py new file mode 100755 index 0000000..b29924b --- /dev/null +++ b/benchmark/pi05/bench_vlacpp_pi05_client.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +"""vla.cpp PI0.5 ZMQ client benchmark using the PhyAI bench runner. + +Start ``vla-server`` separately, then run this client. The common PhyAI runner +handles warmup, timed iterations, JSONL output, and optional profiling. This +script forces CPU timing in the runner because the measured operation is a ZMQ +request to an external server process; CUDA events in the client process would +not cover server-side GPU work. + +Run:: + + python benchmark/pi05/bench_vlacpp_pi05_client.py \ + --vlacpp-root \ + --addr tcp://127.0.0.1:5555 \ + --tokenizer \ + --stats-json \ + --batch-sizes 1 --n-warmup 100 --n-timed 100 \ + --result-file results/vlacpp_pi05.jsonl + +Only batch size 1 is supported because the vla.cpp Python client sends one +request at a time. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import statistics +import sys +from typing import Any + +import numpy as np +import torch + +_BENCHMARK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_BENCHMARK_DIR)) +import bench_n_batch as bnb +from phyai.utils.profile import ( + add_profile_cli_args, + install_profiler, + profile_config_from_args, +) + + +def env_path(name: str) -> Path | None: + value = os.environ.get(name) + return Path(value) if value else None + + +def summarize(values: list[float]) -> dict[str, float] | None: + if not values: + return None + xs = sorted(float(x) for x in values) + return { + "count": len(xs), + "mean_ms": float(statistics.fmean(xs)), + "median_ms": float(statistics.median(xs)), + "p50_ms": float(np.percentile(xs, 50)), + "p90_ms": float(np.percentile(xs, 90)), + "p99_ms": float(np.percentile(xs, 99)), + "min_ms": xs[0], + "max_ms": xs[-1], + } + + +def make_obs(seed: int, num_views: int, prompt: str) -> dict[str, Any]: + """Create deterministic synthetic observations for latency-only requests. + + vla.cpp's official PI0.5 client expects CHW float32 images in [0, 1] and an + unnormalized robot state vector; it converts these into server protobufs. + """ + rng = np.random.default_rng(seed) + obs: dict[str, Any] = { + "observation.images.image": rng.random((3, 224, 224), dtype=np.float32), + "observation.state": rng.standard_normal(8).astype(np.float32), + "task": prompt, + "prompt": prompt, + } + if num_views >= 2: + obs["observation.images.image2"] = rng.random((3, 224, 224), dtype=np.float32) + if num_views >= 3: + obs["observation.images.image3"] = rng.random((3, 224, 224), dtype=np.float32) + return obs + + +def make_setup_fn(args: argparse.Namespace): + eval_root = args.vlacpp_root / "eval" + sys.path.insert(0, str(eval_root)) + sys.path.insert(0, str(eval_root / "client")) + from client.vla_cpp_client import VlaCppClient + + def setup_fn(batch_size: int) -> bnb.BenchSpec: + if batch_size != 1: + raise ValueError("vla.cpp PI0.5 wrapper supports only batch_size=1") + + image_keys = ["observation.images.image"] + if args.num_views >= 2: + image_keys.append("observation.images.image2") + if args.num_views >= 3: + image_keys.append("observation.images.image3") + + client = VlaCppClient( + vla_addr=args.addr, + arch=args.arch, + tokenizer_name=args.tokenizer, + image_keys=image_keys, + max_length=args.max_length, + real_action_dim=args.real_action_dim, + n_action_steps=1, + stats_json=args.stats_json, + ) + obs = make_obs(args.seed, args.num_views, args.prompt) + phase_samples: dict[str, list[float]] = { + "server_total_latency_ms": [], + "server_vision_latency_ms": [], + "server_inference_latency_ms": [], + "server_prefill_latency_ms": [], + "server_denoise_latency_ms": [], + } + phase_summary: dict[str, Any] = {} + call_count = 0 + + def update_phase_summary() -> None: + phase_summary.clear() + for key, values in phase_samples.items(): + phase_summary[key] = summarize(values) + + def step() -> None: + nonlocal call_count + client.get_action(obs) + call_count += 1 + if call_count <= args.n_warmup: + return + r = getattr(client, "_last_response", None) + if r is None: + return + phase_samples["server_total_latency_ms"].append(float(r.latency_ms_total)) + phase_samples["server_vision_latency_ms"].append(float(r.latency_ms_vision)) + phase_samples["server_inference_latency_ms"].append( + float(r.latency_ms_inference) + ) + phase_samples["server_prefill_latency_ms"].append( + float(r.latency_ms_prefill) + ) + phase_samples["server_denoise_latency_ms"].append( + float(r.latency_ms_denoise) + ) + update_phase_summary() + + def teardown() -> None: + sock = getattr(client, "sock", None) + if sock is not None: + sock.close(linger=0) + + spec = bnb.BenchSpec( + name="vlacpp_pi05_zmq_client", + step_callable=step, + teardown_callable=teardown, + ) + # Attach dynamic summary for extras_fn. The runner copies the dict after + # timed steps finish, so it records the final server phase statistics. + spec.vlacpp_phase_summary = phase_summary # type: ignore[attr-defined] + return spec + + return setup_fn + + +def make_extras_fn(args: argparse.Namespace): + def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]: + return { + "runtime": "vla.cpp", + "vlacpp_root": str(args.vlacpp_root), + "addr": args.addr, + "arch": args.arch, + "tokenizer": args.tokenizer, + "stats_json": str(args.stats_json) if args.stats_json else None, + "batch_size_contract": 1, + "num_views": args.num_views, + "chunk_size_metadata": args.chunk_size, + "prompt": args.prompt, + "seed": args.seed, + "server_phase_latency_ms": getattr(spec, "vlacpp_phase_summary", {}), + "timing_scope": "client ZMQ request wall time; server phase timings are copied from PredictResponse extras", + "notes": "vla.cpp server enforces the GGUF chunk size; prefix/expert may be combined in server phase timing.", + } + + return extras_fn + + +def main() -> None: + vlacpp_default = env_path("VLACPP_ROOT") + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--vlacpp-root", + type=Path, + default=vlacpp_default, + required=vlacpp_default is None, + help="Path to the vla.cpp repository clone. Can also be set with VLACPP_ROOT.", + ) + parser.add_argument("--addr", default="tcp://127.0.0.1:5555") + parser.add_argument("--arch", default="pi05") + parser.add_argument( + "--tokenizer", + default="google/paligemma-3b-pt-224", + help="Tokenizer name or local tokenizer path. Prefer a local path if the HF repo is gated.", + ) + parser.add_argument( + "--stats-json", + type=Path, + default=None, + help="Local LIBERO meta/stats.json for arch=pi05; avoids network fetch.", + ) + parser.add_argument("--num-views", type=int, default=2) + parser.add_argument("--chunk-size", type=int, default=50) + parser.add_argument("--prompt", default="do something") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--max-length", type=int, default=200) + parser.add_argument("--real-action-dim", type=int, default=7) + + bnb.add_bench_cli_args(parser) + parser.set_defaults(bench_name="vlacpp_pi05", batch_sizes=[1]) + add_profile_cli_args(parser) + args = parser.parse_args() + + install_profiler(profile_config_from_args(args)) + runner = bnb.NBatchBenchRunner( + setup_fn=make_setup_fn(args), + extras_fn=make_extras_fn(args), + # Force perf-counter timing: the GPU work happens in the vla-server process. + device=torch.device("cpu"), + **bnb.bench_runner_kwargs_from_args(args), + ) + runner.run() + + +if __name__ == "__main__": + main() From 6f635f23799006c64c5fe769b9081fcd693042ee Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 04:41:23 +0000 Subject: [PATCH 16/29] policy: remove pi05 libero adapter files --- phyai/src/phyai/env.py | 194 ----------- phyai/src/phyai/policies/__init__.py | 5 - phyai/src/phyai/policies/pi05_libero.py | 445 ------------------------ 3 files changed, 644 deletions(-) delete mode 100644 phyai/src/phyai/env.py delete mode 100644 phyai/src/phyai/policies/__init__.py delete mode 100644 phyai/src/phyai/policies/pi05_libero.py diff --git a/phyai/src/phyai/env.py b/phyai/src/phyai/env.py deleted file mode 100644 index 08acede..0000000 --- a/phyai/src/phyai/env.py +++ /dev/null @@ -1,194 +0,0 @@ -"""Typed env-var registry for ``PHYAI_*`` overrides. - -Every env var phyai consults goes through one :class:`EnvField` -descriptor in :class:`envs` below — instead of scattered -``os.environ.get(...)`` calls in random modules. Two upsides: - -* one place to document and tune the env-var contract; -* typed parsing (``int`` / ``bool`` / ``torch.dtype``) lives next to - the declaration, so callers always get the right shape back without - rewriting parsing logic per-site. - -Usage ------ -:: - - from phyai.env import envs - - if (raw := envs.PHYAI_ATTN_BACKEND.get()) is not None: - cfg = cfg.replace(backends=cfg.backends.replace(attn=raw)) - -The module is intentionally tiny and dependency-free (no torch, no -phyai imports) so it can be consulted from low-level modules during -their own bootstrap without import cycles. -""" - -from __future__ import annotations - -import os -from typing import Callable, Generic, TypeVar - - -T = TypeVar("T") - - -class EnvField(Generic[T]): - """One typed env-var slot with a parser and an optional default. - - ``get()`` returns ``self.default`` when the variable is unset, and - a parsed value otherwise. Parsing errors raise :class:`ValueError` - with the env-var name attached so the failure points back at the - user's environment rather than the consuming module. - """ - - __slots__ = ("name", "default", "parser") - - def __init__( - self, - name: str, - default: T | None, - parser: Callable[[str], T], - ) -> None: - self.name = name - self.default = default - self.parser = parser - - def is_set(self) -> bool: - """``True`` if the env var is present (even if empty).""" - return self.name in os.environ - - def get(self) -> T | None: - raw = os.environ.get(self.name) - if raw is None: - return self.default - try: - return self.parser(raw) - except (ValueError, TypeError) as e: - raise ValueError(f"{self.name}={raw!r}: {e}") from e - - -def _parse_bool(s: str) -> bool: - """Accept ``1/true/yes/on`` (case-insensitive) -> True; ``0/false/no/off`` -> False.""" - v = s.strip().lower() - if v in ("1", "true", "yes", "on"): - return True - if v in ("0", "false", "no", "off"): - return False - raise ValueError(f"expected a boolean (1/0/true/false/yes/no/on/off), got {s!r}") - - -def _parse_dtype(s: str): - """Map a name like ``"bf16"`` / ``"bfloat16"`` to ``torch.dtype``. - - Imports torch lazily so importing :mod:`phyai.env` stays cheap - and dependency-free at module-import time. - """ - import torch - - table: dict[str, "torch.dtype"] = { - "bf16": torch.bfloat16, - "bfloat16": torch.bfloat16, - "fp16": torch.float16, - "float16": torch.float16, - "half": torch.float16, - "fp32": torch.float32, - "float32": torch.float32, - "float": torch.float32, - "fp64": torch.float64, - "float64": torch.float64, - "double": torch.float64, - } - key = s.strip().lower() - if key not in table: - raise ValueError( - f"expected one of {sorted(table)} (case-insensitive), got {s!r}" - ) - return table[key] - - -def _parse_regex_list(s: str) -> tuple[str, ...]: - """Parse a JSON array of regex strings, e.g. ``'["o_proj$", "\\\\.heads\\\\."]'``. - - JSON (not a comma-split) because regex patterns routinely contain - ``,`` / ``|`` / ``.`` that a naive split would mangle. A single bare - string (not valid JSON, or JSON that isn't a list) is treated as a - one-element list so ``PHYAI_DEBUG_TENSOR_DUMP_FILTER='o_proj$'`` also - works for the common single-pattern case. An empty / whitespace string - yields the empty tuple. Used for the tensor-dump operator filter. - """ - import json - - raw = s.strip() - if not raw: - return () - try: - parsed = json.loads(raw) - except json.JSONDecodeError: - return (raw,) - if isinstance(parsed, str): - return (parsed,) - if isinstance(parsed, list) and all(isinstance(x, str) for x in parsed): - return tuple(parsed) - raise ValueError( - f"expected a JSON array of regex strings (or a single pattern), got {s!r}" - ) - - -class envs: - """Process-level typed env-var registry. - - Read each via ``envs.PHYAI_FOO.get()`` and check ``.is_set()`` when - you need to distinguish "unset" from "set to default". Adding a new - env var means a single new line here and a one-line consumer change. - """ - - # ---------- backend / kernel selection ---------- # - PHYAI_ATTN_BACKEND = EnvField("PHYAI_ATTN_BACKEND", None, str) - PHYAI_NORM_BACKEND = EnvField("PHYAI_NORM_BACKEND", None, str) - PHYAI_LINEAR_BACKEND = EnvField("PHYAI_LINEAR_BACKEND", None, str) - PHYAI_VGPU_BACKEND = EnvField("PHYAI_VGPU_BACKEND", None, str) - - # ---------- device / dtype ---------- # - PHYAI_DEVICE = EnvField("PHYAI_DEVICE", None, str) - PHYAI_PARAMS_DTYPE = EnvField("PHYAI_PARAMS_DTYPE", None, _parse_dtype) - - # ---------- runtime ---------- # - PHYAI_USE_CUDA_GRAPH = EnvField("PHYAI_USE_CUDA_GRAPH", None, _parse_bool) - - # ---------- policy adapters ---------- # - PHYAI_CAMERA_MODE = EnvField("PHYAI_CAMERA_MODE", None, str) - PHYAI_TOKENIZER_PATH = EnvField("PHYAI_TOKENIZER_PATH", None, str) - - # ---------- parallel ---------- # - PHYAI_WORLD_SIZE = EnvField("PHYAI_WORLD_SIZE", None, int) - PHYAI_DP_SIZE = EnvField("PHYAI_DP_SIZE", None, int) - PHYAI_EP_SIZE = EnvField("PHYAI_EP_SIZE", None, int) - PHYAI_SP_SIZE = EnvField("PHYAI_SP_SIZE", None, int) - PHYAI_CP_SIZE = EnvField("PHYAI_CP_SIZE", None, int) - PHYAI_TP_SIZE = EnvField("PHYAI_TP_SIZE", None, int) - - # ---------- low-level tuning ---------- # - PHYAI_FLASHINFER_WORKSPACE_BYTES = EnvField( - "PHYAI_FLASHINFER_WORKSPACE_BYTES", None, int - ) - PHYAI_FLASHINFER_PREFILL_BACKEND = EnvField( - "PHYAI_FLASHINFER_PREFILL_BACKEND", None, str - ) - PHYAI_FORCE_LINEAR_KERNEL = EnvField("PHYAI_FORCE_LINEAR_KERNEL", None, str) - - # ---------- debug / tensor dump ---------- # - # When the dump dir is set the engine runs eager (cuda graph forced - # off) and records every selected leaf operator's output, one .pt per - # step. FILTER is a JSON array of regexes matched against operator - # names (or a single bare pattern); FILTER_FN is a "pkg.mod:func" / - # "/path.py:func" predicate. The two filters are mutually exclusive. - PHYAI_DEBUG_TENSOR_DUMP_DIR = EnvField("PHYAI_DEBUG_TENSOR_DUMP_DIR", None, str) - PHYAI_DEBUG_TENSOR_DUMP_FILTER = EnvField( - "PHYAI_DEBUG_TENSOR_DUMP_FILTER", None, _parse_regex_list - ) - PHYAI_DEBUG_TENSOR_DUMP_FILTER_FN = EnvField( - "PHYAI_DEBUG_TENSOR_DUMP_FILTER_FN", None, str - ) - - -__all__ = ["EnvField", "envs"] diff --git a/phyai/src/phyai/policies/__init__.py b/phyai/src/phyai/policies/__init__.py deleted file mode 100644 index 9a1a0cf..0000000 --- a/phyai/src/phyai/policies/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""High-level policy wrappers.""" - -from phyai.policies.pi05_libero import PI05LiberoPolicy - -__all__ = ["PI05LiberoPolicy"] diff --git a/phyai/src/phyai/policies/pi05_libero.py b/phyai/src/phyai/policies/pi05_libero.py deleted file mode 100644 index bcaea76..0000000 --- a/phyai/src/phyai/policies/pi05_libero.py +++ /dev/null @@ -1,445 +0,0 @@ -"""Thin LIBERO adapter for pi0.5 PhyAI inference.""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any - -import numpy as np -import torch -import torch.nn.functional as F -from safetensors.torch import load_file - -from phyai.engine import Engine, EngineArgs -from phyai.engine_config import BackendConfig, DeviceConfig, EngineConfig, RuntimeConfig -from phyai.env import envs -from phyai.models.pi05.configuration_pi05 import PI05Config -from phyai.models.pi05.main_pi05 import PI05Args -from phyai.models.pi05.scheduler_ws1_pi05 import PI05Request -from phyai_utils_tools.models.pi05 import PI05_DEFAULT_TOKENIZER_NAME, PI05Processor -from phyai_utils_tools.processing.transition import IMAGES, STATE, TASK - -LIBERO_AGENTVIEW_KEYS: tuple[str, ...] = ( - "agentview", - "agentview_image", - "image", - "observation.images.image", -) -LIBERO_WRIST_KEYS: tuple[str, ...] = ( - "wrist", - "robot0_eye_in_hand_image", - "wrist_image", - "image2", - "observation.images.image2", -) - - -def _lerobot_pi05_weight_remap(key: str) -> str | None: - """Strip LeRobot's outer model prefix and drop inference-unused keys.""" - if key.startswith("model."): - key = key[len("model.") :] - if key == "paligemma_with_expert.gemma_expert.lm_head.weight": - return None - return key - - -class PI05LiberoPolicy: - """Adapt vla-evaluation-harness LIBERO observations to ``PI05Processor``.""" - - def __init__( - self, - checkpoint_dir: str | Path, - *, - device: str = "cuda", - params_dtype: torch.dtype = torch.bfloat16, - max_batch_size: int = 1, - use_cuda_graph: bool = True, - attn_backend: str = "flashinfer", - norm_backend: str = "phyai-kernel", - linear_backend: str | None = "flashinfer", - flashinfer_workspace_bytes: int = 512 * 1024 * 1024, - tokenizer_name: str | None = None, - camera_mode: str | None = None, - ) -> None: - self.checkpoint_dir = Path(checkpoint_dir) - self.device = device - self.params_dtype = params_dtype - self.max_batch_size = int(max_batch_size) - self.config = self._read_config() - self.image_size = self._resolve_image_size(self.config) - self._action_dim = self._resolve_action_dim(self.config) - self.max_action_dim = int(self.config.get("max_action_dim", 32)) - self._chunk_size = int(self.config.get("chunk_size", PI05Config().chunk_size)) - self.camera_names = self._resolve_camera_names(camera_mode) - self.tokenizer_name = self._resolve_tokenizer_name(tokenizer_name) - self.prompt_mode = str( - self.config.get("phyai_prompt_mode", "lerobot_state_bins") - ) - self.normalization_mode = str( - self.config.get("phyai_normalization_mode", "mean_std") - ) - self._use_phyai_compat = ( - "phyai_prompt_mode" in self.config - or "phyai_normalization_mode" in self.config - ) - self._normalizer_stats = self._load_processor_state( - "policy_preprocessor.json", "normalizer_processor" - ) - self._unnormalizer_stats = self._load_processor_state( - "policy_postprocessor.json", "unnormalizer_processor" - ) - if self._use_phyai_compat: - self._validate_compat_stats() - self._tokenizer = None - self.processor = PI05Processor.from_pretrained( - self.checkpoint_dir, - tokenizer_name=self.tokenizer_name, - image_size=self.image_size, - num_channels=3, - num_images=len(self.camera_names), - action_dim=self._action_dim, - normalize_pixels=True, - device=device, - params_dtype=params_dtype, - ) - self.engine = Engine( - EngineArgs( - plugin="pi05", - plugin_args=PI05Args( - checkpoint_dir=self.checkpoint_dir, - max_batch_size=self.max_batch_size, - weight_remap=_lerobot_pi05_weight_remap, - inputs_image_shape=[ - [self.image_size, self.image_size, 3] for _ in self.camera_names - ], - ), - config=EngineConfig( - backends=BackendConfig( - attn=attn_backend, norm=norm_backend, linear=linear_backend - ), - device=DeviceConfig(target=device, params_dtype=params_dtype), - runtime=RuntimeConfig( - use_cuda_graph=use_cuda_graph, - flashinfer_workspace_bytes=flashinfer_workspace_bytes, - force_linear_kernel=linear_backend, - ), - ), - ) - ) - - @property - def chunk_size(self) -> int: - return self._chunk_size - - @property - def action_dim(self) -> int: - return int(self.processor.action_dim or self._action_dim) - - @staticmethod - def _resolve_image_size(config: dict[str, Any]) -> int: - resolution = config.get("image_resolution") - if isinstance(resolution, list) and resolution: - return int(resolution[0]) - return PI05Config().vision.image_size - - @staticmethod - def _resolve_action_dim(config: dict[str, Any]) -> int: - shape = config.get("output_features", {}).get("action", {}).get("shape") - if isinstance(shape, list) and shape: - return int(shape[-1]) - return 7 - - def _read_config(self) -> dict[str, Any]: - path = self.checkpoint_dir / "config.json" - if not path.exists(): - return {} - with path.open("r", encoding="utf-8") as f: - return json.load(f) - - def _resolve_camera_names(self, camera_mode: str | None) -> list[str]: - mode = camera_mode or envs.PHYAI_CAMERA_MODE.get() or "three_camera" - if mode == "two_camera": - return ["agentview", "wrist"] - if mode == "three_camera": - return ["agentview", "wrist", "empty"] - raise ValueError(f"Unsupported PHYAI_CAMERA_MODE={mode!r}.") - - def _resolve_tokenizer_name(self, tokenizer_name: str | None) -> str: - if tokenizer_name: - return tokenizer_name - if env_tokenizer := envs.PHYAI_TOKENIZER_PATH.get(): - return env_tokenizer - if config_tokenizer := self.config.get("tokenizer_name"): - return str(config_tokenizer) - return PI05_DEFAULT_TOKENIZER_NAME - - def _load_processor_state( - self, config_name: str, registry_name: str - ) -> dict[str, torch.Tensor]: - path = self.checkpoint_dir / config_name - if not path.exists(): - return {} - with path.open("r", encoding="utf-8") as f: - config = json.load(f) - for step in config.get("steps", []): - if step.get("registry_name") != registry_name: - continue - state_file = step.get("state_file") - if not state_file: - return {} - return load_file(str(self.checkpoint_dir / state_file)) - return {} - - def _validate_compat_stats(self) -> None: - if self.normalization_mode == "openpi_quantile": - normalizer_keys = ("observation.state.min", "observation.state.max") - unnormalizer_keys = ("action.min", "action.max") - else: - normalizer_keys = ("observation.state.mean", "observation.state.std") - unnormalizer_keys = ("action.mean", "action.std") - missing = [ - f"normalizer:{key}" - for key in normalizer_keys - if key not in self._normalizer_stats - ] - missing.extend( - f"unnormalizer:{key}" - for key in unnormalizer_keys - if key not in self._unnormalizer_stats - ) - if missing: - raise ValueError( - f"{self.checkpoint_dir}: compat normalization requires missing stats " - f"{', '.join(missing)}" - ) - - @property - def tokenizer(self): - if self._tokenizer is None: - from transformers import AutoTokenizer - - self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name) - return self._tokenizer - - def observation_to_raw(self, obs: dict[str, Any]) -> dict[str, Any]: - return { - IMAGES: [ - self._extract_camera_tensor(obs, name) for name in self.camera_names - ], - STATE: self._extract_state(obs), - TASK: [self._extract_task(obs)], - } - - def observation_to_request_inputs( - self, obs: dict[str, Any] - ) -> dict[str, torch.Tensor]: - if not self._use_phyai_compat: - processed = self.processor.preprocess(self.observation_to_raw(obs)) - return { - "pixel_values": processed.pixel_values, - "input_ids": processed.input_ids, - "lang_lens": processed.lang_lens, - } - pixel_values = ( - torch.stack( - [ - self._extract_camera_model_tensor(obs, name).squeeze(0) - for name in self.camera_names - ], - dim=0, - ) - .unsqueeze(0) - .to(self.device) - ) - state = self._normalize_state(self._extract_state(obs)) - input_ids, lang_lens = self._tokenize_inputs([self._extract_task(obs)], state) - return { - "pixel_values": pixel_values, - "input_ids": input_ids.to(self.device), - "lang_lens": lang_lens.to(self.device), - } - - def _extract_camera_tensor( - self, obs: dict[str, Any], camera_name: str - ) -> torch.Tensor: - image = self._extract_camera_image(obs, camera_name) - return self._image_to_raw_tensor(image) - - def _extract_camera_model_tensor( - self, obs: dict[str, Any], camera_name: str - ) -> torch.Tensor: - image = self._extract_camera_image(obs, camera_name) - return self._image_to_model_tensor(image) - - def _extract_camera_image( - self, obs: dict[str, Any], camera_name: str - ) -> np.ndarray: - if camera_name == "agentview": - return self._extract_image(obs, LIBERO_AGENTVIEW_KEYS) - if camera_name == "wrist": - return self._extract_image(obs, LIBERO_WRIST_KEYS) - if camera_name == "empty": - return np.zeros((self.image_size, self.image_size, 3), dtype=np.uint8) - raise ValueError(f"Unsupported camera_name={camera_name!r}.") - - @staticmethod - def _extract_image(obs: dict[str, Any], keys: tuple[str, ...]) -> np.ndarray: - candidates: list[Any] = [] - images = obs.get("images") - if isinstance(images, dict): - candidates.extend(images.get(k) for k in keys) - candidates.extend(obs.get(k) for k in keys) - for candidate in candidates: - if candidate is None: - continue - array = np.asarray(candidate) - if array.ndim == 4: - array = array[0] - if array.ndim != 3: - continue - if array.shape[0] == 3 and array.shape[-1] != 3: - array = np.transpose(array, (1, 2, 0)) - if array.shape[-1] == 3: - return array - raise KeyError(f"LIBERO observation does not contain any image keys: {keys}.") - - @staticmethod - def _image_to_raw_tensor(image: np.ndarray) -> torch.Tensor: - array = np.asarray(image, dtype=np.float32) - if array.max(initial=0.0) > 1.0: - array = array / 255.0 - return ( - torch.from_numpy(np.ascontiguousarray(array.transpose(2, 0, 1))) - .unsqueeze(0) - .contiguous() - ) - - def _image_to_model_tensor(self, image: np.ndarray) -> torch.Tensor: - tensor = self._image_to_raw_tensor(image) - if tensor.shape[-2:] != (self.image_size, self.image_size): - tensor = self._resize_with_pad(tensor, self.image_size, self.image_size) - return (tensor * 2.0 - 1.0).contiguous() - - @staticmethod - def _resize_with_pad(images: torch.Tensor, height: int, width: int) -> torch.Tensor: - _, _, cur_height, cur_width = images.shape - ratio = max(cur_width / width, cur_height / height) - resized_height = int(cur_height / ratio) - resized_width = int(cur_width / ratio) - resized = F.interpolate( - images, - size=(resized_height, resized_width), - mode="bilinear", - align_corners=False, - ) - resized = resized.clamp(0.0, 1.0) - pad_h0, rem_h = divmod(height - resized_height, 2) - pad_w0, rem_w = divmod(width - resized_width, 2) - return F.pad( - resized, - (pad_w0, pad_w0 + rem_w, pad_h0, pad_h0 + rem_h), - mode="constant", - value=0.0, - ) - - @staticmethod - def _extract_state(obs: dict[str, Any]) -> torch.Tensor: - state = obs.get("states", obs.get("state")) - if state is None: - raise KeyError("LIBERO observation must contain 'states' or 'state'.") - array = np.asarray(state, dtype=np.float32) - if array.ndim == 1: - array = array[None, :] - return torch.from_numpy(np.ascontiguousarray(array)) - - @staticmethod - def _extract_task(obs: dict[str, Any]) -> str: - task = obs.get("task_description", obs.get("task", "")) - if isinstance(task, (list, tuple)): - task = task[0] if task else "" - return str(task) - - def _normalize_state(self, state: torch.Tensor) -> torch.Tensor: - if self.normalization_mode == "openpi_quantile": - min_v = self._normalizer_stats.get("observation.state.min") - max_v = self._normalizer_stats.get("observation.state.max") - if min_v is None or max_v is None: - return state - return (state - min_v.to(state)) / ( - max_v.to(state) - min_v.to(state) + 1e-6 - ) * 2.0 - 1.0 - mean = self._normalizer_stats.get("observation.state.mean") - std = self._normalizer_stats.get("observation.state.std") - if mean is None or std is None: - return state - return (state - mean.to(state)) / torch.clamp(std.to(state), min=1e-8) - - def _tokenize_inputs( - self, tasks: list[str], states: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.prompt_mode == "openpi_task": - prompts = [ - task.strip().replace("_", " ").replace("\n", " ") + "\n" - for task in tasks - ] - else: - state_np = states.detach().cpu().numpy() - bins = np.linspace(-1.0, 1.0, 257)[:-1] - discretized = np.digitize(state_np, bins=bins) - 1 - discretized = np.clip(discretized, 0, 255) - prompts = [] - for task, state_bins in zip(tasks, discretized): - cleaned = task.strip().replace("_", " ").replace("\n", " ") - state_str = " ".join(map(str, state_bins)) - prompts.append(f"Task: {cleaned}, State: {state_str};\nAction: ") - encoded = self.tokenizer( - prompts, - max_length=int(self.config.get("tokenizer_max_length", 200)), - padding="max_length", - padding_side="right", - truncation=True, - return_tensors="pt", - ) - return encoded["input_ids"].to(torch.int64), encoded["attention_mask"].sum( - dim=-1 - ).to(torch.int64) - - def _postprocess_actions(self, raw_actions: torch.Tensor) -> np.ndarray: - action = raw_actions[..., : self.action_dim].detach().float() - if not self._use_phyai_compat: - actions = self.processor.postprocess(action) - if isinstance(actions, torch.Tensor): - actions = actions.detach().cpu().numpy() - return np.asarray(actions, dtype=np.float32) - action = action.cpu() - if self.normalization_mode == "openpi_quantile": - min_v = self._unnormalizer_stats.get("action.min") - max_v = self._unnormalizer_stats.get("action.max") - if min_v is not None and max_v is not None: - action = (action + 1.0) / 2.0 * ( - max_v.to(action) - min_v.to(action) + 1e-6 - ) + min_v.to(action) - else: - mean = self._unnormalizer_stats.get("action.mean") - std = self._unnormalizer_stats.get("action.std") - if mean is not None and std is not None: - action = action * torch.clamp(std.to(action), min=1e-8) + mean.to( - action - ) - return action.numpy().astype(np.float32) - - def infer( - self, obs: dict[str, Any], *, noise: torch.Tensor | np.ndarray | None = None - ) -> dict[str, np.ndarray]: - request_kwargs = self.observation_to_request_inputs(obs) - if noise is not None: - request_kwargs["noise"] = torch.as_tensor(noise, device=self.device) - request = PI05Request(**request_kwargs) - with torch.inference_mode(): - raw_actions = self.engine.step(request) - actions = self._postprocess_actions(raw_actions) - return {"actions": actions} - - def close(self) -> None: - self.engine.close() From 1889254aa0169f46c984887e486f7cc495262afe Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 04:41:23 +0000 Subject: [PATCH 17/29] policy: add pi05 libero adapter files --- phyai/src/phyai/env.py | 194 +++++++++++ phyai/src/phyai/policies/__init__.py | 5 + phyai/src/phyai/policies/pi05_libero.py | 445 ++++++++++++++++++++++++ 3 files changed, 644 insertions(+) create mode 100644 phyai/src/phyai/env.py create mode 100644 phyai/src/phyai/policies/__init__.py create mode 100644 phyai/src/phyai/policies/pi05_libero.py diff --git a/phyai/src/phyai/env.py b/phyai/src/phyai/env.py new file mode 100644 index 0000000..08acede --- /dev/null +++ b/phyai/src/phyai/env.py @@ -0,0 +1,194 @@ +"""Typed env-var registry for ``PHYAI_*`` overrides. + +Every env var phyai consults goes through one :class:`EnvField` +descriptor in :class:`envs` below — instead of scattered +``os.environ.get(...)`` calls in random modules. Two upsides: + +* one place to document and tune the env-var contract; +* typed parsing (``int`` / ``bool`` / ``torch.dtype``) lives next to + the declaration, so callers always get the right shape back without + rewriting parsing logic per-site. + +Usage +----- +:: + + from phyai.env import envs + + if (raw := envs.PHYAI_ATTN_BACKEND.get()) is not None: + cfg = cfg.replace(backends=cfg.backends.replace(attn=raw)) + +The module is intentionally tiny and dependency-free (no torch, no +phyai imports) so it can be consulted from low-level modules during +their own bootstrap without import cycles. +""" + +from __future__ import annotations + +import os +from typing import Callable, Generic, TypeVar + + +T = TypeVar("T") + + +class EnvField(Generic[T]): + """One typed env-var slot with a parser and an optional default. + + ``get()`` returns ``self.default`` when the variable is unset, and + a parsed value otherwise. Parsing errors raise :class:`ValueError` + with the env-var name attached so the failure points back at the + user's environment rather than the consuming module. + """ + + __slots__ = ("name", "default", "parser") + + def __init__( + self, + name: str, + default: T | None, + parser: Callable[[str], T], + ) -> None: + self.name = name + self.default = default + self.parser = parser + + def is_set(self) -> bool: + """``True`` if the env var is present (even if empty).""" + return self.name in os.environ + + def get(self) -> T | None: + raw = os.environ.get(self.name) + if raw is None: + return self.default + try: + return self.parser(raw) + except (ValueError, TypeError) as e: + raise ValueError(f"{self.name}={raw!r}: {e}") from e + + +def _parse_bool(s: str) -> bool: + """Accept ``1/true/yes/on`` (case-insensitive) -> True; ``0/false/no/off`` -> False.""" + v = s.strip().lower() + if v in ("1", "true", "yes", "on"): + return True + if v in ("0", "false", "no", "off"): + return False + raise ValueError(f"expected a boolean (1/0/true/false/yes/no/on/off), got {s!r}") + + +def _parse_dtype(s: str): + """Map a name like ``"bf16"`` / ``"bfloat16"`` to ``torch.dtype``. + + Imports torch lazily so importing :mod:`phyai.env` stays cheap + and dependency-free at module-import time. + """ + import torch + + table: dict[str, "torch.dtype"] = { + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "fp16": torch.float16, + "float16": torch.float16, + "half": torch.float16, + "fp32": torch.float32, + "float32": torch.float32, + "float": torch.float32, + "fp64": torch.float64, + "float64": torch.float64, + "double": torch.float64, + } + key = s.strip().lower() + if key not in table: + raise ValueError( + f"expected one of {sorted(table)} (case-insensitive), got {s!r}" + ) + return table[key] + + +def _parse_regex_list(s: str) -> tuple[str, ...]: + """Parse a JSON array of regex strings, e.g. ``'["o_proj$", "\\\\.heads\\\\."]'``. + + JSON (not a comma-split) because regex patterns routinely contain + ``,`` / ``|`` / ``.`` that a naive split would mangle. A single bare + string (not valid JSON, or JSON that isn't a list) is treated as a + one-element list so ``PHYAI_DEBUG_TENSOR_DUMP_FILTER='o_proj$'`` also + works for the common single-pattern case. An empty / whitespace string + yields the empty tuple. Used for the tensor-dump operator filter. + """ + import json + + raw = s.strip() + if not raw: + return () + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + return (raw,) + if isinstance(parsed, str): + return (parsed,) + if isinstance(parsed, list) and all(isinstance(x, str) for x in parsed): + return tuple(parsed) + raise ValueError( + f"expected a JSON array of regex strings (or a single pattern), got {s!r}" + ) + + +class envs: + """Process-level typed env-var registry. + + Read each via ``envs.PHYAI_FOO.get()`` and check ``.is_set()`` when + you need to distinguish "unset" from "set to default". Adding a new + env var means a single new line here and a one-line consumer change. + """ + + # ---------- backend / kernel selection ---------- # + PHYAI_ATTN_BACKEND = EnvField("PHYAI_ATTN_BACKEND", None, str) + PHYAI_NORM_BACKEND = EnvField("PHYAI_NORM_BACKEND", None, str) + PHYAI_LINEAR_BACKEND = EnvField("PHYAI_LINEAR_BACKEND", None, str) + PHYAI_VGPU_BACKEND = EnvField("PHYAI_VGPU_BACKEND", None, str) + + # ---------- device / dtype ---------- # + PHYAI_DEVICE = EnvField("PHYAI_DEVICE", None, str) + PHYAI_PARAMS_DTYPE = EnvField("PHYAI_PARAMS_DTYPE", None, _parse_dtype) + + # ---------- runtime ---------- # + PHYAI_USE_CUDA_GRAPH = EnvField("PHYAI_USE_CUDA_GRAPH", None, _parse_bool) + + # ---------- policy adapters ---------- # + PHYAI_CAMERA_MODE = EnvField("PHYAI_CAMERA_MODE", None, str) + PHYAI_TOKENIZER_PATH = EnvField("PHYAI_TOKENIZER_PATH", None, str) + + # ---------- parallel ---------- # + PHYAI_WORLD_SIZE = EnvField("PHYAI_WORLD_SIZE", None, int) + PHYAI_DP_SIZE = EnvField("PHYAI_DP_SIZE", None, int) + PHYAI_EP_SIZE = EnvField("PHYAI_EP_SIZE", None, int) + PHYAI_SP_SIZE = EnvField("PHYAI_SP_SIZE", None, int) + PHYAI_CP_SIZE = EnvField("PHYAI_CP_SIZE", None, int) + PHYAI_TP_SIZE = EnvField("PHYAI_TP_SIZE", None, int) + + # ---------- low-level tuning ---------- # + PHYAI_FLASHINFER_WORKSPACE_BYTES = EnvField( + "PHYAI_FLASHINFER_WORKSPACE_BYTES", None, int + ) + PHYAI_FLASHINFER_PREFILL_BACKEND = EnvField( + "PHYAI_FLASHINFER_PREFILL_BACKEND", None, str + ) + PHYAI_FORCE_LINEAR_KERNEL = EnvField("PHYAI_FORCE_LINEAR_KERNEL", None, str) + + # ---------- debug / tensor dump ---------- # + # When the dump dir is set the engine runs eager (cuda graph forced + # off) and records every selected leaf operator's output, one .pt per + # step. FILTER is a JSON array of regexes matched against operator + # names (or a single bare pattern); FILTER_FN is a "pkg.mod:func" / + # "/path.py:func" predicate. The two filters are mutually exclusive. + PHYAI_DEBUG_TENSOR_DUMP_DIR = EnvField("PHYAI_DEBUG_TENSOR_DUMP_DIR", None, str) + PHYAI_DEBUG_TENSOR_DUMP_FILTER = EnvField( + "PHYAI_DEBUG_TENSOR_DUMP_FILTER", None, _parse_regex_list + ) + PHYAI_DEBUG_TENSOR_DUMP_FILTER_FN = EnvField( + "PHYAI_DEBUG_TENSOR_DUMP_FILTER_FN", None, str + ) + + +__all__ = ["EnvField", "envs"] diff --git a/phyai/src/phyai/policies/__init__.py b/phyai/src/phyai/policies/__init__.py new file mode 100644 index 0000000..9a1a0cf --- /dev/null +++ b/phyai/src/phyai/policies/__init__.py @@ -0,0 +1,5 @@ +"""High-level policy wrappers.""" + +from phyai.policies.pi05_libero import PI05LiberoPolicy + +__all__ = ["PI05LiberoPolicy"] diff --git a/phyai/src/phyai/policies/pi05_libero.py b/phyai/src/phyai/policies/pi05_libero.py new file mode 100644 index 0000000..bcaea76 --- /dev/null +++ b/phyai/src/phyai/policies/pi05_libero.py @@ -0,0 +1,445 @@ +"""Thin LIBERO adapter for pi0.5 PhyAI inference.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from safetensors.torch import load_file + +from phyai.engine import Engine, EngineArgs +from phyai.engine_config import BackendConfig, DeviceConfig, EngineConfig, RuntimeConfig +from phyai.env import envs +from phyai.models.pi05.configuration_pi05 import PI05Config +from phyai.models.pi05.main_pi05 import PI05Args +from phyai.models.pi05.scheduler_ws1_pi05 import PI05Request +from phyai_utils_tools.models.pi05 import PI05_DEFAULT_TOKENIZER_NAME, PI05Processor +from phyai_utils_tools.processing.transition import IMAGES, STATE, TASK + +LIBERO_AGENTVIEW_KEYS: tuple[str, ...] = ( + "agentview", + "agentview_image", + "image", + "observation.images.image", +) +LIBERO_WRIST_KEYS: tuple[str, ...] = ( + "wrist", + "robot0_eye_in_hand_image", + "wrist_image", + "image2", + "observation.images.image2", +) + + +def _lerobot_pi05_weight_remap(key: str) -> str | None: + """Strip LeRobot's outer model prefix and drop inference-unused keys.""" + if key.startswith("model."): + key = key[len("model.") :] + if key == "paligemma_with_expert.gemma_expert.lm_head.weight": + return None + return key + + +class PI05LiberoPolicy: + """Adapt vla-evaluation-harness LIBERO observations to ``PI05Processor``.""" + + def __init__( + self, + checkpoint_dir: str | Path, + *, + device: str = "cuda", + params_dtype: torch.dtype = torch.bfloat16, + max_batch_size: int = 1, + use_cuda_graph: bool = True, + attn_backend: str = "flashinfer", + norm_backend: str = "phyai-kernel", + linear_backend: str | None = "flashinfer", + flashinfer_workspace_bytes: int = 512 * 1024 * 1024, + tokenizer_name: str | None = None, + camera_mode: str | None = None, + ) -> None: + self.checkpoint_dir = Path(checkpoint_dir) + self.device = device + self.params_dtype = params_dtype + self.max_batch_size = int(max_batch_size) + self.config = self._read_config() + self.image_size = self._resolve_image_size(self.config) + self._action_dim = self._resolve_action_dim(self.config) + self.max_action_dim = int(self.config.get("max_action_dim", 32)) + self._chunk_size = int(self.config.get("chunk_size", PI05Config().chunk_size)) + self.camera_names = self._resolve_camera_names(camera_mode) + self.tokenizer_name = self._resolve_tokenizer_name(tokenizer_name) + self.prompt_mode = str( + self.config.get("phyai_prompt_mode", "lerobot_state_bins") + ) + self.normalization_mode = str( + self.config.get("phyai_normalization_mode", "mean_std") + ) + self._use_phyai_compat = ( + "phyai_prompt_mode" in self.config + or "phyai_normalization_mode" in self.config + ) + self._normalizer_stats = self._load_processor_state( + "policy_preprocessor.json", "normalizer_processor" + ) + self._unnormalizer_stats = self._load_processor_state( + "policy_postprocessor.json", "unnormalizer_processor" + ) + if self._use_phyai_compat: + self._validate_compat_stats() + self._tokenizer = None + self.processor = PI05Processor.from_pretrained( + self.checkpoint_dir, + tokenizer_name=self.tokenizer_name, + image_size=self.image_size, + num_channels=3, + num_images=len(self.camera_names), + action_dim=self._action_dim, + normalize_pixels=True, + device=device, + params_dtype=params_dtype, + ) + self.engine = Engine( + EngineArgs( + plugin="pi05", + plugin_args=PI05Args( + checkpoint_dir=self.checkpoint_dir, + max_batch_size=self.max_batch_size, + weight_remap=_lerobot_pi05_weight_remap, + inputs_image_shape=[ + [self.image_size, self.image_size, 3] for _ in self.camera_names + ], + ), + config=EngineConfig( + backends=BackendConfig( + attn=attn_backend, norm=norm_backend, linear=linear_backend + ), + device=DeviceConfig(target=device, params_dtype=params_dtype), + runtime=RuntimeConfig( + use_cuda_graph=use_cuda_graph, + flashinfer_workspace_bytes=flashinfer_workspace_bytes, + force_linear_kernel=linear_backend, + ), + ), + ) + ) + + @property + def chunk_size(self) -> int: + return self._chunk_size + + @property + def action_dim(self) -> int: + return int(self.processor.action_dim or self._action_dim) + + @staticmethod + def _resolve_image_size(config: dict[str, Any]) -> int: + resolution = config.get("image_resolution") + if isinstance(resolution, list) and resolution: + return int(resolution[0]) + return PI05Config().vision.image_size + + @staticmethod + def _resolve_action_dim(config: dict[str, Any]) -> int: + shape = config.get("output_features", {}).get("action", {}).get("shape") + if isinstance(shape, list) and shape: + return int(shape[-1]) + return 7 + + def _read_config(self) -> dict[str, Any]: + path = self.checkpoint_dir / "config.json" + if not path.exists(): + return {} + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + def _resolve_camera_names(self, camera_mode: str | None) -> list[str]: + mode = camera_mode or envs.PHYAI_CAMERA_MODE.get() or "three_camera" + if mode == "two_camera": + return ["agentview", "wrist"] + if mode == "three_camera": + return ["agentview", "wrist", "empty"] + raise ValueError(f"Unsupported PHYAI_CAMERA_MODE={mode!r}.") + + def _resolve_tokenizer_name(self, tokenizer_name: str | None) -> str: + if tokenizer_name: + return tokenizer_name + if env_tokenizer := envs.PHYAI_TOKENIZER_PATH.get(): + return env_tokenizer + if config_tokenizer := self.config.get("tokenizer_name"): + return str(config_tokenizer) + return PI05_DEFAULT_TOKENIZER_NAME + + def _load_processor_state( + self, config_name: str, registry_name: str + ) -> dict[str, torch.Tensor]: + path = self.checkpoint_dir / config_name + if not path.exists(): + return {} + with path.open("r", encoding="utf-8") as f: + config = json.load(f) + for step in config.get("steps", []): + if step.get("registry_name") != registry_name: + continue + state_file = step.get("state_file") + if not state_file: + return {} + return load_file(str(self.checkpoint_dir / state_file)) + return {} + + def _validate_compat_stats(self) -> None: + if self.normalization_mode == "openpi_quantile": + normalizer_keys = ("observation.state.min", "observation.state.max") + unnormalizer_keys = ("action.min", "action.max") + else: + normalizer_keys = ("observation.state.mean", "observation.state.std") + unnormalizer_keys = ("action.mean", "action.std") + missing = [ + f"normalizer:{key}" + for key in normalizer_keys + if key not in self._normalizer_stats + ] + missing.extend( + f"unnormalizer:{key}" + for key in unnormalizer_keys + if key not in self._unnormalizer_stats + ) + if missing: + raise ValueError( + f"{self.checkpoint_dir}: compat normalization requires missing stats " + f"{', '.join(missing)}" + ) + + @property + def tokenizer(self): + if self._tokenizer is None: + from transformers import AutoTokenizer + + self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name) + return self._tokenizer + + def observation_to_raw(self, obs: dict[str, Any]) -> dict[str, Any]: + return { + IMAGES: [ + self._extract_camera_tensor(obs, name) for name in self.camera_names + ], + STATE: self._extract_state(obs), + TASK: [self._extract_task(obs)], + } + + def observation_to_request_inputs( + self, obs: dict[str, Any] + ) -> dict[str, torch.Tensor]: + if not self._use_phyai_compat: + processed = self.processor.preprocess(self.observation_to_raw(obs)) + return { + "pixel_values": processed.pixel_values, + "input_ids": processed.input_ids, + "lang_lens": processed.lang_lens, + } + pixel_values = ( + torch.stack( + [ + self._extract_camera_model_tensor(obs, name).squeeze(0) + for name in self.camera_names + ], + dim=0, + ) + .unsqueeze(0) + .to(self.device) + ) + state = self._normalize_state(self._extract_state(obs)) + input_ids, lang_lens = self._tokenize_inputs([self._extract_task(obs)], state) + return { + "pixel_values": pixel_values, + "input_ids": input_ids.to(self.device), + "lang_lens": lang_lens.to(self.device), + } + + def _extract_camera_tensor( + self, obs: dict[str, Any], camera_name: str + ) -> torch.Tensor: + image = self._extract_camera_image(obs, camera_name) + return self._image_to_raw_tensor(image) + + def _extract_camera_model_tensor( + self, obs: dict[str, Any], camera_name: str + ) -> torch.Tensor: + image = self._extract_camera_image(obs, camera_name) + return self._image_to_model_tensor(image) + + def _extract_camera_image( + self, obs: dict[str, Any], camera_name: str + ) -> np.ndarray: + if camera_name == "agentview": + return self._extract_image(obs, LIBERO_AGENTVIEW_KEYS) + if camera_name == "wrist": + return self._extract_image(obs, LIBERO_WRIST_KEYS) + if camera_name == "empty": + return np.zeros((self.image_size, self.image_size, 3), dtype=np.uint8) + raise ValueError(f"Unsupported camera_name={camera_name!r}.") + + @staticmethod + def _extract_image(obs: dict[str, Any], keys: tuple[str, ...]) -> np.ndarray: + candidates: list[Any] = [] + images = obs.get("images") + if isinstance(images, dict): + candidates.extend(images.get(k) for k in keys) + candidates.extend(obs.get(k) for k in keys) + for candidate in candidates: + if candidate is None: + continue + array = np.asarray(candidate) + if array.ndim == 4: + array = array[0] + if array.ndim != 3: + continue + if array.shape[0] == 3 and array.shape[-1] != 3: + array = np.transpose(array, (1, 2, 0)) + if array.shape[-1] == 3: + return array + raise KeyError(f"LIBERO observation does not contain any image keys: {keys}.") + + @staticmethod + def _image_to_raw_tensor(image: np.ndarray) -> torch.Tensor: + array = np.asarray(image, dtype=np.float32) + if array.max(initial=0.0) > 1.0: + array = array / 255.0 + return ( + torch.from_numpy(np.ascontiguousarray(array.transpose(2, 0, 1))) + .unsqueeze(0) + .contiguous() + ) + + def _image_to_model_tensor(self, image: np.ndarray) -> torch.Tensor: + tensor = self._image_to_raw_tensor(image) + if tensor.shape[-2:] != (self.image_size, self.image_size): + tensor = self._resize_with_pad(tensor, self.image_size, self.image_size) + return (tensor * 2.0 - 1.0).contiguous() + + @staticmethod + def _resize_with_pad(images: torch.Tensor, height: int, width: int) -> torch.Tensor: + _, _, cur_height, cur_width = images.shape + ratio = max(cur_width / width, cur_height / height) + resized_height = int(cur_height / ratio) + resized_width = int(cur_width / ratio) + resized = F.interpolate( + images, + size=(resized_height, resized_width), + mode="bilinear", + align_corners=False, + ) + resized = resized.clamp(0.0, 1.0) + pad_h0, rem_h = divmod(height - resized_height, 2) + pad_w0, rem_w = divmod(width - resized_width, 2) + return F.pad( + resized, + (pad_w0, pad_w0 + rem_w, pad_h0, pad_h0 + rem_h), + mode="constant", + value=0.0, + ) + + @staticmethod + def _extract_state(obs: dict[str, Any]) -> torch.Tensor: + state = obs.get("states", obs.get("state")) + if state is None: + raise KeyError("LIBERO observation must contain 'states' or 'state'.") + array = np.asarray(state, dtype=np.float32) + if array.ndim == 1: + array = array[None, :] + return torch.from_numpy(np.ascontiguousarray(array)) + + @staticmethod + def _extract_task(obs: dict[str, Any]) -> str: + task = obs.get("task_description", obs.get("task", "")) + if isinstance(task, (list, tuple)): + task = task[0] if task else "" + return str(task) + + def _normalize_state(self, state: torch.Tensor) -> torch.Tensor: + if self.normalization_mode == "openpi_quantile": + min_v = self._normalizer_stats.get("observation.state.min") + max_v = self._normalizer_stats.get("observation.state.max") + if min_v is None or max_v is None: + return state + return (state - min_v.to(state)) / ( + max_v.to(state) - min_v.to(state) + 1e-6 + ) * 2.0 - 1.0 + mean = self._normalizer_stats.get("observation.state.mean") + std = self._normalizer_stats.get("observation.state.std") + if mean is None or std is None: + return state + return (state - mean.to(state)) / torch.clamp(std.to(state), min=1e-8) + + def _tokenize_inputs( + self, tasks: list[str], states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.prompt_mode == "openpi_task": + prompts = [ + task.strip().replace("_", " ").replace("\n", " ") + "\n" + for task in tasks + ] + else: + state_np = states.detach().cpu().numpy() + bins = np.linspace(-1.0, 1.0, 257)[:-1] + discretized = np.digitize(state_np, bins=bins) - 1 + discretized = np.clip(discretized, 0, 255) + prompts = [] + for task, state_bins in zip(tasks, discretized): + cleaned = task.strip().replace("_", " ").replace("\n", " ") + state_str = " ".join(map(str, state_bins)) + prompts.append(f"Task: {cleaned}, State: {state_str};\nAction: ") + encoded = self.tokenizer( + prompts, + max_length=int(self.config.get("tokenizer_max_length", 200)), + padding="max_length", + padding_side="right", + truncation=True, + return_tensors="pt", + ) + return encoded["input_ids"].to(torch.int64), encoded["attention_mask"].sum( + dim=-1 + ).to(torch.int64) + + def _postprocess_actions(self, raw_actions: torch.Tensor) -> np.ndarray: + action = raw_actions[..., : self.action_dim].detach().float() + if not self._use_phyai_compat: + actions = self.processor.postprocess(action) + if isinstance(actions, torch.Tensor): + actions = actions.detach().cpu().numpy() + return np.asarray(actions, dtype=np.float32) + action = action.cpu() + if self.normalization_mode == "openpi_quantile": + min_v = self._unnormalizer_stats.get("action.min") + max_v = self._unnormalizer_stats.get("action.max") + if min_v is not None and max_v is not None: + action = (action + 1.0) / 2.0 * ( + max_v.to(action) - min_v.to(action) + 1e-6 + ) + min_v.to(action) + else: + mean = self._unnormalizer_stats.get("action.mean") + std = self._unnormalizer_stats.get("action.std") + if mean is not None and std is not None: + action = action * torch.clamp(std.to(action), min=1e-8) + mean.to( + action + ) + return action.numpy().astype(np.float32) + + def infer( + self, obs: dict[str, Any], *, noise: torch.Tensor | np.ndarray | None = None + ) -> dict[str, np.ndarray]: + request_kwargs = self.observation_to_request_inputs(obs) + if noise is not None: + request_kwargs["noise"] = torch.as_tensor(noise, device=self.device) + request = PI05Request(**request_kwargs) + with torch.inference_mode(): + raw_actions = self.engine.step(request) + actions = self._postprocess_actions(raw_actions) + return {"actions": actions} + + def close(self) -> None: + self.engine.close() From 3f9b9f86f8468bb66c05c6ce54c6e1c2e31ce6f8 Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 04:45:40 +0000 Subject: [PATCH 18/29] docs: remove markdown files from pi05 pr --- .../pi05/README_external_runtime_latency.md | 214 ------- .../pi05/eight_gpu_inference_tutorial.en.md | 193 ------ ...ai_pi05_libero_four_suites_reproduction.md | 592 ------------------ 3 files changed, 999 deletions(-) delete mode 100644 benchmark/pi05/README_external_runtime_latency.md delete mode 100644 benchmark/pi05/eight_gpu_inference_tutorial.en.md delete mode 100644 docs/phyai_pi05_libero_four_suites_reproduction.md diff --git a/benchmark/pi05/README_external_runtime_latency.md b/benchmark/pi05/README_external_runtime_latency.md deleted file mode 100644 index a16a7af..0000000 --- a/benchmark/pi05/README_external_runtime_latency.md +++ /dev/null @@ -1,214 +0,0 @@ -# External PI0.5 runtime latency wrappers - -This document explains how to set up and run the three external PI0.5 latency -wrappers under `benchmark/pi05/`. - -These scripts do not use the PhyAI engine for inference. Each script calls the -target runtime directly, while reusing PhyAI's common benchmark runner for -warmup, timing, and JSONL output. - -| Script | Runtime measured | Timed call | -| --- | --- | --- | -| `bench_flashrt_pi05.py` | FlashRT | `Pi05TorchFrontendRtx.infer(obs)` | -| `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward(...)` | -| `bench_vlacpp_pi05_client.py` | vla.cpp | one ZMQ request to a running `vla-server` | - -## Common setup - -Use one Python environment that can import PhyAI, PyTorch, and -`benchmark/bench_n_batch.py`. - -```bash -cd -python -c "import torch; import phyai; import benchmark.bench_n_batch" -nvidia-smi -``` - -Use the same benchmark settings when comparing runtimes: - -```text -batch size: 1 -views / camera streams: 2 -chunk size: 50 -prompt: keep the same text across runs -warmup / timed iterations: use the same values across runs -precision: label each row by the runtime's real precision path -``` - -The wrappers generate synthetic image/state inputs. They are for latency-only -measurements, not LIBERO accuracy evaluation. - -## Placeholders - -Use your own paths for these placeholders: - -| Placeholder | Meaning | -| --- | --- | -| `` | PhyAI checkout containing `benchmark/pi05/` | -| `` | FlashRT checkout | -| `` | realtime-vla checkout | -| `` | vla.cpp checkout | -| `` | compiled vla.cpp `vla-server` binary | -| `` | PI0.5 safetensors checkpoint directory or file | -| `` | PI0.5 GGUF file for vla.cpp | -| `` | vla.cpp multimodal projector GGUF | -| `` | local tokenizer directory | -| `` | local LIBERO `meta/stats.json` with `observation.state.q01/q99` | - -## FlashRT - -Install FlashRT following its official README, then make the checkout visible: - -```bash -export FLASHRT_ROOT= -cd -python -c "import sys; sys.path.insert(0, '$FLASHRT_ROOT'); from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx; print(Pi05TorchFrontendRtx)" -``` - -Run latency: - -```bash -cd -python benchmark/pi05/bench_flashrt_pi05.py \ - --flashrt-root \ - --checkpoint \ - --precision bf16 \ - --num-views 2 \ - --chunk-size 50 \ - --batch-sizes 1 \ - --n-warmup 100 \ - --n-timed 100 \ - --result-file results/flashrt_pi05.jsonl -``` - -Notes: - -- The script uses FlashRT's direct `Pi05TorchFrontendRtx` API because `chunk_size` - is a frontend constructor argument. -- Do not use `load_model(..., num_steps=50)` to set action chunk size. In - FlashRT, `num_steps` means denoise steps. -- `--precision bf16` sets FlashRT's forced-BF16 PI0.5 RTX path. Use - `--precision fp8_bf16` for FlashRT's optimized FP8/BF16 path and label that - result separately. - -## realtime-vla - -Install realtime-vla following its official README. The wrapper can read a -converted `.pt` / `.pth` checkpoint directly. If you pass a PI0.5 safetensors -checkpoint, also provide FlashRT so the wrapper can reuse FlashRT's PI0.5 -conversion helper. - -```bash -export REALTIME_VLA_ROOT= -export FLASHRT_ROOT= -cd -python -c "import sys; sys.path.insert(0, '$REALTIME_VLA_ROOT'); from pi05_infer import Pi05Inference; print(Pi05Inference)" -``` - -Run latency: - -```bash -cd -python benchmark/pi05/bench_realtime_vla_pi05.py \ - --realtime-vla-root \ - --flashrt-root \ - --checkpoint \ - --num-views 2 \ - --chunk-size 50 \ - --prompt-len 16 \ - --batch-sizes 1 \ - --n-warmup 100 \ - --n-timed 100 \ - --result-file results/realtime_vla_pi05.jsonl -``` - -Notes: - -- The wrapper uses BF16 synthetic inputs. -- For `.pkl` / `.pickle` checkpoints, add `--trust-pickle-checkpoint` only when - the file is trusted. -- If the checkpoint does not contain `language_embeds`, the wrapper creates a - synthetic prompt embedding for latency-only runs. - -## vla.cpp - -vla.cpp uses a server/client flow. Build `vla-server` with CUDA enabled, then -start the server in one shell and run the Python benchmark client in another. - -Basic checks: - -```bash -test -x -test -f -test -f -test -f /tokenizer.json -test -f -``` - -Start server: - -```bash - \ - --bind tcp://127.0.0.1:5555 \ - --timing-detail phase \ - \ - -``` - -Run latency client: - -```bash -cd -python benchmark/pi05/bench_vlacpp_pi05_client.py \ - --vlacpp-root \ - --addr tcp://127.0.0.1:5555 \ - --arch pi05 \ - --tokenizer \ - --stats-json \ - --num-views 2 \ - --chunk-size 50 \ - --batch-sizes 1 \ - --n-warmup 100 \ - --n-timed 100 \ - --result-file results/vlacpp_pi05.jsonl -``` - -Notes: - -- vla.cpp needs GGUF files; a safetensors checkpoint is not enough. -- Prefer local tokenizer and stats files to avoid network or HuggingFace auth - issues during benchmarking. For PI0.5, vla.cpp expects a lerobot-style - `meta/stats.json`; OpenPI-style `norm_stats.json` is not the same format. -- The wrapper records client wall latency. If the server returns phase timing, - it is written under `extras.server_phase_latency_ms`. - -## Quick validation - -After each runtime is installed, reduce iterations to check that the wrapper can -start and write JSONL: - -```bash ---n-warmup 1 --n-timed 1 --result-file results/pi05_smoke.jsonl -``` - -For vla.cpp, keep `vla-server` running before starting the client smoke test. - -## Troubleshooting - -| Symptom | Check | -| --- | --- | -| `No module named flash_rt` | Pass `--flashrt-root` or set `FLASHRT_ROOT`. | -| `No module named pi05_infer` | Pass `--realtime-vla-root` or set `REALTIME_VLA_ROOT`. | -| realtime-vla safetensors conversion fails | Also pass `--flashrt-root`; verify FlashRT import works. | -| vla.cpp client cannot connect | Confirm the server printed that it is ready and the `--addr` matches `--bind`. | -| vla.cpp tokenizer downloads or asks for auth | Use a local tokenizer directory. | -| GPU architecture build error | Check CUDA, PyTorch CUDA, driver, and build flags for the target GPU. | -| Latency is much slower than expected | Check `nvidia-smi`, rerun after warmup/JIT, and make sure no other process is using the GPU. | - -## Timing scope - -- FlashRT: wall time around steady-state `Pi05TorchFrontendRtx.infer(obs)`, - after prompt setup, calibration, and first graph-building call. -- realtime-vla: CUDA-event time around one `Pi05Inference.forward(...)` call. -- vla.cpp: client wall time for one ZMQ request; server phase timing is copied - from the response when available. diff --git a/benchmark/pi05/eight_gpu_inference_tutorial.en.md b/benchmark/pi05/eight_gpu_inference_tutorial.en.md deleted file mode 100644 index ba6d6cd..0000000 --- a/benchmark/pi05/eight_gpu_inference_tutorial.en.md +++ /dev/null @@ -1,193 +0,0 @@ -# PhyAI pi0.5 Eight-GPU Inference Tutorial - -Reproduce **PhyAI `pi05_wn` 8-GPU DP inference + concurrent LIBERO demos**. -Chinese: [`八卡推理从零开始教程.md`](./八卡推理从零开始教程.md) - ---- - -## 1. What you get - -```text -8-GPU PhyAI WebSocket server (port 8000) - ↑ -32 LIBERO shards (4 suites × 8 tasks) - ↓ -Per-shard JSON (success / timing) + optional wait videos -``` - -Recommended settings: - -| Item | Value | -| --- | --- | -| GPUs / batch | 8 GPUs, `MAX_BATCH_SIZE=32` (B=4 per GPU) | -| Chunk | `CHUNK_SIZE=10`, `SEND_ACTION_CHUNKS=1` (true chunk=10) | -| Batching wait | `MAX_WAIT_TIME=0.02` | -| Recording (optional) | `continuous` + 20fps | - -Architecture: all LIBERO clients talk to rank0; after the batch fills (or timeout), DP scatter → each GPU runs its slice → gather and reply. All 8 ranks sync in one step. CUDA graphs use a fixed padded shape, so a partial batch is not proportionally faster. - -Chunk modes: - -| Mode | Flag | Behavior | -| --- | --- | --- | -| True chunk=10 | `SEND_ACTION_CHUNKS=1` | Return 10 actions once; client runs them locally, then requests again | -| Pseudo chunk=1 | `SEND_ACTION_CHUNKS=0` | Model still produces 10; server returns 1 and buffers the rest; request every step | - ---- - -## 2. Setup - -Needs: 8 free GPUs, Docker, `tmux`. - -```bash -export WORKSPACE="$HOME/phyai_workplace" # change me: must contain phyai / vla-evaluation-harness / phyai_models -export PHYAI_ROOT="$WORKSPACE/phyai" -export VLA_ROOT="$WORKSPACE/vla-evaluation-harness" -export MODEL_ROOT="$WORKSPACE/phyai_models" -export DEMO_ROOT="$PHYAI_ROOT/libero_wn_demo" -``` - -Checks: - -```bash -ls "$MODEL_ROOT/pi05_libero_phyai_converted" "$MODEL_ROOT/paligemma-3b-pt-224" -ls "$PHYAI_ROOT/.venv/bin/torchrun" "$VLA_ROOT/.venv/bin/vla-eval" -``` - -Images: - -```bash -sg docker -c 'docker pull nvcr.io/nvidia/pytorch:25.12-py3' -sg docker -c 'docker pull ghcr.io/allenai/vla-evaluation-harness/libero:latest' -``` - -Tokenizer must be offline: place `$MODEL_ROOT/paligemma-3b-pt-224` first (`HF_HUB_OFFLINE=1` is set by the script). -Always pass `PHYAI_ROOT` / `VLA_ROOT` / `MODEL_ROOT` explicitly; do not copy another machine’s absolute paths. - ---- - -## 3. Shortest path (batch32 + true chunk10) - -### A. Confirm idle - -```bash -nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader,nounits -ss -ltnp | grep 8000 || echo '8000 free' -``` - -Cleanup leftovers if needed: - -```bash -tmux kill-session -t phyai_pi05_wn_demo 2>/dev/null || true -sg docker -c 'docker rm -f phyai_pi05_wn_demo' 2>/dev/null || true -``` - -### B. Start server - -```bash -sg docker -c " -SESSION=phyai_pi05_wn_demo CONTAINER=phyai_pi05_wn_demo \ -PHYAI_ROOT=$PHYAI_ROOT VLA_ROOT=$VLA_ROOT MODEL_ROOT=$MODEL_ROOT \ -MAX_BATCH_SIZE=32 MAX_WAIT_TIME=0.02 CHUNK_SIZE=10 SEND_ACTION_CHUNKS=1 \ -MASTER_PORT=29619 \ -bash $DEMO_ROOT/start_pi05_wn_server.sh -" -``` - -After ~1–3 minutes: - -```bash -curl -sS http://127.0.0.1:8000/config -# expect max_batch_size=32 -tmux attach -t phyai_pi05_wn_demo # detach: Ctrl-b d -``` - -### C. Run 32 shards - -```bash -sg docker -c "bash $DEMO_ROOT/run_libero_32_clientchunk10_maxwait002.sh" -``` - -Outputs under `$VLA_ROOT/results/clientchunk10_maxwait002/`: suite `*.json` and `wait_videos/*.mp4` (pixels are only in the videos). - -### D. Cleanup - -```bash -tmux kill-session -t phyai_pi05_wn_demo -sg docker -c 'docker rm -f phyai_pi05_wn_demo' 2>/dev/null || true -``` - ---- - -## 4. Variants - -**Pseudo chunk=1 (batch32)** - -```bash -# same server command with SEND_ACTION_CHUNKS=0 and new SESSION/CONTAINER/MASTER_PORT -sg docker -c "bash $DEMO_ROOT/run_libero_32_serverchunk1_maxwait002.sh" -``` - -**batch16 (B=2 per GPU, 16 shards)** - -```bash -# MAX_BATCH_SIZE=16, SEND_ACTION_CHUNKS=0 or 1 -sg docker -c "bash $DEMO_ROOT/run_libero_batch16_serverchunk1_maxwait002.sh" -``` - -Start batch16 and batch32 as separate servers; do not switch inside one `torchrun`. - -**Custom experiment**: copy `$DEMO_ROOT/clientchunk10_maxwait002/`, edit yaml `output_dir` / suite, then point the client script’s `DEMO_ROOT` and `RESULTS_ROOT` at your dirs. - -**Pure inference latency (no LIBERO)**: - -```bash -sg docker -c "bash $PHYAI_ROOT/benchmark/run_pi05_wn_latency_dp8_docker.sh" -``` - -Reference: batch16 ~60ms, batch32 ~97ms (pure `Engine.step`, not end-to-end). - ---- - -## 5. Recording and results - -For end-to-end stalls use `continuous`, not `step`: - -```yaml -docker: - env: - - VLA_EVAL_WAIT_VIDEO_DIR=/workspace/results/wait_videos - - VLA_EVAL_WAIT_VIDEO_MODE=continuous - - VLA_EVAL_WAIT_VIDEO_FPS=20 - - VLA_EVAL_WAIT_VIDEO_MODEL= -``` - -`continuous` freezes while waiting on the model; freeze length ≈ real wait. `step` removes waits. Changing export fps does not shorten real stalls. - -Useful JSON fields: `metrics.success`, `avg_model_wait_ms`, `model_buffer_hits`, `model_inference_calls`, `wait_video_mode`. -End-to-end wait ≈ queue + predict_batch + ws, often much larger than pure GPU bench. - ---- - -## 6. Troubleshooting - -| Symptom | Fix | -| --- | --- | -| Stuck in setup | Offline tokenizer; change `MASTER_PORT`; remove same-name container; free GPUs | -| `:8000/config` fails | Wait for graph capture; `tmux capture-pane -t -p -S -80` | -| Shard cannot connect | Host networking; URL=`ws://127.0.0.1:8000`; `NO_PROXY='*'` | -| Partial results | Check `$VLA_ROOT/results//*_logs/`; `taskunknown` ≈ task0 | -| FlashInfer/CUDA clash | Use Docker; avoid bare-metal host runs | - ---- - -## 7. File index (under `$DEMO_ROOT`) - -| File | Purpose | -| --- | --- | -| `start_pi05_wn_server.sh` | 8-GPU server | -| `run_libero_32_clientchunk10_maxwait002.sh` | batch32 true chunk10 | -| `run_libero_32_serverchunk1_maxwait002.sh` | batch32 pseudo chunk=1 | -| `run_libero_batch16_serverchunk1_maxwait002.sh` | batch16 pseudo chunk=1 | -| `clientchunk10_maxwait002/` etc. | continuous demo configs | -| `experiment_setup.md` | Historical experiment notes | diff --git a/docs/phyai_pi05_libero_four_suites_reproduction.md b/docs/phyai_pi05_libero_four_suites_reproduction.md deleted file mode 100644 index e405e8d..0000000 --- a/docs/phyai_pi05_libero_four_suites_reproduction.md +++ /dev/null @@ -1,592 +0,0 @@ ---- -title: PhyAI pi0.5 LIBERO four-suite reproduction -description: Run the pi0.5 LIBERO policy with PhyAI on all four LIBERO benchmark suites. ---- - -# PhyAI pi0.5 LIBERO four-suite reproduction - -This guide shows how to run the pi0.5 LIBERO policy with PhyAI and evaluate it with `vla-evaluation-harness` on all four LIBERO suites. -It is written for a fresh machine and avoids local machine-specific paths by using environment variables. - -The four suites are: - -```text -libero_spatial -> configs/benchmarks/libero/spatial.yaml -libero_object -> configs/benchmarks/libero/object.yaml -libero_goal -> configs/benchmarks/libero/goal.yaml -libero_10 -> configs/benchmarks/libero/10.yaml -``` - -The benchmark setup in this guide uses: - -```text -Mode: sync -Chunk size: 10 -Episodes per suite: 10 tasks x 50 episodes = 500 episodes -Total episodes: 4 suites x 500 episodes = 2000 episodes -Model: PhyAI pi0.5 LIBERO converted checkpoint -Simulator: vla-evaluation-harness LIBERO Docker container -Output: one JSON result file per suite with success, steps, timing, and chunk-size fields -``` - -## 1. Prerequisites - -Use a Linux machine with: - -```text -GPU: at least 1 CUDA GPU, 48 GB or more GPU memory recommended -Container runtime: Docker and NVIDIA Container Toolkit -Python environment manager: uv -Utility tools: tmux, nvidia-smi, ss -``` - -Prepare these model resources before you start: - -```text -PhyAI converted checkpoint: pi05_libero_phyai_converted -PaLI-Gemma tokenizer / processor: paligemma-3b-pt-224 -``` - -`paligemma-3b-pt-224` is a gated resource. Prefer syncing it from a machine that already has access instead of downloading it during reproduction. - -## 2. Set environment variables - -Set paths for the target machine: - -```bash -export PHYAI_ROOT=$HOME/phyai -export VLA_ROOT=$HOME/vla-evaluation-harness -export MODEL_ROOT=$HOME/phyai_models -export PHYAI_CONTAINER=phyai_libero_eval - -export PHYAI_CKPT_HOST=$MODEL_ROOT/pi05_libero_phyai_converted -export TOKENIZER_HOST=$MODEL_ROOT/paligemma-3b-pt-224 - -export PHYAI_CKPT_IN_CONTAINER=/data/share/pi05_libero_phyai_converted -export TOKENIZER_IN_CONTAINER=/data/share/paligemma-3b-pt-224 - -export LIBERO_IMAGE=ghcr.io/allenai/vla-evaluation-harness/libero:latest -``` - -Check that the model directories exist: - -```bash -test -d "$PHYAI_CKPT_HOST" -test -d "$TOKENIZER_HOST" -``` - -If the models are on another machine, sync them to the target machine: - -```bash -rsync -azP \ - /path/to/pi05_libero_phyai_converted \ - /path/to/paligemma-3b-pt-224 \ - user@target-host:$MODEL_ROOT/ -``` - -## 3. Clone source code and create environments - -Clone PhyAI: - -```bash -git clone https://github.com/MEmbodied/phyai.git "$PHYAI_ROOT" -cd "$PHYAI_ROOT" -uv sync -``` - -Clone `vla-evaluation-harness`: - -```bash -git clone https://github.com/allenai/vla-evaluation-harness.git "$VLA_ROOT" -cd "$VLA_ROOT" -uv sync -./.venv/bin/vla-eval --help >/tmp/vla_eval_help.log -``` - -If the target machine has no network access, clone both repositories on a networked machine and sync them with `rsync`. -After syncing, still run `uv sync` on the target machine so editable paths, Python versions, CUDA libraries, and local dependencies are resolved correctly. - -## 4. Prepare the LIBERO Docker image - -`vla-evaluation-harness` runs LIBERO inside a benchmark container. -On an `x86_64` machine, pull the official image: - -```bash -docker pull "$LIBERO_IMAGE" -``` - -If the target machine is ARM64 and the official image is only available for `amd64`, build an ARM64 LIBERO image locally: - -```bash -cd "$VLA_ROOT" -export DOCKER_DEFAULT_PLATFORM=linux/arm64 -docker/build.sh libero - -docker image inspect "$LIBERO_IMAGE" \ - --format '{{.Architecture}} {{.Os}}' -``` - -Expected output on ARM64: - -```text -arm64 linux -``` - -Expected output on `x86_64`: - -```text -amd64 linux -``` - -## 5. Create the PhyAI Docker container - -Run the PhyAI server inside a Docker container. -Mount the PhyAI source tree, the `vla-evaluation-harness` source tree, and the model directory: - -```bash -docker run -dit --gpus all \ - -v "$PHYAI_ROOT":/phyai_workspace \ - -v "$VLA_ROOT":/vla-evaluation-harness \ - -v "$MODEL_ROOT":/data/share \ - -w /phyai_workspace \ - --cap-add=SYS_ADMIN \ - --ipc=host \ - --cap-add=SYS_PTRACE \ - --shm-size=4G \ - --security-opt seccomp=unconfined \ - --security-opt apparmor=unconfined \ - --name "$PHYAI_CONTAINER" \ - nvcr.io/nvidia/pytorch:25.12-py3 bash -``` - -Install the PhyAI environment inside the container: - -```bash -docker exec "$PHYAI_CONTAINER" bash -lc ' -cd /phyai_workspace -python3 -m pip install -U uv -uv sync -' -``` - -If `uv sync` produces editable paths that point to the host path instead of the container path, create a compatibility symlink. -Only run this if imports fail because a stale host path is referenced: - -```bash -export HOST_USER=$(id -un) -export COMPAT_PARENT=/compat_mount - -docker exec "$PHYAI_CONTAINER" bash -lc " -mkdir -p $COMPAT_PARENT/$HOST_USER -ln -sfn /phyai_workspace $COMPAT_PARENT/$HOST_USER/phyai -" -``` - -Verify imports inside the container: - -```bash -docker exec "$PHYAI_CONTAINER" bash -lc ' -cd /phyai_workspace -export PYTHONPATH=/phyai_workspace/phyai/src:/phyai_workspace/phyai-kernel:/phyai_workspace/phyai-utils-tools/src:/vla-evaluation-harness/src -/phyai_workspace/.venv/bin/python - <&1 | tee $VLA_ROOT/results/phyai_pi05_libero_server.log -" -``` - -Key settings: - -| Setting | Value | Purpose | -| --- | --- | --- | -| `--checkpoint_path` | `/data/share/pi05_libero_phyai_converted` | PhyAI converted pi0.5 LIBERO checkpoint | -| `PHYAI_TOKENIZER_PATH` | `/data/share/paligemma-3b-pt-224` | Tokenizer and processor directory | -| `PHYAI_CAMERA_MODE` | `two_camera` | LIBERO sends both agent-view and wrist-camera images | -| `--params_dtype` | `bfloat16` | Parameter dtype | -| `--attn_backend` | `flashinfer` | Attention backend | -| `--norm_backend` | `phyai-kernel` | Normalization backend | -| `--linear_backend` | `flashinfer` | Linear backend | -| `--flashinfer_workspace_bytes` | `536870912` | 512 MiB FlashInfer workspace | -| `--chunk_size` | `10` | The policy returns 10 actions per inference call | -| CUDA graph | Enabled by default | Do not pass `--no-use_cuda_graph` | - -Follow the server log: - -```bash -tail -f "$VLA_ROOT/results/phyai_pi05_libero_server.log" -``` - -Wait until the log contains: - -```text -capturing vision-tower CUDA graph -capturing 4 prefix-forward CUDA graph(s) -capturing the full 10-step Euler loop as one CUDA graph -Starting server on ws://0.0.0.0:8000 -``` - -## 8. Run a smoke test - -Run a minimal smoke test before launching the full benchmark. -This checks the model, LIBERO Docker image, WebSocket connection, timing fields, and chunk-size fields. - -```bash -cat > "$VLA_ROOT/configs/benchmarks/libero/smoke_test_phyai_local.yaml" < "$VLA_ROOT/run_phyai_pi05_libero_four_suites.sh" <<'SH' -#!/usr/bin/env bash -set -euo pipefail - -: "${PHYAI_SERVER_URL:?must set PHYAI_SERVER_URL}" -: "${PHYAI_CKPT_IN_CONTAINER:?must set PHYAI_CKPT_IN_CONTAINER}" -: "${LIBERO_IMAGE:=ghcr.io/allenai/vla-evaluation-harness/libero:latest}" - -cd "$(dirname "$0")" - -RUN_ID="phyai_pi05_libero_four_$(date +%Y%m%d_%H%M%S)" -OUT="results/${RUN_ID}" -mkdir -p "$OUT/configs" - -{ - echo "RUN_ID=${RUN_ID}" - echo "START=$(date -Is)" - echo "MODEL=phyai_pi05" - echo "SERVER_URL=${PHYAI_SERVER_URL}" - echo "CHECKPOINT=${PHYAI_CKPT_IN_CONTAINER}" - echo "PHYAI_CAMERA_MODE=two_camera" - echo "MODE=sync" - echo "CHUNK_SIZE=10" - echo "SERVER_CONFIG=use_cuda_graph=True attn=flashinfer norm=phyai-kernel linear=flashinfer workspace=536870912 params_dtype=bfloat16" -} | tee "$OUT/run_summary.log" - -make_cfg() { - local suite_name="$1" - local suite="$2" - local cfg="$3" - cat > "$cfg" <&1 | tee "$log" - status=${PIPESTATUS[0]} - - echo "SUITE_END model=phyai suite=${name} status=${status} log=${log} time=$(date -Is)" | tee -a "$OUT/run_summary.log" - - result_json=$(ls -t "$OUT/${suite}_sync_"*.json "$OUT"/*"${suite}"*_sync_*.json 2>/dev/null | head -1 || true) - if [ -n "$result_json" ]; then - ./.venv/bin/python scripts/summarize_timing.py "$result_json" | sed "s/^/TIMING phyai_${name} /" | tee -a "$OUT/run_summary.log" - else - echo "WARN no result json found for ${name}" | tee -a "$OUT/run_summary.log" - fi - - if [ "$status" -ne 0 ]; then - exit "$status" - fi -done - -echo "ALL_DONE $(date -Is)" | tee -a "$OUT/run_summary.log" -SH -chmod +x "$VLA_ROOT/run_phyai_pi05_libero_four_suites.sh" -``` - -Start the full run: - -```bash -cd "$VLA_ROOT" -tmux new-session -d -s phyai_pi05_libero_four \ - "PHYAI_SERVER_URL=$PHYAI_SERVER_URL PHYAI_CKPT_IN_CONTAINER=$PHYAI_CKPT_IN_CONTAINER LIBERO_IMAGE=$LIBERO_IMAGE ./run_phyai_pi05_libero_four_suites.sh" -``` - -Check progress: - -```bash -tmux ls -cd "$VLA_ROOT" -latest=$(ls -td results/phyai_pi05_libero_four_* | head -1) -sed -n '1,220p' "$latest/run_summary.log" -tail -80 "$latest"/phyai_spatial.log -``` - -## 10. Summarize success rate and timing - -After the run finishes, print the run summary: - -```bash -cd "$VLA_ROOT" -latest=$(ls -td results/phyai_pi05_libero_four_* | head -1) -cat "$latest/run_summary.log" -``` - -Summarize timing from all result JSON files: - -```bash -cd "$VLA_ROOT" -latest=$(ls -td results/phyai_pi05_libero_four_* | head -1) -./.venv/bin/python scripts/summarize_timing.py "$latest"/*.json -``` - -Summarize success rate: - -```bash -cd "$VLA_ROOT" -latest=$(ls -td results/phyai_pi05_libero_four_* | head -1) -./.venv/bin/python - <<'PY' "$latest"/*.json -import json -import sys -from pathlib import Path - -for arg in sys.argv[1:]: - p = Path(arg) - data = json.loads(p.read_text()) - eps = [ep for task in data.get("tasks", []) for ep in task.get("episodes", [])] - succ = sum(1 for ep in eps if ep.get("metrics", {}).get("success")) - total = len(eps) - rate = succ / total * 100.0 if total else 0.0 - steps = sum(int(ep.get("steps", 0)) for ep in eps) - print(f"{p.name}: success={succ}/{total} rate={rate:.1f}% steps={steps} mean_success={data.get('mean_success')}") -PY -``` - -Record these fields for each suite: - -```text -RUN_ID -Result directory -Checkpoint -Server URL -Suite -Success / total -Success rate -Steps -/usr/bin/time -p real -model_wait_sec -model_inference_sec -env_step_sec -obs_sec -avg_model_inference_ms -model_inference_calls -model_buffer_hits -raw_chunk_size_max -served_chunk_size_max -``` - -`raw_chunk_size_max=10` and `served_chunk_size_max=10` are key checks for this four-suite reproduction. - -## 11. Reference results - -Timing depends on the GPU, driver, machine load, and container environment. -Success rate can also vary slightly across runs. -The following results are from a previous run with the same evaluation setup: - -| Suite | Success rate | Success | Steps | Wall time | Model inference time | Env step time | Benchmark action-wait time | Average model inference | Inference calls | Buffer hits | Chunk check | -| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | -| `libero_spatial` | 97.8% | 489/500 | 52,926 | 4,198.73s | 200.56s | 1,239.61s | 2,269.16s | 36.33ms | 5,520 | 47,406 | raw=10, served=10 | -| `libero_object` | 99.8% | 499/500 | 68,712 | 5,199.23s | 256.92s | 1,169.84s | 3,453.50s | 36.18ms | 7,102 | 61,610 | raw=10, served=10 | -| `libero_goal` | 98.0% | 490/500 | 56,292 | 3,981.09s | 210.86s | 1,045.35s | 2,365.71s | 36.10ms | 5,841 | 50,451 | raw=10, served=10 | -| `libero_10` | 94.2% | 471/500 | 134,962 | 9,150.13s | 496.56s | 2,227.75s | 6,237.24s | 36.21ms | 13,713 | 121,249 | raw=10, served=10 | - -Use these numbers as references, not strict pass/fail thresholds. -For reproduction, first confirm: - -```text -All four suites finish 500 episodes each -Chunk check is raw=10 served=10 -The server log confirms CUDA graph capture -The result JSON files contain timing fields -``` - -## 12. Troubleshooting - -### 12.1 LIBERO Docker image architecture mismatch - -If the target machine is ARM64 and the official image is only available for `amd64`, build the image locally: - -```bash -cd "$VLA_ROOT" -export DOCKER_DEFAULT_PLATFORM=linux/arm64 -docker/build.sh libero -``` - -### 12.2 The benchmark cannot connect to the PhyAI server - -If the PhyAI server runs in a bridge Docker container, the host-side `vla-eval` process should connect to the container IP: - -```bash -export PHYAI_CONTAINER_IP=$(docker inspect -f '{{range.NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$PHYAI_CONTAINER") -export PHYAI_SERVER_URL=ws://$PHYAI_CONTAINER_IP:8000 -``` - -### 12.3 The result JSON has no timing fields - -Run the benchmark with `--dev`. -This mounts the host `$VLA_ROOT/src` tree into the LIBERO container and ensures the benchmark uses the runner that records timing fields. - -### 12.4 The PaLI-Gemma tokenizer is missing - -`paligemma-3b-pt-224` is a gated resource. -Sync it from an existing machine to `$MODEL_ROOT/paligemma-3b-pt-224` instead of relying on an on-the-fly download. - -### 12.5 The PhyAI server log has no CUDA graph capture messages - -Check whether the server command accidentally passed `--no-use_cuda_graph` or whether it started the wrong server adapter. -The expected log must include: - -```text -capturing vision-tower CUDA graph -capturing 4 prefix-forward CUDA graph(s) -capturing the full 10-step Euler loop as one CUDA graph -``` - -### 12.6 Release resources - -```bash -tmux kill-session -t phyai_pi05_libero_four || true -tmux kill-session -t phyai_pi05_libero_server || true -docker stop "$PHYAI_CONTAINER" || true -ss -ltnp | grep ':8000' || true -nvidia-smi --query-gpu=index,memory.used,memory.total,utilization.gpu --format=csv,noheader,nounits -``` From 7e8c75c5c24ba3baabc746aca8b3a24c3f6deeeb Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 05:35:35 +0000 Subject: [PATCH 19/29] docs: add pi05 benchmark guide pages --- docs/docs.json | 54 ++++- docs/models/pi05/eight-gpu-inference.mdx | 162 ++++++++++++++ docs/models/pi05/external-runtime-latency.mdx | 196 +++++++++++++++++ docs/models/pi05/libero-four-suites.mdx | 198 ++++++++++++++++++ docs/zh/models/pi05/eight-gpu-inference.mdx | 162 ++++++++++++++ .../models/pi05/external-runtime-latency.mdx | 196 +++++++++++++++++ docs/zh/models/pi05/libero-four-suites.mdx | 198 ++++++++++++++++++ 7 files changed, 1164 insertions(+), 2 deletions(-) create mode 100644 docs/models/pi05/eight-gpu-inference.mdx create mode 100644 docs/models/pi05/external-runtime-latency.mdx create mode 100644 docs/models/pi05/libero-four-suites.mdx create mode 100644 docs/zh/models/pi05/eight-gpu-inference.mdx create mode 100644 docs/zh/models/pi05/external-runtime-latency.mdx create mode 100644 docs/zh/models/pi05/libero-four-suites.mdx diff --git a/docs/docs.json b/docs/docs.json index 91cc5de..4fdce15 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -25,11 +25,20 @@ { "group": "Models", "pages": [ + { + "group": "PI0", + "pages": [ + "models/pi0/ws1" + ] + }, { "group": "PI0.5", "pages": [ "models/pi05/ws1", - "models/pi05/processors" + "models/pi05/processors", + "models/pi05/libero-four-suites", + "models/pi05/eight-gpu-inference", + "models/pi05/external-runtime-latency" ] }, { @@ -37,16 +46,32 @@ "pages": [ "models/cosmos/ws1", "models/cosmos/ws1_policy", + "models/cosmos/wn", + "models/cosmos/wn_policy", "models/cosmos/processors" ] } ] }, + { + "group": "Quantization", + "pages": [ + "quantization/overview", + "quantization/configuration", + "quantization/internals" + ] + }, { "group": "Tools", "pages": [ "tools/tensor-dump" ] + }, + { + "group": "Developer guide", + "pages": [ + "developer_guide/code-of-conduct" + ] } ] }, @@ -64,11 +89,20 @@ { "group": "模型", "pages": [ + { + "group": "PI0", + "pages": [ + "zh/models/pi0/ws1" + ] + }, { "group": "PI0.5", "pages": [ "zh/models/pi05/ws1", - "zh/models/pi05/processors" + "zh/models/pi05/processors", + "zh/models/pi05/libero-four-suites", + "zh/models/pi05/eight-gpu-inference", + "zh/models/pi05/external-runtime-latency" ] }, { @@ -76,16 +110,32 @@ "pages": [ "zh/models/cosmos/ws1", "zh/models/cosmos/ws1_policy", + "zh/models/cosmos/wn", + "zh/models/cosmos/wn_policy", "zh/models/cosmos/processors" ] } ] }, + { + "group": "量化", + "pages": [ + "zh/quantization/overview", + "zh/quantization/configuration", + "zh/quantization/internals" + ] + }, { "group": "工具", "pages": [ "zh/tools/tensor-dump" ] + }, + { + "group": "开发者指南", + "pages": [ + "zh/developer_guide/code-of-conduct" + ] } ] } diff --git a/docs/models/pi05/eight-gpu-inference.mdx b/docs/models/pi05/eight-gpu-inference.mdx new file mode 100644 index 0000000..e144b6c --- /dev/null +++ b/docs/models/pi05/eight-gpu-inference.mdx @@ -0,0 +1,162 @@ +--- +title: "PI0.5 Eight-GPU Inference" +description: "Run the PhyAI pi0.5 WN policy with 8-GPU data parallel inference and concurrent LIBERO clients" +icon: "server" +--- + +# Overview + +This page explains how to reproduce 8-GPU data-parallel inference for the PhyAI `pi05_wn` policy while running multiple LIBERO clients at the same time. + +The current demo uses `chunk_size=1`. Each request returns one action, so the client sends a request at every environment step. + +# Reference Environment + +| Item | Configuration | +| --- | --- | +| OS | Ubuntu 22.04 LTS | +| GPU | 8 x NVIDIA H20, about 96 GB per GPU | +| CPU / memory | 192 cores / about 1.8 TiB | +| Containers | Docker, with both server and LIBERO running in containers | +| Session | `tmux` keeps the server alive | + +Other machines can use the same flow. You need 8 idle CUDA GPUs, Docker, `tmux`, the PI0.5 LIBERO checkpoint, and an offline tokenizer. Different GPU models mainly change latency. + +# Paths + +Set these paths for your machine: + +```bash +export WORKSPACE="$HOME/vla_workplace" +export PHYAI_ROOT="$WORKSPACE/phyai" +export VLA_ROOT="$WORKSPACE/vla-evaluation-harness" +export MODEL_ROOT="$WORKSPACE/phyai_models" +``` + +Check the required files before starting: + +```bash +ls "$MODEL_ROOT/pi05_libero_phyai_converted" +ls "$MODEL_ROOT/paligemma-3b-pt-224" +ls "$PHYAI_ROOT/.venv/bin/torchrun" +ls "$VLA_ROOT/.venv/bin/vla-eval" + +nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader,nounits +ss -ltnp | grep 8000 || echo "8000 free" +``` + +Keep `paligemma-3b-pt-224` local. The scripts run in offline mode, so tokenizer download should not happen during the run. + +# Execution Flow + +```text +32 LIBERO simulation processes + -> connect to the rank0 WebSocket server at ws://127.0.0.1:8000 + -> rank0 collects requests until MAX_BATCH_SIZE is reached or MAX_WAIT_TIME expires + -> 8-GPU DP: scatter -> each GPU runs its batch slice -> gather + -> server returns actions + -> each simulation writes one JSON result +``` + +Demo parameters: + +| Item | Value | +| --- | --- | +| Number of GPUs | 8 | +| `max_batch_size` | 32, B=4 per GPU | +| `chunk_size` | 1 | +| `max_wait_time` | 0.02 | +| dtype | bfloat16 | + +CUDA graph uses a fixed shape. Requests below the full batch size are padded, so a partial batch is not proportionally faster. + +# Start the 8-GPU Server + +Start the server inside a Docker container. The container must mount `PHYAI_ROOT`, `VLA_ROOT`, and `MODEL_ROOT`. The core command is: + +```bash +export PHYAI_TOKENIZER_PATH=/data/share/paligemma-3b-pt-224 +export HF_HUB_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 + +torchrun --nproc_per_node=8 --master_port=29619 \ + -m vla_eval.model_servers.phyai_wn \ + --checkpoint_path /data/share/pi05_libero_phyai_converted \ + --host 0.0.0.0 \ + --port 8000 \ + --chunk_size 1 \ + --max_batch_size 32 \ + --max_wait_time 0.02 \ + --use_cuda_graph \ + --params_dtype bfloat16 \ + --attn_backend flashinfer \ + --norm_backend phyai-kernel \ + --linear_backend flashinfer +``` + +The first start usually takes 1 to 3 minutes for CUDA graph capture. Check readiness with: + +```bash +curl -sS http://127.0.0.1:8000/config +``` + +The response should include `max_batch_size=32`. If the server is running inside `tmux`, inspect logs with: + +```bash +tmux attach -t +# detach: Ctrl-b, then d +``` + +# Start LIBERO Clients + +Start 32 LIBERO clients: 4 suites, 8 shards per suite, all connected to the same server. + +```bash +for suite in spatial object goal libero_10; do + for shard_id in 0 1 2 3 4 5 6 7; do + env NO_PROXY="*" no_proxy="*" \ + "$VLA_ROOT/.venv/bin/vla-eval" run \ + --config "" \ + --server-url ws://127.0.0.1:8000 \ + --output-dir "/$suite" \ + --dev --yes \ + --shard-id "$shard_id" \ + --num-shards 8 & + done +done +wait +``` + +Replace `` with the matching LIBERO config file. Each result directory should contain JSON output after the run. + +# Change Parameters + +| Goal | Change | +| --- | --- | +| Change chunk size | Set server `--chunk_size N` | +| Run batch16 | Set `--max_batch_size 16` and start only 16 clients | +| Change port | Update server `--port`, client `--server-url`, and result directory together | + +Run batch16 and batch32 with separate server launches. Do not switch batch size inside one `torchrun` process. + +Reference pure-model `Engine.step` latency: batch16 is about 60 ms, and batch32 is about 97 ms. This measures model inference only, not LIBERO end-to-end wait time. + +# Cleanup + +After evaluation, stop the server and confirm that GPU memory is released: + +```bash +tmux kill-session -t +docker rm -f 2>/dev/null || true +nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits +``` + +# Troubleshooting + +| Symptom | What to check | +| --- | --- | +| Setup hangs | Offline tokenizer, `master_port`, same-name containers, and GPU memory usage | +| `:8000/config` fails | Wait for graph capture, or inspect the server log | +| Shards cannot connect | Host networking, `ws://127.0.0.1:8000`, and `NO_PROXY="*"` | +| Partial results | Check the logs under the corresponding result directory | +| FlashInfer / CUDA conflicts | Prefer Docker and avoid mixing host CUDA libraries | diff --git a/docs/models/pi05/external-runtime-latency.mdx b/docs/models/pi05/external-runtime-latency.mdx new file mode 100644 index 0000000..bba75f5 --- /dev/null +++ b/docs/models/pi05/external-runtime-latency.mdx @@ -0,0 +1,196 @@ +--- +title: "PI0.5 External Runtime Latency" +description: "Run PI0.5 latency benchmarks for FlashRT, realtime-vla, and vla.cpp" +icon: "timer" +--- + +# Overview + +This page explains how to set up and run the three external PI0.5 latency wrappers under `benchmark/pi05/`. + +These wrappers do not use the PhyAI engine for inference. Each wrapper calls the target runtime directly, while reusing PhyAI's common benchmark runner for warmup, timing, and JSONL output. + +| Script | Runtime | Timed call | +| --- | --- | --- | +| `bench_flashrt_pi05.py` | FlashRT | `Pi05TorchFrontendRtx.infer(obs)` | +| `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward(...)` | +| `bench_vlacpp_pi05_client.py` | vla.cpp | One ZMQ request to a running `vla-server` | + +The wrappers generate synthetic image and state inputs. Use them for latency-only measurements, not LIBERO accuracy evaluation. + +# Common Setup + +Use a Python environment that can import PhyAI, PyTorch, and the common benchmark runner. + +```bash +cd +python -c "import torch; import phyai; import benchmark.bench_n_batch" +nvidia-smi +``` + +Keep these settings aligned when comparing runtimes: + +| Item | Setting | +| --- | --- | +| Batch size | 1 | +| Camera streams | 2 | +| Chunk size | 50 | +| Prompt | Same prompt text across runs | +| Warmup / timed iterations | Same values across runs | +| Precision | Label each row by the runtime's real precision path | + +Use your own paths for these placeholders: + +| Placeholder | Meaning | +| --- | --- | +| `` | PhyAI checkout containing `benchmark/pi05/` | +| `` | FlashRT checkout | +| `` | realtime-vla checkout | +| `` | vla.cpp checkout | +| `` | Compiled vla.cpp `vla-server` binary | +| `` | PI0.5 safetensors checkpoint directory or file | +| `` | PI0.5 GGUF file for vla.cpp | +| `` | vla.cpp multimodal projector GGUF | +| `` | Local tokenizer directory | +| `` | LIBERO `meta/stats.json` with `observation.state.q01/q99` | + +# FlashRT + +Install FlashRT from its official repository, then make the checkout visible to the wrapper. + +```bash +export FLASHRT_ROOT= +cd +python -c "import sys; sys.path.insert(0, '$FLASHRT_ROOT'); from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx; print(Pi05TorchFrontendRtx)" +``` + +Run latency: + +```bash +cd +python benchmark/pi05/bench_flashrt_pi05.py \ + --flashrt-root \ + --checkpoint \ + --precision bf16 \ + --num-views 2 \ + --chunk-size 50 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/flashrt_pi05.jsonl +``` + +Notes: + +| Item | Note | +| --- | --- | +| Chunk size | The wrapper uses FlashRT's direct `Pi05TorchFrontendRtx` API because `chunk_size` is a frontend constructor argument | +| Denoise steps | Do not use `load_model(..., num_steps=50)` to set action chunk size; in FlashRT, `num_steps` means denoise steps | +| Precision | `--precision bf16` uses FlashRT's forced-BF16 PI0.5 RTX path; `--precision fp8_bf16` is a separate optimized FP8/BF16 result | + +# realtime-vla + +Install realtime-vla from its official repository. The wrapper can load a converted `.pt` / `.pth` checkpoint directly. If you pass a PI0.5 safetensors checkpoint, also provide FlashRT so the wrapper can reuse FlashRT's conversion helper. + +```bash +export REALTIME_VLA_ROOT= +export FLASHRT_ROOT= +cd +python -c "import sys; sys.path.insert(0, '$REALTIME_VLA_ROOT'); from pi05_infer import Pi05Inference; print(Pi05Inference)" +``` + +Run latency: + +```bash +cd +python benchmark/pi05/bench_realtime_vla_pi05.py \ + --realtime-vla-root \ + --flashrt-root \ + --checkpoint \ + --num-views 2 \ + --chunk-size 50 \ + --prompt-len 16 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/realtime_vla_pi05.jsonl +``` + +Notes: + +| Item | Note | +| --- | --- | +| Precision | The wrapper uses BF16 synthetic inputs | +| Pickle checkpoints | For `.pkl` / `.pickle`, add `--trust-pickle-checkpoint` only for trusted files | +| Prompt embedding | If the checkpoint does not contain `language_embeds`, the wrapper creates a synthetic prompt embedding for latency-only runs | + +# vla.cpp + +vla.cpp uses a server/client flow. Build `vla-server` with CUDA enabled, then start the server in one shell and run the Python client in another. + +Basic checks: + +```bash +test -x +test -f +test -f +test -f /tokenizer.json +test -f +``` + +Start server: + +```bash + \ + --bind tcp://127.0.0.1:5555 \ + --timing-detail phase \ + \ + +``` + +Run latency client: + +```bash +cd +python benchmark/pi05/bench_vlacpp_pi05_client.py \ + --vlacpp-root \ + --addr tcp://127.0.0.1:5555 \ + --arch pi05 \ + --tokenizer \ + --stats-json \ + --num-views 2 \ + --chunk-size 50 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/vlacpp_pi05.jsonl +``` + +Notes: + +| Item | Note | +| --- | --- | +| Model format | vla.cpp needs GGUF files; a safetensors checkpoint is not enough | +| Local files | Prefer local tokenizer and stats files to avoid network or HuggingFace auth issues | +| Stats format | For PI0.5, vla.cpp expects lerobot-style `meta/stats.json`; OpenPI-style `norm_stats.json` is not the same format | +| Timing detail | If the server returns phase timing, the wrapper writes it under `extras.server_phase_latency_ms` | + +# Timing Scope + +| Runtime | Timing scope | +| --- | --- | +| FlashRT | Wall time around steady-state `Pi05TorchFrontendRtx.infer(obs)`, after prompt setup, calibration, and the first graph-building call | +| realtime-vla | CUDA-event time around one `Pi05Inference.forward(...)` call | +| vla.cpp | Client wall time for one ZMQ request; server phase timing is copied from the response when available | + +# Troubleshooting + +| Symptom | Check | +| --- | --- | +| `No module named flash_rt` | Pass `--flashrt-root` or set `FLASHRT_ROOT` | +| `No module named pi05_infer` | Pass `--realtime-vla-root` or set `REALTIME_VLA_ROOT` | +| realtime-vla safetensors conversion fails | Also pass `--flashrt-root`; verify the FlashRT import works | +| vla.cpp client cannot connect | Confirm the server is ready and client `--addr` matches server `--bind` | +| vla.cpp tokenizer downloads or asks for auth | Use a local tokenizer directory | +| GPU architecture build error | Check CUDA, PyTorch CUDA, driver, and build flags for the target GPU | +| Latency is much slower than expected | Check `nvidia-smi`, rerun after warmup/JIT, and make sure no other process is using the GPU | diff --git a/docs/models/pi05/libero-four-suites.mdx b/docs/models/pi05/libero-four-suites.mdx new file mode 100644 index 0000000..9b9577e --- /dev/null +++ b/docs/models/pi05/libero-four-suites.mdx @@ -0,0 +1,198 @@ +--- +title: "PI0.5 LIBERO Four-Suite Reproduction" +description: "Run the PhyAI pi0.5 policy on LIBERO spatial, object, goal, and libero_10" +icon: "list-checks" +--- + +# Overview + +This page explains how to run the PhyAI PI0.5 LIBERO policy with `vla-evaluation-harness` on the four LIBERO benchmark suites. + +The four suites are: + +| Suite | Config file | +| --- | --- | +| `libero_spatial` | `configs/benchmarks/libero/spatial.yaml` | +| `libero_object` | `configs/benchmarks/libero/object.yaml` | +| `libero_goal` | `configs/benchmarks/libero/goal.yaml` | +| `libero_10` | `configs/benchmarks/libero/10.yaml` | + +Default evaluation settings: + +| Item | Value | +| --- | --- | +| Mode | sync | +| Action chunk size | 10 | +| Episodes per suite | 10 tasks x 50 episodes | +| Total episodes | 4 suites x 500 episodes | +| Model | PI0.5 LIBERO converted checkpoint | +| Simulator | `vla-evaluation-harness` LIBERO Docker container | +| Output | One JSON result file per suite | + +# Resources + +Prepare these resources first: + +| Resource | Notes | +| --- | --- | +| PhyAI repository | Contains the PI0.5 policy and benchmark scripts | +| `vla-evaluation-harness` | LIBERO simulation benchmark framework | +| PI0.5 LIBERO checkpoint | Converted to the format PhyAI can load | +| `paligemma-3b-pt-224` | Tokenizer / processor used by PI0.5 | +| Docker + NVIDIA Container Toolkit | Required for the LIBERO simulation container | + +`paligemma-3b-pt-224` may require HuggingFace access. If possible, sync it from a machine that already has access instead of downloading it during the run. + +# Set Paths + +Set paths on the target machine. These are placeholders; change them for your environment. + +```bash +export PHYAI_ROOT=$HOME/phyai +export VLA_ROOT=$HOME/vla-evaluation-harness +export MODEL_ROOT=$HOME/phyai_models +export PHYAI_CONTAINER=phyai_libero_eval + +export PHYAI_CKPT_HOST=$MODEL_ROOT/pi05_libero_phyai_converted +export TOKENIZER_HOST=$MODEL_ROOT/paligemma-3b-pt-224 + +export PHYAI_CKPT_IN_CONTAINER=/data/share/pi05_libero_phyai_converted +export TOKENIZER_IN_CONTAINER=/data/share/paligemma-3b-pt-224 + +export LIBERO_IMAGE=ghcr.io/allenai/vla-evaluation-harness/libero:latest +``` + +Check the model directories: + +```bash +test -d "$PHYAI_CKPT_HOST" +test -d "$TOKENIZER_HOST" +``` + +# Python Environment + +Create and sync the environment from the PhyAI repository: + +```bash +cd "$PHYAI_ROOT" +uv sync --group cu130 --extra libero +source .venv/bin/activate +``` + +Run a minimal import check: + +```bash +python -c "import torch; import phyai; print(torch.__version__, torch.cuda.get_device_name(0))" +python -c "from phyai.policies.pi05_libero import PI05LiberoPolicy; print(PI05LiberoPolicy)" +``` + +# LIBERO Container + +Pull or build the `vla-evaluation-harness` LIBERO image. After the image is ready, start the container: + +```bash +docker run --gpus all -it --rm \ + --name "$PHYAI_CONTAINER" \ + --ipc=host \ + --network=host \ + -v "$PHYAI_ROOT:$PHYAI_ROOT" \ + -v "$MODEL_ROOT:/data/share" \ + "$LIBERO_IMAGE" \ + bash +``` + +Inside the container, check that the code and models are visible: + +```bash +test -d "$PHYAI_ROOT" +test -d "$PHYAI_CKPT_IN_CONTAINER" +test -d "$TOKENIZER_IN_CONTAINER" +``` + +# Start the Policy Server + +Start the PI0.5 LIBERO policy server in one terminal: + +```bash +cd "$PHYAI_ROOT" +source .venv/bin/activate + +python -m phyai.policies.pi05_libero \ + --checkpoint-dir "$PHYAI_CKPT_IN_CONTAINER" \ + --tokenizer-path "$TOKENIZER_IN_CONTAINER" \ + --host 127.0.0.1 \ + --port 8000 \ + --chunk-size 10 \ + --num-images 2 \ + --dtype bf16 +``` + +Keep the server running. The first request usually triggers initialization and graph capture, so do not include it in benchmark latency. + +# Smoke Test + +In another terminal, run a short check to confirm that the evaluator can reach the policy server: + +```bash +cd "$VLA_ROOT" + +python -m experiments.robot.libero.run_libero_eval \ + --benchmark libero_spatial \ + --policy-host 127.0.0.1 \ + --policy-port 8000 \ + --num-trials-per-task 1 \ + --max-steps 20 \ + --result-dir results/pi05_libero_smoke +``` + +If the smoke test fails, stop there and check the server log, port, checkpoint path, and tokenizer path first. + +# Run All Four Suites + +After confirming that the GPU is idle, run the full evaluation. Save each suite separately. + +```bash +cd "$VLA_ROOT" + +for suite in libero_spatial libero_object libero_goal libero_10; do + python -m experiments.robot.libero.run_libero_eval \ + --benchmark "$suite" \ + --policy-host 127.0.0.1 \ + --policy-port 8000 \ + --num-trials-per-task 50 \ + --result-dir "results/pi05_${suite}" +done +``` + +Record these fields with the result: + +| Item | What to record | +| --- | --- | +| PhyAI commit | `git rev-parse HEAD` | +| Checkpoint | Name and source | +| Tokenizer | Local path or HuggingFace id | +| GPU | `nvidia-smi --query-gpu=name --format=csv,noheader` | +| dtype | For example, `bf16` | +| Chunk size | For example, `10` | +| Number of camera views | For example, `2` | +| Episodes per suite | For example, `500` | + +# Summarize Results + +Each benchmark suite writes a JSON file under its result directory. This helper prints common summary fields when they exist: + +```bash +python - < 连接 rank0 WebSocket，地址 ws://127.0.0.1:8000 + -> rank0 收集请求，凑满 MAX_BATCH_SIZE 或等待 MAX_WAIT_TIME + -> 8 卡 DP：scatter -> 每卡计算自己的 batch slice -> gather + -> server 返回 action + -> 每路仿真写一份 JSON 结果 +``` + +本次 demo 参数： + +| 项 | 值 | +| --- | --- | +| GPU 数 | 8 | +| `max_batch_size` | 32，每卡 B=4 | +| `chunk_size` | 1 | +| `max_wait_time` | 0.02 | +| dtype | bfloat16 | + +CUDA graph 按固定 shape 运行。不足 batch 的请求会 padding，所以 partial batch 不会按比例变快。 + +# 启动八卡 server + +server 在 Docker 容器里启动，容器需要挂载 `PHYAI_ROOT`、`VLA_ROOT` 和 `MODEL_ROOT`。核心启动命令如下： + +```bash +export PHYAI_TOKENIZER_PATH=/data/share/paligemma-3b-pt-224 +export HF_HUB_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 + +torchrun --nproc_per_node=8 --master_port=29619 \ + -m vla_eval.model_servers.phyai_wn \ + --checkpoint_path /data/share/pi05_libero_phyai_converted \ + --host 0.0.0.0 \ + --port 8000 \ + --chunk_size 1 \ + --max_batch_size 32 \ + --max_wait_time 0.02 \ + --use_cuda_graph \ + --params_dtype bfloat16 \ + --attn_backend flashinfer \ + --norm_backend phyai-kernel \ + --linear_backend flashinfer +``` + +第一次启动通常需要 1 到 3 分钟完成 CUDA graph capture。就绪后检查： + +```bash +curl -sS http://127.0.0.1:8000/config +``` + +返回内容里应能看到 `max_batch_size=32`。如果 server 放在 `tmux` 里，可以这样查看日志： + +```bash +tmux attach -t +# 分离：Ctrl-b，然后按 d +``` + +# 启动 LIBERO client + +启动 32 路 LIBERO client：4 个 suite，每个 suite 8 个 shard，全部连接同一个 server。 + +```bash +for suite in spatial object goal libero_10; do + for shard_id in 0 1 2 3 4 5 6 7; do + env NO_PROXY="*" no_proxy="*" \ + "$VLA_ROOT/.venv/bin/vla-eval" run \ + --config "" \ + --server-url ws://127.0.0.1:8000 \ + --output-dir "/$suite" \ + --dev --yes \ + --shard-id "$shard_id" \ + --num-shards 8 & + done +done +wait +``` + +`` 需要换成对应的 LIBERO 配置文件。跑完后，每个结果目录下会生成 JSON。 + +# 改参数 + +| 目标 | 改法 | +| --- | --- | +| 改 chunk | 修改 server 的 `--chunk_size N` | +| batch16 | 改成 `--max_batch_size 16`，并只启动 16 路 client | +| 换端口 | 同时修改 `--port`、client 的 `--server-url` 和结果目录 | + +batch16 和 batch32 建议分两次启动 server，不要在同一个 `torchrun` 进程里切换。 + +纯模型 `Engine.step` 的参考 latency：batch16 约 60 ms，batch32 约 97 ms。这个数只代表模型推理，不等于 LIBERO 端到端等待时间。 + +# 收尾 + +评测结束后关闭 server，并确认 GPU 已释放： + +```bash +tmux kill-session -t +docker rm -f 2>/dev/null || true +nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits +``` + +# 常见问题 + +| 现象 | 处理方式 | +| --- | --- | +| setup 卡住 | 检查离线 tokenizer、`master_port`、同名容器和显存占用 | +| `:8000/config` 不通 | 等待 graph capture 完成，或查看 server 日志 | +| shard 连不上 | 确认 host network、`ws://127.0.0.1:8000` 和 `NO_PROXY="*"` | +| 只有部分结果 | 查看对应结果目录下的日志文件 | +| FlashInfer / CUDA 冲突 | 优先使用 Docker 环境，避免混用宿主机 CUDA 库 | diff --git a/docs/zh/models/pi05/external-runtime-latency.mdx b/docs/zh/models/pi05/external-runtime-latency.mdx new file mode 100644 index 0000000..f1b12df --- /dev/null +++ b/docs/zh/models/pi05/external-runtime-latency.mdx @@ -0,0 +1,196 @@ +--- +title: "PI0.5 外部 Runtime 延迟测试" +description: "使用 FlashRT、realtime-vla 和 vla.cpp 跑 PI0.5 latency benchmark" +icon: "timer" +--- + +# 概述 + +这篇文档说明如何配置并运行 `benchmark/pi05/` 下的三个外部 PI0.5 latency wrapper。 + +这些 wrapper 不使用 PhyAI engine 做推理。每个脚本会直接调用对应 runtime，同时复用 PhyAI 的通用 benchmark runner 来做 warmup、计时和 JSONL 输出。 + +| 脚本 | Runtime | 计时调用 | +| --- | --- | --- | +| `bench_flashrt_pi05.py` | FlashRT | `Pi05TorchFrontendRtx.infer(obs)` | +| `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward(...)` | +| `bench_vlacpp_pi05_client.py` | vla.cpp | 向运行中的 `vla-server` 发一次 ZMQ 请求 | + +这些脚本会生成 synthetic image 和 state 输入，只用于 latency 测试，不用于 LIBERO accuracy 评测。 + +# 通用环境 + +需要一个能 import PhyAI、PyTorch 和通用 benchmark runner 的 Python 环境。 + +```bash +cd +python -c "import torch; import phyai; import benchmark.bench_n_batch" +nvidia-smi +``` + +对比不同 runtime 时，下面这些设置要对齐： + +| 项 | 设置 | +| --- | --- | +| Batch size | 1 | +| Camera streams | 2 | +| Chunk size | 50 | +| Prompt | 各 runtime 使用相同 prompt 文本 | +| Warmup / timed iterations | 各 runtime 使用相同次数 | +| Precision | 按 runtime 的实际精度路径标注结果 | + +下面的命令都使用占位路径： + +| 占位符 | 含义 | +| --- | --- | +| `` | 包含 `benchmark/pi05/` 的 PhyAI 仓库 | +| `` | FlashRT 仓库 | +| `` | realtime-vla 仓库 | +| `` | vla.cpp 仓库 | +| `` | 编译后的 vla.cpp `vla-server` 二进制文件 | +| `` | PI0.5 safetensors checkpoint 目录或文件 | +| `` | vla.cpp 使用的 PI0.5 GGUF 文件 | +| `` | vla.cpp multimodal projector GGUF 文件 | +| `` | 本地 tokenizer 目录 | +| `` | LIBERO `meta/stats.json`，包含 `observation.state.q01/q99` | + +# FlashRT + +先按 FlashRT 官方仓库说明安装，然后把仓库路径暴露给 wrapper。 + +```bash +export FLASHRT_ROOT= +cd +python -c "import sys; sys.path.insert(0, '$FLASHRT_ROOT'); from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx; print(Pi05TorchFrontendRtx)" +``` + +运行 latency： + +```bash +cd +python benchmark/pi05/bench_flashrt_pi05.py \ + --flashrt-root \ + --checkpoint \ + --precision bf16 \ + --num-views 2 \ + --chunk-size 50 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/flashrt_pi05.jsonl +``` + +说明： + +| 项 | 说明 | +| --- | --- | +| Chunk size | wrapper 直接使用 FlashRT 的 `Pi05TorchFrontendRtx` API，因为 `chunk_size` 是 frontend 构造参数 | +| Denoise steps | 不要用 `load_model(..., num_steps=50)` 设置 action chunk size；FlashRT 里的 `num_steps` 表示 denoise steps | +| Precision | `--precision bf16` 使用 FlashRT 的 forced-BF16 PI0.5 RTX 路径；`--precision fp8_bf16` 是单独的 FP8/BF16 优化结果 | + +# realtime-vla + +先按 realtime-vla 官方仓库说明安装。wrapper 可以直接加载转换后的 `.pt` / `.pth` checkpoint。如果传入 PI0.5 safetensors checkpoint，还需要提供 FlashRT 路径，用它的 PI0.5 转换 helper。 + +```bash +export REALTIME_VLA_ROOT= +export FLASHRT_ROOT= +cd +python -c "import sys; sys.path.insert(0, '$REALTIME_VLA_ROOT'); from pi05_infer import Pi05Inference; print(Pi05Inference)" +``` + +运行 latency： + +```bash +cd +python benchmark/pi05/bench_realtime_vla_pi05.py \ + --realtime-vla-root \ + --flashrt-root \ + --checkpoint \ + --num-views 2 \ + --chunk-size 50 \ + --prompt-len 16 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/realtime_vla_pi05.jsonl +``` + +说明： + +| 项 | 说明 | +| --- | --- | +| Precision | wrapper 使用 BF16 synthetic inputs | +| Pickle checkpoint | `.pkl` / `.pickle` 文件只有在可信时才加 `--trust-pickle-checkpoint` | +| Prompt embedding | 如果 checkpoint 没有 `language_embeds`，wrapper 会为 latency-only 测试创建 synthetic prompt embedding | + +# vla.cpp + +vla.cpp 是 server/client 流程。先用 CUDA 编译 `vla-server`，然后在一个 shell 启动 server，在另一个 shell 运行 Python client。 + +基础检查： + +```bash +test -x +test -f +test -f +test -f /tokenizer.json +test -f +``` + +启动 server： + +```bash + \ + --bind tcp://127.0.0.1:5555 \ + --timing-detail phase \ + \ + +``` + +运行 latency client： + +```bash +cd +python benchmark/pi05/bench_vlacpp_pi05_client.py \ + --vlacpp-root \ + --addr tcp://127.0.0.1:5555 \ + --arch pi05 \ + --tokenizer \ + --stats-json \ + --num-views 2 \ + --chunk-size 50 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --result-file results/vlacpp_pi05.jsonl +``` + +说明： + +| 项 | 说明 | +| --- | --- | +| 模型格式 | vla.cpp 需要 GGUF 文件；safetensors checkpoint 不够 | +| 本地文件 | tokenizer 和 stats 建议都用本地文件，避免测试时触发网络或 HuggingFace 权限问题 | +| Stats 格式 | PI0.5 下 vla.cpp 需要 lerobot 风格的 `meta/stats.json`；OpenPI 风格的 `norm_stats.json` 不是同一种格式 | +| 详细计时 | 如果 server 返回 phase timing，wrapper 会写入 `extras.server_phase_latency_ms` | + +# 计时口径 + +| Runtime | 计时范围 | +| --- | --- | +| FlashRT | 对 steady-state `Pi05TorchFrontendRtx.infer(obs)` 计 wall time；不包含 prompt setup、calibration 和第一次 graph-building 调用 | +| realtime-vla | 用 CUDA event 计一次 `Pi05Inference.forward(...)` | +| vla.cpp | 计一次 ZMQ 请求的 client wall time；如果 response 里有 server phase timing，会一并记录 | + +# 常见问题 + +| 现象 | 检查方式 | +| --- | --- | +| `No module named flash_rt` | 传 `--flashrt-root` 或设置 `FLASHRT_ROOT` | +| `No module named pi05_infer` | 传 `--realtime-vla-root` 或设置 `REALTIME_VLA_ROOT` | +| realtime-vla safetensors 转换失败 | 同时传 `--flashrt-root`，并确认 FlashRT import 正常 | +| vla.cpp client 连不上 | 确认 server 已 ready，且 client `--addr` 和 server `--bind` 一致 | +| vla.cpp tokenizer 触发下载或鉴权 | 使用本地 tokenizer 目录 | +| GPU 架构编译报错 | 检查 CUDA、PyTorch CUDA、driver 和目标 GPU 的 build flags | +| latency 明显偏慢 | 看 `nvidia-smi`，排除 GPU 占用；JIT/warmup 后重新跑 | diff --git a/docs/zh/models/pi05/libero-four-suites.mdx b/docs/zh/models/pi05/libero-four-suites.mdx new file mode 100644 index 0000000..604c612 --- /dev/null +++ b/docs/zh/models/pi05/libero-four-suites.mdx @@ -0,0 +1,198 @@ +--- +title: "PI0.5 LIBERO 四套任务复现" +description: "使用 PhyAI pi0.5 policy 跑 LIBERO spatial、object、goal 和 libero_10 四套评测" +icon: "list-checks" +--- + +# 概述 + +这篇文档说明如何用 PhyAI 的 PI0.5 LIBERO policy 跑完 `vla-evaluation-harness` 里的四套 LIBERO benchmark。 + +四套任务对应关系如下： + +| Suite | 配置文件 | +| --- | --- | +| `libero_spatial` | `configs/benchmarks/libero/spatial.yaml` | +| `libero_object` | `configs/benchmarks/libero/object.yaml` | +| `libero_goal` | `configs/benchmarks/libero/goal.yaml` | +| `libero_10` | `configs/benchmarks/libero/10.yaml` | + +默认评测设置： + +| 项目 | 设置 | +| --- | --- | +| 运行模式 | sync | +| Action chunk size | 10 | +| 每套任务 episode 数 | 10 tasks x 50 episodes | +| 总 episode 数 | 4 suites x 500 episodes | +| 模型 | PI0.5 LIBERO converted checkpoint | +| 仿真环境 | `vla-evaluation-harness` LIBERO Docker container | +| 输出 | 每套任务一个 JSON 结果文件 | + +# 准备资源 + +需要先准备这些资源： + +| 资源 | 说明 | +| --- | --- | +| PhyAI 仓库 | 包含 PI0.5 policy 和 benchmark 脚本 | +| `vla-evaluation-harness` | LIBERO 仿真评测框架 | +| PI0.5 LIBERO checkpoint | 已转换成 PhyAI 可加载格式 | +| `paligemma-3b-pt-224` | PI0.5 使用的 tokenizer / processor | +| Docker + NVIDIA Container Toolkit | 用来启动 LIBERO 仿真容器 | + +`paligemma-3b-pt-224` 可能需要 HuggingFace 权限。更稳的做法是从已有权限的机器同步到本机。 + +# 设置路径 + +先在目标机器上设置路径。下面只用占位路径，按实际机器改即可。 + +```bash +export PHYAI_ROOT=$HOME/phyai +export VLA_ROOT=$HOME/vla-evaluation-harness +export MODEL_ROOT=$HOME/phyai_models +export PHYAI_CONTAINER=phyai_libero_eval + +export PHYAI_CKPT_HOST=$MODEL_ROOT/pi05_libero_phyai_converted +export TOKENIZER_HOST=$MODEL_ROOT/paligemma-3b-pt-224 + +export PHYAI_CKPT_IN_CONTAINER=/data/share/pi05_libero_phyai_converted +export TOKENIZER_IN_CONTAINER=/data/share/paligemma-3b-pt-224 + +export LIBERO_IMAGE=ghcr.io/allenai/vla-evaluation-harness/libero:latest +``` + +检查模型目录： + +```bash +test -d "$PHYAI_CKPT_HOST" +test -d "$TOKENIZER_HOST" +``` + +# 准备 Python 环境 + +在 PhyAI 仓库里创建并同步环境： + +```bash +cd "$PHYAI_ROOT" +uv sync --group cu130 --extra libero +source .venv/bin/activate +``` + +做一次最小导入检查： + +```bash +python -c "import torch; import phyai; print(torch.__version__, torch.cuda.get_device_name(0))" +python -c "from phyai.policies.pi05_libero import PI05LiberoPolicy; print(PI05LiberoPolicy)" +``` + +# 准备 LIBERO 容器 + +拉取或构建 `vla-evaluation-harness` 的 LIBERO 镜像。镜像准备好后，启动容器： + +```bash +docker run --gpus all -it --rm \ + --name "$PHYAI_CONTAINER" \ + --ipc=host \ + --network=host \ + -v "$PHYAI_ROOT:$PHYAI_ROOT" \ + -v "$MODEL_ROOT:/data/share" \ + "$LIBERO_IMAGE" \ + bash +``` + +进入容器后，确认能看到代码和模型： + +```bash +test -d "$PHYAI_ROOT" +test -d "$PHYAI_CKPT_IN_CONTAINER" +test -d "$TOKENIZER_IN_CONTAINER" +``` + +# 启动 PhyAI policy server + +在一个终端里启动 PI0.5 LIBERO policy server： + +```bash +cd "$PHYAI_ROOT" +source .venv/bin/activate + +python -m phyai.policies.pi05_libero \ + --checkpoint-dir "$PHYAI_CKPT_IN_CONTAINER" \ + --tokenizer-path "$TOKENIZER_IN_CONTAINER" \ + --host 127.0.0.1 \ + --port 8000 \ + --chunk-size 10 \ + --num-images 2 \ + --dtype bf16 +``` + +server 启动后保持运行。第一次请求通常会触发初始化和图捕获，不建议把这部分算进 benchmark latency。 + +# 跑 smoke test + +在另一个终端里跑一次很短的检查，确认 evaluator 能连到 policy server： + +```bash +cd "$VLA_ROOT" + +python -m experiments.robot.libero.run_libero_eval \ + --benchmark libero_spatial \ + --policy-host 127.0.0.1 \ + --policy-port 8000 \ + --num-trials-per-task 1 \ + --max-steps 20 \ + --result-dir results/pi05_libero_smoke +``` + +如果 smoke test 失败，先不要跑完整四套任务。优先检查 server 日志、端口、模型路径和 tokenizer 路径。 + +# 跑四套任务 + +确认 GPU 空闲后，再跑完整评测。建议每套任务单独保存结果。 + +```bash +cd "$VLA_ROOT" + +for suite in libero_spatial libero_object libero_goal libero_10; do + python -m experiments.robot.libero.run_libero_eval \ + --benchmark "$suite" \ + --policy-host 127.0.0.1 \ + --policy-port 8000 \ + --num-trials-per-task 50 \ + --result-dir "results/pi05_${suite}" +done +``` + +运行时记录这些信息，后续对齐结果会用到： + +| 项目 | 需要记录 | +| --- | --- | +| PhyAI commit | `git rev-parse HEAD` | +| checkpoint | checkpoint 名称和来源 | +| tokenizer | 本地目录或 HF id | +| GPU | `nvidia-smi --query-gpu=name --format=csv,noheader` | +| dtype | 例如 `bf16` | +| chunk size | 例如 `10` | +| camera view 数 | 例如 `2` | +| 每套任务 episode 数 | 例如 `500` | + +# 汇总结果 + +每套 benchmark 结束后，结果目录里会有对应 JSON 文件。可用下面的方式快速看成功率和平均耗时字段： + +```bash +python - < Date: Sun, 2 Aug 2026 08:49:45 +0000 Subject: [PATCH 20/29] docs: fix pi05 libero server command --- docs/models/pi05/libero-four-suites.mdx | 29 ++++++++++++++-------- docs/zh/models/pi05/libero-four-suites.mdx | 29 ++++++++++++++-------- 2 files changed, 36 insertions(+), 22 deletions(-) diff --git a/docs/models/pi05/libero-four-suites.mdx b/docs/models/pi05/libero-four-suites.mdx index 9b9577e..0d62eed 100644 --- a/docs/models/pi05/libero-four-suites.mdx +++ b/docs/models/pi05/libero-four-suites.mdx @@ -111,23 +111,30 @@ test -d "$TOKENIZER_IN_CONTAINER" # Start the Policy Server -Start the PI0.5 LIBERO policy server in one terminal: +`phyai.policies.pi05_libero` provides the PhyAI policy adapter, but it is not a CLI server. Start the WebSocket server through the `vla-evaluation-harness` server entry point: ```bash -cd "$PHYAI_ROOT" +cd "$VLA_ROOT" source .venv/bin/activate -python -m phyai.policies.pi05_libero \ - --checkpoint-dir "$PHYAI_CKPT_IN_CONTAINER" \ - --tokenizer-path "$TOKENIZER_IN_CONTAINER" \ - --host 127.0.0.1 \ - --port 8000 \ - --chunk-size 10 \ - --num-images 2 \ - --dtype bf16 +export PYTHONPATH="$PHYAI_ROOT/phyai/src:$PHYAI_ROOT/phyai-kernel:$PHYAI_ROOT/phyai-utils-tools/src:$VLA_ROOT/src" +export PHYAI_TOKENIZER_PATH="$TOKENIZER_IN_CONTAINER" +export PHYAI_CAMERA_MODE=two_camera + +python -m vla_eval.model_servers.phyai \ + --checkpoint_path "$PHYAI_CKPT_IN_CONTAINER" \ + --device cuda:0 \ + --params_dtype bfloat16 \ + --attn_backend flashinfer \ + --norm_backend phyai-kernel \ + --linear_backend flashinfer \ + --flashinfer_workspace_bytes 536870912 \ + --chunk_size 10 \ + --host 0.0.0.0 \ + --port 8000 ``` -Keep the server running. The first request usually triggers initialization and graph capture, so do not include it in benchmark latency. +Keep the server running. Wait until the log shows that the server is listening on `ws://0.0.0.0:8000`; the first request usually triggers initialization and graph capture. # Smoke Test diff --git a/docs/zh/models/pi05/libero-four-suites.mdx b/docs/zh/models/pi05/libero-four-suites.mdx index 604c612..28ddc48 100644 --- a/docs/zh/models/pi05/libero-four-suites.mdx +++ b/docs/zh/models/pi05/libero-four-suites.mdx @@ -111,23 +111,30 @@ test -d "$TOKENIZER_IN_CONTAINER" # 启动 PhyAI policy server -在一个终端里启动 PI0.5 LIBERO policy server： +`phyai.policies.pi05_libero` 只是 PhyAI policy adapter，不是 CLI server。WebSocket server 需要通过 `vla-evaluation-harness` 的 server 入口启动： ```bash -cd "$PHYAI_ROOT" +cd "$VLA_ROOT" source .venv/bin/activate -python -m phyai.policies.pi05_libero \ - --checkpoint-dir "$PHYAI_CKPT_IN_CONTAINER" \ - --tokenizer-path "$TOKENIZER_IN_CONTAINER" \ - --host 127.0.0.1 \ - --port 8000 \ - --chunk-size 10 \ - --num-images 2 \ - --dtype bf16 +export PYTHONPATH="$PHYAI_ROOT/phyai/src:$PHYAI_ROOT/phyai-kernel:$PHYAI_ROOT/phyai-utils-tools/src:$VLA_ROOT/src" +export PHYAI_TOKENIZER_PATH="$TOKENIZER_IN_CONTAINER" +export PHYAI_CAMERA_MODE=two_camera + +python -m vla_eval.model_servers.phyai \ + --checkpoint_path "$PHYAI_CKPT_IN_CONTAINER" \ + --device cuda:0 \ + --params_dtype bfloat16 \ + --attn_backend flashinfer \ + --norm_backend phyai-kernel \ + --linear_backend flashinfer \ + --flashinfer_workspace_bytes 536870912 \ + --chunk_size 10 \ + --host 0.0.0.0 \ + --port 8000 ``` -server 启动后保持运行。第一次请求通常会触发初始化和图捕获，不建议把这部分算进 benchmark latency。 +server 启动后保持运行。日志出现 `ws://0.0.0.0:8000` 开始监听后，再启动 smoke test；第一次请求通常会触发初始化和图捕获。 # 跑 smoke test From c5a904493b74b29a9126a11caaa72ebd51385169 Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 17:25:29 +0800 Subject: [PATCH 21/29] fix(phyai): allow libero action chunk override --- phyai/src/phyai/policies/pi05_libero.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/phyai/src/phyai/policies/pi05_libero.py b/phyai/src/phyai/policies/pi05_libero.py index bcaea76..3268db7 100644 --- a/phyai/src/phyai/policies/pi05_libero.py +++ b/phyai/src/phyai/policies/pi05_libero.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from dataclasses import replace from pathlib import Path from typing import Any @@ -61,16 +62,18 @@ def __init__( flashinfer_workspace_bytes: int = 512 * 1024 * 1024, tokenizer_name: str | None = None, camera_mode: str | None = None, + chunk_size: int | None = None, ) -> None: self.checkpoint_dir = Path(checkpoint_dir) self.device = device self.params_dtype = params_dtype self.max_batch_size = int(max_batch_size) self.config = self._read_config() + self._engine_config = self._resolve_engine_config(chunk_size) self.image_size = self._resolve_image_size(self.config) self._action_dim = self._resolve_action_dim(self.config) - self.max_action_dim = int(self.config.get("max_action_dim", 32)) - self._chunk_size = int(self.config.get("chunk_size", PI05Config().chunk_size)) + self.max_action_dim = int(self._engine_config.max_action_dim) + self._chunk_size = int(self._engine_config.chunk_size) self.camera_names = self._resolve_camera_names(camera_mode) self.tokenizer_name = self._resolve_tokenizer_name(tokenizer_name) self.prompt_mode = str( @@ -108,6 +111,7 @@ def __init__( plugin="pi05", plugin_args=PI05Args( checkpoint_dir=self.checkpoint_dir, + config=self._engine_config, max_batch_size=self.max_batch_size, weight_remap=_lerobot_pi05_weight_remap, inputs_image_shape=[ @@ -157,6 +161,12 @@ def _read_config(self) -> dict[str, Any]: with path.open("r", encoding="utf-8") as f: return json.load(f) + def _resolve_engine_config(self, chunk_size: int | None) -> PI05Config: + config = PI05Config.from_dict(self.config) + if chunk_size is None: + return config + return replace(config, chunk_size=int(chunk_size)) + def _resolve_camera_names(self, camera_mode: str | None) -> list[str]: mode = camera_mode or envs.PHYAI_CAMERA_MODE.get() or "three_camera" if mode == "two_camera": From 4b201615a85aea3788c5c78441b7216e2a636eec Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 18:32:25 +0800 Subject: [PATCH 22/29] docs(phyai): pin LIBERO evaluation harness revision --- docs/models/pi05/libero-four-suites.mdx | 100 +++++++++++---------- docs/zh/models/pi05/libero-four-suites.mdx | 100 +++++++++++---------- 2 files changed, 102 insertions(+), 98 deletions(-) diff --git a/docs/models/pi05/libero-four-suites.mdx b/docs/models/pi05/libero-four-suites.mdx index 0d62eed..951c21d 100644 --- a/docs/models/pi05/libero-four-suites.mdx +++ b/docs/models/pi05/libero-four-suites.mdx @@ -43,6 +43,32 @@ Prepare these resources first: `paligemma-3b-pt-224` may require HuggingFace access. If possible, sync it from a machine that already has access instead of downloading it during the run. +# Pin Source Revisions + +This reproduction uses immutable commits from two linked pull requests: + +| Repository | Pull request | Commit | +| --- | --- | --- | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`c5a9044`](https://github.com/rebecca26358/phyai/commit/c5a904493b74b29a9126a11caaa72ebd51385169) | +| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`0b2fadb`](https://github.com/rebecca26358/vla-evaluation-harness/commit/0b2fadb3f4b30ad9e2591537120171d9cccee86d) | + +Check out those exact revisions: + +```bash +export PHYAI_ROOT=${PHYAI_ROOT:-$HOME/phyai} +export VLA_ROOT=${VLA_ROOT:-$HOME/vla-evaluation-harness} + +git clone https://github.com/rebecca26358/phyai.git "$PHYAI_ROOT" +git -C "$PHYAI_ROOT" checkout c5a904493b74b29a9126a11caaa72ebd51385169 + +git clone https://github.com/rebecca26358/vla-evaluation-harness.git "$VLA_ROOT" +git -C "$VLA_ROOT" checkout 0b2fadb3f4b30ad9e2591537120171d9cccee86d +``` + +The VLA revision contains `src/vla_eval/model_servers/phyai.py` and +`configs/model_servers/phyai/libero.yaml`. The server script pins the same +PhyAI commit in its PEP 723 dependency metadata. + # Set Paths Set paths on the target machine. These are placeholders; change them for your environment. @@ -59,7 +85,7 @@ export TOKENIZER_HOST=$MODEL_ROOT/paligemma-3b-pt-224 export PHYAI_CKPT_IN_CONTAINER=/data/share/pi05_libero_phyai_converted export TOKENIZER_IN_CONTAINER=/data/share/paligemma-3b-pt-224 -export LIBERO_IMAGE=ghcr.io/allenai/vla-evaluation-harness/libero:latest +export PHYAI_IMAGE=nvcr.io/nvidia/pytorch:25.12-py3 ``` Check the model directories: @@ -69,26 +95,19 @@ test -d "$PHYAI_CKPT_HOST" test -d "$TOKENIZER_HOST" ``` -# Python Environment +# Evaluator Environment -Create and sync the environment from the PhyAI repository: +Create the VLA evaluator environment on the host. This works from a clean checkout; PhyAI does not define `cu130` or `libero` dependency groups. ```bash -cd "$PHYAI_ROOT" -uv sync --group cu130 --extra libero -source .venv/bin/activate -``` - -Run a minimal import check: - -```bash -python -c "import torch; import phyai; print(torch.__version__, torch.cuda.get_device_name(0))" -python -c "from phyai.policies.pi05_libero import PI05LiberoPolicy; print(PI05LiberoPolicy)" +cd "$VLA_ROOT" +uv sync ``` -# LIBERO Container +# PhyAI Runtime Container -Pull or build the `vla-evaluation-harness` LIBERO image. After the image is ready, start the container: +Run the policy server in the CUDA 13 image used by PhyAI. The evaluator starts +the separate LIBERO simulator image declared by each benchmark config. ```bash docker run --gpus all -it --rm \ @@ -96,45 +115,39 @@ docker run --gpus all -it --rm \ --ipc=host \ --network=host \ -v "$PHYAI_ROOT:$PHYAI_ROOT" \ + -v "$VLA_ROOT:$VLA_ROOT" \ -v "$MODEL_ROOT:/data/share" \ - "$LIBERO_IMAGE" \ + "$PHYAI_IMAGE" \ bash ``` -Inside the container, check that the code and models are visible: +Inside the container, sync both clean checkouts and verify the model paths: ```bash -test -d "$PHYAI_ROOT" +cd "$PHYAI_ROOT" +uv sync + +cd "$VLA_ROOT" +uv sync + test -d "$PHYAI_CKPT_IN_CONTAINER" test -d "$TOKENIZER_IN_CONTAINER" ``` # Start the Policy Server -`phyai.policies.pi05_libero` provides the PhyAI policy adapter, but it is not a CLI server. Start the WebSocket server through the `vla-evaluation-harness` server entry point: +`phyai.policies.pi05_libero` is a policy adapter, not a CLI server. Start the WebSocket server through the VLA config added by PR #111: ```bash cd "$VLA_ROOT" -source .venv/bin/activate -export PYTHONPATH="$PHYAI_ROOT/phyai/src:$PHYAI_ROOT/phyai-kernel:$PHYAI_ROOT/phyai-utils-tools/src:$VLA_ROOT/src" +export PHYAI_CHECKPOINT_PATH="$PHYAI_CKPT_IN_CONTAINER" export PHYAI_TOKENIZER_PATH="$TOKENIZER_IN_CONTAINER" -export PHYAI_CAMERA_MODE=two_camera - -python -m vla_eval.model_servers.phyai \ - --checkpoint_path "$PHYAI_CKPT_IN_CONTAINER" \ - --device cuda:0 \ - --params_dtype bfloat16 \ - --attn_backend flashinfer \ - --norm_backend phyai-kernel \ - --linear_backend flashinfer \ - --flashinfer_workspace_bytes 536870912 \ - --chunk_size 10 \ - --host 0.0.0.0 \ - --port 8000 + +uv run vla-eval serve --config configs/model_servers/phyai/libero.yaml ``` -Keep the server running. Wait until the log shows that the server is listening on `ws://0.0.0.0:8000`; the first request usually triggers initialization and graph capture. +Keep the server running while the evaluator connects to `ws://localhost:8000`. # Smoke Test @@ -143,13 +156,7 @@ In another terminal, run a short check to confirm that the evaluator can reach t ```bash cd "$VLA_ROOT" -python -m experiments.robot.libero.run_libero_eval \ - --benchmark libero_spatial \ - --policy-host 127.0.0.1 \ - --policy-port 8000 \ - --num-trials-per-task 1 \ - --max-steps 20 \ - --result-dir results/pi05_libero_smoke +uv run vla-eval run --config configs/benchmarks/libero/smoke_test.yaml ``` If the smoke test fails, stop there and check the server log, port, checkpoint path, and tokenizer path first. @@ -161,13 +168,8 @@ After confirming that the GPU is idle, run the full evaluation. Save each suite ```bash cd "$VLA_ROOT" -for suite in libero_spatial libero_object libero_goal libero_10; do - python -m experiments.robot.libero.run_libero_eval \ - --benchmark "$suite" \ - --policy-host 127.0.0.1 \ - --policy-port 8000 \ - --num-trials-per-task 50 \ - --result-dir "results/pi05_${suite}" +for suite in spatial object goal 10; do + uv run vla-eval run --config "configs/benchmarks/libero/${suite}.yaml" done ``` diff --git a/docs/zh/models/pi05/libero-four-suites.mdx b/docs/zh/models/pi05/libero-four-suites.mdx index 28ddc48..f6bca44 100644 --- a/docs/zh/models/pi05/libero-four-suites.mdx +++ b/docs/zh/models/pi05/libero-four-suites.mdx @@ -43,6 +43,32 @@ icon: "list-checks" `paligemma-3b-pt-224` 可能需要 HuggingFace 权限。更稳的做法是从已有权限的机器同步到本机。 +# 固定源码版本 + +本复现流程使用以下两个 PR 中的不可变提交： + +| 仓库 | Pull request | 提交 | +| --- | --- | --- | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`c5a9044`](https://github.com/rebecca26358/phyai/commit/c5a904493b74b29a9126a11caaa72ebd51385169) | +| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`0b2fadb`](https://github.com/rebecca26358/vla-evaluation-harness/commit/0b2fadb3f4b30ad9e2591537120171d9cccee86d) | + +检出对应版本： + +```bash +export PHYAI_ROOT=${PHYAI_ROOT:-$HOME/phyai} +export VLA_ROOT=${VLA_ROOT:-$HOME/vla-evaluation-harness} + +git clone https://github.com/rebecca26358/phyai.git "$PHYAI_ROOT" +git -C "$PHYAI_ROOT" checkout c5a904493b74b29a9126a11caaa72ebd51385169 + +git clone https://github.com/rebecca26358/vla-evaluation-harness.git "$VLA_ROOT" +git -C "$VLA_ROOT" checkout 0b2fadb3f4b30ad9e2591537120171d9cccee86d +``` + +这个 VLA 版本包含 `src/vla_eval/model_servers/phyai.py` 和 +`configs/model_servers/phyai/libero.yaml`。server 脚本的 PEP 723 +依赖元数据也固定到了同一个 PhyAI 提交。 + # 设置路径 先在目标机器上设置路径。下面只用占位路径，按实际机器改即可。 @@ -59,7 +85,7 @@ export TOKENIZER_HOST=$MODEL_ROOT/paligemma-3b-pt-224 export PHYAI_CKPT_IN_CONTAINER=/data/share/pi05_libero_phyai_converted export TOKENIZER_IN_CONTAINER=/data/share/paligemma-3b-pt-224 -export LIBERO_IMAGE=ghcr.io/allenai/vla-evaluation-harness/libero:latest +export PHYAI_IMAGE=nvcr.io/nvidia/pytorch:25.12-py3 ``` 检查模型目录： @@ -69,26 +95,19 @@ test -d "$PHYAI_CKPT_HOST" test -d "$TOKENIZER_HOST" ``` -# 准备 Python 环境 +# 准备评测环境 -在 PhyAI 仓库里创建并同步环境： +在宿主机同步 VLA 评测环境。这个命令可用于干净 checkout；PhyAI 并没有定义 `cu130` 或 `libero` 依赖组。 ```bash -cd "$PHYAI_ROOT" -uv sync --group cu130 --extra libero -source .venv/bin/activate -``` - -做一次最小导入检查： - -```bash -python -c "import torch; import phyai; print(torch.__version__, torch.cuda.get_device_name(0))" -python -c "from phyai.policies.pi05_libero import PI05LiberoPolicy; print(PI05LiberoPolicy)" +cd "$VLA_ROOT" +uv sync ``` -# 准备 LIBERO 容器 +# 准备 PhyAI 运行容器 -拉取或构建 `vla-evaluation-harness` 的 LIBERO 镜像。镜像准备好后，启动容器： +在 PhyAI 使用的 CUDA 13 镜像中运行 policy server。评测器会根据 benchmark +配置另行启动 LIBERO 仿真镜像。 ```bash docker run --gpus all -it --rm \ @@ -96,45 +115,39 @@ docker run --gpus all -it --rm \ --ipc=host \ --network=host \ -v "$PHYAI_ROOT:$PHYAI_ROOT" \ + -v "$VLA_ROOT:$VLA_ROOT" \ -v "$MODEL_ROOT:/data/share" \ - "$LIBERO_IMAGE" \ + "$PHYAI_IMAGE" \ bash ``` -进入容器后，确认能看到代码和模型： +进入容器后，同步两个干净 checkout，并检查模型路径： ```bash -test -d "$PHYAI_ROOT" +cd "$PHYAI_ROOT" +uv sync + +cd "$VLA_ROOT" +uv sync + test -d "$PHYAI_CKPT_IN_CONTAINER" test -d "$TOKENIZER_IN_CONTAINER" ``` # 启动 PhyAI policy server -`phyai.policies.pi05_libero` 只是 PhyAI policy adapter，不是 CLI server。WebSocket server 需要通过 `vla-evaluation-harness` 的 server 入口启动： +`phyai.policies.pi05_libero` 只是 policy adapter，不是 CLI server。通过 VLA PR #111 新增的配置启动 WebSocket server： ```bash cd "$VLA_ROOT" -source .venv/bin/activate -export PYTHONPATH="$PHYAI_ROOT/phyai/src:$PHYAI_ROOT/phyai-kernel:$PHYAI_ROOT/phyai-utils-tools/src:$VLA_ROOT/src" +export PHYAI_CHECKPOINT_PATH="$PHYAI_CKPT_IN_CONTAINER" export PHYAI_TOKENIZER_PATH="$TOKENIZER_IN_CONTAINER" -export PHYAI_CAMERA_MODE=two_camera - -python -m vla_eval.model_servers.phyai \ - --checkpoint_path "$PHYAI_CKPT_IN_CONTAINER" \ - --device cuda:0 \ - --params_dtype bfloat16 \ - --attn_backend flashinfer \ - --norm_backend phyai-kernel \ - --linear_backend flashinfer \ - --flashinfer_workspace_bytes 536870912 \ - --chunk_size 10 \ - --host 0.0.0.0 \ - --port 8000 + +uv run vla-eval serve --config configs/model_servers/phyai/libero.yaml ``` -server 启动后保持运行。日志出现 `ws://0.0.0.0:8000` 开始监听后，再启动 smoke test；第一次请求通常会触发初始化和图捕获。 +保持 server 运行，评测器会连接 `ws://localhost:8000`。 # 跑 smoke test @@ -143,13 +156,7 @@ server 启动后保持运行。日志出现 `ws://0.0.0.0:8000` 开始监听后 ```bash cd "$VLA_ROOT" -python -m experiments.robot.libero.run_libero_eval \ - --benchmark libero_spatial \ - --policy-host 127.0.0.1 \ - --policy-port 8000 \ - --num-trials-per-task 1 \ - --max-steps 20 \ - --result-dir results/pi05_libero_smoke +uv run vla-eval run --config configs/benchmarks/libero/smoke_test.yaml ``` 如果 smoke test 失败，先不要跑完整四套任务。优先检查 server 日志、端口、模型路径和 tokenizer 路径。 @@ -161,13 +168,8 @@ python -m experiments.robot.libero.run_libero_eval \ ```bash cd "$VLA_ROOT" -for suite in libero_spatial libero_object libero_goal libero_10; do - python -m experiments.robot.libero.run_libero_eval \ - --benchmark "$suite" \ - --policy-host 127.0.0.1 \ - --policy-port 8000 \ - --num-trials-per-task 50 \ - --result-dir "results/pi05_${suite}" +for suite in spatial object goal 10; do + uv run vla-eval run --config "configs/benchmarks/libero/${suite}.yaml" done ``` From 4e95665fae38e7d13a737922398eaadb4db24669 Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 18:58:39 +0800 Subject: [PATCH 23/29] feat(phyai): support batched LIBERO serving --- phyai/src/phyai/policies/pi05_libero.py | 100 ++++++++++++++++++++--- phyai/tests/policies/test_pi05_libero.py | 33 ++++++++ 2 files changed, 121 insertions(+), 12 deletions(-) create mode 100644 phyai/tests/policies/test_pi05_libero.py diff --git a/phyai/src/phyai/policies/pi05_libero.py b/phyai/src/phyai/policies/pi05_libero.py index 3268db7..72550e3 100644 --- a/phyai/src/phyai/policies/pi05_libero.py +++ b/phyai/src/phyai/policies/pi05_libero.py @@ -13,10 +13,17 @@ from safetensors.torch import load_file from phyai.engine import Engine, EngineArgs -from phyai.engine_config import BackendConfig, DeviceConfig, EngineConfig, RuntimeConfig +from phyai.engine_config import ( + BackendConfig, + DeviceConfig, + EngineConfig, + ParallelConfig, + RuntimeConfig, +) from phyai.env import envs from phyai.models.pi05.configuration_pi05 import PI05Config from phyai.models.pi05.main_pi05 import PI05Args +from phyai.models.pi05.main_pi05_wn import PI05WNArgs from phyai.models.pi05.scheduler_ws1_pi05 import PI05Request from phyai_utils_tools.models.pi05 import PI05_DEFAULT_TOKENIZER_NAME, PI05Processor from phyai_utils_tools.processing.transition import IMAGES, STATE, TASK @@ -63,11 +70,17 @@ def __init__( tokenizer_name: str | None = None, camera_mode: str | None = None, chunk_size: int | None = None, + engine_plugin: str = "pi05", + world_size: int = 1, + dp_size: int = 1, ) -> None: self.checkpoint_dir = Path(checkpoint_dir) self.device = device self.params_dtype = params_dtype self.max_batch_size = int(max_batch_size) + self.engine_plugin = engine_plugin + self.world_size = int(world_size) + self.dp_size = int(dp_size) self.config = self._read_config() self._engine_config = self._resolve_engine_config(chunk_size) self.image_size = self._resolve_image_size(self.config) @@ -106,23 +119,45 @@ def __init__( device=device, params_dtype=params_dtype, ) + if self.engine_plugin == "pi05": + plugin_args = PI05Args( + checkpoint_dir=self.checkpoint_dir, + config=self._engine_config, + max_batch_size=self.max_batch_size, + weight_remap=_lerobot_pi05_weight_remap, + inputs_image_shape=[ + [self.image_size, self.image_size, 3] for _ in self.camera_names + ], + ) + parallel = ParallelConfig() + elif self.engine_plugin == "pi05_wn": + plugin_args = PI05WNArgs( + checkpoint_dir=self.checkpoint_dir, + config=self._engine_config, + max_batch_size=self.max_batch_size, + weight_remap=_lerobot_pi05_weight_remap, + inputs_image_shape=[ + [self.image_size, self.image_size, 3] for _ in self.camera_names + ], + ) + parallel = ParallelConfig( + world_size=self.world_size, + dp_size=self.dp_size, + tp_size=1, + ) + else: + raise ValueError(f"Unsupported engine_plugin={self.engine_plugin!r}.") + self.engine = Engine( EngineArgs( - plugin="pi05", - plugin_args=PI05Args( - checkpoint_dir=self.checkpoint_dir, - config=self._engine_config, - max_batch_size=self.max_batch_size, - weight_remap=_lerobot_pi05_weight_remap, - inputs_image_shape=[ - [self.image_size, self.image_size, 3] for _ in self.camera_names - ], - ), + plugin=self.engine_plugin, + plugin_args=plugin_args, config=EngineConfig( backends=BackendConfig( attn=attn_backend, norm=norm_backend, linear=linear_backend ), device=DeviceConfig(target=device, params_dtype=params_dtype), + parallel=parallel, runtime=RuntimeConfig( use_cuda_graph=use_cuda_graph, flashinfer_workspace_bytes=flashinfer_workspace_bytes, @@ -442,7 +477,25 @@ def _postprocess_actions(self, raw_actions: torch.Tensor) -> np.ndarray: def infer( self, obs: dict[str, Any], *, noise: torch.Tensor | np.ndarray | None = None ) -> dict[str, np.ndarray]: - request_kwargs = self.observation_to_request_inputs(obs) + return self.infer_batch([obs], noise=noise) + + def infer_batch( + self, + obs_batch: list[dict[str, Any]], + *, + noise: torch.Tensor | np.ndarray | None = None, + ) -> dict[str, np.ndarray]: + if not obs_batch: + return { + "actions": np.empty( + (0, self.chunk_size, self.action_dim), dtype=np.float32 + ) + } + request_items = [self.observation_to_request_inputs(obs) for obs in obs_batch] + request_kwargs = { + key: torch.cat([item[key] for item in request_items], dim=0) + for key in ("pixel_values", "input_ids", "lang_lens") + } if noise is not None: request_kwargs["noise"] = torch.as_tensor(noise, device=self.device) request = PI05Request(**request_kwargs) @@ -451,5 +504,28 @@ def infer( actions = self._postprocess_actions(raw_actions) return {"actions": actions} + def infer_distributed_worker(self) -> None: + """Participate in one ``pi05_wn`` step on a non-router rank.""" + request = PI05Request( + pixel_values=torch.empty( + 1, + len(self.camera_names), + 3, + self.image_size, + self.image_size, + device=self.device, + dtype=self.params_dtype, + ), + input_ids=torch.empty( + 1, + int(self.config.get("tokenizer_max_length", 200)), + device=self.device, + dtype=torch.int64, + ), + lang_lens=torch.empty(1, device=self.device, dtype=torch.int64), + ) + with torch.inference_mode(): + self.engine.step(request) + def close(self) -> None: self.engine.close() diff --git a/phyai/tests/policies/test_pi05_libero.py b/phyai/tests/policies/test_pi05_libero.py new file mode 100644 index 0000000..b0822ba --- /dev/null +++ b/phyai/tests/policies/test_pi05_libero.py @@ -0,0 +1,33 @@ +"""PI0.5 LIBERO 策略的轻量接口测试。""" + +from __future__ import annotations + +import numpy as np + +from phyai.policies.pi05_libero import PI05LiberoPolicy + + +def test_single_infer_delegates_to_one_item_batch(monkeypatch): + policy = object.__new__(PI05LiberoPolicy) + expected = {"actions": np.zeros((1, 10, 7), dtype=np.float32)} + calls = [] + + def infer_batch(obs_batch, *, noise=None): + calls.append((obs_batch, noise)) + return expected + + monkeypatch.setattr(policy, "infer_batch", infer_batch) + observation = {"state": np.zeros(7, dtype=np.float32)} + + assert policy.infer(observation) is expected + assert calls == [([observation], None)] + + +def test_single_card_constructor_defaults_are_preserved(): + defaults = PI05LiberoPolicy.__init__.__kwdefaults__ + + assert defaults is not None + assert defaults["max_batch_size"] == 1 + assert defaults["engine_plugin"] == "pi05" + assert defaults["world_size"] == 1 + assert defaults["dp_size"] == 1 From 6604081c3501e5c480d65e5f565bfa8eb597fe06 Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 19:13:39 +0800 Subject: [PATCH 24/29] docs(phyai): pin LIBERO server revisions --- docs/models/pi05/eight-gpu-inference.mdx | 13 +++++++++++++ docs/models/pi05/libero-four-suites.mdx | 11 ++++++----- docs/zh/models/pi05/eight-gpu-inference.mdx | 13 +++++++++++++ docs/zh/models/pi05/libero-four-suites.mdx | 11 ++++++----- 4 files changed, 38 insertions(+), 10 deletions(-) diff --git a/docs/models/pi05/eight-gpu-inference.mdx b/docs/models/pi05/eight-gpu-inference.mdx index e144b6c..1c8cd02 100644 --- a/docs/models/pi05/eight-gpu-inference.mdx +++ b/docs/models/pi05/eight-gpu-inference.mdx @@ -22,6 +22,19 @@ The current demo uses `chunk_size=1`. Each request returns one action, so the cl Other machines can use the same flow. You need 8 idle CUDA GPUs, Docker, `tmux`, the PI0.5 LIBERO checkpoint, and an offline tokenizer. Different GPU models mainly change latency. +# Pin Source Revisions + +Use the linked pull requests at these exact commits: + +| Repository | Pull request | Commit | +| --- | --- | --- | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e9566511d4d7ad5ff45ec072c667e93ccc27483) | +| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`6ac5b32`](https://github.com/rebecca26358/vla-evaluation-harness/commit/6ac5b3212a87e93f3d507699fbbbc416d8ed6945) | + +The VLA commit provides both `vla_eval.model_servers.phyai` and +`vla_eval.model_servers.phyai_wn`; the PhyAI commit provides the matching +`PI05LiberoPolicy` batch and distributed-worker interfaces. + # Paths Set these paths for your machine: diff --git a/docs/models/pi05/libero-four-suites.mdx b/docs/models/pi05/libero-four-suites.mdx index 951c21d..b8f29d7 100644 --- a/docs/models/pi05/libero-four-suites.mdx +++ b/docs/models/pi05/libero-four-suites.mdx @@ -49,8 +49,8 @@ This reproduction uses immutable commits from two linked pull requests: | Repository | Pull request | Commit | | --- | --- | --- | -| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`c5a9044`](https://github.com/rebecca26358/phyai/commit/c5a904493b74b29a9126a11caaa72ebd51385169) | -| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`0b2fadb`](https://github.com/rebecca26358/vla-evaluation-harness/commit/0b2fadb3f4b30ad9e2591537120171d9cccee86d) | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e9566511d4d7ad5ff45ec072c667e93ccc27483) | +| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`6ac5b32`](https://github.com/rebecca26358/vla-evaluation-harness/commit/6ac5b3212a87e93f3d507699fbbbc416d8ed6945) | Check out those exact revisions: @@ -59,13 +59,14 @@ export PHYAI_ROOT=${PHYAI_ROOT:-$HOME/phyai} export VLA_ROOT=${VLA_ROOT:-$HOME/vla-evaluation-harness} git clone https://github.com/rebecca26358/phyai.git "$PHYAI_ROOT" -git -C "$PHYAI_ROOT" checkout c5a904493b74b29a9126a11caaa72ebd51385169 +git -C "$PHYAI_ROOT" checkout 4e9566511d4d7ad5ff45ec072c667e93ccc27483 git clone https://github.com/rebecca26358/vla-evaluation-harness.git "$VLA_ROOT" -git -C "$VLA_ROOT" checkout 0b2fadb3f4b30ad9e2591537120171d9cccee86d +git -C "$VLA_ROOT" checkout 6ac5b3212a87e93f3d507699fbbbc416d8ed6945 ``` -The VLA revision contains `src/vla_eval/model_servers/phyai.py` and +The VLA revision contains `src/vla_eval/model_servers/phyai.py`, +`src/vla_eval/model_servers/phyai_wn.py`, and `configs/model_servers/phyai/libero.yaml`. The server script pins the same PhyAI commit in its PEP 723 dependency metadata. diff --git a/docs/zh/models/pi05/eight-gpu-inference.mdx b/docs/zh/models/pi05/eight-gpu-inference.mdx index bac3f27..4f99f1a 100644 --- a/docs/zh/models/pi05/eight-gpu-inference.mdx +++ b/docs/zh/models/pi05/eight-gpu-inference.mdx @@ -22,6 +22,19 @@ icon: "server" 其它机器也可以复现。至少需要 8 张空闲 CUDA GPU、Docker、`tmux`，以及 PI0.5 LIBERO checkpoint 和离线 tokenizer。卡型不同会影响耗时，但流程不变。 +# 固定源码版本 + +使用以下两个 PR 中的精确提交： + +| 仓库 | Pull request | 提交 | +| --- | --- | --- | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e9566511d4d7ad5ff45ec072c667e93ccc27483) | +| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`6ac5b32`](https://github.com/rebecca26358/vla-evaluation-harness/commit/6ac5b3212a87e93f3d507699fbbbc416d8ed6945) | + +VLA 提交同时提供 `vla_eval.model_servers.phyai` 和 +`vla_eval.model_servers.phyai_wn`；PhyAI 提供与之匹配的 +`PI05LiberoPolicy` 批处理与分布式 worker 接口。 + # 路径设置 按机器实际目录修改下面几个路径： diff --git a/docs/zh/models/pi05/libero-four-suites.mdx b/docs/zh/models/pi05/libero-four-suites.mdx index f6bca44..dd1a0bd 100644 --- a/docs/zh/models/pi05/libero-four-suites.mdx +++ b/docs/zh/models/pi05/libero-four-suites.mdx @@ -49,8 +49,8 @@ icon: "list-checks" | 仓库 | Pull request | 提交 | | --- | --- | --- | -| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`c5a9044`](https://github.com/rebecca26358/phyai/commit/c5a904493b74b29a9126a11caaa72ebd51385169) | -| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`0b2fadb`](https://github.com/rebecca26358/vla-evaluation-harness/commit/0b2fadb3f4b30ad9e2591537120171d9cccee86d) | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e9566511d4d7ad5ff45ec072c667e93ccc27483) | +| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`6ac5b32`](https://github.com/rebecca26358/vla-evaluation-harness/commit/6ac5b3212a87e93f3d507699fbbbc416d8ed6945) | 检出对应版本： @@ -59,13 +59,14 @@ export PHYAI_ROOT=${PHYAI_ROOT:-$HOME/phyai} export VLA_ROOT=${VLA_ROOT:-$HOME/vla-evaluation-harness} git clone https://github.com/rebecca26358/phyai.git "$PHYAI_ROOT" -git -C "$PHYAI_ROOT" checkout c5a904493b74b29a9126a11caaa72ebd51385169 +git -C "$PHYAI_ROOT" checkout 4e9566511d4d7ad5ff45ec072c667e93ccc27483 git clone https://github.com/rebecca26358/vla-evaluation-harness.git "$VLA_ROOT" -git -C "$VLA_ROOT" checkout 0b2fadb3f4b30ad9e2591537120171d9cccee86d +git -C "$VLA_ROOT" checkout 6ac5b3212a87e93f3d507699fbbbc416d8ed6945 ``` -这个 VLA 版本包含 `src/vla_eval/model_servers/phyai.py` 和 +这个 VLA 版本包含 `src/vla_eval/model_servers/phyai.py`、 +`src/vla_eval/model_servers/phyai_wn.py` 和 `configs/model_servers/phyai/libero.yaml`。server 脚本的 PEP 723 依赖元数据也固定到了同一个 PhyAI 提交。 From 5fc3288dcd8fe1de0a3859a28cfa7036e42217d8 Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:09:11 +0800 Subject: [PATCH 25/29] test(phyai): cover LIBERO batched inference --- docs/models/pi05/libero-four-suites.mdx | 25 ++++++++++++++-- docs/zh/models/pi05/libero-four-suites.mdx | 24 +++++++++++++-- phyai/tests/policies/test_pi05_libero.py | 35 ++++++++++++++++++++++ 3 files changed, 80 insertions(+), 4 deletions(-) diff --git a/docs/models/pi05/libero-four-suites.mdx b/docs/models/pi05/libero-four-suites.mdx index b8f29d7..2621bb6 100644 --- a/docs/models/pi05/libero-four-suites.mdx +++ b/docs/models/pi05/libero-four-suites.mdx @@ -135,6 +135,19 @@ test -d "$PHYAI_CKPT_IN_CONTAINER" test -d "$TOKENIZER_IN_CONTAINER" ``` +# Validate the Policy Adapter + +Run the lightweight policy tests inside the PhyAI runtime container: + +```bash +cd "$PHYAI_ROOT" +uv run pytest phyai/tests/policies/test_pi05_libero.py -q +``` + +The expected result is `3 passed`. These tests cover the default single-GPU +configuration, single-observation inference, and batched input concatenation. +They do not load a checkpoint or execute CUDA model inference. + # Start the Policy Server `phyai.policies.pi05_libero` is a policy adapter, not a CLI server. Start the WebSocket server through the VLA config added by PR #111: @@ -157,7 +170,11 @@ In another terminal, run a short check to confirm that the evaluator can reach t ```bash cd "$VLA_ROOT" -uv run vla-eval run --config configs/benchmarks/libero/smoke_test.yaml +uv run vla-eval run \ + --config configs/benchmarks/libero/smoke_test.yaml \ + --server-url ws://localhost:8000 \ + --dev \ + --yes ``` If the smoke test fails, stop there and check the server log, port, checkpoint path, and tokenizer path first. @@ -170,7 +187,11 @@ After confirming that the GPU is idle, run the full evaluation. Save each suite cd "$VLA_ROOT" for suite in spatial object goal 10; do - uv run vla-eval run --config "configs/benchmarks/libero/${suite}.yaml" + uv run vla-eval run \ + --config "configs/benchmarks/libero/${suite}.yaml" \ + --server-url ws://localhost:8000 \ + --dev \ + --yes done ``` diff --git a/docs/zh/models/pi05/libero-four-suites.mdx b/docs/zh/models/pi05/libero-four-suites.mdx index dd1a0bd..aae1ab0 100644 --- a/docs/zh/models/pi05/libero-four-suites.mdx +++ b/docs/zh/models/pi05/libero-four-suites.mdx @@ -135,6 +135,18 @@ test -d "$PHYAI_CKPT_IN_CONTAINER" test -d "$TOKENIZER_IN_CONTAINER" ``` +# 验证 policy adapter + +在 PhyAI 运行容器内执行轻量 policy 测试： + +```bash +cd "$PHYAI_ROOT" +uv run pytest phyai/tests/policies/test_pi05_libero.py -q +``` + +预期结果为 `3 passed`。这些测试覆盖默认单卡配置、单条 observation 推理和 +batch 输入拼接，不会加载 checkpoint，也不会执行 CUDA 模型推理。 + # 启动 PhyAI policy server `phyai.policies.pi05_libero` 只是 policy adapter，不是 CLI server。通过 VLA PR #111 新增的配置启动 WebSocket server： @@ -157,7 +169,11 @@ uv run vla-eval serve --config configs/model_servers/phyai/libero.yaml ```bash cd "$VLA_ROOT" -uv run vla-eval run --config configs/benchmarks/libero/smoke_test.yaml +uv run vla-eval run \ + --config configs/benchmarks/libero/smoke_test.yaml \ + --server-url ws://localhost:8000 \ + --dev \ + --yes ``` 如果 smoke test 失败，先不要跑完整四套任务。优先检查 server 日志、端口、模型路径和 tokenizer 路径。 @@ -170,7 +186,11 @@ uv run vla-eval run --config configs/benchmarks/libero/smoke_test.yaml cd "$VLA_ROOT" for suite in spatial object goal 10; do - uv run vla-eval run --config "configs/benchmarks/libero/${suite}.yaml" + uv run vla-eval run \ + --config "configs/benchmarks/libero/${suite}.yaml" \ + --server-url ws://localhost:8000 \ + --dev \ + --yes done ``` diff --git a/phyai/tests/policies/test_pi05_libero.py b/phyai/tests/policies/test_pi05_libero.py index b0822ba..45b30fc 100644 --- a/phyai/tests/policies/test_pi05_libero.py +++ b/phyai/tests/policies/test_pi05_libero.py @@ -3,6 +3,7 @@ from __future__ import annotations import numpy as np +import torch from phyai.policies.pi05_libero import PI05LiberoPolicy @@ -31,3 +32,37 @@ def test_single_card_constructor_defaults_are_preserved(): assert defaults["engine_plugin"] == "pi05" assert defaults["world_size"] == 1 assert defaults["dp_size"] == 1 + + +def test_infer_batch_concatenates_inputs_and_steps_once(monkeypatch): + policy = object.__new__(PI05LiberoPolicy) + policy.device = "cpu" + calls = [] + + def observation_to_request_inputs(obs): + value = int(obs["value"]) + return { + "pixel_values": torch.full((1, 2, 3, 4, 4), value, dtype=torch.float32), + "input_ids": torch.full((1, 3), value, dtype=torch.int64), + "lang_lens": torch.tensor([value], dtype=torch.int64), + } + + class _Engine: + def step(self, request): + calls.append(request) + return torch.zeros(2, 10, 7) + + monkeypatch.setattr( + policy, "observation_to_request_inputs", observation_to_request_inputs + ) + monkeypatch.setattr(policy, "_postprocess_actions", lambda actions: actions.numpy()) + policy.engine = _Engine() + + result = policy.infer_batch([{"value": 1}, {"value": 2}]) + + assert len(calls) == 1 + request = calls[0] + assert request.pixel_values.shape == (2, 2, 3, 4, 4) + assert request.input_ids.tolist() == [[1, 1, 1], [2, 2, 2]] + assert request.lang_lens.tolist() == [1, 2] + assert result["actions"].shape == (2, 10, 7) From b081a005a95f1db9b42ea9d997acf1511ba5433c Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:35:23 +0800 Subject: [PATCH 26/29] docs(phyai): fix pinned LIBERO revisions --- docs/models/pi05/eight-gpu-inference.mdx | 6 +++--- docs/models/pi05/libero-four-suites.mdx | 12 ++++++------ docs/zh/models/pi05/eight-gpu-inference.mdx | 6 +++--- docs/zh/models/pi05/libero-four-suites.mdx | 12 ++++++------ 4 files changed, 18 insertions(+), 18 deletions(-) diff --git a/docs/models/pi05/eight-gpu-inference.mdx b/docs/models/pi05/eight-gpu-inference.mdx index 1c8cd02..428330b 100644 --- a/docs/models/pi05/eight-gpu-inference.mdx +++ b/docs/models/pi05/eight-gpu-inference.mdx @@ -28,12 +28,12 @@ Use the linked pull requests at these exact commits: | Repository | Pull request | Commit | | --- | --- | --- | -| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e9566511d4d7ad5ff45ec072c667e93ccc27483) | -| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`6ac5b32`](https://github.com/rebecca26358/vla-evaluation-harness/commit/6ac5b3212a87e93f3d507699fbbbc416d8ed6945) | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`5fc3288`](https://github.com/rebecca26358/phyai/commit/5fc3288dcd8fe1de0a3859a28cfa7036e42217d8) | +| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`8bb9009`](https://github.com/rebecca26358/vla-evaluation-harness/commit/8bb90099c719ee8d151497c6636e76a9da5d5c6f) | The VLA commit provides both `vla_eval.model_servers.phyai` and `vla_eval.model_servers.phyai_wn`; the PhyAI commit provides the matching -`PI05LiberoPolicy` batch and distributed-worker interfaces. +`PI05LiberoPolicy` batch and distributed-worker interfaces. The VLA server PEP 723 metadata pins the policy implementation commit [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e95665fae38e7d13a737922398eaadb4db24669); the PhyAI checkout above additionally contains the documented three-test validation. # Paths diff --git a/docs/models/pi05/libero-four-suites.mdx b/docs/models/pi05/libero-four-suites.mdx index 2621bb6..c95e743 100644 --- a/docs/models/pi05/libero-four-suites.mdx +++ b/docs/models/pi05/libero-four-suites.mdx @@ -49,8 +49,8 @@ This reproduction uses immutable commits from two linked pull requests: | Repository | Pull request | Commit | | --- | --- | --- | -| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e9566511d4d7ad5ff45ec072c667e93ccc27483) | -| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`6ac5b32`](https://github.com/rebecca26358/vla-evaluation-harness/commit/6ac5b3212a87e93f3d507699fbbbc416d8ed6945) | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`5fc3288`](https://github.com/rebecca26358/phyai/commit/5fc3288dcd8fe1de0a3859a28cfa7036e42217d8) | +| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`8bb9009`](https://github.com/rebecca26358/vla-evaluation-harness/commit/8bb90099c719ee8d151497c6636e76a9da5d5c6f) | Check out those exact revisions: @@ -59,16 +59,16 @@ export PHYAI_ROOT=${PHYAI_ROOT:-$HOME/phyai} export VLA_ROOT=${VLA_ROOT:-$HOME/vla-evaluation-harness} git clone https://github.com/rebecca26358/phyai.git "$PHYAI_ROOT" -git -C "$PHYAI_ROOT" checkout 4e9566511d4d7ad5ff45ec072c667e93ccc27483 +git -C "$PHYAI_ROOT" checkout 5fc3288dcd8fe1de0a3859a28cfa7036e42217d8 git clone https://github.com/rebecca26358/vla-evaluation-harness.git "$VLA_ROOT" -git -C "$VLA_ROOT" checkout 6ac5b3212a87e93f3d507699fbbbc416d8ed6945 +git -C "$VLA_ROOT" checkout 8bb90099c719ee8d151497c6636e76a9da5d5c6f ``` The VLA revision contains `src/vla_eval/model_servers/phyai.py`, `src/vla_eval/model_servers/phyai_wn.py`, and -`configs/model_servers/phyai/libero.yaml`. The server script pins the same -PhyAI commit in its PEP 723 dependency metadata. +`configs/model_servers/phyai/libero.yaml`. Its PEP 723 metadata pins the policy +implementation commit [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e95665fae38e7d13a737922398eaadb4db24669), an ancestor of the PhyAI checkout above. The later PhyAI commits add documentation and tests without changing the policy implementation. # Set Paths diff --git a/docs/zh/models/pi05/eight-gpu-inference.mdx b/docs/zh/models/pi05/eight-gpu-inference.mdx index 4f99f1a..04970c2 100644 --- a/docs/zh/models/pi05/eight-gpu-inference.mdx +++ b/docs/zh/models/pi05/eight-gpu-inference.mdx @@ -28,12 +28,12 @@ icon: "server" | 仓库 | Pull request | 提交 | | --- | --- | --- | -| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e9566511d4d7ad5ff45ec072c667e93ccc27483) | -| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`6ac5b32`](https://github.com/rebecca26358/vla-evaluation-harness/commit/6ac5b3212a87e93f3d507699fbbbc416d8ed6945) | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`5fc3288`](https://github.com/rebecca26358/phyai/commit/5fc3288dcd8fe1de0a3859a28cfa7036e42217d8) | +| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`8bb9009`](https://github.com/rebecca26358/vla-evaluation-harness/commit/8bb90099c719ee8d151497c6636e76a9da5d5c6f) | VLA 提交同时提供 `vla_eval.model_servers.phyai` 和 `vla_eval.model_servers.phyai_wn`；PhyAI 提供与之匹配的 -`PI05LiberoPolicy` 批处理与分布式 worker 接口。 +`PI05LiberoPolicy` 批处理与分布式 worker 接口。VLA server 的 PEP 723 元数据固定到 policy 实现提交 [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e95665fae38e7d13a737922398eaadb4db24669)；上面的 PhyAI checkout 还包含文档中记录的 3 条测试。 # 路径设置 diff --git a/docs/zh/models/pi05/libero-four-suites.mdx b/docs/zh/models/pi05/libero-four-suites.mdx index aae1ab0..e64c8ee 100644 --- a/docs/zh/models/pi05/libero-four-suites.mdx +++ b/docs/zh/models/pi05/libero-four-suites.mdx @@ -49,8 +49,8 @@ icon: "list-checks" | 仓库 | Pull request | 提交 | | --- | --- | --- | -| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e9566511d4d7ad5ff45ec072c667e93ccc27483) | -| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`6ac5b32`](https://github.com/rebecca26358/vla-evaluation-harness/commit/6ac5b3212a87e93f3d507699fbbbc416d8ed6945) | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`5fc3288`](https://github.com/rebecca26358/phyai/commit/5fc3288dcd8fe1de0a3859a28cfa7036e42217d8) | +| `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`8bb9009`](https://github.com/rebecca26358/vla-evaluation-harness/commit/8bb90099c719ee8d151497c6636e76a9da5d5c6f) | 检出对应版本： @@ -59,16 +59,16 @@ export PHYAI_ROOT=${PHYAI_ROOT:-$HOME/phyai} export VLA_ROOT=${VLA_ROOT:-$HOME/vla-evaluation-harness} git clone https://github.com/rebecca26358/phyai.git "$PHYAI_ROOT" -git -C "$PHYAI_ROOT" checkout 4e9566511d4d7ad5ff45ec072c667e93ccc27483 +git -C "$PHYAI_ROOT" checkout 5fc3288dcd8fe1de0a3859a28cfa7036e42217d8 git clone https://github.com/rebecca26358/vla-evaluation-harness.git "$VLA_ROOT" -git -C "$VLA_ROOT" checkout 6ac5b3212a87e93f3d507699fbbbc416d8ed6945 +git -C "$VLA_ROOT" checkout 8bb90099c719ee8d151497c6636e76a9da5d5c6f ``` 这个 VLA 版本包含 `src/vla_eval/model_servers/phyai.py`、 `src/vla_eval/model_servers/phyai_wn.py` 和 -`configs/model_servers/phyai/libero.yaml`。server 脚本的 PEP 723 -依赖元数据也固定到了同一个 PhyAI 提交。 +`configs/model_servers/phyai/libero.yaml`。其 PEP 723 元数据固定到 policy +实现提交 [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e95665fae38e7d13a737922398eaadb4db24669)，它是上面 PhyAI checkout 的祖先；后续 PhyAI 提交只增加文档和测试，没有修改 policy 实现。 # 设置路径 From 3a9a28b21aa53627a07c3be9fcb86e6adb1b4ae8 Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 21:28:03 +0800 Subject: [PATCH 27/29] feat(phyai): add OpenPI pi0.5 checkpoint converter --- docs/models/pi05/libero-four-suites.mdx | 50 +- docs/zh/models/pi05/libero-four-suites.mdx | 48 +- tools/convert_openpi_pi05_to_phyai.py | 856 +++++++++++++++++++++ 3 files changed, 948 insertions(+), 6 deletions(-) create mode 100644 tools/convert_openpi_pi05_to_phyai.py diff --git a/docs/models/pi05/libero-four-suites.mdx b/docs/models/pi05/libero-four-suites.mdx index c95e743..f738bdd 100644 --- a/docs/models/pi05/libero-four-suites.mdx +++ b/docs/models/pi05/libero-four-suites.mdx @@ -37,8 +37,8 @@ Prepare these resources first: | --- | --- | | PhyAI repository | Contains the PI0.5 policy and benchmark scripts | | `vla-evaluation-harness` | LIBERO simulation benchmark framework | -| PI0.5 LIBERO checkpoint | Converted to the format PhyAI can load | -| `paligemma-3b-pt-224` | Tokenizer / processor used by PI0.5 | +| PI0.5 LIBERO checkpoint | OpenPI's official `pi05_libero` checkpoint, converted to the format PhyAI can load | +| [`google/paligemma-3b-pt-224`](https://huggingface.co/google/paligemma-3b-pt-224) | Tokenizer / processor used by PI0.5 | | Docker + NVIDIA Container Toolkit | Required for the LIBERO simulation container | `paligemma-3b-pt-224` may require HuggingFace access. If possible, sync it from a machine that already has access instead of downloading it during the run. @@ -70,6 +70,45 @@ The VLA revision contains `src/vla_eval/model_servers/phyai.py`, `configs/model_servers/phyai/libero.yaml`. Its PEP 723 metadata pins the policy implementation commit [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e95665fae38e7d13a737922398eaadb4db24669), an ancestor of the PhyAI checkout above. The later PhyAI commits add documentation and tests without changing the policy implementation. +# Download and Convert the Checkpoint + +OpenPI publishes the LIBERO checkpoint at +[`gs://openpi-assets/checkpoints/pi05_libero`](https://github.com/Physical-Intelligence/openpi/blob/main/examples/libero/README.md). +Download the checkpoint and tokenizer into `MODEL_ROOT`: + +```bash +export MODEL_ROOT=${MODEL_ROOT:-$HOME/phyai_models} +export OPENPI_CKPT_HOST=$MODEL_ROOT/pi05_libero +export PHYAI_CKPT_HOST=$MODEL_ROOT/pi05_libero_phyai_converted +export TOKENIZER_HOST=$MODEL_ROOT/paligemma-3b-pt-224 + +mkdir -p "$MODEL_ROOT" +gcloud storage rsync --recursive \ + gs://openpi-assets/checkpoints/pi05_libero \ + "$OPENPI_CKPT_HOST" + +uvx --from huggingface_hub hf download \ + google/paligemma-3b-pt-224 \ + --local-dir "$TOKENIZER_HOST" +``` + +The PaliGemma download requires accepting the model terms and authenticating +with Hugging Face. Convert the OpenPI Orbax checkpoint into the safetensors and +processor files loaded by `PI05LiberoPolicy`: + +```bash +cd "$PHYAI_ROOT" +uv run tools/convert_openpi_pi05_to_phyai.py \ + --checkpoint "$OPENPI_CKPT_HOST" \ + --write \ + --out "$PHYAI_CKPT_HOST/model.safetensors" +``` + +The converter uses pinned PEP 723 dependencies. It reads normalization +statistics from the official checkpoint and creates `config.json`, the two +processor JSON files, and their safetensors sidecars next to +`model.safetensors`. + # Set Paths Set paths on the target machine. These are placeholders; change them for your environment. @@ -80,6 +119,7 @@ export VLA_ROOT=$HOME/vla-evaluation-harness export MODEL_ROOT=$HOME/phyai_models export PHYAI_CONTAINER=phyai_libero_eval +export OPENPI_CKPT_HOST=$MODEL_ROOT/pi05_libero export PHYAI_CKPT_HOST=$MODEL_ROOT/pi05_libero_phyai_converted export TOKENIZER_HOST=$MODEL_ROOT/paligemma-3b-pt-224 @@ -92,7 +132,11 @@ export PHYAI_IMAGE=nvcr.io/nvidia/pytorch:25.12-py3 Check the model directories: ```bash -test -d "$PHYAI_CKPT_HOST" +test -d "$OPENPI_CKPT_HOST" +test -f "$PHYAI_CKPT_HOST/model.safetensors" +test -f "$PHYAI_CKPT_HOST/config.json" +test -f "$PHYAI_CKPT_HOST/policy_preprocessor.json" +test -f "$PHYAI_CKPT_HOST/policy_postprocessor.json" test -d "$TOKENIZER_HOST" ``` diff --git a/docs/zh/models/pi05/libero-four-suites.mdx b/docs/zh/models/pi05/libero-four-suites.mdx index e64c8ee..5a6fd23 100644 --- a/docs/zh/models/pi05/libero-four-suites.mdx +++ b/docs/zh/models/pi05/libero-four-suites.mdx @@ -37,8 +37,8 @@ icon: "list-checks" | --- | --- | | PhyAI 仓库 | 包含 PI0.5 policy 和 benchmark 脚本 | | `vla-evaluation-harness` | LIBERO 仿真评测框架 | -| PI0.5 LIBERO checkpoint | 已转换成 PhyAI 可加载格式 | -| `paligemma-3b-pt-224` | PI0.5 使用的 tokenizer / processor | +| PI0.5 LIBERO checkpoint | OpenPI 官方发布的 `pi05_libero` checkpoint，需转换为 PhyAI 可加载格式 | +| [`google/paligemma-3b-pt-224`](https://huggingface.co/google/paligemma-3b-pt-224) | PI0.5 使用的 tokenizer / processor | | Docker + NVIDIA Container Toolkit | 用来启动 LIBERO 仿真容器 | `paligemma-3b-pt-224` 可能需要 HuggingFace 权限。更稳的做法是从已有权限的机器同步到本机。 @@ -70,6 +70,43 @@ git -C "$VLA_ROOT" checkout 8bb90099c719ee8d151497c6636e76a9da5d5c6f `configs/model_servers/phyai/libero.yaml`。其 PEP 723 元数据固定到 policy 实现提交 [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e95665fae38e7d13a737922398eaadb4db24669)，它是上面 PhyAI checkout 的祖先；后续 PhyAI 提交只增加文档和测试，没有修改 policy 实现。 +# 下载并转换 checkpoint + +OpenPI 官方在 +[`gs://openpi-assets/checkpoints/pi05_libero`](https://github.com/Physical-Intelligence/openpi/blob/main/examples/libero/README.md) +发布 LIBERO checkpoint。先把 checkpoint 和 tokenizer 下载到 `MODEL_ROOT`： + +```bash +export MODEL_ROOT=${MODEL_ROOT:-$HOME/phyai_models} +export OPENPI_CKPT_HOST=$MODEL_ROOT/pi05_libero +export PHYAI_CKPT_HOST=$MODEL_ROOT/pi05_libero_phyai_converted +export TOKENIZER_HOST=$MODEL_ROOT/paligemma-3b-pt-224 + +mkdir -p "$MODEL_ROOT" +gcloud storage rsync --recursive \ + gs://openpi-assets/checkpoints/pi05_libero \ + "$OPENPI_CKPT_HOST" + +uvx --from huggingface_hub hf download \ + google/paligemma-3b-pt-224 \ + --local-dir "$TOKENIZER_HOST" +``` + +下载 PaliGemma 前需要在 Hugging Face 接受模型条款并完成登录。接着把 OpenPI +Orbax checkpoint 转成 `PI05LiberoPolicy` 可读取的 safetensors 和 processor 文件： + +```bash +cd "$PHYAI_ROOT" +uv run tools/convert_openpi_pi05_to_phyai.py \ + --checkpoint "$OPENPI_CKPT_HOST" \ + --write \ + --out "$PHYAI_CKPT_HOST/model.safetensors" +``` + +转换脚本通过 PEP 723 固定依赖版本。它从官方 checkpoint 读取 normalization +stats，并在 `model.safetensors` 同级目录生成 `config.json`、两份 processor JSON +和对应的 safetensors sidecar。 + # 设置路径 先在目标机器上设置路径。下面只用占位路径，按实际机器改即可。 @@ -80,6 +117,7 @@ export VLA_ROOT=$HOME/vla-evaluation-harness export MODEL_ROOT=$HOME/phyai_models export PHYAI_CONTAINER=phyai_libero_eval +export OPENPI_CKPT_HOST=$MODEL_ROOT/pi05_libero export PHYAI_CKPT_HOST=$MODEL_ROOT/pi05_libero_phyai_converted export TOKENIZER_HOST=$MODEL_ROOT/paligemma-3b-pt-224 @@ -92,7 +130,11 @@ export PHYAI_IMAGE=nvcr.io/nvidia/pytorch:25.12-py3 检查模型目录： ```bash -test -d "$PHYAI_CKPT_HOST" +test -d "$OPENPI_CKPT_HOST" +test -f "$PHYAI_CKPT_HOST/model.safetensors" +test -f "$PHYAI_CKPT_HOST/config.json" +test -f "$PHYAI_CKPT_HOST/policy_preprocessor.json" +test -f "$PHYAI_CKPT_HOST/policy_postprocessor.json" test -d "$TOKENIZER_HOST" ``` diff --git a/tools/convert_openpi_pi05_to_phyai.py b/tools/convert_openpi_pi05_to_phyai.py new file mode 100644 index 0000000..8f671b3 --- /dev/null +++ b/tools/convert_openpi_pi05_to_phyai.py @@ -0,0 +1,856 @@ +# /// script +# requires-python = ">=3.11,<3.12" +# dependencies = [ +# "jax[cpu]==0.5.3", +# "orbax-checkpoint==0.11.13", +# "numpy==2.2.6", +# "torch==2.7.1", +# "safetensors==0.4.5", +# ] +# /// +"""Convert an OpenPI JAX/Orbax pi0.5 checkpoint to PhyAI native safetensors. + +Source : an OpenPI Orbax/OCDBT ``pi05_libero`` checkpoint directory. +Target : a PhyAI checkpoint directory containing safetensors plus sidecar + config/processor files that ``phyai.weights.load_pretrained`` + consumes for the native pi0.5 engine (see ``modeling_pi05.py``). + +Run with uv so the PEP 723 header installs the pinned conversion dependencies:: + + uv run tools/convert_openpi_pi05_to_phyai.py \ + --checkpoint /path/to/pi05_libero + uv run tools/convert_openpi_pi05_to_phyai.py \ + --checkpoint /path/to/pi05_libero \ + --write \ + --out /path/to/pi05_libero_phyai/model.safetensors + +The default mode is **dry-run**: it loads the source State, builds the full +target tensor dict, and validates its key count and representative shapes. It +writes nothing. Pass ``--reference`` to additionally diff every key and shape +against an existing safetensors file. ``--write`` writes model.safetensors and +generates the PhyAI config/processor sidecars from the official OpenPI +``assets/physical-intelligence/libero/norm_stats.json`` file. + +Conventions confirmed against PhyAI source (read-only) — do NOT change without +re-verifying: + * All PyTorch linears are ``[out, in]`` -> every JAX ``kernel [in, out]`` is + transposed on copy. + * The checkpoint stores SEPARATE q/k/v/o and gate/up/down keys; PhyAI fuses + them at load time. Do NOT pre-fuse here. + * Gemma RMSNorm / AdaRMS use the ``(1 + scale)`` runtime convention on BOTH + sides (JAX stores raw scale, PhyAI's kernel adds 1) -> copy norm scale + verbatim, no ``+1``. + * ``gemma_expert.lm_head.weight`` is a dead weight PhyAI drops at load; the + reference product stores it as ``paligemma.lm_head.weight[:, :1024]``. We + reproduce that for byte-for-byte structural parity. + +Open correctness points that SHAPE validation cannot catch (verify with a +fixed-noise single-step OpenPI-vs-PhyAI action diff after a successful dry-run): + 1. q/k/v einsum axis order (head-major vs hidden-major). + 2. AdaRMS Dense_0 chunk order -> (scale, shift, gate). + 3. gating_einsum gate/up split order (index 0 = gate, 1 = up). +""" + +from __future__ import annotations + +import argparse +import json +import os +import struct +import sys +import tempfile +from pathlib import Path + +import numpy as np + +# --------------------------------------------------------------------------- +# Static dimensions (from modeling_pi05 / configuration_pi05, confirmed against +# the reference product header). Used only to assert shapes — never to fabricate +# data; every tensor's bytes come from the loaded OpenPI State. +# --------------------------------------------------------------------------- +VISION_LAYERS = 27 +GEMMA_LAYERS = 18 # shared count for paligemma text and the action expert + +VIS_HIDDEN = 1152 +VIS_HEADS = 16 +VIS_HEAD_DIM = 72 # 16 * 72 == 1152 +VIS_MLP = 4304 +VIS_PATCH = 14 +VIS_POS = 256 + +TXT_HIDDEN = 2048 +TXT_HEADS = 8 +TXT_HEAD_DIM = 256 +TXT_KV_HEADS = 1 +TXT_MLP = 16384 +VOCAB = 257152 + +EXP_HIDDEN = 1024 +EXP_MLP = 4096 +EXP_ADARMS = 3072 # 3 * 1024 (scale|shift|gate) +JOINT_ATTN = TXT_HEADS * TXT_HEAD_DIM # 2048, shared q/o joint space +KV_DIM = TXT_KV_HEADS * TXT_HEAD_DIM # 256 + +# Target key prefixes (every product key is prefixed with ``model.``). +ROOT = "model." +PALI = ROOT + "paligemma_with_expert.paligemma.model" +EXPERT = ROOT + "paligemma_with_expert.gemma_expert.model" +VIS = PALI + ".vision_tower.vision_model" +PALI_LM_HEAD = ROOT + "paligemma_with_expert.paligemma.lm_head.weight" +EXPERT_LM_HEAD = ROOT + "paligemma_with_expert.gemma_expert.lm_head.weight" + + +# =========================================================================== +# Reference safetensors header reader (shapes only, no tensor bytes). +# =========================================================================== +def read_safetensors_header(path: Path) -> dict[str, dict]: + """Return ``{key: {"shape": [...], "dtype": "..."}}`` without mmapping data.""" + with open(path, "rb") as f: + (n,) = struct.unpack(" dict[tuple, np.ndarray]: + """Load the OpenPI Orbax checkpoint into a flat path->ndarray dict. + + ``checkpoint_dir`` points at the checkpoint root that contains ``params/``. + We restore the full State + via orbax so every tensor is materialised at its true (unsharded) shape. + + OpenPI's public ``pi05_libero`` checkpoint was saved from an 8-device mesh. + A plain Orbax restore may fail on a 1-device conversion machine with + ``sharding ... Got None``. To make the converter portable, we inspect the + checkpoint metadata and request every array to be restored onto the local + first device with ``SingleDeviceSharding`` while preserving each array's + global shape. + """ + import jax + import orbax.checkpoint as ocp + + params_dir = checkpoint_dir / "params" + if not params_dir.exists(): + raise FileNotFoundError(f"no params/ under {checkpoint_dir}") + + ckptr = ocp.PyTreeCheckpointer() + restored = ckptr.restore( + params_dir.resolve(), + args=ocp.args.PyTreeRestore( + restore_args=_single_device_restore_args(ckptr, params_dir, jax, ocp) + ), + ) + + flat: dict[tuple, np.ndarray] = {} + + def _walk(node, prefix: tuple) -> None: + if isinstance(node, dict): + for k, v in node.items(): + _walk(v, prefix + (k,)) + else: + flat[prefix] = np.asarray(node) + + _walk(restored, ()) + return flat + + +def _single_device_restore_args(ckptr, params_dir: Path, jax, ocp): + """Build a restore_args tree that ignores the saved 8-device mesh. + + Orbax metadata leaves expose global ``shape`` / ``dtype``. Mirroring the + metadata tree with ``ArrayRestoreArgs`` lets TensorStore assemble the full + array onto one local device instead of requiring the original save mesh. + """ + + sharding = jax.sharding.SingleDeviceSharding(jax.devices()[0]) + metadata_obj = ckptr.metadata(params_dir.resolve()) + item_metadata = getattr(metadata_obj, "item_metadata", metadata_obj) + metadata = item_metadata.tree + + def _convert(node): + if isinstance(node, dict): + return {k: _convert(v) for k, v in node.items()} + shape = tuple(getattr(node, "shape")) + dtype = getattr(node, "dtype", None) + return ocp.ArrayRestoreArgs( + restore_type=jax.Array, + dtype=dtype, + sharding=sharding, + global_shape=shape, + shape=shape, + strict=False, + ) + + return _convert(metadata) + + +class Source: + """Path-addressable accessor over the flat OpenPI State. + + Keys in the JAX tree end with a trailing ``value`` leaf (as seen in + ``_METADATA``). We normalise lookups so callers pass the human path + (e.g. ``"PaliGemma/img/embedding/kernel"``) and we try both with and + without the trailing ``value``. + """ + + def __init__(self, flat: dict[tuple, np.ndarray]) -> None: + self._flat = flat + # Index by "/"-joined path for both the raw and value-stripped forms. + self._by_str: dict[str, np.ndarray] = {} + for path, arr in flat.items(): + joined = "/".join(str(p) for p in path) + self._by_str[joined] = arr + if path and path[-1] == "value": + self._by_str["/".join(str(p) for p in path[:-1])] = arr + + def get(self, path: str) -> np.ndarray: + key = path.strip("/") + if key in self._by_str: + return self._by_str[key] + # Try the ``params/`` prefix and a trailing ``value`` leaf. + for cand in (f"params/{key}", f"{key}/value", f"params/{key}/value"): + if cand in self._by_str: + return self._by_str[cand] + raise KeyError( + f"source leaf not found: {path!r}\n" + f"available (sample): {sorted(self._by_str)[:8]}" + ) + + +# =========================================================================== +# Transform helpers. +# =========================================================================== +def kT(arr: np.ndarray) -> np.ndarray: + """JAX Linear kernel [in, out] -> PyTorch weight [out, in].""" + return np.ascontiguousarray(np.swapaxes(arr, -1, -2)) + + +# =========================================================================== +# Target tensor-dict builder. Produces ``{target_key: np.ndarray}`` for all +# 812 keys. Layer tensors are unstacked from the leading (stacked) axis. +# =========================================================================== +def build_target(src: Source) -> dict[str, np.ndarray]: + out: dict[str, np.ndarray] = {} + + # ---- action / time heads (root) -------------------------------------- + out[ROOT + "action_in_proj.weight"] = kT(src.get("action_in_proj/kernel")) + out[ROOT + "action_in_proj.bias"] = src.get("action_in_proj/bias") + out[ROOT + "action_out_proj.weight"] = kT(src.get("action_out_proj/kernel")) + out[ROOT + "action_out_proj.bias"] = src.get("action_out_proj/bias") + out[ROOT + "time_mlp_in.weight"] = kT(src.get("time_mlp_in/kernel")) + out[ROOT + "time_mlp_in.bias"] = src.get("time_mlp_in/bias") + out[ROOT + "time_mlp_out.weight"] = kT(src.get("time_mlp_out/kernel")) + out[ROOT + "time_mlp_out.bias"] = src.get("time_mlp_out/bias") + + # ---- vision tower (SigLIP) ------------------------------------------- + _build_vision(src, out) + + # ---- paligemma text LM + tied lm_head -------------------------------- + _build_text(src, out) + + # ---- action expert --------------------------------------------------- + _build_expert(src, out) + + return out + + +def _build_vision(src: Source, out: dict[str, np.ndarray]) -> None: + # patch embedding conv: JAX HWIO [14,14,3,1152] -> torch OIHW [1152,3,14,14] + conv = src.get("PaliGemma/img/embedding/kernel") + out[VIS + ".embeddings.patch_embedding.weight"] = np.ascontiguousarray( + np.transpose(conv, (3, 2, 0, 1)) + ) + out[VIS + ".embeddings.patch_embedding.bias"] = src.get( + "PaliGemma/img/embedding/bias" + ) + # position embedding: JAX [1,256,1152] -> [256,1152] + pos = src.get("PaliGemma/img/pos_embedding") + out[VIS + ".embeddings.position_embedding.weight"] = np.ascontiguousarray(pos[0]) + # final encoder norm -> post_layernorm + out[VIS + ".post_layernorm.weight"] = src.get( + "PaliGemma/img/Transformer/encoder_norm/scale" + ) + out[VIS + ".post_layernorm.bias"] = src.get( + "PaliGemma/img/Transformer/encoder_norm/bias" + ) + # multi_modal_projector (JAX img/head) + out[PALI + ".multi_modal_projector.linear.weight"] = kT( + src.get("PaliGemma/img/head/kernel") + ) + out[PALI + ".multi_modal_projector.linear.bias"] = src.get( + "PaliGemma/img/head/bias" + ) + + eb = "PaliGemma/img/Transformer/encoderblock" + ln0_s = src.get(f"{eb}/LayerNorm_0/scale") + ln0_b = src.get(f"{eb}/LayerNorm_0/bias") + ln1_s = src.get(f"{eb}/LayerNorm_1/scale") + ln1_b = src.get(f"{eb}/LayerNorm_1/bias") + q_k = src.get(f"{eb}/MultiHeadDotProductAttention_0/query/kernel") + q_b = src.get(f"{eb}/MultiHeadDotProductAttention_0/query/bias") + k_k = src.get(f"{eb}/MultiHeadDotProductAttention_0/key/kernel") + k_b = src.get(f"{eb}/MultiHeadDotProductAttention_0/key/bias") + v_k = src.get(f"{eb}/MultiHeadDotProductAttention_0/value/kernel") + v_b = src.get(f"{eb}/MultiHeadDotProductAttention_0/value/bias") + o_k = src.get(f"{eb}/MultiHeadDotProductAttention_0/out/kernel") + o_b = src.get(f"{eb}/MultiHeadDotProductAttention_0/out/bias") + fc1_k = src.get(f"{eb}/MlpBlock_0/Dense_0/kernel") + fc1_b = src.get(f"{eb}/MlpBlock_0/Dense_0/bias") + fc2_k = src.get(f"{eb}/MlpBlock_0/Dense_1/kernel") + fc2_b = src.get(f"{eb}/MlpBlock_0/Dense_1/bias") + + for i in range(VISION_LAYERS): + p = f"{VIS}.encoder.layers.{i}." + out[p + "layer_norm1.weight"] = ln0_s[i] + out[p + "layer_norm1.bias"] = ln0_b[i] + out[p + "layer_norm2.weight"] = ln1_s[i] + out[p + "layer_norm2.bias"] = ln1_b[i] + # attention: per-head kernel [in=1152, heads=16, head_dim=72] -> + # flatten head-major to [1152, 1152] then transpose to [out, in]. + out[p + "self_attn.q_proj.weight"] = _vis_qkv_w(q_k[i]) + out[p + "self_attn.q_proj.bias"] = q_b[i].reshape(VIS_HIDDEN) + out[p + "self_attn.k_proj.weight"] = _vis_qkv_w(k_k[i]) + out[p + "self_attn.k_proj.bias"] = k_b[i].reshape(VIS_HIDDEN) + out[p + "self_attn.v_proj.weight"] = _vis_qkv_w(v_k[i]) + out[p + "self_attn.v_proj.bias"] = v_b[i].reshape(VIS_HIDDEN) + # out: JAX [heads=16, head_dim=72, out=1152] -> [in=1152, out=1152] -> T + out[p + "self_attn.out_proj.weight"] = np.ascontiguousarray( + o_k[i].reshape(VIS_HEADS * VIS_HEAD_DIM, VIS_HIDDEN).T + ) + out[p + "self_attn.out_proj.bias"] = o_b[i] + out[p + "mlp.fc1.weight"] = kT(fc1_k[i]) + out[p + "mlp.fc1.bias"] = fc1_b[i] + out[p + "mlp.fc2.weight"] = kT(fc2_k[i]) + out[p + "mlp.fc2.bias"] = fc2_b[i] + + +def _vis_qkv_w(per_layer: np.ndarray) -> np.ndarray: + """SigLIP q/k/v kernel [in=1152, heads=16, head_dim=72] -> weight [1152,1152].""" + flat = per_layer.reshape(VIS_HIDDEN, VIS_HEADS * VIS_HEAD_DIM) # [in, out] + return np.ascontiguousarray(flat.T) # [out, in] + + +def _build_text(src: Source, out: dict[str, np.ndarray]) -> None: + # tied embedder -> paligemma.lm_head (full vocab) + embed = src.get("PaliGemma/llm/embedder/input_embedding") # [257152, 2048] + out[PALI_LM_HEAD] = embed + # final norm + out[PALI + ".language_model.norm.weight"] = src.get( + "PaliGemma/llm/final_norm/scale" + ) + + q = src.get("PaliGemma/llm/layers/attn/q_einsum/w") # [18, 8, 256, 2048] + kv = src.get("PaliGemma/llm/layers/attn/kv_einsum/w") # [18, 2, 1, 256, 2048] + o = src.get("PaliGemma/llm/layers/attn/attn_vec_einsum/w") # [18, 8, 256, 2048] + gate_up = src.get("PaliGemma/llm/layers/mlp/gating_einsum") # [18, 2, 2048, 16384] + down = src.get("PaliGemma/llm/layers/mlp/linear") # [18, 16384, 2048] + pre_attn = src.get("PaliGemma/llm/layers/pre_attention_norm/scale") # [18, 2048] + pre_ffw = src.get("PaliGemma/llm/layers/pre_ffw_norm/scale") # [18, 2048] + + for i in range(GEMMA_LAYERS): + p = f"{PALI}.language_model.layers.{i}." + out[p + "input_layernorm.weight"] = pre_attn[i] + out[p + "post_attention_layernorm.weight"] = pre_ffw[i] + # q: [heads=8, head_dim=256, hidden=2048] -> [2048, 2048] head-major + out[p + "self_attn.q_proj.weight"] = _gemma_q_w(q[i], TXT_HIDDEN) + # kv: [2, 1, head_dim=256, hidden] -> k=idx0, v=idx1 -> [256, hidden] + out[p + "self_attn.k_proj.weight"] = _gemma_kv_w(kv[i, 0], TXT_HIDDEN) + out[p + "self_attn.v_proj.weight"] = _gemma_kv_w(kv[i, 1], TXT_HIDDEN) + # o: [heads=8, head_dim=256, out=2048] -> [in=2048, out=2048] -> T + out[p + "self_attn.o_proj.weight"] = _gemma_o_w(o[i], TXT_HIDDEN) + # mlp gate/up: gating_einsum [2, in, inter] -> gate=idx0, up=idx1 + out[p + "mlp.gate_proj.weight"] = kT(gate_up[i, 0]) + out[p + "mlp.up_proj.weight"] = kT(gate_up[i, 1]) + out[p + "mlp.down_proj.weight"] = kT(down[i]) + + +def _build_expert(src: Source, out: dict[str, np.ndarray]) -> None: + # expert lm_head: dead weight PhyAI drops; product stores paligemma[:, :1024] + embed = src.get("PaliGemma/llm/embedder/input_embedding") # [257152, 2048] + out[EXPERT_LM_HEAD] = np.ascontiguousarray(embed[:, :EXP_HIDDEN]) + + # final norm is AdaRMS (Dense_0) + fn_k = src.get("PaliGemma/llm/final_norm_1/Dense_0/kernel") # [1024, 3072] + fn_b = src.get("PaliGemma/llm/final_norm_1/Dense_0/bias") # [3072] + out[EXPERT + ".norm.dense.weight"] = kT(fn_k) + out[EXPERT + ".norm.dense.bias"] = fn_b + + q = src.get("PaliGemma/llm/layers/attn/q_einsum_1/w") # [18, 8, 256, 1024] + kv = src.get("PaliGemma/llm/layers/attn/kv_einsum_1/w") # [18, 2, 1, 256, 1024] + o = src.get("PaliGemma/llm/layers/attn/attn_vec_einsum_1/w") # [18, 8, 256, 1024] + gate_up = src.get("PaliGemma/llm/layers/mlp_1/gating_einsum") # [18, 2, 1024, 4096] + down = src.get("PaliGemma/llm/layers/mlp_1/linear") # [18, 4096, 1024] + pre_attn_k = src.get( + "PaliGemma/llm/layers/pre_attention_norm_1/Dense_0/kernel" + ) # [18, 1024, 3072] + pre_attn_b = src.get("PaliGemma/llm/layers/pre_attention_norm_1/Dense_0/bias") + pre_ffw_k = src.get("PaliGemma/llm/layers/pre_ffw_norm_1/Dense_0/kernel") + pre_ffw_b = src.get("PaliGemma/llm/layers/pre_ffw_norm_1/Dense_0/bias") + + for i in range(GEMMA_LAYERS): + p = f"{EXPERT}.layers.{i}." + out[p + "input_layernorm.dense.weight"] = kT(pre_attn_k[i]) + out[p + "input_layernorm.dense.bias"] = pre_attn_b[i] + out[p + "post_attention_layernorm.dense.weight"] = kT(pre_ffw_k[i]) + out[p + "post_attention_layernorm.dense.bias"] = pre_ffw_b[i] + # q: [heads=8, head_dim=256, hidden=1024] -> [2048, 1024] head-major + out[p + "self_attn.q_proj.weight"] = _gemma_q_w(q[i], EXP_HIDDEN) + out[p + "self_attn.k_proj.weight"] = _gemma_kv_w(kv[i, 0], EXP_HIDDEN) + out[p + "self_attn.v_proj.weight"] = _gemma_kv_w(kv[i, 1], EXP_HIDDEN) + # o: ASYMMETRIC [heads=8, head_dim=256, out=1024] -> [in=2048, out=1024] -> T + out[p + "self_attn.o_proj.weight"] = _gemma_o_w(o[i], EXP_HIDDEN) + out[p + "mlp.gate_proj.weight"] = kT(gate_up[i, 0]) + out[p + "mlp.up_proj.weight"] = kT(gate_up[i, 1]) + out[p + "mlp.down_proj.weight"] = kT(down[i]) + + +def _gemma_q_w(per_layer: np.ndarray, in_dim: int) -> np.ndarray: + """q_einsum -> PyTorch weight ``[heads * head_dim, in]``. + + OpenPI checkpoints have appeared with both ``[heads, head_dim, in]`` and + ``[heads, in, head_dim]`` metadata layouts across versions. Normalize to + head-major ``[heads, head_dim, in]`` before flattening. + """ + heads, a, b = per_layer.shape + if b == in_dim: + head_dim = a + normalized = per_layer + elif a == in_dim: + head_dim = b + normalized = np.swapaxes(per_layer, 1, 2) + else: + raise AssertionError( + f"q shape {per_layer.shape} incompatible with in_dim={in_dim}" + ) + return np.ascontiguousarray(normalized.reshape(heads * head_dim, in_dim)) + + +def _gemma_kv_w(per_kv: np.ndarray, in_dim: int) -> np.ndarray: + """kv_einsum slice -> PyTorch weight ``[head_dim, in]``. + + Accept both ``[kv_head, head_dim, in]`` and ``[kv_head, in, head_dim]``. + """ + kv_head, a, b = per_kv.shape + assert kv_head == TXT_KV_HEADS, per_kv.shape + if b == in_dim: + normalized = per_kv + head_dim = a + elif a == in_dim: + normalized = np.swapaxes(per_kv, 1, 2) + head_dim = b + else: + raise AssertionError( + f"kv shape {per_kv.shape} incompatible with in_dim={in_dim}" + ) + return np.ascontiguousarray(normalized.reshape(head_dim, in_dim)) + + +def _gemma_o_w(per_layer: np.ndarray, out_dim: int) -> np.ndarray: + """attn_vec_einsum [heads, head_dim, out] -> weight [out, heads*head_dim]. + + JAX o-projection contracts over (heads, head_dim) producing ``out``; the + leaf is ``[heads, head_dim, out_dim]``. PyTorch o_proj is ``[out, in]`` with + ``in = heads*head_dim`` -> reshape to ``[in, out]`` then transpose. + """ + heads, head_dim, jout = per_layer.shape + assert jout == out_dim, f"o out-dim {jout} != {out_dim}" + flat = per_layer.reshape(heads * head_dim, out_dim) # [in, out] + return np.ascontiguousarray(flat.T) # [out, in] + + +# =========================================================================== +# Diff (dry-run) and write. +# =========================================================================== +NP_TO_ST = { + "float32": "F32", + "float16": "F16", + "bfloat16": "BF16", # only reached after torch cast; numpy has no bf16 +} + + +def diff_against_reference(built: dict[str, np.ndarray], reference: Path) -> int: + ref = read_safetensors_header(reference) + ref_keys = set(ref) + built_keys = set(built) + + missing = sorted(ref_keys - built_keys) # in product, not produced + extra = sorted(built_keys - ref_keys) # produced, not in product + mismatched = [] + for k in sorted(ref_keys & built_keys): + rshape = list(ref[k]["shape"]) + bshape = list(built[k].shape) + if rshape != bshape: + mismatched.append((k, bshape, rshape)) + + print(f"reference keys : {len(ref_keys)}") + print(f"built keys : {len(built_keys)}") + print(f"missing : {len(missing)}") + print(f"extra : {len(extra)}") + print(f"mismatched : {len(mismatched)}") + + def _show(title, items, fmt): + if items: + print(f"\n-- {title} (first 20) --") + for it in items[:20]: + print(" " + fmt(it)) + + _show("MISSING", missing, lambda k: k) + _show("EXTRA", extra, lambda k: k) + _show( + "MISMATCHED", + mismatched, + lambda t: f"{t[0]} built={t[1]} ref={t[2]}", + ) + + ok = not missing and not extra and not mismatched + print(f"\nRESULT: {'OK — shapes match' if ok else 'FAIL — see above'}") + return 0 if ok else 1 + + +def validate_target_structure(built: dict[str, np.ndarray]) -> int: + """Validate the fixed pi05_libero target structure without a reference file.""" + expected = { + PALI_LM_HEAD: (VOCAB, TXT_HIDDEN), + EXPERT_LM_HEAD: (VOCAB, EXP_HIDDEN), + ROOT + "action_in_proj.weight": (EXP_HIDDEN, 32), + ROOT + "action_out_proj.weight": (32, EXP_HIDDEN), + f"{VIS}.embeddings.patch_embedding.weight": ( + VIS_HIDDEN, + 3, + VIS_PATCH, + VIS_PATCH, + ), + } + errors = [] + if len(built) != 812: + errors.append(f"expected 812 tensors, built {len(built)}") + for key, shape in expected.items(): + actual = built.get(key) + if actual is None: + errors.append(f"missing {key}") + elif actual.shape != shape: + errors.append(f"{key}: expected {shape}, got {actual.shape}") + + if errors: + print("RESULT: FAIL - invalid target structure") + for error in errors: + print(f" {error}") + return 1 + print("RESULT: OK - 812 tensors and representative shapes match") + return 0 + + +def write_safetensors(built: dict[str, np.ndarray], out_path: Path, dtype: str) -> None: + import torch + from safetensors.torch import save_file + + torch_dtype = { + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "fp16": torch.float16, + "float16": torch.float16, + "fp32": torch.float32, + "float32": torch.float32, + }[dtype] + + tensors: dict[str, torch.Tensor] = {} + for k, arr in built.items(): + t = torch.from_numpy(np.ascontiguousarray(arr)) + # F32 params (norms, action/time heads) stay float32 to match product; + # everything else casts to the requested dtype. + if t.dtype == torch.float32 and _keep_fp32(k): + tensors[k] = t.contiguous() + else: + tensors[k] = t.to(torch_dtype).contiguous() + + out_path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=str(out_path.parent), suffix=".tmp") + os.close(fd) + tmp_path = Path(tmp) + try: + save_file(tensors, tmp, metadata={"format": "pt"}) + os.replace(tmp, out_path) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + print(f"wrote {len(tensors)} tensors -> {out_path}") + + +def _keep_fp32(key: str) -> bool: + """Keys the reference product stores as float32 (norms + action/time heads).""" + return ( + key.endswith(".dense.weight") + or key.endswith(".dense.bias") + or "patch_embedding" in key + or "position_embedding" in key + or "language_model.norm.weight" in key + or "action_in_proj" in key + or "action_out_proj" in key + or "time_mlp_in" in key + or "time_mlp_out" in key + ) + + +# =========================================================================== +# Complete checkpoint sidecars (config + processor stats). +# =========================================================================== +SIDECAR_FILES = ( + "config.json", + "policy_preprocessor.json", + "policy_postprocessor.json", + "policy_preprocessor_step_2_normalizer_processor.safetensors", + "policy_postprocessor_step_0_unnormalizer_processor.safetensors", +) + + +DEFAULT_NORM_STATS = Path("assets/physical-intelligence/libero/norm_stats.json") +TOKENIZER_NAME = "google/paligemma-3b-pt-224" + + +def _processor_stats(norm_stats_path: Path) -> dict[str, "torch.Tensor"]: + import torch + + with norm_stats_path.open("r", encoding="utf-8") as f: + raw = json.load(f).get("norm_stats", {}) + required = {"state", "actions"} + if not required.issubset(raw): + raise ValueError( + f"{norm_stats_path} must contain norm_stats.state and " + "norm_stats.actions" + ) + + tensors: dict[str, torch.Tensor] = {} + for source_name, target_name in ( + ("state", "observation.state"), + ("actions", "action"), + ): + source = raw[source_name] + for source_stat, target_stat in ( + ("mean", "mean"), + ("std", "std"), + ("q01", "min"), + ("q99", "max"), + ): + if source_stat not in source: + raise ValueError( + f"{norm_stats_path}: {source_name}.{source_stat} is missing" + ) + tensors[f"{target_name}.{target_stat}"] = torch.tensor( + source[source_stat], dtype=torch.float32 + ) + return tensors + + +def _write_json(path: Path, value: dict) -> None: + with path.open("w", encoding="utf-8") as f: + json.dump(value, f, indent=2) + f.write("\n") + + +def write_checkpoint_sidecars( + out_dir: Path, checkpoint_dir: Path, norm_stats_path: Path | None = None +) -> None: + """Generate the config and processor sidecars required by PhyAI.""" + from safetensors.torch import save_file + + stats_path = norm_stats_path or checkpoint_dir / DEFAULT_NORM_STATS + if not stats_path.exists(): + raise FileNotFoundError( + f"normalization stats not found at {stats_path}; pass --norm-stats" + ) + stats = _processor_stats(stats_path) + + features = { + "observation.state": {"type": "STATE", "shape": [8]}, + "action": {"type": "ACTION", "shape": [7]}, + } + norm_map = { + "ACTION": "QUANTILES", + "STATE": "QUANTILES", + "VISUAL": "IDENTITY", + } + config = { + "type": "pi05", + "input_features": {"observation.state": features["observation.state"]}, + "output_features": {"action": features["action"]}, + "chunk_size": 10, + "max_state_dim": 32, + "max_action_dim": 32, + "num_inference_steps": 10, + "image_resolution": [224, 224], + "tokenizer_max_length": 200, + "tokenizer_name": TOKENIZER_NAME, + "phyai_prompt_mode": "openpi_task", + "phyai_normalization_mode": "openpi_quantile", + "phyai_image_mask_mode": "openpi_libero", + } + preprocessor = { + "name": "policy_preprocessor", + "steps": [ + { + "registry_name": "normalizer_processor", + "config": { + "eps": 1e-8, + "features": features, + "norm_map": norm_map, + }, + "state_file": SIDECAR_FILES[3], + }, + { + "registry_name": "pi05_prepare_state_tokenizer_processor_step", + "config": {}, + }, + { + "registry_name": "tokenizer_processor", + "config": { + "max_length": 200, + "task_key": "task", + "padding_side": "right", + "padding": "max_length", + "truncation": True, + "tokenizer_name": TOKENIZER_NAME, + }, + }, + { + "registry_name": "device_processor", + "config": {"device": "cuda", "float_dtype": None}, + }, + ], + } + postprocessor = { + "name": "policy_postprocessor", + "steps": [ + { + "registry_name": "unnormalizer_processor", + "config": { + "eps": 1e-8, + "features": {"action": features["action"]}, + "norm_map": norm_map, + }, + "state_file": SIDECAR_FILES[4], + }, + { + "registry_name": "device_processor", + "config": {"device": "cpu", "float_dtype": None}, + }, + ], + } + + out_dir.mkdir(parents=True, exist_ok=True) + _write_json(out_dir / SIDECAR_FILES[0], config) + _write_json(out_dir / SIDECAR_FILES[1], preprocessor) + _write_json(out_dir / SIDECAR_FILES[2], postprocessor) + save_file(stats, str(out_dir / SIDECAR_FILES[3])) + save_file(stats, str(out_dir / SIDECAR_FILES[4])) + print(f"generated {len(SIDECAR_FILES)} sidecar files -> {out_dir}") + + +# =========================================================================== +# CLI. +# =========================================================================== +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--checkpoint", + type=Path, + required=True, + help="OpenPI checkpoint root containing params/", + ) + ap.add_argument( + "--reference", + type=Path, + help="optional reference safetensors for a full key/shape diff", + ) + ap.add_argument( + "--norm-stats", + type=Path, + help=( + "OpenPI norm_stats.json override; defaults to the pi05_libero asset " + "under --checkpoint" + ), + ) + ap.add_argument( + "--write", + action="store_true", + help="actually write the converted checkpoint (default: dry-run only)", + ) + ap.add_argument( + "--out", + type=Path, + help="output safetensors path for --write", + ) + ap.add_argument( + "--allow-overwrite", + action="store_true", + help="required to let --out point at an existing file (e.g. the product)", + ) + ap.add_argument("--dtype", default="bf16", help="write dtype for non-fp32 params") + args = ap.parse_args(argv) + + if args.write and args.out is None: + ap.error("--out is required when --write is used") + + # Header reader works without OpenPI; load source only when needed (always, + # since both dry-run and write build the target dict). + print(f"loading OpenPI State from {args.checkpoint} ...", file=sys.stderr) + flat = load_openpi_state(args.checkpoint) + src = Source(flat) + print(f" loaded {len(flat)} source leaves", file=sys.stderr) + + built = build_target(src) + + # Sanity: embedder must be full vocab, not a shard. + embed_shape = built[PALI_LM_HEAD].shape + assert tuple(embed_shape) == (VOCAB, TXT_HIDDEN), ( + f"embedder shape {embed_shape} != ({VOCAB}, {TXT_HIDDEN}); " + "the OpenPI State looks sharded - check the Orbax restore and " + "checkpoint format." + ) + + if args.reference is None: + rc = validate_target_structure(built) + else: + rc = diff_against_reference(built, args.reference) + + if not args.write: + print("\n(dry-run — nothing written; pass --write to produce a checkpoint)") + return rc + + if rc != 0: + print("\nrefusing to --write: shape diff is non-empty (fix mapping first)") + return rc + + out = args.out + assert out is not None + if ( + args.reference is not None + and out.resolve() == args.reference.resolve() + and not args.allow_overwrite + ): + print( + f"\nrefusing to overwrite the reference product at {out};\n" + "pass --allow-overwrite to override (NOT recommended)." + ) + return 2 + existing = [out, *(out.parent / name for name in SIDECAR_FILES)] + existing = [path for path in existing if path.exists()] + if existing and not args.allow_overwrite: + print( + "\nrefusing to overwrite existing output files:\n " + + "\n ".join(str(path) for path in existing) + + "\nchoose a new --out or pass --allow-overwrite." + ) + return 2 + + write_safetensors(built, out, args.dtype) + write_checkpoint_sidecars(out.parent, args.checkpoint, args.norm_stats) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 4b94ae9c83f086ec7cd2aaf8cafd8f923e49a7bd Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 21:30:35 +0800 Subject: [PATCH 28/29] docs(phyai): pin checkpoint conversion workflow --- docs/models/pi05/libero-four-suites.mdx | 6 +++--- docs/zh/models/pi05/libero-four-suites.mdx | 6 +++--- tools/convert_openpi_pi05_to_phyai.py | 3 +-- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/docs/models/pi05/libero-four-suites.mdx b/docs/models/pi05/libero-four-suites.mdx index f738bdd..1a8b51a 100644 --- a/docs/models/pi05/libero-four-suites.mdx +++ b/docs/models/pi05/libero-four-suites.mdx @@ -49,7 +49,7 @@ This reproduction uses immutable commits from two linked pull requests: | Repository | Pull request | Commit | | --- | --- | --- | -| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`5fc3288`](https://github.com/rebecca26358/phyai/commit/5fc3288dcd8fe1de0a3859a28cfa7036e42217d8) | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`3a9a28b`](https://github.com/rebecca26358/phyai/commit/3a9a28b21aa53627a07c3be9fcb86e6adb1b4ae8) | | `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`8bb9009`](https://github.com/rebecca26358/vla-evaluation-harness/commit/8bb90099c719ee8d151497c6636e76a9da5d5c6f) | Check out those exact revisions: @@ -59,7 +59,7 @@ export PHYAI_ROOT=${PHYAI_ROOT:-$HOME/phyai} export VLA_ROOT=${VLA_ROOT:-$HOME/vla-evaluation-harness} git clone https://github.com/rebecca26358/phyai.git "$PHYAI_ROOT" -git -C "$PHYAI_ROOT" checkout 5fc3288dcd8fe1de0a3859a28cfa7036e42217d8 +git -C "$PHYAI_ROOT" checkout 3a9a28b21aa53627a07c3be9fcb86e6adb1b4ae8 git clone https://github.com/rebecca26358/vla-evaluation-harness.git "$VLA_ROOT" git -C "$VLA_ROOT" checkout 8bb90099c719ee8d151497c6636e76a9da5d5c6f @@ -68,7 +68,7 @@ git -C "$VLA_ROOT" checkout 8bb90099c719ee8d151497c6636e76a9da5d5c6f The VLA revision contains `src/vla_eval/model_servers/phyai.py`, `src/vla_eval/model_servers/phyai_wn.py`, and `configs/model_servers/phyai/libero.yaml`. Its PEP 723 metadata pins the policy -implementation commit [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e95665fae38e7d13a737922398eaadb4db24669), an ancestor of the PhyAI checkout above. The later PhyAI commits add documentation and tests without changing the policy implementation. +implementation commit [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e95665fae38e7d13a737922398eaadb4db24669), an ancestor of the PhyAI checkout above. The later PhyAI commits add documentation, tests, and the checkpoint converter without changing the policy implementation. # Download and Convert the Checkpoint diff --git a/docs/zh/models/pi05/libero-four-suites.mdx b/docs/zh/models/pi05/libero-four-suites.mdx index 5a6fd23..6c489c9 100644 --- a/docs/zh/models/pi05/libero-four-suites.mdx +++ b/docs/zh/models/pi05/libero-four-suites.mdx @@ -49,7 +49,7 @@ icon: "list-checks" | 仓库 | Pull request | 提交 | | --- | --- | --- | -| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`5fc3288`](https://github.com/rebecca26358/phyai/commit/5fc3288dcd8fe1de0a3859a28cfa7036e42217d8) | +| PhyAI | [mingti-org/phyai#21](https://github.com/mingti-org/phyai/pull/21) | [`3a9a28b`](https://github.com/rebecca26358/phyai/commit/3a9a28b21aa53627a07c3be9fcb86e6adb1b4ae8) | | `vla-evaluation-harness` | [allenai/vla-evaluation-harness#111](https://github.com/allenai/vla-evaluation-harness/pull/111) | [`8bb9009`](https://github.com/rebecca26358/vla-evaluation-harness/commit/8bb90099c719ee8d151497c6636e76a9da5d5c6f) | 检出对应版本： @@ -59,7 +59,7 @@ export PHYAI_ROOT=${PHYAI_ROOT:-$HOME/phyai} export VLA_ROOT=${VLA_ROOT:-$HOME/vla-evaluation-harness} git clone https://github.com/rebecca26358/phyai.git "$PHYAI_ROOT" -git -C "$PHYAI_ROOT" checkout 5fc3288dcd8fe1de0a3859a28cfa7036e42217d8 +git -C "$PHYAI_ROOT" checkout 3a9a28b21aa53627a07c3be9fcb86e6adb1b4ae8 git clone https://github.com/rebecca26358/vla-evaluation-harness.git "$VLA_ROOT" git -C "$VLA_ROOT" checkout 8bb90099c719ee8d151497c6636e76a9da5d5c6f @@ -68,7 +68,7 @@ git -C "$VLA_ROOT" checkout 8bb90099c719ee8d151497c6636e76a9da5d5c6f 这个 VLA 版本包含 `src/vla_eval/model_servers/phyai.py`、 `src/vla_eval/model_servers/phyai_wn.py` 和 `configs/model_servers/phyai/libero.yaml`。其 PEP 723 元数据固定到 policy -实现提交 [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e95665fae38e7d13a737922398eaadb4db24669)，它是上面 PhyAI checkout 的祖先；后续 PhyAI 提交只增加文档和测试，没有修改 policy 实现。 +实现提交 [`4e95665`](https://github.com/rebecca26358/phyai/commit/4e95665fae38e7d13a737922398eaadb4db24669)，它是上面 PhyAI checkout 的祖先；后续 PhyAI 提交增加了文档、测试和 checkpoint converter，没有修改 policy 实现。 # 下载并转换 checkpoint diff --git a/tools/convert_openpi_pi05_to_phyai.py b/tools/convert_openpi_pi05_to_phyai.py index 8f671b3..7bb8868 100644 --- a/tools/convert_openpi_pi05_to_phyai.py +++ b/tools/convert_openpi_pi05_to_phyai.py @@ -620,8 +620,7 @@ def _processor_stats(norm_stats_path: Path) -> dict[str, "torch.Tensor"]: required = {"state", "actions"} if not required.issubset(raw): raise ValueError( - f"{norm_stats_path} must contain norm_stats.state and " - "norm_stats.actions" + f"{norm_stats_path} must contain norm_stats.state and " "norm_stats.actions" ) tensors: dict[str, torch.Tensor] = {} From b275371c2e6496bcb7c3c92e2e72b470debf5534 Mon Sep 17 00:00:00 2001 From: rebecca26358 <108570141+rebecca26358@users.noreply.github.com> Date: Sun, 2 Aug 2026 13:41:40 +0000 Subject: [PATCH 29/29] benchmark: improve pi05 latency reproducibility --- benchmark/pi05/bench_realtime_vla_pi05.py | 2 +- benchmark/pi05/bench_vlacpp_pi05_client.py | 17 +- docs/models/pi05/external-runtime-latency.mdx | 280 ++++++++++-------- .../models/pi05/external-runtime-latency.mdx | 280 ++++++++++-------- 4 files changed, 310 insertions(+), 269 deletions(-) mode change 100755 => 100644 benchmark/pi05/bench_realtime_vla_pi05.py mode change 100755 => 100644 benchmark/pi05/bench_vlacpp_pi05_client.py diff --git a/benchmark/pi05/bench_realtime_vla_pi05.py b/benchmark/pi05/bench_realtime_vla_pi05.py old mode 100755 new mode 100644 index f91a6e4..8cb9217 --- a/benchmark/pi05/bench_realtime_vla_pi05.py +++ b/benchmark/pi05/bench_realtime_vla_pi05.py @@ -81,6 +81,7 @@ def setup_fn(batch_size: int) -> bnb.BenchSpec: if batch_size != 1: raise ValueError("realtime-vla PI0.5 wrapper supports only batch_size=1") + torch.manual_seed(args.seed) checkpoint = load_checkpoint( args.checkpoint, args.flashrt_root, args.trust_pickle_checkpoint ) @@ -97,7 +98,6 @@ def setup_fn(batch_size: int) -> bnb.BenchSpec: tokenizer_path=str(args.tokenizer) if args.tokenizer else None, discrete_state_input=args.discrete_state_input, ) - torch.manual_seed(args.seed) input_image = torch.randn( args.num_views, 224, 224, 3, dtype=torch.bfloat16, device="cuda" ) diff --git a/benchmark/pi05/bench_vlacpp_pi05_client.py b/benchmark/pi05/bench_vlacpp_pi05_client.py old mode 100755 new mode 100644 index b29924b..403fdd5 --- a/benchmark/pi05/bench_vlacpp_pi05_client.py +++ b/benchmark/pi05/bench_vlacpp_pi05_client.py @@ -35,8 +35,8 @@ _BENCHMARK_DIR = Path(__file__).resolve().parents[1] sys.path.insert(0, str(_BENCHMARK_DIR)) -import bench_n_batch as bnb -from phyai.utils.profile import ( +import bench_n_batch as bnb # noqa: E402 +from phyai.utils.profile import ( # noqa: E402 add_profile_cli_args, install_profiler, profile_config_from_args, @@ -121,11 +121,6 @@ def setup_fn(batch_size: int) -> bnb.BenchSpec: phase_summary: dict[str, Any] = {} call_count = 0 - def update_phase_summary() -> None: - phase_summary.clear() - for key, values in phase_samples.items(): - phase_summary[key] = summarize(values) - def step() -> None: nonlocal call_count client.get_action(obs) @@ -146,9 +141,11 @@ def step() -> None: phase_samples["server_denoise_latency_ms"].append( float(r.latency_ms_denoise) ) - update_phase_summary() def teardown() -> None: + phase_summary.clear() + for key, values in phase_samples.items(): + phase_summary[key] = summarize(values) sock = getattr(client, "sock", None) if sock is not None: sock.close(linger=0) @@ -158,8 +155,8 @@ def teardown() -> None: step_callable=step, teardown_callable=teardown, ) - # Attach dynamic summary for extras_fn. The runner copies the dict after - # timed steps finish, so it records the final server phase statistics. + # The runner writes JSONL after teardown, so the final summary is available + # without adding percentile calculations to the timed step. spec.vlacpp_phase_summary = phase_summary # type: ignore[attr-defined] return spec diff --git a/docs/models/pi05/external-runtime-latency.mdx b/docs/models/pi05/external-runtime-latency.mdx index bba75f5..a9d1b2e 100644 --- a/docs/models/pi05/external-runtime-latency.mdx +++ b/docs/models/pi05/external-runtime-latency.mdx @@ -1,196 +1,218 @@ --- -title: "PI0.5 External Runtime Latency" -description: "Run PI0.5 latency benchmarks for FlashRT, realtime-vla, and vla.cpp" +title: "PI0.5 Latency Benchmark" +description: "Reproduce native PhyAI PI0.5 latency and run optional external runtime comparisons" icon: "timer" --- # Overview -This page explains how to set up and run the three external PI0.5 latency wrappers under `benchmark/pi05/`. +This page uses the native PhyAI benchmark as the primary PI0.5 latency path: -These wrappers do not use the PhyAI engine for inference. Each wrapper calls the target runtime directly, while reusing PhyAI's common benchmark runner for warmup, timing, and JSONL output. +```text +benchmark/bench_n_batch_ws1_pi05.py +``` + +The script builds a PhyAI `Engine` with the `pi05` plugin, creates synthetic inputs, runs warmup, and records steady-state `Engine.step(request)` latency. It is a performance benchmark, not a LIBERO accuracy evaluation. -| Script | Runtime | Timed call | -| --- | --- | --- | -| `bench_flashrt_pi05.py` | FlashRT | `Pi05TorchFrontendRtx.infer(obs)` | -| `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward(...)` | -| `bench_vlacpp_pi05_client.py` | vla.cpp | One ZMQ request to a running `vla-server` | +The FlashRT, realtime-vla, and vla.cpp wrappers are listed at the end as optional comparison paths. -The wrappers generate synthetic image and state inputs. Use them for latency-only measurements, not LIBERO accuracy evaluation. +# What the PhyAI benchmark measures + +| Item | Behavior | +| --- | --- | +| Timed call | One `Engine.step(request)` | +| Input | Synthetic images and a one-token synthetic prompt | +| Included | Vision tower, language-model prefix, all 10 Euler expert steps, scheduler work, and action output | +| Excluded | Environment setup, checkpoint loading, engine construction, CUDA graph capture, request creation, and warmup | +| GPU timing | CUDA events around each step, followed by CUDA synchronization | +| Output | One JSONL row per batch size | -# Common Setup +The request is created on the target device before timing. Image decoding, tokenization, simulator communication, and host-to-device preprocessing are therefore outside this benchmark. -Use a Python environment that can import PhyAI, PyTorch, and the common benchmark runner. +# Prepare the environment + +From a clean checkout: ```bash -cd -python -c "import torch; import phyai; import benchmark.bench_n_batch" +git clone https://github.com/MEmbodied/phyai.git +cd phyai +uv sync +``` + +Confirm that the environment can import PhyAI and see the GPU: + +```bash +uv run python -c "import torch, phyai; print(torch.__version__, torch.version.cuda); print(torch.cuda.get_device_name(0))" nvidia-smi ``` -Keep these settings aligned when comparing runtimes: +Record the source revision and GPU state with each result: -| Item | Setting | -| --- | --- | -| Batch size | 1 | -| Camera streams | 2 | -| Chunk size | 50 | -| Prompt | Same prompt text across runs | -| Warmup / timed iterations | Same values across runs | -| Precision | Label each row by the runtime's real precision path | +```bash +git rev-parse HEAD +nvidia-smi --query-gpu=name,driver_version,memory.used,utilization.gpu --format=csv +nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv +``` -Use your own paths for these placeholders: +Run latency measurements only when the GPU and CPU are idle. -| Placeholder | Meaning | -| --- | --- | -| `` | PhyAI checkout containing `benchmark/pi05/` | -| `` | FlashRT checkout | -| `` | realtime-vla checkout | -| `` | vla.cpp checkout | -| `` | Compiled vla.cpp `vla-server` binary | -| `` | PI0.5 safetensors checkpoint directory or file | -| `` | PI0.5 GGUF file for vla.cpp | -| `` | vla.cpp multimodal projector GGUF | -| `` | Local tokenizer directory | -| `` | LIBERO `meta/stats.json` with `observation.state.q01/q99` | +# Prepare the checkpoint -# FlashRT +`--checkpoint` must point to one directory containing: + +```text +/ + config.json + model.safetensors +``` -Install FlashRT from its official repository, then make the checkout visible to the wrapper. +A sharded checkpoint may use `model.safetensors.index.json` and its shard files instead. + +Check the files before starting: ```bash -export FLASHRT_ROOT= -cd -python -c "import sys; sys.path.insert(0, '$FLASHRT_ROOT'); from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx; print(Pi05TorchFrontendRtx)" +test -f /config.json +test -f /model.safetensors || \ + test -f /model.safetensors.index.json ``` -Run latency: +The action chunk size and denoise-step count come from `config.json`. The standard PI0.5 configuration uses `chunk_size=50` and `num_inference_steps=10`; the benchmark has no separate `--chunk-size` flag. + +Use the same PI0.5 checkpoint family when comparing runtimes. Converted formats may be used when a runtime requires them, but they should originate from the same checkpoint. + +# Run the PhyAI BF16 benchmark + +The command below is the recommended single-batch, two-camera LIBERO-shaped latency run: ```bash cd -python benchmark/pi05/bench_flashrt_pi05.py \ - --flashrt-root \ +mkdir -p results + +uv run python benchmark/bench_n_batch_ws1_pi05.py \ --checkpoint \ - --precision bf16 \ - --num-views 2 \ - --chunk-size 50 \ + --dtype bf16 \ + --vision-dtype bfloat16 \ + --num-images 2 \ --batch-sizes 1 \ --n-warmup 100 \ --n-timed 100 \ - --result-file results/flashrt_pi05.jsonl + --run-name pi05_bf16_views2_bs1 \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_views2_bs1.jsonl ``` -Notes: +This configuration means: -| Item | Note | +| Argument | Value | | --- | --- | -| Chunk size | The wrapper uses FlashRT's direct `Pi05TorchFrontendRtx` API because `chunk_size` is a frontend constructor argument | -| Denoise steps | Do not use `load_model(..., num_steps=50)` to set action chunk size; in FlashRT, `num_steps` means denoise steps | -| Precision | `--precision bf16` uses FlashRT's forced-BF16 PI0.5 RTX path; `--precision fp8_bf16` is a separate optimized FP8/BF16 result | +| Engine precision | BF16 | +| Vision tower precision | BF16 | +| Batch size | 1 | +| Camera streams | 2 | +| Action chunk | Read from checkpoint, normally 50 | +| Euler steps | Read from checkpoint, normally 10 | +| CUDA graph | Enabled by default on CUDA | +| Prompt shape | One synthetic token | +| Warmup / timed iterations | 100 / 100 | + +The PI0.5 plugin also applies the runtime recommendations stored in its model configuration when the user has not overridden them. Record any `PHYAI_*` environment overrides because they can change backend selection or workspace settings. -# realtime-vla +## Three-camera PI0.5 base shape -Install realtime-vla from its official repository. The wrapper can load a converted `.pt` / `.pth` checkpoint directly. If you pass a PI0.5 safetensors checkpoint, also provide FlashRT so the wrapper can reuse FlashRT's conversion helper. +The public `pi05_base` contract uses three camera streams. Change only `--num-images` and the output label: ```bash -export REALTIME_VLA_ROOT= -export FLASHRT_ROOT= -cd -python -c "import sys; sys.path.insert(0, '$REALTIME_VLA_ROOT'); from pi05_infer import Pi05Inference; print(Pi05Inference)" +uv run python benchmark/bench_n_batch_ws1_pi05.py \ + --checkpoint \ + --dtype bf16 \ + --vision-dtype bfloat16 \ + --num-images 3 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --run-name pi05_bf16_views3_bs1 \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_views3_bs1.jsonl ``` -Run latency: +Do not compare two-camera and three-camera rows as if they had the same workload. + +## FP32 vision parity row + +To compare with an implementation whose vision tower cannot run in BF16, keep the language and expert stacks in BF16 and run SigLIP plus its projector in FP32: ```bash -cd -python benchmark/pi05/bench_realtime_vla_pi05.py \ - --realtime-vla-root \ - --flashrt-root \ +uv run python benchmark/bench_n_batch_ws1_pi05.py \ --checkpoint \ - --num-views 2 \ - --chunk-size 50 \ - --prompt-len 16 \ + --dtype bf16 \ + --vision-dtype float32 \ + --num-images 2 \ --batch-sizes 1 \ --n-warmup 100 \ --n-timed 100 \ - --result-file results/realtime_vla_pi05.jsonl + --run-name pi05_bf16_vision_fp32_views2_bs1 \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_vision_fp32_views2_bs1.jsonl ``` -Notes: +Label this row as `BF16 language/expert + FP32 vision`, not as full BF16. -| Item | Note | -| --- | --- | -| Precision | The wrapper uses BF16 synthetic inputs | -| Pickle checkpoints | For `.pkl` / `.pickle`, add `--trust-pickle-checkpoint` only for trusted files | -| Prompt embedding | If the checkpoint does not contain `language_embeds`, the wrapper creates a synthetic prompt embedding for latency-only runs | - -# vla.cpp +## Batch-size sweep -vla.cpp uses a server/client flow. Build `vla-server` with CUDA enabled, then start the server in one shell and run the Python client in another. - -Basic checks: +Use one command when measuring scaling under otherwise identical settings: ```bash -test -x -test -f -test -f -test -f /tokenizer.json -test -f +uv run python benchmark/bench_n_batch_ws1_pi05.py \ + --checkpoint \ + --dtype bf16 \ + --vision-dtype bfloat16 \ + --num-images 2 \ + --batch-sizes 1 2 4 \ + --n-warmup 100 \ + --n-timed 100 \ + --run-name pi05_bf16_views2_batch_sweep \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_views2_batch_sweep.jsonl ``` -Start server: +The engine is rebuilt for each batch size so that `max_batch_size`, scheduler buffers, and captured graphs match that row. -```bash - \ - --bind tcp://127.0.0.1:5555 \ - --timing-detail phase \ - \ - -``` +# Read the result + +The result file is append-only. Use a new path for a new experiment or remove an old result deliberately before rerunning. + +Each JSONL row records: + +- `latency_ms_mean`, `latency_ms_median`, `latency_ms_p50` +- `latency_ms_p90`, `latency_ms_p99` +- `latency_ms_stdev`, `latency_ms_min`, `latency_ms_max` +- `throughput_samples_per_s` +- batch size, warmup count, timed count, dtype, device, and CUDA graph state -Run latency client: +Inspect the latest row with: ```bash -cd -python benchmark/pi05/bench_vlacpp_pi05_client.py \ - --vlacpp-root \ - --addr tcp://127.0.0.1:5555 \ - --arch pi05 \ - --tokenizer \ - --stats-json \ - --num-views 2 \ - --chunk-size 50 \ - --batch-sizes 1 \ - --n-warmup 100 \ - --n-timed 100 \ - --result-file results/vlacpp_pi05.jsonl +tail -n 1 results/phyai_pi05_bf16_views2_bs1.jsonl | python -m json.tool ``` -Notes: +For a comparable report, retain the command, Git commit, checkpoint identity, GPU and driver, precision, number of images, batch size, warmup count, timed count, and whether CUDA graph was enabled. -| Item | Note | -| --- | --- | -| Model format | vla.cpp needs GGUF files; a safetensors checkpoint is not enough | -| Local files | Prefer local tokenizer and stats files to avoid network or HuggingFace auth issues | -| Stats format | For PI0.5, vla.cpp expects lerobot-style `meta/stats.json`; OpenPI-style `norm_stats.json` is not the same format | -| Timing detail | If the server returns phase timing, the wrapper writes it under `extras.server_phase_latency_ms` | +# Optional external runtime wrappers -# Timing Scope +These wrappers reuse PhyAI's common warmup and result format but do not use the PhyAI engine for inference. Prepare each runtime using its official repository before running the corresponding script. -| Runtime | Timing scope | -| --- | --- | -| FlashRT | Wall time around steady-state `Pi05TorchFrontendRtx.infer(obs)`, after prompt setup, calibration, and the first graph-building call | -| realtime-vla | CUDA-event time around one `Pi05Inference.forward(...)` call | -| vla.cpp | Client wall time for one ZMQ request; server phase timing is copied from the response when available | +| Runtime | Script | Timed scope | Model format | +| --- | --- | --- | --- | +| FlashRT | `benchmark/pi05/bench_flashrt_pi05.py` | Steady-state `Pi05TorchFrontendRtx.infer(obs)` wall time | PI0.5 safetensors | +| realtime-vla | `benchmark/pi05/bench_realtime_vla_pi05.py` | One `Pi05Inference.forward(...)` call | Converted checkpoint or PI0.5 safetensors | +| vla.cpp | `benchmark/pi05/bench_vlacpp_pi05_client.py` | One client ZMQ request; server phases are recorded separately | PI0.5 GGUF plus multimodal projector | -# Troubleshooting +Show the runtime-specific arguments with: -| Symptom | Check | -| --- | --- | -| `No module named flash_rt` | Pass `--flashrt-root` or set `FLASHRT_ROOT` | -| `No module named pi05_infer` | Pass `--realtime-vla-root` or set `REALTIME_VLA_ROOT` | -| realtime-vla safetensors conversion fails | Also pass `--flashrt-root`; verify the FlashRT import works | -| vla.cpp client cannot connect | Confirm the server is ready and client `--addr` matches server `--bind` | -| vla.cpp tokenizer downloads or asks for auth | Use a local tokenizer directory | -| GPU architecture build error | Check CUDA, PyTorch CUDA, driver, and build flags for the target GPU | -| Latency is much slower than expected | Check `nvidia-smi`, rerun after warmup/JIT, and make sure no other process is using the GPU | +```bash +uv run python benchmark/pi05/bench_flashrt_pi05.py --help +uv run python benchmark/pi05/bench_realtime_vla_pi05.py --help +uv run python benchmark/pi05/bench_vlacpp_pi05_client.py --help +``` + +For comparisons, align checkpoint origin, batch size, camera count, action chunk size, prompt shape, warmup, timed iterations, and actual runtime precision. Keep external-runtime rows separate when their model format or timing boundary differs from PhyAI. diff --git a/docs/zh/models/pi05/external-runtime-latency.mdx b/docs/zh/models/pi05/external-runtime-latency.mdx index f1b12df..ab78649 100644 --- a/docs/zh/models/pi05/external-runtime-latency.mdx +++ b/docs/zh/models/pi05/external-runtime-latency.mdx @@ -1,196 +1,218 @@ --- -title: "PI0.5 外部 Runtime 延迟测试" -description: "使用 FlashRT、realtime-vla 和 vla.cpp 跑 PI0.5 latency benchmark" +title: "PI0.5 延迟测试" +description: "复现 PhyAI 原生 PI0.5 延迟，并按需运行外部 runtime 对比" icon: "timer" --- # 概述 -这篇文档说明如何配置并运行 `benchmark/pi05/` 下的三个外部 PI0.5 latency wrapper。 +PI0.5 延迟测试以 PhyAI 原生脚本为主： -这些 wrapper 不使用 PhyAI engine 做推理。每个脚本会直接调用对应 runtime，同时复用 PhyAI 的通用 benchmark runner 来做 warmup、计时和 JSONL 输出。 +```text +benchmark/bench_n_batch_ws1_pi05.py +``` + +脚本使用 `pi05` plugin 构造 PhyAI `Engine`，生成合成输入，完成预热后统计 steady-state `Engine.step(request)` 延迟。它只用于性能测试，不评估 LIBERO 准确率。 -| 脚本 | Runtime | 计时调用 | -| --- | --- | --- | -| `bench_flashrt_pi05.py` | FlashRT | `Pi05TorchFrontendRtx.infer(obs)` | -| `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward(...)` | -| `bench_vlacpp_pi05_client.py` | vla.cpp | 向运行中的 `vla-server` 发一次 ZMQ 请求 | +FlashRT、realtime-vla 和 vla.cpp 放在文档末尾，作为可选对比路径。 -这些脚本会生成 synthetic image 和 state 输入，只用于 latency 测试，不用于 LIBERO accuracy 评测。 +# PhyAI 测了什么 + +| 项目 | 说明 | +| --- | --- | +| 计时调用 | 一次 `Engine.step(request)` | +| 输入 | 合成图像和单 token 合成 prompt | +| 包含 | 视觉塔、LLM prefix、10 次 Euler expert step、scheduler 逻辑和 action 输出 | +| 不包含 | 环境安装、checkpoint 加载、Engine 构造、CUDA Graph capture、request 构造和 warmup | +| GPU 计时 | 每一步使用 CUDA event，并在读取结果前同步 CUDA | +| 输出 | 每个 batch size 写一行 JSONL | -# 通用环境 +request 在计时前直接创建在 GPU 上。因此，图像解码、tokenization、仿真通信和 CPU 到 GPU 的预处理不在这项 latency 里。 -需要一个能 import PhyAI、PyTorch 和通用 benchmark runner 的 Python 环境。 +# 配置环境 + +从仓库源码安装： ```bash -cd -python -c "import torch; import phyai; import benchmark.bench_n_batch" +git clone https://github.com/MEmbodied/phyai.git +cd phyai +uv sync +``` + +确认当前环境可以导入 PhyAI 并识别 GPU： + +```bash +uv run python -c "import torch, phyai; print(torch.__version__, torch.version.cuda); print(torch.cuda.get_device_name(0))" nvidia-smi ``` -对比不同 runtime 时，下面这些设置要对齐： +每次实验都应记录代码版本和 GPU 状态： -| 项 | 设置 | -| --- | --- | -| Batch size | 1 | -| Camera streams | 2 | -| Chunk size | 50 | -| Prompt | 各 runtime 使用相同 prompt 文本 | -| Warmup / timed iterations | 各 runtime 使用相同次数 | -| Precision | 按 runtime 的实际精度路径标注结果 | +```bash +git rev-parse HEAD +nvidia-smi --query-gpu=name,driver_version,memory.used,utilization.gpu --format=csv +nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv +``` -下面的命令都使用占位路径： +确认 GPU 和 CPU 空闲后再运行正式测试。 -| 占位符 | 含义 | -| --- | --- | -| `` | 包含 `benchmark/pi05/` 的 PhyAI 仓库 | -| `` | FlashRT 仓库 | -| `` | realtime-vla 仓库 | -| `` | vla.cpp 仓库 | -| `` | 编译后的 vla.cpp `vla-server` 二进制文件 | -| `` | PI0.5 safetensors checkpoint 目录或文件 | -| `` | vla.cpp 使用的 PI0.5 GGUF 文件 | -| `` | vla.cpp multimodal projector GGUF 文件 | -| `` | 本地 tokenizer 目录 | -| `` | LIBERO `meta/stats.json`，包含 `observation.state.q01/q99` | +# 准备 checkpoint + +`--checkpoint` 必须指向一个目录，至少包含： + +```text +/ + config.json + model.safetensors +``` -# FlashRT +分片 checkpoint 可以使用 `model.safetensors.index.json` 和对应的 shard 文件。 -先按 FlashRT 官方仓库说明安装，然后把仓库路径暴露给 wrapper。 +运行前检查： ```bash -export FLASHRT_ROOT= -cd -python -c "import sys; sys.path.insert(0, '$FLASHRT_ROOT'); from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx; print(Pi05TorchFrontendRtx)" +test -f /config.json +test -f /model.safetensors || \ + test -f /model.safetensors.index.json ``` -运行 latency： +action chunk size 和 denoise step 数来自 `config.json`。标准 PI0.5 配置是 `chunk_size=50`、`num_inference_steps=10`，PhyAI benchmark 没有单独的 `--chunk-size` 参数。 + +对比不同 runtime 时，应使用同一 PI0.5 checkpoint family。某个 runtime 如果需要转换格式，转换前的权重来源仍应相同。 + +# 运行 PhyAI BF16 latency + +下面是推荐的单 batch、双相机 LIBERO shape 测试命令： ```bash cd -python benchmark/pi05/bench_flashrt_pi05.py \ - --flashrt-root \ +mkdir -p results + +uv run python benchmark/bench_n_batch_ws1_pi05.py \ --checkpoint \ - --precision bf16 \ - --num-views 2 \ - --chunk-size 50 \ + --dtype bf16 \ + --vision-dtype bfloat16 \ + --num-images 2 \ --batch-sizes 1 \ --n-warmup 100 \ --n-timed 100 \ - --result-file results/flashrt_pi05.jsonl + --run-name pi05_bf16_views2_bs1 \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_views2_bs1.jsonl ``` -说明： +这条命令对应的设置如下： -| 项 | 说明 | +| 参数 | 设置 | | --- | --- | -| Chunk size | wrapper 直接使用 FlashRT 的 `Pi05TorchFrontendRtx` API，因为 `chunk_size` 是 frontend 构造参数 | -| Denoise steps | 不要用 `load_model(..., num_steps=50)` 设置 action chunk size；FlashRT 里的 `num_steps` 表示 denoise steps | -| Precision | `--precision bf16` 使用 FlashRT 的 forced-BF16 PI0.5 RTX 路径；`--precision fp8_bf16` 是单独的 FP8/BF16 优化结果 | +| Engine 精度 | BF16 | +| 视觉塔精度 | BF16 | +| Batch size | 1 | +| 相机数 | 2 | +| Action chunk | 从 checkpoint 读取，通常为 50 | +| Euler steps | 从 checkpoint 读取，通常为 10 | +| CUDA Graph | CUDA 环境下默认开启 | +| Prompt shape | 单个合成 token | +| Warmup / timed iterations | 100 / 100 | -# realtime-vla +如果用户没有显式覆盖，PI0.5 plugin 还会应用模型配置中的推荐 runtime 参数。提交结果时应记录所有 `PHYAI_*` 环境变量，因为它们可能改变 attention backend 或 workspace 设置。 -先按 realtime-vla 官方仓库说明安装。wrapper 可以直接加载转换后的 `.pt` / `.pth` checkpoint。如果传入 PI0.5 safetensors checkpoint，还需要提供 FlashRT 路径，用它的 PI0.5 转换 helper。 +## 三相机 PI0.5 base shape + +公开的 `pi05_base` 默认使用三路相机。只修改 `--num-images` 和结果名称： ```bash -export REALTIME_VLA_ROOT= -export FLASHRT_ROOT= -cd -python -c "import sys; sys.path.insert(0, '$REALTIME_VLA_ROOT'); from pi05_infer import Pi05Inference; print(Pi05Inference)" +uv run python benchmark/bench_n_batch_ws1_pi05.py \ + --checkpoint \ + --dtype bf16 \ + --vision-dtype bfloat16 \ + --num-images 3 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --run-name pi05_bf16_views3_bs1 \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_views3_bs1.jsonl ``` -运行 latency： +双相机和三相机的工作量不同，不能放在同一行直接比较。 + +## 视觉塔 FP32 对齐组 + +如果对比对象的视觉塔无法使用 BF16，可以保持 language/expert 为 BF16，只把 SigLIP 和 projector 改为 FP32： ```bash -cd -python benchmark/pi05/bench_realtime_vla_pi05.py \ - --realtime-vla-root \ - --flashrt-root \ +uv run python benchmark/bench_n_batch_ws1_pi05.py \ --checkpoint \ - --num-views 2 \ - --chunk-size 50 \ - --prompt-len 16 \ + --dtype bf16 \ + --vision-dtype float32 \ + --num-images 2 \ --batch-sizes 1 \ --n-warmup 100 \ --n-timed 100 \ - --result-file results/realtime_vla_pi05.jsonl + --run-name pi05_bf16_vision_fp32_views2_bs1 \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_vision_fp32_views2_bs1.jsonl ``` -说明： - -| 项 | 说明 | -| --- | --- | -| Precision | wrapper 使用 BF16 synthetic inputs | -| Pickle checkpoint | `.pkl` / `.pickle` 文件只有在可信时才加 `--trust-pickle-checkpoint` | -| Prompt embedding | 如果 checkpoint 没有 `language_embeds`，wrapper 会为 latency-only 测试创建 synthetic prompt embedding | - -# vla.cpp +这组结果应标为“BF16 language/expert + FP32 vision”，不能写成全 BF16。 -vla.cpp 是 server/client 流程。先用 CUDA 编译 `vla-server`，然后在一个 shell 启动 server，在另一个 shell 运行 Python client。 +## Batch size sweep -基础检查： +需要测试 batch scaling 时，在其他设置不变的情况下使用一条命令： ```bash -test -x -test -f -test -f -test -f /tokenizer.json -test -f +uv run python benchmark/bench_n_batch_ws1_pi05.py \ + --checkpoint \ + --dtype bf16 \ + --vision-dtype bfloat16 \ + --num-images 2 \ + --batch-sizes 1 2 4 \ + --n-warmup 100 \ + --n-timed 100 \ + --run-name pi05_bf16_views2_batch_sweep \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_views2_batch_sweep.jsonl ``` -启动 server： +脚本会为每个 batch size 重新构造 Engine，使 `max_batch_size`、scheduler buffer 和 CUDA Graph 与该行结果一致。 -```bash - \ - --bind tcp://127.0.0.1:5555 \ - --timing-detail phase \ - \ - -``` +# 查看结果 -运行 latency client： +结果文件采用追加写入。新的实验应使用新的文件名；如果要复用旧文件名，应先明确处理旧结果，避免把两次实验混在一起。 + +每行 JSONL 包含： + +- `latency_ms_mean`、`latency_ms_median`、`latency_ms_p50` +- `latency_ms_p90`、`latency_ms_p99` +- `latency_ms_stdev`、`latency_ms_min`、`latency_ms_max` +- `throughput_samples_per_s` +- batch size、warmup 次数、timed 次数、dtype、device 和 CUDA Graph 状态 + +查看最新一行： ```bash -cd -python benchmark/pi05/bench_vlacpp_pi05_client.py \ - --vlacpp-root \ - --addr tcp://127.0.0.1:5555 \ - --arch pi05 \ - --tokenizer \ - --stats-json \ - --num-views 2 \ - --chunk-size 50 \ - --batch-sizes 1 \ - --n-warmup 100 \ - --n-timed 100 \ - --result-file results/vlacpp_pi05.jsonl +tail -n 1 results/phyai_pi05_bf16_views2_bs1.jsonl | python -m json.tool ``` -说明： +一份可比较的结果至少要保留：完整命令、Git commit、checkpoint 标识、GPU 与 driver、精度、相机数、batch size、warmup 次数、timed 次数和 CUDA Graph 开关。 -| 项 | 说明 | -| --- | --- | -| 模型格式 | vla.cpp 需要 GGUF 文件；safetensors checkpoint 不够 | -| 本地文件 | tokenizer 和 stats 建议都用本地文件，避免测试时触发网络或 HuggingFace 权限问题 | -| Stats 格式 | PI0.5 下 vla.cpp 需要 lerobot 风格的 `meta/stats.json`；OpenPI 风格的 `norm_stats.json` 不是同一种格式 | -| 详细计时 | 如果 server 返回 phase timing，wrapper 会写入 `extras.server_phase_latency_ms` | +# 可选外部 runtime -# 计时口径 +下面三个 wrapper 复用 PhyAI 的 warmup 和结果格式，但推理过程不经过 PhyAI Engine。运行前需要按各自官方仓库配置好对应环境。 -| Runtime | 计时范围 | -| --- | --- | -| FlashRT | 对 steady-state `Pi05TorchFrontendRtx.infer(obs)` 计 wall time；不包含 prompt setup、calibration 和第一次 graph-building 调用 | -| realtime-vla | 用 CUDA event 计一次 `Pi05Inference.forward(...)` | -| vla.cpp | 计一次 ZMQ 请求的 client wall time；如果 response 里有 server phase timing，会一并记录 | +| Runtime | 脚本 | 计时范围 | 模型格式 | +| --- | --- | --- | --- | +| FlashRT | `benchmark/pi05/bench_flashrt_pi05.py` | steady-state `Pi05TorchFrontendRtx.infer(obs)` wall time | PI0.5 safetensors | +| realtime-vla | `benchmark/pi05/bench_realtime_vla_pi05.py` | 一次 `Pi05Inference.forward(...)` | 转换后的 checkpoint 或 PI0.5 safetensors | +| vla.cpp | `benchmark/pi05/bench_vlacpp_pi05_client.py` | 一次 client ZMQ 请求，server phase 另行记录 | PI0.5 GGUF 和 multimodal projector | -# 常见问题 +通过 `--help` 查看每个 runtime 的必要参数： -| 现象 | 检查方式 | -| --- | --- | -| `No module named flash_rt` | 传 `--flashrt-root` 或设置 `FLASHRT_ROOT` | -| `No module named pi05_infer` | 传 `--realtime-vla-root` 或设置 `REALTIME_VLA_ROOT` | -| realtime-vla safetensors 转换失败 | 同时传 `--flashrt-root`，并确认 FlashRT import 正常 | -| vla.cpp client 连不上 | 确认 server 已 ready，且 client `--addr` 和 server `--bind` 一致 | -| vla.cpp tokenizer 触发下载或鉴权 | 使用本地 tokenizer 目录 | -| GPU 架构编译报错 | 检查 CUDA、PyTorch CUDA、driver 和目标 GPU 的 build flags | -| latency 明显偏慢 | 看 `nvidia-smi`，排除 GPU 占用；JIT/warmup 后重新跑 | +```bash +uv run python benchmark/pi05/bench_flashrt_pi05.py --help +uv run python benchmark/pi05/bench_realtime_vla_pi05.py --help +uv run python benchmark/pi05/bench_vlacpp_pi05_client.py --help +``` + +做横向对比时，需要对齐 checkpoint 来源、batch size、相机数、action chunk size、prompt shape、warmup、timed iterations 和 runtime 实际精度。如果模型格式或计时边界不同，应把结果单独标注。