UniZero: Generalized and Efficient Planning with Scalable Latent World Models

2024-06-15

arXiv.org

AI Summary

To address poor scalability of world models and inefficient long-horizon planning in heterogeneous multi-task environments, this paper proposes a modular Transformer-based unified latent world model. The model jointly models dynamics prediction and decision-oriented representation within a shared latent space, integrating latent-space Monte Carlo Tree Search (MCTS) with the value-equivalence principle to unify long-term memory modeling and scalable multi-task planning. Its core innovation is a novel co-prediction architecture that overcomes the limitations of MuZero-style methods in handling task diversity and heterogeneous dependencies. Experiments demonstrate that the method significantly outperforms existing approaches on benchmarks requiring long-term memory; achieves superior scalability in Atari multi-task learning; and matches or exceeds state-of-the-art performance on single-task Atari and DeepMind Control Suite benchmarks.

Technology Category

• Multiagent Systems: Multiagent Planning
• Humans and AI: Human-Aware Planning and Behavior Prediction
• Machine Learning: Mixture of Experts (MoE)

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Learning predictive world models is crucial for enhancing the planning capabilities of reinforcement learning (RL) agents. Recently, MuZero-style algorithms, leveraging the value equivalence principle and Monte Carlo Tree Search (MCTS), have achieved superhuman performance in various domains. However, these methods struggle to scale in heterogeneous scenarios with diverse dependencies and task variability. To overcome these limitations, we introduce UniZero, a novel approach that employs a modular transformer-based world model to effectively learn a shared latent space. 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 the latent space. We show that UniZero significantly outperforms existing baselines in benchmarks that require long-term memory. Additionally, UniZero demonstrates superior scalability in multitask learning experiments conducted on Atari benchmarks. In standard single-task RL settings, such as Atari and DMControl, UniZero matches or even surpasses the performance of current state-of-the-art methods. Finally, extensive ablation studies and visual analyses validate the effectiveness and scalability of UniZero's design choices. Our code is available at extcolor{magenta}{https://github.com/opendilab/LightZero}.

Problem

• AI planning
• long-term memory
• multi-task learning

Innovation

• Shared Scalable Hidden Map Model
• Long-term Memory Optimization
• Multi-task Learning