[Paper Review] TransDreamer: Reinforcement Learning with Transformer World Models

Chang Chen, Yifu Wu+2 more|arXiv (Cornell University)|Feb 19, 2022

Reinforcement Learning in Robotics8 citations

TL;DR

TransDreamer introduces a fully transformer-based model-based reinforcement learning agent that replaces the recurrent world model in Dreamer with a novel Transformer State-Space Model (TSSM), enabling improved long-range memory and planning. It outperforms Dreamer on complex 2D and 3D visual control tasks requiring long-term reasoning, while matching its performance on simpler tasks with short-term dependencies.

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.

Motivation & Objective

• To investigate whether transformer-based world models can improve sample efficiency and long-term reasoning in model-based reinforcement learning.
• To address the challenge of training stable transformer-based policies in MBRL, where prior work shows difficulty in training pure transformer policies from sparse rewards.
• To develop a stochastic, action-conditioned world model compatible with transformers that supports parallel training and effective imagination.
• To evaluate whether the improved world model leads to better performance on tasks requiring complex, long-range memory interactions.
• To validate that the proposed framework maintains strong performance on simpler tasks where RNN-based models are expected to succeed.

Proposed method

• Proposes the Transformer State-Space Model (TSSM), a stochastic, action-conditioned world model that uses self-attention to capture long-range dependencies in latent space.
• Adapts the TSSM to support both prior and posterior inference for stochastic dynamics prediction, enabling accurate rollouts in imagination.
• Shares the TSSM world model between the policy and world model components, enabling end-to-end training with consistent latent representations.
• Uses a variational inference objective (ELBO) to train the TSSM, with stochastic latent states and deterministic hidden states updated via multi-head self-attention.
• Employs a transformer-based actor-critic policy trained on imagined trajectories generated by the TSSM, enabling sample-efficient learning.
• Integrates the TSSM into the Dreamer framework, replacing the RSSM while preserving the overall MBRL training loop.

Experimental results

Research questions

• RQ1 Can a transformer-based world model outperform recurrent models in long-horizon visual control tasks requiring complex memory interactions?
• RQ2 What are the key architectural and training challenges in building a stable, end-to-end transformer-based MBRL agent?
• RQ3 Does a transformer-based world model improve the quality of imagined trajectories, particularly in image generation and reward prediction?
• RQ4 How does the performance of a transformer-based MBRL agent compare to a recurrent baseline (Dreamer) on tasks with short-term dependencies?
• RQ5 Can the TSSM world model generalize to diverse visual environments while maintaining training stability and computational efficiency?

Key findings

• TransDreamer outperforms Dreamer on the 5-Ball Dense environment, a complex 2D visual control task requiring long-term memory and precise ball-order reasoning.
• TransDreamer generates sharper, more accurate imagined frames and correctly predicts reward spikes at ball collection events, while Dreamer produces blurry images and incorrect reward predictions.
• The TSSM world model generates more accurate future frame predictions and reward estimates than Dreamer’s RSSM, especially in later imagination steps where long-range dependencies matter.
• On short-term memory tasks (e.g., DMC Cheetah Run and Atari games), TransDreamer achieves comparable final performance to Dreamer, though with slower convergence.
• TransDreamer shows improved reward prediction accuracy and visual imagination fidelity, indicating a more robust and precise world model than Dreamer’s RNN-based counterpart.
• The TSSM enables stable training of a fully transformer-based RL agent, demonstrating that transformer-based MBRL is viable and effective when proper architectural design is applied.