1. VLA模型基础架构搭建

任务说明

构建一个最小可运行的VLA（Vision-Language-Action）模型基础框架，包含视觉编码器、语言编码器和动作解码器三个核心模块，并实现完整的前向传播流程。输入为图像和文本指令，输出为连续动作向量。

环境依赖

需要 Python 3.9+、PyTorch 2.0+、torchvision、transformers。若使用预训练权重需联网下载，本案例为演示自包含实现，不依赖外部预训练模型。

核心实现

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer

class VisionEncoder(nn.Module):
    """视觉编码器：将图像转换为视觉序列"""
    def __init__(self, img_size=224, patch_size=16, in_ch=3, embed_dim=512):
        super().__init__()
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_ch, embed_dim, kernel_size=patch_size, stride=patch_size)
        self.cls_ = nn.Parameter(torch.randn(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.randn(1, self.n_patches + 1, embed_dim))
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):
        # x: (B, C, H, W)
        B = x.shape[0]
        x = self.proj(x)  # (B, D, H/P, W/P)
        x = x.flatten(2).transpose(1, 2)  # (B, N, D)
        x = torch.cat([self.cls_.expand(B, -1, -1), x], dim=1)
        x = x + self.pos_embed
        return self.norm(x)  # (B, N+1, D)

class LanguageEncoder(nn.Module):
    """语言编码器：将文本指令编码为语言序列"""
    def __init__(self, vocab_size=30522, embed_dim=512, max_len=64):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.pos_embed = nn.Parameter(torch.randn(1, max_len, embed_dim))
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, input_ids, attention_mask=None):
        # input_ids: (B, L)
        B, L = input_ids.shape
        x = self.embed(input_ids)  # (B, L, D)
        x = x + self.pos_embed[:, :L, :]
        x = self.norm(x)
        if attention_mask is not None:
            x = x * attention_mask.unsqueeze(-1)
        return x  # (B, L, D)

class ActionDecoder(nn.Module):
    """动作解码器：融合视觉与语言特征，输出连续动作"""
    def __init__(self, embed_dim=512, action_dim=7, nhead=8, num_layers=2):
        super().__init__()
        self.action_dim = action_dim
        # 可学习的动作查询向量
        self.query = nn.Parameter(torch.randn(1, 1, embed_dim))
        decoder_layer = nn.TransformerDecoderLayer(d_model=embed_dim, nhead=nhead, batch_first=True)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.fc_out = nn.Linear(embed_dim, action_dim)

    def forward(self, vis_feat, lang_feat):
        # vis_feat: (B, Nv, D), lang_feat: (B, Nl, D)
        B = vis_feat.shape[0]
        query = self.query.expand(B, -1, -1)  # (B, 1, D)
        # 将视觉和语言特征拼接作为decoder的memory
        memory = torch.cat([vis_feat, lang_feat], dim=1)  # (B, Nv+Nl, D)
        out = self.decoder(query, memory)  # (B, 1, D)
        action = self.fc_out(out).squeeze(1)  # (B, action_dim)
        return action

class VLA(nn.Module):
    """VLA模型：视觉编码器 + 语言编码器 + 动作解码器"""
    def __init__(self, img_size=224, patch_size=16, embed_dim=512,
                 vocab_size=30522, action_dim=7, max_len=64):
        super().__init__()
        self.vision_encoder = VisionEncoder(img_size, patch_size, 3, embed_dim)
        self.language_encoder = LanguageEncoder(vocab_size, embed_dim, max_len)
        self.action_decoder = ActionDecoder(embed_dim, action_dim)

    def forward(self, image, input_ids, attention_mask=None):
        vis_feat = self.vision_encoder(image)      # (B, Nv, D)
        lang_feat = self.language_encoder(input_ids, attention_mask)  # (B, Nl, D)
        action = self.action_decoder(vis_feat, lang_feat)  # (B, action_dim)
        return action

# ========== 最小可运行示例 ==========
if __name__ == "__main__":
    # 初始化模型
    model = VLA(img_size=224, patch_size=16, embed_dim=512,
                vocab_size=30522, action_dim=7, max_len=64)

    # 构造输入：随机图像 (B=2, C=3, H=224, W=224)
    image = torch.randn(2, 3, 224, 224)

    # 构造文本输入：使用transformers的tokenizer
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    texts = ["move the arm forward", "grasp the red block"]
    encodings = tokenizer(texts, padding=True, truncation=True, max_length=64, return_tensors="pt")
    input_ids = encodings["input_ids"]
    attention_mask = encodings["attention_mask"]

    # 前向传播
    with torch.no_grad():
        action = model(image, input_ids, attention_mask)

    print(f"输入图像形状: {image.shape}")
    print(f"输入文本形状: {input_ids.shape}")
    print(f"输出动作形状: {action.shape}")
    print(f"输出动作值范围: [{action.min().item():.3f}, {action.max().item():.3f}]")

运行方式

pip install torch torchvision transformers
python vla_model.py

验证方法

运行后观察输出：

• 输出动作张量形状应为 (2, 7)，对应2个样本各7维动作向量（如机械臂的关节角度或末端位姿）；
• 前向传播无报错，且输出值在合理数值范围内（非NaN/Inf）；
• 修改输入文本或图像，输出动作应随之变化，说明模型成功融合了视觉与语言信息。

关键说明

• 本案例为教学演示，使用随机初始化的编码器，未加载预训练权重，因此输出动作无实际物理意义，仅验证框架正确性；
• 实际应用中，需将视觉编码器替换为预训练ViT/ResNet，语言编码器替换为预训练BERT，并在机器人数据集上训练动作解码器；
• 动作维度 action_dim=7 对应常见机械臂的7自由度关节角度，可根据实际机器人调整。