4. VLA策略网络生成动作指令

任务说明

本案例实现一个视觉-语言-动作（VLA）策略网络的核心推理路径：输入一张仿真环境图像和一条自然语言指令，网络输出对应的离散动作指令（例如前进、左转、抓取等）。重点在于构建跨模态融合模块与动作解码头，并验证在给定输入下网络能产生与指令语义一致的动作分布。

环境依赖

需要 Python 3.9+、PyTorch 2.0+、torchvision、transformers（仅用于加载预训练CLIP文本编码器）。若本机无预训练权重，代码会自动降级为随机初始化文本编码器，不影响网络结构验证。

核心实现

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextModel, CLIPTokenizer

class VLA_Policy(nn.Module):
    """视觉-语言-动作策略网络：图像+文本 → 离散动作分布"""
    def __init__(self, num_actions=6, embed_dim=256, text_dim=512):
        super().__init__()
        # 视觉编码器（轻量CNN，输入3x64x64）
        self.vision_encoder = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=5, stride=2, padding=2),  # 32x32x32
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # 64x16x16
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),# 128x8x8
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),                          # 128x1x1
            nn.Flatten(),                                          # 128
            nn.Linear(128, embed_dim)
        )
        # 文本编码器（使用预训练CLIP，若加载失败则随机初始化）
        try:
            self.text_encoder = CLIPTextModel.from_pretrained("/clip-vit-base-patch32")
            self.text_proj = nn.Linear(512, embed_dim)
            self.use_pretrained = True
        except Exception:
            self.text_encoder = nn.LSTM(input_size=text_dim, hidden_size=embed_dim,
                                        batch_first=True, bidirectional=True)
            self.text_proj = nn.Linear(2*embed_dim, embed_dim)
            self.use_pretrained = False
        # 跨模态融合 + 动作解码
        self.fusion_gate = nn.Linear(embed_dim*2, embed_dim)
        self.action_head = nn.Linear(embed_dim, num_actions)

    def encode_text(self, text_list, device):
        """将文本列表编码为embedding向量"""
        if self.use_pretrained:
            tokenizer = CLIPTokenizer.from_pretrained("/clip-vit-base-patch32")
            inputs = tokenizer(text_list, return_tensors="pt", padding=True,
                               truncation=True, max_length=32).to(device)
            with torch.no_grad():
                last_hidden = self.text_encoder(**inputs).last_hidden_state
            # 取[CLS]位置向量
            text_feat = last_hidden[:, 0, :]
        else:
            # 随机初始化LSTM回退方案：将字符映射为向量
            char_embed = nn.Embedding(128, text_dim).to(device)
            seqs = []
            for t in text_list:
                seqs.append(torch.tensor([min(ord(c), 127) for c in t], device=device))
            padded = nn.utils.rnn.pad_sequence(seqs, batch_first=True)
            emb = char_embed(padded)
            _, (h_n, _) = self.text_encoder(emb)
            text_feat = torch.cat([h_n[0], h_n[1]], dim=1)  # 双向拼接
        return self.text_proj(text_feat)

    def forward(self, image, text_list):
        """
        image: (B, 3, 64, 64) 像素值归一化到[0,1]
        text_list: 长度为B的字符串列表
        返回: (B, num_actions) 动作logits
        """
        device = image.device
        vis_feat = self.vision_encoder(image)                          # (B, embed_dim)
        txt_feat = self.encode_text(text_list, device)                 # (B, embed_dim)
        # 门控融合
        combined = torch.cat([vis_feat, txt_feat], dim=1)              # (B, 2*embed_dim)
        gate = torch.sigmoid(self.fusion_gate(combined))               # (B, embed_dim)
        fused = gate * vis_feat + (1 - gate) * txt_feat                # (B, embed_dim)
        logits = self.action_head(fused)                               # (B, num_actions)
        return logits

# 动作标签映射（0-5）
ACTION_NAMES = ["前进", "左转", "右转", "抓取", "放下", "停止"]

def generate_action(net, image_tensor, instruction):
    """输入单张图像和指令，返回动作名称"""
    net.eval()
    with torch.no_grad():
        logits = net(image_tensor.unsqueeze(0), [instruction])
        action_id = torch.argmax(logits, dim=1).item()
    return ACTION_NAMES[action_id]

运行与验证

将上述代码保存为 vla_policy.py，然后运行以下验证脚本：

import torch
from vla_policy import VLA_Policy, generate_action, ACTION_NAMES

# 初始化网络
net = VLA_Policy(num_actions=6)
# 构造模拟输入：3通道64x64随机图像（实际场景中为仿真器渲染帧）
dummy_image = torch.rand(3, 64, 64)
# 测试指令
test_instructions = ["请向前移动", "向左转", "抓取红色物体"]
for instr in test_instructions:
    action = generate_action(net, dummy_image, instr)
    print(f"指令: {instr} → 动作: {action}")

# 验证网络输出维度正确
batch_images = torch.rand(4, 3, 64, 64)
batch_texts = ["前进", "左转", "右转", "停止"]
logits = net(batch_images, batch_texts)
assert logits.shape == (4, 6), f"输出维度错误: {logits.shape}"
print("维度验证通过: (4, 6)")

通过条件：

  1. 脚本无异常运行结束，输出4条动作预测结果（由于随机初始化，动作内容不要求语义正确）；
  2. 维度断言通过，即网络对批量输入返回 (batch_size, 6) 的logits张量；
  3. 若要验证语义对齐能力，需加载预训练权重并输入真实仿真图像，观察输出动作是否与指令匹配——本案例仅验证网络结构与前向通路正确性。

关键说明

本实现采用门控融合机制，视觉特征与文本特征通过可学习的门控权重动态组合，动作头输出6类离散动作的logits。预训练CLIP文本编码器提供语义理解能力；若网络不可用则自动回退到字符级LSTM，保证代码在任何环境可运行。实际部署时需将图像替换为仿真器真实观测，并配合行为克隆损失进行端到端训练。

← 上一章 📖 返回目录 下一章 →