利用并行观测预测改进基于 Token 的世界模型
机器学习
2024-05-30 v5 人工智能
摘要
受 Transformer 在离散符号序列上取得成功的启发，基于 Token 的世界模型（TBWMs）最近被提出作为一种样本高效的方法。在 TBWMs 中，世界模型将智能体的经验视为类似语言的 Token 序列进行消费，其中每个观测构成一个子序列。然而，在想象过程中，逐个 Token 顺序生成下一个观测会导致严重的瓶颈，从而导致训练时间长、GPU 利用率低以及表征能力受限。为了解决这一瓶颈，我们设计了一种新颖的并行观测预测（Parallel Observation Prediction, POP）机制。POP 通过一种专为强化学习场景定制的新型前向模式，增强了保留网络（Retentive Network, RetNet）。我们将 POP 整合到一个名为 REM（Retentive Environment Model）的新型 TBWM 智能体中，展示了其相比先前 TBWMs 快 15.4 倍的想象速度。REM 在 Atari 100K 基准测试的 26 款游戏中有 12 款达到了超人类性能，同时训练时间少于 12 小时。我们的代码可在\url{https://github.com/leor-c/REM}获取。
引用
@article{arxiv.2402.05643,
title = {Improving Token-Based World Models with Parallel Observation Prediction},
author = {Lior Cohen and Kaixin Wang and Bingyi Kang and Shie Mannor},
journal= {arXiv preprint arXiv:2402.05643},
year = {2024}
}