WAM-Trainer世界动作模型训练实战：IDM/FDM模块化架构与DeepSpeed ZeRO-2多卡训练

最近把WAM-Trainer这个世界动作模型训练平台完整跑了一遍，从5.18亿帧第一视角数据的清洗对齐，到IDM逆动力学和FDM前动力学的模块化训练，再到8×A100上用DeepSpeed ZeRO-2跑通，最后WebSocket策略服务器把推理延迟压到150ms。LIBERO跑到99.3%，这里把工程上几个关键模块的代码和配置记录下来，方便做具身智能方向的同学参考。

整个项目解决的核心问题是：真机机器人数据太贵，怎么用世界模型生成想象视频来扩充训练数据。思路是视频生成模型先生成未来帧，IDM从想象帧反推动作，FDM做前向一致性验证，Qwen3-VL负责语言指令理解。下面按模块拆代码。

数据集适配器基类

我们接了12个数据集，每个数据集的观测空间、动作空间、相机配置都不一样，所以先抽了一个基类，所有数据集适配器继承它，统一输出到80维动作空间：

class BaseDatasetAdapter:
    """所有机器人数据集的统一适配器基类"""
    def __init__(self, dataset_name, action_dim=80):
        self.dataset_name = dataset_name
        self.action_dim = action_dim  # 统一80维动作空间
        self.obs_keys = []
        self.action_scale = 1.0

    def load_episode(self, episode_path):
        """加载单个episode，返回统一格式的trajectory"""
        raw_data = self._load_raw(episode_path)
        frames = self._extract_frames(raw_data)       # B,T,3,H,W
        actions = self._align_actions(raw_data)        # B,T,action_dim
        lang = self._parse_language(raw_data)          # str
        frames, actions = self._sync_timestamp(frames, actions)
        return {"frames": frames, "actions": actions, "lang": lang}

    def _align_actions(self, raw_actions):
        """将不同本体的动作映射到80维统一空间，缺失补零"""
        T, orig_dim = raw_actions.shape
        aligned = np.zeros((T, self.action_dim), dtype=np.float32)
        copy_dim = min(orig_dim, self.action_dim)
        aligned[:, :copy_dim] = raw_actions[:, :copy_dim] * self.action_scale
        return aligned

    def _sync_timestamp(self, frames, actions, fps_video=30, hz_action=10):
        """视频fps和动作hz不一致时做时间戳重采样"""
        ratio = fps_video / hz_action
        action_resampled = np.interp(
            np.arange(len(frames)) / ratio,
            np.arange(len(actions)),
            actions[:, 0]
        )
        return frames, action_resampled

这里踩过最大的坑是时间戳对齐。Open X-Embodiment里不同子数据集的视频帧率和动作频率都不一样，不对齐直接训，模型学到的动作和画面就是错位的。上面_sync_timestamp用线性插值把动作重采样到视频帧率，一开始没做这步，LIBERO只有82%左右。

IDM逆动力学模型结构

IDM吃进去当前帧和未来想象帧，输出中间的动作序列。结构上用一个双流编码器分别编码两帧，然后在token维度交叉注意力，最后接动作头输出未来16步的80维动作：

import torch
import torch.nn as nn

class InverseDynamicsModule(nn.Module):
    """IDM逆动力学：给定当前帧和未来帧，反推动作"""
    def __init__(self, frame_encoder, action_dim=80, chunk_size=16, hidden=1024):
        super().__init__()
        self.frame_encoder = frame_encoder  # 视频帧编码器，来自Wan2.2主干
        self.cross_attn = nn.MultiheadAttention(hidden, num_heads=8, batch_first=True)
        self.action_head = nn.Sequential(
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, chunk_size * action_dim)
        )
        self.chunk_size = chunk_size
        self.action_dim = action_dim
        self.smooth_weight = 0.1  # 时序平滑loss权重

    def forward(self, frame_curr, frame_future, lang_embed=None):
        tok_curr = self.frame_encoder(frame_curr)     # B, N, D
        tok_future = self.frame_encoder(frame_future)  # B, N, D
        fused, _ = self.cross_attn(tok_future, tok_curr, tok_curr)
        pooled = fused.mean(dim=1)                     # B, D
        if lang_embed is not None:
            pooled = pooled + lang_embed
        action_chunk = self.action_head(pooled)        # B, chunk*dim
        return action_chunk.view(-1, self.chunk_size, self.action_dim)

    def action_smooth_loss(self, action_chunk):
        """相邻步动作变化的平滑约束"""
        diff = action_chunk[:, 1:, :] - action_chunk[:, :-1, :]
        return diff.pow(2).mean()

一开始只加了MSE回归loss，动作预测出来老是抖。加上action_smooth_loss之后，相邻步动作变化被约束住，连续操作任务成功率涨了8个点。这个平滑loss的权重0.1是调了好几组消融定的，太大了动作僵硬，太小了没效果。

DeepSpeed ZeRO-2 训练配置

8×A100训练，视频模型参数量大，用ZeRO-2切优化器状态。完整的ds_config.json如下：

{
  "train_micro_batch_size_per_gpu": 2,
  "gradient_accumulation_steps": 4,
  "gradient_clipping": 1.0,
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "none",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "activation_checkpointing": {
    "partition_activations": true,
    "cpu_checkpointing": true
  },
  "wall_clock_breakdown": false
}

一开始用ZeRO-3直接OOM，因为视频模型的参数也要切分到各卡，通信量太大。换成ZeRO-2之后优化器状态切分、参数保留在本地，通信量降了一截。micro batch size一开始设4，跑到一半报CUDA out of memory. Tried to allocate 2.37 GiB，降到2加gradient checkpointing才稳住。

训练循环

Dual范式的训练循环：先冻住视频生成模型，只训IDM和FDM：

def train_step(batch, idm, fdm, optimizer, engine):
    frames = batch["frames"].cuda()        # B, T, 3, H, W
    actions = batch["actions"].cuda()     # B, T, 80
    lang_embed = batch["lang_embed"].cuda()

    frame_curr = frames[:, 0]    # B, 3, H, W
    frame_future = frames[:, -1]

    # IDM预测动作
    pred_actions = idm(frame_curr, frame_future, lang_embed)

    # IDM loss = MSE + 时序平滑
    loss_mse = nn.functional.mse_loss(pred_actions, actions[:, :16])
    loss_smooth = idm.action_smooth_loss(pred_actions)
    loss_idm = loss_mse + 0.1 * loss_smooth

    # FDM前向一致性：用真实动作预测下一帧，和真实帧对比
    pred_next_frame = fdm(frame_curr, actions[:, 0])
    loss_fdm = nn.functional.l1_loss(pred_next_frame, frames[:, 1])

    total_loss = loss_idm + 0.5 * loss_fdm

    engine.backward(total_loss)
    engine.step()
    return {"loss_idm": loss_idm.item(), "loss_fdm": loss_fdm.item()}

FDM的权重0.5是消融出来的，太大了IDM被带偏，太小了想象推演的一致性约束不够。Tri范式的话三个模块分别建optimizer，这里就不展开了。

WebSocket推理服务

部署到真机上用WebSocket做策略服务器，端到端延迟压到150ms：

import asyncio
import websockets
import torch
import json

class PolicyServer:
    def __init__(self, idm_model, video_model, port=8765):
        self.idm = idm_model.cuda().half().eval()
        self.video = video_model.cuda().half().eval()
        self.port = port

    async def handle(self, websocket):
        async for message in websocket:
            data = json.loads(message)
            frame = self._decode_frame(data["frame"])   # 30ms
            with torch.no_grad():
                future_frame = self.video.predict(frame, steps=2)  # 60ms
            with torch.no_grad():
                action = self.idm(frame, future_frame, data["lang_embed"])  # 40ms
            await websocket.send(json.dumps({"action": action.cpu().tolist()}))

    def run(self):
        start_server = websockets.serve(self.handle, "0.0.0.0", self.port)
        asyncio.get_event_loop().run_until_complete(start_server)
        asyncio.get_event_loop().run_forever()

延迟分解：图像预处理30ms，视频模型推理60ms（从4步砍到2步+半精度才压下来），IDM推理40ms，网络和序列化30ms，合计约150ms。视频模型一开始要120ms，是延迟大头，砍未来预测步数和开fp16之后改善明显。

整个项目跑下来最深的体会是，世界动作模型这个方向，论文里讲的是方法，工程上全是脏活------时间戳对齐、动作空间映射、显存切分、延迟拆解，每一个都是踩坑踩出来的。LIBERO 99.3%不是调一个参数调出来的，是上面这些细节一个一个打磨的结果。

后续把完整配置、数据处理脚本和评测记录都整理成了文档资料，做这个方向的同学可以交流。