Awesome World Models / OpenReview

Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining

作者：作者信息待提取 · 单位：State Key Laboratory of Multimodal Artificial Intelligence Systems, Institute of Automation, Chinese Academy of Sciences, School of Artificial Intelligence, University of Chinese Academy of Sciences, Alibaba Group · 来源：General Approaches to World Models · 发布日期：日期待提取 · 分类：8. World Models in the context of Model-Based RL

视觉语言动作/策略学习后训练/规模化基于模型的强化学习

精读笔记

Problem Setting

这篇论文解决的不是传统单任务 offline model-based RL，而是一个更具体也更困难的问题：用静态、多游戏、像素级 Atari replay 数据，训练一个单一 agent，同时获得跨任务共享表征、共享决策能力和少样本迁移能力。

关键矛盾在于：scaling 在监督学习里通常靠更大的模型和更多数据带来更好的表征，但 TD-based RL 的 bootstrapping、off-policy 数据和分布外动作会让大模型更容易不稳定。尤其在 offline RL 中，模型容量越大，不一定越好，可能只是更强地拟合数据分布偏差和价值误差。

以前方法大致卡在两端：Decision Transformer / conditional BC 类方法容易 scale，但本质是轨迹建模和行为复制，依赖专家或高回报轨迹，决策能力来自数据中已有策略；Scaled-QL 类方法保留 TD 学习，但多任务共享更多是 representation sharing，通常还需要 task-specific Q head，迁移到新任务时决策层并没有真正被预训练。

所以 JOWA 实际要证明的是：TD-based offline RL 也可以通过 world model regularization 和 shared sequence backbone 进入 scaling regime，而不是只能靠 BC-style sequence modeling。

Motivation

作者的核心观察是：world model 在视频生成领域已经显示出跨场景时间序列建模能力，而 offline RL 的多任务预训练恰好缺一个能稳定吸收大规模像素序列的结构性监督信号。单纯 TD loss 太稀疏、太 bootstrapped、太容易受离线分布偏差影响；单纯行为建模又无法真正优化 return。

因此缺口不是“有没有 world model”，而是 world model 是否能成为 large-scale TD learning 的稳定器，并且是否能把 dynamics knowledge 直接服务于 decision-time improvement。JOWA 的动机可以理解为：把 video-model-style temporal prediction 当作 RL scaling 的基础设施，把 TD learning 当作让该表征可行动化的控制信号。

这也是它和很多 model-based offline RL 的差别：不是为了生成大量 rollout 再训练 policy，而是用 world modeling 帮 critic 学得更稳，再在推理时用很短的 model rollout 修正 critic。

Core Idea

JOWA 的核心思想是 joint world-action pretraining：同一个 Transformer 同时承担环境序列预测和 Q-value estimation。world model 不是独立预训练后冻结的表征模块，也不是只用于 synthetic data augmentation；critic 也不是在一个纯 RL feature extractor 上单独训练。两类损失共同回传到同一个 temporal representation，使得表示空间同时服务于“下一帧/奖励/终止可预测”和“动作价值可区分”。

直觉上这会有效，是因为多任务离线 TD 学习最缺的是稳定、密集、跨任务共享的训练信号。world modeling loss 提供了几乎每个 token 都有监督的预测目标，迫使 Transformer 学到通用的视觉动态和历史状态压缩；TD loss 则把其中与控制相关的维度拉出来。相比纯 Decision Transformer，它不把策略质量完全绑定到数据中的行为；相比纯 Scaled-QL，它不是只让表征共享、Q 头分裂，而是让 decision head 也建立在统一的历史动态模型上。

引入的新 inductive bias 是：一个好的 offline generalist RL backbone 应该首先是一个 conditional dynamics model，其次才是一个 value estimator。这个 bias 对像素 Atari 很自然，因为不同游戏虽然 reward/action semantics 不同，但图像时序规律、短期因果结构和局部状态转移可以共享。

Method

1. 共享 backbone 的 world-action model：解决 TD 大模型不稳定和多任务表征不统一的问题。world loss 给 dense supervised gradients，TD/CQL loss 给 control gradients。核心变化是让 world representation 和 value representation 不再解耦。

2. Distributional TD + CQL：解决 offline critic 的 extrapolation 和 MSE value regression 的训练病态。文中 ablation 显示 MSE TD 在该设置下容易过度优化 CQL 项，导致价值分布畸形。这里 distributional value learning 不只是性能 trick，更像是让 value learning 与 classification-style world losses 在优化形态上更兼容。

3. 短 horizon Q-guided beam planning：解决单步 greedy Q 估计误差。planner 不在全动作空间里搜索，而是用 Q top-K 约束 action candidates，再用 learned reward/dynamics rollout 计算短期 return + terminal Q。核心变化是把 world model 用作 inference-time reranker，给 critic 增加 test-time compute。

4. 两阶段训练：先让 tokenizer/world-part 具备基本序列预测能力，再 joint optimize。机制上这是为了避免 RL loss 一开始污染尚未成形的视觉 token dynamics。它重要但更偏 training stabilization，不是概念核心。

Key Insight / Why It Works

我认为最关键的贡献不是“用了 world model”，而是把 world modeling 变成 large-scale offline TD learning 的正则器。TD learning 在大模型上失败，往往不是容量不够，而是 bootstrapped target 噪声、OOD action 误差和稀疏 reward signal 让 representation drift。JOWA 用下一 token / reward / done prediction 把 representation 锚定在数据支持的真实时序结构上，使 critic 不至于只追逐自举目标。

第二个有效点是 representation alignment。纯 world model 可能学到与控制无关的像素细节，纯 Q-learning 可能学到脆弱的 reward shortcut。joint training 让共享 backbone 同时受 dynamics 和 value 约束，这比“先预训练 world model，再训练 policy”更可能学到 action-relevant dynamics。

第三个有效点是 test-time compute。beam planning 本质上不是强规划，而是局部搜索 + Q reranking。它利用 Q 函数缩小搜索空间，再用 world model 检查短期后果。这个机制尤其适合 Atari：短期奖励和状态变化通常足够信息密集，H=2 的 rollout 已经能修正一些明显的 Q ranking 错误。但这也说明它不是长期推理能力的证据。

哪些可能只是辅助：VQ-VAE tokenization、GPT-2 backbone、task embedding、两阶段训练、FlashAttention 等主要是让系统可训练；它们重要但不是核心思想。CQL 和 distributional TD 是必要 stabilizer，属于已有组件的正确组合。

哪些可能主要来自 scaling / data：6B tokens、150M 参数、多游戏预训练和 Atari 任务族共享结构很可能解释了相当一部分泛化收益。文中虽有 scaling curve，但没有完全隔离“更大模型”与“world-action joint objective”的交互。所谓 OOD 泛化也主要是在 Atari 内部 held-out games，不应过度解读为跨域泛化。

需要警惕的是，planning 增益可能部分来自 evaluation-time hyperparameter tuning 和 action-space prior，而不是 learned model 真正具备长程模拟能力。文中对 world model error 与 control gain 的关系没有充分展开，增益来源仍有不清楚之处。

Relation To Prior Work

JOWA 最接近三条线的交汇：multi-game Decision Transformer、Scaled-QL 式多任务 offline Q-learning、以及 IRIS/FICC/TD-MPC 类 world-model-based RL。

相对 MGDT / Elastic DT，本质差异是它不把 RL 简化为条件行为建模。DT 系方法的泛化通常依赖 return prompt 和数据中存在的高质量行为，容易变成 expert trajectory retrieval。JOWA 保留 TD backup 和 conservative value learning，因此理论上可以从非专家混合数据中做 policy improvement。

相对 Scaled-QL，差异在于 JOWA 不只是共享视觉 encoder，而是共享一个带历史的 dynamics-aware Transformer，并且不依赖每个任务独立 Q head。这个点对迁移很关键：如果 Q head 是 task-specific 的，那么预训练更多是 representation pretraining；JOWA 试图让 action-value mapping 本身也部分共享。

相对 FICC 或传统 model-based offline RL，JOWA 不强调用 synthetic rollout 扩充训练数据。恰恰相反，文中发现 synthetic data 在多任务预训练中负增益。这说明它的 model-based 成分主要体现在 representation regularization 和 inference-time search，而不是 model rollout policy optimization。

看似新的部分很多是已有思想重组：VQ token world model、CQL、distributional RL、beam search 都不是新组件。实质创新在于把这些组件组织成一个 scalable offline TD pretraining recipe，并证明在 multi-game Atari low-data setting 下 joint objective 比单独 model-free 或 BC-style scaling 更有效。

Dataset / Evaluation

评估覆盖了 15 个 Atari 训练游戏和 5 个 held-out 游戏，任务间视觉、动态和动作语义有差异，但仍属于同一个 Atari 生态。它验证的是“within-family multi-task generalization”，不是跨机器人、跨物理、跨 embodied setting 的泛化。

低数据设置使用 DQN-Replay 的 10% 子采样，确实比 full dataset 更能测试 sample efficiency。但数据仍来自 DQN replay，包含大量与 DQN 学习过程相关的行为分布；fine-tuning 的 5k transitions 还使用 expert-level replay 子集。这里的少样本迁移更像在相似 benchmark family 内快速适配，而不是在真实未知环境中从弱数据中学习。

核心 claim 中“scaling trend”主要由 IQM aggregate 支撑。作者也承认单个游戏存在反 scaling。由于 Atari 分数和 reward structure 高度非线性，aggregate scaling 可能掩盖 game-level instability。文中对这种不一致归因为 offline RL 固有问题，这个判断合理但没有解决。

无 sticky action 的 evaluation 沿用部分 prior，但会降低环境随机性，对 planning 和 memorization 都更友好。benchmark 支持 JOWA 在该协议下优于 baselines，但不足以证明 robust deployment 或强 OOD generalization。

Limitation

第一，方法强依赖数据覆盖。CQL 只能在离线数据支持附近保守估值，world model 也只能在数据分布附近短程预测。beam search 如果走出数据支持，model error 会迅速放大。JOWA 的 planner 通过 Q top-K 限制动作，实际上是在避免真正开放搜索。

第二，planning 的上限很明显。H=2 的 beam search 更像 local correction，不是长期 credit assignment 或 abstract planning。它可以修正短期 reward/dynamics 下的 Q ranking mistake，但不能说明模型学会了长程策略推理。

第三，泛化可能主要来自 Atari family 的 shared structure。预训练和 held-out 任务之间共享输入格式、动作接口、模拟器统计规律和奖励设计风格。迁移到动作空间不一致、观测模态变化、reward semantics 完全不同的任务时，joint Q head 是否仍有意义，文中未充分说明。

第四，增益归因仍不完全清晰。JOWA 同时改变了模型规模、world loss、distributional TD、CQL、planning、task embedding 和训练流程。虽然 ablation 做了一些拆解，但没有完全回答：如果给 Scaled-QL 同等 Transformer history backbone 和同等训练预算，会差多少；如果只做 world-regularized representation without planner，会保留多少迁移收益。

第五，synthetic data 负增益暴露了 model-based offline RL 的核心难题并没有被解决：learned model 可以帮助推理和正则化，但一旦用于训练数据生成，多步误差和保守惩罚会恶化。也就是说，JOWA 并没有解决 offline model rollout 的根本问题，而是绕开了它。

Takeaway

  • 1. 对 scalable offline RL 来说，world modeling 更有价值的角色可能不是生成 synthetic data，而是作为 TD learning 的 dense regularizer 和 representation anchor。
  • 2. 多任务 offline TD learning 要想 scale，关键不是简单加大 Q network，而是给 critic 一个稳定的、跨任务共享的时序建模目标；否则大模型会放大 bootstrapping 和 OOD error。
  • 3. Inference-time planning 在当前这类系统中更现实的形态可能是 Q-guided short-horizon reranking，而不是完整树搜索。
  • test-time compute 可以显著提升性能，但要避免把它误读为长期推理。

一句话总结

JOWA 是一篇把 world-model-style sequence prediction、conservative distributional Q-learning 和短程 test-time search 组合成 scalable multi-task offline TD pretraining recipe 的论文，真正贡献在于证明 world modeling 可以作为大规模离线价值学习的稳定器，而不只是 rollout generator。