{"id":"f9ca92f7-ed34-4b30-8faf-eab31553f0d6","arxiv_id":"2502.01366","paper_version":2,"verdict":"CONDITIONAL","confidence":"HIGH","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"A pre-trained trajectory world model with interleaved temporal and variate attention achieves positive transfer across heterogeneous control environments, improving transition prediction, off-policy evaluation, and model predictive control.","lead":"The paper introduces UniTraj, a dataset of 1.3 million trajectories from 80 diverse control environments, and TrajWorld, a Transformer architecture that learns world models across these heterogeneous environments by interleaving temporal and variate attention.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"UniTraj includes DMC analogues of the downstream Gym robots; the ablation excludes only TD-MPC2/Modular RL, so the heterogeneous-transfer claim rests on an untested 'unseen' assumption.","rationale":"The reader's weakest-assumption analysis correctly identifies the most load-bearing issue: the downstream D4RL environments are not genuinely disjoint from the pre-training corpus because DMC analogues of HalfCheetah, Hopper, and Walker2D are present in multiple UniTraj components. My reading of the manuscript confirms this: Section 3 explicitly excludes only the OpenAI Gym versions, and Appendix C.7 removes only TD-MPC2 and Modular RL, which are the most obviously aligned sources but not the only ones containing DMC robots. The claim of 'heterogeneous transfer' therefore depends on the untested premise that DMC-to-Gym differences in state/action dimensions, observation order, and parameters are sufficient to make the downstream environments 'unseen'. The proposed removal experiment directly tests that premise. I do not see a second concern more load-bearing: the architecture and dataset construction are internally consistent, the empirical comparisons are reasonably controlled, and the bounded-discretization limitation is transparently acknowledged. The OPE 'state-of-the-art' claim is partly qualified by Regret@1, but that is a secondary claim; even if it were softened, the central transfer claim would still stand or fall on the heterogeneity question. Because the current conditional verdict already requires clarification of the evaluation protocol and stronger evidence of true cross-environment transfer, my review does not change the verdict; it sharpens the specific experiment that should be demanded before acceptance.","tokens_in":1053,"tokens_out":1089,"duration_ms":51503,"concrete_test":"Pre-train TrajWorld on a variant of UniTraj with every environment sharing the downstream morphology removed—i.e., delete all DMC Cheetah, Hopper, and Walker data plus Cheetah-*/Hopper-*/Walker-* variants from ExORL, RL Unplugged, DB-1, TD-MPC2, and Modular RL—then fine-tune on the 15 D4RL datasets and compare aggregated transition MAE and OPE MAE/rank against from-scratch and against full-UniTraj. If the pre-training gain over scratch shrinks to a small fraction of its full-UniTraj value, the heterogeneous-transfer claim fails; if gains persist, the concern is resolved. A complementary control: pre-train only on the removed DMC Hopper/Walker/Cheetah data and fine-tune; if this alone reproduces most of the full-UniTraj gain, the same-morphology explanation is confirmed.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central claim is that pre-training on 80 heterogeneous environments transfers to complex 'previously unseen' environments. The main downstream evidence is fine-tuning on D4RL HalfCheetah, Hopper, and Walker2D. But Table 2 shows that DMC equivalents of these very robots are inside UniTraj: DMC Walker appears in ExORL, RL Unplugged, DB-1, and TD-MPC2; DMC Hopper appears in DB-1 and TD-MPC2; DMC Cheetah variants appear in DB-1, TD-MPC2, and Modular RL. Section 3 only excludes OpenAI Gym versions of the three robots, and the appendix ablation (Appendix C.7) removes only TD-MPC2 and Modular RL, leaving DMC analogues in ExORL/RL Unplugged/DB-1. If a substantial share of the reported gains comes from seeing the same or near-same morphology (same joints, same torque/position semantics) rather than from a general ability to model arbitrary sensor/actuator sets, then the headline 'first successful transfer across heterogeneous complex environments' is not established. The note in Table 2 that Gym and DMC versions differ in state/action definitions and parameters is an assumption, not a control. The zero-shot cross-environment demo (Cart-2-Pole) is qualitative only, and the parameter-transfer experiments in Section 5.1 vary physics within a single body, so they do not close this gap.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper introduces UniTraj, a dataset of 1.3M trajectories from 80 control environments with varying state/action dimensions, and TrajWorld, a transformer that scalarizes each state/action/reward component and applies interleaved temporal and variate attention. TrajWorld is pre-trained on UniTraj and then fine-tuned on D4RL datasets of Gym HalfCheetah, Hopper, and Walker2D. The authors report that pre-training reduces transition-prediction MAE, improves off-policy evaluation (OPE) on the DOPE benchmark, and improves model-predictive-control (MPC) returns, with qualitative zero-shot predictions on Cart-2-Pole/Cart-3-Pole. They claim the first successful transfer of world models across heterogeneous complex control environments.","tokens_in":42797,"tokens_out":6850,"duration_ms":59214,"significance":"The dataset and architecture are useful contributions, and the experimental protocol is extensive (75 train-test pairs, three seeds, multiple baselines, scale/diversity ablations). If the transfer results hold for genuinely unseen environments, this would be a notable step toward generalist world models for low-dimensional continuous control. The central caveat is that the downstream Gym environments are not established to be genuinely unseen, because DMC analogues of the same robots appear in UniTraj and the appendix ablation does not remove that overlap. The zero-shot evidence is qualitative and concerns morphology-near environments. The OPE state-of-the-art claim is stronger than the per-task results. These issues affect the headline claim and need to be addressed before publication.","major_comments":[{"comment":"The central claim that fine-tuning targets are \"previously unseen\" environments is not established. UniTraj contains DMC analogues of the same robots (Walker, Hopper, Cheetah) in ExORL, RL Unplugged, DB-1, and TD-MPC2; the paper excludes only the OpenAI Gym versions (Section 3) and the ablation in Appendix C.7 removes only TD-MPC2 and Modular RL, so these DMC analogues remain in the pre-training set. Because the DMC and Gym versions share the same body morphology and state/action semantics (joint angles, velocities, torques), the reported gains may largely reflect near-duplicate morphology rather than general cross-environment transfer. The Table 2 note that Gym and DMC versions differ in state/action definitions and parameters is an assumption, not a control. Please re-run the main transition-prediction, OPE, and MPC experiments with a pre-training set that excludes all DMC/Modular-RL analogues of the three test robots, or otherwise demonstrate that the gains persist on environments with non-overlapping morphology.","section":"§5.2, Table 2, Appendix A.2 and C.7"},{"comment":"The zero-shot cross-environment evidence is qualitative and restricted to Cart-2-Pole/Cart-3-Pole, which are morphologically very close to Cartpole environments present in UniTraj. Visual inspection of a small number of trajectories does not substantiate the claim of zero-shot transfer to arbitrary unseen sensor/actuator sets. Please provide quantitative prediction errors on a held-out set of test trajectories for these environments, and preferably add evaluation on a held-out embodiment that is more distant from the training morphologies.","section":"§5.1, Figures 4b, 13, 14"},{"comment":"The abstract's phrase \"achieves a new state-of-the-art for off-policy evaluation\" is stronger than the reported results. In Table 6, TrajWorld (w/ PT) is worse than ETM on several environment-level tasks (e.g., Hopper-medium raw absolute error 127 vs. 47, Hopper-medium-replay 73 vs. 29, HalfCheetah-random 1059 vs. 842) and it underperforms on Regret@1 as acknowledged in the text. The SOTA claim should be qualified to the averaged normalized absolute error and rank-correlation metrics, and the per-environment results should be discussed in the main text.","section":"§5.3, Table 6"}],"minor_comments":[{"comment":"The sentence \"evaluting both both from-scratch and fine-tuned variants\" contains a duplicated \"both\" and a typo in \"evaluting\".","section":"§5.4"},{"comment":"In the Related Work section, \"no for computational efficiency\" appears to be a typo for \"not only for computational efficiency\".","section":"§6"},{"comment":"The bullet \"resulting in simultaneous and significant improvements in transition prediction, off-policy policy, and model predictive control\" should say \"off-policy evaluation\" instead of \"off-policy policy\".","section":"Contributions bullet, §1"},{"comment":"The caption reads \"Y-axis at log scale\"; it should be \"Y-axis on a log scale\".","section":"Figure 1 caption"},{"comment":"The text \"two Gaussian distributions with the same standard derivation\" should be \"standard deviation\".","section":"Appendix B.1"}],"recommendation":"major_revision","confidential_remarks":"The paper is well-executed: the dataset and architecture are useful contributions, and the experimental effort is substantial. The main risk is that the headline claim of 'first successful transfer across heterogeneous complex environments' is vulnerable to the morphological-overlap critique. The authors should be asked to complete the ablation excluding all DMC analogues, or soften the claim. The OPE 'state-of-the-art' phrasing also needs qualification. These are fixable with additional analysis or revised wording, so I recommend major revision rather than rejection."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"First thing you should know: this is a solid empirical paper that makes a real step forward in pre-training world models across heterogeneous low-dimensional control environments. It ships a large dataset (UniTraj, 80 envs, roughly 700M steps) and an architecture (TrajWorld) that handles varying state/action dimensions via interleaved temporal and variate attention. The headline result—pre-training helps fine-tuned transition prediction on D4RL HalfCheetah, Hopper, Walker2D—is credible and reproduced across 75 train-test pairs with three seeds. That is genuinely new: TDM failed to transfer to Walker2D, and here pre-training consistently helps TrajWorld, while hurting an MLP baseline.\n\nThe OPE results are good, but the 'state-of-the-art' claim is softer than the abstract suggests: TrajWorld leads on MAE and rank correlation but underperforms ETM on Regret@1. The authors attribute this to bounded discretization, which they acknowledge in the limitations. That is honest and plausible, but it does cut the headline OPE claim.\n\nThe main soft spot is the heterogeneity claim. The three downstream environments are held out from UniTraj, but DMC analogues of those same robots are in the pre-training data (ExORL, RL Unplugged, DB-1, TD-MPC2). A true test of transfer across disjoint morphology would hold out all DMC variants of the downstream body. The appendix ablation removes only TD-MPC2 and Modular RL, which is not the strongest control. The note that Gym and DMC versions differ in state/action definitions is an assumption, not a controlled experiment. So the paper establishes positive transfer from data that includes close morphological relatives, but it does not yet establish the stronger 'general ability to model arbitrary sensor/actuator sets' claim. The zero-shot Cart-2-Pole demo is qualitative only.\n\nThat said, the work is well-executed and transparent about its limitations. The architecture is sensible, the dataset will be useful to the community, and the negative-transfer result for MLP is a useful caution. This deserves a serious referee. I'd recommend conditional acceptance: ask for an ablation that excludes all DMC analogues of the downstream robots (or at least ExORL, RL Unplugged, and DB-1), and for a more grounded zero-shot evaluation. If the transfer signal survives that, the heterogeneity claim is solid. If it doesn't, the paper is still a useful empirical study but should soften the 'first heterogeneous transfer' framing.","headline":"Solid empirical step forward in world-model pre-training, but the 'heterogeneous transfer' claim is weakened by DMC analogues of the downstream robots in the pre-training data.","tokens_in":43370,"tokens_out":1961,"would_cite":true,"duration_ms":18145,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Pre-training a world model on 1.3M trajectories from 80 heterogeneous environments transfers to unseen robots, improving prediction, off-policy evaluation, and model predictive control.","keywords":["world models","pre-training","heterogeneous environments","cross-environment transfer","in-context learning","off-policy evaluation","model predictive control","trajectory datasets"],"falsifier":"Re-train TrajWorld on UniTraj after deleting every pre-training environment that resembles the test robots—all DMC walker, cheetah, and hopper bodies, plus their Modular RL and TD-MPC2 variants—and then fine-tune on the Gym HalfCheetah, Hopper, and Walker2D D4RL datasets; if the gains over from-scratch training largely vanish, the claimed heterogeneous transfer is mostly near-duplicate morphology recall, while persistent gains would support the paper's interpretation. A complementary check is to rerun the zero-shot Cart-2-Pole and Cart-3-Pole rollouts with history context removed and confirm that prediction error rises, which would verify that in-context inference, not memorized per-environment statistics, carries the transfer.","tokens_in":42317,"feed_emoji":"🤖","tokens_out":14386,"duration_ms":107596,"temperature":0.7,"pith_summary":"This paper claims that a world model—a learned predictor of how actions change a system's state—can be pre-trained on trajectories from many different control environments and then transferred to new environments with different sensors, actuators, and dynamics. To make that possible, it curates UniTraj (1.3 million trajectories, 719 million steps, from 80 environments) and builds TrajWorld, a Transformer that reads each scalar state, action, and reward channel as an individual token rather than as a fixed-length vector. After pre-training and fine-tuning, TrajWorld reduces transition-prediction error on the HalfCheetah, Hopper, and Walker2D robots of the D4RL benchmark, sets a new state of the art on the DOPE off-policy evaluation benchmark, and improves online model predictive control. The payoff if true is practical: a single pre-trained model carries general knowledge of physical dynamics, so a new robot can get an accurate predictive model with far fewer costly interactions.","feed_headline":"One model learns 80 environments, predicts unseen robots better","feed_subtitle":"Pre-training on diverse sensors and actuators cuts prediction error and boosts offline evaluation and control.","key_machinery":"The mechanism is scalar-level tokenization plus interleaved two-axis attention. TrajWorld flattens a trajectory into a two-dimensional matrix of timesteps by variates, where each variate is a single scalar dimension of the state, action, or reward; every scalar is discretized into one of $B$ uniform bins, and the resulting categorical tokens are embedded with learned timestep, variate, and prediction-target embeddings. Each transformer block then alternates causal attention along the time axis (per variate) with unmasked attention along the variate axis (per timestep), so any predicted quantity can aggregate information from all variates at all previous timesteps. The model never receives an environment ID: it must identify which environment it is in, and how actions map to state changes, from the recent trajectory history, which the authors call capturing environment dynamics in-context. Predictions are produced as a next-step categorical distribution over the discretized state and reward variates, trained with cross-entropy, so all variates are predicted jointly rather than one after another.","core_discovery":"The paper's central claim is that low-dimensional sensor readings, despite looking different across robots, contain enough shared structure for large-scale pre-training—provided the model treats every scalar sensor or actuator channel as a token and infers each environment's dynamics from its own trajectory history, with no environment identifier given. Concretely, the authors report that pre-training TrajWorld on UniTraj and then fine-tuning it on downstream data yields substantial gains in transition prediction, a new state of the art for off-policy evaluation, and superior online performance of model predictive control, and they describe this as the first demonstrated transfer of world models across heterogeneous, complex control environments. The claim is carried jointly by the data and the architecture: UniTraj is scaled while keeping diversity across data sources, policy levels, and 80 embodiments, and TrajWorld's interleaved temporal and variate attention lets it identify how actions drive state changes from context, which is what the zero-shot experiments on cart-like environments are meant to show.","pith_inferences":["The paper's strongest evidence for true heterogeneity would be a pre-training set with every near-relative of the test robots removed; since DMC versions of walker, cheetah, and hopper bodies appear in UniTraj, part of the reported gain could reflect near-duplicate morphology, and an exclusion study would settle how much.","If the in-context mechanism is general, TrajWorld should handle entirely new sensor channels (touch, IMU, camera-derived features) as just new variate tokens, a testable prediction beyond anything the paper runs.","The learned continuity of bin orderings suggests a direct fix for the paper's acknowledged fixed-range limitation: extending bin boundaries beyond the observed data range should let the discretized model extrapolate like a regression model, which is a concrete experiment the authors flag but do not run.","A scalar-as-token world model points toward a shared dynamics prior across simulation, robotics, and teleoperation data, where the same pre-trained model seeds many downstream task-specific models; the authors gesture at this direction with multimodal vision-plus-proprioception pre-training."],"forward_implications":["A newly encountered environment needs far fewer collected trajectories: fine-tuning the pre-trained model beats training from scratch on the same data, and the gap widens as data grows scarcer.","Off-policy evaluation—scoring candidate policies by rolling them out in the learned model—becomes accurate enough to set a new state of the art on the DOPE benchmark, which makes offline policy selection cheaper and safer.","Model predictive control using the pre-trained world model improves online returns on fragile robots like Hopper and Walker2D, because more accurate predictions keep planned actions away from unsafe states.","The dataset-scale ablations show a consistent scaling trend: larger and more diverse pre-training data yields better downstream prediction and control, so the gains are expected to grow if UniTraj is extended.","Architecture choice is part of the result: the same pre-training data hurts an MLP ensemble (negative transfer) and transfers less through the one-dimensional sequence baseline TDM, showing that the two-axis attention design is load-bearing rather than incidental."],"supporting_citations":[{"why":"Trains a generalist transformer dynamics model (TDM) on 80 heterogeneous environments but only observes transfer to a simple cart-pole; it is the direct predecessor and baseline this paper extends to complex environments.","marker":"(Schubert et al., 2023)"},{"why":"TD-MPC2 contributes 30 environments of replay-buffer data to UniTraj and exemplifies the zero-padding approach to heterogeneous dimensions that TrajWorld replaces.","marker":"(Hansen et al., 2024)"},{"why":"ExORL contributes exploratory offline data (541K episodes) to UniTraj, supplying the distribution diversity the paper argues is essential for transfer.","marker":"(Yarats et al., 2022)"},{"why":"Modular RL supplies the customized multi-limb embodiments that round out UniTraj's morphology diversity and are used to collect additional pre-training trajectories.","marker":"(Huang et al., 2020)"},{"why":"DB-1 contributes expert trajectories across 58 environments, extending UniTraj's environment count and adding an expert-rollout data distribution.","marker":"(Wen et al., 2022)"},{"why":"D4RL provides the fifteen downstream datasets (HalfCheetah, Hopper, Walker2D at five policy levels) on which transition prediction and fine-tuning are evaluated.","marker":"(Fu et al., 2020)"},{"why":"DOPE is the off-policy evaluation benchmark that supplies the metrics (absolute error, rank correlation, Regret@k) where TrajWorld reports state-of-the-art results.","marker":"(Fu et al., 2021)"},{"why":"The MLP ensemble dynamics model is the main from-scratch baseline and the one that suffers negative transfer under pre-training, which the paper credits to architectural inflexibility.","marker":"(Chua et al., 2018)"},{"why":"ETM is the previous state of the art on the DOPE benchmark that TrajWorld's OPE results are compared against and claimed to surpass.","marker":"(Chen et al., 2024)"}],"fun_headline_variants":["World model pre-trained on 1M trajectories transfers to unseen robots","No environment IDs: one model masters 80 control domains","TrajWorld: tokenized sensors unify 80 heterogeneous environments","Cross-robot zero-shot world model trained on 1M trajectories","One world model for 80 environments, no labels needed"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The transfer claim depends on the three test robots being genuinely unseen to the model, even though very similar versions of those same robot bodies were included in its pre-training data, and the paper assumes the differences in state definition, action dimension, and dynamics make the transfer non-trivial.","fun_headline_variants_meta":{"raw":{"variants":["World model pre-trained on 1M trajectories transfers to unseen robots","No environment IDs: one model masters 80 control domains","TrajWorld: tokenized sensors unify 80 heterogeneous environments","Cross-robot zero-shot world model trained on 1M trajectories","One world model for 80 environments, no labels needed"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.001122,"raw_usage":{"total_tokens":4647,"prompt_tokens":904,"completion_tokens":3743,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":520,"completion_tokens_details":{"reasoning_tokens":3656}},"tokens_in":520,"tokens_out":3743,"duration_ms":27645,"temperature":1.0,"reasoning_tokens":3656,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-09T15:31:37.997043+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Re-train TrajWorld on UniTraj after deleting every pre-training environment that resembles the test robots—all DMC walker, cheetah, and hopper bodies, plus their Modular RL and TD-MPC2 variants—and then fine-tune on the Gym HalfCheetah, Hopper, and Walker2D D4RL datasets; if the gains over from-scratch training largely vanish, the claimed heterogeneous transfer is mostly near-duplicate morphology recall, while persistent gains would support the paper's interpretation. A complementary check is to rerun the zero-shot Cart-2-Pole and Cart-3-Pole rollouts with history context removed and confirm that prediction error rises, which would verify that in-context inference, not memorized per-environment statistics, carries the transfer.","supporting_citations":[],"review_version":1}