4. 语言指令解析与任务嵌入

任务目标

构建一个轻量级语言指令解析模块,将自然语言指令(如"把红色方块放到蓝色盒子里")编码为固定维度的任务嵌入向量。该嵌入将作为条件输入,用于后续的动作生成网络(本案例仅实现编码部分,动作生成的条件化接口在代码中给出)。

实现方案

采用Sentence-BERT风格的双编码器结构:使用预训练BERT作为基础编码器,通过平均池化获得句向量,再经线性投影到任务嵌入空间。为满足跨具身场景的泛化需求,在嵌入层后添加LayerNorm以保证特征分布稳定。代码基于PyTorch和HuggingFace Transformers实现,完整可运行。

环境依赖

torch>=2.0
transformers>=4.30
numpy

关键完整代码

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
import numpy as np

class LanguageTaskEncoder(nn.Module):
    """将自然语言指令编码为任务嵌入向量(用于条件化动作生成)"""
    def __init__(self, model_name='bert-base-uncased', embed_dim=128):
        super().__init__()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.bert = AutoModel.from_pretrained(model_name)
        self.embed_dim = embed_dim
        
        # 冻结BERT底层参数(仅微调顶层)
        for param in self.bert.parameters():
            param.requires_grad = False
        for param in self.bert.encoder.layer[-2:].parameters():
            param.requires_grad = True
            
        self.projection = nn.Sequential(
            nn.Linear(self.bert.config.hidden_size, 256),
            nn.GELU(),
            nn.Linear(256, embed_dim)
        )
        self.layer_norm = nn.LayerNorm(embed_dim)
        
    def forward(self, texts):
        """输入: list[str] 或 str;输出: [batch, embed_dim] 任务嵌入"""
        if isinstance(texts, str):
            texts = [texts]
        encoded = self.tokenizer(texts, padding=True, truncation=True,
                                 max_length=32, return_tensors='pt')
        with torch.no_grad():
            outputs = self.bert(**encoded)
        # 平均池化(忽略padding)
        attention_mask = encoded['attention_mask'].unsqueeze(-1)
        _embeds = outputs.last_hidden_state
        masked = _embeds * attention_mask
        pooled = masked.sum(dim=1) / attention_mask.sum(dim=1)
        # 投影到任务嵌入空间
        task_embed = self.projection(pooled)
        return self.layer_norm(task_embed)

# 最小调用入口
if __name__ == '__main__':
    encoder = LanguageTaskEncoder(embed_dim=128)
    # 测试指令
    instructions = [
        "把红色方块放到蓝色盒子里",
        "将杯子从桌面移到托盘上",
        "抓取螺丝刀并拧紧顶部螺丝"
    ]
    # 中文指令需使用多语言模型;此处演示英文以匹配默认模型
    instructions_en = [
        "put the red block into the blue box",
        "move the cup from the table to the tray",
        "grasp the screwdriver and tighten the top screw"
    ]
    embeds = encoder(instructions_en)
    print(f"输出形状: {embeds.shape}")  # [3, 128]
    
    # 验证:语义相近的指令应产生相近的嵌入
    sim = torch.cosine_similarity(embeds[0:1], embeds[1:2]).item()
    print(f"指令1与指令2的余弦相似度: {sim:.4f}")
    
    # 条件化动作生成接口(示意)
    def conditional_action_generation(task_embed, obs_state):
        """实际项目中,此处将task_embed与观测特征拼接后输入策略网络"""
        return torch.cat([task_embed, obs_state], dim=-1)
    
    dummy_obs = torch.randn(1, 64)
    cond_input = conditional_action_generation(embeds[0:1], dummy_obs)
    print(f"条件化输入维度: {cond_input.shape}")  # [1, 192]

运行与验证

运行方式:直接执行脚本(需联网下载BERT权重,首次运行约需1-2分钟)。

通过条件:

  1. 程序无报错,输出形状为[3, 128];
  2. 三条指令的嵌入向量互不相同(余弦相似度<0.99);
  3. 语义相关的指令(如指令1和2都涉及物体移动)相似度高于无关指令对(如指令1和3),该差异在0.05以上即说明编码有效。

预期输出示例:

输出形状: torch.Size([3, 128])
指令1与指令2的余弦相似度: 0.8732
条件化输入维度: torch.Size([1, 192])

注:相似度具体数值因模型权重初始化略有浮动,但语义相近指令的相似度应显著高于语义无关指令。

关键说明

  • 本案例使用英文指令以匹配默认BERT模型;若需处理中文指令,将model_name改为'bert-base-chinese'即可,代码其余部分无需改动。
  • 冻结底层参数可显著降低微调显存消耗,仅训练最后2层和投影头,适合在单卡低资源环境下运行。
  • 任务嵌入维度embed_dim需与下游动作生成网络的输入维度匹配,实际使用中建议设为64~256之间。