TransDreamer: Reinforcement Learning with Transformer World Models

2022/02/19 by Chang Chen, Chen, Chang, Yifu Wu +5 · 21 citations
Computer Science · #Anomaly Detection Techniques and Applications #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning and Data Classification #Reinforcement Learning in Robotics

paper · pdf · doi:10.48550/arxiv.2202.09481

openalex publication_date 2022/02/19 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

The Dreamer agent provides various benefits of Model-Based Reinforcement Learning (MBRL) such as sample efficiency, reusable knowledge, and safe planning. However, its world model and policy networks inherit the limitations of recurrent neural networks and thus an important question is how an MBRL framework can benefit from the recent advances of transformers and what the challenges are in doing so. In this paper, we propose a transformer-based MBRL agent, called TransDreamer. We first introduce the Transformer State-Space Model, a world model that leverages a transformer for dynamics predictions. We then share this world model with a transformer-based policy network and obtain stability in training a transformer-based RL agent. In experiments, we apply the proposed model to 2D visual RL and 3D first-person visual RL tasks both requiring long-range memory access for memory-based reasoning. We show that the proposed model outperforms Dreamer in these complex tasks.

Cited by

• Bridging MARL to SARL: An Order-Independent Multi-Agent Transformer via Latent Consensus
• Biased Dreams: Limitations to Epistemic Uncertainty Quantification in Latent Dynamics Models
• Spectral Representation-based Reinforcement Learning
• World Models Can Leverage Human Videos for Dexterous Manipulation
• RynnVLA-002: A Unified Vision-Language-Action and World Model
• Counterfactual World Models via Digital Twin-conditioned Video Diffusion
• Object-Centric World Models for Causality-Aware Reinforcement Learning
• Learning to Focus: Prioritizing Informative Histories with Structured Attention Mechanisms in Partially Observable Reinforcement Learning
• When Object-Centric World Models Meet Policy Learning: From Pixels to Policies, and Where It Breaks
• Off-policy Reinforcement Learning with Model-based Exploration Augmentation
• Compositional Monte Carlo Tree Diffusion for Extendable Planning
• A Comprehensive Survey on World Models for Embodied AI
• Vector Quantization in the Brain: Grid-like Codes in World Models
• MultiModal Action Conditioned Video Generation
• DyMoDreamer: World Modeling with Dynamic Modulation
• Embodied AI: From LLMs to World Models
• Remote Sensing-Oriented World Model
• TransZero: Parallel Tree Expansion in MuZero using Transformer Networks
• STORI: A Benchmark and Taxonomy for Stochastic Environments
• Learning Primitive Embodied World Models: Towards Scalable Robotic Learning
• Edge General Intelligence Through World Models and Agentic AI: Fundamentals, Solutions, and Challenges

Related

• TransDreamerV3: Implanting Transformer In DreamerV3
• Transformer-based World Models Are Happy With 100k Interactions
• On Transforming Reinforcement Learning by Transformer: The Development Trajectory
• Facing Off World Model Backbones: RNNs, Transformers, and S4
• Learning Transformer-based World Models with Contrastive Predictive Coding