WIMLE: Uncertainty‑Aware World Models with IMLE for Sample‑Efficient Continuous Control
Authors: Mehran Aghabozorgi
Alireza Moazeni
Yanshu Zhang
Ke Li
Venue: ICLR 2026 Poster (2026)
Keywords: Reinforcement Learning
Model-based RL
Primary area: reinforcement learning
TL;DR: Uncertainty aware IMLE world models can outperform strong model-based and model-free baselines in terms of both asymptotic performance and sample efficiency.

Abstract: Model-based reinforcement learning promises strong sample efficiency but often underperforms in practice due to compounding model error, unimodal world models that average over multi-modal dynamics, and overconfident predictions that bias learning. We introduce WIMLE, a model-based method that extends Implicit Maximum Likelihood Estimation (IMLE) to the model-based RL framework to learn stochastic, multi-modal world models without iterative sampling and to estimate predictive uncertainty via ensembles and latent sampling. During training, WIMLE weights each synthetic transition by its predicted confidence, preserving useful model rollouts while attenuating bias from uncertain predictions and enabling stable learning. Across $40$ continuous-control tasks spanning DeepMind Control, MyoSuite, and HumanoidBench, WIMLE achieves superior sample efficiency and competitive or better asymptotic performance than strong model-free and model-based baselines. Notably, on the challenging Humanoid-run task, WIMLE improves sample efficiency by over $50$\% relative to the strongest competitor, and on HumanoidBench it solves $8$ of $14$ tasks (versus $4$ for BRO and $5$ for SimbaV2). These results highlight the value of IMLE-based multi-modality and uncertainty-aware weighting for stable model-based RL.

Published as a conference paper at ICLR 2026
WIMLE: UNCERTAINTY-AWARR WORLD MODELS WITH IMLE FOR SAMPLE-EFFICIENT CONTINUOUS CONTROL
Mehran Aghabozorgi, Alireza Moazeni, Yanshu Zhang, Ke APEX Lab, School of Computing Science Simon Fraser University (mehran aghabozorgi, sam62, yanshu zhang, keli)
ABSTRACT
Model-based reinforcement learning promises strong sample efficiency but of- ten underperforms in practice due to compounding model eπror, unimodal world models that average that bias learning. We introduce WIMLE, a model-based method that extends Implicit Maximum Likelihood Estimation (IMLE) to the model-based RI. frame- work to learn stochastic, multi-modal world models without iterative sampling and estimate predictive uncertainty via ensembles and latent sampling. During training, WIMLE weights each synthetic transition by its predicted confidence, preserving useful model rollouts while attenuating bias from uncertain predic- tions and enabling stable learning. Across 40 continuous-control tasks spanning DeepMind Control, MyoSuite, and HumanoidBench, WIMLE achieves superior sample efficiency and competitive or better asymptotic performance than strong model-free and model-based baselines. Notably, on the challenging Humanoid- run task, WIMLE improves sample efficiency by over 50% relative to the strongest competitor, and on HumanoidBench solves 8 of 14 tasks (versus 4 for BRO and 5 for V2) These results highlight the value of IMLE-based multi- multi-modality and uncertainty-aware weighting for stable model-based RL.
multi-modal dynamics, and overconfident predictions
INTRODUCTION
Humanoid-rru run (DMC)
HI -ss----eeed (Humanoidlenchh)
Myo-key-tum peoo-ke--tum--aaa (MyoSuite)
MB.
0.
Ss etione
SIrBAIZ
Drcamery
g04
310
TOWPC2
8.2
0.2
0.1
1200
(20
$00
1033
200
400
640
$0.0
Figure 1: Sample efficiency on challenging tasks from each benchmark suite. WIMLE achieves superior sample efficiency and asymptotic performance strong model-free and model-based baselines. Y-axes show interquartile mean. Shaded areas indicate 95% confidence intervals.
Reinforcement learning has become powerful framework for solving complex decision-making problems across diverse domains such as autonomous control (Kiumarsi et al. 2018), strategic game playing (Hosu & Rebedeaj 2016), and natural language processing (Lambert) [025 Cetina et al.)
2021). However, a significant challenge in RL is the need for substantial number of interactions with the environment to learn good policy. Without simulator, learning requires world trials,
which are costly, slow, and risky (Chen et al. 2022 Weng et al. 2023 Hessel et al. 2018 Schulman al. 2017b).
Masel-based RL (MBRL) methods aim to address this sample efficiency issue by first learning parametric world model from collected environment interactions, then leveraging this learned

---

Published as a conference paper at ICLR 2026
model reduce real environment samples and accelerate policy learning (Hatner al. 2020
2021
2023
Ye et al.
2021 Laskin et al. 2020). Common uses of the learned world model include (1)
generating synthetic rollouts that augment training data for policy learning (Janer et al. 2019 Ha &
Schmidhuber
2018
Hafner et al. 2020 Clavera et al. (2000) and (2) planning by simulating future
trajectories guide action selection (Zhu 2023] Frauenknechh al. 2025 Janner al.
2019
Lowrey ct al., 2019 Hafner et al. 2019 Argenson & Dulac- 202 In this work, we focus
on the former.
Historically, MBRL has struggled to surpass strong model-free baselines, largely because com- pounding rollout errors bias training and mislead the policy (Janner al. 2019 Xiao et 2019 Talvitic 2017; Frauenknecht et al., [2025] Venkatraman et al. [0155] Asadi et al. 2018b a We at-
tribute this to two key issues: (1) standard predictive models struggle when the same state-action pair yields different, conflicting supervision due to partial observability, contact-rich dynamics, or inherent stochasticity (Kurutach et al.] 2018); and (2) a lack of uncertainty awareness in model predictions (Frauenknecht et al.
2025 which leads overconfidence in regions with complex dy-
namics or limited data. Despite attempts to address these issues (Janner et al. 2019 Zhu al. [023] Frauenknecht et al., 2025 Somalwarr et al., 2025 Hansen al. (0222), MBRL methods have yet to
consistently outperform strong model-free baselines in practice (Nauman et al. 2024 Lee et al, 2025b).
To address these issues, we propose WIMLE (World models with IMLE) an uncertainty-awar model-based RL approach. We integrate IMLE (Li & Malik 2018), mode-covering generative model with demonstrated success in low-data regimes (Aghabozorgi et al. 2023] Vashist et al,, 2024), into the MBRL framework. This allows us to learn world models that handle different, conflicting supervision and from which we extract predictive uncertainty estimates. We incorporate these uncertainty estimates into the RL objective to prevent overconfident predictions from biasing learning. To the best of our knowledge, this is the first work to extend IMLE for uncertainty-awwat aware world models in MBRL..
We evaluate WIMLE on 40 tasks across DMC, HumanoidBench, and MyoSuite. WIMLE deliv- ers considerable gains in sample efficiency and asymptotic performance over strong model-free and model-based baselines, Notably, on the notoriously challenging Humanoid-run task, WIMLE im- proves the sample efficiency of the most competitive method by over 50%. On HumanoidBench,, WIMLE successfully solves of 14 tasks, compared to 4 for BRO and for S2 (Figure I0) Across suites, Figure [
shows one example task per benchmark, each showing more than 50% sample-efficiency improvement for WIMLE over the strongest competing method.
2
PRELIMINARIES
2.
RL
discounted Markov decision process (MDP) (Sr..Prr,,)) (Bell- man 1957) with initial state distribution Pu- At time the agent observes S, selects Or
We consider an
Tol St). receives reward r, = r(*,, 4x), and the environment transitions as 8+41 P(. 81, at))))
The objective is to learn a policy that maximizes the expected discounted return
(TT) = Err(po.P.r P.To))
Σ En'r(s Crrr( t=0
(1)
In the continuous-control settings considered here, states and actions are real-valued. The discount factor satisfies 7 E (0, 1). We denote the action-value function under policy æ by Q=(s,a):
(" (TN(T.P)
Σ (r((s, az) so =
(2)
2.2 MODEL-BASED RL
In model-based RL, we learn a parametric world model with parameters 0 that approximates the unknown environment transition dynamics P(s s,,a;) through learned conditional distri- bution
Po (8441, Pt Sc;
(3)
2

---

Published as a conference paper at ICLR 2026
trained from limited environment interactions. A model rollout (prediction) of horizon H under policy T from So is the sequence = (8o, Ou, To: $1y..., generated by
at To 8t); (8)+1, po |s,, at).
Such rollouts are commonly used for planning or to provide synthetic transitions for RL training (Janner et al. 2019 Zhu et al. 2023).
2.3
IMPLICIT MAXIMUM LIKELIHOOD ESTIMATION
Implicit Maximum Likelihood Estimation (IMLE) (Li Malik 2018 learns latent-variable gen- erator (z(2) that maps noise 2 I), to data space. Given data {a; ) the IMLE objective is:
Ø* = arg min E (
Σ min 13/5m
[190(2)) - x: m:|1?
(5)
In practice, given current parameters 0, we realize this objective by drawing pool of candidate latents {zg 1m ii.d from (00, I) per data point and selecting the nearest generated sample; this step gradient-free and fully parallelizable,
arg min m5jsm
[[90(2)) mi
(6)
and minimizing the resulting empirical loss using stochastic gradient descent.
レ 0
[ao(2;) - 2z
KEB
Optimizing Eq. equation 5 yields maximum likelihood estimation (MLE) of 0 and ensures mode coverage (Aghabozorgi et al. 2023): each data point is represented by least generated sam- ple. In practice, IMLE is sample efficient and effective for modeling multi-modal distributions. Conditional IMLE (Li et al. 2020) (o 2) models multi-modal conditional distributions; adopt this form in WIMLE. For a more detailed discussion of IMLE, we refer readers to (Li & Malik
2018
Aghabozorgi al. 2023).
3 WIMLE
WIMLE addresses the key limitations of traditional MBRL through three main components: (1)
IMLE-trained stochastic world models that capture complex multi-modal transition dynamics, (2)
predictive uncertainty estimation that reflects the model's confidence in its predictions, and (3)
sicertainty-weightee learning that scales the influence of synthetic data based on model confidence.
We detail each component below.
3.1
IMLE WORLD MODEL.
Recent MBRL approaches span autoregressive sequence models, latent-variable generators, diffu- sion models, and planning-centric objectives (Ha & Schmidhuber 2018 Hafner et al.
2019 2020
2021 Robine al. 2023 Micheli al. 2023 Zhang al. [023]
2023] Hansen al. 2024
Huang al.
2024 Diffusion models are effective but rely on iterative sampling, which limits their usage in
the online RL setting where rollout throughput critical (Huang et al., 2024 Karras et al. 2022).
Despite progress, these methods often require substantial data and still struggle to consistently sur- pass strong model-free baselines (Nauman et al. 2024 Lee et al. (0255)). Moreover, simple and sample-efficient unimodal Gaussian world models underfit inherently multi modli, complex dynam- ies in partially observable or contact-rich settings, exacerbating model bias and compounding errors
Janner et al. 2019, Zhu et al. 2023).
On the other hand, we leverage IMLE learn transitions, a one-step generative method that-unliy diffusion models-avoie iterative sampling and enables fast online rollouts. In practice, IMLE yields strong rollout throughput; Figure reports wall-clock time among model-based methods. We represent the world model conditional stochastic generator go that maps state catte-action pair and latent noise to the next outcome:
reports
(st-1; 1; fz)
(1,, 2~ N/O, T).
(8)
3

---

Published as a conference paper at ICLR 2026
Infoprop
Dreamerv3
TD MPC2
0.4
0.0
Wallclock Time thours)
Figure 2: Wall-clock comparison among model-based methods seeds) on single NVIDIA L40S GPU for the humanoid-run task. Y-axis shows interquartile mean; shaded areas indicate 95% confi- dence intervals.
Here, the latent variable induces distribution over next outcomes for the same sate-action pair, capturing inherent stochasticity and multi-modality in the dynamics.
IMLE Training Procedure. Given dataset of transitions ((st, de, r2, (++1))(+ we form targets y: = (rt, 8t-1] and train go using the IMLE objective. The training proceeds in two alternating steps:
Assignment Step: For each data point Yi, sample candidate latents {zg 1 and assign the
nearest candidate that minimizes the prediction error:
arg min [190(s;, 2z) y/ 1| 15/5m
This assignment step is computationally efficient as it requires no gradient computation, is fully parallelizable across data points, and typically uses small values of m (e.g., 5-10) in conditional IMLE settings.
Update Step: We then perform gradient descent on the empirical loss using the assigned latents:
- 0 Vo
Σ
1190 a, 27) vill?
(10)
B
KEB
where is a minibatch of indices and v) > 0 is the learning rate.
This procedure ensures mode coverage by matching each data point to at least generated sam- ple, avoiding collapse to a single mean prediction. In contrast, a standard Gaussian regression model trained with least squares (Janner al. 2019) predicts the conditional mean; in multi-modal settings this falls between modes and produces averaged, often implausible next states-knov as regression to the mean (Bruna et al.) [016]] that compound over rollouts. IMLE's per-sample la- tent ussignment avoids this averaging and yields sharper, mode-consistent predictions (Aghabozorgi
ct al. 2023 Vashist et al. 2024).
Inference and Rollouts. After training, we generate rollouts following the procedure described in Section 2.2] Multi-step rollouts of horizon H are generated by initializing from real state Su and iteratively applying: TT(-(1E), 2, ~ ~O0 I), T,) = ga,,, az 2,) for = O.... H 1.
3.2 UNCERTAINTY ESTIMATION
Reliable uncertainty estimation is crucial for deciding when to trust model predictions. We there- fore compute predictive uncertainty for each synthetic transition and use it reweight the RL objective. Each transition's contribution is scaled by its estimated confidence. This preserves useful rollouts and reduces bias from uncertain predictions without changing the underlying al- gorithm. Alternative integrations exist. For example, Infoprop (Frauenknecht et (025) computes an information- -iformation-theorttic corruption measure and uses to truncate rollouts during generation. In contrast, we integrate uncertainty directly into the learning objective via confidence weights.
We use a single predictive uncertainty measure (s, a) that reflects the model's confidence in its next-step prediction. Concretely, maintain ensemble of K IMLE world models (see Sec- tion I and, for each model, draw m latent samples, yielding predictions:
{ge, (s, a, 2z) k.=.kk
(11)
We define (s, a) the standard deviation across these predictions-a direct measure of model agreement (see Algorithm
lines 13
14]:
a (8, a) = stdk. [90, (s, a, 23)]
(12)
4

---

Published as a conference paper at ICLR 2026
In practice, we compute per-dimension standard deviations for the predicted reward and next state and average them to obtain a single scalar for the transition. o (s, , a) decreases when models agree and increases when predictive uncertainty is high (e.g., limited data complex dynamics). By the law of total variance we can decompose (²(s,)) a) into epistemic and aleatoric components,
(s, a)
Vark E2 [90, (8, =)]]
E3 [Vrr [90, (s, a, =)]1
(13)
ensemble /epistemic
latent allattric
Following recent insights uncertainty from Smith al. 2024)) this total predictive variance is the Bayes risk of acting under squared-error loss.
3.3 UNCERTAINTY-WFIGHTT LEARNING
Having defined a single predictive uncertainty ass, a), we now describe how it enters learning. The key idea is simple: weight each synthetic transition by the model's confidence reliable pre- dictions contribute more and uncertain ones less. Because uncertainty typically grows with roll- horizon due error accumulation, later steps the rollout receive smaller weights, naturally down-weighting distant, noisier predictions.
To choose these weights, invoke a standard heteroscedastic regression result (formalized in Sec- tion 3.4p: when estimating a scalar quantity from independent noisy observations with different noise variances, the unique minimum-variance linear unbiased estimator weights each observation inversely to its variance. We use this principle as guidance and map each predictive standard devia- tion (s, a) a bounded weight
1
w(s, a) a) +1
(14)
which preserves the inverse-variance ordering and keeps w E (0, 1], avoiding exploding gradients.
Our analysis suggests the idealized scaling w(,aa) x 1/o0(s, a). In practice, we use the bounded proxy wss, a) 1/(1 (s, a)), which we find empirically be more stable and work better.
During rollout generation, we compute per-transition weights w; = w(s; a,) for each synthetic tran- sition (st, ai ri, s{) using the predictive uncertainty defined above. We incorporate these weights into the RL objective by modifying the temporal difference (TD) loss:
Laitii E ((r
[1q: 60],
(15)
where 8, = (Q0 0:41) Qe a;) is the TD error for transition i, with a,+1 a. s1+1), Qs is a parameterized Q-function, and we is the corresponding uncertainty weight. For real environment data, we simply use w; 1.
This approach lets us leverage synthetic rollouts while keeping TD updates well-conditioned: high- variance, uncertain transitions have smaller impact the stochastic gradients, while the Bellman fixed point remains unchanged, so we tame noisy (including purely stochastic) updates without discarding or biasing them.
3.4 THEORETICAL ANALYSIS
To make the effect of our weighting concrete, we provide theoretical analysis of its properties. First, we establish that multiplying each squared Bellman error (y Q(s, )² by positive weight w(s, a) preserves the Bellman fixed point. Second, in a tractable setting where the critic is linear in a feature representation, we demonstrate that choosing weights inversely proportional to the (total) target noise variance minimizes the variance of the learned parameters, thereby improving conver- gence rate and sample efficiency. The lemmas below formalize these results: the first shows that any strictly positive reweighting leaves the Bellman target Q* unchanged, and the second shows that, in the linear-critic regime, inverse-variance weighting w; x 1/0 yields the ninimum-covariance linear unbiased estimator. Full proofs are provided in Appendix B
Lemma (Positive weights preserve the Bellman target).
y denote the one-step Bellman target for a value function V with conditional mean Ju(s, a) = Ely s, a]. Given a strictly positive weight function w(s, > the corresponding population weighted Bellman squared loss
Lyp Q) E(ssa)yar P( [M (v a(
(16)
5

---

Published as a conference paper at ICLR 2026
Then the unique minimizer of Lw (Q) over all action-value functions Q is Q*(ss = pu(s a) for all (s, a). ie., any strictly positive reweighting leaves the Bellman fixed point unchanged and only re-emphasizes different regions of (s, a). A full proof is given in Appendix B
Lemma (Linear critics, inverse-variance weighting). Consider the setting where (1) the action- value function linear in some feature representation o(s, a) R i.c. Qs a) o(s, a) for parameters 0 E gd and (2) for a batch of m transitions = (s;, a:) the TD targets satisfy Y: ((x;)) + E, with Eley zy] 0. Var(e; æ;) of and independent noise terms EN- Among all linear unbiased estimators of 0, the inverse-variance choice x 1/0 yields the minimum
covariance matrix for Q. A proof follows the classical Garko- theorem and is provided in Appendix B
Proof sketch. Stacking targets into vector y and features into design matrix 4 minimizing w,yyy - Qo(x,))² yields the weighted least-squares estimator d, ( w) - Wy with W diag(w;). Choosing W proportional to the inverse noise covariance Σ-1 diag(1/o{)) yields an estimator whose covariance Cov(Ow) is minimal among all linear unbiased estimators by the Garko- theorem. We refer to Appendix for a full derivation.
Implications for WIMLE. Here of is the total predictive variance at (5,, az), which can include both epistemic and aleatoric components (Eq. [3). The lemma above shows that, in the linear-critic regime, weighting each transition inversely to this total variance minimizes the covariance of the learned parameters, independently of the source of the noise. Lower parameter variance means fewer samples are needed to reach a given accuracy, i.e., inverse-variance weighting is provably more sample-efficient in this regime. Combined with the Bellman fixed-point lemma above (proved in Appendix B). this means our choice of wnss, a) shrinks update variance without ever changing the Bellman solution-
in the limit of a perfect world model where all uncertainty (and thus down-weighted noise) is purely aleatoric.
3.5 ALGORITHM
Algorithm presents the overall WIMLE procedure. For more complete implementation of the algorithm, including training frequencies and hyperparameters, see Algorithm 3 in Appendix A
Algorithm WIMLE: World Models with Implicit Maximum Likelihood Estimation
1: Input: Rollout horizon H, ensemble size K, number of rollouts M, number of latent codes m
2: (Red text indicates steps that fundamentally differ from MBPO.)
3: initialize ensemble world models {ge, J environment and model datasets Denys D model
4: for environment steps do
5: Collect environment transitions using T&: add to Dem
6: // IMLE World Model Training
7: Train ensemble {ge, K 1==1 in parallel on bootstrap samples of Deuv using IMLE (Eqs. 10
8: for M model rollouts do
9:
Sample starting state so from Denv
for 0 H 1
11:
a (-1st) Sample m latents (2,17= N(O,
13:
Generate predictions {(5141, (
K.m Ik=1j=1 {go, (5c, az, 2,) K,)) from all ensem-
ble members
14:
Compute predictive uncertainty: or = stdk.s Igo, (se, aes 21)] (aggregated over ensem- bles and latents)
15:
Set weight w; = 1/ 1)
Add weighted transition (st, de; Te, 8141, we) to Dmodel Policy Learning
17:
// Uncertainty-Weigee
18:
Sample batch from Danv J Dmoddl (real data has w = 1)
19:
Update policy using weighted RL objective:
20:
Es ""s batch (RL(S, a,

---

Published as a conference paper at ICLR 2026
Figure 3: WIMLE world model architecture.
The algorithm maintains the underlying RI method as a black box through the weighted loss func- tion L. where fRL represents any standard RI. objective (e.g., TD error for critics, policy gradient for actors). The key insight is that uncertainty weights wi automatically scale the contribution of each synthetic transiton--high-coonfe high-confidence predictions receive higher weights while uncertain predictions contribute proportionally less.
3.6
DESIGN CHOICES
3.6.1 TRAINING
Model Rollouts. We experiment with synthetic rollouts using horizons up to H = 8, All roll- outs are initialized from real environment states sampled uniformly from the environment dataset Dene. following standard practice in model-based RL to ensure rollouts start from the data distri- bution. We select task-specific rollout horizons through empirical experimentation as described in Appendix D.I
RI. Training. We use Soft Actor-Critic (SAC) (Haarnoja al. 2018) with distributional Q-learning
[Bellemare et al.) 2017) as underlying RL algorithm. Following recent work that has demon-
strated the effectiveness of distributional RL for continuous control (Nauman et al.] 2024 Lee et al.
20256 Dabney et al., 2018a), specifically adapt quantile Q-learning (Dabney al. 2018h Nau-
man et al. 2024).
Ensemble Training. We train ensemble of K 7 IMLE world models in parallel improve predictive uncertainty estimation and calibration (Section 3.3). Each ensemble member is initial- ized with different random parameters and trained bootstrap samples of the environment data. The parallel training of ensemble members computationally efficient and scales well with avail- able compute resources, enabling reliable predictive uncertainty without significant computational overhead.
3.6.2 ARCHITECTURE
Figure illustrates the WIMLE world model architecture. The network takes as input state St action Ge. and latent variable followed dense layer that maps representation. The core of the architecture consists of three residual blocks, each containing dense layers with ReLU activations and L2 normalization. Following recent findings LeLee al. (2025b), we employ L2 normalization within the residual blocks, which has been shown to improve stability and performance in RL settings. The network outputs separate predictions for rewards and next states through dedicated dense heads.
312-dimensional hidden
4 EXPERIMENTS
We evaluate WIMLE across diverse continuous-comte benchmarks DeeedsineeMii Control Suite (in- cluding Dog and Humanoid), MyoSuite, and HumanoidBench TTassa al.) 2018 Caggiano al,, 2022 Sferrazza et al. 2024). Across 40 tasks spanning locomotion and dexterous manipulation with high-dimensional state/action spaces and sparse rewards, we compare against strong model- free and model-based methods, including MR.Q, PPO, SAC, Simba, Simba BRO, TD-MPC2,
and DreamerV3 (Fujimoto et al. 20255, Schulman et al. 2017a Haarnoja et al. 2018 Lee et al, Nauman et al. [2024] Hansen et al. [2224] Hafner al. 2023). and present per- benchmarka results. Through our experiments, we aim to answer: (i) How does WIMLE compare to strong
2025a Nauman
model-free and model-based methods? (ii) How does IMI.-bassed multi-modality in the world

---

Published as a conference paper at ICLR 2026
model affect results compared to standard unimodal Gaussian models? (iii) How do uncertainty estimates evolve during training, and how do they affect performance?
4.1
EXPERIMENTAL SETUP
All experiments are run for IM environment steps with 10 random seeds unless otherwise specified. We report the interquartile mean (IQM) and 95% confidence intervals computed with RLiable (Agar- wal et al. 2021). 1) using stratified bootstrap across tasks and seeds. Following BRO and SimbaV2 NNumann al. 2024 Lee et al., 2025b), aggregate normalized scores per BRO/Simba V2 proto- col (DMC [o MyoSuite success, HumanoidBench success-normalizedl Where official baseline results available, we report the authors' numbers; otherwise, we run public implementations with their recommended settings. We provide full details about the experimental setup, hyperparameters, and baselines in Section D and of the appendix.
4.2 COMPARISON TO BASELINES
& Humanoid (7 tusks)
DMC (16 tasks)
MyuSuite (10 tasks)
HumanoidBench (14 tasks)
Figure 4: Aggregate results across benchmarks. WIMLE outperforms strong model-free and model- based baselines overall. Gains are most pronounced on the challenging Dog & Humanoid subset, where it achieves superior sample efficiency and asymptotic performance. On MyoSuite, it performs asymptotically on par with strong baselines that are already near the maximum score (1.0), and on HumanoidBench it significantly outperforms the baselines, solving 8/14 tasks versus BRO 4 and SimbaV2 5. Y-axes show interquartile mean; shaded areas denote 95% confidence intervals.
We summarize aggregate performance across benchmarks in Figure and provide detailed per-task results in Section C of the appendix. WIMLE consistently leads among strong model-free and model-based methods on Dog & Humanoid, the full DMC suite, and HumanoidBench, while per- forming asymptotically on par with strong MyoSuite baselines that are already close to the maximum score (1.0). Notably, gains are largest on the high high-dimensional and challenging Dog Humanoid tasks (Dog: S|-223,, |A|=38; Humanoid: |S|=67, [A|-24). On HumanoidBench, WIMLE sig- nificantly outperforms baselines, solving of tasks versus BRO and SimbaV2 5 (Figure TO). We summarize performance across timesteps in Section E. where WIMLE performs best or com- petitively across most evaluations. We attribute these improvements to IMLE-driven multi modality in the world model and uncertainty-weightedd learning that scales the influence of synthetic roll- outs by model confidence, mitigating bias from overconfident predictions while preserving useful signal, which we discuss in more detail in the next section. Per-task performance is reported in Figures
g 9 and 10
4.3 METHOD ANALYSIS
We analyze how uncertainty-awar weighting and multi-modal dynamics modeling impact perfor- mance and how predictive uncertainty evolves during training.
Effect of uncertainty- aware weighting to an unweighted variant that identical in every respect except that all per- transition weights are fixed to w,=1.0. The unweighted curve lags and can even underperform strong model-free baseline early on, indicating that ignoring predictive uncertainty significantly bi- ases learning and hinders performance. Figure 5 (Right) studies rollout sensitivity. Increasing the model rollout horizon from ==1 to ===4 to =6 improves performance, and extending to H= maintains performance rather than showing the severe degradation typically observed in model- based methods when increasing rollout length (Janner et al. 2019). This improved stability at longer horizons demonstrates that uncertainty-awe aware weighting reduces the model bias typically introduced by longer horizon errors, enabling us to leverage longer synthetic rollouts without considerable per- formance degradation.
weighting Figure 5 (Left) compares WIMLE with uncertainty-
8

---

Published as a conference paper at ICLR 2026
Impact of IMLE-based multi-modality 6 contrasts WIMLE (IMLE world model) with an otherwise identical unimodal Gaussian world model (MBPO-style; (Janner et al. 2019). ) with both variants using uncertainty-awar weighting. The IMLE variant significantly outperforms the Gaussian, underscoring the value of modeling multi-modal transition dynamics for uncertainty estimation in complex, contact-rich control. Figure O (Right) shows how weights evolve. During a brief warm-up with limited environment samples and training, both models are uncalibrated and weights can appear transiently high; as data accumulates and the estimators calibrate, weights drop to reflect high uncertainty and low confidence. As training progresses and more data are collected, IMLE's weights increase to reflect higher confidence in the predictions, whereas the Gaussian's remain relatively flat, indicating limited calibration. Together, these results show that multi-modal modeling improves both performance and the quality of uncertainty estimates, reducing the risk of overconfident, biased predictions misleading the policy.
Uncertainty- vs, unweighted
Rollout sensitivity
0..3
Ours
1.0
Hel
Ro 1o
Hed
Free
0.0
0.29
0.0
60.23
0.3
0.2
0.23
0.51
0.0g
202
reeeeneee Stepe txtree)
Feutstraer Strp4 1xh00
Figure 5: Uncertainty-awwe weighting reduces model bias and enables stable training at longer horizons on Humanoid-run. Left: Uncertainty-aware WIMLE compared to an unweighted variant that identical except all per-transition weights are fixed to w, = 1.0 and model-free variant that is identical except that it does not use the model; the unweighted curve lags and can even underperform the model-free variant early on, indicating that ignoring uncertainty will bias learning and hinder performance. Right: Rollout ablation (H = 1, 4, 6, 8) for WIMLE: increasing the model rollout horizon from = to H=4 to == improves performance, and extending to H-8 does not substantially degrade performance, suggesting that sicertainty-awan weighting mitigates harm from error accumulation at longer horizons. All variants use the same SAC backbone and distributional critics; only the ablated components differ. All plots on DMC's Humanoid-run task with 5 seeds.
IMLF vs. Gaussian (performance)
IMLF vs. Gaussian (weights)
PLI
BLE Weights (A Geundin Welgen (Wo)
Goutukan
0.30
3025
0.20
0.23
0.0
0.5
0.00
Divsormet mord Stans (0100))
Eivaurmtint Stapo 10000
Figure 6: Multi-modality strengthens model-based learning. Left: WIMLE (IMLE world model) versus an otherwise identical unimodal Gaussian world model (MBPO-style; (Janner et al.) 2019); both use uncertainty-aware weighting): the IMLE variant significantly outperforms the Gaussian, highlighting the value of multi-modal modeling and IMLE's efficacy. Right: Weight dynamics: After a brief warm-up phase, IMLE's weights are lower when uncertainty high and increase as training progresses and more data are collected, reflecting growing model confidence; the unimodal Gaussian fails to capture this evolution, yielding relatively flat weights over time. All plots are on Humanoid-run.
5 RELATED WORK
Model-free RL. Foundational model-free methods such PPO (Schulman et al., 2017a) and
SAC (Haarnoja et al. 2018) remain strong references for continuous control. Recent advances focus on scaling and regularization: BRO (Nauman et al., 2024) scales critic networks to 5M parameters with strong regularization and optimistic exploration, achieving sfate-of-the-a performance. Simba [Lee et al. 2025a) introduces an architecture that embeds simplicity bias through running statistics
9

---

Published as a conference paper at ICLR 2026
normalization, residual feedforward blocks, and layer normalization, enabling effective parameter scaling; SimbaV2 (Lee et al. 2025b) further constrains feature and weight norms via hyperspheri- cal normalization. Contemporary work like MR.Q (Fujimoto al. 2025) explores improved value estimation for better sample efficiency. Collectively, these methods provide strong model-free base- lines.
Model-based RL. Model-based RL methods learn world models to improve sample efficiency via synthetic rollouts and planning. DreamerV3 (Hafner et al., 2023) learns a latent world model and achieves strong performance in continuous control with large-scale training. MBPO (Janner aal.|2019) uses short moodll-genertted rollouts branched from real data avoid model exploitation while maintaining sample efficiency. TD-MPCC (Hansen et al. 2024) learns implicit world models through joint-embedding prediction and performs local trajectory optimization in latent space for scalable multi-task learning. STORM (Zhang et al.] 2023) combines Transformer-bant
sequence modeling with categorical VAEs for efficient world model learning in visual domains. Diffusion- based world models generate trajectories via iterative denoising and incur high inference cost, which hinders online RL (Janner et al. [0222] Ajay et al., [2023: He et al.) 2023). Despite these algorithmic advances, model-basee methods have struggled to consistently surpass recent model-free approaches like BRO (Nauman et al. 2024 and Simba V2 Lee et al. 2025b).
Model Bias. Model bias and error accumulation remain fundamental challenges in MBRL. Tra- jectory models (Asadi et al. 2019 Lambert et al.] (202)) address the compounding-erre problem by learning multi-step models that directly predict outcomes of action sequences, avoiding the accumu- lation of one-step prediction errors. Self-correcting models (Talvitie) (017) train models to correct themselves when producing errors. Infoprop (Frauenknecht truncating rollouts using information-thece theorretc corruption measures, but is not competitive on com- plex, high-dimensional tasks such as Humanoid-run (see Figure TT. In contrast, we estimate a single predictive uncertainty and weight each synthetic transition accordingly, integrating this directly into the learning objective to preserve useful synthetic data while reducing the influence of uncertain predictions. This yields state-of-the-ar art results on challenging tasks (Figures
al.) (025) integrates uncertainty by
[and 4).
Implicit Maximum Likelihood Estimation. IMLE (Li Malik 2018) trains implicit generative models by minimizing the expected distance from each data point to its nearest generated sam- plc, avoiding mode-collapse and GAN (Goodfellow et al., 2014) training issues. Adaptive IMLE (Aghabozorgi al. 2023) extends this with adaptive thresholding and curriculum learning for bet- ter few-shot performance. These methods demonstrate that likelihood- objectives can achieve good sample quality without adversarial training on low-data settings.
6 LIMITATIONS AND FUTURE WORK
WIMLE uses world models solely to generate synthetic rollouts. Other uses, such as planning with the model or integrating the model into policy-gradient formulations, remain unexplored here. Future work should evaluate WIMLE in these settings. Our experiments use proprioceptive state observations only. Extending WIMLE to image-based control is important direction, especially since IMLE has been shown to be effective in few-shot image synthesis (Aghabozorgi et al. 2023 Vashist et al., 2024). Finally, similar to MBPO (Janner et al.) 2019) and POMP (Zhu et al. 2023), the rollout horizon is a desk-dependent hyperparameter. Learning to adapt the horizon online based on model confidence is a promising avenue for future research.
7 CONCLUSION
WIMLE advances model-based reinforcement learning by extending IMLE learn stochastic, multi-modal world models and by weighting synthetic data with predictive confidence. This re- duces model bias and stabilizes learning while retaining the benefits of synthetic rollouts. Across 40 continuous-controll tasks in DMC, MyoSuite, and HumanoidBench, WIMLE achieves superior sample efficiency and competitive or higher asymptotic performance than strong model-free and model-based baselines. Gains are largest on challenging Dog and Humanoid locomotion tasks. On HumanoidBench, WIMLE significantly outperforms baselines and solves 8 of 14 tasks. The ap- proach integrates cleanly with standard RL objectives and scales with compute through ensembles and parallel latent sampling. We hope these results renew interest in practical world models for challenging continuous control.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This research was enabled in part by support provided by NSERC, the BC DRI Group and the Digital Research Alliance of Canada. We are also grateful to James A. Peltier for his support on the computing infrastructure that made this work possible.
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural informa- tion processing systems, 34:29304-2993 293200, 2021.
Mehran Aghabozorgi, Shichong Peng. and Ke Li. Adaptive IMLE for few-shot pretraining free generative modelling. 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. 248 264. PMLR, 2023. URL BRLLhtpss: httlattaaapprcccci.ccoo mer...mrr.... mwt.ppess/vv222. aghabozorgi22a.. html
Anurag Ajay. Yilun Du, Abhi Gupta, Joshua Tenenbaum. Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision decision-making?? International Conference on Learning Representations, 2023.
Arthur Argenson and Gabriel Dulac-Arold. Model-based offline planning. ArXiv, abs/2008.05556 2021.
Kavosh Asadi, Evan Cater, Dipendra Misra, and Michael L. Littman. Towards a Simple Approach to Multi-step Model-based Reinforcement Learning. arXiv, 2018a.
Kavosh Asadi, Dipendra Misra, and Michael L. Littman. Lipschitz Continuity in Model-based Reinforcement Learning. arXiv, 2018b.
Kavosh Asadi, Dipendra Misra, Seungchan Kim, and Michel L. Littman. Combating Compounding-Errot Problem with Multi-step Model. arXiv, 2019.
Marc G. Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforce-
ment learning. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th Interna- tional Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-1 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 149-458, PMLR, 2017. URL
http::/proceed: ngs mmm...p... www..press/vee /v/ssvv1//// be lemare re 7a htm
Richard Bellman. A markovian decision process. Indiana Univ. Math. J., 6:679-684, 1957. ISSN 0022-2518.
Joan Bruna, Pablo Sprechmann, and Yann LeCun. Super-resolution with deep convolutional suffi- cient statistics. International Conference on Learning Representations, 2016.
Vittorio Caggiano, Huawei Wang. Guillaume Durandau, Massimo Sartori, and Vikash Kumar. Myosuite: A contact-rich simulation suite for musculoskeletal motor control. In Roya Firoozi, Negar Mehr, Esen Yel, Rika Antonova, Jeannette Bohg. Mac Schwager, and Mykel Kochen- derfer (eds.), Learning for Dynamics and Control Conference, L4DC 2022, 23-24 June 2022, Stanford University, Stanford, CA, USA, volume 168 of Proceedings of Machine Learning Re- search, pp. 492-507. PMLR, 2022. URL htttps://procee//// ngs.m ml www.mmlr.....ll..mmmmmmmm www.ss/vl6.... Eaggiano22a.hn html
Victor Ue Cetina, Nicolás Navarro-Guerrero, Ana Martin-González, Cornelius Weber, and Stefan Wermter. Survey on reinforcement learning for language processing. Artificial Intelligence Re- view. 16:151115575, 2021. URL http::/oppismmntcccocooooooo.o. wttp://laaaarg//ooooooooooo 233210638
Jie Chen, Jian Sun, and Gang Wang. From unmanned systems to autonomous intelligent systems. Engineering, 12:16-19, 2022.
11

---

Published as a conference paper at ICLR 2026
Ignasi Clavera, Yao Fu, and P. Abbeel. Model-augmented actor-critic: Backpropagating through paths. ArXiv, 1ww/2005508063 2020.
Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for dis- wwwp:/s//ggaaooooe. 8068 06923
tributional reinforcement learning, 20. URL
Will Dabney, Mark Rowland, Marc G. Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Sheila A. Mcllraith and Kilian Q. Weinberger (eds.), Pro- ceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium on Ed- ucational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, Febru- URL.
ary 2-7, 2018, pp. 2892-290200 AAAI Press, 2018b. doi: (1600//AAAV.31.... https: //doi.::/dd.... wwg. / 10. 1609/6600 aaai. 112211.1179)
Bernd Frauenknecht, Devdutt Subhasish, Friedrich Solowjow, and Sebastian Trimpe. On rollouts in model-based reinforcement learning. In The Thirteenth International Conference on Learn- ing Representations, ICLR 2025, Singapore, April 24-28, 2025. OpenReview.net, 2025. URL https://<penrevi// net fert//ooum?ta fit.ormmnid-hhssnmn...
Scott Fujimoto, Pierluca D' O'Oro, Amy Zhang. Yuandong Tian, and Michael Rabbat. Towards general-purpose model-free reinforcement learning. arXiv preprint arXiv:2501.16142, 2025.
Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks, 2014. URL https:
/ //axxiv.erg/abs///....
2661
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems 31, pp. 2451-2463 Curran Associates, Inc., 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement leaming with stochastic actor. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmäässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1865-1865. PMLR, 2018. URL mlr. www.press/vee v80 wttp:///v/aaaarooooomm haaaaoo ja Bb. html
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference Machine Learning, pp. 2555-2565, 2019.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
URL
net net forum iat.rrummid-11111144240
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with
discrete world models. In International Conference on Learning Representations, 2021. URL
ltttps://ppereviiie net/
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint 2tti::2030.04000. 2023.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predic- tive control. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Re- search, pp. 8387-8406. PMLR, 2022. URL harsen22a html
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. URL ttttp:://opereeviee
12

---

Published as a conference paper at ICLR 2026
Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xue- long Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. Advances in Neural Information Processing Systems, 2023.
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
lonel-Alexandru Hosu and Traian Rebedea. Playing atari games with deep reinforcement learning and human checkpoint replay. CoRR, abs/1 607.05077, 2016. URI. http:://rxxv..rrr
abs/ 05077
Renming Huang. Yunqiang Pei, Guoqing Wang. Yangming Zhang, Yang Yang, Peng Wang, and Hengtao Shen. Diffusion models as optimizers for efficient planning in offline rl, 2024. URL
hittss://arxt org/ alss/ 2207 16142
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine When trust your model: Model-based policy optimization. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelz- imer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett (eds.). Advances in Neu- ral Information Processing Systems 32: Annual Conference on Neural Information Pro- cessing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 12498 11498-12508 12509, 2019. URL https::/ppooceeeiiiiiiccc Absssarattttttttttttta....aa...
neurips. octpee/eeece....e..rroeccco/pppp
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. International Conference on Machine Learning, 2022.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- 00302000600000
based generative models, 2022.
Bahare Kiumarsi, Kyriakos G. Vamvoudakis, Hamidreza Modares, and Frank L. Lewis. Op- timal and autonomous control using reinforcement learning: A survey. IEEE Transactions on Neural Networks and Learning Systems, 19:2042-2062 2018. URL https://api.
Gemanticscholar..c org/ httpg/g/ocooousus0000..5511 ott.arr.....c.....oooooooooaaraa
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. In 6th International Conference Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Truck Proceedings. net OpenReview.net, 2018. URL htttps://opezeeeeeeee httpp//cwwerrrteeeeeeo/ooooodaaa foreeeffoorummmmmss.nnnie.
Nathan Lambert. Reinforcement learning from human feedback, 2025. URL https://aaxxx
org/ borg/bb//20 2504/ 12501
Nathan O. Lambert, Albert Wilcox, Howard Zhang, Kristofer S. Pister, and Roberto Calandra. Learning Accurate Long-term Dynamics for Model-based Reinforcement Learning. IEEE Conf on Decision and Control, 2021.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: tions for reinforcement learning. In International Conference on Machine Learning, pp. 5639- 5650. PMLR, 2020.
Contrastive unsupervised representa-
Hojoon Lee, Dongyoon Hwang, Donghu Kim, Hyunseung Kim, Jun Jet Tai, Kaushik Subramanian, Peter R. Wurman, Jaegul Choo, Peter Stone, and Takuma Seno. Simba: Simplicity bias for scaling up parameters in deep reinforcement learning. In The Thirteenth International Conference Learning Representations, ICLR 2025, Singapore, April 24-28, 2025. OpenReview.ned 2025a. URL htttps://openrevieess net. forum? id= DKsuDo
Hojoon Lee, Youngdo Lee, Takuma Seno, Donghu Kim, Peter Stone, and Jaegul Choo. Hyper- spherical normalization for scalable deep reinforcement learning. CoRR, 052.222222212210.. 2025b. doi: 10448550AAXXI..550..2222 URL lhttps://dol wwwpo..oo......oooooo.oo..o..ooo 48550///1...00///00// awv.vv 2502. 15280
13

---

Published as a conference paper at ICLR 2026
Ke Li and Jitendra Malik. Implicit maximum likelihood estimation. CoRR, URL http:://axxii......//// org/ 09087
2018.
Ke Li, Shichong Peng. Tianhao Zhang, and Jitendra Malik. Multimodal image synthesis with conditional implicit maximum likelihood estimation. International Journal of Computer Vi- sion, 128:2607 2628, 2020. URL/tttps://apl. wwwantecscoonaaataiaa. httpa//ssssppp..smmmiicccooo..r. org org/CorpusID:
2154 16300
Kendall Lowrey, Aravind Rajeswaran, Sham M. Kakade, Emanuel Todorov, and Igor Mordatch. Plan online, learn offline: Efficient learning and exploration via model-based control. ArXiv, abs/1811, 11.0188,, 2019.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world mod- cls. In The Eleventh International Conference on Learning Representations, ICLR 2023, Ki- gali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https:://opnnrevier Forum?id whhumidd-wFulAAbbo..
Michal Nauman, Mateusz Ostaszewski, Krzysztof Jankowski, Piotr Milos, and Marek Cy- gan. Bigger, regularized, optimistic: scaling for compute and sample efficient con- tinuous control. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang (eds.), Advances in Neu- ral Information Processing Systems 38: Anmual Conference on Neural Information Pro- cessing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024. URL http:/ /pater/ nips. htwpapperr.tlleeepper////// htw.://2aaaaaccccaaaaaccaooonaae
Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world mod- els happy with 100k interactions. In The Eleventh International Conference Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL httpss:/cppereeye
net/ wwtt/eeeeeiidateeeeeeeeeeereeeee forummid==dBaDDGeC
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: DrxXiv:/07..00 1707.06347 2017a.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347, 2017b.
Carmelo Sferrazza, Dun-Ming Huang, Xingyu Lin, Youngwoon Lee, and Pieter Abbeel. Humanoid- Simulated humanoid benchmark for whole-body locomotion and manipulation. In Dana Kulie, Gentiane Venture, Kostas F. Bekris, and Enrique Coronado (eds.), Robotics: Science and Systems XX. Delft, The Netherlands, July 15-19, 2024, 2024. doi: 10.156677RSS5.224.XX... 00..0..56677788 2024. XX. 061
bench:
URL https:://oi..oo.
Freddie Bickford Smith, Jannik Kossen, Eleanor Trollope, Mark der Wilk, Adam Foster, and Tom Rainforth. Rethinking aleatoric and epistemic uncertainty. CoRR, abs/2412.20892, 2024. doi: 1044550/ARXXI.2412222555 URL ltttt:::/ http://ti..ogo/////// /10.//....... 48550//rxxv..2411
20:92
Anne Somalwar, Bruce D. Lee, George J. Pappas, and Nikolai Matni. Learning with imperfect models: When multi-step prediction mitigates compounding error. CoRR, abs/2504.01766, 2025. doi: 1048550/RXXV.2504.000 01766
URI. httpt::////o htt.::////0...og/0/0045555///111 488500/axxxx 2504.
Erik Talvitie. Self-correcting models for model-based reinforcement learning. Satinder Singh
and Shaul Markovitch (eds.), Proceedings of the Thirty-Fürst AAAI Conference on Artificial Intel- ligence, Februury 4-9, 2017, San Francisco, California, USA, pp. 2597-2603, AAAI Press, 2017.
doi: 001609/AAALL31I1 108500.
URL sttps::/doi.ooe wwt.cii..ggg/oooooooooo wtt.:ooo/11..110/aaaa... a0a/ v31il.
10850
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden. Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq. Timothy P. Lillicrap, and Mar- tin A. Riedmiller. Deepmind control suite. CoRR, abs/1801.00690, 2018. URL http:
7/arxiv wwwp://.wrabbba/l...11
14

---

Published as a conference paper at ICLR 2026
Chirag Vashist, Shichong Peng, and Ke Li. Rejection sampling IMLE: designing priors for better few-shot image synthesis. 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 XXI, volume 15079 of Lecture Notes in Computer Science, pp. 441 456. Springer, 2024. doi: 10.
25. URL
https:://oi.org// 1tt.///...1001///1000
1007/ 978 3-01-- 72664-662222 25
Arun Venkatraman, Martial Hebert, and J.. Bagnell. Improving Multi-Step Prediction of Learned Time Series Models. AAAI, 2015.
Boxi Weng, Jian Sun, Gao Huang, Fang Deng. Gang Wang, and Jie Chen. Competitive meta- learning. IEEE/CAA Journal of Automatica Sinica, 1009):1902-1904 2023.
Chenjun Xiao, Yifan Wu, Chen Ma, Dale Schuurmans, and Martin Müller. Learning combat compounding-error in model-based reinforcement learning, CoRR, abs/1912.11206,, 2019. URL
wwtp://////a///ooo http://arxiv.ort att//bb///999 11206
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 34:25476- 25488, 2021.
Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. STORM: efficient stochastic transformer based world models for reinforcement learning. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL 564 266778842222222220003881110-0023
piap nips httpp/sccccppppppepeaeeeeeee
Jinhua Zhu, Yue Wang. Lijun Wu, Tao Qin, Wengang Zhou, Tie- Yan Liu, and Houqiang Li. Making better decision by directly planning in continuous control. In The Eleventh International Confer- ence on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URLbttpps opwnrevieveeeeeeeeee net/ fort/form??? dur 7idxyF
15

---

Published as a conference paper at ICLR 2026
AUTHOR STATEMENT: USE OF LANGUAGE MODELS
We used large language models to help polish writing and improve clarity. All ideas, methods, ex- periments, and analyses were created and verified by the authors. Any suggested text reviewed and edited by the authors for accuracy and originality.
A ALGORITHM DETAILS
Algorithm B the detailed implementation of WIMLE, including training frequencies, batch sizes, and other practical considerations omitted from the main algorithm for clarity. The IMLE training procedure is detailed in Algorithm Hyperparameters are provided in Section D
Algorithm 2 IMLE World Model Training
1: Input: Environment dataset Deny, ensemble {ge, 1==1, number of latent codes m, number of updates U, learning rate
2: for u = 1 U do
3: Sample minibatch {(s,, as. 91+1))ee with replacement from D
4: Form targets y/: = [re. $;41 for all E
5: // Assignment Step (Eq.9)
6: Sample m candidate latents { ~ I)
7: for k 1 K in parallel do
8:
N = arg min; <jsm l1ge, (st, a2 2z) - yill2 for all B
9:
// Update Step (Eq. 10)
for k 1 K in parallel do
11:
02 - 02 Vo B EiEB 1190, (8,, 04 zi,e) - -yill2
Algorithm 3 WIMLE: Detailed Implementation
1: Input: Rollout horizon H, ensemble size K 7, batch size B, model training frequency train_freq, number of latent codes m. number model updates U
2: Initialize policy Tg. ensemble of IMLE world models {ge, J environment dataset Dew. model dataset model
3: for environment steps do
4: Collect environment transition using Tे. add to Dan
5: if step mod train freq then
6: // IMLE World Model Training
7:
Train ensemble (go,) K parallel using Algorithm) // Uncertainty-Awarr Rollout Generation
8:
9:
Clear Dmodel
10:
Sample batch of starting states {000) ) with replacement from Deny for to H 11-1 do
a, (-(sss (4) for all i (1....,B) B}
(0)
12:
13:
Sample m latents {zilE1 ~ ~)
14:
Generate predictions {go, (st de 22)) = from all ensemble members for all i (i)
Compute predictive uncertainty: o = stdy. 990. (2) a 12)] (aggregated over en-
15:
sembles and latents}
16:
Set weight w¹"" 10 +1
17:
Select transitions (s ( from predictions
Add weighted transitions {(s).) aa 8111 (1) to D
(6) (1)
19:
// Ancertainty-Weighee
Policy Learning
20:
Sample batch from Danv J Dmoddel (real data has w = 1)
21:
Update policy using weighted RL objective:
22:
Es "ss. bbach (RL(S, a, r, s')]
16

---

Published as a conference paper at ICLR 2026
B THEORETICAL ANALYSIS OF RISK-WEIGHTED BELLMAN ESTIMATION
Bellman fixed point under positive reweighting. Recall the one-step Bellman target
Fix policy π and target value function V.
y v=r+aV(s))
(17)
with conditional mean
(1ss, a) = Ely s, a] = (() +VV(s))
(18)
Given a strietly positive weight function w A (0, xc), the corresponding population weighted Bellman regression loss for arbitrary action-value function Q is
L (Q) Els,a)ydd ((ra")) [w(s.
(v [ a)
Q(s.a))²].
(19)
Lemma (Weights do not change the Bellman target). Assume that w(s, a) 0 for all (s, ) in the support of d², Then the unique minimizer of (Q) over all action-vale functions Q SSx R
(s(s.a) = p(s, a)
for all (s, as)
(20)
which is the same minimizer as for the unweighted objective (obtained by setting w = 1).
Proof. Fix (s, a) and consider the conditional risk function of a scalar q E
(1) E w(s, a) (y x)² s. al w(s a) E (1 x)² 8
(21)
Since w(s.a) a) is a strictly positive constant with respect to the inner expectation, it does not affect the minimizer in q. The derivative of the unweighted term with respect q is
a [[v q)² s, a] 2(q Ely s, al) 2(q - ((s,)))))) 0q
(22)
which vanishes if and only if y = u(s a). Thus, for each fixed (s,4). the unique minimizer of (g) q a), independently of the choice of w(s a) >
The global objective (Q) in Eq. 19 is the expectation of la ((, a)) under d". For any action- value function Q, we can write
(Q.a(Q(s, - (u a)) w(s B(((s))[((((so) - u(s, a(), s.a] 2 0,
(23)
with equality if and only if a) p(s, a). Integrating this inequality with respect to d shows that (Q) > (H), with strict inequality whenever Q differs from on set of positive d=- measure. Hence, Q = is the unique minimizer of Cw,, and this minimizer does not depend on the particular choice of strictly positive weights.
Proof of the linear-critic GLS lemma. We restate the setting from the main text. Let (ss, a,) denote m nate-action pairs with feature vectors Φ(x;)) E Rdd, and let y be the corresponding one- step TD targets. Assume linear value model
y: = (zx;) 8* +
(24)
where 0* I is the true parameter vector, the noise satisfies Ele; x.]] 0, and the conditional variances are Var(e; x.)) = o? Collect the targets into a vector y E R" the features into design matrix O E mm whose i-th row is T, and define the noise vector E = (E1,..., Then
h +
(25)
with E(s] and covariance
Σ= Cov(s) diag ....g. ..)
(26)
We assume that & has full column rank that all normal equations below solvable.
17

---

Published as a conference paper at ICLR 2026
Linear unbiased estimators and GLS. Consider linear estimators of H* of the form
i = -Ay,
(27)
for some matrix A E Rdxm Unbiasedness for all 6* requires
E E[A($0* + +
A@6 0*
for all 0*
(28)
which is equivalent the constraint
AΦ = Ia-
(29)
Under this constraint, the covariance of θ is
Cov(() A Cov(y) A AΣAT.
(30)
The generalized least-squares (GLS) estimator corresponds to the specific choice
ACI. = Σ-> -
(31)
so that
0GTS- ACTSV- It is immediate that AGI. = Ia. so OGL is linear and unbiased. Its covariance is
(32)
Cov(fGls) = AGLLE = (o -/)
(33)
Optimality of GLS.
Let 0 Ay be any other linear unbiased estimator, so that A I7 by
equation 29 Define
C= ACLS.
(34)
Then
C$ AΦ ACIS
la la 0.
(35)
The covariance of á can expanded as
Cōv( AEAT
(36)
(AciS + C + CCT AGIS AGISEC + CEAGLS + C.CT.
(37)
(38)
We now show that the cross terms vanish. Using the definition of AGLS.
AGLS
$ -I
<x Φ
(39)
and similarly
EAILS = ΣΣー 1 --00--10---000-
(40)
Therefore,
AGLSECT ( Σ -1 $ CT -===aT>>> (TIΣ- -1 ((C T 0.
(41)
CEATLS C+ ( Σ 1 = 0,
(42)
where we used equation 35 Thus
Cov(ū) Cov(Dors) CΣCT.
(43)
For any vector Rdd
v = (CTT)TE(CTO) 2 0,
(44)
because Σ is positive semidefinite. Hence C is positive semidefinite, and we have
Cov(êGLs) 1 Cov((),
(45)
with equality if and only if C 0, i.c., A = AGu.S This establishes that GLS has minimum covariance among all linear unbiased estimators.
18

---

Published as a conference paper at ICLR 2026
Connection to inverse-variance weighting. In setting, weighted least squares with per- sample weights Wi > corresponds to minimizing
Σ w (zx)) T))² = (y- TW W 20),
(46)
where W diag(u) ......... wm). The normal equations yield the estimator @w ( _ ( - Wy.
(47)
Choosing weights wi x 1/03 makes W proportional to Σ- so that
,
Ow - + Σ Σ ÔGLS-
(48)
Thus inverse-varianee weights w; 1/ recover the GLS estimator and, by the argument above, minimize the covariance matrix among all linear unbiased estimators.
C PER-TASK RESULTS
We present detailed per-task performance results for WIMLE and other baselines across all bench- marks. The performance on each individual task is shown in Figures 7 and I0
TOMPE
Rrpestadd
Rcy Naww
Spra
Figure 7: Per-task results for high-dimensional Dog Humanoid tasks from DeepMind Control Suite. We present the IQM of rewards and 95% confidence intervals for BRO and other baselines run for 1M steps.
Figure 8: Per-task results for DeepMind Control Suite tasks with low-dimensional statc/action spaces. We present the IQM of rewards and 95% confidence intervals for WIMLE and other base- lines run for 1M steps.
19

---

Published as a conference paper at ICLR 2026
Figure 9: Per-task results for MyoSuite tasks. We present the IQM of success rate and 95% confi- dence intervals for WIMLE and other baselines run for 1M steps.
Figure 10: Per-task results for HumanoidBench tasks. We present the IQM of rewards and 95% confidence intervals for WIMLE and other baselines run for 1M steps. The red dashed line indicates the success threshold for each task.
D HYPERPARAMETERS
Table lists the common hyperparameters used across all tasks. These parameters were selected through hyperparameter tuning based on standard practices in RL.
Based on our empirical evaluations, we found that increasing the model batch size to the maximum extent allowed by available GPU resources while proportionally decreasing the number of model
20

---

Published as a conference paper at ICLR 2026
Table 1: Common hyperparameters used across all tasks.
Parameter
Value
SAC Parameters
Batch size
128
Actor learning rate
3 x 10-4
Critic learning rate
3 10
Number of quantiles
100
Updates per step
World Model Parameters
Model learning rate
10-3
Model batch size
512
Model updates
200
Number of latent codes
4
Model training frequency
1000
Number of rollouts
300
Number of ensembles
updates can achieve similar performance with improved training speed. When scaling the batch size in this manner, the model learning rate should be adjusted accordingly following standard Machine Learning practices.
D.1 ROLLOUT LENGTH SELECTION
We select task-specific rollout horizons H through experimentation. For easier tasks where base- lines already saturate near the maximum score (e.g., MyoSuite manipulation tasks), we start with short horizons (H=1-2) and increase only if performance benefits observed, longer horizons may still introduce slight performance tegradation-ttouuh not to the extent seen in traditional MBRL methods. For harder tasks requiring longer-term planning (c.g., HumanoidBench. Dog & Humanoid), we begin with longer horizons (H=8) and decrease only if performance gains are seen empirically. However, we that even simpler tasks may benefit from longer horizons in some cases, reflecting the task-specific nature of optimal rollout length. This selection balances the ben- efits of synthetic data augmentation with computational cost, as the marginal benefit of additional rollout steps diminishes beyond each task's optimal horizon. We cap // at 8 to maintain rollout throughput and because we observe diminishing returns beyond task-specific optima; Tables 6 7 and [8] report the chosen H per task. An interesting future direction would be to dynamically adjust rollout horizons based on the model's uncertainty level, potentially allowing for adaptive rollout lengths that scale with model confidence.
21

---

Published as a conference paper at ICLR 2026
0.9
H-10
0.8
11=20
0.7
H-40
0.6
0.5
0.4
0.3
0.2
0.1
0.0
100
200
300
400
500
Environment Steps (x1000)
Figure 11: Rollout length sensitivity on hl-hurdle-v0 vO task of HumanoidBench up to H 40. WIMLE remains stable longer horizons without rollout scheduling.
EXPERIMENT DETAILS
This section provides detailed descriptions of the benchmark environments used in our evaluation. We explain the task suites, their characteristics, and the normalization procedures used for fair com- parison across different score scales. The following subsections describe each benchmark suite with complete task lists and their state/action dimensions.
E.1 DETAILED RESULTS
We present comprehensive IQM results for WIMLE and baseline methods across all benchmark suites at 100k, 200k, 500k, and 1M environment steps. The best performing method for each step count is highlighted in bold and the second best are underlined. WIMLE performs better across most evaluations.
Table 2: IQM results for DMC suite. Best scores are highlighted in bold, second best are underlined.
Method
100k
200k
500k
1M
MR.Q
0.153
0.362
0.714
0.830
SAC
0.037
0.082
0,210
0.326
SimBA
0.120
0.263
0.522
0.691
SimBAV2
0.235
0.495
0.730
0.845
BRO
0.294
0.519
0.542
0.846
DreamerV3
0.051
0.075
0.165
0.286
TD-MPC.
0.152
0.374
0.566
0.696
WIMLE
0.332
0.575
0.812
0.871
F.2 DEEPMIND CONTROL SUITE
DeepMind Control Suite TTssa et al. 2018 DMC) is a standard continuous control benchmark en- compassing locomotion and manipulation tasks with varying complexity. We evaluate 16 tasks from DMC, focusing the most challenging locomotion tasks including Dog and Humanoid embodi- ments. All returns are normalized by dividing by 1000 to scale performance to [0,1]. The complete list of tasks with their observation and action dimensions is provided in Table 0
22

---

Published as a conference paper at ICLR 2026
Table 3: IQM results for Dog & Humanoid suite. Best scores are highlighted in bold, second best are underlined.
Method
100k
200k
500k
1M
MR.Q
0.042
0.127
0.557
0.796
SAC
0.007
0.008
0.043
0.069
SimBA
0.067
0.173
0.533
0.773
SimBAV2
0.082
0.200
0.601
0.808
BRO
0.086
0.290
0.355
0.864
DreamerV3
0.006
0.006
0.007
0.010
TD-MPC2
0.014
0.058
0.302
0.527
WIMLE
0.140
0.389
0.803
0.897
Table 4: IQM results for MyoSuite. Best scores highlighted in bold, second best are underlined.
Method
100k
200k
500k
1M
SAC
0.038
0.350
0.622
0.714
SimBA
0,566
0.728
0.912
0.952
SimBAV2
0.724
0.830
0.956
0.990
BRO
0,440
0.736
0.816
0.980
DreamerV3
0.028
0.044
0.181
0.466
TD-MPC2
0.088
0.394
0.688
0.775
WIMLE
0.460
0.620
0.928
0.980
Table 5: IQM results for HumanoidBench. Best scores highlighted in bold, second best underlined.
Method
100k
200k
500k
1M
SAC
0.008
0.020
0.060
0.168
SimBA
0.070
0.164
0.322
0.521
SimBAV2
0.059
0.179
0.488
0.799
BRO
0.064
0.127
0.100
0.530
DreamerV3
0.003
0.003
0.005
0.007
TD-MPC2
0.023
0.064
0.382
0.734
WIMLE
0.169
0.393
0.715
0.898
E.3 MYOSUITE
MyoSuite (Caggiano al. 2022) provides high-fidelity musculoskeletal simulations for dexterous manipulation tasks. We evaluate 10 tasks including both fixed-goal and randomized-goa (hard) settings. Performance is measured using success rates, which naturally scale to [0.1]. The complete list of tasks with their observation and action dimensions is provided in Table
E.4 HUMANOIDBENCH
HumanoidBench Sferrazza al. 2024 provides locomotion tasks for the UniTree HI humanoid robot. We evaluate 14 tasks spanning balance, locomotion, and manipulation. For fair comparison across tasks with different score scales, all HumanoidBench scores normalized using each task ' target success score and random score following the same procedure as in (Lee et al. 2025a b):
random score
Success-Normalizedd:)) Target success score random score
The complete list of tasks with their observation and action dimensions is provided in Table X and the random scores and target success scores used for normalization listed in Table
23

---

Published as a conference paper at ICLR 2026
Table 6: DMC Tasks. Complete list of 16 DMC tasks evaluated, with state and action dimensions and rollout lengths.
Task
State dim [S]
Action dim
|
H
acrobet-sswinuup
8
cheetah-run
finger-turn.hard
12
1
fish-swim
8
hopper-hop
15
1
pendulum-sswingup
quadruped-run
78
12
2
reacher-hard
walker-run
24
1
humancid-stand
67
24
2
humanoid-walk
67
24
humancid-run
67
24
6
dog-stand
223
dog-walk
223
38
4
dog-rum run
dog-trot
223
38
4
Table 7: MyoSuite Tasks. Complete list of 10 MyoSuite tasks evaluated, with state and action dimensions and rollout lengths.
Task
State dim [8]
Action dim A
H
myo-key-tura
93
39
myo-key-tura-hard
93
39
1
myo-obj-hold
91
myo-obj-hold-harrd
91
39
myo-pen-twirl
83
39
myo-pen-twiil-hard
83
39
1
myo-pose
108
39
2
myo-pose-hard
108
39
myo eyyoreeac
115
39
4
avoorreah-han
115
39
Table 8: HumanoidBench Tasks. Complete list of 14 HumanoidBench tasks evaluated, with state and action dimensions and rollout lengths.
Task
State dim
[S]
Action dim
[A]
H
hl-balance simple-v0
h1 h1-balance.har har-vvo
77
19
6
h1-crawl-v0
51
19
10
h1-hurdle-v0
51
19
10
h1-maze-v0
51
19
10
h1-pole-v0
19
10
h1-reach-v0
57
19
10
h1-run-v0
51
19
10
h1-slide-v0
51
19
10
h1-sit.hard-v0
h1-sit.simplee-v0
51
19
10
h1-stair-v0
19
h1-stand-v0
51
19
10
h1-walk-vo
19
10
24

---

Published as a conference paper at ICLR 2026
Table 9: HumanoidBench Normalization Scores. Random scores and target success scores used for normalization.
Task
Random Score
Target Success Score
h1 pil-bllancessiplle
9.391
800
h1-ballacc-hhard
9.044
800
h1 hl-crawl
272.658
700
h1 h1-hurdle
2.214
700
hl-maze
106.441
1200
hl-pole
20.09
700
h1 h1-reach
260.302
12000
h1 h1-run
2.02
700
h1 h1-slide-v0
2.02
700
hl-slide-vi
2.02
700
hl-sit-hard
10.545
800
hl-stair
2.214
700
hl-stand
10.545
800
h1-walk
2.377
700
F ADDITIONAL RESULTS
Figure 12 compares standard WIMLE (mixed real imagined data) against an "imagined-only" vari- ant whose critic and actor are trained exclusively on -ooeel-generated rollouts. Removing real tran- sitions does reduce performance slightly, but the imagined-only curve remains close the original WIMLE results, underscoring that our synthetic trajectories are strong enough to sustain competitive learning.
Figure 14 and Table TO illustrate the wall-clock efficiency of WIMLE. Despite the overhead of ensemble training, WIMLE's superior sample efficiency results in significantly lower total time- to-solution compared to baselines.
1.0
Imagined Only WIMLE
0.8
0.6
1QM
0.2
0.0
200
400
600
800
1000
Environment Steps (x (x1000)
Figure 12: Imagined-only WIMLE real data in critic/actor updates) achieves performance com- parable to the standard configuration, highlighting the strength of the generated rollouts.
25

---

Published as a conference paper at ICLR 2026
8000
InfoProp
7000
6000
5000
4000
3000
2000
1000
0
25
100 Environment Steps (x1000)
175
200
Figure 13:
Humanoid (Mujoco) comparison against InfoProp for 4 seeds. Unlike InfoProp, which uses task-specific termination functions for synthetic rollouts, WIMLE runs with the generic setup used for our DMC experiments yet still matches or exceeds InfoProp.
Method
Step 300k
Step 500k
Step 1000k
Time (h)
IQM
Time (h)
IQM
Time (h)
IQM
Infoprop
11.49
0.044
19.15
0.102
38.30
0,077
DreamerV3
3,90
0,001
6.50
0.001
13,00
0,001
TD-MPC2
3.43
0.001
5.72
0.038
11.43
0.190
Ours
4.74
0.168
7.91
0.345
15.81
0.561
Table 10: Performance and time different training milestones.
1.0
BRD
0.8
0.6
IQM
0.4
0.2
0.0
0
10
20
30
40
50
60
70
Data Collection + Compute Time (hours)
Figure 14: Projected total wall-clock time (algorithm compute + data collection) on Humanoid-rum assuming a 5Hz control rate. WIMLE reaches asymptotic performance significantly faster than BRO (the fastest model-free baseline) because superior sample efficiency drastically reduces the time spent collecting real-world data, outweighing its higher per-update compute cost.
26


Reviews:
[{"review_id": "nx1E3brAET", "reviewer": "Reviewer_FqDz", "summary": "The paper proposes WIMLE, a model-based reinforcement learning (MBRL) algorithm that uses Implicit Maximum Likelihood Estimation (IMLE) from Li & Malik 2018 to learn multi-modal, uncertainty-aware world models for continuous control. WIMLE leverages conditional IMLE to generate stochastic transitions without iterative sampling, enabling fast rollouts while avoiding mode collapse typical of unimodal Gaussian models. An ensemble of IMLE world models provides predictive uncertainty estimates, which are then incorporated into the RL objective through an uncertainty-weighted TD loss. Empirically, WIMLE is evaluated across 40 tasks from different benchmark suites, demonstrating substantial gains in sample efficiency and strong asymptotic performance relative to baselines. The paper positions WIMLE as a practical alternative to diffusion- or transformer-based world models for online RL, claiming competitive performance with reduced inference cost.", "questions": "1. You convincingly argue that dynamics are multi-modal; why not extend the IMLE-based formulation to multi-modal rewards or even policy/value functions?\n2. If a non-multimodal model achieves the same asymptotic performance as WIMLE, does that imply that the underlying task is not truly multimodal? How do you detect multimodality in practice?\n3. Comparison to diffusion or transformer world models - since you position WIMLE as a more efficient alternative, could you report asymptotic and wall-time comparisons against diffusion-based and transformer-based world models?\n4. Why only update the critic uncertainty-weighted (Eq. 13)? Would uncertainty-aware plicy gradients offer additional stability?\n5. Figure 2 suggests favorable throughput. Can you provide wall-time comparisons against the baselines in Figure 4? Effectively swap out the x-axis for time.\n6. Could WIMLE be used as a plug-in world model within planning-based algorithms (e.g., TD-MPC2 or MPPI)? How would uncertainty weighting interact with such planners?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "1. **Timely and relative contribution** - addresses two active challenges in model-based RL - multi-modal dynamics modeling and uncertainty calibration through an elegant adaptation of IMLE.\n2. **Strong experimental coverage** - Evaluation spans three modern continuous-control suites and 40 tasks with statistically rigorous reporting from Agarwal et al., 2021\n3. **Conceptual generality** - the proposed uncertainty-aware weighting is algorithm-agnostic and could integrate with diverse RL backbones.\n4. **Empirical clarity** - Ablations (Fig 5-6) isolate the impact of uncertainty weighting and IMLE-based multimodality, demonstrating that multi-modal model predictions both improve calibration and stability at longer horizons.", "weaknesses": "1. **Writing and organization** - Section 3 (Method) is dense and at times difficult to follow. The algorithmic code could be clarified by referencing a baseline (e.g., MBPO) and then enumerating modifications. Sections 3.2-3.3 overlap conceptually and might be merged for concision. Figure placement and scaling should be revised - several are too small to read; figure text should roughly match main text size.\n2. **Missing multimodal context** - While the paper emphasizes multimodal dynamics, it remains unclear how well IMLE captures multimodality relative to more contemporary models, e.g., transformer-based world models (Micheli et al., 2023; Robine et al., 2023) or diffusion models (Ajay et al., 2023). A pedagogical analysis (similar to Figure 3, Chi et al., 2024) illustrating where multimodality matters would strengthen the argument.\n3. **Missing baselines** - Although MR.Q and PPO are mentioned as baselines, their results appear absent from key plots. Including them, especially PPO for wall-time comparison, would provide a clearer model-free reference.\n4. **Insufficient exploration of uncertainty weighting** - the uncertainty metric $w=\\dfrac{1}{\\sigma + 1}$ is intuitive but somewhat ad hoc. Alternative formulations (e.g., softmax or learned calibration of confidence) would help validate robustness and necessity.\n5. **Presentation and formatting** - The manuscript has several typographic and formatting inconsistencies. Figures (especially 3, 5, and 6) and algorithm boxes should be cleaned up for readability. The main algorithm should be self-contained enough to implement directly.\n6. **Scope of multimodality** - the paper restricts multimodal modeling to transition dynamics. Given that the reward and value functions can also be multimodal, extending the IMLE framework to those components could be a natural next step.", "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": "JQdHn7PqUr", "reviewer": "Reviewer_GhD7", "summary": "This paper proposes WIMLE, a model-based reinforcement learning (MBRL) framework that combines IMLE with uncertainty-aware weighting of rollouts to downweight uncertain counterfactuals in a dyna-style framework. The method trains an ensemble of IMLE world models to capture multi-modal dynamics and computes predictive uncertainty by aggregating variance across ensemble members and latent samples. Synthetic rollouts are then weighted by uncertainty to downweight uncertain transitions in the critic update, reducing bias from unreliable predictions. Using SAC with distributional and quantile Q-learning, WIMLE achieves improved sample efficiency and asymptotic performance on 40 continuous-control tasks from DMC, MyoSuite, and HumanoidBench, outperforming strong model-free and model-based baselines and notably solving tasks in the HumanoidBenchmark that comparable methods could not.", "questions": "1. How does WIMLE perform on standard Gym benchmarks? And as a follow on, why were InfoProp and Macura-style adaptive rollout methods excluded from direct comparison, given their conceptual similarity in uncertainty modulation? I saw the note about InfoProp not performing in humanoid-run, but as I mentioned in the weaknesses section, the explanation doesn't seem sufficient with the added context of Ant and Humanoid from Gym.\n2. Have you tested rollout horizons beyond H = 8? Would uncertainty weighting still stabilize longer synthetic rollouts, as MBPO demonstrated for horizons up to 25 or more? Are there asymptotic max performance gains to be had with longer rollouts due to the down weighting? Or does the training dataset get saturated with bad/down weighted synthetic data if you opt for progressively longer rollouts and that leads to degraded performance since entire batches are down weighted?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is the first to extend Implicit Maximum Likelihood Estimation (IMLE) to the conditional, stochastic setting of model-based reinforcement learning. This adaptation allows the world model to represent multi-modal transitions, mitigating the regression-to-the-mean problem common in Gaussian dynamics models like Model-Based Policy Optimization (MBPO). This improvement is shown through ablations wherein they compare their method with and without IMLE.\n\n- The world model choice of IMLE is simple to swap in place of prior methods and relatively straightforward making the wall-clock (Figure 2) and performance gains (Figure 4) attractive alongside the relative simplicity of implementation.\n\n- WIMLE introduces a clean and general mechanism to handle model uncertainty once IMLE is introduced and is comparable to prior words like MBPO: ensemble variance aggregated over latent samples is used to compute per-transition confidence weights. These weights directly scale the critic loss, reducing the influence of unreliable synthetic transitions. This simple yet effective weighting stabilizes learning, mitigates autoregressive errors across a few short horizon lengths.\n\n- The authors evaluate on 40 continuous-control tasks spanning three suites (DMC, MyoSuite, and HumanoidBench), which together cover locomotion and dexterous manipulation at different difficulty levels. WIMLE achieves strong sample efficiency improvements and solves more HumanoidBench tasks (8) than competing baselines (4 and 5 for Bro and SimbaV2, respectively). The diversity of evaluated settings and consistent advantage over both model-free and model-based methods lend credibility to the reported improvements.\n\n- The paper is well-organized, and the visualizations are effective in communicating the author's desired goals. Figures 5 and 6 clearly illustrate the practical effects of uncertainty weighting and IMLE-based multi-modality on performance and confidence calibration. The authors’ explanations are intuitive and technically sound, making the concepts introduced in this work accessible to a broader RL audience and relatively easy to implement.", "weaknesses": "- Comparisons to other uncertainty-aware methods such as InfoProp [1] and Macura [2] are mostly missing, even though these are conceptually the most similar baselines which is noted frequently by the authors. There is a note in the model-bias section of the paper saying that InfoProp fails in high dimensional complex domains like DMC humanoid-run as well as a subfigure (Figure 1)(a)), but InfoProp was not deployed in DMC in the original paper (just Gym) and likewise with Macura. Was any hyperparameter tuning done to show these methods don't work? Furthermore, the Ant and Humanoid tasks in OpenAI Gym are both complex, high-dimensional tasks, and their state and action spaces are larger than many in DMC other than dog. Both Macura and InfoProp outperform prior work on those tasks, contradicting this paper’s claim that such methods fail to scale to high-dimensional domains. Even without direct evaluation against other uncertainty-based methods, it is at least clear that WIMLE is competitive or better than SOTA model-based and model-free methods. However the relative performance when compared to InfoProp and Macura is not a fair comparison when neither was tested in DMC, Myosuite, etc and leaves out an important comparison to prior similar works.\n\n- Following on the prior point, the paper does not include any results on OpenAI Gym, making it difficult to assess cross-domain generalization, which has been shown to be a challenge for Dyna-style methods. Prior works [3] and [4] show that performance often fails to transfer between Gym and DMC for Dyna-style methods (e.g., MBPO, ALM), so the omission of deployment in this benchmark weakens the claim of broad generality across benchmarks despite including 40 tasks across 3 benchmarks.  \n\n- The performance attribution between WIMLE’s IMLE weighting and its upgraded SAC backbone is unclear. The authors use a distributional, quantile-based version of SAC without explaining why, but then go on to compare against SAC without these additional modifications. As a result, some of the observed gains might come from the introduction of these modifications to SAC, but be attributed to WIMLE. It seems necessary to at minimum include further ablations showing performance after the removal of the model beyond what is shown in Figure 5(a) for humanoid-run for the model-free case. \n\n- Reported rollout horizons are limited to H = 8, whereas MBPO and related methods demonstrate stability at horizons of 15–25 (when combined with heuristic or algorithmic iteratively rollout length selection) and then show performance collapse for horizons of 500 in MBPO (See Figure 3 (b) in the MBPO paper). The claim that uncertainty-aware weighting mitigates long-horizon bias is therefore untested at scales where such bias typically becomes problematic enough to cause performance collapse.  \n\n- As mentioned previously, several design choices are under-explained, including the decision to use quantile Q-learning and distributional critics. These components can independently improve stability and learning dynamics, but the paper does not disentangle their contributions from those of WIMLE’s IMLE-based uncertainty weighting. \n\n- Unless I missed it, the wall-clock comparison in Figure 2 lacks key details such as the benchmark environment and codebase implementations (e.g., JAX, PyTorch, TensorFlow) used for each baseline. Without these specifics, it is difficult to assess whether the timing advantage reflects algorithmic efficiency or implementation-level differences.\n\n[1] Frauenknecht, B., Subhasish, D., Solowjow, F., & Trimpe, S. (2025). On rollouts in model-based reinforcement learning. In Proceedings of the 2025 International Conference on Learning Representations (ICLR 2025). arXiv preprint arXiv:2501.16918. Available at https://arxiv.org/abs/2501.16918\n\n[2] Frauenknecht, B., Eisele, A., Subhasish, D., Solowjow, F., & Trimpe, S. (2024). Trust the model where it trusts itself: Model-based actor-critic with uncertainty-aware rollout adaption. In Proceedings of the 41st International Conference on Machine Learning (ICML 2024) (Vol. 235, pp. 13973–14005). Available at https://arxiv.org/abs/2405.19014\n\n[3] Voelcker, C. A., Hussing, M., & Eaton, E. (2024). Can we hop in general? A discussion of benchmark selection and design using the Hopper environment. In Finding the Frame: An RLC Workshop for Examining Conceptual Frameworks (RLC 2024). arXiv preprint arXiv:2410.08870. Available at https://arxiv.org/abs/2410.08870\n\n[4] Barkley, B., & Fridovich-Keil, D. (2025). Stealing that free lunch: Exposing the limits of Dyna-style reinforcement learning. In Proceedings of the 42nd International Conference on Machine Learning (Vol. 267, pp. 2978–3002). Proceedings of Machine Learning Research. Available at https://proceedings.mlr.press/v267/barkley25a.html", "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": "emYmo7ythi", "reviewer": "Reviewer_8Fgq", "summary": "The paper proposed to leverage Implicit Maximum Likelihood Estimation (IMLE) as base architecture for model-based reinforcement learning.\n\nWhile the use of IMLE is to some extent incremental, leveraging its uncertainty estimates within TD learning is interesting.\nOverall the empirical results of this paper are encouraging: on most tasks the performance generally modest, however on several harder tasks (Humanoid Run, H1-slide and Myo-key-turn-hard), the proposed algorithm is roughly two to three times more sample efficient.", "questions": "* What UTD values did you use when comparing the model-free methods (e.g. BRO)? Did you try to ablate their performance and sample efficiency when increasing UTD?\n* How does WIMLE compares to model-free methods in terms of wall-clock time? It would be great to discuss/demonstrate empirically the tradeoff between performance and wall-clock time.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "* Weighting Bellman backups by model uncertainty is an interesting idea.\n\n* The experiments are convincing: the paper compares WIMLE across several hard continuous RL benchmarks against what is (to my knowledge) the state-of-the-art algorithms (model-free and model-based)", "weaknesses": "**The usage of uncertainty is not well grounded**\n* For instance, equation 12 does not distinguish between the aleatoric and epistemic uncertainty. This can lead to low credit assignment in states with high stochasticity, rather than states with large model error. I believe this is not observed in the experiments as most of these Mujoco tasks are mostly deterministic.\n* In addition, rewarding states based on their epistemic uncertainty (“intrinsic rewards”) is known to accelerate learning in sparse reward tasks. A discussion/study on how WIMLE works even in light of this contradiction would improve the paper.\n* Aside from very high-level intuition about penalizing states where model accuracy is low, there is no grounding for why this may work. Building some theoretical intuition (even if not at the strictest rigorousness) would improve the paper.\n\n**Other**\n\n* Minor: adding vision tasks would make the contribution of this paper more impactful.", "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": "kzOgLulVFB", "reviewer": "Reviewer_z1Wb", "summary": "This paper proposes WIMLE, a Model-Based Reinforcement Learning (MBRL) approach that constructs an uncertainty-aware world model by integrating the IMLE generative model. The method aims to address compound errors and the issue of overconfident predictions prevalent in existing MBRL algorithms. The authors evaluate WIMLE on several vector-state benchmarks, including DMC, HumanoidBench, and MyoSuite. The results indicate that WIMLE achieves significant improvements in both sample efficiency and asymptotic performance.", "questions": "1.  As noted in Weakness 1, could the experimental scope be extended to image-state environments (e.g., a subset of DMC tasks using raw pixel observations or the Atari 100k benchmark) to evaluate WIMLE’s effectiveness in such contexts?\n2.  As noted in Weakness 2, given that the comparison algorithms (DreamerV3, TDMPC2) primarily use model-generated trajectories for policy training in their original implementations, please clarify whether this variable was unified in the current comparative experiments. If not, to ensure a fair comparison, could the authors provide an ablation experiment where the WIMLE policy is trained exclusively on model-generated trajectories (i.e., without real environment data)?\n3. WIMLE uses uncertainty-based weighted updates within the Actor-Critic architecture. Is the weighting factor applied only to the loss update of the Value Network, or is it applied to the loss updates of both the Policy Network and the Value Network?\n4. It is recommended that the results presented in Figure 2 be converted into a table format. The current graphical representation is poor for clearly depicting the performance (e.g., time efficiency) of lower-performing algorithms like DreamerV3.\n5. Code and Implementation Details: To ensure research reproducibility, can the authors publicly release the complete experimental code and model implementation upon acceptance of the paper? This is critical for the community to verify the claimed performance improvements.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper successfully extends and integrates IMLE (a mode-covering generative model proven effective in low-data regimes) into the MBRL framework for world model learning.\n- The method effectively integrates uncertainty estimates from model predictions into the Reinforcement Learning objective function, which helps mitigate policy misalignment caused by model bias and overconfident predictions.\n- WIMLE demonstrates substantial improvements over strong model-free and model-based baselines across multiple continuous control tasks, particularly showcasing outstanding performance on the challenging Humanoid-run task.", "weaknesses": "- Limitation in Observation Space: The experimental evaluation is limited to vector-state (low-dimensional) environments. Many state-of-the-art MBRL algorithms have demonstrated strong performance in high-dimensional (pixel/image) state spaces. The paper currently lacks verification of WIMLE’s performance in these crucial high-dimensional settings.\n- Potential for Unfair Comparison: The WIMLE algorithm uses a mixture of real and imaginary data for policy training. In contrast, the original paper implementations of comparable baselines (e.g., DreamerV3, TDMPC2) typically rely solely on imaginary trajectories. This difference in data usage introduces a potential source of unfairness in the comparison and requires clarification from the authors.", "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)