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..8cb9217 --- /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") + + torch.manual_seed(args.seed) + 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, + ) + 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..403fdd5 --- /dev/null +++ b/benchmark/pi05/bench_vlacpp_pi05_client.py @@ -0,0 +1,238 @@ +#!/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 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: + 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) + + spec = bnb.BenchSpec( + name="vlacpp_pi05_zmq_client", + step_callable=step, + teardown_callable=teardown, + ) + # 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 + + 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/docs.json b/docs/docs.json index 3486462..2013a8f 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -36,7 +36,10 @@ "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" ] }, { @@ -106,7 +109,10 @@ "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" ] }, { diff --git a/docs/models/pi05/eight-gpu-inference.mdx b/docs/models/pi05/eight-gpu-inference.mdx new file mode 100644 index 0000000..428330b --- /dev/null +++ b/docs/models/pi05/eight-gpu-inference.mdx @@ -0,0 +1,175 @@ +--- +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. + +# 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) | [`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. 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 + +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..a9d1b2e --- /dev/null +++ b/docs/models/pi05/external-runtime-latency.mdx @@ -0,0 +1,218 @@ +--- +title: "PI0.5 Latency Benchmark" +description: "Reproduce native PhyAI PI0.5 latency and run optional external runtime comparisons" +icon: "timer" +--- + +# Overview + +This page uses the native PhyAI benchmark as the primary PI0.5 latency path: + +```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. + +The FlashRT, realtime-vla, and vla.cpp wrappers are listed at the end as optional comparison paths. + +# 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 | + +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. + +# Prepare the environment + +From a clean checkout: + +```bash +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 +``` + +Record the source revision and GPU state with each result: + +```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 +``` + +Run latency measurements only when the GPU and CPU are idle. + +# Prepare the checkpoint + +`--checkpoint` must point to one directory containing: + +```text +/ + config.json + model.safetensors +``` + +A sharded checkpoint may use `model.safetensors.index.json` and its shard files instead. + +Check the files before starting: + +```bash +test -f /config.json +test -f /model.safetensors || \ + test -f /model.safetensors.index.json +``` + +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 +mkdir -p results + +uv run python benchmark/bench_n_batch_ws1_pi05.py \ + --checkpoint \ + --dtype bf16 \ + --vision-dtype bfloat16 \ + --num-images 2 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --run-name pi05_bf16_views2_bs1 \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_views2_bs1.jsonl +``` + +This configuration means: + +| Argument | Value | +| --- | --- | +| 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. + +## Three-camera PI0.5 base shape + +The public `pi05_base` contract uses three camera streams. Change only `--num-images` and the output label: + +```bash +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 +``` + +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 +uv run python benchmark/bench_n_batch_ws1_pi05.py \ + --checkpoint \ + --dtype bf16 \ + --vision-dtype float32 \ + --num-images 2 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --run-name pi05_bf16_vision_fp32_views2_bs1 \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_vision_fp32_views2_bs1.jsonl +``` + +Label this row as `BF16 language/expert + FP32 vision`, not as full BF16. + +## Batch-size sweep + +Use one command when measuring scaling under otherwise identical settings: + +```bash +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 +``` + +The engine is rebuilt for each batch size so that `max_batch_size`, scheduler buffers, and captured graphs match that row. + +# 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 + +Inspect the latest row with: + +```bash +tail -n 1 results/phyai_pi05_bf16_views2_bs1.jsonl | python -m json.tool +``` + +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. + +# Optional external runtime wrappers + +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 | 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 | + +Show the runtime-specific arguments with: + +```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/models/pi05/libero-four-suites.mdx b/docs/models/pi05/libero-four-suites.mdx new file mode 100644 index 0000000..1a8b51a --- /dev/null +++ b/docs/models/pi05/libero-four-suites.mdx @@ -0,0 +1,273 @@ +--- +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 | 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. + +# 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) | [`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: + +```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 3a9a28b21aa53627a07c3be9fcb86e6adb1b4ae8 + +git clone https://github.com/rebecca26358/vla-evaluation-harness.git "$VLA_ROOT" +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, tests, and the checkpoint converter 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. + +```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 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 + +export PHYAI_CKPT_IN_CONTAINER=/data/share/pi05_libero_phyai_converted +export TOKENIZER_IN_CONTAINER=/data/share/paligemma-3b-pt-224 + +export PHYAI_IMAGE=nvcr.io/nvidia/pytorch:25.12-py3 +``` + +Check the model directories: + +```bash +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" +``` + +# Evaluator Environment + +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 "$VLA_ROOT" +uv sync +``` + +# PhyAI Runtime 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 \ + --name "$PHYAI_CONTAINER" \ + --ipc=host \ + --network=host \ + -v "$PHYAI_ROOT:$PHYAI_ROOT" \ + -v "$VLA_ROOT:$VLA_ROOT" \ + -v "$MODEL_ROOT:/data/share" \ + "$PHYAI_IMAGE" \ + bash +``` + +Inside the container, sync both clean checkouts and verify the model paths: + +```bash +cd "$PHYAI_ROOT" +uv sync + +cd "$VLA_ROOT" +uv sync + +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: + +```bash +cd "$VLA_ROOT" + +export PHYAI_CHECKPOINT_PATH="$PHYAI_CKPT_IN_CONTAINER" +export PHYAI_TOKENIZER_PATH="$TOKENIZER_IN_CONTAINER" + +uv run vla-eval serve --config configs/model_servers/phyai/libero.yaml +``` + +Keep the server running while the evaluator connects to `ws://localhost:8000`. + +# Smoke Test + +In another terminal, run a short check to confirm that the evaluator can reach the policy server: + +```bash +cd "$VLA_ROOT" + +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. + +# 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 spatial object goal 10; do + uv run vla-eval run \ + --config "configs/benchmarks/libero/${suite}.yaml" \ + --server-url ws://localhost:8000 \ + --dev \ + --yes +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..ab78649 --- /dev/null +++ b/docs/zh/models/pi05/external-runtime-latency.mdx @@ -0,0 +1,218 @@ +--- +title: "PI0.5 延迟测试" +description: "复现 PhyAI 原生 PI0.5 延迟，并按需运行外部 runtime 对比" +icon: "timer" +--- + +# 概述 + +PI0.5 延迟测试以 PhyAI 原生脚本为主： + +```text +benchmark/bench_n_batch_ws1_pi05.py +``` + +脚本使用 `pi05` plugin 构造 PhyAI `Engine`，生成合成输入，完成预热后统计 steady-state `Engine.step(request)` 延迟。它只用于性能测试，不评估 LIBERO 准确率。 + +FlashRT、realtime-vla 和 vla.cpp 放在文档末尾，作为可选对比路径。 + +# 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 里。 + +# 配置环境 + +从仓库源码安装： + +```bash +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 +``` + +每次实验都应记录代码版本和 GPU 状态： + +```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 空闲后再运行正式测试。 + +# 准备 checkpoint + +`--checkpoint` 必须指向一个目录，至少包含： + +```text +/ + config.json + model.safetensors +``` + +分片 checkpoint 可以使用 `model.safetensors.index.json` 和对应的 shard 文件。 + +运行前检查： + +```bash +test -f /config.json +test -f /model.safetensors || \ + test -f /model.safetensors.index.json +``` + +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 +mkdir -p results + +uv run python benchmark/bench_n_batch_ws1_pi05.py \ + --checkpoint \ + --dtype bf16 \ + --vision-dtype bfloat16 \ + --num-images 2 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --run-name pi05_bf16_views2_bs1 \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_views2_bs1.jsonl +``` + +这条命令对应的设置如下： + +| 参数 | 设置 | +| --- | --- | +| 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 | + +如果用户没有显式覆盖，PI0.5 plugin 还会应用模型配置中的推荐 runtime 参数。提交结果时应记录所有 `PHYAI_*` 环境变量，因为它们可能改变 attention backend 或 workspace 设置。 + +## 三相机 PI0.5 base shape + +公开的 `pi05_base` 默认使用三路相机。只修改 `--num-images` 和结果名称： + +```bash +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 +``` + +双相机和三相机的工作量不同，不能放在同一行直接比较。 + +## 视觉塔 FP32 对齐组 + +如果对比对象的视觉塔无法使用 BF16，可以保持 language/expert 为 BF16，只把 SigLIP 和 projector 改为 FP32： + +```bash +uv run python benchmark/bench_n_batch_ws1_pi05.py \ + --checkpoint \ + --dtype bf16 \ + --vision-dtype float32 \ + --num-images 2 \ + --batch-sizes 1 \ + --n-warmup 100 \ + --n-timed 100 \ + --run-name pi05_bf16_vision_fp32_views2_bs1 \ + --bench-name pi05_ws1 \ + --result-file results/phyai_pi05_bf16_vision_fp32_views2_bs1.jsonl +``` + +这组结果应标为“BF16 language/expert + FP32 vision”，不能写成全 BF16。 + +## Batch size sweep + +需要测试 batch scaling 时，在其他设置不变的情况下使用一条命令： + +```bash +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 +``` + +脚本会为每个 batch size 重新构造 Engine，使 `max_batch_size`、scheduler buffer 和 CUDA Graph 与该行结果一致。 + +# 查看结果 + +结果文件采用追加写入。新的实验应使用新的文件名；如果要复用旧文件名，应先明确处理旧结果，避免把两次实验混在一起。 + +每行 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 +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 开关。 + +# 可选外部 runtime + +下面三个 wrapper 复用 PhyAI 的 warmup 和结果格式，但推理过程不经过 PhyAI Engine。运行前需要按各自官方仓库配置好对应环境。 + +| 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 的必要参数： + +```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 实际精度。如果模型格式或计时边界不同，应把结果单独标注。 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..6c489c9 --- /dev/null +++ b/docs/zh/models/pi05/libero-four-suites.mdx @@ -0,0 +1,270 @@ +--- +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 | 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 权限。更稳的做法是从已有权限的机器同步到本机。 + +# 固定源码版本 + +本复现流程使用以下两个 PR 中的不可变提交： + +| 仓库 | Pull request | 提交 | +| --- | --- | --- | +| 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) | + +检出对应版本： + +```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 3a9a28b21aa53627a07c3be9fcb86e6adb1b4ae8 + +git clone https://github.com/rebecca26358/vla-evaluation-harness.git "$VLA_ROOT" +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 提交增加了文档、测试和 checkpoint converter，没有修改 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。 + +# 设置路径 + +先在目标机器上设置路径。下面只用占位路径，按实际机器改即可。 + +```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 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 + +export PHYAI_CKPT_IN_CONTAINER=/data/share/pi05_libero_phyai_converted +export TOKENIZER_IN_CONTAINER=/data/share/paligemma-3b-pt-224 + +export PHYAI_IMAGE=nvcr.io/nvidia/pytorch:25.12-py3 +``` + +检查模型目录： + +```bash +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" +``` + +# 准备评测环境 + +在宿主机同步 VLA 评测环境。这个命令可用于干净 checkout；PhyAI 并没有定义 `cu130` 或 `libero` 依赖组。 + +```bash +cd "$VLA_ROOT" +uv sync +``` + +# 准备 PhyAI 运行容器 + +在 PhyAI 使用的 CUDA 13 镜像中运行 policy server。评测器会根据 benchmark +配置另行启动 LIBERO 仿真镜像。 + +```bash +docker run --gpus all -it --rm \ + --name "$PHYAI_CONTAINER" \ + --ipc=host \ + --network=host \ + -v "$PHYAI_ROOT:$PHYAI_ROOT" \ + -v "$VLA_ROOT:$VLA_ROOT" \ + -v "$MODEL_ROOT:/data/share" \ + "$PHYAI_IMAGE" \ + bash +``` + +进入容器后，同步两个干净 checkout，并检查模型路径： + +```bash +cd "$PHYAI_ROOT" +uv sync + +cd "$VLA_ROOT" +uv sync + +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： + +```bash +cd "$VLA_ROOT" + +export PHYAI_CHECKPOINT_PATH="$PHYAI_CKPT_IN_CONTAINER" +export PHYAI_TOKENIZER_PATH="$TOKENIZER_IN_CONTAINER" + +uv run vla-eval serve --config configs/model_servers/phyai/libero.yaml +``` + +保持 server 运行，评测器会连接 `ws://localhost:8000`。 + +# 跑 smoke test + +在另一个终端里跑一次很短的检查，确认 evaluator 能连到 policy server： + +```bash +cd "$VLA_ROOT" + +uv run vla-eval run \ + --config configs/benchmarks/libero/smoke_test.yaml \ + --server-url ws://localhost:8000 \ + --dev \ + --yes +``` + +如果 smoke test 失败，先不要跑完整四套任务。优先检查 server 日志、端口、模型路径和 tokenizer 路径。 + +# 跑四套任务 + +确认 GPU 空闲后，再跑完整评测。建议每套任务单独保存结果。 + +```bash +cd "$VLA_ROOT" + +for suite in spatial object goal 10; do + uv run vla-eval run \ + --config "configs/benchmarks/libero/${suite}.yaml" \ + --server-url ws://localhost:8000 \ + --dev \ + --yes +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 - < 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, + 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) + self._action_dim = self._resolve_action_dim(self.config) + 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( + 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, + ) + 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=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, + 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_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": + 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]: + 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) + with torch.inference_mode(): + raw_actions = self.engine.step(request) + 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..45b30fc --- /dev/null +++ b/phyai/tests/policies/test_pi05_libero.py @@ -0,0 +1,68 @@ +"""PI0.5 LIBERO 策略的轻量接口测试。""" + +from __future__ import annotations + +import numpy as np +import torch + +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 + + +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) diff --git a/tools/convert_openpi_pi05_to_phyai.py b/tools/convert_openpi_pi05_to_phyai.py new file mode 100644 index 0000000..7bb8868 --- /dev/null +++ b/tools/convert_openpi_pi05_to_phyai.py @@ -0,0 +1,855 @@ +# /// 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())