UniZero: Generalized and Efficient Planning with Scalable Latent World Models

学习预测世界模型对于增强强化学习代理的规划能力至关重要。值得注意的是，基于价值等价原则和蒙特卡罗树搜索（MCTS）的MuZero风格算法在各个领域都取得了超人类的表现。然而，在需要捕捉长期依赖关系的环境中，MuZero的性能会迅速恶化。我们确定这部分是由于潜在表示与历史信息的“缠结”部分导致的，这导致与辅助自监督状态正则化不兼容。为了克服这个限制，我们提出了一种新的方法UniZero，它使用基于Transformer的潜在世界模型将潜在状态与隐含的潜在历史分离开来。通过同时预测与学习的潜在历史相关的潜在动态和决策导向量，UniZero使得长期规划和策略的联合优化成为可能，从而在潜在空间中实现更广泛和更高效的规划。我们证明了即使使用单帧输入，UniZero在Atari 100k基准测试中也能与MuZero风格的算法相匹配或超越其性能。此外，它在需要长期记忆的基准测试中显著优于以前的基线。最后，我们通过广泛的消融研究、视觉分析和多任务学习结果验证了我们设计选择的有效性和可扩展性。代码可在\textcolor{magenta}{https://github.com/opendilab/LightZero}上获得。

解决问题：论文旨在通过提出一种新的方法UniZero来解决MuZero在捕捉长期依赖方面的性能下降问题。

关键思路：UniZero通过使用基于Transformer的潜在世界模型将潜在状态与历史信息解耦，从而实现了长期依赖的规划。

其他亮点：论文在Atari 100k基准测试中证明了UniZero即使使用单帧输入也能匹配或超越MuZero风格算法的性能，并且在需要长期记忆的基准测试中明显优于之前的基线。此外，论文还进行了大量消融研究、可视化分析和多任务学习结果验证其设计选择的有效性和可扩展性。代码已经开源。

相关研究：最近在这个领域中，还有一些相关的研究，如MuZero、World Models等。