Improving Transformer World Models for Data-Efficient RL

• URL: http://arxiv.org/abs/2502.01591v1
• Date: Mon, 03 Feb 2025 18:25:17 GMT
• Title: Improving Transformer World Models for Data-Efficient RL
• Authors: Antoine Dedieu, Joseph Ortiz, Xinghua Lou, Carter Wendelken, Wolfgang Lehrach, J Swaroop Guntupalli, Miguel Lazaro-Gredilla, Kevin Patrick Murphy,
• Abstract summary: We present an approach to model-based RL that achieves a new state of the art performance on the challenging Craftax-classic benchmark.<n>Our algorithm achieves a reward of 67.4% after only 1M environment steps, significantly outperforming DreamerV3, which achieves 53.2%, and, for the first time, exceeds human performance of 65.0%.
• Score: 5.920669613621277
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: We present an approach to model-based RL that achieves a new state of the art performance on the challenging Craftax-classic benchmark, an open-world 2D survival game that requires agents to exhibit a wide range of general abilities -- such as strong generalization, deep exploration, and long-term reasoning. With a series of careful design choices aimed at improving sample efficiency, our MBRL algorithm achieves a reward of 67.4% after only 1M environment steps, significantly outperforming DreamerV3, which achieves 53.2%, and, for the first time, exceeds human performance of 65.0%. Our method starts by constructing a SOTA model-free baseline, using a novel policy architecture that combines CNNs and RNNs. We then add three improvements to the standard MBRL setup: (a) "Dyna with warmup", which trains the policy on real and imaginary data, (b) "nearest neighbor tokenizer" on image patches, which improves the scheme to create the transformer world model (TWM) inputs, and (c) "block teacher forcing", which allows the TWM to reason jointly about the future tokens of the next timestep.

Related papers

• Towards foundational LiDAR world models with efficient latent flow matching [9.86884512471034]
  Existing LiDAR world models are narrowly trained; each model excels only in the domain for which it was built.<n>We conduct the first systematic domain transfer study across three demanding scenarios.<n>Given different amounts of fine-tuning data, our experiments show that a single pre-trained model can achieve up to 11% absolute improvement.
  arXiv  Detail & Related papers  (2025-06-30T00:16:55Z)
• Enter the Void - Planning to Seek Entropy When Reward is Scarce [6.208654695856247]
  We propose a novel approach that anticipates and actively seeks out high-entropy states using short-horizon latent predictions.<n>We present a hierarchical planner that dynamically decides when to replan, planning horizon length, and the weighting between reward and entropy.<n>Our method finishes the Miniworld procedurally generated mazes 50% faster than base Dreamer at convergence and the policy trained in imagination converges in only 60% of the environment steps that base Dreamer needs.
  arXiv  Detail & Related papers  (2025-05-22T15:28:50Z)
• Output Scaling: YingLong-Delayed Chain of Thought in a Large Pretrained Time Series Forecasting Model [55.25659103706409]
  This framework achieves state-of-the-art performance for our designed foundation model, YingLong.<n>YingLong is a non-causal, bidirectional attention encoder-only transformer trained through masked token recovery.<n>We release four foundation models ranging from 6M to 300M parameters, demonstrating superior results in zero-shot tasks.
  arXiv  Detail & Related papers  (2025-05-20T14:31:06Z)
• Learning Transformer-based World Models with Contrastive Predictive Coding [58.0159270859475]
  We show that the next state prediction objective is insufficient to fully exploit the representation capabilities of Transformers. We propose to extend world model predictions to longer time horizons by introducing TWISTER, a world model using action-conditioned Contrastive Predictive Coding. TWISTER achieves a human-normalized mean score of 162% on the Atari 100k benchmark, setting a new record among state-of-the-art methods that do not employ look-ahead search.
  arXiv  Detail & Related papers  (2025-03-06T13:18:37Z)
• Masked Generative Priors Improve World Models Sequence Modelling Capabilities [19.700020499490137]
  Masked Generative Modelling has emerged as a more efficient and superior inductive bias for modelling.<n>GIT-STORM demonstrates substantial performance gains in RL tasks on the Atari 100k benchmark.<n>We apply Transformer-based World Models to continuous action environments for the first time, addressing a significant gap in prior research.
  arXiv  Detail & Related papers  (2024-10-10T11:52:07Z)
• 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)
• DigiRL: Training In-The-Wild Device-Control Agents with Autonomous Reinforcement Learning [61.10299147201369]
  This paper introduces a novel autonomous RL approach, called DigiRL, for training in-the-wild device control agents. We build a scalable and parallelizable Android learning environment equipped with a VLM-based evaluator. We demonstrate the effectiveness of DigiRL using the Android-in-the-Wild dataset, where our 1.3B VLM trained with RL achieves a 49.5% absolute improvement.
  arXiv  Detail & Related papers  (2024-06-14T17:49:55Z)
• STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning [82.03481509373037]
  Recently, model-based reinforcement learning algorithms have demonstrated remarkable efficacy in visual input environments. We introduce Transformer-based wORld Model (STORM), an efficient world model architecture that combines strong modeling and generation capabilities. Storm achieves a mean human performance of $126.7%$ on the Atari $100$k benchmark, setting a new record among state-of-the-art methods.
  arXiv  Detail & Related papers  (2023-10-14T16:42:02Z)
• Aligning Large Multimodal Models with Factually Augmented RLHF [176.54751941088819]
  Large Multimodal Models (LMM) are built across modalities and misalignment between two modalities can result in "hallucination" We adapt the Reinforcement Learning from Human Feedback (RLHF) from the text domain to the task of vision-language alignment. We propose a new alignment algorithm called Factually Augmented RLHF that augments the reward model with additional factual information. Our approach achieves remarkable improvement on the LLaVA-Bench dataset with the 94% performance level of the text-only GPT-4.
  arXiv  Detail & Related papers  (2023-09-25T20:59:33Z)
• Emergent Agentic Transformer from Chain of Hindsight Experience [96.56164427726203]
  We show that a simple transformer-based model performs competitively with both temporal-difference and imitation-learning-based approaches. This is the first time that a simple transformer-based model performs competitively with both temporal-difference and imitation-learning-based approaches.
  arXiv  Detail & Related papers  (2023-05-26T00:43:02Z)
• Mastering the Unsupervised Reinforcement Learning Benchmark from Pixels [112.63440666617494]
  Reinforcement learning algorithms can succeed but require large amounts of interactions between the agent and the environment. We propose a new method to solve it, using unsupervised model-based RL, for pre-training the agent. We show robust performance on the Real-Word RL benchmark, hinting at resiliency to environment perturbations during adaptation.
  arXiv  Detail & Related papers  (2022-09-24T14:22:29Z)
• Generative Multi-Stream Architecture For American Sign Language Recognition [15.717424753251674]
  Training on datasets with low feature-richness for complex applications limit optimal convergence below human performance. We propose a generative multistream architecture, eliminating the need for additional hardware with the intent to improve feature convergence without risking impracticability. Our methods have achieved 95.62% validation accuracy with a variance of 1.42% from training, outperforming past models by 0.45% in validation accuracy and 5.53% in variance.
  arXiv  Detail & Related papers  (2020-03-09T21:04:51Z)

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.