`modeling_prismatic.py` — 这是 OpenVLA-OFT 模型的核心实现文件（约 2000 行）¶

文件路径: verl/experimental/vla/models/openvla_oft/modeling_prismatic.py 模块路径: verl.experimental.vla.models.openvla_oft.modeling_prismatic

文件概述¶

这是 OpenVLA-OFT 模型的核心实现文件（约 2000 行），包含视觉骨干、多模态投影器、条件生成模型和动作预测模型。这是整个 VLA 系统中最复杂的文件。

架构概览¶

输入: 图像 + 文本提示 ("What action should the robot take to ...?")
      ↓
┌─────────────────────────────────────────┐
│  PrismaticVisionBackbone (SigLIP/DINOv2)│  ← 视觉特征提取
│  ↓                                       │
│  PrismaticProjector (MLP)                │  ← 投影到语言空间
│  ↓                                       │
│  LLM (Llama/Mistral)                    │  ← 多模态推理
│  ↓                                       │
│  动作 Token Logits → 离散化 → 连续动作  │  ← 动作预测
└─────────────────────────────────────────┘

一、PrismaticVisionBackbone - 视觉骨干¶

支持单骨干和双骨干融合（SigLIP + DINOv2）：

class PrismaticVisionBackbone(nn.Module):
    def __init__(self, use_fused_vision_backbone, image_sizes,
                 timm_model_ids, timm_override_act_layers):
        # 主骨干（SigLIP）
        self.featurizer = timm.create_model(timm_model_ids[0], ...)
        self.embed_dim = self.featurizer.embed_dim

        # 融合骨干（DINOv2，可选）
        if use_fused_vision_backbone:
            self.fused_featurizer = timm.create_model(timm_model_ids[1], ...)
            self.embed_dim += self.fused_featurizer.embed_dim

    def forward(self, pixel_values):
        if not self.use_fused_vision_backbone:
            return self.featurizer(pixel_values)

        # 融合模式：分割通道 -> 分别提取 -> 拼接
        img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
        patches = self.featurizer(img)
        patches_fused = self.fused_featurizer(img_fused)
        return torch.cat([patches, patches_fused], dim=2)

二、PrismaticProjector - 多模态投影器¶

将视觉特征投影到语言模型的嵌入空间：

class PrismaticProjector(nn.Module):
    def __init__(self, use_fused_vision_backbone, vision_dim, llm_dim):
        if not use_fused_vision_backbone:
            # 简单 MLP: vision_dim -> llm_dim
            self.fc1 = nn.Linear(vision_dim, llm_dim)
            self.fc2 = nn.Linear(llm_dim, llm_dim)
        else:
            # 更大的 MLP: vision_dim -> 4*vision_dim -> llm_dim -> llm_dim
            self.fc1 = nn.Linear(vision_dim, 4 * vision_dim)
            self.fc2 = nn.Linear(4 * vision_dim, llm_dim)
            self.fc3 = nn.Linear(llm_dim, llm_dim)

三、OpenVLAForActionPrediction - 动作预测模型¶

这是最终的用户接口类，继承自 PrismaticForConditionalGeneration。

动作离散化¶

OpenVLA 将连续动作离散化为 256 个 bin：

class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
    def __init__(self, config):
        super().__init__(config)
        # 创建 256 个均匀分布的 bin，范围 [-1, 1]
        self.bins = np.linspace(-1, 1, config.n_action_bins)
        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
        # 词表大小（去掉填充部分）
        self.vocab_size = config.text_config.vocab_size - config.pad_to_multiple_of

verl 推理方法 generate_action_verl¶

这是 verl 框架使用的主要推理方法，支持批量和采样：

def generate_action_verl(self, input_ids, pixel_values, attention_mask,
                         padding_idx, do_sample, unnorm_key, temperature):
    # 1. 创建 labels（用于动作 mask 计算）
    labels = input_ids.clone()
    labels[:] = IGNORE_INDEX

    # 2. 添加占位动作 token 和停止 token
    input_ids, attention_mask = self._prepare_input_for_action_prediction(
        input_ids, attention_mask)

    # 3. 处理左填充
    # ... (重新排列 padding 到左侧)

    # 4. 提取视觉特征并投影
    patch_features = self.vision_backbone(pixel_values)
    projected_patch_embeddings = self.projector(patch_features)

    # 5. 构建多模态嵌入
    # [BOS] + [视觉 patches] + [文本 tokens] + [动作 tokens] + [STOP]

    # 6. 前向传播获取 logits
    language_model_output = self.language_model(inputs_embeds=...)

    # 7. 从 logits 中提取动作预测
    # 关键切片：只取词表末尾 256+64 中的 256 个 token
    action_logits = logits[..., -256-64:-64]

动作反归一化¶

def _unnormalize_actions(self, normalized_actions, unnorm_key):
    """将 [-1, 1] 范围的归一化动作转换回原始尺度"""
    action_stats = self.get_action_stats(unnorm_key)
    action_high = np.array(action_stats["q99"])
    action_low = np.array(action_stats["q01"])
    actions = 0.5 * (normalized_actions + 1) * (action_high - action_low) + action_low
    return actions

四、多模态嵌入构建流程¶

输入序列: [BOS] [text_tokens...] [action_placeholders] [STOP]
              ↓
嵌入层:   [BOS_emb] [text_embs...] [action_embs(清零)] [stop_emb]
              ↓ 插入视觉嵌入（在 BOS 后面）
最终:     [BOS_emb] [patch_embs...] [text_embs...] [action_embs] [stop_emb]
              ↓
attention_mask 也相应扩展

核心类列表¶

名称	说明
`PrismaticVisionBackbone`	视觉骨干（TIMM，支持融合）
`PrismaticProjector`	视觉到语言的投影器
`PrismaticPreTrainedModel`	HF 预训练模型基类
`PrismaticForConditionalGeneration`	条件生成模型
`OpenVLAForActionPrediction`	动作预测模型（最终接口）

与其他模块的关系¶

使用 constants.py 的动作常量
使用 train_utils.py 的 mask 计算函数
使用 configuration_prismatic.py 的配置类
被 naive_rollout_rob.py 调用做推理
被 register_vla_models.py 注册到 HuggingFace

小结¶

这个文件实现了从图像+文本到机器人动作的完整推理链。核心创新是将连续动作离散化为词表中的 token，使得可以用标准的语言模型架构来预测机器人动作。理解 vocab_size-256-64:-64 这个切片范围是理解整个模型的关键。

modeling_prismatic.py — 这是 OpenVLA-OFT 模型的核心实现文件（约 2000 行）¶