Learning to Be Uncertain: Pre-training World Models with Horizon-Calibrated Uncertainty
Authors: Shenghua Wan
Le Gan
De-Chuan Zhan
Venue: ICLR 2026 Poster (2026)
Keywords: World Models
Unsupervised Pre-training
Temporal Relative Embeddings
Horizon-Calibrated Uncertainty
Primary area: reinforcement learning
TL;DR: We propose a new world models pre-training framework that explicitly modeling the uncertainty grows with the temporal horizon to learn more robust dynamic representations for downstream RL tasks.

Abstract: Pre-training world models on large, action-free video datasets offers a promising path toward generalist agents, but a fundamental flaw undermines this paradigm. Prevailing methods train models to predict a single, deterministic future, an objective that is ill-posed for inherently stochastic environments where actions are unknown. We contend that a world model should instead learn a structured, probabilistic representation of the future where predictive uncertainty correctly scales with the temporal horizon. To achieve this, we introduce a pre-training framework, **H**orizon-c**A**librated
**U**ncertainty **W**orld **M**odel (HAUWM), built on a probabilistic ensemble that predicts frames at randomly sampled future horizons. The core of our method is a Horizon-Calibrated Uncertainty (HCU) loss, which explicitly shapes the latent space by encouraging predictive variance to grow as the model projects further into the future. This approach yields a latent dynamics model that is not only predictive but also equipped with a reliable measure of temporal confidence. When fine-tuned for downstream control, our pre-trained model significantly outperforms state-of-the-art methods across a diverse suite of benchmarks, including MetaWorld, the DeepMind Control Suite, and RoboDesk. These results highlight the critical role of structured uncertainty in robust decision-making.

Published as a conference paper at ICLR 2026
PRE-TRAINING WORLD MODELS WITH HORIZON-CALIBRATED UN- CERTAINTY
LEARNING TO BE UNCERTAIN:
Shenghua, Wan Le Gan", De-Chuan Z.a
School of Artificial Intelligence, Nanjing University, Nanjing, China
"National Key Laboratory for Novel Software Technology, Nanjing University, Nanjing, China "School of Computer Science and Technology. Nanjing University of Science and
Technology, Nanjing, China
wanshe 1 amda. n ju. edu. en, ganle@njusee edu. en, zhandc@n ji judduc edu. cn
ABSTRACT
Pre-training world models large, action-free video datasets offers promising path toward generalist agents, but a fundamental flaw undermines this paradigm. Prevailing methods train models to predict single, deterministic future, an ob- jective that ill-posed for inherently stochastic environments where actions are unknown. We contend that a world model should instead learn structured, prob- abilistic representation of the future where predictive uncertainty correctly scales with the temporal horizon. To achieve this, introduce pre-training framework, Corizon-cAAliratte Uncertainty World Model (HAUWM), built on a probabilis- tic ensemble that predicts frames at randomly sampled future horizons. The core of our method is Corizon-Calibrati Uncertainty (HCU) loss, which explicitly shapes the latent space by encouraging predictive variance grow the model projects further into the future. This approach yields a latent dynamics model that is not only predictive but also equipped with reliable measure temporal con- fidence. When fine-tuned for downstream control, our pre-trained model signifi- cantly outperforms state-of-the-ar art methods across a diverse suite of benchmarks, including the DeepMind Control Suite, Meta World, and RoboDesk. These results underscore the crucial role of structured uncertainty in informed decision-making.
INTRODUCTION
World models, which learn compressed representations of environmental dynamics, have become cornerstone of modern reinforcement learning (RI.), boosting agent sample efficiency and perfor- mance in complex domains (Ha & Schmidhuber 2018 Hafner al. [2020] Hu et al., 2023 Wan et al. 2024 Assran et al.) 2025; Agarwal et al. (025). A promising frontier is to preitrain these models on large, action-free video datasets (Seo et al. 2022). allowing them build foundational understanding of dynamics from passive observation that can be rapidly fine-tuned for downstream action.
control tasks. This paradigm enables versatile agents to learn without extensive task-specific inter-
However, a critical flaw undermines current approaches. Video-based pre-training frameworks typ- ically optimize for single objective: maximizing predictive accuracy under the assumption of deterministic future (Seo et al. 2022, Wu et al. 2023). This objective is fundamentally misaligned with the nature of action-free video. Without action labels, any observed future is merely one of many possible outcomes, as illustrated in Figure [ model trained on this data is erroneously compelled to predict one "correct" future from state that could unfold in myriad ways.
This deterministic bias creates a paradox. By forcing the model to predict a single future with high precision, pre-training penalizes any representation of environmental stochasticity- the very property that makes real-world dynamics so challenging. Consequently, models learn to suppress
"Corresponding Author

---

Published as a conference paper at ICLR 2026
0.39
Future
APV
without action
with action
Past
Future
HAUUM (oun))
0233
Uttemtaitty
Past
Future
0.20
Past
Future
0.19
8072
Pre-training Phase
Fine-tuning Phase
Seginent Length 03)
(a) Pre-training and Fine-tuning Paradigm
(b) Dynamics Uncertainty Video Segment Length
Figure 1: (a) Prevailing pre-training methods erroneously compel world model predict single deterministic outcome from action-free video, ignoring the multiple futures that could unfold from the same past. This deterministic bias creates fundamental conflict with the fine-tuning stage, in which an agent must learn to use actions to actively select a desired future from among these pos- sibilities. (b) The dynamics uncertainty of various methods pretrained on the video dataset evolves with the temporal interval between predicted frames (the experimental details are in appendix [..3) Notably, HAUWM consistently maintains well-calibrated long-term predictive uncertainty across extended prediction horizons.
ambiguity rather than represent it, fostering false certainty and losing the capacity for diverse forecasting-as quantitatively evidenced in fig. (b). where standard RSSM-based baselines (APV and Contex(WM),, even when equipped with ensemble heads, exhibit artificially low and nearly flat predictive uncertainty that fails grow with the prediction horizon, in stark contrast to the naturally increasing uncertainty of HAUWM. This limitation becomes a significant liability during fine-tuning, when the agent must navigate the action-conditioned dynamics that its pre-trained model was never equipped to handle.
This suggests an opportunity to reframe the objective action-free pre-training. We contend that instead of prioritizing predictive accuracy alone, more robust approach is to develop structured representation of temporal uncertainty. Specifically, world models should learn that predictive confi- dence decays with the horizon, principle that is not only under-emphasized but actively suppressed in current methods (Seo et al.]
2022; Wu et al.) 2023 Gao et al., 2025). leading to the observed pre-training frame- work named Horizon-cA Aibratted Uncertainty World Model (HAUWM) that redefines the model's method trains probabilistic semble to forecast states across variable, randomly sampled time horizons. The foundation of our approach is the Korizon-Calibrale Uncertainty (HCU) loss, which explicitly enforces that the ensemble's predictive variance grows with the prediction horizon.
uncertainty collapse shown in fig. (b). To realize this, introduce
relationship with the future. Instead of single-step prediction,
Our framework produces world models that not only predict future states but also quantify their own confidence across different timescales, a capability vital for robust decision- When fine- tuned on benchmarks from the DeepMind Control Suite (Tassa et al., 2018), MetaWorld (Yu et al.] 2020a), and RoboDesk (Kannan et al., 2021), our models significantly outperform state-of-the-art methods. This demonstrates that explicitly modeling temporal uncertainty not just beneficial essential for building truly generalizable agents. Our contributions are:
1. We identify and analyze the key limitation of deterministic prediction in action-free world model pre-training: it suppresses environmental stochasticity rather than representing it.
2. We propose a novel framework using variable-horizon prediction and introduce the Horizon-Calibrated Uncertainty (HCU) loss learn the relationship between time and pre- dictive uncertainty explicitly.
3. We provide comprehensive empirical evidence that our uncertainty- aware pre-training leads to superior performance across diverse and challenging control benchmarks.
2

---

Published as a conference paper at ICLR 2026
2 RELATED WORK
2.1 WORLD MODELS FOR REINFORCEMENT LEARNING
Model-based reinforcement learning (MBRL) seeks to improve the sample efficiency of agents by learning a model of the environment's dynamics (Janner et al. 2019 Yu et al. 2020b 2021). This learned model, often referred to as "world model", enables through simulated experience, thereby reducing the need for costly real-world interactions (Sutton, 1991 Kaiser et al. 2020)) Early approaches often operated in low-dimensional state spaces; the ad- vent of deep learning enabled the creation of world models that learn directly from high-dimensional sensory inputs, such as images (Janner et al.
agent plan or learn behaviors
2019 Yu al,, 2020b 2021 Hansen et al. 2024).
A seminal work by Ha & Schmidhuber (2018) demonstrated that compact, recurrent neural net- work could capture the essential dynamics of an environment, enabling an agent to solve tasks en- tirely within its dreamed latent space. This concept was significantly advanced by the development of latent dynamics models, most notably the Recurrent State-Space Model (RSSM) introduced in PlaNet (Hafner et al.) (011)) and its successors, the Dreamer family of agents (Hafner et al.) 22220) 2021 2025). These models learn a probabilistic representation of the environment state by combin- ing a deterministic recurrent pathway with stochastic latent variable. This structure allows them to model complex dynamics and uncertainty while remaining computationally tractable. The policy is then trained efficiently on imagined trajectories generated by rolling out the learned latent dynam- ies model. While highly effective, the success of these models hinges on their ability to accurately capture the actual environmental dynamics from scratch, process that can still require substantial in-domain experience, especially in visually complex or diverse settings (Fu et al. [202]] Wan et al., 2023). Our work leverages the power of these latent dynamics models but alleviates the burden of learning from scratch by pre-training on external data.
2.2 PRE-TRAINING ON VIDEO DATASETS
The paradigm of pre-training on large, unlabeled datasets has revolutionized fields like natural lan- guage processing (Devlin et al. 2019 Radford et al., 2018: Brown et al. 2020) and computer vi- sion (Murahari et al. 2002 He et al. 2022). where models first learn pencrrl-purpose representations that are later fine-tuned for specific tasks. Recently, this approach has gained significant traction in reinforcement learning, with a focus on using vast, readily available video data as a source of prior knowledge about the physical world (Seo al.) [002]] Wu al., 2023 Zhang et al.) [024] Wu al,, 2024 Gao et al. 2025). The core idea is that by observing wide range of dynamic scenes, model can learn a foundational understanding of physics, object interactions, and temporal coherence be- fore taking an action in target environment.
Several works have explored this direction. Action-Free Pre-training from Vidcos (APV) (Sco et al,, 2022) demonstrated that world model pre-trained on collection of action-free robotic videos could be effectively fine-tuned for new, unseen manipulation and locomotion tasks. To bridge the gap between the action-fred free pre-training phase and the introduced a stacked architecture that preserves the pre-trained representations while integrating ac- tion inputs. Other approaches, such as ContextWM (Wu al. 2023) and iVideoGPT (Wu al.) 2024]. have focused on disentangling static and contextual information from dynamic information in complex in- in-the-wild videos to improve knowledge transfer. Similarly, some research has explored using contrastive learning objectives on video data to learn useful representations for downstream control (Nair et al. 2022, Xiao et al. 2022). Despite their success, a common limitation unites many of these generative pre-training methods: they typically rely on deterministic, single-step predic- tion objective (i.e., predicting frame t + 1 from frame t). As we argue in this paper, this objective is fundamentally misaligned with the stochastic nature of action-free dynamics (Babaeizadeh et al.) 2018), where multiple plausible futures exist. This forces the model to average over all possibilities, often resulting in blurry predictions and failure to capture the multi-modal nature of the future. Our work directly addresses this limitation by designing a framework that not only predicts the future but also explicitly models and structures its uncertainty over time.
fine-tuning phase, APV
3

---

Published as conference paper at ICLR 2026
Pretraining Video Datasets
Finctuning Downstream Tasks
.B8
D
Alg
A
Pz
0:
A6
Temporal Embeddings 8
Temporal Interval
Segment Length K
ArE.
AC
Embbedding Length d,
Figure 2: Our pre-training framework (left) uses dynamics ensemble predict states at variable horizons, conditioned by temporal embedding, learning a structured representation of future un- certainty via our HCU loss. During finettuning (right), this pre-trained model is frozen to provide an action-agnostic foundation, while a new, lightweight module is trained to learn sask-specific, action- conditioned dynamics.
PRELIMINARIES
Recent work in model-based reinforcement learning has established a powerful two-phase paradigm: unsupervised pre- tre-training on large video datasets followed by task-specific fine-tuning (Seo al. 2022 Wu et al., 2023; 2024). This approach, exemplified by Action-Free Pre-training from Videos (APV) (Seo al.) 2022). aims learn generalizable world knowledge from passive, action-free videos, which can then be leveraged to significantly improve the sample efficiency and performance
on downstream control tasks (Seo et al., 20022 Wu et al. 2023 2024).
The first phase involves action-free pre-training. Here, latent video prediction model is trained on a diverse corpus of videos without corresponding action labels. This model is typically variant of latent dynamics model, such as Recurrent State-Space Model (RSSM) (Hafner et al. 2019), which learns to predict future frames by operating in a compressed latent space. consists of three core components: (i) representation model 2, qu q0(2e/2e11,0) ot) that encodes the current observation or into a latent state 2,, (ii) a latent transition model & ~ Po (2 (2111) that predicts the next state without access the observation, and (iii) image decoder ô ( 22) that reconstructs the observation from the latent state (Hafner et al. 2019 2020). The entire model trained to reconstruct the video sequence by optimizing the variational lower bound (ELBO).
The second phase is aetion-conditioned fine-tuning. To adapt the pre-trained model for a specific RI. task, agent actions must be incorporated. key challenge to integrate this information without overwriting the valuable priors learned during pre-traaeee traininiiig APV addresses this by introduc- ing a stacked latent prediction model, where new condinnooniiiinnall dynamics model is stacked on top of the frozen or gently-tune state 8; conditioned on both the agent's 's action a;-1 and the underlying action-fred -fre state 2,. This composite world model is then used to learn policy, typically via an actor-critic algorithm that trains on imagined trajectories generated by the model.
action-free model. This new model layer produces task-specific
METHODOLOGY
Our methodology is structured around a swo-phase framework as shown in fig. 2 (1) uncertainty- aware pre-training action-free video, followed by (2) efficient fine-tuning for downstream
4

---

Published as a conference paper at ICLR 2026
trol tasks. In the pre-training phase, we introduce variable-horizon prediction task, where an ensemble of dynamics models learns to forecast future states over randomly sampled time inter- vals. The comerstone of this phase is novel CoorioonCCalibrtte Uncertainty (HCU) loss, which explicitly encourages the predictive variance of the ensemble to increase with the temporal hori- zon. This procedure endows the model with structured understanding of temporal uncertainty. For fine-tuning, we freeze the pre-trained uncertainy--warr model and augment it with a lightweight, action-conditioned dynamics module, enabling rapid adaptation specific control objectives while retaining the foundational world knowledge.
4.1 UNCERTAINTY-AWARR PRE-TRAINING ON ACTION-FREE VIDEO SETS
Conventional video-based world model pre-training often erroneously assumes deterministic futures by optimizing single-step prediction accuracy (Seo et al.) [022 Wu et al.) 2023). However, video data inherently lacks explicit action labels and permits multiple plausible futures from identical states. This deterministic bias suppresses environmental stochasticity and compromises the model's ability to predict diverse futures a critical capability for robust decision-making To resolve this tension, our pre-training framework simultancously optimizes three key objectives: (I) learning meaningful visual representations from unlabeled video data, (2) maintaining predictive fidelity to observed trajectories, and (3) preserving the capacity represent multiple plausible futures.
We achieve these objectives using an ensemble of M independent dynamics heads. Each head ( models the stochastic transition future latent state St+k, conditioned on the current latent state 8, and temporal horizon Atr (for the consistency of expression, the variable-horizz & and temporal Formally, each ensemble head outputs the Gaussian embedding Atf, detailed in section 4.2
parameters for the future latent state:
Pg (#L+* 81, A = NV (ua, (8c, A()) ob (st. A(G))
To ensure effective extraction of observational information, we employ a widely used encoder- decoder architecture of Dreamer style (Hafner et al. 2020; Sco et al. 2022). The encoder qa (5t4k st, Atg. o++)) infers the latent state St+k from st, Atg. g. and future observation Ot+-k. Under POMDP assumptions, the image decoder P& acts latent states. The set of predicted means {no, Jim represents discrete distribution over plausible futures. For image reconstruction, we compute the ensemble mean ji++ ΣM No, (st, Atj) as S1+k and decode it as Ô (p($++)).
emission function, reconstructing observations from
We optimize model parameters by minimizing combined objective. The primary predictive loss is the negative variational lower bound:
Spee 3 DRI. [uo(xrs Ari, o.+1) (8 A[] ln po
(2)
where 3 weights the KL-divergence term. This loss forces the model to learn effective representa- tions while maintaining predictive accuracy.
To explicitly preserve future-state diversity, we introduce a
Uncertainty (HCU) loss. This loss encourages the ensemble's predictive dispersion to grow monotonically with the temporal horizon Atg. To implement this, we adopt widely-used approach based on the model
disagreement (Pathak et al. 2019 Sekar et al. 2020 Seyde et al. (021):
M
1
LHCU - M
Σ
(wo, (St, A B24R)
(3)
The HCU loss maximizes the model disagreement scaled by horizon length k. Minimizing Licu explicitly encodes the inductive bias that uncertainty should increase with prediction horizon.
To dynamically balance predictive accuracy against uncertainty representation, formulate dual optimization objective:
Lisal Lpeed XLUCU
Here, serves as an adaptive counterweight that maintains the essential tension between two com- peting objectives: minimizing reconstruction error (Cpre) to retain predictive fidelity and maximiz- ing uncertainty diversity (Chct) to capture environmental stochasticity simultaneously. This self- regulating equilibrium allows the model to maintain high reconstruction quality while progressively
5

---

Published as a conference paper at ICLR 2026
capturing necessary stochasticity avoiddi degenerate solutions where either extreme uncertainty or artificial determinism would degrade downstream performance.
4.2
VARIABLE-HORIZON PREDICTION AND RELATIVE TEMPORAL EMBEDDING
Conventional video-based world models predominantly focus on single-step next-frame prediction, neglecting the inherent multi-scale temporal relationships in unlabeled video data (Hong et al. Brooks et al., 2024). To make visual states exhibit predictable transitions across arbitrary time intervals when equipped with explicit temporal conditioning, we implement variable-horizon prediction during pre-training. For each training instance, we sample a random prediction horizon Uniform (1, 2..... Kax] and construct observation pairs (Ot, Ot++). ) To condition the dynamics model on the relative temporal relationship, we generate sinusoidal positional embeddings (Vaswani et al. 2017) for the entire video segment E E T xde, where T is the segment length. We then ex- tract the specific temporal embedding Atg E[k] R corresponding to horizon k. This relative encoding strategy provides two critical advantages: it inherently normalizes temporal relationships across videos of varying lengths by representing k as proportion of T, and preserves event dynam- ies by encoding the relative position k/T rather than absolute time steps.
2022
The embedding Aig thereby captures whether k represents a short or long interval relative to the video context, enabling the model to distinguish between rapid and gradual state transitions. This temporal conditioning input to each ensemble dynamics head, allowing unified modeling of transi- tions across diverse timescales. Unlike standard positional encodings in Transformers (Vaswani et al., 2017). which primarily index absolute sequence positions to preserve order, our relative temporal embedding A/S explicitly encodes the physical time gap itself. This forces each dynamics head to directly learn the transition &t -> St-k in single forward pass, rather than relying on recursive single-step rollouts. More critically, this design is tightly coupled with the certainty (HCU) loss: the embedding serves as a direct conditioning signal for the expected degree of stochasticity, enabling the HCU loss to reliably enforce monotonic variance growth with hori- zon length a mechanism absent in prior temporal conditioning approaches that target only mean accuracy. Thus, variable-horize prediction is not merely positional encoding variant, but essen- tial architectural choice deliberately engineered to make structured, horizon-dependent uncertainty emerge naturally and controllably. Implementation details are provided in appendix C.2
Un-
4.3 FINE-TUNING FOR DOWNSTREAM CONTROT.
Having established an uncertainty- aware world model through action-free video pre-training, we adapt this foundation to task-specific control via a structured fine-tuning phase. This transition preserves the rich dynamics and stochasticity learned during pre-training while integrating action- conditioned transitions essential for decision-making. We achieve this by augmenting the pre- trained model with a dedicated task-oriented pathway, following established setraining-finetuninn paradigms (Seo et al.) 2022 WW et al.) 2023) but critically extending them maintain uncertainty calibration. Specifically. we initialize the visual encoder and ensemble dynamics heads from pre- training, freezing their parameters to retain general visual and temporal representations, while intro- ducing a new cction-conditioned dynamics stream trained from scratch.
Our fine-tuning architecture processes two parallel latent streams. The pre-trained stream-unmodified from pre-training-reccinn only temporal embeddings and generates latent states & that capture action-agnostic environmental dynamies. To maintain compatibility with this stream's 's original design, we condition it on relative temporal embeddings Atr (as defined in Sec- tion (.2). injecting a Gaussian noise (a 0.01) enhance robustness against temporal discretiza- tion artifacts. Concurrently, a new action-conditionnd stream processes agent actions d and the pre-trained latent state & to produce task- -pecific predictions & via (5: 8; 4;11, &c). This dual-stream design ensures that the model leverages pre-trained uncertainty awareness while learn- ing real action effects. The composite latent state [&r; 5.] then drives policy optimization, with an auxiliary reward predictor Rg trained task-specific rewards to guide imagination- based planning.
Using this integrated world model, we execute standard model-based reinforcement learning: the agent learns actor-critic policy by generating imagined trajectories entirely within the calibrated latent space. Crucially, the pre-trained ensemble dynamics provide uncertainty- state pre- dictions during imagination, enabling the policy account for environmental stochasticity while

---

Published as a conference paper at ICLR 2026
Button Pres lapdown Wall
Drawer Open
Coffee Peal
Bessh
Doar Lock
509
Dval Reaur
2004
22:0
2000
200
190
186
Step (X10%)
Sep (x10")
Step (X10%)
Stop (x10*)
(x20
Dial
Lones
Dowe Open
Pub Green
Retrm
Retun
10:00
2009
1908
1960-
Skep (M105)
Sep (x10F)
Step (x10))
Stop (x110)
Step (x10
Choash Itre
Walker Burn
Horperr Hap
Queadrped Rum
Oper Siiee
Era RetuE
02r
(X
(K10%)
Step (X109)
(X10F)
Step (x20
[AUWW (eam)
APV
CortectWM
PRLAR
(VidarGrr
Performance comparison of HAUWM against state-of-the-ar baselines on a suite of down- stream manipulation and locomotion tasks. Solid curves represent the mean evaluation return across four random seeds, and shaded regions denote the 95% confidence interval.
Figure 3:
exploiting task-specific action knowledge from the fine-tuned stream. The overall algorithm is in ap- pendix B and full architectural details appear in appendix C.I
5 EXPERIMENTS
We conduct series of experiments to validate our proposed framework and answer the following key research questions:
1. Performance (RQ1): Does HAUWM lead to improved sample efficiency and final perfor- mance on downstream RL tasks compared to sfuee-of-the-e methods?
2. Ablation (RQ2): What are the relative contributions of the core components of HAUWM?
3. Analysis (RQ3): Does HAUWM successfully estimate the uncertainty in different training stage?
4. Robustness (RQ4): Can our pre-training world model generalize to diverse downstream learning paradigms?
5.1 EXPERIMENTAL SETUP
Benchmark Environments. We evaluate our method on diverse suite of challenging, vision- based continuous control benchmarks. These include several locomotion tasks from the Deed Control Suite (DMC) (Tassa et 2018). set of distinct robotic manipulation tasks from Meta- World (Yu et al. 2020a), and complex, long-horizon tasks from RoboDesk (Kannan et al. 2021) These environments test the agent's 's ability to learn complex motor skills, generalize across tasks, and operate from high-ddmensional pixel inputs. All observations are rendered as 64 x 64 x 3 images.
Pre-training Data. For the pre-training phase, we utilize a large, publicly available video dataset, Something-Something (Goyal et al. 2017), for HAUWM and all baselines. The key charac-
7

---

Published as a conference paper at ICLR 2026
teristic is that these videos action-free and sourced from domains that may differ visually and dynamically from the downstream fine-tuning tasks. More details are in appendix A
Baselines. We compare method, which we denote HAUWM, against several strong base-
lines to provide a comprehensive evaluation:
APV (Seo et al. 2022): A leading method that first establishes foundational framework for leveraging unlabeled videos through stacked latent prediction.
2023): A state of the art method that effectively handles complex in-the-wild videos by disentangling contextual and dynamic information.
ContextWM (Wu al.
PreLAR (Zhang al., 2024): A method that learns implicit action representations directly
from observation parrs during pre-training and fine-tuning with real actions.
iVideoGPT (Wu et al.) 2024]; A flexible framework that shows scalability through
transformer-bssed architectures trained on over one million manipulation trajectories.
5.2 DOWNSTREAM TASK PERFORMANCE (RQ1)
As illustrated in fig. 3 our method, HAUWM, achieves state-of-the-arr sample efficiency and final performance the majority of the tested bench- marks. The advantage is particularly pronounced in dynamically complex locomotion tasks such as Walker Run, and Hopper Hop. In these environ- ments, HAUWM consistently learns faster and con- verges to a higher final return than all baselines. We attribute this strong performance to our core contri- bution: by pre- pre-training model that explicitly repre- sents temporal uncertainty, the agent builds a more robust and realistic internal model of the world. This foundation allows it to adapt more effectively during fine-tuning where it must control its actions amidst environmental stochasticity.
Table 1: Ablation study of HAUWM.
Method
DMC
MetaWorld
RoboDesk
À = 10.0 0.67
0.77=0005
0.61 + 0.09
10-1 0.69
0.80=0.10
0.60 + 0.05
= 10-2 0,70+0.04
0.76++0.07
0.65 + 0.07
w/o HCU 0.6410.11
0.73
0.55 1 0.08
HAUWM 0.74.10.03 0.85
0.71 1 0.05
Conversely, on the Push Green task, HAUWM is outperformed by ContextWM. We hypothesize this is because the task's dynamics are more deterministic and narrowly goal-oriented, reducing the relative benefit of sophisticated uncertainty modeling. In such settings, methods that learn direct, implicit action-to-outcomme mapping may have an advantage. The somewhat larger performance variance of HAUWM on tasks like Dial Turn is an expected trade-off of our wosemmllebbsse archi- tecture, which, while enabling superior average performance, can naturally lead to greater diversity across independent training seeds.
5.3 ABLATION STUDIES (RQ2)
We conduct series of ablation studies to dissect the contributions of HAUWM's ' key components and validate our design choices, and record results in table and fig. 4 Each value is normalized against the random and expert returns and averaged across all correspond-
Normalized Retur
0.8-
ing tasks.
0.6
First, to verify the necessity of our core contribution, we trained variant without the Horizon-Calibrated Uncertainty loss (w/o HCU). This led to significant performance degradation, particu- larly on the DMC benchmark, confirming that explicitly modeling structured temporal uncertainty is critical for learning robust dy- namics representations. Next, we examined the sensitivity to the initial weight of HCU loss, A. The results indicate that performance suffers with both overly large (A = 10.0) and excessively small (A (A= 10-2) values. A large À forces the model to prioritize uncer- tainty at the expense of predictive fidelity, while a small À provides an insufficient training signal. Our final model employs moderate weight (A = 1.0), striking an effective balance between these competing objectives. The episodic returns of each task are in table 4
0.4
0.2
0.0
MetWoold RoboDesk
Kmecl1
Km
Kmsc=2
Kma=10
Figure K Kmax-
Ablation results on
8

---

Published as a conference paper at ICLR 2026
Es4
Egu
Es3
Eo
ground trath
ground truth
insagination
imsagination
dytamic
dynamics
f19
E21
E26
Eso
dyuamis 2
dyuamics 2
0.182
0.162
0.14
146
6,01K
0.011
dymamics
(a) Pre-training Phase (action-free)
(b) Fine-tuning Phase (action-conditooned)
Figure 5: Imagined future rollouts from three randomly selected dynamics heads during action- free pre-training (a) and cotion-conditioned fine-tuning (b). The model produces diverse, high- uncertainty futures when action is given, which converge to single, low-uncertainty outcome when conditioned on a specific action, with uncertainty values noted below each frame.
Finally, we analyzed the impact of the maximum prediction horizon, maxx- As shown in fig. 4 moderate value of Kmax 5 yields the best performance. small horizon like Kmax provides an insufficient range of temporal intervals to learn meaningful uncertainty curve. Con- versely, a large horizon like Kmax 10 can make the long-range prediction task overly complex, introducing noisy signals that may destabilize training. For additional ablation studies examining the scaling function and the number of ensemble dynamics heads M, refer to the appendix E.4 and appendix E.5
5.4 ANALYSIS OF LEARNED UNCERTAINTY (RQ3)
To address RQ3, perform qualitative analysis of the model's uncertainty estimation capabilities at different stages. During the action-free pre-training phase (left in fig. 5). HAUWM successfully preserves predictive diversity, demonstrating its capacity to represent environmental stochasticity. Different heads of the dynamics ensemble forecast distinct, plausible futures from the same ini- tial state; for instance, dynamics model I imagines the robotic arm lifting away from the surface, whereas model 3 predicts it continuing to push rightward. The model correctly learns that uncer- tainty increases with the prediction horizon, which aligns with our design of the HCU loss. As shown by the quantitative values below the frames, the uncertainty associated with predictions fur- ther the future (e.g., frame predicted at 134 from (2o) is consistently higher than for nearer frames (e.g., frame predicted at t43 from t4i), validating the efficacy of our HCU loss.
During fine-tuning (right in fig. 5). this learned agent actions appropriately constrain stochasticity. When conditioned verge on a single, consistent outcome. Concurrently, the associated uncertainty values drop dra- matically compared to the action-free stage. This behavior demonstrates that HAUWM learns well-calibrated representation of temporal uncertainty, fulfilling our primary motivation.
specific and same action sequence, the different dynamics models
5.5 PERFORMANCES UNDER DIVERSE DOWNSTREAM LEARNING PARADIGMS (RQ4)
A crucial test for pre-trained foundation model is its versatility across different downstream learn- ing paradigms, not just standard online RL. Therefore, we fine-tune our pre-trained model using established algorithms for imitation learning (II.) and offline RL on several benchmark tasks, with standard VMAII. (Rafailov et al., 20216) and LOMPO (Rafailov et al.) 2021a) algorithms, respec- tively (extra experimental details in appendix D The results, presented in table 2 demonstrate the broad applicability of our approach. the imitation learning setting. HAUWM achieves state- of-the-at performance, outperforming baselines in several tasks and performing competitively in Drawer Close. This suggests that the structured uncertainty learned during pre-training provides a robust foundation for mimicking experts, where understanding plausible future states is crucial.
Furthermore, HAUWM shows strong performance in the offline RL, where learning from a fixed dataset without online interaction is required. It significantly surpasses the baselines almost all

---

Published as a conference paper at ICLR 2026
Table Performance of HAUWM and baselines on diverse downstream learning paradigms, includ- ing imitation and offline RL. All scores are normalized returns, reported as the mean and standard deviation across 10 evaluation trajectories (details in appendix
Cheetah Run
Hopper Hop
Dial Tum
Drawer Close
Push Green
APV
0.78 0.11
0.83 0.07
0.55
0.13
0.89 +: 0.21
0.71 + 0.05
Contex(WM
0,91 + 0.04
0.85 + 0.05
0.69 = 0.09
0.98 + 0.03
0.69 = 0.05
Imitation
PreLAR
0.82 4: 0.07
0.81 + 0.11
0.72 = 0.06
0.93 + 0.05
0.69 = 0.05
HAUWM
0.94 0.03
0.89 = 0.05
0.77 = 0.06
0.97 0.02
0.63 = 0.06
APV
0.56 0.06
0.66 - 0.12
0.74 0.13
0
0.63 = 0.04
Offline RL ContextWM
0.61 + 0.06
0.69 + 0.06
0.69 + 0.07
0.70 0.05
0.66 = 0.03
PreLAR
0.71 + 0.08
0.65
= 0.05
0.69 = 0.02
0.68 0.03
0.69 = 0.09
HAUWM
0.81 + 0.13
0.71 = 0.08
0.67 = 0.12
0.74
0.70 = 0.04
tasks. While APV performs better the Dial Turn task, HAUWM's overall significant perfor- mance across both IL and offline RL highlights its versatility. Our uncertainty-aware pre-training strategy equips the model with a greater representational capacity under diverse downstream leamn- ing paradigms. This udaptability validates that HAUWM is not just pre-training solution for online RL but a powerful and versatile foundation model for wide array of decision-making tasks.
6 CONCLUSION
In this work, we address critical limitation in pre-training world models from action-free video: the reliance on deterministic prediction objectives, which are ioo-p for inherently stochastic environ- ments. We posit that models must instead explicitly learn temporal uncertainty-the principle that long-term futures are less predictable than near-term ones. Our proposed framework achieves this through a variable-horizon prediction task and novel Horizon-Calibrated Uncertainty (HCU) loss. The HCU loss incentivizes an ensemble's predictions to diverge systematically as the time horizon increases. Across multiple benchmarks, our method achieves wate-of-the-ane sample efficiency and final performance. This work demonstrates that making structured uncertainty a primary objective of pre-training is crucial for building capable and ceneral-purpone world models.
Future work will focus on developing more pappuutaattinlllllffiiieeee probabilistic models, such as diffusion models (Ho et al., 2020). and improving their scalability with large-scale Transformer architectures (Vaswani et al., 2017: Brown et al. (020). Scaling our approach to massive, in-the- wild video datasets remains promising direction for creating truly generalist agents.
ACKNOWLEDGMENTS
We thank Yucen Wang, Xingye Xu. and Jiabei Lv for their valuable discussions. This work was partially supported by the Young Scientists Fund of the National Natural Science Foundation of China (Ph.D) Candidate) under Grant No. 624B200197, the National Science and Technology Ma- jor Project under Grant No. 2022ZD01 14805, Collaborative Innovation Center of Novel Software Technology and Industrialization, and NSFC (62376118, 62006112, 62250069, 61921006).
REFERENCES
Niket Agarwal, Arslan Ali, Maciej Bala, Yogesh Balaji, Erik Barker, Tiffany Cai, Prithvijit Chat- topadhyay, Yongxin Chen, Yin Cui, Yifan Ding, Daniel Dworakowski. Jiaojiao Fan, Michele Fenzi, Francesco Ferroni, Sanja Fidler, Dieter Fox, songwei Ge. Yunhao Ge. Jinwei Gu. Sid- dharth Gururani, Ethan He, Jiahui Huang, Jacob Samuel Huffman, Pooya Jannaty, Jingyi Jin, Seung Wook Kim, Gergely Klár, Grace Lam, Shiyi Lan, Laura Leal-Taixé Anqi Li, Zhaoshuo Li, Chen-Hsuan Lin, Tsung-Yi Yi Lin, Huan Ling, Ming-Yu Liu, Xian Liu, Alice Luo, Qianli Ma, Hanzi Mao, Kaichun Mo, Arsalan Mousavian, Seungjun Nah, Sriharsha Niverty, David Page. De- spoina Paschalidou, Zeeshan Patel, Lindsey Pavao, Morteza Ramezanali, Fitsum Reda, Xiaowei Ren, Vasanth Rao Naik Sabavat, Ed Schmerling, Stella Shi, Bartosz Stefaniak, Shitao Tang, Lyne Tchapmi, Przemek Tredak, Wei-Cheng Tseng, Jibin Varghese, Hao Wang, Haoxiang Wang, Heng Wang, Ting-Chur Wang, Fangyin Wei, Xinyue Wei, Jay Zhangjie Wu, Jiashu Xu, Wei Yang.
10

---

Published as a conference paper at ICLR 2026
Lin Yen-Chen, Xiaohui Zeng, Yu Zeng. Jing Zhang, Qinsheng Zhang, Yuxuan Zhang, Qingqing Zhao, and Artur Zólkowski. Cosmos world foundation model platform for physical AL. ArXiv, foss/5011.0355.. 2025.
Mido Assran, Adrien Bardes, David Fan, Quentin Garrido, Russell Howes, Mojtaba Komeili, Matthew J. Muckley, Ammar Rizvi, Claire Roberts, Koustuv Sinha, Artem Zholus, Sergio Ar- naud, Abha Gejji, Ada Martin, Francois Robert Hogan, Daniel Dugas, Piotr Bojanowski. Vasil Khalidov, Patrick Labatut, Francisco Massa, Mare Szafraniec, Kapil Krishnakumar, Yong Li. Xi- aodong Ma, Sarath Chandar, Franziska Meier, Yann LeCun, Michael Rabbat, and Nicolas Ballas. V-JEPA 2: Self-supervised video models enable understanding, prediction and planning. ArXiv, (bss/200609988. 2025.
Mohammad Babacizadch, Chelsea Finn, Dumitru Erhan, Roy H. Campbell, and Sergey Levine. Stochastic variational video prediction. In 6th International Conference on Learning Represen- tations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceed- ings. OvenRevie 2018.
Tim Brooks, Bill Peebles. Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, et al. Video generation models as world simulators. OpenAl Blog, 18))1, 2024.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voos Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hisuan-Tier Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Pro- cessing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.). Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pp. 4171-4186. Association for Computational Linguistics, 2019.
Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi S. Jaakkola. Learning task informed abstractions. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 3480-3491, PMLR, 2021.
Shenyuan Gao, Siyuan Zhou, Yilun Du, Jun Zhang, and Chuang Gan. Adaworld: Learning adaptable 2025.
world models with latent actions. ArXiv,
Raghav Goyal, Samira Ebrahimi Kahou, Vincent Michalski, Joanna Materzynska, Susanne West- phal, Heuna Kim, Valentin Haenel, Ingo Fründ, Peter Yianilos, Moritz Mueller-Freitag, Florian Hoppe, Christian Thurau, Ingo Bax, and Roland Memisevic. The "something something" video database for learning and evaluating visual common sense. In IEEE International Conference on Computer Vision, ICCV 2017, Venice, Italy, October 22-29, 2017, pp. 5843-5851. IEEE Com- puter Society, 2017
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman. Nicolò Cesa-Bianchi, and Ro- man Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018 Montréal, Canada, pp. 2455-2467, 2018.
Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine
11

---

Published as a conference paper at ICLR 2026
Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings PMLR, 2019.
of Machine Learning Research, pp.
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learn- ing behaviors by latent imagination. In 8th International Conference on Learning Representa- tions, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. 2021.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy P. Lillicrap. Mastering diverse control tasks through world models. Nat., 64008059):647-6333 2025.
Nicklas Hansen, Hao Su, and Xiaolong Wang. TD-MPC2: scalable, robust world models for contin- uous control. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick. Masked autoencoders are scalable vision learners. In IEEE/CVF Conference Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 15979-15988. IEEE, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Mare Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
Wenyi Hong, Ming Ding, Wendi Zheng, Xinghan Liu, and Jie Tang. Cogvideo: Large-scale pre- training for text-to-video generation via transformers. arXiv preprint arXiv:2205. 15868, 2022.
Anthony Hu, Lloyd Russell, Hudson Yeo, Zak Murez, George Fedoseev, Alex Kendall, Jamie Shot- ton, and Gianluca Corrado. GAIA-I: A generative world model for autonomous driving. ArXiv, abs/2309. 170800 17080, 2023.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-BBu,, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Process- ing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 22498-12500 2019.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Harini Kannan, Danijar Hafner, Chelsea Finn, and Dumitru Erhan. Robodesk: A multi-takk rein- forcement learning benchmark. htttp:://githbb 2021.
Vishvak Murahari, Dhruv Batra, Devi Parikh, and Abhishek Das. Large-scale pretraining for visual dialog: A simple ofate-of-the--rt baseline. In Andrea Vedaldi, Horst Bischof, Thomas Brox, and Jan-Michael Frahm (eds.), Computer Vision - ECCV 2020 - 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XVIII, volume 12363 of Lecture Notes in Computer Science, pp. 332-352. Springer, 2020.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3M: A uni- versal visual representation for robot manipulation. Karen Liu, Dana Kulic, and Jeffrey lch- nowski (eds.), Conference on Robot Learning, CoRI. 2022, 14-18 December 2022, Auckland, New Zealand, volume 205 of Proceedings of Machine Learning Research, pp. 892-909. PMLR, 2022.
Deepak Pathak, Dhiraj Gandhi, and Abhinav Kumar Gupta. Self-supervised exploration via dis- agreement. ArXiv, abs/1 990.00116, 2019.
12

---

Published as a conference paper at ICLR 2026
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language under- standing by generative pre- training. 2018.
Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Offline reinforcement learning from images with latent space models. In Ali Jadbabaie. John Lygeros, George J. Pappas, Pablo A. Parrilo, Benjamin Recht, Claire J. Tomlin, and Melanie N. Zeilinger (eds.), Proceedings of the 3rd Annual Conference on Learning for Dynamics and Control, LADC 2021, 7-8 June 2021, I, Virtual Event, Switzerland, volume 144 of Proceedings of Machine Learning Research, pp. 1154-1168. PMLR, 2021a.
Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Visual adversarial imita- tion learning using variational models. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 3016-3028, 2021b.
Ramanan Sekar, Oleh Rybkin, Kostas Danilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In Proceedings of the 37th International Conference on Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 3583-8592 PMLR, 2020.
Younggyo Seo, Kimin Lee, Stephen James, and Pieter Abbeel. Reinforcement learning with action-free pre-training from videos. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, 2C21, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 19561-19579, PMLR, 2022.
Tim Seyde, Wilko Schwarting, Sertac Karaman, and Daniela Rus. Learning to plan optimistically: Uncertainty-guude deep exploration via latent model ensembles. In Aleksandra Faust, David Hsu, and Gerhard Neumann (eds.), Conference on Robot Learning. 8-11 November 1021, London, UK, volume 164 of Proceedings of Machine Learning Research, pp. 1156- 1167. PMLR, 2021.
Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting, SIGART Bull., 164):160-163, 1991.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy P. Lillicrap, and Martin A. Riedmiller. DeepMind control suite. ArXiv, tts//801..0099 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp- 5998 2008-608, 2017.
Shenghua Wan, Yueen Wang, Minghao Shao, Ruying Chen, and De-Chuan Zhan. Semail: Elimi- nating distractors in visual imitation via separated models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning. ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 35426 35443. PMLR, 2023.
Shenghua Wan, Hai- Hai-Hang Sun, Le Gan, and De-Chuan Zhan. Moser: learning sensory policy for task-specific viewpoint via view-conditional world model. In Proceedings of the Thirdy International Joint Conference on Artificial Intelligence, pp. 5046-5054, 2024.
Jialong Wu, Haoyu Ma, Chaoyi Deng, and Mingsheng Long. Pre-training contextualized world models with in-the-will wild videos for reinforcement learning. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Anmual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
13

---

Published as a conference paper at ICLR 2026
Jialong Wu, Shaofeng Yin, Ningya Feng. Xu He, Dong Li, Jianye Hao, and Mingsheng Long. Interactive videogpts are scalable world models. Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang (eds.), Ad- vances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
ivideogpt:
Tete Xiao, Lija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. ArXiv, abs/2203.06173, 2022.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100, PMLR, 2020a.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y. Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: model-based offline policy optimization. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Floring Balcan, and Hsuan-Tien Lin (cds.),.vvaneee in Neural Information Processing Systems 33: Annual Conference on Neural Information Pro- cessing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020b.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. COMBO: conservative offline model-based policy optimization. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Process- ing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 28954-28967, 2021.
Lixuan Zhang, Meina Kan, Shiguang Shan, and Xilin Chen. Prelar: World model pre- pre-training with learnable action representation. In Ales Leonardis, Elisa Ricci, Stefan Roth, Olga Russakovsky, Torsten Sattler, and Gül Varol (eds.), Computer Vision ECCV 2024- 18th European Conference, Milan, Italy September 29-October 4, 2024, Proceedings, Part XXIII, volume 15081 of Lecture Notes in Computer Science, pp. 185-201. Springer, 2024.
THE USE OF LARGE LANGUAGE MODELS
In preparing this work, we utilized Large Language Model (LLM) solely for limited, auxiliary pur- poses. Specifically, the LLM was employed to assist with refining the English phrasing of certain paragraphs, correcting grammatical errors, and, in the Related Work section, for generating ideas and searching for potentially relevant literature. The model did not contribute to the core research ideation, mathematical formulation, theoretical analysis, algorithm development, experimental de- sign, or interpretation of results. All authors take full responsibility for the entire content of this work, including any text generated with the aid of the L.M..
A ENVIROMENTS AND DATASETS
We evaluate our method across three representative benchmarks that collectively span diverse do- mains of visual reinforcement learning. The DeepMind Control Suite (DMC) (Tassa et al., 2018) provides standardized testbed for locomotion tasks with physically realistic dynamics, featur- ing continuous control challenges that require precise motor coordination and long-term planning. Specifically, we examine walker_rum, which tests bipedal locomotion stability; cheetah_run run, evalu- ating high-speed quadrupedal movement; hopper assessing single-legged balance and propul- sion; and quadruped run, measuring complex multi-joint coordination for agile movement. Meta- World (Yu al. (2020a) offers a comprehensive benchmark of robotic manipulation tasks with vary- ing geometric and physical properties, enabling rigorous evaluation of transfer learning capabiliiiess across structurally different but conceptually related challenges. Our evaluation includes surn, requiring precise rotational control; lever_pull. testing force application at mechanical advantage points; door door_open, .open, evaluating sequential manipulation of articulated objects; and drawer_closs as- sessing ine-grained spatial reasoning for object containment. Finally, RoboDesk (Kannan et al. 202I) presents real-world inspired manipulation scenarios with increased visual complexity and longer horizons, where push green challenges object relocation with color-based discrimination,
14

---

Published as a conference paper at ICLR 2026
Bution Press Topdown
Drawer Open
Coffee Push
Door
Dial Turn
Lever Pull
Door Open
Drawer Close
Push Green
Cheetah Run
Walker Run
Hopper Hop
Quadraped Run
Open Slide
Figure 6: The observation example of each task in the environments: Meta-World (top left), DMC (bottom left), and RoboDesk (right).
while open_slide demands precise linear actuation of constrained mechanisms. We display the ob- servation example of each task in fig. 6
For pre-training, we leverage the Something- Something V2 (SSv2) dataset (Goyal al. 2017
large-scale collection of 220,000 short video clips capturing humans performing basic physical in- teractions with everyday objects under natural lighting conditions. This dataset provides diverse action demonstrations -nniruudionsnncuuiiiin "Putting something on surface", "Moving something up", and "Covering something with something"-thh despite originating from non-robotic contexts, contain fundamental physical dynamics transferable to robotic manipulation. By utilizing this in-the-will video resource, our approach demonstrates that meaningful world models can emerge from general human-object interactions rather than domain-specttn robotic data, significantly reducing the need for task specific pre- vre-trinning collections while maintaining strong transfer performance to down- stream robotic control tasks.
B PSEUDO CODES
We list the pseudo-codes of HAUWM in algorithm
C MODEL DETAILS
C.1 MODEL ARCHITECTURE
Our model architecture builds upon the Recurrent State-Space Model (RSSM) framework (Hafner al. 2019 2020) and composed of several key modules: shared image encoder, corresponding decoder, and dynamics models for the pre-training and fine-tuning phases.
Encoder and Decoder We use a ResNet-basse based architecture for both the visual encoder and de- coder, enabling the processing of 64 x 64 x 3 image observations.
15

---

Published as a conference paper at ICLR 2026
Algorithm Uncertainty-Aware Pre-training and Fine-tuning
1: Input: Action-free video dataset Viieo- task interaction budget
2: Initialize: World model parameters 0, o; adaptive weight A; task replay buffer Dinkk +-
3:
4: Phase 1: Uncertainty-aware Pre-training with Dual Optimization
5: for each pre-training step do
6: Sample video clip from Dvideo and a horizon & Uniform(1,.... Kmax
7: Get temporal embedding Arg for horizon k.
8:
Compute ensemble predictions { (St, Atg))s=1 and model uncertainty.
9:
Calculate total loss Lwas = pedd + ACHCU-
10:
Update world model parameters by minimizing Liaal-
11: Update À via dual optimization to balance Lpred and LucU-
end for
13:
Phase 2: Fine-tuning for Downstream Control
15: Initialize action-condinn conditionnd dynamics Pv and policy mT. Freeze 0, ø. for each environment step up to budget do // World Model Adaptation and Policy Learning
17:
Sample trajectories from Dankk to update Py and reward predictor Rg.
19:
Imagine trajectories using the policy ey and the composite world model (pe, p Ro).
20:
Update policy मę on imagined trajectories.
21:
// Environment Interaction
22:
Execute action ag ~ TT(()) observe (o1+1, (+41) and add to D
23:
end for
24:
return Trained policy
The encoder a convolutional neural network (CNN) with 3 residual blocks. It processes input image to produce a www-dimeniional embedding. This embedding serves as the primary input for the dynamics model.
Encoder:
Decoder:
The decoder is transposed convolutional neural network, also with 3 residual blocks. It reconstructs the image observation 6ş from a given latent state
Pre-training During the action-free pre-training phase, the world model learns environmental dy- namics using variable-horizon prediction task.
Dynamics Model: We use an ensemble of M 5 dynamics heads built upon an EnsembleRSSM The model features deterministic state of size 200 and stochastic state of size 30. It takes the image embedding and a sinusoidal temporal embedding A/E to predict future latent state 8
For downstream tasks, we freeze the pre-trained components and introduce new action-conditioned module.
Fine-tuning
Dynamics Model: The fine-tuning architecture employs dual-stream design.
The frozen, pre-trained En.sembl 1 eRRSM generates an action-agnostic latent state &
conditioned on the image embedding and fixed temporal embedding (Atg..). )).
new, lightweight contiincondiiooned RSSM, trained from scratch, produces task- specific latent state 8, from the agent's 's action 0:-1 and the action-agnostic state š,
The policy and value functions are MLPs that take the composite latent state [s;; &,] input. The reward predictor is separate MLP that estimates the task reward from the task-specific state š, Both consist of 2 hidden layers with 400 units each and employ the ELU activation function.
Actor-Critic and Reward Predictor:
16

---

Temporal Embedding: To inform the model of the time gap, generate relative tem- poral embedding AIy. for each sampled gap k. This is achieved by creating sinusoidal positional encoding matrix for the full length of the source video; the embedding for spe- Published as a conference paper at ICLR 2026
C.2 TRAINING DETAILS
We implement our framework in PyTorch and use the Adam optimizer for all model components. All experiments are run on NVIDIA A6000 GPUs with 16-bit precision (FP16) to accelerate training and take about 2000 GPU hours.
Pre-training Phase The goal of the pre-training phase is to learn a robust, uncertainty- aware world model from action-free video data.
Data and Batching: We pre-train our model on the
(SSv2)
dataset (Goyal et al. 2017). To implement our hcriablehorrzon prediction strategy, we process each video into training sequences. For cach sample, extract sequence of frames with length L. where the temporal gap & between any two consecutive frames is randomly sampled from a uniform distribution, ~ Uniform (1.... Kmar). )- We set the maximum horizon to Kmaz = 5. This creates batches where the model must predict across varied and non-contiguous time steps.
cific gap k is then looked up from this matrix. This embedding is supplied as input to the dynamics model, analogous to an action.
Optimization: The world model is trained for 600k gradient steps. We use the Adam optimizer with learning rate of 1 x 10 and a batch size of 32. The adaptive weight for the HCU loss is also optimized with Adam, using a learning rate of 10-1
Fine-tuning Phase In the fine-tuning phase, the pre-trained model is adapted for downstream control tasks.
Initialization: We initialize the visual encoder and the action-agnostic dynamics model with the weights from the pre-training phase and freeze them. The new action-conditioned dynamics model and the actor-critic policy are trained from scratch.
Training Loop: The agent first collects an initial dataset 5,000 steps by executing a random policy. It then enters the main training loop, alternating between environment interaction to collect new data into a replay buffer and updating the model parameters. The agent is trained for total of million steps in the DMC environment and 250 thousand steps in the Meta-World and RoboDesk environments.
Temporal Conditioning: To maintain compatibility with the pre-trained components, the frozen action-agnostic dynamics stream conditioned on fixed temporal embedding for k = L, with minor Gaussian noise (a = 0.01) added to enhance robustness.
Optimization:
The new world model components (action-conditionnd stream and reward predictor) are trained with the Adam optimizer using a learning rate of x 10-4, The actor and critic components are trained with a learning rate of 8 x 10 -5
C.3 HYPER PARAMETERS
We summarize the key hyperparameters used for the pre-training and fine-tuning phases in table
D DETAILED BASELINES
APV (Action-Free Pre-training from Videos) (Seo et al.] 2022) introduces a two-phase paradigm that first pre-trains world model on unlabeled videos without action information, then adapts to downstream reinforcement learning tasks through fine-tuning, During pre-training, APV learns a generative video prediction model by optimizing a standard single-step predictive objective on action-free video sequences, effectively capturing visual dynamics without requiring action labels. To bridge the gap between action-free pre-training and condiiiond reinforcement learning, APV employs stacked architecture that preserves pre-trained representations while integrating
17

---

Published as a conference paper at ICLR 2026
Phase
Hyperparameter
Value
Batch Size
32
Sequence Length
len(wideo)
Max Prediction Horizon (max))
Total Gradient Steps
600,000
Pre-training
Ensemble Size (M)
KL Loss Balance
0.8
Initial HCU Weight (A)
1.0
World Model Learning Rate
x 10 -4
HCU À Learning Rate
10
Optimizer
Adam
Batch Size
32
Sequence Length
50
Replay Buffer Capacity
10 10,000
Initial Random Steps (Prefill)
Total Environment Steps
DMC
x 10
Fine-tuning
Meta-Worla & RoboDesk
2.5 x 105
Imagination Horizon Discount Factor (7)
15
0.99
GAE Parameter (AGAE)
0.95
World Model Learning Rate
x 10 44
Actor-Critic Learning Rate
x 10
Optimizer
Adam
Table 3: Core hyperparameters for HAUWM pre-training and fine-tuning.
action inputs during fine-tuning. This approach significantly improves sample efficiency in down- stream tasks by transferring knowledge from diverse video datasets to specific control problems, establishing an important foundation for world model pre-training that subsequent methods have built upon.
ContextWM (Comtextualized World Models) (Wu et al. 2023) specifically addresses the challenges of pre-training world models on complex, diverse in-the-wild videos that contain numerous contex- tual factors irrelevant to task dynamics. Rather than treating all visual information equally, Con- textWM explicitly disentangles static contextual elements (such as backgrounds and object appear- ances) from dynamic state transitions through a specialized architecture. During pre-training. lears separate these components, preserving only the Synammics-reeeann information while fil- tering out task-irrelevant variations. This contextual disentanglement enables more effective knowl- edge transfer when fine-tuning on downstream tasks, particularly when pre-training data comes from visually diverse sources that differ significantly from target environments. By focusing on the core dynamics rather than superficial visual variations, ContextWM achieves superior performance wide range of control tasks despite the domain gap between pre-training and target environments.
PreLAR (World Model Pre-training with Learnabl Action Representation) (Zhang al. 2024) addresses the fundamental limitation of previous approaches by incorporating conttiacon learning during the pre-training phase, even when explicit action labels are unavailable. Instead of treating pre-training as a pure video prediction task, PreLAR extracts implicit action representations by encoding observations from two consecutive time steps, effectively learning the causal transi- tions between states from unlabeled videos. The method introduces action-state consistency loss to ensure these learned representations align with actual dynamics, creating more seamless transition between pre-training and fine-tuning phases. This innovative approach eliminates the chitectural discrepancy between pre-training (video prediction) and fine-tuning world model), resulting in more effective knowledge transfer and significantly improved sample
18

---

Published as a conference paper at ICLR 2026
efficiency across various visual control tasks, particularly in robotic manipulation scenarios where action understanding is critical.
iVideoGPT (Interactive VideoGPTs) (Wu et al., 2024) represents significant advancement in world model scalability by leveraging transformer architectures to handle long video sequences while maintaining interactivity. Unlike RNN-based approaches that process sequences step-by- step, iVideoGPT employs a token-based video prediction framework that processes entire video sequences in parallel, enabling efficient training on massive datasets comprising over one million trajectories from both robotic and human manipulation. The model uses a cross-entropy loss to pre- dict subsequent video tokens autoregressively. creating versatile foundation that can be adapted to various downstream tasks through domain-specifh fine-tuning, Crucially, iViddeoPTT demonstrates that world models can be both highly scalable (processing long sequences efficiently) and interac- tive (supporting action conditioning for control), overcoming fundamental limitation in previous approaches that often sacrificed one capability for the other. This dual capability makes it partic- ularly effective for complex manipulation tasks requiring both long-term planning and responsive interaction.
E EXTRA EXPERIMENTAL DETAILS
E.1 DIVERSE DOWNSTREAM LEARNING PARADIGMS
To rigorously assess the versatility of our pre-trained world model, we fine-tune it on downstream tasks using two distinct and challenging learning paradigms: model-based imitation learning (IL) and offline reinforcement learning (RL). This tests the model's ability to serve as a robust foundation not only for online exploration but also for learning from expert demonstrations and static, pre- collected datasets.
For the imitation learning paradigm, we adapt the Visual Model-based Adversarial Imitation Learn- ing (VMAIL) framework (Rafailov et al., (2221b. Instead of training a variational dynamics model from scratch as in the original work, we leverage the fine-tuned, action-conditioned world models (APV (Seo et al., 2022). ContextWM (Wu et al.] (023), HAUWM) as the latent environment for policy optimization. Expert demonstrations are first encoded into latent trajectories using the frozen pre-trained encoder. A discriminator, Da, is then trained to distinguish these latent expert trajecto- ries from those generated by the agent's 's policy, Ty. during imagined rollouts within our world model. The policy and discriminator are optimized via a minimax game where the discriminator's 's output provides a dense reward signal. This process minimizes the Jensen-Shannon divergence between the expert and policy visitation distributions in the latent space, governed by the objective:
min max T De
Erg-Del loll-loggg log D (s. a)]
(5)
+ E log(1 (w
By training the policy entirely on imagined rollouts, this approach effectively capitalizes on the pre-trained dynamics to achieve high sample efficiency in the target imitation task.
For the offline RL paradigm, we address the critical challenge of learning from a fixed dataset where distributional shift can lead to catastrophic failures. We integrate our framework with the princi- ples of Latent Offline Model-Based Policy Optimization (LOMPO) (Rafailov et al.] 2021a). Our uncertainty-awar pre-trained world model is uniquely suited for this paradigm. The ensemble dy- namics, a core component of HAUWM, provides natural and robust mechanism for quantifying epistemic uncertainty through model disagreement. For APV (Seo et al. 2022) and ContextWM (Wu) et al. 2023), we also equipped them with ensemble dynamics models. To mitigate the risk of the pol- icy exploiting out-of-distributtio states where the model is inaccurate, we incorporate a pessimism principle directly into the reward function for imagined trajectories. Specifically, the reward is pe- nalized by the variance of the predictions across the dynamics ensemble heads, discouraging the policy from venturing into uncertain regions of the state space. The uncertainty-penalizee reward F(st, de) is formulated as:
F(s,, a.) = r(st, az) - Xu(s,, as)
(6)
Here, u(s,, a.) is the uncertainty measured as the variance of the ensemble's 's predictive distributions, and À is a hyperparameter controlling the degree of conservatism. This ensures that the policy learss
19

---

Published as a conference paper at ICLR 2026
Table 4: The absolute episodic return values for the random agent, the expert, and all ablated versions on all tasks averaged by four seeds.
Task
Random
Expert
x= 10.0
A= 10
== 10- w/o HCU
HAUWM
Dial Turn
687 + 189 4232 270
3417 177
1-23.1:55 355 3381 1 248 3258 709 3225 1: 281
Lever Pull
239 + 215 4 3838 +
3978
#778:466
3791
= 327
3723
935 4022 = 288
Door Open
277 + 158 4637 478 3634
3765
136 3591 = 305
3609
872
3811 = 502
Drawer Close
221 + 198 4895 + 578
3820 + 234
3960
+ 467
3773 = 327
3707 + 935
4736 = 393
Cheetah Run
71 38
821 + 44
574
589 45
596 = 30
544 75
626 = 23
Walker
120
664:100
+
688 = 32
81
721
Hopper Hop
25 12
478 : 101
329
338 27
342 + 18
310 45
360 - 14
Quadruped Run
+ 19
795 + 96
544 99
559 + 46
567 + 30
514 + 76
597 + 23
Push Green
5
378 77
235 33
232 + 18
250 26
213 33
272 = 18
Open Slide
7
111
109 9
118 +
15
128
effective behaviors while remaining confined to the regions of the state space well-supported by the offline data. We choose the "medexp" dataset for each task as the offline dataset, a combination of medium replay and expert demonstrations.
E.2 DETAILED EPISODIC RETURNS
To facilitate clear and intuitive comparison, we report all performance metries in the ablation (RQ2) and extension (RQ4) experiments as normalized returns. These returns are normalized for each task based on corresponding random and expert performance. The absolute return values for the random agent, the expert, and all ablated versions are detailed in table 4
E.3 UNCERTAINTY COLLAPSE IN STANDARD RSSM-BASED BASELINES
To provide quantitative evidence that prevailing methods suffer from "false certainty", we conducted the following controlled experiment. We retrained the APV (Seo et al. (2022) and ContextWM (Wul ctt al. 2023) and replaced their single deterministic dynamics heads with an ensemble of M 7 independent Gaussian heads with the same architecture used in HAUWM, which ensures a strictly fair comparison among all methods. We then sampled 5,000 random frames from the Something- Something-v' v2 pre-training dataset (the same dataset used for all methods in study). For each frame, we performed open-loop rollouts with horizons k (1,2,3.4.5) steps and measured the dynamics uncertainty:
M M Σ
Uncertainty(h
(na, Atz) - )&+x)²
The results shown in fig. I (b) (orange and green for APV and ContextWM. yellow curve for HAUWM). Both baselines exhibit near-zero growth in disagreement as the horizon in- creases, confirming that the standard single-step ELBO objective actively suppresses representation of stochasticity and induces deterministic, low-variance predictions even when the model architec- ture capable of expressing uncertainty (i.c., when given ensemble heads). contrast, HAUWM shows the desired monotonic increase, validating that the HCU loss is directly responsible for the calibrated growth of predictive uncertainty over time. This controlled experiment provides clear em- pirical evidence of the claimed defect in baseline methods and rigorously justifies the introduction of the HCU loss.
E.4 SCALING FUNCTION IN THE HORIZON-CALIBRANN UNCERTAINTY LOSS
We adopt linear scaling & in the HCU loss because it most directly encodes the physically intuitive prior that predictive uncertainty accumulates linearly with time-a simple, stable, and effective inductive bias for real-world dynamics. To confirm this design choice, we compared it against two
alternatives: exponential scaling & and periodic scaling sin
k Rmne ) where Kmax 5 is the
20

---

Published as a conference paper at ICLR 2026
maximum horizon used in pre-training. Performance four representative tasks (mean = std over four seeds) reported in table 5
Table 5: Ablation on the scaling function in the HCU loss (eq.
3
Linear scaling achieves the best
performance and training stability across all tasks.
Scaling function Cheetah Run Hopper Hop
Dial Turn
Drawer Close
k (lincar, ours)
626 + 23
360 14
3225 # 281
4736 + 393
(exponential)
587 42
341
3102 = 156
4265 349
|sin(:: 2m
421 # 87
104 + 55
2106 # 333
3089 + 796
Linear scaling consistently outperforms both alternatives. The exponential variant causes uncer- tainty to grow too aggressively, which disproportionately penalizes long-horizon predictions and destabilizes optimization, particularly in tasks with extended temporal dependencies. The periodic variant violates the monotonicity assumption and produces oscillating uncertainty signals that con- fuse the dynamics learning. These results validate linear scaling as the optimal balance between physical plausibility and training robustness.
E.5 THE EFFECT OF THE NUMBER OF ENSEMBLE DYNAMICS HEADS
The ensemble size M controls the trade-off between predictive diversity and mean accuracy. When M the disagreement term in eq. (3) vanishes, making the model equivalent to the "w/o HCU" ablation in table of the main paper. We evaluated M {1,3, 7,9} on the same four tasks (mean + std over four seeds). Results are shown in table
Table 6: Ablation ensemble size M. Performance peaks at M 7, confirming that moderate ensemble diversity yields the best downstream control performance.
Number heads Cheetah Run Hopper Hop
Dial Tum
Drawer Close
M=9
579 28
342 22
3377 298
4522 299
M=7 (ours)
626
340
3225 281
4736 + 393
M=5
589 = 43
379
3113 + 257
4411 # 330
M=3
550 + 94
327
3177 294
4087 + 498
M=1 (w/o HCU)
544 +
310
3258 709
3707 935
Performance improves steadily from M 1 to M = 7 as greater ensemble diversity better captures
multi-modal futures. Beyond M marginal returns diminish and training becomes slightly noisier without meaningful gains, confirming M = as a sweet spot that balances representational capacity, calibration quality, and computational efficiency.
21


Reviews:
[{"review_id": "Gq654H3KCC", "reviewer": "Reviewer_mHFR", "summary": "The paper introduces the Horizon-cAlibrated Uncertainty World Model (HAUWM) framework to address the limitation of \"false certainty\" in world models pre-trained on action-free video, a flaw stemming from deterministic prediction objectives that suppress environmental stochasticity. HAUWM employs an ensemble of probabilistic dynamics heads and a novel Horizon-Calibrated Uncertainty (HCU) loss that explicitly enforces predictive uncertainty to scale monotonically with the temporal horizon. The authors claim that HAUWM significantly outperforms state-of-the-art baselines across diverse downstream control benchmarks (DMC, MetaWorld, RoboDesk) and demonstrates versatility in Imitation Learning and Offline RL. Qualitative evidence supports the model's well-calibrated uncertainty: the ensemble predicts diverse, high-uncertainty futures during pre-training, which converge to a low-uncertainty outcome when conditioned on a specific action during fine-tuning.", "questions": "1.\tThe authors claim that prevailing methods predict a 'single, deterministic future.' However, established baselines such as APV, ContextWM, and PreLAR are founded on DreamerV2's Recurrent State Space Model (RSSM), which explicitly models state uncertainty using a Gaussian distribution. Given this, what is the specific technical advantage of employing the proposed multi-head dynamics prediction ensemble over the existing Gaussian uncertainty modeling within the standard RSSM latent state? Please provide a detailed analysis of this distinction and how the ensemble mitigates the alleged 'deterministic bias.'\n\n2.\tIn the image reconstruction phase, the predicted future latent state $s_{t+k}$ is set to the ensemble mean, $\\overline{\\mu}_{t+k} = \\frac{1}{M}\\sum_{i=1}^{M}\\mu_{\\theta_{i}}(s_{t}, \\Delta t_{k}^{e})$. This mean $\\overline{\\mu}_{t+k}$ is then decoded to reconstruct the observation $\\hat{o}_{t+k}$. Using the ensemble mean $\\overline{\\mu}_{t+k}$ directly, instead of sampling $s_{t+k}$ from the distribution represented by the ensemble, introduces determinism into the observation reconstruction process. Could the authors explain why they chose this deterministic decoding approach, and how this design choice—which sacrifices the model's stochasticity at the observation level—affects the overall goal of learning an uncertainty-calibrated world model?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "1.\tNovel and Well-Formulated Solution (HCU Loss): The proposed Horizon-Calibrated Uncertainty (HCU) loss is a novel contribution. It enforces an intuitive and realistic inductive bias: predictive uncertainty should grow monotonically with the temporal horizon. This is achieved by maximizing the model disagreement of a probabilistic ensemble, scaled by the horizon length.\n\n2.\tStrong Empirical Results on Diverse Benchmarks (RQ1): The method (HAUWM) achieves state-of-the-art performance and sample efficiency across a broad and challenging suite of benchmarks, including locomotion (DMC), diverse manipulation (MetaWorld), and complex, long-horizon tasks (RoboDesk). The advantage is particularly noticeable in dynamically complex tasks like Walker Run and Hopper Hop.", "weaknesses": "1.\tMotivation for HCU Loss Lacks Direct Quantitative Evidence: The core innovation of the paper—the Horizon-Calibrated Uncertainty (HCU) Loss—is strongly motivated by the claim that prevailing RSSM-based models (like APV) suffer from \"false certainty,\" where predictive uncertainty is suppressed or collapses over long temporal horizons. However, the paper does not present direct empirical evidence (e.g., a dedicated graph showing the variance/disagreement of APV or ContextWM decreasing as the prediction horizon $k$ increases) to substantiate this foundational claim. Without quantitative data explicitly showing this defect in the baseline, the necessity and direct corrective action of the HCU loss are less rigorously established.\n\n2.\tNovelty of Variable-Horizon Prediction: While the HCU loss is novel, the concept of variable-horizon prediction or conditioning dynamics on a temporal embedding is not entirely new (e.g., it is related to positional encodings used in Transformers or general long-term planning models). The authors should more clearly contextualize their specific variable-horizon implementation and its necessity, distinguishing it from prior art that uses temporal embeddings for sequence modeling. \n\n3.\tLimited Task Generalization Scope: Although HAUWM shows strong performance across 10 diverse tasks in three major benchmarks (DMC, MetaWorld, RoboDesk) and generalizes well to Imitation Learning and Offline RL, the total number of tasks tested is a relatively small subset of the established control environments. To fully validate HAUWM as a versatile general-purpose foundation model, the empirical validation should be expanded to include a wider variety of structurally distinct tasks, particularly from the extensive MetaWorld suite, to more robustly substantiate the model's claim of broad generalizability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "x4VXpqAK89", "reviewer": "Reviewer_jamQ", "summary": "The author proposes an uncertainty-based World Model (WM) for model-based reinforcement learning in decision-making. The training is in two stages: the first is uncertainty-based pretraining, where the first stage aims to learn action-free WM, and the second stage is for task-specific fine-tuning. The paper overall presents a clear and sound storyline about their action-free WM and has adequate experiments to support their evidence (e.g. Fig. 5). However, I'm still sceptical about the effectiveness and ablation study. I tend to accept this paper at this point, but I may change my view based on the author's responses.", "questions": "see weakness", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The writing is clear and easy to follow; the claim and presentation are reasonable and clear.\n2. The HCU loss is straightforward and theoretically sound.", "weaknesses": "**Major Concern**\n1. The used benchmark is too old, making the actual performance hard to appreciate. The author should try some latest benchmarks, e.g. SimpleEnv [1], CALVIN [2], BridgeData V2[3], to make their findings more convincing. \n2. Does the second fine-tuning phase largely depend on different tasks under one specific setting? For example, should Push Green and Open Slide on Fig.6 be trained separately or not? If so, I'm curious about the performance under cross setting: finetune WM on push green and perform RL on open slide and vice versa, which will better support the author's claims.\n3. I'm not very sure about the ablation study on the uncertainty assumption: what if you just pre-train the WM using reconstruction loss and still fine-tune it on different tasks? Will it be really bad?\n\n\n**Minor**\n1. The notation is misaligned: $T$ is the segment length in Sec. 4.2, but stands for embedding length in Fig.2.\n\n2. On Eq.3, the model's disagreement is linearly degraded with $k$; how about the other function? e.g. exponential or sinusoidal.\n\n**Ref.**\n\n[1] Evaluating Real-World Robot Manipulation Policies in Simulation CoRL 2024\n\n[2] CALVIN - A benchmark for Language-Conditioned Policy Learning for Long-Horizon Robot Manipulation Tasks R-AL 2022\n\n[3] BridgeData V2: A Dataset for Robot Learning at Scale CoRL 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ceB8sTbNzA", "reviewer": "Reviewer_M6Jh", "summary": "The paper argues that action‑free video pre‑training for world models is insufficient if the objective forces a single deterministic future. It proposes HAUWM, a pre‑training framework that (i) predicts variable‑horizon futures using a temporal embedding, and (ii) encourages uncertainty to increase with prediction horizon via a Horizon‑Calibrated Uncertainty (HCU) loss. Concretely, an ensemble of dynamics heads outputs Gaussian latent predictions; the HCU loss promotes model disagreement scaled by the time gap, combined with a standard ELBO‑style predictive loss in a total loss. During fine‑tuning, the pre‑trained uncertainty‑aware stream is stacked with a lightweight action‑conditioned stream to learn control.", "questions": "Please see the weakness.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The paper identifies an insightful problem of action‑free video pre‑training, where deterministic next‑frame targets suppress stochasticity. The proposed solution (model horizon‑dependent uncertainty) is well motivated. \n\n2. HAUWM integrates variable‑horizon prediction with sinusoidal temporal embeddings and an uncertainty‑increasing prior. The approach is compatible with RSSM/Dreamer‑style world models and standard stacked fine‑tuning. \n\n3. The result is compelling: in pre‑training, ensemble heads diverge for longer horizons, whereas in fine‑tuning, they converge under fixed actions and show reduced uncertainty. This directly supports the paper’s claims.", "weaknesses": "1. It would be interesting and more convincing to see quantitative calibration metrics (coverage) under the influence of different numbers of dynamic heads. These coverage metrics effectively reflect the quality of the pre-trained model.  \n\n2. Each head outputs Gaussian parameters, but the HCU loss uses only squared deviations of head means. Wouldn’t this conflate diversity with calibration?  \n\n3. The reconstruction loss uses mean of ensembles. For multiple possible futures, this choice tends to average them, which would contradict with the diversity encouraged by HCU.  \n\n4. Can authors also ablate the number of ensembles, especially when M=1?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)