Learning predictive world models is essential for enhancing the planning
capabilities of reinforcement learning agents. Notably, the MuZero-style
algorithms, based on the value equivalence principle and Monte Carlo Tree
Search (MCTS), have achieved superhuman performance in various domains.
However, in environments that require capturing long-term dependencies,
MuZero's performance deteriorates rapidly. We identify that this is partially
due to the \textit{entanglement} of latent representations with historical
information, which results in incompatibility with the auxiliary
self-supervised state regularization. To overcome this limitation, we present
\textit{UniZero}, a novel approach that \textit{disentangles} latent states
from implicit latent history using a transformer-based latent world model. By
concurrently predicting latent dynamics and decision-oriented quantities
conditioned on the learned latent history, UniZero enables joint optimization
of the long-horizon world model and policy, facilitating broader and more
efficient planning in latent space. We demonstrate that UniZero, even with
single-frame inputs, matches or surpasses the performance of MuZero-style
algorithms on the Atari 100k benchmark. Furthermore, it significantly
outperforms prior baselines in benchmarks that require long-term memory.
Lastly, we validate the effectiveness and scalability of our design choices
through extensive ablation studies, visual analyses, and multi-task learning
results. The code is available at
\textcolor{magenta}{this https URL}.

UniZero 是一种基于 Transformer 的新增方法，通过解耦潜在状态和历史信息，能够更好地捕捉长期依赖关系，提供了更加广泛和高效的潜在空间规划，相比 MuZero-style 算法在 Atari 100k 基准测试中表现更好，且在需要长期记忆的基准测试中显著优于之前的基线模型。