1. 搭建VLA与世界模型联合训练的基础框架

任务目标

搭建一个最小可运行的VLA（视觉-语言-动作）与世界模型联合训练框架。框架包含：一个接收图像和语言指令并输出动作的VLA策略网络、一个预测下一时刻图像特征的世界模型，以及一个交替更新两组参数的联合训练循环。通过本案例，你将得到一个可直接扩展的骨架代码，并验证前向传播、损失计算和参数更新链路完整可用。

环境依赖

需要 Python 3.9+、PyTorch 2.0+、torchvision。无需额外数据集，使用随机张量模拟输入。

核心实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import models

# ---------- 1. VLA策略网络 ----------
class VLA_Policy(nn.Module):
    """输入: 图像(3,64,64) + 语言指令(seq_len, 32) -> 输出: 连续动作(4,)"""
    def __init__(self, lang_dim=32, action_dim=4):
        super().__init__()
        # 视觉编码器: 使用预训练ResNet18的conv层(冻结)
        resnet = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
        self.vision_encoder = nn.Sequential(*list(resnet.children())[:-2])
        for p in self.vision_encoder.parameters():
            p.requires_grad = False
        # 语言编码器: 简单LSTM
        self.lang_lstm = nn.LSTM(lang_dim, 128, batch_first=True)
        # 动作头: 融合视觉+语言特征
        self.fusion = nn.Sequential(
            nn.Linear(512 + 128, 256),
            nn.ReLU(),
            nn.Linear(256, action_dim)
        )

    def forward(self, image, lang_seq):
        # image: (B,3,64,64), lang_seq: (B,T,32)
        vis_feat = self.vision_encoder(image)          # (B,512,2,2)
        vis_feat = vis_feat.mean(dim=[2,3])            # (B,512)
        _, (h_n, _) = self.lang_lstm(lang_seq)         # h_n: (1,B,128)
        lang_feat = h_n.squeeze(0)                     # (B,128)
        fused = torch.cat([vis_feat, lang_feat], dim=1) # (B,640)
        action = self.fusion(fused)                    # (B,4)
        return action

# ---------- 2. 世界模型 ----------
class WorldModel(nn.Module):
    """输入: 当前视觉特征 + 动作 -> 输出: 预测的下一时刻视觉特征"""
    def __init__(self, vis_dim=512, action_dim=4, hidden_dim=256):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(vis_dim + action_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, vis_dim)
        )

    def forward(self, vis_feat, action):
        # vis_feat: (B,512), action: (B,4)
        x = torch.cat([vis_feat, action], dim=1)       # (B,516)
        next_vis_feat = self.fc(x)                     # (B,512)
        return next_vis_feat

# ---------- 3. 联合训练循环 ----------
def joint_training_step(policy, world_model, opt_p, opt_w,
                        image, lang_seq, action_target, next_image):
    # 1) 前向: VLA输出动作
    pred_action = policy(image, lang_seq)              # (B,4)
    # 2) 前向: 世界模型预测下一视觉特征
    with torch.no_grad():
        vis_feat = policy.vision_encoder(image).mean(dim=[2,3])  # (B,512)
        next_vis_feat_target = policy.vision_encoder(next_image).mean(dim=[2,3])
    pred_next_vis_feat = world_model(vis_feat, pred_action.detach())

    # 3) 计算损失
    loss_action = F.mse_loss(pred_action, action_target)
    loss_wm = F.mse_loss(pred_next_vis_feat, next_vis_feat_target)

    # 4) 更新VLA策略(仅动作损失)
    opt_p.zero_grad()
    loss_action.backward()
    opt_p.step()

    # 5) 更新世界模型(仅世界模型损失)
    opt_w.zero_grad()
    loss_wm.backward()
    opt_w.step()

    return loss_action.item(), loss_wm.item()

# ---------- 4. 最小验证入口 ----------
if __name__ == "__main__":
    torch.manual_seed(42)
    B, T = 4, 8
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    policy = VLA_Policy().to(device)
    world_model = WorldModel().to(device)
    opt_p = optim.Adam(policy.fusion.parameters(), lr=1e-4)  # 只训练融合层
    opt_w = optim.Adam(world_model.parameters(), lr=1e-3)

    # 模拟输入
    image = torch.randn(B, 3, 64, 64).to(device)
    next_image = torch.randn(B, 3, 64, 64).to(device)
    lang_seq = torch.randn(B, T, 32).to(device)
    action_target = torch.randn(B, 4).to(device)

    # 执行3步训练
    for step in range(3):
        la, lw = joint_training_step(policy, world_model, opt_p, opt_w,
                                     image, lang_seq, action_target, next_image)
        print(f"Step {step+1}: action_loss={la:.4f}, wm_loss={lw:.4f}")

    # 验证前向输出形状
    with torch.no_grad():
        act = policy(image, lang_seq)
        vf = policy.vision_encoder(image).mean(dim=[2,3])
        nvf = world_model(vf, act)
    print(f"Action shape: {act.shape}, Next vis feat shape: {nvf.shape}")
    assert act.shape == (B, 4), "动作输出维度错误"
    assert nvf.shape == (B, 512), "世界模型输出维度错误"
    print("✓ 前向传播与训练循环验证通过")

运行方式

将上述代码保存为 vla_wm_joint.py，在终端执行：

python vla_wm_joint.py

验证方法与通过标准

• 运行成功标志：程序无异常退出，打印3行训练损失（数值会因随机初始化而变化），并输出 ✓ 前向传播与训练循环验证通过。
• 损失下降观察：由于输入为随机噪声且仅训练3步，损失不要求严格下降；若你增加训练步数（如50步），可观察到 action_loss 和 wm_loss 呈下降趋势，说明两个网络均在学习。
• 形状断言：代码末尾的 assert 语句确保动作输出维度为 (B,4)、世界模型输出维度为 (B,512)，这是框架正确性的基本保证。

关键说明

• 视觉编码器使用冻结的ResNet18，仅训练融合层和世界模型，以最小化计算开销。若需端到端微调，将 requires_grad 改为 True 并将视觉编码器参数加入 opt_p 即可。
• 世界模型预测的是视觉特征而非原始像素，这是VLA联合训练中的常见设计，可降低预测难度并加速收敛。
• 本框架刻意省略了数据加载器、真实环境交互和奖励计算，聚焦于联合训练的核心参数更新逻辑。接入真实数据时，只需替换 image/lang_seq/action_target/next_image 的来源。