PaperlayerView exact revision on arXiv

Source-linked AI summary

TransDreamer: Reinforcement Learning with Transformer World Models

Chang Chen, Yi-Fu Wu, Jaesik Yoon, Sungjin Ahn

arXiv:2202.09481v2cs.LG

TL;DR

Dreamer’s recurrent components limit direct access to long-term memory, motivating a transformer-based MBRL agent. The paper introduces TSSM and TransDreamer, then shows improved performance over Dreamer on complex memory tasks while remaining comparable on simpler tasks. It also reports better world-model image and reward prediction and stable transformer RL training.

  • Problem

    RNN-based MBRL models limit long-term memory access, while training transformer-based RL policies from rewards alone is difficult.

  • Method

    TransDreamer combines a transformer-based stochastic world model, TSSM, with a transformer-based policy within the Dreamer framework.

  • Results

    TransDreamer outperforms Dreamer on long-term complex-memory tasks, is comparable on simple DMC and Atari tasks, and has better image generation and reward prediction.

  • Takeaways & Limitations

    Transformer world modeling can extend Dreamer to tasks requiring long-term complex memory interactions while preserving comparable performance on simpler tasks.

  • Takeaways & Limitations

    The posterior representation model omits h_t and past stochastic states, relying on deterministic states to encode temporal information.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper asks how MBRL can gain transformers’ long-range memory benefits while addressing the difficulty of training transformer-based RL agents. It proposes TransDreamer and reports stronger performance than Dreamer on complex memory tasks, with comparable performance on simpler tasks.

  • MBRL offers sample-efficient imagination training, reusable environment knowledge, and planning through a learned dynamics model.
  • Dreamer uses latent-space world modeling and outperforms model-free Rainbow and IQN in visual MBRL.
  • RNN-based MBRL dynamics models are limited for long-term dependencies, whereas transformers provide direct memory access for memory-based reasoning.
  • TransDreamer is a transformer-based MBRL agent built on Dreamer.
  • TSSM and TransDreamer address the need for stochastic action-conditioned transitions, parallel trainability, and workable transformer-based RL design.
  • TransDreamer outperforms Dreamer on long-term, complex-memory tasks while remaining comparable on simple DMC and Atari tasks.

2 PRELIMINARIES

Dreamer learns a stochastic latent world model, trains its policy on imagined trajectories, and uses differentiable dynamics for lower-variance policy optimization.

  • Dreamer cycles through world-model learning, policy learning, and environment interaction using replay-buffer experience.
  • RSSM represents each latent state as stochastic z_t and deterministic h_t components updated by probabilistic sampling and recurrent dynamics.
  • The stochastic state supports multiple hypothetical futures, while latent rollouts avoid generating observation images during imagination.
  • RSSM learning uses an evidence lower bound with representation, observation, and reward models.
  • Dreamer trains its policy from imagined trajectories generated by the learned RSSM rather than by interacting with the environment.
  • Backpropagation through the differentiable world model provides lower-variance gradients than REINFORCE for policy learning.

3 TRANSDREAMER

TransDreamer replaces Dreamer’s recurrent world model with TSSM, a stochastic transformer state-space model designed for direct historical access, parallel training, and sequential imagination. Its fixed transformer parameters also avoid the reported RL stability issues, while memory demands require fewer imagined trajectories.

  • Transformers directly access historical states and model complex long-term temporal dependencies, motivating their use in TransDreamer.
  • The transformer world model is designed to access past states, update time steps in parallel, roll out sequentially, and remain stochastic.
  • 3.1 TRANSFORMER STATE SPACE MODEL (TSSM): TSSM replaces RSSM’s recurrent update with a transformer, enabling direct historical interaction and removing the main sequential computation during training.
  • 3.1 TRANSFORMER STATE SPACE MODEL (TSSM): TSSM removes h_t from the posterior representation model, using q(z_t|x_t), so transformer inputs and state updates can be computed in parallel.
  • 3.1 TRANSFORMER STATE SPACE MODEL (TSSM): The posterior simplification assumes temporal information can be represented in deterministic states, while imagination still uses the h_t-conditioned prior.
  • 3.2 POLICY LEARNING AND IMPLEMENTATION DETAILS: TransDreamer retains Dreamer’s policy-learning framework while replacing RSSM with TSSM for stronger long-term temporal modeling.
  • 3.2 POLICY LEARNING AND IMPLEMENTATION DETAILS: Holding transformer parameters fixed during agent training avoids the stability issues observed in transformer-based RL, without additional gating.
  • 3.2 POLICY LEARNING AND IMPLEMENTATION DETAILS: Transformer memory requirements make it infeasible to generate imagined trajectories from every sampled replay-buffer state, so TransDreamer uses a smaller subset.

4 RELATED WORKS

The related work connects TransDreamer to transformer applications, stochastic transformer models, and earlier model-based reinforcement-learning frameworks, especially Dreamer and RSSM.

  • Prior work applies transformers across language, vision, and video generation, while transformer RL methods address stability with gating layers.
  • Stochasticity has been incorporated into transformer architectures for generation, translation, completion, layout, and uncertainty quantification.
  • TransDreamer builds on Dyna, Dreamer, and RSSM-based latent-space planning within the broader MBRL literature.

5 EXPERIMENTS

Experiments compare TransDreamer with Dreamer across long-term-memory tasks, world-model prediction, imagined trajectories, and short-term-memory benchmarks. TransDreamer performs better on complex memory-dependent tasks and world-model predictions, while eventually reaching comparable performance on simpler tasks.

  • Evaluation scope: The experiments evaluate long-term memory and reasoning, world-model quality, and short-term-memory performance against Dreamer.The evaluation includes 2D and 3D Hidden Order Discovery tasks, quantitative and qualitative world-model comparisons, and DMC and Atari sanity checks.
  • 2D Hidden Order Discovery: TransDreamer outperforms Dreamer across 4-, 5-, and 6-ball 2D Hidden Order Discovery configurations.In the 4-ball setting, TransDreamer reaches about 7 average episode reward versus Dreamer’s about 4.
  • 2D Hidden Order Discovery: 23% versus 7% success rates show that TransDreamer more often collects all four balls in the correct order than Dreamer.Performance degrades for both agents as the number of balls increases, because more balls create more possible orders to test.
  • 3D Hidden Order Discovery: The 3D Unity task increases partial observability and navigation distance, providing a more demanding test of long-term dependency and complex reasoning.Dense and sparse settings vary the number of balls and their separation.
  • World-model quantitative results: TransDreamer generally achieves lower or comparable image MSE and more accurate reward prediction than Dreamer, especially with longer contexts.The reward analysis focuses on nonzero +3 rewards, and TransDreamer continues improving with additional context in the 5-ball dense setting.
  • World-model qualitative results: TransDreamer predicts later imagined states, ball disappearances, and rewards more accurately, whereas Dreamer’s imagined trajectories become increasingly incorrect.The qualitative comparison attributes TransDreamer’s stronger world-model quality particularly to later imagination steps where long-term memory matters more.
  • Short-term-memory tasks: On DMC and Atari short-term-memory tasks, both agents eventually reach comparable performance, but TransDreamer generally saturates more slowly.TransDreamer is slightly better and faster on DMC Cheetah Run.

6 CONCLUSION

TransDreamer combines a transformer-based stochastic world model with Dreamer and is comparable on short-memory tasks while outperforming Dreamer on long-term memory tasks. Its world model also improves image generation and reward prediction, while exploration-heavy Atari games remain outside scope.

  • TransDreamer uses the Transformer State-Space Model, described as the first transformer-based stochastic world model.
  • TransDreamer performs comparably to Dreamer on DMC and Atari tasks without long-term memory, but outperforms Dreamer on Hidden Order Discovery tasks requiring complex long-term memory.
  • The TSSM achieves better image generation and reward prediction than Dreamer, qualitatively and quantitatively.
  • The evaluation excludes some Atari games requiring both long-term memory and a good exploration policy because exploration is not addressed.

A.1 DREAMERV2

DreamerV2 modifies Dreamer for Atari by using discrete states and gradient estimators that address differentiation through those states and actions.

  • DreamerV2 replaces the continuous stochastic hidden state with a discrete state for Atari environments.
  • DreamerV2 uses straight-through gradients and REINFORCE or mixed gradients to differentiate through discrete states and actions.

A.2 TRANSDREAMER LOSS FUNCTION

The TransDreamer loss optimizes a negative ELBO for an action-conditioned model while also predicting observations, rewards, and discounts. Separate hyperparameters scale these loss terms.

  • The objective is the negative ELBO of the action-conditioned model with additional terms for reward and discount prediction.
  • The loss includes terms for predicting observations, rewards, and discounts from the latent model.
  • ηx, ηr, and ηγ scale the observation, reward, and discount loss terms, respectively.

A.2.1 ELBO

The ELBO uses a transformer-conditioned generative model and an approximate variational posterior, then expands into reconstruction, reward, discount, and KL terms.

  • The generative model predicts observations from latent states and transformer summaries conditioned on preceding latent states and actions.
  • The variational posterior approximates each latent state using q(zt|xt).
  • The expanded objective contains log-likelihood terms for observations, rewards, and discounts.

A.3 DMC AND ATARI

The DMC and Atari experiments use configurations largely matched to Dreamer and DreamerV2, with TSSM-specific transformer hyperparameters introduced for TransDreamer.

  • Configuration: Dreamer and DreamerV2 use almost identical configurations, while TransDreamer changes the number of imagined trajectories because imagining from every batch state is too computationally expensive.The modified trajectory count is necessary for TransDreamer’s imagination procedure.
  • TSSM hyperparameters: TSSM-specific choices include gating, identity map reordering, transformer depth and heads, MLP hidden size, and relational positional embeddings.For DMC and Atari, the typical configuration uses a 2-layer transformer without dropout, gating, or identity map reordering.

A.4 HIDDEN ORDER DISCOVERY

Hidden Order Discovery evaluates memory-based reasoning and long-term knowledge through increasingly difficult 2D and 3D tasks. TransDreamer generally achieves stronger completion, reward prediction, and imagination quality than Dreamer, including on challenging contexts.

  • A.4 Hidden Order Discovery: The 2D and 3D tasks increase object count, inter-object distance, or both to test complex memory-based reasoning and long-term knowledge extraction.In 2D, the number of balls increases from 4 to 6 while distance remains unchanged.
  • A.4 Hidden Order Discovery: The Hidden Order Discovery experiments use a 6-layer TSSM with identity map reordering and imagine from one randomly sampled starting state until the trajectory limit.Other hyperparameters follow the DreamerV2 Crafter configuration.
  • Success rate: TransDreamer performs better than Dreamer on the ratio of trajectories completing at least one hidden-order round, while episode length is limited to 100 steps.As the number of balls increases from 4 to 5, TransDreamer’s performance changes from 23% to 5%, whereas Dreamer nearly fails.
  • Image generation: More than half of the overall image-generation MSE gap occurs in the foreground, which is isolated using a Unity-rendered foreground mask.The full MSE results are reported in Table 4.
  • Reward prediction: TransDreamer has a large reward-prediction advantage over Dreamer on nonzero-reward timesteps in the 4-Ball Dense and 5-Ball Dense settings, while both generally perform well on zero-reward timesteps.The gap is smaller in the 4-Ball Sparse setting, where foreground balls are seen less frequently.
  • World-model imagination: Given the same context, TransDreamer correctly imagines the balls and rewards for two targets even on an out-of-distribution sample, though it mispredicts one reward and still surpasses Dreamer’s imagination quality.The comparison is illustrated in Figure 8 for a 5-Ball Dense sample.
Loading 2202.09481v2…