# Trajectory World Models for Heterogeneous Environments

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 66 of its 75 references that are in the map.
- First posted on arXiv 3 Feb 2025.

## Placed in

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

## Builds on (most cited first, 30 of 66)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [OpenVLA: An Open-Source Vision-Language-Action Model](https://mlsanity.com/paper/2406.09246/) (2024)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [ViViT: A Video Vision Transformer](https://mlsanity.com/paper/2103.15691/) (2021)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Cosmos World Foundation Model Platform for Physical AI](https://mlsanity.com/paper/2501.03575/) (2025)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Genie: Generative Interactive Environments](https://mlsanity.com/paper/2402.15391/) (2024)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Implicit Behavioral Cloning](https://mlsanity.com/paper/2109.00137/) (2021)
- [DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning](https://mlsanity.com/paper/2411.04983/) (2024)
- [Axial Attention in Multidimensional Transformers](https://mlsanity.com/paper/1912.12180/) (2019)
- [Doubly Robust Off-policy Value Evaluation for Reinforcement Learning](https://mlsanity.com/paper/1511.03722/) (2015)
- [Latent Action Pretraining from Videos](https://mlsanity.com/paper/2410.11758/) (2024)
- [Behavior Transformers: Cloning k modes with one stone](https://mlsanity.com/paper/2206.11251/) (2022)

## Built on by (most cited first, 4 of 4)

- [Vid2World: Crafting Video Diffusion Models to Interactive World Models](https://mlsanity.com/paper/2505.14357/) (2025)
- [A Comprehensive Survey on World Models for Embodied AI](https://mlsanity.com/paper/2510.16732/) (2025)
- [RLVR-World: Training World Models with Reinforcement Learning](https://mlsanity.com/paper/2505.13934/) (2025)
- [WestWorld: A Knowledge-Encoded Scalable Trajectory World Model for Diverse Robotic Systems](https://mlsanity.com/paper/2603.14392/) (2026)

## Abstract (arXiv)

Heterogeneity in sensors and actuators across environments poses a significant challenge to building large-scale pre-trained world models on top of this low-dimensional sensor information. In this work, we explore pre-training world models for heterogeneous environments by addressing key transfer barriers in both data diversity and model flexibility. We introduce UniTraj, a unified dataset comprising over one million trajectories from 80 environments, designed to scale data while preserving critical diversity. Additionally, we propose TrajWorld, a novel architecture capable of flexibly handling varying sensor and actuator information and capturing environment dynamics in-context. Pre-training TrajWorld on UniTraj yields substantial gains in transition prediction, achieves a new state-of-the-art for off-policy evaluation, and also delivers superior online performance of model predictive control. To the best of our knowledge, this work, for the first time, demonstrates the transfer benefits of world models across heterogeneous and complex control environments. Code and data are available at https://github.com/thuml/TrajWorld.