Trajectory World Models for Heterogeneous Environments

2025-02-03

AI Summary

To address the weak generalization and poor transferability of world models in multi-source heterogeneous sensor-actuator environments, this paper proposes TrajWorld—a trajectory-based world model capable of cross-environment generalization. Our method introduces three key innovations: (1) UniTraj, the first unified trajectory dataset supporting modality-agnostic representation and environment-aware adaptation; (2) a dynamic adaptation architecture integrating modality-adaptive encoders, in-context environmental modeling, and cross-environment data standardization and alignment; and (3) the first successful transfer of a world model across 80 heterogeneous control environments. Experiments demonstrate state-of-the-art performance on both state-transition prediction and off-policy evaluation tasks, significantly enhancing the generalization capability and practical utility of large-scale pre-trained world models—particularly under low-dimensional perceptual inputs.

Technology Category

• Intelligent Robots: Multimodal Perception & Sensor Fusion
• Multiagent Systems: Adversarial Agents
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

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 demonstrates significant improvements in transition prediction and achieves a new state-of-the-art for off-policy evaluation. 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.

Problem

• World Model
• Pre-training
• Model Generalization

Innovation

• UniTraj
• TrajWorld
• Environment Prediction