arXiv:2410.00564cs.LGcs.AI2024-10ICLR被引 16

用60亿token数据预训练世界-动作模型,实现高效离线强化学习与跨任务泛化。

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

  • 通过共享Transformer架构联合优化世界-动作模型,稳定大模型训练过程。
  • 150万参数模型仅用10%数据即达78.9%人类水平性能,优于现有方法31.6%。
  • 每新游戏仅需约4条轨迹即可高效迁移,适合资源受限的通用智能体开发。

离线强化学习的重要目标是利用大规模异构数据构建具备强泛化能力的通用智能体。然而,已有方法要么高度依赖专家轨迹,要么难以在未见任务上泛化。受条件视频生成中世界模型优异泛化能力的启发,本文探索基于图像观测的世界模型在离线强化学习中的潜力。提出JOWA:联合优化的世界-动作模型,基于60亿标记数据在多个Atari游戏中进行预训练,以学习通用表征与决策能力。该方法通过共享Transformer骨干网络联合优化世界-动作模型,在预训练阶段稳定大型模型的时序差分学习。此外,设计了一种可证明高效且可并行化的规划算法,以补偿Q值估计误差,从而搜索更优策略。实验表明,最大规模模型(150万参数)仅使用10%子采样离线数据,在预训练游戏中达到78.9%人类水平性能,平均优于现有最先进大规模离线强化学习基线31.6%。此外,JOWA在模型容量增大时仍表现良好,可在每个新游戏中仅用5000条离线微调数据(约4条轨迹)实现高效迁移,展现出卓越泛化能力。代码与模型权重将开源于https://github.com/CJReinforce/JOWA。

原文摘要 · Abstract (English)

A significant aspiration of offline reinforcement learning (RL) is to develop a generalist agent with high capabilities from large and heterogeneous datasets. However, prior approaches that scale offline RL either rely heavily on expert trajectories or struggle to generalize to diverse unseen tasks. Inspired by the excellent generalization of world model in conditional video generation, we explore the potential of image observation-based world model for scaling offline RL and enhancing generalization on novel tasks. In this paper, we introduce JOWA: Jointly-Optimized World-Action model, an offline model-based RL agent pretrained on multiple Atari games with 6 billion tokens data to learn general-purpose representation and decision-making ability. Our method jointly optimizes a world-action model through a shared transformer backbone, which stabilize temporal difference learning with large models during pretraining. Moreover, we propose a provably efficient and parallelizable planning algorithm to compensate for the Q-value estimation error and thus search out better policies. Experimental results indicate that our largest agent, with 150 million parameters, achieves 78.9% human-level performance on pretrained games using only 10% subsampled offline data, outperforming existing state-of-the-art large-scale offline RL baselines by 31.6% on averange. Furthermore, JOWA scales favorably with model capacity and can sample-efficiently transfer to novel games using only 5k offline fine-tuning data (approximately 4 trajectories) per game, demonstrating superior generalization. We will release codes and model weights at https://github.com/CJReinforce/JOWA

离线RL世界模型迁移学习Atari

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。