微软研究院×清华大学：靠“潜在动作建模”领跑机器人操作！平均成功率达 90.1%

class ActorModule(nn.Module):    def __init__(self, latent_dim=128, action_dim=8, vision_dim=768, lang_dim=768):        super().__init__()        # 1. VLM特征提取 (PaliGemma)        self.vlm = AutoModelForVision2Seq.from_pretrained("google/paligemma-3b-vision-text")        self.vision_proj = nn.Linear(vision_dim, latent_dim)  # 视觉特征投影        self.lang_proj = nn.Linear(lang_dim, latent_dim)       # 语言特征投影         # 2. ACT-latent (潜在动作扩散模型)        self.act_latent = GaussianDiffusion(            model=TransfusionTransformer(  # 参考Transfusion的扩散Transformer                dim=latent_dim,                depth=6,                heads=8,                input_dim=latent_dim + vision_dim + lang_dim  # 融合特征            ),            loss_type="flow_matching"  # 流匹配损失        )         # 3. ACT-robot (机器人动作扩散模型)        self.act_robot = GaussianDiffusion(            model=TransfusionTransformer(                dim=action_dim,                depth=6,                heads=8,                input_dim=latent_dim + vision_dim + lang_dim  # 含潜在动作特征            ),            loss_type="flow_matching"        )         # 4. 本体感觉上下文嵌入 (适配不同机器人)        self.embodiment_emb = nn.Embedding(num_embeddings=10, embedding_dim=latent_dim)  # 机器人类型嵌入     def forward(self, vision_inputs, lang_inputs, proprio, embodiment_id, actions=None):        # 1. VLM特征提取        vlm_output = self.vlm(            pixel_values=vision_inputs.pixel_values,            input_ids=lang_inputs.input_ids,            return_dict=True        )        vision_feat = self.vision_proj(vlm_output.vision_outputs.last_hidden_state.mean(dim=1))        lang_feat = self.lang_proj(vlm_output.language_outputs.last_hidden_state.mean(dim=1))         # 2. 融合特征 (视觉+语言+本体感觉嵌入)        emb = self.embodiment_emb(embodiment_id)  # 机器人结构嵌入        fused_feat = torch.cat([vision_feat, lang_feat, emb], dim=-1)         # 3. ACT-latent生成潜在动作序列        latent_actions = self.act_latent(fused_feat, cond=fused_feat, labels=actions)         # 4. ACT-robot生成机器人动作        robot_actions = self.act_robot(            torch.cat([latent_actions, fused_feat], dim=-1),            cond=fused_feat,            labels=actions        )        return latent_actions, robot_actions

def train_actor(module: ActorModule, dataloader: DataLoader, is_pretrain=True):    optimizer = optim.AdamW(        module.parameters(),        lr=3e-4,        weight_decay=0.05,        betas=(0.9, 0.95)    )    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10000)    loss_fn = FlowMatchingLoss()  # 流匹配损失函数     module.train()    for epoch in range(100):        total_loss = 0.0        for batch in dataloader:            # 输入解析            vision = batch["vision"]            lang = batch["lang"]            proprio = batch["proprio"]            actions = batch["actions"]            embodiment_id = torch.zeros(len(batch), dtype=torch.long)  # 机器人ID             # 前向传播            latent_actions, robot_actions = module(                vision, lang, proprio, embodiment_id, actions            )             # 损失计算 (联合潜在动作和机器人动作损失)            loss_latent = loss_fn(latent_actions, actions)            loss_robot = loss_fn(robot_actions, actions)            loss = loss_latent + loss_robot  # 联合损失             # 反向传播            optimizer.zero_grad()            loss.backward()            optimizer.step()            scheduler.step()             total_loss += loss.item()        print(f"Epoch {epoch}, Loss: {total_loss/len(dataloader):.4f}")     # 预训练后冻结trunk，微调仅更新stem和head    if not is_pretrain:        for param in module.vlm.parameters():            param.requires_grad = False