[vla] fix: pi05 models does not apply scale factor for language embeddings (#33367)

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/benchmarks/bench_pi05_openpi.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 import argparse
 import asyncio
+import copy
 import dataclasses
 import json
 import sys
@@ -53,7 +54,7 @@ PROFILES = {
         openpi_checkpoint="gs://openpi-assets/checkpoints/pi05_libero",
         prompt="pick up the object",
         sglang_action_horizon=50,
-        openpi_action_horizon=10,
+        openpi_action_horizon=50,
         action_dim=32,
         output_action_dim=7,
     ),
@@ -231,15 +232,13 @@ def _make_aloha_observation(
         },
         "prompt": prompt,
     }
-    sglang_state = np.zeros((32,), dtype=np.float32)
-    sglang_state[: state.shape[0]] = state
     sglang_observation = {
         "images": {
             "base_0_rgb": np.transpose(cam_high, (1, 2, 0)),
             "left_wrist_0_rgb": np.transpose(cam_left, (1, 2, 0)),
             "right_wrist_0_rgb": np.transpose(cam_right, (1, 2, 0)),
         },
-        "state": sglang_state,
+        "state": state,
     }
     return openpi_obs, sglang_observation
 
@@ -264,6 +263,33 @@ def build_observations(
     return openpi_observations, sglang_observations
 
 
+def build_openpi_model_inputs(
+    policy,
+    observations: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    """Convert raw robot observations to the exact inputs consumed by OpenPI."""
+    model_inputs = []
+    for observation in observations:
+        transformed = policy._input_transform(copy.deepcopy(observation))
+        images = {
+            name: np.asarray(value) for name, value in transformed["image"].items()
+        }
+        model_inputs.append(
+            {
+                "images": images,
+                "image_masks": {
+                    name: bool(value)
+                    for name, value in transformed["image_mask"].items()
+                },
+                "camera_order": tuple(images),
+                "state": np.asarray(transformed["state"]),
+                "tokens": np.asarray(transformed["tokenized_prompt"]),
+                "token_masks": np.asarray(transformed["tokenized_prompt_mask"]),
+            }
+        )
+    return model_inputs
+
+
 def _json_tensor(array: np.ndarray) -> dict[str, Any]:
     return {
         "dtype": str(array.dtype),
@@ -290,6 +316,13 @@ def build_sglang_payload(
         "images": encoded_images,
         "state": _json_tensor(np.asarray(observation["state"], dtype=np.float32)),
     }
+    if "tokens" in observation:
+        encoded_observation["tokens"] = np.asarray(observation["tokens"]).tolist()
+        encoded_observation["token_masks"] = np.asarray(
+            observation["token_masks"]
+        ).tolist()
+        encoded_observation["image_masks"] = observation["image_masks"]
+        encoded_observation["camera_order"] = list(observation["camera_order"])
     if noise is not None:
         encoded_observation["noise"] = _json_tensor(noise.astype(np.float32))
     return {
@@ -326,6 +359,11 @@ def build_sglang_python_payload(
         },
         "state": np.asarray(observation["state"], dtype=np.float32),
     }
+    if "tokens" in observation:
+        encoded_observation["tokens"] = np.asarray(observation["tokens"])
+        encoded_observation["token_masks"] = np.asarray(observation["token_masks"])
+        encoded_observation["image_masks"] = observation["image_masks"]
+        encoded_observation["camera_order"] = observation["camera_order"]
     if noise is not None:
         encoded_observation["noise"] = noise.astype(np.float32)
     return {
@@ -366,6 +404,11 @@ def build_sglang_openpi_ws_payload(
     }
     for key, value in observation["images"].items():
         payload[f"observation.images.{key}"] = np.asarray(value)
+    if "tokens" in observation:
+        payload["tokens"] = np.asarray(observation["tokens"])
+        payload["token_masks"] = np.asarray(observation["token_masks"])
+        payload["image_masks"] = observation["image_masks"]
+        payload["camera_order"] = observation["camera_order"]
     if noise is not None:
         payload["observation.noise"] = noise.astype(np.float32)
     return payload
@@ -791,11 +834,11 @@ def _openpi_infer(policy, observation: dict[str, Any], noise: np.ndarray | None)
     return policy.infer(observation, noise=noise)
 
 
-def _openpi_direct_batch(
+def _openpi_model_batch(
     policy,
     observations: list[dict[str, Any]],
     noises: list[np.ndarray] | None,
-):
+) -> tuple[np.ndarray, np.ndarray]:
     import jax
     import numpy as onp
     from openpi.models import model as openpi_model
@@ -845,6 +888,15 @@ def _openpi_direct_batch(
         actions_np = onp.asarray(actions)
         states_np = onp.asarray(inputs["state"])
 
+    return actions_np, states_np
+
+
+def _openpi_direct_batch(
+    policy,
+    observations: list[dict[str, Any]],
+    noises: list[np.ndarray] | None,
+):
+    actions_np, states_np = _openpi_model_batch(policy, observations, noises)
     outputs = []
     for idx in range(actions_np.shape[0]):
         outputs.append(
@@ -858,6 +910,15 @@ def _openpi_direct_batch(
     return outputs
 
 
+def _openpi_model_actions(
+    policy,
+    observation: dict[str, Any],
+    noise: np.ndarray,
+) -> np.ndarray:
+    actions, _ = _openpi_model_batch(policy, [observation], [noise])
+    return actions[0]
+
+
 def run_openpi_policy(
     policy,
     observations: list[dict[str, Any]],
@@ -922,6 +983,12 @@ def run_openpi_policy(
         precision["output_action_dtype"] = str(first_actions.dtype)
         precision["output_action_shape"] = list(first_actions.shape)
 
+    first_model_actions = (
+        _openpi_model_actions(policy, observations[0], noise)
+        if observations and noise is not None
+        else None
+    )
+
     return {
         "single": _stats_ms(single_latencies),
         "batch": _stats_ms(batch_latencies),
@@ -930,6 +997,7 @@ def run_openpi_policy(
             key: _stats_ms(values) for key, values in policy_timings.items()
         },
         "first_output": single_outputs[0] if single_outputs else None,
+        "first_model_actions": first_model_actions,
         "batch_mode": batch_mode,
         "precision": precision,
     }
@@ -951,10 +1019,18 @@ def _openpi_actions(output: dict[str, Any]) -> np.ndarray | None:
 
 def compare_first_actions(
     sglang_output: dict[str, Any] | None,
-    openpi_output: dict[str, Any] | None,
+    openpi_output: dict[str, Any] | np.ndarray | None,
 ) -> dict[str, Any]:
     sglang_actions = _sglang_actions(sglang_output)
-    openpi_actions = _openpi_actions(openpi_output)
+    openpi_actions = (
+        _openpi_actions(openpi_output)
+        if isinstance(openpi_output, dict)
+        else (
+            np.asarray(openpi_output, dtype=np.float32)
+            if openpi_output is not None
+            else None
+        )
+    )
     if sglang_actions is None or openpi_actions is None:
         return {"available": False}
     horizon = min(sglang_actions.shape[0], openpi_actions.shape[0])
@@ -1051,6 +1127,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--disable-prefix-cache", action="store_true")
     parser.add_argument("--disable-cuda-graph", action="store_true")
     parser.add_argument("--deterministic-noise", action="store_true")
+    parser.add_argument("--action-max-abs-diff", type=float, default=0.05)
+    parser.add_argument("--action-mean-abs-diff", type=float, default=0.005)
     parser.add_argument("--skip-sglang", action="store_true")
     parser.add_argument("--skip-openpi", action="store_true")
     parser.add_argument("--output-file", default="")
@@ -1096,6 +1174,24 @@ def main() -> None:
             dtype=np.float32,
         )
 
+    openpi_policy = None
+    if not args.skip_openpi:
+        openpi_policy = create_openpi_policy(
+            openpi_config,
+            openpi_checkpoint,
+            pytorch_device=args.openpi_device,
+            num_inference_steps=args.num_inference_steps,
+            pytorch_compile_mode=(
+                None
+                if args.openpi_pytorch_compile_mode == "none"
+                else args.openpi_pytorch_compile_mode
+            ),
+        )
+        if not args.skip_sglang and args.deterministic_noise:
+            sglang_observations = build_openpi_model_inputs(
+                openpi_policy,
+                openpi_observations,
+            )
     payloads = []
     if args.skip_sglang:
         pass
@@ -1138,20 +1234,6 @@ def main() -> None:
             for observation in sglang_observations
         ]
 
-    openpi_policy = None
-    if not args.skip_openpi:
-        openpi_policy = create_openpi_policy(
-            openpi_config,
-            openpi_checkpoint,
-            pytorch_device=args.openpi_device,
-            num_inference_steps=args.num_inference_steps,
-            pytorch_compile_mode=(
-                None
-                if args.openpi_pytorch_compile_mode == "none"
-                else args.openpi_pytorch_compile_mode
-            ),
-        )
-
     sglang_result = None
     if args.skip_sglang:
         pass
@@ -1210,9 +1292,20 @@ def main() -> None:
         "repeats": args.repeats,
         "warmup": args.warmup,
         "deterministic_noise": args.deterministic_noise,
-        "action_diff": compare_first_actions(
-            None if sglang_result is None else sglang_result.get("first_output"),
-            None if openpi_result is None else openpi_result.get("first_output"),
+        "action_diff": (
+            compare_first_actions(
+                None if sglang_result is None else sglang_result.get("first_output"),
+                (
+                    None
+                    if openpi_result is None
+                    else openpi_result.get("first_model_actions")
+                ),
+            )
+            if args.deterministic_noise
+            else {
+                "available": False,
+                "reason": "use --deterministic-noise for action comparison",
+            }
         ),
         "sglang": (
             None
@@ -1229,7 +1322,7 @@ def main() -> None:
             else {
                 key: value
                 for key, value in openpi_result.items()
-                if key not in ("first_output",)
+                if key not in ("first_output", "first_model_actions")
             }
         ),
     }
@@ -1238,6 +1331,19 @@ def main() -> None:
             json.dump(result, f, indent=2, sort_keys=True)
     print_summary(result)
 
+    action_diff = result["action_diff"]
+    if action_diff.get("available") and (
+        action_diff["max_abs_diff"] > args.action_max_abs_diff
+        or action_diff["mean_abs_diff"] > args.action_mean_abs_diff
+    ):
+        raise AssertionError(
+            "Pi0.5 action mismatch: "
+            f"max_abs_diff={action_diff['max_abs_diff']:.6f} "
+            f"(threshold {args.action_max_abs_diff:.6f}), "
+            f"mean_abs_diff={action_diff['mean_abs_diff']:.6f} "
+            f"(threshold {args.action_mean_abs_diff:.6f})"
+        )
+
 
 if __name__ == "__main__":
     main()

python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py

@@ -1336,6 +1336,8 @@ class Pi05CoreModel(nn.Module):
             att_masks += [0] * num_image_embs
 
         lang_emb = self.paligemma_with_expert.embed_language_tokens(tokens)
+        # Match OpenPI's Pi0.5 prefix embedding semantics.
+        lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])
         embs.append(lang_emb)
         pad_masks.append(token_masks)
         att_masks += [0] * lang_emb.shape[1]

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/pi05_preprocess.py

@@ -14,7 +14,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConf
 from sglang.multimodal_gen.runtime.vla.observation import VLAObservationBatch
 
 
-def _tensor_from_image(value: Any) -> torch.Tensor:
+def _tensor_from_image(value: Any) -> tuple[torch.Tensor, bool, bool]:
     if isinstance(value, torch.Tensor):
         tensor = value.detach()
         if tensor.ndim == 4:
@@ -33,14 +33,16 @@ def _tensor_from_image(value: Any) -> torch.Tensor:
             )
         is_integer = not tensor.is_floating_point()
         tensor = tensor.to(dtype=torch.float32)
-        if is_integer or tensor.max() > 2.0:
+        is_byte_scaled = is_integer or tensor.max() > 2.0
+        is_normalized = not is_byte_scaled and tensor.min() < 0.0
+        if is_byte_scaled:
             tensor = tensor / 255.0
-        return tensor
+        return tensor, is_byte_scaled, is_normalized
 
     if isinstance(value, Image.Image):
         image = value.convert("RGB")
         arr = np.asarray(image, dtype=np.float32) / 255.0
-        return torch.from_numpy(arr).permute(2, 0, 1)
+        return torch.from_numpy(arr).permute(2, 0, 1), True, False
 
     if isinstance(value, (np.ndarray, list)):
         arr = np.asarray(value)
@@ -57,15 +59,21 @@ def _tensor_from_image(value: Any) -> torch.Tensor:
             )
         is_integer = not tensor.is_floating_point()
         tensor = tensor.to(dtype=torch.float32)
-        if is_integer or tensor.max() > 2.0:
+        is_byte_scaled = is_integer or tensor.max() > 2.0
+        is_normalized = not is_byte_scaled and tensor.min() < 0.0
+        if is_byte_scaled:
             tensor = tensor / 255.0
-        return tensor
+        return tensor, is_byte_scaled, is_normalized
 
     raise TypeError(f"Unsupported Pi05 image type: {type(value)}")
 
 
 def _resize_with_pad_image_tensor(
-    tensor: torch.Tensor, size: tuple[int, int]
+    tensor: torch.Tensor,
+    size: tuple[int, int],
+    *,
+    round_to_uint8: bool = False,
+    pad_value: float = 0.0,
 ) -> torch.Tensor:
     height, width = size
     if tensor.shape[-2:] == (height, width):
@@ -80,16 +88,32 @@ def _resize_with_pad_image_tensor(
         mode="bilinear",
         align_corners=False,
     )[0]
+    if round_to_uint8:
+        # openpi rounds resized byte images before mapping them to [-1, 1]
+        tensor = torch.round(tensor * 255.0).clamp_(0.0, 255.0) / 255.0
     pad_h0, rem_h = divmod(height - resized_height, 2)
     pad_w0, rem_w = divmod(width - resized_width, 2)
     return F.pad(
         tensor,
         (pad_w0, pad_w0 + rem_w, pad_h0, pad_h0 + rem_h),
         mode="constant",
-        value=0.0,
+        value=pad_value,
     )
 
 
+def _preprocess_image(value: Any, size: tuple[int, int]) -> torch.Tensor:
+    tensor, is_byte_scaled, is_normalized = _tensor_from_image(value)
+    tensor = _resize_with_pad_image_tensor(
+        tensor,
+        size,
+        round_to_uint8=is_byte_scaled,
+        pad_value=-1.0 if is_normalized else 0.0,
+    )
+    if is_normalized:
+        return tensor.clamp_(-1.0, 1.0)
+    return tensor * 2.0 - 1.0
+
+
 class Pi05Preprocessor:
     def __init__(self, config: Pi05PipelineConfig):
         self.config = config
@@ -142,9 +166,7 @@ class Pi05Preprocessor:
             value = raw_images.get(key)
             is_present = value is not None and bool(image_masks_in.get(key, True))
             if is_present:
-                tensor = _tensor_from_image(value)
-                tensor = _resize_with_pad_image_tensor(tensor, self.config.image_size)
-                tensor = tensor * 2.0 - 1.0
+                tensor = _preprocess_image(value, self.config.image_size)
             else:
                 channels = 3
                 height, width = self.config.image_size

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -370,7 +370,7 @@ PI05_ACTION_CI_sampling_params = DiffusionSamplingParams(
     extras={
         "action_horizon": 50,
         "action_dim": 32,
-        "state_dim": 32,
+        "state_dim": 14,
         "image_size": 64,
         "num_inference_steps": 2,
         "seed": 0,

python/sglang/multimodal_gen/test/single_test_file/test_pi05_e2e.py

@@ -58,7 +58,7 @@ def _action_request_kwargs(tag: str) -> dict:
         "state": np.linspace(
             -0.5,
             0.5,
-            _env_int("SGLANG_PI05_E2E_STATE_DIM", 32),
+            _env_int("SGLANG_PI05_E2E_STATE_DIM", 14),
             dtype=np.float32,
         ),
         "noise": rng.standard_normal((action_horizon, action_dim)).astype(np.float32),
@@ -138,9 +138,9 @@ def test_pi05_python_action_e2e(pi05_generator):
     _assert_action_output(output)
 
 
-def test_pi05_python_action_consistency(pi05_generator):
-    first = pi05_generator.generate_action(_action_request_kwargs("consistency"))
-    second = pi05_generator.generate_action(_action_request_kwargs("consistency"))
+def test_pi05_python_action_repeatability_and_cache(pi05_generator):
+    first = pi05_generator.generate_action(_action_request_kwargs("repeatability"))
+    second = pi05_generator.generate_action(_action_request_kwargs("repeatability"))
     _assert_action_output(first, expect_cache_hit=False)
     _assert_action_output(second, expect_cache_hit=True)
 

python/sglang/multimodal_gen/test/test_utils.py

@@ -39,7 +39,7 @@ logger = init_logger(__name__)
 # NPU/ascend) is read from sgl-project/ci-data-diffusion, where the GT-gen workflows
 # publish.
 SGL_TEST_FILES_CI_DATA_REPO = "sgl-project/ci-data-diffusion"
-SGL_TEST_FILES_CI_DATA_REVISION = "739c6c9b7cb972149cc3472cf19fe4bb29cf15c3"
+SGL_TEST_FILES_CI_DATA_REVISION = "d05810e3ea3eff1d137dec723f6e66d9c11b470f"
 
 # The NPU pin is kept as a separate branch so ascend GT can be bumped independently
 # when it's regenerated on its own cadence.

python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py

@@ -2,12 +2,14 @@
 
 from types import SimpleNamespace
 
+import numpy as np
 import torch
 from torch import nn
 
 import sglang.multimodal_gen.runtime.models.vlas.pi05_policy as pi05_policy_module
 from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
 from sglang.multimodal_gen.runtime.models.vlas.pi05_core import (
+    Pi05CoreModel,
     Pi05SiglipAttention,
     patch_siglip_vision_attention_to_native,
 )
@@ -15,6 +17,10 @@ from sglang.multimodal_gen.runtime.models.vlas.pi05_policy import (
     Pi05CheckpointManifest,
     Pi05PolicyModel,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.pi05_preprocess import (
+    _preprocess_image,
+    _resize_with_pad_image_tensor,
+)
 from sglang.multimodal_gen.runtime.vla.denoise_cuda_graph import (
     VLADenoiseGraphRunner,
     _CapturedDenoiseGraph,
@@ -172,3 +178,64 @@ def test_siglip_attention_patch_uses_native_wrapper_once():
 
     assert isinstance(first, Pi05SiglipAttention)
     assert layer.self_attn is first
+
+
+def test_prefix_language_embedding_matches_openpi_scale():
+    image_embedding = torch.ones(1, 2, 8)
+    language_embedding = torch.full((1, 3, 8), 0.25)
+    model = SimpleNamespace(
+        paligemma_with_expert=SimpleNamespace(
+            embed_images=lambda images: [image_embedding],
+            embed_language_tokens=lambda tokens: language_embedding,
+        )
+    )
+
+    embeddings, _, _ = Pi05CoreModel.embed_prefix(
+        model,
+        images=[torch.zeros(1, 3, 4, 4)],
+        image_masks=[torch.ones(1, dtype=torch.bool)],
+        tokens=torch.ones(1, 3, dtype=torch.long),
+        token_masks=torch.ones(1, 3, dtype=torch.bool),
+    )
+
+    torch.testing.assert_close(
+        embeddings[:, 2:],
+        language_embedding * (language_embedding.shape[-1] ** 0.5),
+    )
+
+
+def test_uint8_resize_rounds_before_normalization():
+    image = torch.tensor([[[0.0, 1.0], [2.0, 3.0]]]) / 255.0
+
+    resized = _resize_with_pad_image_tensor(
+        image,
+        (3, 3),
+        round_to_uint8=True,
+    )
+
+    expected = (
+        torch.round(
+            torch.nn.functional.interpolate(
+                image[None], size=(3, 3), mode="bilinear", align_corners=False
+            )[0]
+            * 255.0
+        )
+        / 255.0
+    )
+    torch.testing.assert_close(resized, expected, rtol=0.0, atol=0.0)
+
+
+def test_normalized_float_image_is_not_normalized_twice():
+    image = np.full((2, 4, 3), -0.5, dtype=np.float32)
+
+    preprocessed = _preprocess_image(image, (4, 4))
+
+    assert preprocessed.shape == (3, 4, 4)
+    torch.testing.assert_close(
+        preprocessed[:, 1:3],
+        torch.full((3, 2, 4), -0.5),
+    )
+    torch.testing.assert_close(
+        preprocessed[:, (0, 3)],
+        torch.full((3, 2, 4), -1.0),
+    )