WIMLE: Model-Based RL for Continuous Control

• WIMLE is a model-based reinforcement learning approach for continuous control that integrates IMLE, stochastic latent-variable world modeling, ensemble uncertainty estimation, and confidence-based weighting.
• It addresses key challenges such as compounding model error, unimodal averaging in multi-modal environments, and overconfident predictions by selecting optimal latent assignments and attenuating uncertain transitions.
• Empirical evaluations on benchmarks like DeepMind Control, MyoSuite, and HumanoidBench demonstrate that WIMLE improves sample efficiency, with over 50% gains in some tasks while matching or exceeding asymptotic performance.

WIMLE is a model-based reinforcement learning method for continuous control that combines Implicit Maximum Likelihood Estimation (IMLE), stochastic latent-variable world modeling, ensemble-based uncertainty estimation, and confidence-based weighting of synthetic transitions. It is introduced as a response to three recurrent failure modes in model-based RL: compounding model error, unimodal world models that average over multi-modal dynamics, and overconfident predictions that bias learning. The method learns stochastic, multi-modal world models without adversarial training or diffusion-style iterative sampling, and then uses predictive confidence to attenuate the effect of uncertain model rollouts during policy and critic updates (Aghabozorgi et al., 15 Feb 2026).

1. Problem setting and conceptual scope

Model-based reinforcement learning is motivated by sample efficiency, but the formulation underlying WIMLE emphasizes that practical underperformance often arises from model bias rather than insufficient planning or optimization capacity alone. The motivating claim is that compounding model error, unimodal world models that regress toward averages in genuinely multi-modal environments, and overconfident predictions can systematically distort downstream learning (Aghabozorgi et al., 15 Feb 2026).

Within that framing, WIMLE targets continuous-control domains and is evaluated on DeepMind Control, MyoSuite, and HumanoidBench. Its central design choice is to treat the world model not as a deterministic next-state predictor nor as a unimodal Gaussian regressor, but as a one-step conditional latent-variable generator trained with IMLE. This suggests a deliberate shift from likelihood models that can exhibit regression-to-the-mean toward a training rule in which each observed transition is matched to a latent realization that best explains it.

A plausible implication is that WIMLE is best understood as an overview of three mechanisms rather than a single architectural novelty: IMLE-based mode coverage, ensemble-plus-latent predictive uncertainty, and inverse-uncertainty weighting in synthetic data usage. The paper’s empirical claims are presented as evidence that these mechanisms jointly improve stability and sample efficiency across a diverse set of continuous-control problems (Aghabozorgi et al., 15 Feb 2026).

2. World-model parameterization and IMLE objective

WIMLE’s world model is a one-step conditional latent-variable generator

$$g_\theta : (s_t,a_t,z) \to ([r_t, s_{t+1}]), \qquad z \sim \mathcal N(0,I).$$

The training dataset is

$\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$

with

$y_i = [r_i, s_{i+1}].$

For each data point, the method samples $m$ latent candidates

$\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$

Training proceeds by two alternating steps. In the assignment step, performed without gradients, the selected latent is

$$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$

In the update step, $\theta$ is optimized by SGD:

$$\theta \leftarrow \theta - \eta \nabla_\theta \frac{1}{|B|}\sum_{i \in B}\|g_\theta(s_i,a_i,z_i^*) - y_i\|^2.$$

In compact form, the IMLE objective is

$$\min_\theta\;\frac1N\sum_{i=1}^N\;\min_{1\le j\le m}\;\Big\|\,g_\theta(s_i,a_i,z_{i,j})\;-\;y_i\Big\|^2,\quad z_{i,j}\sim\mathcal N(0,I).$$

The method characterizes this per-sample nearest-latent assignment as ensuring mode coverage and avoiding regression-to-the-mean. In the terminology of the source, this yields stochastic, multi-modal dynamics modeling without adversarial training or diffusion-style iterative sampling, while still enabling fast one-step rollouts (Aghabozorgi et al., 15 Feb 2026).

The significance of this objective lies in what it excludes as much as in what it includes. The paper explicitly contrasts the IMLE update with the “averaging” bias of Gaussian MLE. This suggests that WIMLE’s multi-modality is not implemented through explicit mixture-density estimation but through latent assignment under an implicit generator, preserving one-step inference while permitting multiple plausible transition realizations.

3. Predictive uncertainty from ensembles and latent sampling

To estimate predictive uncertainty, WIMLE maintains an ensemble of $K$ independently initialized IMLE models,

$\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$0

At inference time, for a given $\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$1, it draws $\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$2 latents per ensemble member and collects the $\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$3 predictions

$\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$4

with all $\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$5.

The empirical predictive distribution is approximated as

$\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$6

From this sample cloud, WIMLE defines the predictive mean

$\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$7

and the total predictive variance

$\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$8

The variance may further be decomposed by the law of total variance into epistemic and aleatoric components:

$\mathcal D = \{(s_i,a_i,r_i,s_{i+1})\}_{i=1}^N,$9

This decomposition is central to how WIMLE operationalizes uncertainty. Ensemble disagreement contributes epistemic uncertainty, while within-model variation across latent samples contributes aleatoric uncertainty. The paper states that predictive uncertainty is estimated via ensembles and latent sampling, and later argues that this provides a reliable confidence signal (Aghabozorgi et al., 15 Feb 2026).

A plausible implication is that WIMLE does not treat uncertainty as a secondary diagnostic but as a quantity directly coupled to synthetic data curation. In that sense, the predictive distribution is not only descriptive of model ambiguity; it is part of the control mechanism for how imagined experience enters RL optimization.

4. Confidence-based weighting and weighted RL updates

Each synthetic transition $y_i = [r_i, s_{i+1}].$0 is assigned a weight inversely proportional to predictive uncertainty:

$y_i = [r_i, s_{i+1}].$1

The source notes that the “$y_i = [r_i, s_{i+1}].$2” or, more generally, “$y_i = [r_i, s_{i+1}].$3” ensures numerical stability and bounds $y_i = [r_i, s_{i+1}].$4 away from zero. High-variance predictions are therefore down-weighted, while low-variance predictions receive near-full weight (Aghabozorgi et al., 15 Feb 2026).

The weighting is not an isolated heuristic; it is inserted into critic and actor learning. For a mixed batch $y_i = [r_i, s_{i+1}].$5 drawn from $y_i = [r_i, s_{i+1}].$6, the weighted TD-error for critic $y_i = [r_i, s_{i+1}].$7 is

$y_i = [r_i, s_{i+1}].$8

with critic loss

$y_i = [r_i, s_{i+1}].$9

The actor is updated by maximizing the weighted expected value

$m$0

plus an entropy regularizer if using SAC.

The paper explicitly identifies several consequences of this design. First, the rollout loop computes per-step weights $m$1 decreasing with $m$2. Second, both the critic loss and the policy objective incorporate these weights. Third, any off-the-shelf RL backbone, including “SAC + distributional Q,” may be used under the same scheme (Aghabozorgi et al., 15 Feb 2026).

The empirical interpretation offered in the source is that this weighting preserves useful model rollouts while attenuating bias from uncertain predictions and enabling stable learning. It further states that inverse-variance weighting preserves the Bellman fixed point while provably minimizing update variance (Gauss–Markov) (Aghabozorgi et al., 15 Feb 2026). This positions the weighting mechanism as both a practical robustness device and a statistically motivated estimator design.

5. End-to-end algorithmic workflow

WIMLE is presented as a single loop combining real data collection, world-model training, uncertainty estimation, rollout generation, and policy update. The initialization includes policy $m$3, replay buffer $m$4, model buffer $m$5, ensemble $m$6, rollout horizon $m$7, and latents per step $m$8 (Aghabozorgi et al., 15 Feb 2026).

During real interaction, the policy samples actions from $m$9, observes $\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$0, and stores $\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$1 in $\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$2. Periodically, world-model training is performed independently for each ensemble member using minibatches from $\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$3, latent resampling, nearest-latent assignment, and SGD updates under the IMLE objective.

Synthetic rollouts then begin from $\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$4 random start states sampled from $\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$5. At each imagined step, the policy chooses an action, each ensemble member and latent sample predicts $\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$6, the method computes $\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$7, and assigns

$\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$8

One predicted transition, for example from a random $\{z_{i,j} \sim \mathcal N(0,I)\}_{j=1}^m.$9 or from the mean, is added to $$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$0, after which rollout continues from the predicted next state (Aghabozorgi et al., 15 Feb 2026).

The following table condenses the workflow components explicitly described in the source.

Stage	Operation	Stored or updated quantity
Real data collection	Execute $$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$1 and observe transition	Add $$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$2 to $$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$3
World-model training	Sample minibatch, sample latents, assign $$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$4, update $$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$5	Ensemble $$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$6
Synthetic rollout generation	Predict with all $$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$7 pairs, compute $$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$8, set $$z_i^* = \arg\min_{1 \le j \le m} \|g_\theta(s_i,a_i,z_{i,j}) - y_i\|^2.$$9	Add weighted synthetic transitions to $\theta$0
Policy and critic update	Sample mixed batch and optimize weighted objectives	Update $\theta$1 and $\theta$2

The paper’s “Key points” highlight three intended properties of this loop: the IMLE update avoids the averaging bias of Gaussian MLE; the rollout loop computes per-step weights decreasing with uncertainty; and the critic and policy objectives both incorporate those weights (Aghabozorgi et al., 15 Feb 2026). In effect, WIMLE treats model learning, uncertainty quantification, and policy optimization as tightly coupled rather than modularly separable.

6. Empirical evaluation across continuous-control benchmarks

WIMLE was evaluated on 40 continuous-control tasks spanning three benchmark families: DeepMind Control Suite with 16 tasks, including the 4 high-dim Dog / Humanoid tasks; MyoSuite dexterous-manipulation benchmarks with 10 tasks; and HumanoidBench locomotion suite with 14 tasks (Aghabozorgi et al., 15 Feb 2026).

Across all domains, the reported comparison set includes SAC, SimBA/SimBA V2, BRO, TD-MPC2, and Dreamer V3. The paper states that WIMLE achieved both superior sample efficiency and competitive or better asymptotic performance versus these strong baselines (Aghabozorgi et al., 15 Feb 2026).

The benchmark-specific highlights reported in the source are summarized below.

Benchmark result	WIMLE outcome	Comparator outcome
Humanoid-run (DMC)	Reached 80% of final score in 300 K steps	500 K for the best model-based baseline
Dog & Humanoid aggregate (4 tasks) at 500 K steps	Interquartile mean score $\theta$3	BRO’s $\theta$4
MyoSuite at 1 M steps	98% success	SimBA V2 $\theta$5
MyoSuite to reach 90% success	Often required 30–50% fewer environment samples	Top model-free baseline
HumanoidBench (14 tasks) by 1 M steps	Solved 8/14 tasks	BRO 4/14; SimBA V2 5/14
HumanoidBench average at 500 K steps	70% aggregated success	TD-MPC2 38%; SAC 10%

The abstract separately emphasizes two aggregate claims: on the challenging Humanoid-run task, WIMLE improves sample efficiency by over $\theta$6 relative to the strongest competitor, and on HumanoidBench it solves $\theta$7 of $\theta$8 tasks versus $\theta$9 for BRO and $$\theta \leftarrow \theta - \eta \nabla_\theta \frac{1}{|B|}\sum_{i \in B}\|g_\theta(s_i,a_i,z_i^*) - y_i\|^2.$$0 for SimbaV2 (Aghabozorgi et al., 15 Feb 2026). The detailed summary also states that, on the most challenging Humanoid-run task, WIMLE improved sample efficiency by $$\theta \leftarrow \theta - \eta \nabla_\theta \frac{1}{|B|}\sum_{i \in B}\|g_\theta(s_i,a_i,z_i^*) - y_i\|^2.$$1 relative to the next-best method and matched or exceeded asymptotic performance across all benchmarks.

These results are interpreted in the source as demonstrating three points: IMLE-based multi-modal world models capture rich stochastic dynamics and avoid mean-regression artifacts; predictive uncertainty from ensembles plus latent sampling provides a reliable confidence signal; and inverse-variance weighting preserves the Bellman fixed point while provably minimizing update variance (Gauss–Markov) (Aghabozorgi et al., 15 Feb 2026). Taken together, the empirical section presents WIMLE as a stable and highly sample-efficient model-based RL method for diverse continuous-control settings.

7. Relation to baseline families and interpretive boundaries

WIMLE is evaluated against both model-free and model-based baselines. The named comparison methods are SAC, SimBA/SimBA V2, BRO, TD-MPC2, and Dreamer V3 (Aghabozorgi et al., 15 Feb 2026). The significance of this choice is that the claimed gains are not restricted to one comparison class: some reported advantages concern sample efficiency relative to strong model-based baselines, whereas others compare favorably to model-free methods on asymptotic or aggregate task performance.

The paper’s own explanatory account attributes performance to the interaction of three components rather than to any single benchmark-specific engineering choice. First, IMLE is used to learn stochastic, multi-modal world models without iterative sampling. Second, ensembles and latent sampling provide predictive uncertainty. Third, uncertainty is converted into per-transition weights used in critic and actor optimization (Aghabozorgi et al., 15 Feb 2026).

Several misconceptions are implicitly addressed by this formulation. WIMLE is not presented as a deterministic world model with post hoc uncertainty estimates; uncertainty enters through the empirical predictive distribution over ensemble members and latent samples. Nor is it presented as relying on adversarial training or diffusion-style iterative sampling; the source explicitly characterizes IMLE as avoiding both. Likewise, the role of weighting is not merely to discard model rollouts wholesale. The weighting scheme preserves useful model rollouts while attenuating bias from uncertain predictions, so low-confidence synthetic transitions remain in the learning process with reduced influence rather than necessarily being removed (Aghabozorgi et al., 15 Feb 2026).

A plausible implication is that WIMLE occupies a distinct position among model-based RL approaches that use synthetic data: it neither assumes a unimodal parametric transition law nor treats all imagined samples as equally trustworthy. In the source’s formulation, that distinction is the basis for its claim of stable learning under multi-modal dynamics and overconfidence-prone model regimes (Aghabozorgi et al., 15 Feb 2026).