Trajectory World Models for Heterogeneous Environments
Authors: Shaofeng Yin
Jialong Wu
Siqiao Huang
Xingjian Su
Xu He
Jianye HAO
Mingsheng Long
Venue: ICML 2025 poster (2025)
Keywords: world models
pre-training
heterogeneous environments
Primary area: reinforcement_learning->deep_rl
TL;DR: We introduce UniTraj and TrajWorld, a unified dataset and flexible architecture to enable positive transfer when pre-training world models across heterogeneous environments.

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 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.

Trajectory World Models for Heterogeneous Environments
Shaofeng Yin
Jialong Wu Siqiao Huang' Xingjian Su
Xu He
Jianye Hao
Mingsheng Long'
Abstract
D4RL (75 train-test dataset pairs)
Heterogeneity in sensors and actuators across en- vironments poses a significant challenge to build- ing large-scale pre-trained world models on top 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 pre- serving critical diversity. Additionally, we pro- pose TrriWo novel architecture capable of flexibly handling varying sensor and actuator in- formation and capturing environment dynamics in-context. Pre-training TrajWorld on UniTraj yields substantial gains in transition prediction, achieves a new state-of-the-an for off-policy eval- uation, and also delivers superior online perfor- mance 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 envi- ronments. Code and data are available at https: //git com/ whtul/Traajoo
Normalized Prediction Error
91%
wlo Pre-Training
18.85
w/ Pre- Pre-Training
2.08
2x10"
+35%
+9%
10'
1.00
0.91
MLP Ensemble
TDM
TrajWorld (Ours)
Figure I. Aggregated transition prediction error (MAE) across 75 train-tes: dataset pairs, comparing MLP Ensemble (Chua et al., 2018), TDM (Schubert et al., 2023), and proposed TrajWorld, with and without pre-training on UniTraj dataset. Y-axis at log scale.
existing methods often learn world models tabula rasa, rely- ing on data from single, specific environment. This limits their ability to generalize to catoo--distribution transitions, demanding substantial number of costly interactions with the environment.
In recent years, machine learning has been revolutionized by foundation models pre-trained on large-scale, diverse data (Achiam et al., 2023; Oquab et al., 2024; Kirillov et al., 2023). General world models have also been realized through pre-training, enabled by the homogeneity present within massive and diverse datasets specific modalities, such text (Wang et al., 2024b; Gu al., 2024; Chae et al., 2025; Wu al., 2025), images (Zhou al., 2024), and videos (Seo al., 2022; Wu et al., 2024a; Bruce et al., 2024; Wu et al., 2024b; Agarwal et al., 2025). However, a unique challenge of world models from Internet Al commonly overlooked or circumvented: the heterogeneity inherent in sensor and actuator information, which means proprioceptive data, such as joint positions and velocities, well as optional target positions, vary significantly across environments. Failing properly address this heterogeneity can result in no transfer even negative transfer.
1. Introduction
World models (Ha & Schmidhuber, 2018; LeCun, 2022) have made remarkable progress in addressing sequential decision- deeision-making problems (Hafner al., 2020; Schrit- twieser al., 2020; Hansen et al., 2024). Trained on trajec- tory data, these models simulate environments and are leveraged to either evaluate complex actions (Chua et al., 2018; Ebert et al., 2018; Tian et al., 2023) or optimize poli- cies (Janner et al., 2019; Kurutach et al., 2018). However,
'Equal contribution 'Tsinghua University. Huuawi Noah's Ark Lab.
Shaofeng Yin <ysf22 @ @em22mmail.cem wwaalls.sinngu...u.... Jialong Wu <wujjalong02296 ema.l.gcmm...oo..... com>. Correspondence Mingsheng Long <mingsheng @tsinghua.cdu.cn>>.
We argue that no modality in world models should be left behind, including essential sensor information represented as low-dimensional vectors. In this work, we take first step
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Trajectory World Models for Heterogeneous Environments
(15,
(11,3)
(18
25
Trajectory World Model
(17. 6)
(17.6)
C10 73
(42.9)
(13,4)
(17. 6)
(4,1)
(9.2)
(4,3)
Pre-training
Fine-tuning
G2,
(slmit)) diun(A)
Figure 2. Illustration pre-training a world model from heterogencous environments, with each environment labeled by its state and action dimensions. A Trajectory World Model, designed for flexibility in handling divergent state and action definitions, demonstrates effective positive transfer across distinct, heterogeneous, and complex control environments.
to bridge this gap by exploring the potential of pre-training a world model to extract shared knowledge from trajectories across heterogeneous environments (illustrated in Figure 2). To this end, it is essential to overcome the transfer barriers from both data and model architecture perspectives.
is enabled to naturally accommodate varying numbers of sensors and actuators through variate attention and, more importantly, to capture their relationships in-context through temporal attention. This in-context learning capability goes beyond learning specific environment dynamics and thus enhances the model's generalizability across environments.
Scaling data. To achieve strong generalization through pre-training, access to vast and diverse data is essential (Team et al., 2021). While scaling data is straightforward, the real challenge lies in scaling data while preserving di- versity. Diversity in our work has two key aspects. First, it refers to the data sources, i.e., the environments from which the data is collected. Second, it concerns the data proper- ties, specifically the distribution of the data itself. Even within the same environment, different policies various levels can produce significantly different data distributions. To tackle these challenges, we curate the UniTraj dataset, including over one million trajectories collected from vari- ous distributions from 80 heterogeneous environments. By scaling data while maintaining these diversities, we ensure that the model focuses on the core knowledge shared across environments, thereby enabling successful transferability.
By pre-training our flexible Traj architecture on the di- verse and massive UniTraj dataset, we demonstrate, for the first time, the transfer benefits of world models across het- erogeneous and complex control environments. Fine-tuning TrajWorke on 15 datasets from three previously unseen en- vironments (Fu al., 2020) significantly reduces transi- tion prediction errors for both in-distribution and out-of- distribution actions (as shown in Figure 1). This improved predictive accuracy also translates to our state-of-the-art performance on off-policy evaluation (OPE) tasks (Fu et al., 2021), enabling the offline evaluation and selection set of complex policies for best performance. Furthermore, also manifests in superior online performance with model predictive control (MPC).
The main contributions can be summarized as follows:
Flexible architecture. Previous approaches often address size variations in state and action spaces by applying zero- padding to match a maximum length (Yu et al., 2020a: Hansen et al., 2024) or employing separate input and output heads for each environment (Wang al., 2024a; E'rao al., 2020). However, zero-padding imposes a dimension limit and adds training overheads, while the separate head approach requires training new heads for new environments, hindering zero-shot transfer. A truly capable model for heterogeneous environments requires a more flexible archi- tecture. To address this, we propose the Trajectory World Model World), a novel architecture that integrates in- terleaved variate and temporal attention mechanisms. It
We investigate an under-explored world model pre- training paradigm across heterogeneous environments.
We curate UniTraj, a unified trajectory dataset, en- abling large-scale pre-training of world models.
We propose Traj World, a novel architecture to facilitate transfer between heterogeneous environments.
For the first time, our experiments demonstrate posi- tive world model transfer across diverse and complex environments, resulting in simultaneous and signifi- cant improvements in transition prediction, off-policy policy, and model predictive control.
2

---

Trajectory World Models for Heterogeneous Environments
Table Statistics for six components of the UniTraj dataset. The checkmark (V) represents a dataset collected or curated by ourselves.
Dataset
#Env
#Episodes
#Steps
State dim.
Action dim.
Characteristic
ExORL (Yarats al., 2022)
4
541,336
330,985,000
78
1 12
Exploratory
RL Unplugged (Gulcehre et al., 2020)
7
5,841
5,841,000
~ 67
1 ~
Experience replay
JAT (Gallouédec et al., 2024)
5
50,000
26,238,954
4~ 23
1~ 7
Expert
DB-1 (Wen et al., 2022)
58
290
19,320
5 67
1~
Expert: Diversity
TD-MPC2 (Hansen et al., 2024)
30
672.000
336,000,000
3~ 24
1~6
Experience replay
Modular RI. (Huang al., 2020)
20
37,199
19,996,902
7 23
1~6
Experience replay
UniTraj (Ours)
80
1,306,666
719,081,176
3 78
21
Omnifarious
2. Problem Formulation
teristics. To further enhance diversity, we also by ourselves collect the training buffer of agents from a set of diverse morphologies (Huang et al., 2020). As result, Uni oc- cupies total of 1.3M trajectories (or 719M steps) from 80 distinct environments, as summarized in Table 1. A detailed list of dataset information can be found in Appendix A.
An environment is typically described by a Markov decision process (MDP) M = {S, AA, P. r. ()), specified by the state space S (of sensors), the action space A (of actuators), the transition function P: S A + A(S), the reward function r : x A + R, and the initial state distribution μ A(S).
Beyond its unprecedented scale, the collected UniTraj rep- resents diversity in several aspects:
Given an MDP, a trajectory of length T:
r= (su. ag. 1', $1, 1.. 2r--2 11 11)
Environment diversity. UniTraj encompasses a wide range of control environments. These include not only widely-used environments from the DeepMind Control Suite (DMC) (Tassa al., 2018) and OpenAI Gym (Brock- man, 2016), but also customized embodiments and tasks proposed in Modular RL and TD-MPC2. Notably, we pur- posely exclude all trajectories from the HalfCheetah, Hop- per, and Walker2D environment of OpenAI Gym, which are held out as our downstream test environments.
is recorded as interactions between the environment and an agent, according the following protocol: starting from an initial state Su ~ 14, at each discrete time step = 0, 1,..., the agent performs an action a; E A according to its pol- icy, receives an immediate reward r++1 = r(8,, az), and observes the next state after transition St+1 ~ P(st, az).
A world model po(st+1, rtt1 az), or more generally (8(+1+1, re+11/1:5 d1=)), learns its parameter 0 from dataset of recorded trajectories D = {Ti} to approximate the underlying transition probability and reward function, thus serving as an alternative of the environment.
Distribution diversity, The dataset contains data col- lected from various distributions, resulting from different collection methods and policies. Specifically, data from RL Unplugged, TD-MPC2, and Modular RL are gathered by recording the training agent's replay buffer, while JAT and DB-1 data are collected through expert policies roll- outs. Additionally, EXORL data are collected by storing the transitions from running unsupervised exploration al- gorithms (Laskin et al., 2021). The policies cover a range of approzches, including a wide range of reinforcement learning algorithms (e.g., D4PG (Barth-Maron et al., 2018), 8), PPO (Schulman al., 2017)) and state-of-the-art model predictive control algorithms, TD-MPC2.
Our work. While most literature learns a world model on target environment M from scratch, we investi- gate an under-explored paradigm of pre-training a world model from a family of heterogeneous' environments LM,MM, ......,,MM. MK.. Through learning from mixed tra- jectory data {D>, DD DD,...,,DAk, }, we obtain a good starting- point of the model Oo. ready for either zero-shot generaliza- tion to unseen M or fine-tuning to obtain a world model of M with strong generalization given limited data. We elab- orate on the intuition behind this paradigm in Section 1.
By scaling up the dataset while preserving diversity, we empower the model with the potential to generalize across varied environments.
3. UniTraj Dataset
We introduce UniTraj a large-scale unified trajectory dataset from heterogeneous environments, to support the pre- training of a trajectory world model. To ensure diversity, we merge five publicly available datasets with different charac-
4. Trajectory World Models
In this section, we first explain the intuition behind the pro- posed Trajectory World Models (TrajWorld) (Section 4.1), then provide a detailed overview of the architecture imple- mentation (Section 4.2), and conclude with a discussion of the pre-training and fine-tuning paradigm (Section 4.3).
'ee use the term "heterogeneous" to highlight that different environments not only feature varying transition and reward func- tions but also, more challengingly, possess distinct state and action spaces tied to unique sets of sensors and actuators.
3

---

Trajectory World Models for Heterogeneous Environments
Projection
Prediction
Feed- Forward
M
Feed- Porward
Temporal Attention
Variate Attention
Discretization
Temporal Attention
Variate Attention
Cross Entss-nn Loss
XTM
x.
xT
Figure Architecture of Trajectory World Models. A trajectory flattened into scalars, organized into two dimensions by timesteps and variates (each variate corresponds to single dimension in the state, action, and reward), and then discretized into categorical representations, Transformer with interleaved temporal and variate attentions processes the inputs to prodict the categorical distribution for the next timestep autoregressively. Layer normalizations and residual connections omitted for simplicity.
4.1. Intuition
transitions from the current state to the next. Instead of using simple one-dimensional attention over flattened se- quences, explicitly modeling these two dimensions has the potential to enhance transferability in downstream tasks, as it guides the model learn in tematic manner. This is supported by empirical results in Section 5.2. In short, we use two-way attention mechanism instead of one-dimensional attention on sequences.
To address the challenges of heterogeneity and promote knowledge transfer, we make three key observations:
Rediscovering homogeneity in scalars. While hetero- geneity often arises in differently sized vector information, there exists inherent homogeneity at the scalar level. Each variate-4 single scalar dimension in the state, ac- tion, or reward-rrprresets a fundamental quantity with its own physical meanings of the environment, e.g., position or torque of single joint, and can be consistently modeled, re- gardless of the shape of the whole vector information. This insight leads information whole, we break it into the scalar level for processing and prediction.
structured and sys-
4.2. Architecture
Building on the above intuitions, we realize Transformer- based architecture for TrajWorld (see Figure 3).
design choice: Instead of treating vector
Scalarization.
To exploit the inherent homogeneity at the scalar level, we flatten a trajectory (Equation (1)) from the spaces S T A R into two-dimensional representation organized by timesteps and variates:
Identifying environment through historical context. Unlike single-nnviromment scenarios with fixed state and action definitions, in our setting, variants can represent dif- ferent quantities across environments despite the same index in the vector. While environment IDs are typically included as inputs to distinguish environments, we instead leverage the in-context learning ability of Transformers (Brown et al., 2020): history transitions can provide the context needed for the model to infer relationships between variants. This makes pre-training even more critical. By exposing the model to diverse data across environments, we encourage it to learn "how to learn environment dynamics,"-4 more generalizable knowledge rather than solely focusing on specific enviromments. This ability is demonstrated in Sec- tion 5. where our pre-trained model has satisfactory zero- shot performance. summary, we provide historical con- text instead of environment identities, guiding the model to learn infer dynamics through context.
TO
-=
(m)
(1)
(n)
8T
PT-1
aT
ar-
(2)
where s denotes the i-th dimension of st. Padding is ap- plied ro and @T-1 zeros. This transformation converts heterogeneous trajectories of varying lengths and dimen- sions into matrices X E RTXM where M + 1, which can be flexibly processed by the attention mechanism.
Discretization and embeddings. Transformers excel in processing discrete inputs, so we further convert scalars into categorical representations. For each variate s(i) or a(i), we define B uniform bins with boundaries by < b1 where bn and by represent the minimum and maximum values of the variate in the training data. Scalars are then mapped to these bins using one-hot encoding or Gaussian histograms (Imani & White, 2018; Farebrother et al., 2024).
bg,
Inductive bias for two-dimensional representations. So far, our modeling for heterogeneous dynamics involves two dimensions: one focuses on capturing the relationships among variants, and the other models how actions drive
The resulting discrete representation Q E [0, [T is

---

Trajectory World Models for Heterogeneous Environments
linearly projected to match the Transformer's hidden size d. Additionally, apply three learned embeddings- timestep- embedding (TE), variate embedding (VE), and prediction embedding (PE)- to capture timestep indices, variate iden- tities, and whether a variate is a target for prediction. For- mally, for each i E [2'] and j E [M]]
Prediction Error (MSF)
Method
Pendulum
Walker2D
Last-step Mirroring
1.3 10-3
1.7 x 10
FrajWWold (w/o history)
1.7 x 10-6
2.1 x 10-3
Traj World (w/ history)
2.9 10-6
7.2 x 10-*
(a) Environment parameters transfer.
Z Win Qiy + TE(1) + VEG) PE((I & m 1J). (3)
Zero-Shot Prediction on Cart-2-Pole
1 (Pred)
2 (Pred)
Interleaved temporal-variate attentions. The input 7 E R2''xxxé is processed through a series of L trans- former blocks, adapted for the two-dimensional input struc- ture. In each block 1 1,..., L, we first apply temporal attention, processing each variate independently:
Traj 1 (Reel)
Traj 2 (Reel)
Position (m)
0.0
-0.2
U UkiT = CausalAtttention/ g
Vj E [M]
(4)
3
9
12
15
18
Time Step
followed by feedforward network (FFN): tr FFN(U*).
(b) Cross coossennvir transfer.
Afterwards, variale attention is applied at each timestep:
Figure 4. Zero-shot generalization. (a) Mean squared error of zero- shot transition predictions in modified Gym Pendulum (holdout gravity) and Walker2D (holdout friction etc.). (b) TrajWorld zero-shot predictions for two Cart-2-Pos 2-Pole trajectories, which share 10 context steps but diverge due to differing subsequent actions.
v 1,1:M = Attention(U(.1..)))
Vi [T].
(5)
Since there no causal dependencies between variates at the same timestep, causal mask is applied during variate attention. Finally, another FFN is applied: Z' = FF)(
5. Experiments
Through interleaved temporal and variate attentions, each entry in our model efficiently aggregates information from all variates across all previous timesteps. As previously discussed, this enables the model to infer environment dy- namics in-context for transition prediction.
In this section, we test the following hypotheses:
Large-scale trajectory pre-training can generalize ef- fectively and even enable zero-shot generalization, con- trary to the common belief (Section 5.1).
Traj outperforms alternative architectures for transition prediction when transferring dynamics knowledge to new environments (Section 5.2).
Prediction and objective. A linear prediction head, fol- lowed by a softmax operation, produces the prediction dis- (entmaxx(Wnnt22 E [0, 1 Our
tribution P- model is trained using next-step prediction objective to match the categorical representation of the inputs:
TrajWorld leverages the general dynamics knowledge acquired from pre-training to improve performance in downstream tasks (Section 5.3).
T-1m+1 B
C(P.Q) = - ΣV Σ Qu+1. log Pij.k-
5.1. Zero-shot Generalization
2=1
We first demonstrate that through in-context learning ability, TrriW exhibits favorable generalization across heteroge- neous environments, which differ not only in their transition dynamics but also in state and action spaces.
During inference, the next-step prediction can be obtained by sampling from or taking the expectation of the predicted categorical distribution over bin centers.
Environment parameter transfer. We pre-train a Traj- World model on data from Gym Pendulum environments with varying gravity values and evaluate its transition predic- tion error holdout gravity values. As shown in Table 4a, World achieves significantly lower prediction error in zero-shot settings compared to a naive bascline that simply mimics the last timestep. Moreover, the performance of Traj World deteriorates noticeably when historical informa- tion excluded, highlighting the critical role of contexts for
4.3. Towards General Trajectory World Model
We pre-train general Trajectory World Model on offline datasets from diverse environments. This same pre-train pre-trained model can then be applied to all downstream tasks for fine- tuning. Thanks to the Transformer's flexible architecture design and in-context learning capabilities, the pre-trained knowledge becomes more transferable, benefiting a wide range of heterogeneous and complex control environments.
5

---

Trajectory World Models for Heterogeneous Environments
nattchastat-rrnddaaw2
hafthnestah-meduuww22
balthhselah-oxpertiiv
10'
Predicion
10"
hoppennaandomww22
hoppen-medummrer ay
happer-mediume V22
hoppen-eedurme dumpeeetttt--
hopper-exp expertb222
Predictoon Entor
10"
10
10"
8xx10
10
10
10
BK110
vadiarrad-randoom--222
wwikar26-medlmm-22
walkke226-xparr-222
TrsfArrs (wis PTT) TrafWors (we PT)
10
10"
Predkzion
10"
10
4xx1
3x10
2x10
10
10
Figure 5. Mean absolute errors (MAE) of transition prediction for Traj World, with and without pre-training (PT), across different train-test dataset pairs. Each subplot corresponds to a distinct training dataset, with the test datasets shown on the x-axis (r=random, m-r=medium- replay, m=medium, ree=mmddium-xppn e=expert). Error bars represent the standard deviation across three random seeds.
the model to effectively infer environment parameters. The results are consistent in similar experiment conducted on Gym Walker2D, where friction, mass, etc., are varied.
adopted for dynamics modeling, and TDM (Schubert et al., 2023), which similar our model but flattens inputs and uses one-dimensional attention. Each baseline is evaluated both for training from scratch and fine-tuning pre-trained ones on the same UniTraa dataset as Traj World. To enable pre-training, we pad the state and action vectors with zeros to match the same dimensionality for MLP. Additionally, we compare with our model trained from scratch.
Cross-environment transfer. We further find that Tra- jWorld, when trained on the large-scale UniTraj dataset, also capable of zero-shot generalizing to unseen envi- ronments, Cart-2-Pole and Cart-3-Pole from DMC (Figure 4b and 13). Specifically, TrajWorld successfully infers the influence of the action value (pushing force) the state di- mension (cart position) and accurately predicts the outcomes for different action sequences performed subsequently.
Results. Figure/ 1 presents the aggregated mean absolute error of 75 train-test dataset pairs for various models. Tra- jWorld outperforms all baselines, highlighting the effec- tiveness of its pre-training strategy and architecture design. Notably, MLP Ensemble with pre-training performs worse than its con-pre-trainee counterpart, emphasizing the im- portance of careful model design for world modeling across heterogeneous environments. While TDM also benefits sig- nificantly from pre-training, it still lags behind Traj World. This is likely because TDM naively treats everything as a 1D sequence, neglecting the unique problem structures. In contrast, TrajWorld explicitly models variate relationships and temporal transitions, leveraging different facets of dy- namics knowledge from the pre-training. Moreover, TDM predicts variants sequentially, which may accumulate errors and lead to less accurate results, whereas TrajWorld predicts all variables jointly, mitigating compounding errors.
5.2. Transition Prediction
We then evaluate how different world models benefit from pre-training for transition prediction, particularly for out- of-distribution queries, when fine-tuned to more complex, standard environments.
Setup. We use datasets of three environments HalfCheetah, Hopper, and Walker2D-froe D4RL (Fu et al., 2020) as our testbed. Each environment in D4RL is provided with five datasets of different distributions from policies of varying performance levels. We train world models in each of the fifteen datasets and test prediction errors of states and rewards across all five datasets under the same environment, resulting in 75 train-test dataset pairs.
In Figure 5, we further show detailed prediction error re- sults for TrajWorld compared to its non-pre- con-pre-te terparts. In 12 out of 15 training datasets, fine-tuned Tra-
Baselines. We compare our approach against two base- lines: an ensemble of MLPs (Chua et al., 2018), widely
trained coun-
6

---

Trajectory World Models for Heterogeneous Environments
Mean Rank Correiaton
Regrosg21
0.8
08
0.5
0.3
03
-0.2
-04
0.0
FOE
Eisenble det
TOM PT)
TrafWorrd (we
Figure 6. Overall off-policy evaluation (OPE) results across datascts of environments, averaged across three random seeds.
Worldd achieves a lower average prediction error across test datasets, further validating the effectiveness of pre- training. Moreover, the transfer benefits are evident in both in an-distribution and that the model generalizes well even when trained and tested on transitions collected from different policies.
Results. Figure 6 shows that TrajWorld significantly im- proves OPE compared to its non-pre-trained variant and outperforms all baselines in both average normalized ab- solute error and rank correlation. TrajWorld slightly un- derperforms on Regret@ 1, likely due to bounded reward prediction (see discussion in Appendix D). Consistent with Section 5.2, MLP Ensemble with pre-training suffers from negative transfer, showing a notable drop in performance compared to the non-pre-traine model. Although TDM also benefits from pre-training, does not reach the same level of performance as Traj TrajWorld. We attribute this the same reason discussed in Section 5.2.
scenarios, indicating
5.3. Off-Policy Evaluation
Off-policy evaluation (OPE) estimates the value of tar- get policy using an offline transition dataset collected by a separate behavior policy. It is commonly used to select the most performant policy from a set of candidates when online evaluation is too costly be practical. This task pro- vides an ideal evaluation scenario for world models, value estimation can be acquired by rolling out the target policy within the learned world model. This is particularly advanta- geous for evaluating long-hoe horgzhorizon predictions, where direct environment interaction is infeasible and model accuracy over extended timeframes critical.
5.4. Model Predictive Control
Model predictive control (MPC) selects actions by optimiz- ing predicted future rewards over finite horizon using learned world model, making it well-suited for evaluating world model performance in online control settings.
Setup. We evaluate MPC performance in practical scenario where world models trained medium- replayy datasets are used to enhance medium-level proposal policies through model predictive control. Specifically, we utilize three medium-replay datasets from D4RL and medium-level policies from DOPE. Additionally, we experiment with MPC using a random shooting planner. Implementation details are provided in Appendix B.5.
Setup. We adopt the DOPE benchmark (Fu et al., 2021) over various D4RL environments. The tasks in this bench- mark are particularly challenging, as the target policies are of different levels and may differ significantly from the be- havior policy. To perform well these tasks, the world model must generalize well across all possible state-action distributions. Evaluation metrics include mean absolute error comparing estimated vs. ground-truth policy values, rank correlation between estimated and actual policy rank- ings, and Regret@/ measuring accuracy in selecting the best policy, as detailed Appendix B.4.3.
Baselines. As in the previous section, we compare our method against the MLP Ensemble and TDM baselines, evaluting both both from-scratch and fine-tuned variants.
Results. Figure 7 presents results for MPC with proposal policies. Overall, MPC using TrajWorld yields the highest- performing agents, outperforming both baseline models and its from-scratch counterpart. We find that MPC leads to significant gains in the Hopper and Walker2D environments, but has limited effects in HalfCheetah, likely due to its in- herent stability and lower risk of failure. contrast, Hopper and Walker2D are fragile, and our world models help pre- vent unsafe actions, leading better planning. Notably,
Baselines. In addition to the MLP Ensemble and TDM models mentioned earlier, we compare our approach against several other baselines. Notably, Energy- -based Transition Models (ETM) (Chen et al., 2024) currently sets the otate-of- the-art on this benchmark, outperforming prior methods by a significant margin. We also include the classical methods from the original DOPE paper (Fu et al., 2021) for a more comprehensive comparison.

---

Trajectory World Models for Heterogeneous Environments
Adgregated Result Envionments)
Enviroment:
Enviooment: Walker
Envfronment: HalCChetahh
Prapasa Pollcy
Fall
Nemultz Eprodd Ream
wia the Traning
Training
1180-
Eptrodo Riturn
1.00
Epikoda Ratumn
Epiedde Relurn
183+
MLP Ercemble
TrafWods
MLP Encemble
TraWWordd
MLP Ensemble
TraWord
Figure 7. Model predictive control (MPC) results with proposal policies across three environments, averaged over three random seeds.
the TDM model exhibits negative transfer in the MPC with proposal policies setting, despite showing positive transfer in transition prediction and off off-policy evaluation.
Effects of pre-training scale and diversity We assess the impact of dataset scale and diversity by pre-training three TrajWorld variants on subsets of UniTraj: a 1/10 sample, 1/100 sample, and the JAT subset (purely expert trajec- tories from five environments), followed by fine-tuning on downstream environments. For transition prediction, adopt challenging setup where models are fine-tuned on expert data per environment but evaluated across all data levels. Model-predictive control (MPC) results are also reported. As shown in Figure 9, all subset-pretrained mod- els outperform training from scratch, but fall short of the model trained on the full UniTraj dataset. This reveals a desirable scaling trend: larger and more diverse pre-training data consistently lead to better generalization. These find- ings highlight the value of heterogeneous pre-training on large-scale datasets.
In the random shooting setting, the planner's limited abil- ity to sample high-quality actions hinders its effective uti- lization of model predictions, leading consistently poor performance across all world models. Nevertheless, Traj- World demonstrates comparatively better results under these limitations. See Appendix C.4 for detailed results.
5.5. Analysis
Few-shot adaptation. TrajWorld presents pre-training benefits in few-shot scenarios. In Figure 8a, we show the prediction eтror across varying levels of data scarcity and compare TrajWorld with and without pre-training. These results highlight that the advantages of pre-training become increasingly pronounced data becomes more limited.
6. Related Work
Trajectory dataset. Data-driven approaches for control like imitation learning (Florence et al., 2022; Shafiullah et al., 2022; Gallouédec et al., 2024) and offline reinforce- ment learning (Fu et al., 2020; Rafailov et al., 2024; Gul- cehre et al., 2020; Qin et al., 2022) have promoted the public availability of trajectory datasets. However, these datasets are rarely utilized as unified big data for foundation models, likely due their isolated characteristics, such as differ- ences in policy levels, observation spaces, and action spaces. In fact, the largest robotics dataset, Open X- X-Embodiment (O' (eileill et al., 2024), is typically used for imitation learning with homogeneous visual observations and end-effector ac- tions (Team et al., 2024; Kim et al., 2024). Gato (Reed et al., 2022) collects a large-scale dataset across diverse environ- ments for generalist agent, but it not publicly available. In contrast, we curate public heterogeneous datasets, target- ing more capable trajectory world model.
Discretization visualization. We use t-SNE (van der Maaten & Hinton, 2008) to visualize the linear weights of our model's s prediction head for each category. The mapped weights exhibit strong continuity in Figure 8b. Since the output categories' indices aligned with the bins in in- creasing order, this indicates that our model has learned the ordering of bins shared by variants, despite being trained via an unordered classification objective. This suggests the model's potential for fine interpolation between existing bins and extrapolation to unseen ranges of variant values.
Variate attention visualization. We visualize the variate attention maps of our fine-tuned model in the Walker2D en- vironment, whose states are ordered with joint positions first, followed by their velocities. As shown Figure 8c, the at- tention map exhibits prominent diagonal patterns that focus on the corresponding joint's position and velocity, suggest- ing the model's 'ss understanding of each variate's semantics. Additionally, the strong attention between neighboring vari- atc, such as physically linked joints, further confirms the model's 's grasp of joint relationships. We also observe no- table attention patterns between states and actions, and these additional results are available in Appendix C.6.
Cross-enviromment architecture. Zero- Zero-padding to fit maximum length (Yu et al., 2020a; Hansen et al., 2024; Schmied et al., 2024; Seo et al., 2022) or using separate neu- ral network heads (Wang et al., 2024a: D' E'Ermmo et al., 2020) hinders knowledge transfer between heterogeneous environ-
8

---

Trajectory World Models for Heterogeneous Environments
Error
19NE Prediction Head Weigtt
0.13
Postions
Watoond
22
8.155
3:00-
1.1
30
Halic Dasa Prondded
(a) Fho-shot adaptation.
(b) Discretization visualization.
(c) Variate attention visualization.
Figure 8. Model analysis. (a) Downstream prediction error of Traj World under varying data scarcity levels. (b) t-SNE visualization of the linear weights the model's prediction head. (c) Variate attention map from the third layer of Traj World fine-tuned on Walker2D.
Apgregated Results (15 train-test sataset pains) JAT Oity
Envirommnn
Hoppor
Nomme Prediction
130
1900
-4
9.73
None 1%
10%
100% LnnTTey
Nome
JAT Only
15
10% 105% kiiTTN
None
JATT Onty
1%
10%
107% linit
(a) Transition prediction error.
(b) Model predictive control.
Figure 9. Effects of pre-training scale and diversity. (a) Aggregated transition prediction error. (b) Model predictive control performance on Walker2D and Hopper. All results are obtained from models fine-tuned from pre-trained TrajWorld on different subsets of UniTraj.
ments with mismatched or differently sized state and action spaces. Previous work has resorted to flexible architectures like graph neural networks (Huang et al., 2020; Kurin et al., 2021) and Transformers (Gupta et al., 2022; Hong et al., 2021) for policy learning. Our method leverages similar architecture for world modeling (Janner et al., 2021; Zhang et al., 2021), but introduces the wwo-dimensional attention design for the first time in this context. More importantly, no for computational efficiency, as in prior work of other fields (Ho et al., 2019; Arnab et al., 2021; Nayakanti et al., 2023), we show its benefits for cross- cross-environment transfer.
7. Conclusion
We address the challenge of building large-scale pre-trained world models for heterogeneous environments with distinct sensors, actuators, and dynamics. Our contributions include UniTraj, dataset of over one million trajectories from 80 environments, and TrajWorld, a flexible architecture for cross coos-environmeet transfer. Pre-training TrajWork on Uni- Traj achieves superior results in transition prediction and off- policy evaluation, demonstrating the first successful transfer of world models across complex control environments.
Limitations and future work. While this work takes a successful first step. there is significant room for further study. Despite the strong practical performance, tion of our architecture is that the discretization scheme con- strains predictions to fixed range, making it theoretically difficult to model extremely out-of-distribution transitions beyond these bounds. Additionally, our model, designed for scalable pre- pre-training, has a larger capacity compared to classic MLPs, which poses challenges in model calibration (Guo et al., 2017), particularly in scenarios where uncer- tainty quantification is critical, such offline RL (Yu et al., 2020b). This increased complexity also comes with addi- tional computational costs. For future work, we envision that pre-training multimodal world models incorporating both visual and proprioceptive observations could lead to models with a deeper understanding of the physical world.
World model pre-training. The homogeneity videos across diverse tasks, environments, and even embodiments has driven rapid advancements in large-scale video pre- training for world models (Seo et al., 2022; Wu et al., 2024a;b; Ye et al., 2024; Cheang et al.. 2024). However, heterogeneity across different sets of sensors and actuators poses significant challenges to developing general world models based on low-dimensional sensor information.
limita-
Our work is particularly relevant to Schubert et al. (2023), which trains generalist transformer dynamics model from 80 heterogeneous environments. Still, they only observe positive transfer when adapting to simple cart-pole envi- ronment and fail for a more complex walker environment. In contrast, our work, for the first time, validates the positive transfer benefits across such more complex environments.

---

Trajectory World Models for Heterogeneous Environments
Acknowledgements
Cheang, C.-L., Chen, G., Jing, Y., Kong, T., Li, H., Li, Y., Liu, Y., Wu, H., Xu, J., Yang, Y., et al. Gr-2: A generative video-langunge-actioo model with web-scale knowledge for robot manipulation. arXiv preprint arXiv:2410.06155 2024.
This work was supported by the National Natural Science Foundation of China (U2342217 and 62021002), the BNRist Innovation Fund (BNR2024RC01010) and the National Engineering Research Center for Big Data Software.
Chen, R., Jía, C.,, Huang, Z., Liu, T.-S., Liu, X.-H., and Yu, Y. Offline transition modeling via contrastive energy learning. In ICML, 2024.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Chua, K., Calandra, R., McAllister, R., and Levine, S. Deep reinforcement learning in a handful of trials using proba- bilistic dynamics models. In NeurIPS, 2018
D D'Eramo, C., Tateo, D., Bonarini, A., Restelli, M., and Peters, J. Sharing knowledge in multi-task deep rein- forcement learning. In ICLR, 2020.
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L.,, Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.0877 2023.
Ebert, F., Finn, C., Dasari, S., Xie, A., Lee, A., and Levine, S. Visual foresight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint urXiv:/8122 18/2.0058, 2018
Agarwal, N., Ali, A., Bala, M.. Balaji, Y., Barker, E., Cai, T., Chattopadhyay, P., Chen, Y., Cui, Y., Ding, Y., et al. Cosmos world foundation model platform for physical ai. arXiv preprint arXiv:2501.03572 2025.
Farebrother, J., Orbay, J., Vuong, Q., Taïga, A. A., Chebotar, Y., Xiao, T., Irpan, A., Levine, S., Castro, P. S., Faust, A., et al. Stop regressing: Training value functions via classification for scalable deep rl. In ICML, 2024.
Arnab, A., Dehghami, M., Heigold, G., Sun, C., Lučić, M., and Schmid, C. Vivit: A video vision transformer. In ICCV, 2021.
Florence, P., Lynch, C., Zeng, A., Ramirez, O. A., Wahid, A., Downs, L., Wong, A., Lee, J., Mordatch, I., and Tompson, J. Implicit behavioral eloning. In CoRL, 2022.
Barth-Maron, G., Hoffman, M. W., Budden, D., Dabney, W., Horgan, D., Dhruva, T., Muldal, A., Heess, N., and Lil- licrap, T. Distributed distributional deterministic policy gradients. In ICLR, 2018.
Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, reinforcement
S. D4rl: Datasets for deep learning. arXiv preprint arXiv:2004.07219, 2020.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http://g thub. com/ jax-ml/jax
Fu, J.. Norouzi, M., Nachum, O., Tucker, G., Wang, Z. Novikov, A., Yang, M., Zhang, M. R., Chen, Y., Kumar, A., et al. Benchmarks for deep off- fff-policy evaluation. In ICLR, 2021.
Fujimoto, S., Hoof, H., and Meger, D. Addressing function approximation error in actor-critic methods. In ICML, 2018.
Brockman, G. Openai gym. 1:X1v:6666001444 2016.
arXiv preprint
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neclakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. In NewrlPS, 2020.
Gallouédec, Q.. Beeching, E.,, Romac, C., and Dellandréa, E. Jack of all trades, master of some, muuli-pprpor trans- former agent. arXiv preprint arXiv:2402.00984. 2024.
Gu, Y., Zheng, B., Gou, B., Zhang, K., Chang, C., Srivas- tava, S., Xie, Y., Qi, P., Sun, H., and Su, Y. Is your I1m secretly world model of the internet? model- plan- ning for web agents. arXiv preprint arXiv:241 1.06559, 2024.
Bruce, J., Dennis, M. D., Edwards, A., Parker-Holder, J., Shi, Y., Hughes, E., Lai, M., Mavalankar, A., Steiger- wald, R., Apps, C., et al. Genie: Generative interactive environments. In JCML, 2024.
Chae, H., Kim, N., Ong, K. T.-i., Gwak, M., Song, G., Kim, J., Kim, S., Lee, D., and Yeo, J. Web agents with world models: Learning and leveraging environment dynamics in web navigation, In ICLR, 2025.
Gulcehre, C., Wang, Z., Novikov, A., Paine, T., Gómez, S., Zolna, K., Agarwal, R., Merel, J. S., Mankowitz, D. J., Paduraru, C., et al. RI umplugged: A suite of benchmarks for offline reinforcement learning. NeurlPS, 2020.
10

---

Trajectory World Models for Heterogeneous Environments
Guo, C., Pleiss, G., Sun, Y., and Weinberger, K. Q. On calibration of modern neural networks. In ICMI., 2017.
Kurutach, T., Clavera, Duan, Y., Tamar, A., and Abbeel, P. Model-ensemble trust-region policy optimization. In ICLR, 2018.
Gupta, A., Fan, L.., Ganguli, S., and Fei-Fei, Metamorph: Learning universal controllers with transformers. In ICLR, 2022.
Laskin, M., Yarats, D., Liu, H., Lee, K., Zhan, A., Lu, K., Cang, C.,, Pinto, L., and Abbeel, P. Urlb: Unsupervised reinforcement learning benchmark. In NeurlPS, 2021.
Ha, and Schmidhuber, J. Recurrent world models facili- tate policy evolution. In NeurIPS, 2018.
Le, H., Voloshin, C., and Yue, Y. Batch policy learning under constraints. In ICMI,, 2019.
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination. In ICLR, 2020.
LeCun, Y. A path towards autonomous machine intelligence version 0.9, 2, 2022-06-27. Open Review, 62(1):1-62, 2022.
Hansen, N., Su, H., and Wang, X. Td-mpc2: Scalable, robust world models for continuous control. In ICLR, 2024.
Mitchell, E., Rafailov, R., Peng, X. B., Levine, S., and Finn, C. Offline meta-reinforcement learning with advantage weighting, In ICML, 2021.
Ho, J., Kalchbrenner, N., Weissenborn, D., and Salimans, T. Axial attention in multidimensional transformers. arXiv preprint 1rXiv:191..1221 12100, 2019.
Nayakanti, N., Al-Rfou, R., Zhou, A., Goel, K., Refaat, K. S., and Sapp, B. Wayformer: Motion forecasting via simple & efficient attention networks. In ICRA, 2023.
Hong, S., Yoon, D., and Kim, K.-E. Structure-aware trans- former policy for inhomogeneous multi-task reinforce- ment learning. In ICLR, 2021.
Oquab, M., Darcet, T., Moutakanni, T. Vo, H., Szafraniec, M., Khalidov, V., Fernandez, P., Haziza, D., Massa, F,, El- Nouby, A., et al. Dinov2: Learning robust visual features without supervision. TMLR, 2024.
Huang. W., Mordatch, L., and Pathak, D. One policy to con- trol them all: Shared modular policies for agent-agnostic control. In ICML, 2020.
Imani, E. and White, M. Improving regression performance with distributional losses. In ICML, 2018.
O'Neill, A., Rehman, A., Maddukuri, A., Gupta, A., Padalkar, A., Lee, A., Pooley, A., Gupta, A., Mandlekar, A., Jain, A., et Open x-embodiment: Robotic learning datasets and rt-x models: Open x-embodiment collabora- tion O. In ICRA, 2024.
Janner, M., Fu, J., Zhang, M., and Levine, S. When to trust your model: Model-based policy optimization. NeurlPS, 2019.
Petrenko, A., Huang, Z., Kumar, T., Sukhatme, G., and Koltun, V. Sample factory: Egocentric 3d control from pixels at 100000 fps with asynchronous reinforcement learning. In ICML, 2020.
Janner, M., Li, Q., and Levine, S. Offline reinforcement learning as one big sequence modeling problem. In NeurlPS, 2021.
Jiang, N. and Li, L.. Doubly robust off-policy value evalua- tion for reinforcement learning. ICML, 2016.
Qin, R.-J., Zhang, X., Gao, S., Chen, X.-.., Li, 7.,, Zhang, W., and Yu, Y. Neorl: A near real-world benchmark for offline reinforcement learning. In NeurlPS, 2022.
Kim, M. J., Pertsch, K., Karamcheti, S., Xiao, T., Balakr- ishna, A., Nair, S., Rafailov, R., Foster, E., Lam, G., San- keti, P., et Openvla: An open-source vision-language- action model. In CoRI., 2024.
Rafailov, R., Hatch, K., Singh, A., Smith, L., Kumar, A., Kostrikov, I., Hansen-Estruch, P., Kolev, V., Ball, P., Wu, J., et al. D5rl: Diverse datasets for data-driven deep reinforcement learning. In RLC. 2024.
Kirillov, A., Mintun, E., Ravi, N., Mao, H., Rolland, C., Gustafson, L., Xiao, T., Whitchead, S., Berg, A. C., Lo, W.-Y. al. Segment anything. ICCV, 2023.
Reed, S., Zolna, K., Parisotto, E., Colmenarejo, S. G., Novikov, A., Barth-Maron, G., Gimenez, M., Sulsky, Y., Kay, J., Springenberg, J. T., et al. A generalist agent. TMLR, 2022.
Kostrikov, and Nachum, O. Statistical bootstrapping for uncertainty estimation off-policy evaluation. arXiv preprint arXiv:2007.13600 2020.
Kurin, V., Igl, M., Rocktäschel, T., Boehmer, W., and White- son, S. My body is cage: the role of morphology in graph-based incompatible control. In ICLR, 2021.
Schmied, T., Hofmarcher, M., Paischer, F., Pascanu, R., and Hochreiter, S. Learning to modulate pre-trained models in rl. NeurIPS, 2024.

---

Trajectory World Models for Heterogeneous Environments
Schrittieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L.,, Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839): 604-609, 2020..
Wen, J., Dai, B., Li, L.., and Schuurmans, D. Batch station- ary distribution estimation. In ICMI., 2020.
Wen, Y., Wan, Z., Zhou, M., Hou, S., Cao. Z., Le. C., Chen, J., Tian, Z., Zhang, W., and Wang. J. On realiza- tion of intelligent decision-making in the real world: A foundation decision model perspective. arXiv preprint arXiv:22122 2. 12669, 2022.
Schubert, I., Zhang, J., Bruce, J., Bechtle, S., Parisotto, E., Riedmiller, M., Springenberg, J. T., Byravan, A., Hasen- clever, L., and Heess, N. A generalist dynamics model for control. arXiv preprint arXiv:2305. 2023.
Wu, J., Ma, H., Deng, C., and Long, M. Pre-training con- textualized world models with in-the-wild videos for re- inforcement learning. In NeurIPS, 2024a.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347, 2017.
Wu, J., Yin, S., Feng, N., He, X., Li, D., Hao, J., and Long, M. ivideogpt: Interactive videogpts are scalable world models. In NeurIPS, 2024b.
Seo, Y., Lee, K., James, S. L., and Abbeel, P. Reinforcement learning with action-free pre-training from videos. In ICML, 2022.
Ju, Yin, S., Feng, N., and Long, M. Rlvr-world: Training world models with reinforcement learning. arXiv preprint arXiv:2505. 2025.
Shafiullah, N. M., Cui, Z., Altanzaya, A. A., and Pinto, L.. Behavior transformers: Cloning modes with one stone. In NewrlPS, 2022.
Yang. M., Nachum, O., Dai, B., L., and Schuurmans, D. Off off-po evaluation via the regularized lagrangian. In NeurlPS, 2020.
Song. H. F., Abdolmaleki, A., Springenberg. J. T., Clark, A., Soyer, H., Rae, J. W., Noury, S., Ahuja, A., Liu, S., Tirumala, D., Heess, N., Belov, D., Riedmiller, M., and Botvinick, M. M. V-mpo: On-policy maximum posteriori policy optimization for discrete and continuous control. In ICLR, 2020.
Yarats, D.,. Brandfonbrener, D., Liu, H.,, Laskin, M., Abbeel, P., Lazaric, A., and Pinto, L. Don't 't change the algorithm, change the data: Exploratory data for offline reinforce- ment learning. arXiv preprint arXiv:2201 20t::2201.13222 13425, 2022.
Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., de Las Casas, D., Budden, D., Abdolmalcki, A., Merel, J., Lefrancq, A., Lillicrap, T., and Riedmiller, M. Deepmind control suite, arXiv preprint arXiv: 1801.00690, 2018.
Ye, S., Jang, J., Jeon, B., Joo, S., Yang, J., Peng, B., Mandlekar, A., Tan, R., Chao, Y.W.,, Lin, B. Y,, et al. Latent action pretraining from videos. arXiv preprint arXiv:2410.11755 2024.
Team, O. E. L., Stooke, A., Mahajan, A., Barros, C., Deck, C., Bauer, J., Sygnowski, J., Trebacz, M., Jaderberg, M., Mathieu, M., et al. Open-ended learning leads to gener- ally capable agents. arXiv preprint 11Xiv:2107.12800 2021.
Yu, T., Quillen, D., He, Z., Julian, R., Hausman, K., Finn, C., and Levine, S. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In CoRI., 2020a.
Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J. Y., Levine, S., Finn, C., and Ma, T. Mopo: Model-based offline policy optimization. In NeurIPS, 2020b.
Team, O. M., Ghosh, D., Walke, H., Pertsch, K., Black, K., Mees, O., Dasari, S., Hejna, J., Kreiman, T., Xu, C., et al. Octo: open-source generalist robot policy. In RSS, 2024.
Zhang, M. R., Paine, T. L.. Nachum, O., Paduraru, C., Tucker, G., Wang, 7., and Norouzi, M. Autoregressive dynamics models for offline policy evaluation and opti- mization. In ICLR, 2021.
Tian, S., Finn, C., and Wu, J. A control-centric benchmark for video prediction. ICLR, 2023.
van der Maaten, L. and Hinton, G. Visualizing data using t-sne. JMLR, 2008
Zhou, G., Pan, H., LeCun, Y., and Pinto, Dino-wm: World models on pre-trained visual features enable zero- shot planning. arXiv preprint 2xXi:2411..4998 2024.
Wang, I.,, Chen, X., Zhao, J., and He, K. Scaling proprioceptive-viuur learning with heterogeneous pre- trained transformers. In NeurlPS, 2024a.
Wang, R., Todd, G., Xiao, Z., Yuan, X., Côté, M.-A., Clark, P., and Jansen, P. Can language models serve as text- based world simulators? ACI, 2024b.
12

---

Trajectory World Models for Heterogeneous Environments
A. UniTraj Dataset Details
A.I. Overview of UniTraj Components
In this part, we provide a brief overview of each component of the UniTraj dataset.
ExORL (Yarats et al., 2022). Exploratory Data for Offline RL (ExORL) follows a two-step data collection protocol. First, data is generated in reward-free environments using unsupervised exploration strategies (Laskin al., 2021). Next, this data is relabeled with either standard or hand-designed reward function specific to each environment's 's task. This procedure leads to data with broader state-action space coverage, which benefits generalizatioon-emmaniii scenarios like offline RL.
RL Unplugged (Gulcehre et al., 2020). We incorporate RL Unplugged's dataset from the DeepMind Control Suite domains. Most of the data collected in this domain are generated by recording D4PG's training runs (Barth-Maron et al., 2018), while Manipulator insert ball and Manipulator insert peg's 's data is collected using V-MPO (Song et al., 2020).
JAT (Gallouédec et al., 2024). We utilize Jack of All Trades (JAT)'s released dataset, which is collected using expert RL agent's rollouts. These agents are trained using asynchronous PPO (Schulman al., 2017), following the Sample Factory implementation (Petrenko et al., 2020). Specifically, we only use the subset of the dataset that was collected in the OpenAl Gym environments, excluding data collected in Walker2D, HalfCheetah, and Hopper.
DB-1 (Wen et al., 2022). The dataset for Digital Brain-1 (DB-1), a reproduction of Gato (Reed et al., 2022), also consists solely of expert policy rollouts. Although the released dataset contains only five expert episodes per domain, it spans multiple environments, including various DeepMind Control Suite environments and custom ones from Modular RL.
TD-MPC2 (Hansen al., 2024). TD-MPC2 is state-of-the-an model-based RL algorithm. We include released data from single-task TD-MPC2 agents' replay buffers, collected from DeepMind Control Suite environments.
Modular RL (Huang et al., 2020). The Modular RL environments introduced by Huang et al. (2020) feature customizable embodiments with varying limb and joint configurations. We collected the data on these environments by ourselves. Specifically, we used the provided XMI. files to define different embodiment structures and followed the original reward function designs. We ran the TD3 algorithm (Fujimoto et al., 2018) and stored all episodes until the policy began to converge. The hyperparameters for TD3 are kept consistent with the default settings provided in the official repository repository?.
A.2. List of Environments
The curated UniTraj dataset spans a diverse range of environments from multiple sources, including DeepMind Control Suite, OpenAl Gym, and various customized environments. In Table 2, we provide a detailed list of environments used in each component of UniTraj.
A.3. Sampling Weights
We manually weighted different subsets, trying to balance size and diversity. The sample weights are shown Table
B. Experimental Details
B.1. Model Implementation
TrajWorld. For discretization, described in Section 4.2, we can employ two methods: one-hot encoding and Gaussian histograms. Specifically, the Gaussian histogram method is utilized for input discretization, while the one-hot encoding is applied for target discretization. Compared to one-hot encoding, Gaussian histograms provide a more fine- meaigrai representation of value information. While we can also use Gaussian histograms for target discretization, one-hot encoding more suitable for uncertainty quantization future applications such as offline RL. This is because two Gaussian distributions with the same standard derivation can yield different entropy when discretized into histograms.
Antpps: :tttps:://iiii thub. com/ s su fm//DDD
13

---

Trajectory World Models for Heterogeneous Environments
Component
Environments
EXORL
Cartpole, Jaco, Quadruped, Walker!
RL Unplugged
Cartpole, Fish, Humanoid, Manipulator, Walker!
JAT
Double Pendulum", Pendulum", Pusher", Reacher", Swimmer"
Acrobat, Ball In Cup, Cartpole, Cheetah-2-back, Cheetah-2-front, Cheetah-3-back,
Cheetah-3-front, Cheetah-4-allback Cheetah-4-allfront Cheetah-4-back, Cheetah-4-front,
Cheetah-5-back, Cheetah-S-balanncel Cheetah-S-front, Cheetah-6-back,
Cheetah-6-front, Finger, Fish, Hoppert, Hopper-3, Hopper55, Humanoid,
Elumanoid-2d---icffl-ee-
DB-1
Humanoid -a-anoid-2d-7-raee 2d-7-right-amm Humanoid-244 2---riihid,,----ggt--lle..g Humanoid 22--8leftt-neee lefe-knee,
Humanoid-2d-8--ige 2---rightt-e,, Humanoid-244-9-flll,
Manipulator, Reacher, Swimmer6, Swimmer15, Walker",
Walker-2-llippee, Walker-2-main, Walker-3-fippeed, Walker-3-main,
Walker-4-fipped, Walker-4-main, Walker-5-flipped, Walker-5-main, Walker-6-flipped Walker-6-main
Acrobot, Ball In Cup, Cartpole Cheetah Finger, Fish, Hopper',
TD-MPC2
Pendulum Reacher Walker
Cheetah-2-back, Cheetah-2-front, Cheetah-3-back
Cheetah-4-allback, Cheetah-4-back, Cheetah-4-front Cheetah-5-back
Modular RL
Micetah-5-balancee Cheetah-5-front Cheetah-6-back. 0heetah-6 6-front, Hopper-3, Hopper-5, Walker-2-flipped, Walker-3-fipped, Falker-4-flippede Walker -lllpe---flipped, falker---flipped Walker-7-flppee
Table 2. A detailed list of environments used in the UniTraj dataset. For environments sharing the same name, we mark those from OpenAl Gym with an asterisk (+) and those from DeepMind Control Suite with a dagger (t). Notably, the Gym Hopper, Walker2D, and HalfCheetah environments used for evaluating our methods and baselines differ from their DeepMind Control Suite counterparts, exhibiting variations in state/action definitions and environment parameters.
Subsets
ExoRL
RLU
JAT
DB-1
TD-MPC2
Modular RL
(Unnormalized) sampling weight
75
5
90
90
30
Table 3. Sampling weights subsets for pre-training with UniTraj dataset.
For prediction, each bin [b b,] is represented by its center C, (b;-1 + b:)/2. Given the predicted bin probability Pt-
the output value distribution can be expressed as P(X = x) = E.=p,, == q) or P(X = x) = Em (: < < b,)/(b, 11- We use the former for simplicity.
When pre-training with data from heterogeneous environments, for practical reasons, each batch is made up of data from a single environment.
We provide the hyperparameters used in pre-training and fine-tuning in Table 4. On transition prediction and OPE experiments, the environment-specific models trained from scratch use the same set of hyperparameters as fine-tuning.
Baseline: Transformer Dynamics Model (TDM). TDM (Schubert et al., 2023) does not provide an official implementa- tion. To enable a fair comparison, we adapt our TrajWorld implementation to reproduce TDM while maintaining consistency in discretization and embedding methods. Furthermore, when trained using cross-entropy loss, we mask actions and require the model to only predict the next states and During inference, the model predicts each scalar dimension of the state sequentially, followed by setting each scalar of the action (c.g., provided by the policy in off-policy evaluation) one at a time. The hyperparameters for pre-training and fine-tuning are kept consistent with those used in TrajWorld (Table 4), except for the batch size for pre-training Due to GPU memory constraints, the batch size for pre- pre-trainne training, originally set to 64, is reduced to 16. Like Traj World, we use the same hyperparameters as fine-tuning tne-tuning for environment specific models trained from scratch.
the TDM paper, where all variates are predicted.
14

---

Trajectory World Models for Heterogeneous Environments
Hyperparameter
Value
Input discretization
Gasss-hist
Target discretization
One-hot
Transformer blocks number
6
Attention heads number
4
Architecture
Transformer context length
20
Hidden dimension
256
MLP hidden
[1024,256]
MLP activation
GeLU
Total gradient steps
1M
Batch size
64
Learning rate
x 10
Dropout rate
0.05
Pre-training
Optimizer
Adam
Weight decay
10-5
Gradient clip norm
0.25
Scheduler
Warmup cosine decay
Scheduler warmup steps
10000
Total max gradient steps
1.5M
Max epochs
300
Steps per epoch
5000
Batch size
64
Learning rate
x 10-5
Fine-tuning
Dropout rate
0.05
Optimizer
Adam
Weight decay
10 10-5
Gradient clip norm Scheduler
0.25
Warmup cosine decay
Scheduler warmup steps
10000
Table Hyperparameters for TrajWorld.
Baseline: MLP Ensemble. Following prior work (Chua et al., 2018; Janner al., 2019; Yu et al., 2020b), we train an ensemble of transition models, parameterized diagonal Gaussian distribution of the next state and reward, implemented using MLPs. These models are trained with bootstrapped training samples, and optimized via negative log-likelihood. After training, we select an clite subset of models based on validation loss, and during inference, a model from this subset is randomly sampled for predictions. For pre-training on heterogeneous environments, we implement the MLP Ensemble baseline by padding each state vector to dimensions and each action vector to dimensions, resulting in 120-dimensional input the MI.. The model outputs the distribution over 91 -l-dimensional vector (90 for the next state and I for the reward). To ensure a fair comparison with other methods, match the parameter count of the ensemble to TralWorld, and no environment identities are provided to this baseline. The hyperparameters are listed in Table 5, Enviromment-specifee models trained from scratch use the same hyperparameters as in fine-tuning.
B.2. Zero-Shot Generalization
B.2.1. ENVIRONMENT PARAMETER TRANSFER
Pendulum. We pre-train the TrajWorld model on 60 Gym Pendulum environments, where the gravity values range from 8 m/s² to 12 m/s?. The pre-training dataset is collected by running the TD3 algorithm (Fujimoto et al., 2018) and storing all episodes until the policy converges. For evaluation, we use five holdout environments with gravity values between 5.m//2 and 7.5m/s?, collecting data in the same manner as the training datasets. The zero-shot results are reported as the average prediction error on these holdout datasets.
15

---

Trajectory World Models for Heterogeneous Environments
Hyperparameter
Value
MLP hidden
[640, 640, 640, 640]
Architecture
Ensemble number
Ensemble elite Number
Total gradient steps Batch size
1M
256
Pre-training
Learning rate
x 10
Optimizer
Adam
Total max gradient steps
1.5M
Max epochs
300
Steps per epoch Batch size
5000
Fine-tuning
256
Learning rate
x 10
Optimizer
Adam
Table 5. Hyperparameters for MLP Ensemble.
Walker2D. We pre-train four-layer Traj World model using 45 training datasets provided by MACAW (Mitchell et al., 2021) and evaluate it on separate dataset also from MACAW. The datasets in MACAW are collected under varying physical conditions, including differences in body mass, friction, damping, and inertia.
B.2.2. CROSS-ENVIROMAS ENVIROMENT TRANSFER
We evaluate the model pre-trained on UniTraj by performing a ten-step rollout in the Cart-2-Pole and Cart-3-Pole environment from DeepMind Control Suite. The rollout conditioned on history of ten prior timesteps. After this initial context, actions are applied simple predefined manner: either continuously pushing to the right (a 0.5) or the left (a -0.5). The action repeat for the Cart-2-Pole and Cart-3-Pole environment is set to 4.
B.3. Transition Prediction
The model is trained on a dataset using this dataset's 's training set and tested on five test datasets that come from the same environment. The evaluation for each test set is based the model's prediction across the entire test dataset. We the Mean Absolute Error (MAE) as the evaluation metric. The prediction of TrajWorld World is done by maintaining a history context window of 19 to predict the 20th state and reward.
In Figure 1, the prediction errors for each train-test dataset pair are normalized by dividing them by the MAE of the TrajWorld model without pre-training. The final result is then obtained by averaging across all environments.
B.4. Off-Policy Evaluation
B.4.1. IMPLEMENTATION: MODEL-BASED OPE
Given world model, the most direct method for off-policy evaluation (OPE) is Monte Carlo policy evaluation. This involves starting from a set of initial states, performing policy rollouts within the learned model, and averaging the accumulated rewards to estimate the policy value. The procedure is summarized Algorithm 1.
In practice,
discount factor of 7 0.995 and horizon length of h 2000. The number of samples N is set such that each trajectory's initial state from the behavior dataset is used exactly once, resulting in approximately N z 1000. We use KV cache to accelerate the rollouts of our TrrjW.rld.
B.4.2. BASELINES
We primarily compare against model-based OPE with Energy-based Transition Models (ETM) (Chen et al., 2024), a strong baseline that significantly outperforms previous methods and represents state-of-the-an on the DOPE benchmark (Fu al., 2021).
16

---

Trajectory World Models for Heterogeneous Environments
Algorithm 1 Model-Based OPE
Input: learned world model Po (a(s Te41 at), test policy я, samples number N, initial state distribution So. discount
factor 7, horizon length h.
for i to N do
R; +-
Sample imitial state so So
for 0 do
0r (-16((
St+1, r141 ()))))
R₂ +- R2 + 7/ro+1
end for
end for
Return ((T) N Ei= R
We also include five classic OPE methods as baselines from the DOPE benchmark: Fitted Q-Evaluation (FQE) (Le et al., 2019), Doubly Robust (DR) (Jiang & Li, 2016), Importance Sampling (IS), (Kostrikov & Nachum, 2020) DICE (Yang ct al., 2020), and Variational Power Method (VPM) (Wen et al., 2020).
B.4.3. METRICS
We adopt the evaluation metries used in the DOPE benchmark.
Mean Absolute Error The absolute error quantifies the deviation between the true value and the estimated value of policy, defined as:
AbsErr = IV - - VT
where V represents the true value of the policy, and / denotes its estimated value. The Mean Absolute Error (MAE) is computed as the average absolute error across all evaluated policies. To aggregate results, these values are normalized by the difference between the maximum and minimum true policy values.
Rank correlation. Rank correlation, also known Spearman's rank correlation coefficient (p), measures the ordinal correlation between the estimated policy values and their true values. It is given by:
RankCorr Cov(Vi.N ViAN) @VV (V IN)
(8)
where 1 : N represents the indices of the evaluated policies.
Regret@k.. Regret@k quantifies the performance gap between the actual best policy and the best policy selected from the top-k candidates (ranked by estimated values). formally defined as:
Regret @ max V
max V
(9)
:EI:N
(EIo))(1:NN
where topk(1 N) denotes the indices of the top l policies based on estimated values V In our experiments, we specifically use normalized Regret@) @ as the evaluation metric.
B.5. Model Predictive Control
We evaluate model predictive control (MPC) performance under two planning settings: policy proposal and random shooting.
Policy proposal setting. Action candidates are generated by perturbing the output of a learned action policy with Gaussian noise. Specifically, we first query the policy to obtain mean action sequence and then add zero-mean Gaussian noise to each action in the sequence. This results of diverse trajectories centered around the policy's 's behavior.
17

---

Trajectory World Models for Heterogeneous Environments
Random shooting setting. Candidate action sequences are sampled directly from Gaussian distribution without guidance from learned policy. Each trajectory is independently sampled by drawing actions from zero-mean Gaussian distribution with fixed standard deviation.
Hyperparameters. For both settings, sample size of 128 candidate trajectories per MPC rollout, across all environments. The best-performing action sequence is selected based predicted cumulative reward computed using the world model.
The planning horizon is set based on the characteristics each environment. Specifically, we use a horizon of 25 steps for both HalfCheetah and Walker2D, while a longer horizon of 50 steps is adopted for Hopper. This extended horizon for Hopper helps mitigate short-sighted planning behavior, which particularly detrimental this more fragile environment.
To ensure optimal performance across different world models, the standard deviation of the Gaussian noise used for action sampling is tuned individually for each environment. The noise level is set to 0.05 for Hopper, 0.2 for Walker2D, and 0.025 for HalfCheetah. These values were empirically selected to balance exploration and stability during trajectory sampling.
These settings are used consistently in all experiments involving MPC in this work. The same configurations are applied for evaluating all world models, ensuring fair comparison.
B.6. Computational Cost
Our implementation, built upon JAX (Bradbury et al., 2018), benefits from significant computational efficiency. Both pre-training and fine-tuning of the Traj World model can be conducted on single 24GB NVIDIA RTX 4090 GPU. For comparison, the computational cost for 1.5M training steps in our implementations of the MI.P Ensemble, TDM, and TrajWorld is 1.5, 36, and 28 hours, respectively. This highlight that TrajWorld achieves strong performance with lower computational cost than TDM.
C. Extended Experimental Results
C.1. Detailed Prediction Error for Baselines
We report the prediction error for MLP Ensemble and TDM in Figure 10 and 11, respectively.
C.2. Quantitative Results for Off-Policy Evaluation
We report the raw absolute error, rank correlation and regret@ for each OPE method and task in Table 6.
C.3. Off-Policy Evaluation with Pre-trained Models on Varamater-Variant Variant Environments.
In addition to the zero-shot prediction error reported in Section 5. we further investigate our four-layer model pre-trained on Walker2D with variant friction, mass, etc. Specifically, we evaluate the model's performance by fine-tuning it and testing on downstream off-policy evaluation tasks on standard Walker2I). The results are summarized Table 7. This provides additional evidence, beyond the zero-shot prediction error, demonstrating that TrajWorld exhibits strong capability for transfer to environments with varying parameters.
Env.
Level
Traj World (w/o PT)
TrajWorld (w/ PT)
random
262 + 34
76
medium
68 + 2
40 4
Walker2D
m-replay
71 + 11
46 1
m-expert
49 +
76 + 1
expert
281 + 8
186 +
Table Raw absolute error of off off-policy evaluation four-layer World model trained from scratch compared a model fine-tuned from pre-trained version on the Walker2D dataset with variant environment parameters with holdout onest, averaged over two seeds.
18

---

Trajectory World Models for Heterogeneous Environments
hattineddah-medmmy-
heltiheetah- expertt-222
10
hoppar--andoom-222
hoppit-mmlllm--22
hoppait-oxpatt-22
10"
10"1
60-
1H-0
81-0
m
60-7
11-0
m-
61-0
waller225-randsmm-222
We 1akko222-maaiium--pl----222
wakkr26-modiumm
walkk
wilka26--opettt-22
Mi Feueell (A
Predation Enar
10"
10'
10"
10
10
10"
Figure 10. Mean absolute errors (MAE) of transition prediction for MLP Ensemble, with and without pre-training (PT), across different train-test dataset pairs. Each subplot corresponds to distinct training dataset, with the test datasets shown the x-axis (r=random, m-r=medium-replay. m-medium, m-c=medium-expert, e=expert). Error hars represent the standard deviation across three random seeds.
meduumv22
halidheeddhh-oxpert-v
10"
10"
Predcclon
1O
10
hopppri-randoom-122
hoppet-medlum-02
hopper-medduucooppreee
hogper-export--22
10
10"
Bx110"
mhe
velker2d-randione
walker2ddmedummrrpzdee
waaker2-mmdimmwz
walker255experr
FIM TDM (we
10a
10
4x110
$x10
Figure 11. Mean absolute emors (MAE) of transition prediction for TDM, with and without pre-training (PT), across different train-tes datasct pairs. Each subplot corresponds distinct training dataset, with the test datasets shown the x-axis (r=random, m-r=medium- replay, m=medium, m-e=medium-expert, e=expert). Error bars represent the standard deviation across three random seeds.
19

---

Trajectory World Models for Heterogeneous Environments
Env.
Level
FTM
MI.P (w/o PT)
MI.P (w/ PT)
TDM (w/o PT)
TDM (w/ PT)
TW (w/o PT)
TW (w/ PT)
random
236 + 15
245 #
307 = 15
79 19
160 = 17
259 27
98 =
medium
47 21
149 +:
181 # 18
140 $ 10
145 7
81 #
127
Hopper
m-replay
8
2
56
+
m-expert
32
87 3:
173 = 15
116 21
79
48 +
69
expert
71 16
167 + 36
218 29
2
100
4
random
339 10
356 1
3 2
291 + 40
264 2
312 19
269 +:
medium
159
181 + 10
371
104
123 12
61
101 #7
Walker2D
m-replay
132
131
313
14
14
54
182 10
m-expert
9
210
= 19
24
137 = 17
60 11
72
expert
3
344 + 20
368 + 15
403 1 141
458 + 19
272 + 124
100 +
random
842
965
1137 27
1079
1050
1028 17
1059 =
medium
655 114
734 244
973 + 91
54
1312
568 23
444 =
Halfcheetah
m-replay
727 119
59
993 - 41
961
730
540
540 + 16
m-expert
689 203
65
1117
923
1319 23
+ 150
528 10
expert
758 + 116
973 175
1243 + 36
1273 = 158
646 + 50
1013 # 246
841 + 14
(a) Raw absolute error
Eny.
Level
ETM
MLP (w/ PT)
MLP (w/ PT)
TDM (w/o PT)
TDM (w/ PT)
TW (w/o PT)
TW (w/ PT)
random random
0.15
0.15
0.43 0.09
0,90 = 0.05
0.82
0.56
0.81 = 0,01
medium
0,94 0.04
0.81
0.72 0.03
= 0.10
0.06
0.81
0.31 0.10
Hopper
m-replay
0.97 + 0.02
0.98 + 0.00
0.96 + 0.01
0.89 0.89
0.86 0.86
0.61 # 0.36
m-expert
0.01 0.01
0.90
0.79
0.55 = 0.32
0.00
0.90
0.87 0.02
expert
0.85 + 0.05
0.02
0.42
-0.34
0.04
0.86
0.95 0.00
random
-0.12 0.32
0.03
0.17
0.73 0.10
0.02
0.00
0.78
medium
0.78 0.12
0.90 0.03
0.44 +
0.91
0.001
Walker2D
m-replay
0.77
0.01
0.72 0.08
= 0.03
0.93
0.07
m-expert
0.67 + 0.14
0.02
0.74 0.06
0.91 + 0.04
0.79
0.95 0.001
0.96 0.01
expert
0
0.36
0.11 0.42
0,36 + 0.42
0.80
0.330
0,94 0.01
random
0.76 0.10
0.90
0.84 + 0.12
0.93 0.00
0.00
0.900
0.94 0.00
medium
0.78 0.12
0.93
0.01
0.32
0.14
0.00
0.98 0.00
Halfcheetah
m-replay
0.77
0.00
0.88
0.93 0.03
0.82
0.00
m-expert
0.03
0.02
0.90 0.00
+ 0.10
0.07
0.06
expert
0.10
0.90 0.06
0.24 = 0.44
0.24 + 0.20
0.81 = 0.02
0.29
0.94 0.01
(b) Rank correlation
Eny.
Level
ETM
MLP (w/o PT)
MLP Ens.ww/ PT)
TDM (w/o PT)
TDM (w/ PT)
TW (w/o PT)
TW (w/ PT)
random
0.20 0.10
0.30 0.28
0.62 0.00
0.20 0.15
0.25 0.08
0.39 0.32
0.37 = 0.00
medium
0.05 0.04
0.03 0.04
0.05 0.04
0.17 # 0.17
0.16 0.00
0.22 0.13
0.11 + 0.06
Hopper
m-replay
0.00 0.00
0.08 0.00
0.08 0.00
0.09 0.02
0.08
0.05
0.14
m-expert
0.08 0.07
0.09 0.02
0.12 0.21
0.37 0.00
0.08
0.15 0.13
+
expert
0.08 0.02
0.17 0.17
0.17
0.08
0.12
random
0.16 0.09
0.05 + 0.01
0.41 + 0.40
0.10 0.04
0.05 + 0.00
0.17 0.00
0.08 + 0.00
medium
0.00 0.00
0.08 0.00
0.28 0.00
0.08 + 0.07
0.00
0.04 0.07
0.08 + 0.00
Walker2D
m-replay
0.00 0.00
0.07 = 0.02
0.19
0.08 = 0.04
0.10
0.03 0.05
0.00 0.00
m-expert
0.03 : 0.02
0.08 = 0.00
0.16
0.06 = 0.10
0.00
0.05 0.05
0.08 = 0.00
expert
0.05 0.05
0.19 0.26
0.12 0.14
0.28 0.00
0.28 0.28
0.17 0.10
random
0.20 1 0.10
0.15
0.11
0.04
0.008
0.10
= 0.02
medium
0.08 0.08
0.18 0.00
0.17 + 0.02
0.70 0.52
0.23 0.007
0.12 0.11
0.12 0.02
Halfcheetah
m-replay
0.16 : 0.12
0.15 0.10
0.23 = 0.07
0.16 =
0.25
0.18 = 0.00
0.18 0.00
m-expert
0.11 + 0.10
0.06
0.18 0.00
0.16 = 0.05
0.37
0.14
0.00 0.00
expert
0.12 0.07
0.06
0.04 0.00
0.27 0.17
0.14
0.18 0.03
0.20 0.19
(c) Regret@1
Table 6. Quantitative results of all model-based methods (TW=TrajWorld) for OPE, averaged seeds.
20

---

Trajectory World Models for Heterogeneous Environments
Env.
MLPP (w/o PT)
MLP (w/ PT)
TDM (w/o PT)
TDM (w/ PT)
TW (w/o PT)
TW (w/ PT)
Proposal
Hopper
948++61
1091--125
1287+26
1117+145
1090+-225
1401+236
1078+143
Walker
3353+83
3465+20
3056+236
2619+33
2422+-455
3427+370
3049-104
HalfCheetah
5645+10
5692+19
5611
5647+-25
5858+177
5809+11
5697+33
Table 8. Quantitative results of all model-based methods (TW=TrajWorld)) for MPC with action proposal, averaged over 3 seeds.
datbset:
datsset
Proponal Trainng
Eptocd Rolum
Aggregated Results (6 Training Sets)
Eplocote
Trakning Pre Tratiing
Creent
Trilettt (Ount)
Enswmble
TraWWrd (Dum)
Cheart
TralWide (Oung
30
Train datamal: walkr22k4aaddrmme2
Trai dicaset: karldtmedd.nmag
Trann ditioal: walker2d-osput--222
198-
14--
Eppood foeum
Ephoodo Rooum
TrafWard (Durs)
TrafAbse (Durs)
(Our)
TiglWbe (Ous)
Figure 12. Model predictive control (MPC) results using random shooting planner, averaged across three random seeds. The proposal policy line indicates the performance of a random action-sampling strategy.
C.4. Additional Model Predictive Control Results
Quantitative results with proposal policies. We report quantitative results MPC with action proposal in Table 8.
MPC with random shooting planner. Figure 12 presents MPC results using a random shooting planner with models trained on different datasets.
Computational efficiecny. TrajWorld predicts variates jointly, unlike TDM which processes them sequentially. This leads to a major speedup: MPC for 1000 environment steps in HalfCheetah takes 40 minutes with TDM, but only 3 minutes with Torl..
C.5. Additional Zero-shot Cross-Environment Transfer
Prediction
0.2
(Prec)
(Pred) Tra (.Rell)
Comparison with baselines. We also provide zero-shot prediction from other baselines Figure 14. As shown, in an unseen environ- ment, both TDM and MLP baselines fail to generalize, producing incorrect predictions and failing to capture the underlying state-action relationship all. Specifically, TDM fails to predict how push forces from two opposite directions lead to different positions. On the other hand, MLP fails to produce reasonable results with extreme error accumulation.
Traj (Real)
Position
-0.2
Time Stop
Figure 13. TrajWorld's zero-shot predictions for two Cart-3-Pole trajectories, which share 10 context steps but diverge due to differing subsequent actions.
Cart-3-pole environment. We also test Worldss'''s zero-shot pre- diction on the more challenging Cart-3-pole environment, which has an 1 1-dimensional state space. Surprisingly, TrajWorld can still give cart's position predictions roughly aligned with the ground truth, despite not seeing this embodiment before. The action sequence is depicted in Section B.2.2.
21

---

Trajectory World Models for Heterogeneous Environments
Zero-Shot Prodiction on Cart-2-Pole
Zaro-Shot Prodiction on Carl-2-Pole
0.2
Tra (Rea))
Tra (TDM)
Traj (Real)
Traj (TDM)
Try (TeiWond))
Tra (MLP)
Traj (TnijWond)
Tra (MLP)
Postion (m)
0.0
Positio
-0.2
-0.2
6
12
15
18
3
9
12
15
18
Time Stop
Time Stop
Figure 14. Zero-shot predictions from different pre-trained models on two Cart-2-Pole trajectories that share the same context steps but diverge thereafter due different future actions.
C.6. Additional Variate Attention Visualization
We present the variate attention maps of our Traj World model across all six layers, comparing fine-tuned model and a model trained from scratch, in Figures 15 and 16.
For the fine-tuned model, the early layers (Layer 0 and 1), attention is more scattered and less structured, likely capturing broad and low-level features. In contrast, later layers (Layer and 5) exhibit more focused attention, suggesting the model concentrating specific relationships entities. The prominent diagonal patterns and neighboring attentions diseussed in Section 5.5 can also be clearly observed in Layer 2. Additionally, diagonal patterns linking joint velocities and actions appear in Layers and 2. Such diagonal patterns are also observed in the attention maps of the model trained from scratch.
notable difference between the attention maps the fine-tuned model and the model trained from scratch the earlier emergence of diagonal patterns in the layers of the model trained from scratch. Specifically, while the first two layers of the fine-tuned model exhibit more scattered and less interpretable attention, the scratch-trained model immediately begins capturing structured diagonal patterns, particularly between positions and velocities, well as velocities and actions. This probably suggests that pre- pre-training changes the model's 's behavior. The model without pre-tree pre-training tends to focus environment-specife patterns and more localized features for prediction. In contrast, the fine-tuned model seems to dedicate its first two layers to extracting more semantically meaningful and generalizable features, encouraging the model to perform inference through in-context learning from these environment-agnosie representations.
Layer
Layer
Layer
Layer
Layer
Layer
Figure 1.5. Variate attention maps of our pre-trained Traj World Model, fine-tuned under Walker2D environment.
Layer 0
Layer
Layer 2
Layer
Layer
Layer 5
Figure 16, Variate attention maps of our TrajWorld Model in the Walker2D environment, trained from scratch.
C.7. Additional Ablation Study on Pre-training Dataset
To investigate the contributions of different components of the UniTraj dataset to the pre- pre-training process, we conduct an ablation study by training four-layer TrajWorld model on modified version of the UniTraj dataset, excluding two data
22

---

Trajectory World Models for Heterogeneous Environments
sources more closely aligned with the target environments: Modular RL and TD-MPC2. The results presented in Table indicate that the advantages of pre-training stem from the diversity encompassed within the complete UniTraj dataset, rather than relying solely on data from domains closely resembling the target environments.
Env.
Level
TrajWorld (w/ PT)
TrajWorld (w/ PT)
Halfcheetah
medium
491 94
468 + 40
Walker2D
medium
88 = 21
54 = 2
Walker2D
expert
141 + 15
116 + 12
(a) Raw absolute error
Env.
Level
TrajWorld (w/o PT)
TrajWorld (w/ PT)
Halfcheetah
medium
0.95 # 0.00
0.97 + 0.01
Walker2D
medium
0.93 = 0.01
0.97 + 0.02
Walker2I
expert
0.56 + 0.20
0.81 0.04
(b) Rank Correlation
Env.
Level
Traj World (w/o PT)
TrajWorld (w/ PT)
Halfcheetah
medium
0.18 = 0.00
0.05 0.07
Walker2D
medium
0.04 + 0.06
0.00 0.00
Walker2D
expert
0.34 # 0.31
0.16 + 0.16
(c) Regret@ @ 1
Table OPE results for four-layer Wrrld model trained from scratch compared to model fine-tuned from a pre-trained version on the ablation dataset, averaged over two seeds.
D. Extended Discussion
Limitations of bounded prediction. Our discretization scheme (Section 4.2) has the drawback that it can only represent variate values within the bounded range [bo. bB), restricted by the maximum and minimum in training data. This can lead to inaccurate predictions. For example, a model trained on trajectories from low-performing policies, may underestimate the reward a high-rewarding transition. This may explain why our model slightly underperforms Regret@ for off- off-policy evaluation tasks. Since all variates share the same bin embeddings, promising way to address this issue to simply extend the value range of bins beyond the observed data limits for variates with narrow coverages. Although the model would not have encountered those out-of-range values for specific variate during training, we hypothesize it could extrapolate similarly to regression models (e.g., MLPs), leveraging learned bin ordering shared with other variates. This hypothesis is supported by the bin continuity observed in Figure 8b. Further exploration and improvement of this approach are left for future work.
Discussion with Schubert et al. (2023). We demonstrate positive transfer to complex downstream environments such as Walker2D, not only for offline transition prediction and policy evaluation, but also for online MPC, which Schubert et al. (2023) did not. Our work differentiate from theirs in: (1) Setting: Instead of fine-tuning with 10a episodes for MPC with random shooting, we more practically fine-tune with 102 episodes for MPC with proposal policies; (2) Data diversity: Our UniiTraj dataset emphasizes distribution diversity, rather than using pure expert trajectories; (3) Architecture: TrajWorld incorporates inductive biases tailored to the 2D structure of trajectory data for enhanced transfcrability. Notably, TDM exhibits negative transfer in our practical MPC setting. We believe our work complements and extends Schubert et al. (2023), offering new insights to the community.
23


Reviews:
[{"review_id": "Ilj33e2pUP", "reviewer": "Reviewer_2GyG", "summary": "This paper presents a trajectory world model that handles varying sensor and actuator information across different environments. To support the generalization of the world model, this work composes a large dataset, UniTraj, comprising over one million trajectories from 80 environments. The key ingredient of the proposed method is to model the dynamic transitions on discretized variates over the temporal horizon. The learned world model demonstrates effective positive transfer across heterogeneous and complex control environments while achieving a new state-of-art for off-policy evaluation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer 2GyG for the thoughtful review and valuable comments, especially the recognition of our idea of unifying world modeling across heterogeneous environments.\n\n## Q1: Dataset contribution\n\n**Dataset construction**: We respectfully disagree with the assessment that our dataset is a simple composition. We elaborate on the insights behind UniTraj as follows:\n\n1. **Selection**: We carefully **selected** the subsets to assemble UniTraj. Unlike Schubert et al. (2023), which utilizes only expert or near-expert trajectories, our dataset emphasizes distribution diversity beyond sole environment diversity, as detailed in Sec 3.\n2. **Additional collection**: To further enrich environment diversity, we **collected** new trajectories ourselves from Modular RL, going beyond existing datasets.\n3. **Filtering**: As noted on Line 124, we **filtered** all trajectories from three downstream environments, contributing a reasonable testbed for cross-environment transfer to the community. \n4. **Weighting**: We manually **weighted** different subsets, trying to balance size and diversity. For example, DB-1 is oversampled due to its too small size. We apologize for not including these weights in the appendix and provide them below:\n\n| Subsets                        | ExoRL | RLU | JAT | DB-1 | TD-MPC2 | Modular RL |\n| ------------------------------ | ----- | --- | --- | ---- | ------- | ---------- |\n| (Unnormalized) sampling weight | 75    | 5   | 90  | 1    | 90      | 30         |\n\n**Analysis of dataset scales and diversity**: We appreciate the suggestion to analyze the impact of pre-training data. We conducted new experiments by pre-training three versions of TrajWorld on different subsets of UniTraj—namely, 1/10 size, 1/100 size, and the JAT subset (with only expert trajectories from 5 environments)—followed by fine-tuning on downstream tasks.\n\nFor **transition prediction**, we adopt a challenging setting: for each environment, we train models on the expert dataset, and test them on datasets of all levels. We also provide results for **model-predictive control** (see experimental details below).\n\nThe results, shown in this [anonymous figure](https://anonymous.4open.science/r/TrajWorld/data-analysis.pdf), compare these new models with one pre-trained on the full UniTraj and another trained from scratch. We observe that all subset pre-trained models underperform the fully pre-trained one, revealing a scaling law with respect to data size. These findings underscore the **importance of both scale and diversity** in pre-training data, and strengthen our contribution in advocating for large-scale, heterogeneous pre-training and in constructing the UniTraj dataset to support it.\n\n## Q2: Architecture novelty\n\nWe respectfully believe there may have been a misunderstanding regarding our architectural contributions in comparison to TDM:\n\n- TDM actually models tokens across both the variate and temporal dimensions, similar to ours. But it **flattens them into a one-dimensional sequence** and applies the original GPT architecture. This approach discards the inherent 2D structure of trajectory data and may 'not bring further insights to the community'.\n- Compared to TDM, our TrajWorld does not 'newly introduce the temporal dimension', but rather **preserves and exploits the natural 2D structure** by employing a two-way attention mechanism, with each one capturing relationships within its respective dimension.\n\nThe superior performance of TrajWorld underscores the importance of appropriate inductive biases for enhancing transferability in trajectory world models.\n\n## Q3: Experiments on decision-making (MPC)\n\nWe have added experiments on improving policies via sampling-based optimization (**model-predictive control**). In this setting, **TrajWorld outperforms both baseline models and its counterpart trained from scratch** (see this [anonymous figure](https://anonymous.4open.science/r/TrajWorld/mpc.pdf)).\n\nDue to space limitation, we kindly refer the reviewer to **response to W1 for Reviewer puvo** for detailed experimental setup and results.", "claims_and_evidence": "Yes, the claims made in this submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Issues of methods:\n\n- **Limited contribution regarding the UniTraj dataset**. This paper introduces the UniTraj dataset by composing large-scale data from different environments. However, there are limited contributions in this regard. First, the dataset seems to be a simple composition without any rigorous curation, selection, or filtering process. Further, even though the paper demonstrated that using UniTraj in pre-training can improve downstream performances as shown in Fig.1 and Fig.5, there is a lack of critical studies on how different data scales and diversity affect model performance. For instance, how would the world model generalize when pre-trained on a subset of UniTraj, such as 1/2 or 1/10? Given these issues, the introduction of this dataset cannot form a valid technical contribution.\n\n\n- **Limited novelty of the proposed TrajWorld architecture compared to TDM**. According to Fig.3 and descriptions in L322-324, the main difference between TDM[1] and the proposed TrajWorld is the newly introduced temporal dimension. However, leveraging the temporal dimension with additional attention modules has been widely studied in the literature on world models[2] and video generation models[3, 4]. Such an adaptation does not bring further insights to the community.\n\n[1] A Generalist Dynamics Model for Control\n\n[2] Generalized Predictive Model for Autonomous Driving\n\n[3] AnimateDiff: Animate Your Personalized Text-to-Image Diffusion Models without Specific Tuning\n\n[4] Align your Latents: High-Resolution Video Synthesis with Latent Diffusion Models", "theoretical_claims": "No issues with the theoretical aspect of this work.", "experimental_designs_or_analyses": "Issues of experimental designs:\n\n\n- **Lack of experiments on decision-making**. The paper did provide one means of utilizing the pre-trained world model, which is off-policy evaluation in Sec.5.3.  However, investigations on using this model to improve policy through reinforcement learning or sampling-based optimizations are underexplored.", "relation_to_broader_scientific_literature": "The key contributions of the paper relate to the areas of world models, reinforcement learning, and pre-training on large-scale dataset.", "essential_references_not_discussed": "N/A.", "other_strengths_and_weaknesses": "Strengths:\n- The paper is well-organized and easy to follow\n- Experiments on how the pre-training influences downstream fine-tuning are meaningful. As shown in Fig.5, the authors conducted a wide range of experiments to demonstrate that large-scale pre-training is helpful for downstream fine-tuning with various train-test pairs. \n- The idea of unifying world modeling on heterogeneous environments through in-context variates is intriguing and worth exploring.\n\nOther weakness:\nSee weaknesses in the above sections.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "N/A.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lbu3k6WAEB", "reviewer": "Reviewer_puvo", "summary": "The paper aims to tackle the heterogeneity issue in world model pretraining. To achieve this, the authors curate a unified trajectory dataset from 80 control environments. Based on this dataset, they introduce TrajWorld, a world model architecture that naturally accommodates varying sensors and actuators, thereby enabling efficient knowledge transfer across environments. The effectiveness of TrajWorld is validated across three unseen environments in terms of prediction error and policy evaluation reliability.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer puvo for the thorough review, insightful questions, and a positive evaluation of our work.\n\n## W1: Model-predictive control (MPC) evaluation\n\nWe have conducted MPC experiments comparing different world models. \n\n**Setup**: Following Schubert et al, we first attempted *MPC with random shooting planner*, but found it ineffective in our high-dimensional target environments, with any world models. We then adopt the *MPC with proposal* setting, also from Schubert et al. We set a practical scenario where world models trained on medium replay datasets are used to improve the medium-level trained proposal policies via MPC. \n\n**Implementation**: We use three medium-replay datasets from D4RL and medium-level policies from DOPE. The sample size is fixed at 128. The planning horizon is set to 25 steps for HalfCheetah and Walker, and 50 steps for Hopper to avoid myopic behavior in this fragile environment. Sampling noise is tuned for optimal performance across all world models: 0.05 (Hopper), 0.2 (Walker), 0.025 (HalfCheetah).\n\n**Results**: Our main results for MPC with proposal are shown in the table below and this [anonymous figure](https://anonymous.4open.science/r/TrajWorld/mpc.pdf). We find MPC improves proposal polices in Hopper and Walker but has little effect on HalfCheetah, which is more stable and less prone to failure. In contrast, Hopper and Walker are fragile, and the model helps prevent unsafe actions, leading to better planning.\n\nOverall, **MPC with TrajWorld delivers the best performing agents** compared to baseline models or its counterpart trained from scratch. \n\n|MPC w/ proposal & | MLP-Ensemble (w/o pt) | MLP-Ensemble (w/ pt) | TDM (w/o pt) | TDM (w/ pt) | TrajWorld (w/o pt) | TrajWorld (w/ pt) |*Proposal only*|\n|-|-|-|-|-|-|-|-|\n|Hopper|948±61|1091±125|1287±26|1117±145|1090±225|**1401**±236|*1078±143*|\n|Walker|3353±83|**3465**±20|3056±236|2619±36|2422±455|**3427**±370|*3049±104*|\n|HalfCheetah|5645±10|5692±19|5611±85|5647±25|**5858**±17|5809±15|*5697±30*|\n\nFor MPC with random shooting, no world model provides successful agents. However, TrajWorld still performs relatively best among them (see [anonymous figure](https://anonymous.4open.science/r/TrajWorld/mpc.pdf)).\n\n**Efficiency**: TrajWorld predicts all variates jointly, unlike TDM which processes them sequentially. This leads to a major speedup: MPC for 1000 environment steps in HalfCheetah takes **40 minutes with TDM**, but only **3 minutes with TrajWorld**.\n\nWe thank all reviewers for encouraging us to add these experiments, which further strengthen the contribution of our work. We will include them in the final version.\n\n## W2: Long-horizon comparison\n\nWe remark that our TrajWorld has a flexible architecture similar to TDM and can also handle arbitrary prediction horizons (only constrained by its training context length).\n\nBelow, we summarize our experimental settings regarding horizons:\n\n- **Zero-shot generalization** (Fig. 4b): rollouts for **10 steps**. We have expanded the baseline results for comparison (see Q2 below).\n- **Transition prediction** (Sec. 5.2): short horizon, report **one-step** prediction error.\n- **Off-policy evaluation** (Sec. 5.3) involves rollouts of an extremely long horizon (**2000 steps** as mentioned in App B.4.1). Due to our model's short context length of 20 (limited by our computational resources), this is done by sliding windows.\n- **Model-predictive control**: rollouts over relatively long horizons of 25 or 50 steps.\n\nAcross all these scenarios, TrajWorld provides superior performance compared to baselines.\n\n## Q1: Scaling laws regarding data sizes and diversity\n\nWe have conducted an analysis on the effects of different data sizes and diversity for pre-training. Our results in this [anonymous figure](https://anonymous.4open.science/r/TrajWorld/data-analysis.pdf) validate that both the large scale and diversity contribute to the effectiveness of pre-training.\n\nDue to space limitation, we kindly refer the reviewer to **response to Q1 for Reviewer 2GyG** for experimental details and results.\n\n## Q2: Straight zero-shot prediction\n\nWe suppose that the overly straight predictions are likely due to the short context window (10 steps) and slow motion speed. Through this context, the zero-shot model is not able to precisely capture the quadratic relationship between state and action, and is only able to reflect coarse directional changes.\n\nFor comparison, we also provide zero-shot prediction from other baselines in this [anonymous figure](https://anonymous.4open.science/r/TrajWorld/zero-shot-prediction.pdf). As shown, in an unseen environment, both TDM and MLP baselines fail to generalize, producing incorrect predictions and failing to capture the underlying state-action relationship at all. Specifically, TDM fails to predict how push forces from two opposite directions lead to different x positions. On the other hand, MLP fails to produce any reasonable results with extreme error accumulation.", "claims_and_evidence": "Yes, there are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The effectiveness of the proposed world model is assessed through prediction error and off-policy evaluation. However, an evaluation of model predictive control performance would be more convincing, as did by the closely related work TDM [1].\n\n[1] A Generalist Dynamics Model for Control. Ingmar Schubert, et al.", "theoretical_claims": "I have checked the theoretical claims in this paper.", "experimental_designs_or_analyses": "The paper provides extensive experiments with sufficient details.", "relation_to_broader_scientific_literature": "The paper makes two key contributions in comparison to the related literature. First, it curates a large-scale trajectory datasets sourced from 80 control enviroments. Second, it presents a unified architecture to efficiently extract transferable knowledge from the heterogeneous data sources. These contributions could inspire future research on developing generalist world models.", "essential_references_not_discussed": "As far as I know, all closely related works are cited appropriately.", "other_strengths_and_weaknesses": "W1) **Missing MPC results.** While the proposed recipe improves zero-shot prediction and policy evaluation, its advantages are not demonstrated in model predictive control performance, which would provide an intuitive comparison to TDM.\n\nW2) **Long-horizon comparison.** I agree that the sequential prediction of TDM may lead to error accumulation, but its flexible architecture could offer advantages for planning over long horizons. The paper does not specifically mention the experimental settings regarding horizons, so I am curious whether the proposed joint prediction scheme still dominates in long-horizon settings.", "other_comments_or_suggestions": "A typo in Line 637: \"the official repository repository\" should be \"the official repository\".", "questions_for_authors": "Q1) **Scaling law.** As the authors have collected a large-scale dataset from 80 heterogeneous environments and demonstrated the benefits of pretraining, it would be interesting to see an analysis of the scaling laws regarding data diversity.\n\nQ2) **Straight zero-shot prediction results.** The zero-shot prediction results in Figure 4 appear overly straight. Could you briefly discuss possible reasons for this phenomenon? Besides, I wonder if it is possible to visualize the predictions of other baselines for comparison? I believe it would be beneficial to demonstrate the advantages of TrajWorld.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xMbMZvCGeW", "reviewer": "Reviewer_8E9V", "summary": "This paper introduces the UniTraj dataset, which contains a large set of trajectories collected from 80 heterogeneous environments. It also presents a world model, TrajWorld, pretrained on this dataset. The pretrained world model demonstrates positive transferability to new environments in zero or few-shot settings. The paper primarily evaluates the trained world model in an off-policy evaluation setting for transition prediction policy evaluation.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer 8E9V for the thorough review and valuable questions. \n\n## Q1: MPC evaluation\n\nFollowing Schubert et al., we have added **online MPC experiments**. In this setting, TrajWorld outperforms both baselines and its counterpart trained from scratch (see [anonymous figure](https://anonymous.4open.science/r/TrajWorld/mpc.pdf)). Due to space limitation, please refer to **response to W1 for Reviewer puvo** for details.\n\n**Discussions on transferability compared to Schubert et al.**: Given these new results, we can expand the discussions on Schubert et al.:\n\n> We demonstrate positive transfer to complex downstream environments such as Walker2D, not only for offline transition prediction/policy evaluation, but also for online MPC, which Schubert et al. did not. Our work differentiate from theirs in: (1) Setting: Instead of finetuning with $10^4$ episodes for MPC with random shooting, we more practically finetune with $10^2$ episodes for MPC with proposal policy; (2) Data diversity: Our UniTraj dataset emphasizes distribution diversity, rather than using pure expert trajectories; (3) Architecture: TrajWorld incorporates inductive biases tailored to the 2D structure of trajectory data for enhanced transferability. Notably, TDM exhibits negative transfer in our practical MPC setting. We believe our work complements and extends Schubert et al., offering new insights to the community.\n\n## Q2: Difference with factorized attention in Wayformer\n\nWe appreciate the feedback and acknowledge the relevance of Wayformer. We will include it as related work in the final version.\n\nWhile TrajWorld and Wayformer both adopt two-way attention (a.k.a. axial/factorized attention [1,2,3]) to handle a heterogeneous set of inputs with various numbers of dimensions, our work differs in several key aspects to preserve novelty:\n\n1. **Tasks & architectures**: Wayformer targets motion forecasting, a regression task with heterogeneity in the multimodal **input** space. TrajWorld is designed for world modeling, an autoregressive task involving heterogeneity in both **input and output** space. This leads to significantly different macro and micro designs: **decoder-only with causal attention vs. encoder-decoder with bidirectional attention**.\n2. **Homogeneous basics**: Wayformer handles various numbers of contextual objects across modalities via attention, but still relies on modality-specific projections to unify embedding dimensions. More thoroughly, TrajWorld achieves scalar-level homogeneity, **without the need for modality-specific projections** and is capable of zero-shot generalization to unseen input/output spaces.\n3. **Transferability beyond efficiency**: Wayformer uses factorized attention mainly for efficiency, not showing a performance boost over full attention in its target task. In contrast, TrajWorld shows that inductive biases for the 2D structure **enhance transferability**, outperforming its 1D counterpart, TDM.\n\nTo our knowledge, TrajWorld is the first to apply two-way attention in trajectory world modeling, and we hope it provides valuable insights to the community.\n\n[1] ViViT: A Video Vision Transformer.\n\n[2] Axial Attention in Multidimensional Transformers.\n\n[3] CCNet: Criss-Cross Attention for Semantic Segmentation.\n\n## Q3: Other questions/clarifications\n\n- **Error accumulation in TDM**: TDM predicts sequentially along both variate and temporal dimensions due to its 1D architecture. TrajWorld, with its 2D architecture, predicts sequentially over time but jointly across variates at each timestep, which helps mitigate error accumulation.\n- **Worse results of TDM compared to TrajWorld**: As discussed above, we attribute TDM’s performance gap to error accumulation along sequences of variates and the lack of appropriate inductive biases for the 2D structure of trajectory data. Both models use the same parameter count.\n- **Discretization strategies**: Uniform discretization is widely used (e.g., in Gato, TDM, Farebrother et al.) and performs well in our experiments. Thus, we did not explore more complex methods like quantile-based binning.\n- **In- vs. out-of-distribution**: For transition prediction, we test trained models on both the same and different datasets. For instance, a model trained on hopper-medium-replay is tested on both hopper-medium-replay (in-distribution) and hopper-expert (out-of-distribution). Pre-training benefits TrajWorld in both scenarios.\n- **Which policy would give a more generalized world model**: We believe that diverse data from a range of policies (random, medium, expert, exploratory, etc.) yields a more generalizable model than trajectories from any single policy. We have conducted experiments pre-training our model on JAT, a subset of UniTraj with only expert data, which underperforms those pre-trained on the full data. Please refer to **response to Q1 for Reviewer 2GyG** for details.", "claims_and_evidence": "The main claim is that by pretraining on a diverse set of trajectory environments, the pretrained world model can adapt to unseen environments with zero-shot or few-shot transfer, demonstrating improved transition prediction ability. Compared to previous models, the proposed model achieves better transferability in an offline setting. However, the authors did not test its ability to utilize the pretrained world model online, as was done in previous work with MPC.  There are also some questions for the experiment setup that need to be clarified.", "methods_and_evaluation_criteria": "Yes, for the most part, the proposed methods and evaluation criteria align with the problem and application at hand. The use of the UniTraj dataset for pretraining and the off-policy evaluation (OPE) setup are reasonable choices for assessing transferability in heterogeneous environments. However, there are some concerns relate to the experiment design which will be discussed in that section.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "After closely looking at the experimental design, I have the following concern regarding its setup. \n\nCompared to Schubert et al. (2023), why did the authors choose to evaluate using off-policy evaluation rather than a more realistic setting where the learned WM is directly used to solve the task online through MPC in the real environment?\n\nIn the related work section, the main point authors is trying to argue is that Schubert et al. (2023) did not show a positive transfer in walker 2D environment but Schubert et al. (2023) evaluates differently compared to this work. I consider the evaluation setup of Schubert et al. (2023) is more challenging as they need to solve the problem with the real environment through MPC.", "relation_to_broader_scientific_literature": "Modeling a world model is important for creating realistic simulations, especially when running such simulations is costly. This paper demonstrates that with a transformer-based world model and a diverse set of related environments, one can produce an adaptive world model prior.", "essential_references_not_discussed": "The close related work is discussed Schubert et al. (2023), though regarding the novelty of the transformer architecture, factorized transformer [1] was not mentioned. \n\n[1] Nayakanti, N., Al-Rfou, R., Zhou, A., Goel, K., Refaat, K.S. and Sapp, B., 2023, May. Wayformer: Motion forecasting via simple & efficient attention networks. In 2023 IEEE International Conference on Robotics and Automation (ICRA)(pp. 2980-2987). IEEE.", "other_strengths_and_weaknesses": "Strength:\n\nThe analysis in the experiment section is detailed and clear. \nProposed architecture performs better than the re-implementation of the baselines and demonstrate good transfer ability. \n\nWeakness: \nThe proposed transform architecture look similar to the factorized attention in previous work, which makes the novelty a bit questionable. \nThe usefulness of the learned WM prior is questionable given the author did not use it online with the environment.", "other_comments_or_suggestions": "Figure 5 missing labels for each row.", "questions_for_authors": "In 5.2 the authors mention \"Moreover, TDM predicts variants sequentially, which may accumulate errors and lead to less accurate results.\", but TrajWorld also operates sequentially which can accumulate errors , could the author clarify? \n\nTDM seems to produce a lot worse results than TW in Table 5, do the authors have any insights on this? Are the parameters count equivalent between the two models?\n\nCould the author identify the major difference between the proposed interleaved temporal-variate attentions versus factorized attention introduced in Wayformer?\n\nDid the author explore the different discretization strategy? \n\n“Moreover, the transfer benefits are evident in both in-distribution and out-of-distribution scenarios.”  what does this sentence referring to? \n\nWhich policy would give you a more generalized world model, the expert or random policy or something in-between?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "txxt7RRW2K", "reviewer": "Reviewer_bd1u", "summary": "This manuscript has 2 contributions: \n1. A trajectory dataset UniTraj, a large-scale dataset including over one million trajectories collected from various distributions from 80 heterogeneous environments.\n2. A Transformer-based architecture TrajWorld integrates interleaved variate and temporal attention mechanisms, aiming at transition prediction.\nThe authors first introduced the current challenge for trajectory prediction in a heterogeneous environment and their motivation. Then introduced the building process of dataset UniTraj, the architecture of TrajWorld model. The proposed model was tested on 15 datasets of 3 environments, which validated its performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate Reviewer bd1u's strongly positive feedback on our work. Your recognition of our clear writing, well-motivated approach, and the effectiveness of both our UniTraj dataset and TrajWorld architecture is truly encouraging. We are grateful for your support and share your belief that our work will contribute meaningfully to the broader scientific community.", "claims_and_evidence": "Yes, I think this manuscript is well-organized and clearly stated.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes. I think the authors clearly described the problem, related works, and their ideas.", "experimental_designs_or_analyses": "Yes. This proposed model used a two-way attention mechanism. The authors set a similar but with a one-dimensional attention model as the baseline, validating the good performance of the two-way attention mechanism", "relation_to_broader_scientific_literature": "The large-scale dataset can be a helpful tool for broader scientific literature.", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "I think the proposed dataset can be helpful for other researchers. And the proposed model can be a good baseline to study with.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "n/a", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes UniTraj, a dataset of many offline RL tasks, and TrajWorld, a world model that can train on all of them. 

First the dataset, it seems to be a concat of many existing datasets. The world model (TrajWorld) is a fairly off the shelf architecture, with a spatiotemporal factorization like is done in many video models circa 2022-2024.

The overall premise is promising, "building a foundation world model", yet, the approach doesn't quite live up to it. It is quite likely that to really do this would require leveraging internet videos, and it seems very unlikely we would want a foundation model for small scale proprioceptive tasks.
Decision: Accept (poster)