Saved in:
Bibliographic Details
Main Authors: Yin, Shaofeng, Wu, Jialong, Huang, Siqiao, Su, Xingjian, He, Xu, Hao, Jianye, Long, Mingsheng
Format: Preprint
Published: 2025
Subjects:
Online Access:https://arxiv.org/abs/2502.01366
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908398328479744
author Yin, Shaofeng
Wu, Jialong
Huang, Siqiao
Su, Xingjian
He, Xu
Hao, Jianye
Long, Mingsheng
author_facet Yin, Shaofeng
Wu, Jialong
Huang, Siqiao
Su, Xingjian
He, Xu
Hao, Jianye
Long, Mingsheng
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2502_01366
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Trajectory World Models for Heterogeneous Environments
Yin, Shaofeng
Wu, Jialong
Huang, Siqiao
Su, Xingjian
He, Xu
Hao, Jianye
Long, Mingsheng
Machine Learning
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.
title Trajectory World Models for Heterogeneous Environments
topic Machine Learning
url https://arxiv.org/abs/2502.01366