diff --git a/benchmark/pi05/README_external_pi05_benchmarks.md b/benchmark/pi05/README_external_pi05_benchmarks.md new file mode 100644 index 0000000..fa08f53 --- /dev/null +++ b/benchmark/pi05/README_external_pi05_benchmarks.md @@ -0,0 +1,168 @@ +# 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 new file mode 100755 index 0000000..9d390a2 --- /dev/null +++ b/benchmark/pi05/bench_flashrt_pi05.py @@ -0,0 +1,229 @@ +#!/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 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..6e07a67 --- /dev/null +++ b/benchmark/pi05/bench_vlacpp_pi05_client.py @@ -0,0 +1,249 @@ +#!/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/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()