# One-shot World Models Using a Transformer Trained on a Synthetic Prior

arXiv: https://arxiv.org/abs/2409.14084 · HTML: https://mlsanity.com/paper/2409.14084/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 23 of its 24 references that are in the map.
- First posted on arXiv 21 Sep 2024.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

## Builds on (most cited first, 23 of 23)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Learning Interactive Real-World Simulators](https://mlsanity.com/paper/2310.06114/) (2023)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [TransDreamer: Reinforcement Learning with Transformer World Models](https://mlsanity.com/paper/2202.09481/) (2022)
- [Transformer-based World Models Are Happy With 100k Interactions](https://mlsanity.com/paper/2303.07109/) (2023)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [Augmented World Models Facilitate Zero-Shot Dynamics Generalization From a Single Offline Environment](https://mlsanity.com/paper/2104.05632/) (2021)
- [Unsupervised Zero-Shot Reinforcement Learning via Functional Reward Encodings](https://mlsanity.com/paper/2402.17135/) (2024)
- [Context is Everything: Implicit Identification for Dynamics Adaptation](https://mlsanity.com/paper/2203.05549/) (2022)
- [Pre-training with Synthetic Data Helps Offline Reinforcement Learning](https://mlsanity.com/paper/2310.00771/) (2023)
- [Learning Synthetic Environments and Reward Networks for Reinforcement Learning](https://mlsanity.com/paper/2202.02790/) (2022)
- [MDP Playground: An Analysis and Debug Testbed for Reinforcement Learning](https://mlsanity.com/paper/1909.07750/) (2019)

## Abstract (arXiv)

A World Model is a compressed spatial and temporal representation of a real world environment that allows one to train an agent or execute planning methods. However, world models are typically trained on observations from the real world environment, and they usually do not enable learning policies for other real environments. We propose One-Shot World Model (OSWM), a transformer world model that is learned in an in-context learning fashion from purely synthetic data sampled from a prior distribution. Our prior is composed of multiple randomly initialized neural networks, where each network models the dynamics of each state and reward dimension of a desired target environment. We adopt the supervised learning procedure of Prior-Fitted Networks by masking next-state and reward at random context positions and query OSWM to make probabilistic predictions based on the remaining transition context. During inference time, OSWM is able to quickly adapt to the dynamics of a simple grid world, as well as the CartPole gym and a custom control environment by providing 1k transition steps as context and is then able to successfully train environment-solving agent policies. However, transferring to more complex environments remains a challenge, currently. Despite these limitations, we see this work as an important stepping-stone in the pursuit of learning world models purely from synthetic data.