摘要
离线强化学习(RL)的一个重要目标是从大型异构数据集中训练出具有高能力的通用智能体。然而,以往扩展离线RL的方法要么严重依赖专家轨迹,要么难以泛化到不同的未见任务。受条件视频生成中世界模型出色泛化能力的启发,我们探索了基于图像观测的世界模型在扩展离线RL和增强新任务泛化能力方面的潜力。本文介绍了JOWA:联合优化的世界-动作模型,这是一个基于离线模型的RL智能体,在包含60亿token数据的多个Atari游戏中进行预训练,以学习通用表示和决策能力。我们的方法通过共享的Transformer主干网络联合优化世界-动作模型,这在预训练过程中稳定了大型模型的时序差分学习。此外,我们提出了一种可证明高效且可并行的规划算法来补偿Q值估计误差,从而搜索出更好的策略。实验结果表明,我们最大的智能体(拥有1.5亿个参数)仅使用10%的子采样离线数据,在预训练游戏中就达到了78.9%的人类水平性能,平均比现有的最先进的大规模离线RL基线高出31.6%。此外,JOWA能够很好地扩展模型容量,并且可以使用每个游戏仅5k的离线微调数据(大约4条轨迹)高效地迁移到新的游戏中,展现了优越的泛化能力。我们将发布代码和模型权重到https://github.com/CJReinforce/JOWA。