PI05 模型架构与推理流程详解

1
2
3
4

class ActionSelectKwargs(TypedDict, total=False):
    inference_delay: int | None
    prev_chunk_left_over: Tensor | None
    execution_horizon: int | None

1
2
3
4
5
6
7
8
9

def get_safe_dtype(target_dtype, device_type):
    if device_type == "mps" and target_dtype == torch.float64:
        return torch.float32
    if device_type == "cpu":
        if target_dtype == torch.bfloat16:
            return torch.float32
        if target_dtype == torch.float64:
            return torch.float64
    return target_dtype

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

def create_sinusoidal_pos_embedding(
    time: torch.Tensor, dimension: int,
    min_period: float, max_period: float, device="cpu"
) -> Tensor:
    if dimension % 2 != 0:
        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
    if time.ndim != 1:
        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")

    dtype = get_safe_dtype(torch.float64, device.type)
    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
    period = min_period * (max_period / min_period) ** fraction

    scaling_factor = 1.0 / period * 2 * math.pi
    sin_input = scaling_factor[None, :] * time[:, None]
    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)

1
2
3
4
5

def sample_beta(alpha, beta, bsize, device):
    alpha_t = torch.tensor(alpha, dtype=torch.float32)
    beta_t = torch.tensor(beta, dtype=torch.float32)
    dist = torch.distributions.Beta(alpha_t, beta_t)
    return dist.sample((bsize,)).to(device)

1
2
3
4
5

def make_att_2d_masks(pad_masks, att_masks):
    cumsum = torch.cumsum(att_masks, dim=1)
    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
    return att_2d_masks & pad_2d_masks

1
2
3
4

cumsum = torch.cumsum(att_masks, dim=1)                   # [B, N] 累积和
att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]   # [B, N, N] 广播比较
pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]  # padding token 完全不可见
return att_2d_masks & pad_2d_masks                         # 交集

1
2
3
4
5
6
7

def clone_past_key_values(past_key_values):
    return DynamicCache(
        tuple(
            (keys.clone(), values.clone(), sliding_window)
            for keys, values, sliding_window in past_key_values
        )
    )

1
2
3
4

def pad_vector(vector, new_dim):
    if vector.shape[-1] >= new_dim:
        return vector
    return F.pad(vector, (0, new_dim - vector.shape[-1]))

1

def resize_with_pad_torch(images, height, width, mode="bilinear"):

1

def compute_layer_complete(inputs_embeds, attention_mask, position_ids, adarms_cond, layers, rotary_emb):

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

query_states = []; key_states = []; value_states = []; gates = []
for i, hidden_states in enumerate(inputs_embeds):
    layer = layers[i]
    hidden_states, gate = layernorm_forward(layer.input_layernorm, hidden_states, adarms_cond[i])
    gates.append(gate)
    input_shape = hidden_states.shape[:-1]
    hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
    query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    key_state   = layer.self_attn.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    query_states.append(query_state); key_states.append(key_state); value_states.append(value_state)

1
2
3

query_states = torch.cat(query_states, dim=2)   # cat on seq dim → [B, num_heads, prefix_len + suffix_len, head_dim]
key_states   = torch.cat(key_states,   dim=2)
value_states = torch.cat(value_states, dim=2)

1
2
3

dummy_tensor = torch.zeros(query_states.shape[0], query_states.shape[2], query_states.shape[-1], ...)
cos, sin = rotary_emb(dummy_tensor, position_ids)
query_states, key_states = modeling_gemma.apply_rotary_pos_emb(query_states, key_states, cos, sin, unsqueeze_dim=1)

1
2
3
4
5
6
7
8

batch_size = query_states.shape[0]
paligemma_layer = layers[0]
scaling = paligemma_layer.self_attn.scaling
att_output, _ = modeling_gemma.eager_attention_forward(
    paligemma_layer.self_attn, query_states, key_states, value_states, attention_mask, scaling,
)
head_dim = paligemma_layer.self_attn.head_dim
att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim)

1
2
3
4
5
6
7
8

outputs_embeds = []
start_pos = 0
for i, hidden_states in enumerate(inputs_embeds):
    layer = layers[i]
    end_pos = start_pos + hidden_states.shape[1]
    if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
        att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
    out_emb = layer.self_attn.o_proj(att_output[:, start_pos:end_pos])

1
2
3

    out_emb = _gated_residual(hidden_states, out_emb, gates[i])  # 第一个残差
    after_first_residual = out_emb.clone()
    out_emb, gate = layernorm_forward(layer.post_attention_layernorm, out_emb, adarms_cond[i])

1
2
3
4
5
6

    if layer.mlp.up_proj.weight.dtype == torch.bfloat16:
        out_emb = out_emb.to(dtype=torch.bfloat16)
    out_emb = layer.mlp(out_emb)
    out_emb = _gated_residual(after_first_residual, out_emb, gate)  # 第二个残差
    outputs_embeds.append(out_emb)
    start_pos = end_pos

1
2
3
4
5
6

if use_gradient_checkpointing:
    inputs_embeds = torch.utils.checkpoint.checkpoint(
        compute_layer_complete, inputs_embeds, attention_mask,
        position_ids, adarms_cond, use_reentrant=False,
        preserve_rng_state=False, layers=layers, rotary_emb=rotary_emb,
    )

1
2
3
4

class GemmaConfig:
    def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim):
        self.width = width; self.depth = depth; self.mlp_dim = mlp_dim
        self.num_heads = num_heads; self.num_kv_heads = num_kv_heads; self.head_dim = head_dim

1
2
3
4

class PaliGemmaWithExpertModel(nn.Module):
    def __init__(self, vlm_config, action_expert_config, use_adarms=None,
                 precision="bfloat16", image_size=DEFAULT_IMAGE_SIZE,
                 freeze_vision_encoder=False, train_expert_only=False):

1
2
3
4
5
6
7
8

vlm_config_hf = CONFIG_MAPPING["paligemma"]()
vlm_config_hf._vocab_size = 257152              # Gemma tokenizer vocab size
vlm_config_hf.image_token_index = 257152        # <image> 特殊 token
vlm_config_hf.text_config.hidden_size = vlm_config.width
vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim
# ...
vlm_config_hf.text_config.use_adarms = use_adarms[0]
vlm_config_hf.text_config.adarms_cond_dim = vlm_config.width if use_adarms[0] else None

1
2
3

action_expert_config_hf = CONFIG_MAPPING["gemma"](...)
action_expert_config_hf.use_adarms = use_adarms[1]
action_expert_config_hf.adarms_cond_dim = action_expert_config.width if use_adarms[1] else None

1
2
3

self.paligemma = PaliGemmaForConditionalGenerationWithPiGemma(config=vlm_config_hf)
self.gemma_expert = PiGemmaForCausalLM(config=action_expert_config_hf)
self.gemma_expert.model.embed_tokens = None   # Expert 不需要 embedding 层

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

def to_bfloat16_for_selected_params(self, precision):
    if precision == "bfloat16":
        self.to(dtype=torch.bfloat16)
    elif precision == "float32":
        self.to(dtype=torch.float32); return

    params_to_keep_float32 = [
        "vision_tower", "multi_modal_projector",
        "input_layernorm", "post_attention_layernorm", "model.norm",
    ]
    for name, param in self.named_parameters():
        if any(selector in name for selector in params_to_keep_float32):
            param.data = param.data.to(dtype=torch.float32)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

def _set_requires_grad(self):
    if self.freeze_vision_encoder:
        for param in self.paligemma.model.vision_tower.parameters():
            param.requires_grad = False
    if self.train_expert_only:
        for param in self.paligemma.parameters():
            param.requires_grad = False

def train(self, mode=True):
    super().train(mode)
    if self.freeze_vision_encoder: self.paligemma.model.vision_tower.eval()
    if self.train_expert_only:     self.paligemma.eval()

1
2
3
4
5
6
7
8
9

def embed_image(self, image):
    out_dtype = image.dtype
    if image.dtype != torch.float32:
        image = image.to(torch.float32)
    image_outputs = self.paligemma.model.get_image_features(image)
    features = image_outputs.pooler_output
    if features.dtype != out_dtype:
        features = features.to(out_dtype)
    return features

1
2

def embed_language_tokens(self, tokens):
    return self.paligemma.model.language_model.get_input_embeddings()(tokens)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

if inputs_embeds[1] is None:
    prefix_output = self.paligemma.model.language_model.forward(
        inputs_embeds=inputs_embeds[0],
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_values=past_key_values,
        use_cache=use_cache,              # use_cache=True → 返回 KV Cache
        adarms_cond=adarms_cond[0],       # None（Prefix 塔无 AdaRMS）
    )
    prefix_past_key_values = prefix_output.past_key_values
    prefix_output = prefix_output.last_hidden_state
    suffix_output = None

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

elif inputs_embeds[0] is None:
    suffix_output = self.gemma_expert.model.forward(
        inputs_embeds=inputs_embeds[1],
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_values=past_key_values,  # 携带 Prefix KV Cache
        use_cache=False,                   # 不需要再缓存
        adarms_cond=adarms_cond[1],        # time_emb → AdaRMS
    )
    suffix_output = suffix_output.last_hidden_state
    prefix_output = None
    prefix_past_key_values = None

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

else:
    # 逐层联合计算
    for layers in zip(paligemma_layers, gemma_expert_layers, strict=True):
        inputs_embeds = compute_layer_complete(
            inputs_embeds, attention_mask, position_ids,
            adarms_cond, layers=layers, rotary_emb=rotary_emb,
        )
    # 分别最终 norm
    outputs_embeds = []
    for i, hidden_states in enumerate(inputs_embeds):
        out_emb, _ = layernorm_forward(final_norms[i], hidden_states, adarms_cond[i])
        outputs_embeds.append(out_emb)
    prefix_output = outputs_embeds[0]
    suffix_output = outputs_embeds[1]
    prefix_past_key_values = None

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

class PI05Pytorch(nn.Module):
    def __init__(self, config: PI05Config, rtc_processor=None):
        self.config = config
        paligemma_config = get_gemma_config(config.paligemma_variant)
        action_expert_config = get_gemma_config(config.action_expert_variant)

        self.paligemma_with_expert = PaliGemmaWithExpertModel(
            paligemma_config, action_expert_config,
            use_adarms=[False, True],      # ← PI05 硬编码：Prefix无AdaRMS，Suffix有
            precision=config.dtype,
            image_size=config.image_resolution[0],
            freeze_vision_encoder=config.freeze_vision_encoder,
            train_expert_only=config.train_expert_only,
        )

1
2
3
4

        self.action_in_proj  = nn.Linear(config.max_action_dim, action_expert_config.width)
        self.action_out_proj = nn.Linear(action_expert_config.width, config.max_action_dim)
        self.time_mlp_in     = nn.Linear(action_expert_config.width, action_expert_config.width)
        self.time_mlp_out    = nn.Linear(action_expert_config.width, action_expert_config.width)

1
2
3
4

        if config.compile_model:
            torch.set_float32_matmul_precision("high")
            self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
            self.forward = torch.compile(self.forward, mode=config.compile_mode)

1
2
3
4
5

def gradient_checkpointing_enable(self):
    self.gradient_checkpointing_enabled = True
    self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = True
    self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = True
    self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True

1
2
3
4
5

def _apply_checkpoint(self, func, *args, **kwargs):
    if self.gradient_checkpointing_enabled and self.training:
        return torch.utils.checkpoint.checkpoint(
            func, *args, use_reentrant=False, preserve_rng_state=False, **kwargs)
    return func(*args, **kwargs)

1
2
3

def _prepare_attention_masks_4d(self, att_2d_masks):
    att_2d_masks_4d = att_2d_masks[:, None, :, :]
    return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)

1
2

def sample_noise(self, shape, device):
    return torch.normal(mean=0.0, std=1.0, size=shape, dtype=torch.float32, device=device)

1
2
3
4
5
6
7

def sample_time(self, bsize, device):
    time_beta = sample_beta(
        self.config.time_sampling_beta_alpha,
        self.config.time_sampling_beta_beta, bsize, device
    )
    time = time_beta * self.config.time_sampling_scale + self.config.time_sampling_offset
    return time.to(dtype=torch.float32, device=device)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

def embed_prefix(self, images, img_masks, tokens, masks):
    embs = []; pad_masks = []; att_masks = []

    # 逐个图像编码
    for img, img_mask in zip(images, img_masks, strict=True):
        img_emb = self._apply_checkpoint(image_embed_func, img)    # SigLIP + Projector
        bsize, num_img_embs = img_emb.shape[:2]
        embs.append(img_emb)
        pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs))  # [B] → [B, N_img]
        att_masks += [0] * num_img_embs                                   # 图像 token 全在同一个 attn group

1
2
3
4
5

    # 语言 token 编码
    lang_emb = self._apply_checkpoint(lang_embed_func, tokens)
    embs.append(lang_emb)
    pad_masks.append(masks)                            # 来自 tokenizer 的 attention_mask
    att_masks += [0] * num_lang_embs                   # 语言 token 也在同一个 attn group

1
2
3
4
5

    embs = torch.cat(embs, dim=1)        # [B, N_img1 + N_img2 + ... + N_lang, D]
    pad_masks = torch.cat(pad_masks, dim=1)
    att_masks = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)
    att_masks = att_masks[None, :].expand(bsize, len(att_masks))  # [1, N] → [B, N]
    return embs, pad_masks, att_masks

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

def embed_suffix(self, noisy_actions, timestep):
    # 1. 时间步 → 正弦编码
    time_emb = create_sinusoidal_pos_embedding(
        timestep,
        self.action_in_proj.out_features,   # = Expert hidden size D
        min_period=self.config.min_period,
        max_period=self.config.max_period,
        device=timestep.device,
    )
    time_emb = time_emb.type(dtype=timestep.dtype)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

    # 2. 动作 → 线性投影
    action_emb = self._apply_checkpoint(action_proj_func, noisy_actions)  # [B, chunk, D]

    # 3. 时间嵌入 → 2层 MLP
    def time_mlp_func(time_emb):
        x = self.time_mlp_in(time_emb)    # D → D
        x = F.silu(x)
        x = self.time_mlp_out(x)          # D → D
        return F.silu(x)                  # 最终激活
    time_emb = self._apply_checkpoint(time_mlp_func, time_emb)
    adarms_cond = time_emb                # ← 这就是 AdaRMS 的条件向量

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

    # 4. Suffix = 动作嵌入（不含时间嵌入的显式拼接！）
    action_time_emb = action_emb
    embs.append(action_time_emb)

    bsize, action_time_dim = action_time_emb.shape[:2]
    action_time_mask = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)
    pad_masks.append(action_time_mask)

    # 5. Attention mask: 第一个 token att_mask=1（可看到 Prefix），其余 att_mask=0（因果链）
    att_masks += [1] + ([0] * (self.config.chunk_size - 1))

1
2
3
4
5

def forward(self, images, img_masks, tokens, masks, actions, noise, time) -> Tensor:
    # Flow Matching 加噪
    time_expanded = time[:, None, None]                    # [B] → [B, 1, 1]
    x_t = time_expanded * noise + (1 - time_expanded) * actions   # 线性插值
    u_t = noise - actions                                  # 目标速度场

1
2
3
4
5
6
7
8

    # Prefix + Suffix 嵌入
    prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, tokens, masks)
    suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, time)

    # 精度对齐
    if self.paligemma_with_expert.paligemma.model.language_model.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16:
        suffix_embs = suffix_embs.to(dtype=torch.bfloat16)
        prefix_embs = prefix_embs.to(dtype=torch.bfloat16)

1
2
3
4
5
6

    # 拼接 Pad/Att masks + 构造全局 attention mask
    pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1)
    att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)
    att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
    position_ids = torch.cumsum(pad_masks, dim=1) - 1
    att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

    # 联合前向（Prefix + Suffix 同时进模型）
    def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond):
        (_, suffix_out), _ = self.paligemma_with_expert.forward(
            attention_mask=att_2d_masks_4d,
            position_ids=position_ids,
            past_key_values=None,                       # 训练不缓存
            inputs_embeds=[prefix_embs, suffix_embs],    # 联合模式
            use_cache=False,
            adarms_cond=[None, adarms_cond],             # Prefix无AdaRMS, Suffix有
        )
        return suffix_out

    suffix_out = self._apply_checkpoint(forward_func, ...)
    suffix_out = suffix_out[:, -self.config.chunk_size :]   # 只取 suffix 部分
    suffix_out = suffix_out.to(dtype=torch.float32)          # 恢复到 fp32 计算 loss

1
2

    v_t = self._apply_checkpoint(action_out_proj_func, suffix_out)
    return F.mse_loss(u_t, v_t, reduction="none")   # [B, chunk, action_dim]

1
2
3
4
5
6
7

@torch.no_grad()
def sample_actions(self, images, img_masks, tokens, masks, noise=None, num_steps=None, **kwargs):
    bsize = tokens.shape[0]; device = tokens.device

    if noise is None:
        actions_shape = (bsize, self.config.chunk_size, self.config.max_action_dim)
        noise = self.sample_noise(actions_shape, device)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

    prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, tokens, masks)
    prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
    prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
    prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks)

    # 强制 eager attention（torch.compile 下 SDPA 可能与 DynamicCache 不兼容）
    self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager"

    _, past_key_values = self.paligemma_with_expert.forward(
        attention_mask=prefix_att_2d_masks_4d,
        position_ids=prefix_position_ids,
        past_key_values=None,
        inputs_embeds=[prefix_embs, None],  # ← Prefix-Only 模式
        use_cache=True,
    )

1
2
3
4
5

    dt = -1.0 / num_steps                       # 负步长：从 t=1 走到 t=0
    x_t = noise
    for step in range(num_steps):
        time = 1.0 + step * dt                  # 1.0, 1-dt, 1-2dt, ..., 0
        time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)

1
2
3
4
5
6
7

        def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
            return self.denoise_step(
                prefix_pad_masks=prefix_pad_masks,
                past_key_values=past_key_values,
                x_t=input_x_t,
                timestep=current_timestep,
            )

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

        if self._rtc_enabled():
            v_t = self.rtc_processor.denoise_step(
                x_t=x_t, prev_chunk_left_over=..., inference_delay=...,
                time=time, original_denoise_step_partial=denoise_step_partial_call,
                execution_horizon=...,
            )
        else:
            v_t = denoise_step_partial_call(x_t)

        x_t = x_t + dt * v_t                   # Euler 步进 {#fn-euler-step}

1
2
3
4

        if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
            self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)

    return x_t                                  # 形状 [B, chunk_size, max_action_dim]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

def denoise_step(self, prefix_pad_masks, past_key_values, x_t, timestep):
    # 1. Suffix 嵌入
    suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, timestep)

    # 2. 构造 full attention mask（Suffix attend to Prefix + Suffix causal）
    batch_size = prefix_pad_masks.shape[0]
    suffix_len = suffix_pad_masks.shape[1]
    prefix_len = prefix_pad_masks.shape[1]

    prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
    suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
    full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)

1
2
3

    # 3. Position IDs = Prefix 总长度 + Suffix 内的 cumsum
    prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]   # [B, 1]
    position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1

1
2
3

    # 4. 强制 eager attention + 深拷贝 KV Cache
    self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"
    past_key_values = clone_past_key_values(past_key_values)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

    # 5. Suffix-Only 前向
    outputs_embeds, _ = self.paligemma_with_expert.forward(
        attention_mask=full_att_2d_masks_4d,
        position_ids=position_ids,
        past_key_values=past_key_values,        # 携带 Prefix KV Cache
        inputs_embeds=[None, suffix_embs],       # Suffix-Only 模式
        use_cache=False,                         # 不需要缓存 Suffix KV
        adarms_cond=[None, adarms_cond],
    )
    suffix_out = outputs_embeds[1]               # Suffix 塔的输出
    suffix_out = suffix_out[:, -self.config.chunk_size :]   # 截取 suffix 序列
    suffix_out = suffix_out.to(dtype=torch.float32)
    return self.action_out_proj(suffix_out)       # D → action_dim

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

class PI05Policy(PreTrainedPolicy):
    config_class = PI05Config
    name = "pi05"

    def __init__(self, config):
        require_package("transformers", extra="pi")
        super().__init__(config)
        config.validate_features()
        self.init_rtc_processor()
        self.model = PI05Pytorch(config, rtc_processor=self.rtc_processor)
        if config.gradient_checkpointing:
            self.model.gradient_checkpointing_enable()
        self.model.to(config.device)
        self.reset()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

@classmethod
def from_pretrained(cls, pretrained_name_or_path, *, config=None, force_download=False, ..., strict=True, **kwargs):
    # 1. 打印 disclaimer
    print("The PI05 model is a direct port of the OpenPI implementation...")

    # 2. 获取/创建 config
    if config is None:
        config = PreTrainedConfig.from_pretrained(pretrained_name_or_path, ...)

    # 3. 创建模型实例（随机权重）
    model = cls(config, **kwargs)

    # 4. 下载 safetensors 文件
    from transformers.utils import cached_file
    resolved_file = cached_file(pretrained_name_or_path, "model.safetensors", ...)
    from safetensors.torch import load_file
    original_state_dict = load_file(resolved_file)

1

    fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config)

1
2
3
4
5
6
7

    remapped_state_dict = {}
    for key, value in fixed_state_dict.items():
        if not key.startswith("model."):
            new_key = f"model.{key}"          # ← 加 "model." 前缀
            remapped_state_dict[new_key] = value
        else:
            remapped_state_dict[key] = value

1

    missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)

1
2
3
4
5
6

if re.match(r"paligemma_with_expert\.gemma_expert\.model\.layers\.\d+\."
            r"(input_layernorm|post_attention_layernorm)\.weight", key):
    expert_uses_adarms = getattr(self.model.paligemma_with_expert.gemma_expert.config, "use_adarms", False)
    if expert_uses_adarms:
        logging.warning(f"Skipping layer norm key (adaRMS mismatch): {key}")
        continue

1
2
3
4

if key.startswith("action_time_mlp_in."):
    new_key = key.replace("action_time_mlp_in.", "time_mlp_in.")
elif key.startswith("action_time_mlp_out."):
    new_key = key.replace("action_time_mlp_out.", "time_mlp_out.")

1
2
3

if key.startswith("state_proj."):
    logging.warning(f"Skipping state_proj key in pi05 mode: {key}")
    continue

1
2
3

if key in ("model.paligemma_with_expert.paligemma.lm_head.weight",
           "paligemma_with_expert.paligemma.lm_head.weight"):
    fixed_state_dict["model.paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"] = value.clone()

1
2
3
4

def _preprocess_images(self, batch):
    device = next(self.parameters()).device
    present_img_keys = [key for key in self.config.image_features if key in batch]
    missing_img_keys = [key for key in self.config.image_features if key not in batch]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

    for key in present_img_keys:
        img = batch[key].to(device).to(torch.float32)

        # 格式检测 + 转换
        is_channels_first = img.shape[1] == 3
        if is_channels_first:
            img = img.permute(0, 2, 3, 1)        # [B,C,H,W] → [B,H,W,C]

        # Resize + Pad
        if img.shape[1:3] != self.config.image_resolution:
            img = resize_with_pad_torch(img, *self.config.image_resolution)

        # 归一化: [0,1] → [-1,1]
        img = img * 2.0 - 1.0

        if is_channels_first:
            img = img.permute(0, 3, 1, 2)        # [B,H,W,C] → [B,C,H,W]

        images.append(img)
        img_masks.append(torch.ones(bsize, dtype=torch.bool, device=device))

1
2
3
4
5

    for _num_empty_cameras in range(len(missing_img_keys)):
        img = torch.ones_like(img) * -1          # 全 -1 填充（SigLIP 预处理后的"黑"）
        mask = torch.zeros_like(mask)            # mask = 0（不可见）
        images.append(img)
        img_masks.append(mask)

1
2

def prepare_action(self, batch):
    return pad_vector(batch[ACTION], self.config.max_action_dim)

1
2
3
4
5
6
7
8

@torch.no_grad()
def select_action(self, batch):
    assert not self._rtc_enabled()  # RTC 不兼容单步动作选择

    if len(self._action_queue) == 0:
        actions = self.predict_action_chunk(batch)[:, :self.config.n_action_steps]
        self._action_queue.extend(actions.transpose(0, 1))  # (B, steps, dim) → 按 step 拆入队列
    return self._action_queue.popleft()

1
2
3
4
5
6
7

@torch.no_grad()
def predict_action_chunk(self, batch, **kwargs):
    images, img_masks = self._preprocess_images(batch)
    tokens, masks = batch["observation_language_tokens"], batch["observation_language_attention_mask"]
    actions = self.model.sample_actions(images, img_masks, tokens, masks, **kwargs)
    actions = actions[:, :, :original_action_dim]   # 解 pad 到真实动作维度 {#fn-action-proj}
    return actions