arXiv:2502.01591cs.LGcs.AI2025-02ICML被引 16

改进Transformer世界模型,让强化学习在少量数据下表现更优

Improving Transformer World Models for Data-Efficient RL

  • 用动态预热策略,先用真实数据训练,再逐步引入想象数据
  • 采用静态词元化方法,提升图像分块编码的稳定性
  • 通过块教师强制机制,实现未来状态的联合推理,提升生成效率

我们针对基于Transformer的世界模型(TWM)提出了三项改进:(a) “Dyna with warmup”——先用真实数据训练策略,待世界模型充分训练后再引入虚拟数据;(b) “最近邻词元化器”,用于图像分块,确保词元在创建后保持静态,为TWM提供稳定目标;(c) “块教师强制”,使世界模型能联合推理下一时刻的多个未来令牌,而非逐个生成。实验表明,该方法在多个环境中显著优于现有方法。尤其在挑战性极高的Craftax-classic基准上,仅需100万环境步数即达69.66%奖励,远超DreamerV3的53.2%,首次超越人类65.0%的表现。初步结果还涵盖Craftax-full、MinAtar及三款双人游戏,验证了方法的通用性。

原文摘要 · Abstract (English)

We present three improvements to the standard model-based RL paradigm based on transformers: (a) "Dyna with warmup", which trains the policy on real and imaginary data, but only starts using imaginary data after the world model has been sufficiently trained; (b) "nearest neighbor tokenizer" for image patches, which improves upon previous tokenization schemes, which are needed when using a transformer world model (TWM), by ensuring the code words are static after creation, thus providing a constant target for TWM learning; and (c) "block teacher forcing", which allows the TWM to reason jointly about the future tokens of the next timestep, instead of generating them sequentially. We then show that our method significantly improves upon prior methods in various environments. We mostly focus on the challenging Craftax-classic benchmark, where our method achieves a reward of 69.66% after only 1M environment steps, significantly outperforming DreamerV3, which achieves 53.2%, and exceeding human performance of 65.0% for the first time. We also show preliminary results on Craftax-full, MinAtar, and three different two-player games, to illustrate the generality of the approach.

强化学习Transformer世界模型数据高效

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