sgl-project/sglang · error · ValueError

Pi05 noise must have shape {expected}, got {tuple(noise_tens

Error message

Pi05 noise must have shape {expected}, got {tuple(noise_tensor.shape)}

What it means

When explicit flow-matching noise is supplied to the Pi05 stage, it must match the expected action-noise layout of exactly [1, action_horizon, action_dim] (a 2-D [H, D] input is auto-unsqueezed to [1, H, D]). Any other shape is rejected because the denoiser expects noise aligned token-for-token with the action trajectory.

Source

Thrown at python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/pi05_preprocess.py:200

            if state_tensor.ndim == 1:
                state_tensor = state_tensor.unsqueeze(0)
            if state_tensor.shape[0] != 1:
                raise ValueError("Pi05 v1 expects one state vector per request")
            if state_tensor.shape[-1] > self.config.state_dim:
                raise ValueError(
                    f"Pi05 state dim must be <= {self.config.state_dim}, "
                    f"got {state_tensor.shape[-1]}"
                )

        noise = raw_observation.get("noise")
        noise_tensor = None
        if noise is not None:
            noise_tensor = torch.as_tensor(noise, dtype=torch.float32)
            if noise_tensor.ndim == 2:
                noise_tensor = noise_tensor.unsqueeze(0)
            expected = (1, self.config.action_horizon, self.config.action_dim)
            if tuple(noise_tensor.shape) != expected:
                raise ValueError(
                    f"Pi05 noise must have shape {expected}, "
                    f"got {tuple(noise_tensor.shape)}"
                )

        tokens = raw_observation.get("tokens")
        if tokens is None:
            tokens = raw_observation.get("tokenized_prompt")
        token_masks = raw_observation.get("token_masks")
        if token_masks is None:
            token_masks = raw_observation.get("tokenized_prompt_mask")
        if tokens is not None:
            tokens_tensor = torch.as_tensor(tokens, dtype=torch.long)
            if tokens_tensor.ndim == 1:
                tokens_tensor = tokens_tensor.unsqueeze(0)
            if token_masks is None:
                token_masks_tensor = tokens_tensor != self.tokenizer.pad_token_id
            else:
                token_masks_tensor = torch.as_tensor(token_masks, dtype=torch.bool)

View on GitHub (pinned to 0132848349)

Solutions

  1. Generate noise as torch.randn(1, cfg.action_horizon, cfg.action_dim) (or let the stage sample it by omitting 'noise').
  2. Verify config.action_horizon / config.action_dim match the checkpoint you serve.
  3. If you have [H, D], it is auto-unsqueezed — but never pass batched noise.

Example fix

# before
obs = {"noise": torch.randn(1, cfg.action_dim), ...}

# after
obs = {"noise": torch.randn(1, cfg.action_horizon, cfg.action_dim), ...}
Defensive patterns

Strategy: validation

Validate before calling

if noise is not None:
    n = torch.as_tensor(noise, dtype=torch.float32)
    if n.ndim == 2: n = n.unsqueeze(0)
    assert tuple(n.shape) == (1, cfg.action_horizon, cfg.action_dim)

Type guard

def noise_shape_ok(noise, cfg) -> bool:
    if noise is None: return True
    n = torch.as_tensor(noise)
    if n.ndim == 2: n = n.unsqueeze(0)
    return tuple(n.shape) == (1, cfg.action_horizon, cfg.action_dim)

Prevention

When it happens

Trigger: Passing raw_observation['noise'] with shape [1, D] (missing horizon axis), [B, H, D] with B>1, [H, D, W], or any tensor whose dims don't equal (1, config.action_horizon, config.action_dim).

Common situations: Sampling noise with the wrong horizon/dim from a config mismatch (action_dim or action_horizon changed between training and serving); reusing cached noise tensors after changing chunk length; batched generation attempts.

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/3e5ee8a3a0df280d. Report an issue: GitHub.