anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
codex-dlo-lifecycle-optimization
vllm-omni/tests/diffusion/models/pi0/test_pi0_units.py
T
Yichen WangandDidan Deng 2a9f6aaa8e Add π0 VLA model support (#4222) 
Signed-off-by: Yichen Wang <yichen.wang@bytedance.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
2026-08-17 14:29:17 +08:00

586 lines
22 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CPU unit tests for the π0 (Pi-Zero) VLA model.
Deliberately lightweight: pure utility code (config plumbing, mask building,
sinusoidal embeddings, normalization, image preprocessing, weight-load remap)
that runs on CPU with no weights and no lerobot. The full LeRobot parity check
lives in ``tests/diffusion/models/pi0/test_pi0_parity.py`` and requires a model download.
pytest tests/diffusion/models/pi0/test_pi0_units.py -v
"""
from __future__ import annotations
import logging
import math
import pytest
import torch
from vllm_omni.diffusion.models.pi0.config import Pi0Config
from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
OPENPI_ATTENTION_MASK_VALUE,
Pi0ForActionPrediction,
_apply_norm,
_build_norm_buffers,
create_sinusoidal_pos_embedding,
get_gemma_config,
make_att_2d_masks,
prepare_attention_masks_4d,
)
from vllm_omni.diffusion.models.pi0.processor_pi0 import (
Pi0ImageProcessor,
build_model_inputs,
pil_image_to_tensor,
resize_with_pad,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# ----------------------------------------------------------------------------
# Pi0Config dataclass + checkpoint resolver
# ----------------------------------------------------------------------------
_LEROBOT_CFG = {
"type": "pi0",
"paligemma_variant": "gemma_2b",
"action_expert_variant": "gemma_300m",
"chunk_size": 50,
"max_action_dim": 32,
"max_state_dim": 32,
"num_inference_steps": 10,
"image_resolution": [224, 224],
"tokenizer_max_length": 48,
"dtype": "float32",
# Training-only keys that must be filtered out of the dataclass.
"optimizer_lr": 2.5e-5,
"scheduler_warmup_steps": 1000,
"input_features": {
"observation.images.base_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
"observation.images.left_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
"observation.images.right_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
"observation.state": {"type": "STATE", "shape": [32]},
},
}
_EXPECTED_CAMERA_ORDER = [
"observation.images.base_0_rgb",
"observation.images.left_wrist_0_rgb",
"observation.images.right_wrist_0_rgb",
]
def test_config_camera_order_from_input_features():
c = Pi0Config.from_model_config(_LEROBOT_CFG)
assert c.image_feature_keys == _EXPECTED_CAMERA_ORDER
def test_config_filters_training_only_keys():
c = Pi0Config.from_model_config(_LEROBOT_CFG)
assert not hasattr(c, "optimizer_lr")
assert not hasattr(c, "scheduler_warmup_steps")
def test_config_coerces_image_resolution_to_tuple():
c = Pi0Config.from_model_config(_LEROBOT_CFG)
assert c.image_resolution == (224, 224)
assert isinstance(c.image_resolution, tuple)
def test_config_non_square_resolution_raises():
with pytest.raises(ValueError):
Pi0Config(image_resolution=(224, 256))
@pytest.mark.parametrize("bad", [(224,), (224, 224, 3), 224])
def test_config_malformed_resolution_raises(bad):
with pytest.raises(ValueError):
Pi0Config(image_resolution=bad)
def test_config_defaults_and_empty():
c = Pi0Config.from_model_config(None)
assert c.chunk_size == 50
assert c.max_action_dim == 32
assert c.max_state_dim == 32
assert c.num_inference_steps == 10
assert c.image_resolution == (224, 224)
assert c.tokenizer_max_length == 48
assert c.max_cameras == 3
assert c.norm_stats is None
assert c.image_key_map == {}
def test_config_deploy_yaml_override():
"""Deploy-yaml model_config is authoritative (e.g. force dtype/variant)."""
c = Pi0Config.from_model_config({**_LEROBOT_CFG, "dtype": "bfloat16", "num_inference_steps": 4})
assert c.dtype == "bfloat16"
assert c.num_inference_steps == 4
# ----------------------------------------------------------------------------
# Kernel utilities: sinusoidal timestep embedding
# ----------------------------------------------------------------------------
def test_sinusoidal_embedding_shape():
emb = create_sinusoidal_pos_embedding(torch.tensor([0.0, 0.5, 1.0]), dimension=64)
assert emb.shape == (3, 64)
def test_sinusoidal_embedding_values_at_zero():
emb = create_sinusoidal_pos_embedding(torch.tensor([0.0]), dimension=4)
assert torch.allclose(emb[0, :2].float(), torch.zeros(2), atol=1e-6)
assert torch.allclose(emb[0, 2:].float(), torch.ones(2), atol=1e-6)
def test_sinusoidal_embedding_odd_dim_raises():
with pytest.raises(ValueError):
create_sinusoidal_pos_embedding(torch.tensor([1.0]), dimension=3)
def test_sinusoidal_embedding_non_1d_raises():
with pytest.raises(ValueError):
create_sinusoidal_pos_embedding(torch.tensor([[0.0, 0.5]]), dimension=4)
def test_sinusoidal_embedding_distinct_timesteps_differ():
emb = create_sinusoidal_pos_embedding(torch.tensor([0.1, 0.2, 0.5, 0.9]), dimension=16)
for i in range(emb.shape[0]):
for j in range(i + 1, emb.shape[0]):
assert not torch.allclose(emb[i], emb[j])
# ----------------------------------------------------------------------------
# Kernel utilities: attention mask construction
# ----------------------------------------------------------------------------
def test_make_att_2d_masks_bidirectional():
result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.zeros(1, 4, dtype=torch.long))
assert result.all()
def test_make_att_2d_masks_causal():
result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.ones(1, 4, dtype=torch.long))
expected = torch.tril(torch.ones(4, 4, dtype=torch.bool))
assert torch.equal(result[0], expected)
def test_make_att_2d_masks_prefix_lm():
result = make_att_2d_masks(torch.ones(1, 5, dtype=torch.bool), torch.tensor([[0, 0, 0, 1, 1]], dtype=torch.long))
assert result[0, 0, 0]
assert not result[0, 0, 3]
assert result[0, 3, 0]
assert not result[0, 3, 4]
def test_make_att_2d_masks_respects_padding():
pad = torch.tensor([[1, 1, 1, 0]], dtype=torch.bool) # last token is padding
att = torch.zeros(1, 4, dtype=torch.long) # full bidirectional
result = make_att_2d_masks(pad, att)
assert not result[0, :, 3].any(), "Nothing should attend to padded key"
assert not result[0, 3, :].any(), "Padded query should attend to nothing"
def test_prepare_attention_masks_4d():
mask_2d = torch.tensor([[[True, False], [True, True]]])
mask_4d = prepare_attention_masks_4d(mask_2d)
assert mask_4d.shape == (1, 1, 2, 2)
assert mask_4d[0, 0, 0, 0].item() == 0.0
assert mask_4d[0, 0, 1, 1].item() == 0.0
expected = torch.tensor(OPENPI_ATTENTION_MASK_VALUE, dtype=mask_4d.dtype)
assert torch.equal(mask_4d[0, 0, 0, 1], expected)
def test_attention_mask_value_finite_and_negative():
assert math.isfinite(OPENPI_ATTENTION_MASK_VALUE)
assert OPENPI_ATTENTION_MASK_VALUE < -1e30
# ----------------------------------------------------------------------------
# Normalization round-trips
# ----------------------------------------------------------------------------
def test_build_norm_buffers_missing():
assert _build_norm_buffers(None, "state") is None
assert _build_norm_buffers({}, "state") is None
assert _build_norm_buffers({"state": None}, "state") is None
def test_build_norm_buffers_mean_std():
stats = {"state": {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 0.5]}}
out = _build_norm_buffers(stats, "state")
assert out["mode"] == "mean_std"
assert torch.equal(out["mean"], torch.tensor([1.0, 2.0]))
assert torch.equal(out["std"], torch.tensor([0.5, 0.5]))
def test_build_norm_buffers_min_max():
stats = {"action": {"mode": "min_max", "min": [-1.0, 0.0], "max": [1.0, 2.0]}}
out = _build_norm_buffers(stats, "action")
assert out["mode"] == "min_max"
assert torch.equal(out["min"], torch.tensor([-1.0, 0.0]))
assert torch.equal(out["max"], torch.tensor([1.0, 2.0]))
def test_apply_norm_mean_std_forward_inverse():
stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.5, -0.2], "std": [2.0, 0.1]}}, "s")
x = torch.tensor([[1.0, 0.3]])
z = _apply_norm(x, stats, inverse=False)
assert torch.allclose(z, torch.tensor([[0.25, 5.0]]), atol=1e-5)
x_back = _apply_norm(z, stats, inverse=True)
assert torch.allclose(x_back, x, atol=1e-5)
def test_apply_norm_min_max_forward_inverse():
stats = _build_norm_buffers({"s": {"mode": "min_max", "min": [0.0, -1.0], "max": [1.0, 1.0]}}, "s")
x = torch.tensor([[0.25, 0.5]])
z = _apply_norm(x, stats, inverse=False)
assert torch.allclose(z, torch.tensor([[-0.5, 0.5]]), atol=1e-5)
x_back = _apply_norm(z, stats, inverse=True)
assert torch.allclose(x_back, x, atol=1e-5)
def test_apply_norm_preserves_padded_tail():
stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.0, 0.0], "std": [1.0, 1.0]}}, "s")
x = torch.tensor([[1.0, 2.0, 99.0, -99.0]])
z = _apply_norm(x, stats, inverse=False)
assert torch.allclose(z[:, :2], x[:, :2])
assert torch.allclose(z[:, 2:], x[:, 2:])
def test_apply_norm_noop_when_stats_none():
x = torch.randn(2, 5)
assert torch.equal(_apply_norm(x, None, inverse=False), x)
assert torch.equal(_apply_norm(x, None, inverse=True), x)
# ----------------------------------------------------------------------------
# Gemma variant configs
# ----------------------------------------------------------------------------
def test_gemma_2b_dims():
c = get_gemma_config("gemma_2b")
assert c.width == 2048
assert c.depth == 18
def test_gemma_300m_dims():
c = get_gemma_config("gemma_300m")
assert c.width == 1024
assert c.depth == 18
def test_gemma_unknown_variant_raises():
with pytest.raises(ValueError):
get_gemma_config("gemma_7b")
# ----------------------------------------------------------------------------
# load_weights remap rules (tiny model; needs torch + transformers)
# ----------------------------------------------------------------------------
def _tiny_pi0_model():
"""Construct a tiny π0 model for remap tests (full Gemma depth but small
action/state dims — the PaliGemma backbone is always full-size)."""
cfg = Pi0Config(max_action_dim=8, max_state_dim=8, chunk_size=2)
return Pi0ForActionPrediction(cfg)
@pytest.fixture(autouse=True)
def _silence_pi0_loader(caplog):
# The remap tests feed a single synthetic weight; the loader's
# "missing params" warning is expected noise here.
logging.getLogger("vllm_omni.diffusion.models.pi0.modeling_pi0").setLevel(logging.ERROR)
yield
@pytest.mark.slow
def test_remap_paligemma_submodules_to_nested_layout():
model = _tiny_pi0_model()
params = dict(model.named_parameters())
target = next(k for k in params if k.startswith("paligemma_with_expert.paligemma.model.vision_tower."))
flat_key = "model." + target.replace(
"paligemma_with_expert.paligemma.model.vision_tower.",
"paligemma_with_expert.paligemma.vision_tower.",
)
payload = torch.full(params[target].shape, 0.12345)
before = params[target].detach().clone()
model.load_weights([(flat_key, payload)])
after = dict(model.named_parameters())[target]
assert not torch.equal(before, after), "flat→nested remap broken"
assert torch.allclose(after, payload)
@pytest.mark.slow
def test_lm_head_remap_to_embed_tokens():
model = _tiny_pi0_model()
params = dict(model.named_parameters())
target = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
assert target in params
payload = torch.full(params[target].shape, 0.5)
before = params[target].detach().clone()
model.load_weights([("model.paligemma_with_expert.paligemma.lm_head.weight", payload)])
after = dict(model.named_parameters())[target]
assert not torch.equal(before, after), "lm_head→embed_tokens remap broken"
assert torch.allclose(after, payload)
@pytest.mark.slow
def test_vision_tower_version_robust_remap():
"""transformers >=5.4 flattens SigLIP from vision_tower.vision_model.* to
vision_tower.*; the lerobot/pi0_base checkpoint uses the vision_model infix.
load_weights must adapt so vision weights aren't silently dropped (which
would leave the vision encoder at random init)."""
model = _tiny_pi0_model()
params = dict(model.named_parameters())
# A representative vision-tower param as it exists on the model.
target = next(
k
for k in params
if k.startswith("paligemma_with_expert.paligemma.model.vision_tower.") and "patch_embedding.weight" in k
)
vt = "paligemma_with_expert.paligemma.model.vision_tower."
rest = target[len(vt) :]
# Build the checkpoint key in whichever layout differs from the model's.
if rest.startswith("vision_model."):
ckpt_key = vt + rest[len("vision_model.") :]
else:
ckpt_key = vt + "vision_model." + rest
payload = torch.full(params[target].shape, 0.073)
before = params[target].detach().clone()
model.load_weights([("model." + ckpt_key, payload)])
after = dict(model.named_parameters())[target]
assert not torch.equal(before, after), "vision_tower version-robust remap broken"
assert torch.allclose(after, payload)
# ----------------------------------------------------------------------------
# Image processor
# ----------------------------------------------------------------------------
def test_resize_with_pad_square():
out = resize_with_pad(torch.ones(1, 3, 224, 224), 224, 224)
assert out.shape == (1, 3, 224, 224)
def test_resize_with_pad_wide():
"""Wide source (H < W): padding applied to top/bottom (rows)."""
out = resize_with_pad(torch.zeros(1, 3, 100, 200), 224, 224)
assert out.shape == (1, 3, 224, 224)
assert out[0, 0, 0, 112].item() == -1.0 # top row is padded
def test_resize_with_pad_tall():
"""Tall source (H > W): padding applied to left/right (cols)."""
out = resize_with_pad(torch.zeros(1, 3, 400, 100), 224, 224)
assert out.shape == (1, 3, 224, 224)
assert out[0, 0, 112, 0].item() == -1.0 # left-most column padded
def test_resize_with_pad_rejects_non_4d():
with pytest.raises(ValueError):
resize_with_pad(torch.zeros(3, 100, 200), 224, 224)
def test_pil_to_tensor():
import numpy as np
from PIL import Image
img = Image.fromarray(np.full((10, 10, 3), 128, dtype=np.uint8))
t = pil_image_to_tensor(img)
assert t.shape == (1, 3, 10, 10)
expected = 128.0 / 255.0 * 2.0 - 1.0
assert abs(t[0, 0, 5, 5].item() - expected) < 1e-4
def test_pil_to_tensor_rgba_to_rgb():
import numpy as np
from PIL import Image
rgba = np.full((8, 8, 4), 200, dtype=np.uint8)
img = Image.fromarray(rgba, mode="RGBA")
t = pil_image_to_tensor(img)
assert t.shape == (1, 3, 8, 8)
def test_make_empty_image():
proc = Pi0ImageProcessor(image_size=224)
empty = proc.make_empty_image()
assert empty.shape == (1, 3, 224, 224)
assert (empty == -1.0).all()
def test_preprocess_single_resizes_ndarray():
import numpy as np
proc = Pi0ImageProcessor(image_size=224)
arr = np.full((120, 120, 3), 100, dtype=np.uint8)
t = proc.preprocess_single(arr)
assert t.shape == (1, 3, 224, 224)
# ----------------------------------------------------------------------------
# build_model_inputs: camera ordering, padding, masks, state
# ----------------------------------------------------------------------------
class _FakeTokenizer:
"""Minimal stand-in for the PaliGemma tokenizer (no download)."""
def __call__(
self, text, padding=None, max_length=None, truncation=None, add_special_tokens=None, return_tensors=None
):
ids = [2] + [ord(c) % 100 for c in text][: max_length - 1]
ids = ids[:max_length]
attn = [1] * len(ids)
while len(ids) < max_length:
ids.append(0)
attn.append(0)
return {"input_ids": ids, "attention_mask": attn}
def test_build_model_inputs_camera_order_and_padding():
import numpy as np
config = Pi0Config.from_model_config(_LEROBOT_CFG) # 3-camera config
device = torch.device("cpu")
# Provide only 2 of the 3 declared cameras.
robot_obs = {
"images": {
"observation.images.base_0_rgb": np.full((224, 224, 3), 50, dtype=np.uint8),
"observation.images.left_wrist_0_rgb": np.full((224, 224, 3), 80, dtype=np.uint8),
},
"state": [0.1] * 14, # raw state shorter than max_state_dim=32
"prompt": "pick up the block",
}
images, image_masks, lang_tokens, lang_masks, state = build_model_inputs(
robot_obs, config, _FakeTokenizer(), device
)
assert len(images) == 3
assert all(img.shape == (1, 3, 224, 224) for img in images)
masks = [bool(m.item()) for m in image_masks]
assert masks == [True, True, False]
# The padded (3rd) camera must be all -1.
assert (images[2] == -1.0).all()
assert lang_tokens.shape == (1, 48)
assert lang_masks.shape == (1, 48)
assert state.shape == (1, 32)
# State padded with zeros beyond the 14 real entries.
assert torch.allclose(state[0, 14:], torch.zeros(18))
def test_build_model_inputs_state_truncated():
config = Pi0Config.from_model_config(_LEROBOT_CFG)
robot_obs = {"images": {}, "state": [0.5] * 40, "prompt": "x"}
_, _, _, _, state = build_model_inputs(robot_obs, config, _FakeTokenizer(), torch.device("cpu"))
assert state.shape == (1, 32)
def test_build_model_inputs_image_key_map():
import numpy as np
cfg = Pi0Config.from_model_config({**_LEROBOT_CFG, "image_key_map": {"cam_high": "observation.images.base_0_rgb"}})
robot_obs = {
"images": {"cam_high": np.full((224, 224, 3), 60, dtype=np.uint8)},
"prompt": "go",
}
_, image_masks, _, _, _ = build_model_inputs(robot_obs, cfg, _FakeTokenizer(), torch.device("cpu"))
masks = [bool(m.item()) for m in image_masks]
assert masks == [True, False, False] # only base_0_rgb (mapped) is present
# ----------------------------------------------------------------------------
# Version-stability / self-consistency (gated on a real pi0_base checkpoint)
#
# The LeRobot parity gate (test_pi0_parity.py) only validates correctness at
# transformers 5.3.0 (where LeRobot loads pi0_base cleanly). Production vllm-omni
# ships transformers 5.6.2, and LeRobot is not a usable reference there (it
# incorrectly loads the SigLIP vision tower at >=5.4). So we pin the kernel's deterministic
# output to golden values established at 5.3.0 and re-check them here: running on
# the production transformers proves the port is version-stable and guards the
# vision_tower / embed_scale version-robust paths against regressions.
#
# PI0_SELFCONSIST_CKPT=/path/to/pi0_base \
# pytest tests/diffusion/models/pi0/test_pi0_units.py -k version_stable -v -s
# ----------------------------------------------------------------------------
import os # noqa: E402
_SELFCONSIST_CKPT = os.environ.get("PI0_SELFCONSIST_CKPT")
# Golden values established at transformers 5.3.0 (where parity = 7.15e-07).
_GOLDEN_SUM = -67.808113
_GOLDEN_STD = 0.252789
_GOLDEN_MEAN = -0.021190
_GOLDEN_EMBED_IMAGE_STD = 2.17603
def _selfconsist_inputs(device):
# NB: draw order is load-bearing — the golden values were established with
# state drawn first, then the 3 images, then lang (shared RNG, seed 0).
g = torch.Generator(device="cpu").manual_seed(0)
b = 2
state = torch.randn(b, 32, generator=g)
images = [torch.rand(b, 3, 224, 224, generator=g) for _ in range(3)]
masks = [torch.ones(b, dtype=torch.bool) for _ in range(3)]
lang = torch.randint(0, 1000, (b, 48), generator=g)
lmask = torch.ones(b, 48, dtype=torch.bool)
ng = torch.Generator(device="cpu").manual_seed(42)
noise = torch.randn(b, 50, 32, generator=ng)
return (
[im.to(device) for im in images],
[m.to(device) for m in masks],
lang.to(device),
lmask.to(device),
state.to(device),
noise.to(device),
)
def _selfconsist_model(device):
import safetensors.torch
model = Pi0ForActionPrediction(
Pi0Config(max_action_dim=32, max_state_dim=32, chunk_size=50, num_inference_steps=10, dtype="float32")
)
state = safetensors.torch.load_file(os.path.join(_SELFCONSIST_CKPT, "model.safetensors"))
filled = model.load_weights(list(state.items()))
expected = {n for n, _ in model.named_parameters() if "rotary_emb" not in n and not n.endswith(".inv_freq")}
missing = expected - (filled or set())
assert not missing, f"{len(missing)} params received no weight (e.g. {sorted(missing)[:3]})"
return model.to(device).eval()
@pytest.mark.slow
@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
@torch.no_grad()
def test_version_stable_actions():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = _selfconsist_model(device)
images, masks, lang, lmask, state, noise = _selfconsist_inputs(device)
actions = (
model.sample_actions(
images=images,
image_masks=masks,
lang_tokens=lang,
lang_masks=lmask,
state=state,
noise=noise,
num_steps=10,
)
.float()
.cpu()
)
s, m, std = actions.sum().item(), actions.mean().item(), actions.std().item()
assert abs(s - _GOLDEN_SUM) < 0.5, f"action sum drifted: {s} vs {_GOLDEN_SUM}"
assert abs(std - _GOLDEN_STD) < 0.01, f"action std drifted: {std} vs {_GOLDEN_STD}"
assert abs(m - _GOLDEN_MEAN) < 0.01, f"action mean drifted: {m} vs {_GOLDEN_MEAN}"
@pytest.mark.slow
@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
@torch.no_grad()
def test_version_stable_embed_image():
"""Directly guards the SigLIP vision-tower load across the >=5.4 flattening."""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = _selfconsist_model(device)
g = torch.Generator(device="cpu").manual_seed(7)
img = torch.rand(1, 3, 224, 224, generator=g, dtype=torch.float32).to(device)
std = model.paligemma_with_expert.embed_image(img).float().std().item()
assert abs(std - _GOLDEN_EMBED_IMAGE_STD) < 0.05, (
f"embed_image std drifted: {std} vs {_GOLDEN_EMBED_IMAGE_STD} — "
"likely a vision_tower weight-load regression under this transformers version."
)
Reference in new issue View Git Blame Copy Permalink