UniZero: Generalized and Efficient Planning with Scalable Latent World Models

• URL: http://arxiv.org/abs/2406.10667v2
• Date: Fri, 03 Jan 2025 08:57:39 GMT
• Title: UniZero: Generalized and Efficient Planning with Scalable Latent World Models
• Authors: Yuan Pu, Yazhe Niu, Zhenjie Yang, Jiyuan Ren, Hongsheng Li, Yu Liu,
• Abstract summary: UniZero is a novel approach that employs a modular transformer-based world model to effectively learn a shared latent space. We show that UniZero significantly outperforms existing baselines in benchmarks that require long-term memory. 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.
• Score: 29.648382211926364
• License:
• 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 \textcolor{magenta}{https://github.com/opendilab/LightZero}.

Related papers

• Vintix: Action Model via In-Context Reinforcement Learning [72.65703565352769]
  We present the first steps toward scaling ICRL by introducing a fixed, cross-domain model capable of learning behaviors through in-context reinforcement learning. Our results demonstrate that Algorithm Distillation, a framework designed to facilitate ICRL, offers a compelling and competitive alternative to expert distillation to construct versatile action models.
  arXiv  Detail & Related papers  (2025-01-31T18:57:08Z)
• Mamba-CL: Optimizing Selective State Space Model in Null Space for Continual Learning [54.19222454702032]
  Continual Learning aims to equip AI models with the ability to learn a sequence of tasks over time, without forgetting previously learned knowledge. State Space Models (SSMs) have achieved notable success in computer vision. We introduce Mamba-CL, a framework that continuously fine-tunes the core SSMs of the large-scale Mamba foundation model.
  arXiv  Detail & Related papers  (2024-11-23T06:36:16Z)
• Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient [9.519619751861333]
  We propose a state space model (SSM) based world model based on Mamba. It achieves $O(n)$ memory and computational complexity while effectively capturing long-term dependencies. This model is accessible and can be trained on an off-the-shelf laptop.
  arXiv  Detail & Related papers  (2024-10-11T15:10:40Z)
• Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining [49.730897226510095]
  We introduce JOWA: Jointly-Reinforced World-Action model, an offline model-based RL agent pretrained on Atari games with 6 billion tokens data. Our largest agent, with 150 million parameters, 78.9% human-level performance on pretrained games using only 10% subsampled offline data, outperforming existing state-of-the-art large-scale offline RL baselines by 31.6% on averange.
  arXiv  Detail & Related papers  (2024-10-01T10:25:03Z)
• LightZero: A Unified Benchmark for Monte Carlo Tree Search in General Sequential Decision Scenarios [32.83545787965431]
  Building agents based on tree-search planning capabilities with learned models has achieved remarkable success in classic decision-making problems, such as Go and Atari. It has been deemed challenging or even infeasible to extend Monte Carlo Tree Search (MCTS) based algorithms to diverse real-world applications. In this work, we introduce LightZero, the first unified benchmark for deploying MCTS/MuZero in general sequential decision scenarios.
  arXiv  Detail & Related papers  (2023-10-12T14:18:09Z)
• Maximize to Explore: One Objective Function Fusing Estimation, Planning, and Exploration [87.53543137162488]
  We propose an easy-to-implement online reinforcement learning (online RL) framework called textttMEX. textttMEX integrates estimation and planning components while balancing exploration exploitation automatically. It can outperform baselines by a stable margin in various MuJoCo environments with sparse rewards.
  arXiv  Detail & Related papers  (2023-05-29T17:25:26Z)
• Equivariant MuZero [14.027651496499882]
  We propose improving the data efficiency and generalisation capabilities of MuZero by explicitly incorporating the symmetries of the environment in its world-model architecture. We prove that, so long as the neural networks used by MuZero are equivariant to a particular symmetry group acting on the environment, the entirety of MuZero's action-selection algorithm will also be equivariant to that group.
  arXiv  Detail & Related papers  (2023-02-09T17:46:29Z)
• Efficient Offline Policy Optimization with a Learned Model [83.64779942889916]
  MuZero Unplugged presents a promising approach for offline policy learning from logged data. It conducts Monte-Carlo Tree Search (MCTS) with a learned model and leverages Reanalyze algorithm to learn purely from offline data. This paper investigates a few hypotheses where MuZero Unplugged may not work well under the offline settings.
  arXiv  Detail & Related papers  (2022-10-12T07:41:04Z)
• Exploration Policies for On-the-Fly Controller Synthesis: A Reinforcement Learning Approach [0.0]
  We propose a new method for obtaining unboundeds based on Reinforcement Learning (RL) Our agents learn from scratch in a highly observable partially RL task and outperform existing overall, in instances unseen during training.
  arXiv  Detail & Related papers  (2022-10-07T20:28:25Z)
• Transformers are Sample Efficient World Models [1.9444242128493845]
  We introduce IRIS, a data-efficient agent that learns in a world model composed of a discrete autoencoder and an autoregressive Transformer. With the equivalent of only two hours of gameplay in the Atari 100k benchmark, IRIS achieves a mean human normalized score of 1.046, and outperforms humans on 10 out of 26 games.
  arXiv  Detail & Related papers  (2022-09-01T17:03:07Z)
• Meta-Learned Attribute Self-Gating for Continual Generalized Zero-Shot Learning [82.07273754143547]
  We propose a meta-continual zero-shot learning (MCZSL) approach to generalizing a model to categories unseen during training. By pairing self-gating of attributes and scaled class normalization with meta-learning based training, we are able to outperform state-of-the-art results.
  arXiv  Detail & Related papers  (2021-02-23T18:36:14Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.