CoLA-World：将 latent action 模型与 world model 协同联合训练CoLA-World: co-evolving joint training of the latent action model and the world model
Yucen Wang, Fengming Zhang, De-Chuan Zhan（南京大学）· Li Zhao, Kaixin Wang, Jiang Bian（Microsoft Research Asia）Yucen Wang, Fengming Zhang, De-Chuan Zhan (Nanjing University) · Li Zhao, Kaixin Wang, Jiang Bian (Microsoft Research Asia)
CoLA-World 首次打破"先训 LAM、再训 world model"的两阶段范式，将二者联合端到端训练——通过一个关键的 warm-up 阶段防止 codebook collapse，最终实现两者相互促进的协同进化，在视频预测质量和视觉规划成功率上全面超越两阶段基线。CoLA-World is the first to break the two-stage paradigm of "train the LAM first, then train the world model" by training the two jointly and end-to-end: a critical warm-up stage prevents codebook collapse, ultimately yielding a mutually reinforcing co-evolution that surpasses the two-stage baseline in both video prediction quality and visual planning success rate.
机器人和视频生成领域普遍采用"latent action model（LAM）+ world model（WM）"的两阶段流水线：先用 inverse dynamics model（IDM）从无标注视频中学习 latent action codebook，再将 codebook 固定，训练以 latent action 为条件的 world model。这种分离训练方式存在根本性冗余——IDM 内部的 forward dynamics model（FDM）与 WM 本质上做的是同一件事（预测下一帧），却各自为政、无法相互促进。Robotics and video generation widely adopt a two-stage pipeline of "latent action model (LAM) + world model (WM)": an inverse dynamics model (IDM) first learns a latent action codebook from unlabeled video, then the codebook is frozen and a world model conditioned on latent actions is trained. This separated training carries a fundamental redundancy — the forward dynamics model (FDM) inside the IDM and the WM essentially do the same thing (predicting the next frame), yet they work in isolation and cannot reinforce each other.
"We argue that there is an inherent redundancy in the design: the FDM and the world model perform almost identical functions, both modeling the transition dynamics of the environment."
Figure 1. (a) 先前方法：两阶段流水线，先固定训练 LAM，再用其 latent actions 训练 world model。(b) CoLA-World：直接用 world model 替代 IDM 中的 FDM，实现一阶段联合训练——world model 充当"知识丰富的导师"，持续为 LAM 提供梯度，而不断演进的 LAM 又反过来提升 world model 的控制精度。Figure 1. (a) Prior methods: a two-stage pipeline that first trains the LAM and freezes it, then trains a world model with its latent actions. (b) CoLA-World: the world model directly replaces the FDM inside the IDM, enabling one-stage joint training — the world model acts as a "knowledgeable tutor" that keeps supplying gradients to the LAM, while the continually evolving LAM in turn improves the control precision of the world model.
2.73×视觉规划平均成功率提升（joint vs two-stage，VP2 benchmark）Average visual planning success rate gain (joint vs two-stage, VP2 benchmark)
158.36FVD（joint）vs 167.06（two-stage），LIBERO OOD 视频预测FVD (joint) vs 167.06 (two-stage), LIBERO OOD video prediction
38K步联合训练（计算量更少）即可达到或超越 60K 步两阶段基线Joint training (with less compute) already matches or surpasses the 60K-step two-stage baseline
02方法Method
CoLA-World 将 IDM 中原有的 FDM 替换为预训练 video generation world model，然后对整个系统端到端联合训练。核心挑战是直接联合训练会导致 latent action codebook 的"表征崩塌"（representational collapse），通过引入 warm-up 阶段可以优雅地解决这一问题。CoLA-World replaces the original FDM inside the IDM with a pretrained video generation world model, then trains the whole system jointly end-to-end. The core challenge is that naive joint training causes representational collapse of the latent action codebook, which is elegantly resolved by introducing a warm-up stage.
Figure 2. 直接从头联合训练时的 codebook 指标：utilization 降至 0，单个 code 使用率接近 100%，entropy 归零——即所有 latent action 退化为同一个 code，丧失表达能力。使用预训练权重初始化（"pre"曲线）虽可短暂延缓，但最终仍不可避免地崩塌。Figure 2. Codebook metrics when training jointly from scratch: utilization drops to 0, the usage rate of a single code approaches 100%, and entropy falls to zero — that is, all latent actions degenerate into the same code and lose expressive power. Initializing from pretrained weights (the "pre" curves) delays this for a while, but collapse remains inevitable.
解决方案：Warm-Up + 端到端联合训练Solution: Warm-Up + End-to-End Joint Training
CoLA-World 的训练分两阶段：Training in CoLA-World proceeds in two stages:
Warm-Up 阶段：固定 world model 参数，仅让其提供梯度来更新 LAM（IDM）。此时 world model 充当"冻结的导师"，帮助 LAM 建立稳健的 codebook 表示，防止崩塌。Warm-up stage: the world model parameters are frozen and it only supplies gradients to update the LAM (IDM). Here the world model acts as a "frozen tutor", helping the LAM build a robust codebook representation and preventing collapse.
End-to-End 联合训练阶段：解冻 world model，两者共同更新。LAM 持续产生更高质量的 latent actions，world model 以此为条件提升视频预测精度，形成正向反馈循环——即"co-evolution"（协同进化）。End-to-end joint training stage: the world model is unfrozen and both are updated together. The LAM keeps producing higher-quality latent actions, on which the world model conditions to improve video prediction accuracy, forming a positive feedback loop — namely "co-evolution".
Figure 3. Warm-up 阶段不同步数对应的 codebook 指标变化。warm-up 步数越充分，后续联合训练时 codebook 越稳定，utilization 保持在健康水平（>20%），entropy 不降为零。Figure 3. Codebook metrics under different numbers of warm-up steps. The more sufficient the warm-up, the more stable the codebook during subsequent joint training, with utilization kept at a healthy level (>20%) and entropy never falling to zero.
Co-Evolution 证据Evidence of Co-Evolution
Figure 4. (a) 在联合训练中，LAM 的 linear probing loss（衡量 latent action 质量）下降更快——对比 world model 固定时的情形，说明 world model 的持续进化正在帮助 LAM 学到更有意义的表示。(b) World model 在联合训练下视频预测性能持续提升，而固定 LAM 时出现明显瓶颈。Figure 4. (a) Under joint training the linear probing loss of the LAM (a proxy for latent action quality) drops faster than when the world model is frozen, showing that the continually evolving world model is helping the LAM learn more meaningful representations. (b) The world model keeps improving in video prediction under joint training, whereas a frozen LAM leads to a clear bottleneck.
Real Action 适配：跨域迁移Real Action Adaptation: Cross-Domain Transfer
学习到的 latent action codebook 可通过轻量级 adapter 与真实机器人动作对齐，实现从无标注视频到有标注机器人数据的迁移。实验表明，CoLA-World 的 codebook 在适配过程中保持健康的 utilization（>20%），而两阶段基线的 codebook 在适配阶段发生崩塌（utilization 跌至 ~10%，单个 code 占用率激增至 0.5）。The learned latent action codebook can be aligned with real robot actions through a lightweight adapter, transferring from unlabeled video to labeled robot data. Experiments show that the codebook of CoLA-World keeps a healthy utilization (>20%) throughout adaptation, whereas the codebook of the two-stage baseline collapses during adaptation (utilization drops to ~10% and the occupancy of a single code surges to 0.5).
03实验Experiments
评估涵盖四个机器人/视频数据集（OXE、EgoCentric、AgiBot、LIBERO），以及 VP2 benchmark 上的视觉规划任务（RoboDesk）。主要指标：视频预测用 FVD（越低越好），视觉规划用成功率（越高越好）。对比基线为相同计算预算（60K steps）的两阶段方法。The evaluation covers four robot/video datasets (OXE, EgoCentric, AgiBot, LIBERO) and the visual planning tasks of the VP2 benchmark (RoboDesk). Main metrics: FVD for video prediction (lower is better) and success rate for visual planning (higher is better). The baseline is a two-stage method under the same compute budget (60K steps).
Table 1：视频预测性能（FVD，越低越好）Table 1: Video prediction performance (FVD, lower is better)
数据集Dataset
两阶段（LAM30K+WM30K）Two-stage (LAM30K+WM30K)
CoLA-World（Warm8K+E2E52K）
OXE
281.05
278.90
EgoCentric
259.33
252.45
AgiBot
180.45
174.93
LIBERO（OOD）
167.06
158.36
在所有数据集上，相同计算预算下联合训练均优于两阶段方法，OOD 场景（LIBERO）提升尤为显著。On every dataset, joint training beats the two-stage method under the same compute budget, and the gain is especially pronounced in the OOD setting (LIBERO).
Table 2：Real Action 适配后的视频预测性能（FVD，越低越好）Table 2: Video prediction performance after real action adaptation (FVD, lower is better)
视觉规划任务上，CoLA-World 的平均成功率为两阶段方法的 2.73×，验证联合训练带来的 codebook 质量优势可直接转化为下游控制性能提升。On visual planning tasks the average success rate of CoLA-World is 2.73× that of the two-stage method, confirming that the codebook quality advantage brought by joint training translates directly into downstream control performance.
数据效率与 Codebook 健康性Data Efficiency and Codebook Health
Figure 5. Real action 适配过程中的 codebook 指标对比。两阶段方法（橙色）的 codebook utilization 跌至 ~10%，最大 code 使用率激增，entropy 崩塌；CoLA-World（蓝色）在整个适配过程中保持 utilization >20% 的健康状态，说明联合训练带来的表示质量优势具备对分布偏移的鲁棒性。Figure 5. Comparison of codebook metrics during real action adaptation. For the two-stage method (orange) codebook utilization drops to ~10%, the maximum code usage surges and entropy collapses; CoLA-World (blue) stays healthy with utilization >20% throughout adaptation, showing that the representation quality gained from joint training is robust to distribution shift.
Ablation：数据效率Ablation: Data Efficiency
计算量更少的配置（Warm8K + E2E30K，共 38K 步）已接近甚至超越 60K 步两阶段基线，表明联合训练收敛更快。此外，更大的 world model backbone（更多 DiT blocks）可持续提升 LAM 的 linear probing loss，验证 world model 能力与 LAM 质量之间的正向关联；增大 batch size（64→128）同样改善两者性能。A cheaper configuration (Warm8K + E2E30K, 38K steps in total) already approaches or even surpasses the 60K-step two-stage baseline, showing that joint training converges faster. Moreover, a larger world model backbone (more DiT blocks) keeps improving the linear probing loss of the LAM, confirming the positive link between world model capability and LAM quality; enlarging the batch size (64→128) likewise improves both.
Figure 7. Action transfer 定性结果：source 和 target 视频来自不同数据集。CoLA-World 学到的 latent actions 能够跨数据集迁移并合理重现目标场景中的动作语义。Figure 7. Qualitative action transfer results: the source and target videos come from different datasets. The latent actions learned by CoLA-World transfer across datasets and plausibly reproduce the action semantics of the target scene.
04局限性Limitations
说明：论文未设置单独的 Limitations 章节。以下条目中，标注"（论文明确指出）"者来自原文，其余为根据方法设计推断（inferred）。Note: the paper has no dedicated Limitations section. Among the items below, those marked "(stated in the paper)" come from the original text; the rest are inferred from the method design.
依赖大规模预训练 video generation model（论文明确指出）Depends on a large-scale pretrained video generation model (stated in the paper)
CoLA-World 需要以预训练的 video generation world model 为起点，方能在 warm-up 阶段为 LAM 提供有意义的梯度。若预训练模型质量不足或领域差异过大，效果将受限。CoLA-World has to start from a pretrained video generation world model in order to supply meaningful gradients to the LAM during the warm-up stage. If the pretrained model is of insufficient quality or the domain gap is too large, the benefit will be limited.
计算成本显著高于两阶段方法（论文明确指出）Compute cost is markedly higher than the two-stage method (stated in the paper)
联合训练（Warm8K + E2E52K）约需 ~100 小时，而两阶段（LAM30K + WM30K）约需 ~75 小时（Table 8）。尽管收敛更快、性能更优，但绝对计算成本仍然较高，对资源受限的场景形成门槛。Joint training (Warm8K + E2E52K) takes about ~100 hours, whereas the two-stage pipeline (LAM30K + WM30K) takes about ~75 hours (Table 8). Despite faster convergence and better performance, the absolute compute cost remains high and forms a barrier for resource-constrained settings.
视觉规划绝对成功率仍然偏低（inferred）Absolute visual planning success rate is still low (inferred)
VP2 benchmark 上，CoLA-World 的 5-task 平均成功率为 21.20%，虽远优于两阶段的 7.73%，但绝对数值仍然偏低，距离实际机器人部署所需的可靠性尚有差距。On the VP2 benchmark the 5-task average success rate of CoLA-World is 21.20%, far better than the 7.73% of the two-stage method, yet the absolute value is still low and falls short of the reliability required for real robot deployment.
Warm-up 步数的超参敏感性（inferred）Hyperparameter sensitivity to the number of warm-up steps (inferred)
实验显示 warm-up 步数对 codebook 稳定性有显著影响，需要一定调参成本；在新的数据集或 backbone 上应用时，最优 warm-up 时长可能需要重新搜索。Experiments show that the number of warm-up steps has a significant effect on codebook stability and demands a certain amount of tuning; when applying the method to a new dataset or backbone, the optimal warm-up length may have to be searched again.
