title: Improving Transformer World Models for Data-Efficient RL
summary: 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.
paper_authors._id: 67a1a4b72bf092a7612b36eb
paper_authors.name: Antoine Dedieu
paper_authors._id: 67a1a4b72bf092a7612b36ec
paper_authors.name: Joseph Ortiz
paper_authors._id: 67a1a4b72bf092a7612b36ed
paper_authors.name: Xinghua Lou
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-02-04T21:08:45.474Z
paper_authors._id: 665c8e13a1ff38df9706379e
paper_authors.avatarUrl: /avatars/e0aecfac58ff98c628fd57afee53f791.svg
paper_authors.fullname: xinghua Lou
paper_authors.type: user
paper_authors.user: nickname-xingxing
paper_authors._id: 67a1a4b72bf092a7612b36ee
paper_authors.name: Carter Wendelken
paper_authors._id: 67a1a4b72bf092a7612b36ef
paper_authors.name: Wolfgang Lehrach
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-02-04T21:08:36.618Z
paper_authors._id: 6310ec5d64939fabc00aea54
paper_authors.avatarUrl: /avatars/d3fde9392fd30a6d80e2e7989ed4db17.svg
paper_authors.fullname: Wolfgang Lehrach
paper_authors.type: user
paper_authors.user: wpl
paper_authors._id: 67a1a4b72bf092a7612b36f0
paper_authors.name: J Swaroop Guntupalli
paper_authors._id: 67a1a4b72bf092a7612b36f1
paper_authors.name: Miguel Lazaro-Gredilla
paper_authors._id: 67a1a4b72bf092a7612b36f2
paper_authors.name: Kevin Patrick Murphy
paper_title: Improving Transformer World Models for Data-Efficient RL
paper_summary: 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.
paper_ai_summary: A model-based RL approach using CNNs, RNNs, Dyna with warmup, nearest neighbor tokenizer, and block teacher forcing achieves superior performance on the Craftax-classic benchmark with sample efficiency improvements.
paper_ai_keywords: CNNs RNNs Dyna with warmup nearest neighbor tokenizer transformer world model (TWM) block teacher forcing