Research map / Reinforcement learning / Visual and representation RL
Trajectory World Models for Heterogeneous Environments
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Pre-trains a flexible Transformer world model (TrajWorld) with interleaved temporal and variate attention on UniTraj, a 1.3M-trajectory, 80-environment heterogeneous dataset, showing positive transfer for transition prediction, off-policy evaluation, and MPC.
Problem
Pre-training world models across heterogeneous control environments whose sensor and actuator dimensions and semantics differ, where naive padding or per-environment heads hinder transfer.
Approach
Flatten trajectories into a timestep-by-variate scalar grid, discretize each scalar into bins, and process with interleaved causal temporal attention (per variate) and variate attention (per timestep), so environment identity is inferred in-context from history rather than given as an ID.
Key results
- TrajWorld with UniTraj pre-training achieves a new state of the art on DOPE off-policy evaluation, outperforming ETM in average normalized absolute error and rank correlation, while slightly underperforming on Regret@1.
- Fine-tuned TrajWorld achieves lower average prediction error than its non-pre-trained counterpart in 12 of 15 training datasets across the 75 train-test pairs.
- MLP Ensemble suffers negative transfer from pre-training, and TDM shows negative transfer in MPC with proposal policies despite positive transfer in other tasks.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Pre-train TrajWorld on UniTraj before downstream fine-tuning | TrajWorld · Transition prediction across 75 D4RL train-test dataset pairs | MAE | Helped |
| Pre-train the MLP Ensemble on UniTraj | MLP Ensemble · Transition prediction across D4RL datasets | MAE | Hurt |
| Pre-train TDM on UniTraj | TDM · Transition prediction across D4RL datasets | MAE | Helped |
| Use pre-trained TrajWorld instead of its non-pre-trained variant | TrajWorld · DOPE benchmark | normalized absolute error | Helped |
| Use pre-trained TrajWorld instead of its non-pre-trained variant | TrajWorld · DOPE benchmark | rank correlation | Helped |
| Pre-train the MLP Ensemble on UniTraj | MLP Ensemble · DOPE benchmark | OPE performance | Hurt |
| Use TrajWorld for MPC with proposal policies | MPC · HalfCheetah, Hopper, and Walker2D | online performance | Helped |
| Apply MPC to the proposal policy | MPC with proposal policies · HalfCheetah | online performance | No clear effect |
| Use random shooting rather than proposal policies | MPC · Three D4RL environments | online performance | Hurt |
| Pre-train TrajWorld on a larger and more diverse UniTraj subset, rather than a 1/10 sample, 1/100 sample, or JAT-only subset (ablation) | TrajWorld · Transition prediction and MPC | generalization performance | Helped |
| Use historical context rather than excluding it (ablation) | TrajWorld · Zero-shot Gym Pendulum parameter transfer | prediction error | Helped |
Surprising findings
- MLP Ensemble with pre-training performs worse than its non-pre-trained counterpart (negative transfer).
- TDM exhibits negative transfer in MPC with proposal policies despite benefiting in transition prediction and OPE.
Limitations
- Discretization constrains predictions to a fixed range, making extreme out-of-distribution transitions hard to model.
- Larger capacity than MLPs poses calibration challenges for uncertainty-critical uses such as offline RL.
- Regret@1 slightly underperforms ETM, attributed to bounded reward prediction.
Experimental setup
Datasets: UniTraj ExORL RL Unplugged JAT DB-1 TD-MPC2 replay buffers Modular RL (self-collected) D4RL HalfCheetah/Hopper/Walker2D DOPE MACAW Walker2D
Benchmarks: D4RL HalfCheetah/Hopper/Walker2D transition prediction (75 train-test pairs) DOPE off-policy evaluation MPC on D4RL Hopper/Walker2D/HalfCheetah DMC Cart-2-Pole/Cart-3-Pole zero-shot
Metrics: mean absolute error rank correlation Regret@1 average return under MPC
Baselines: MLP Ensemble TDM ETM FQE DR IS DICE VPM
Builds on: Transformer Dynamics Model (TDM) DOPE benchmark D4RL
Closest papers by meaning
- UniTraj: Learning a Universal Trajectory Foundation Model from Billion-Scale Worldwide Traces (NeurIPS 2025 · similarity 0.86)
- Time-Aware World Model for Adaptive Prediction and Control (ICML 2025 · similarity 0.86)
- From Forecasting to Planning: Policy World Model for Collaborative State-Action Prediction (NeurIPS 2025 · similarity 0.86)
- STraj: Self-training for Bridging the Cross-Geography Gap in Trajectory Prediction (AAAI 2025 · similarity 0.85)
- COME: Adding Scene-Centric Forecasting Control to Occupancy World Model (NeurIPS 2025 · similarity 0.85)
- Markovian Dynamics Enforcer: Feasibility Preserving Correction on Learned Dyanmics Manifolds (NeurIPS 2026 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.