Bootstrap Off-policy with World Model

AI summary

Overview

• Research area: Reinforcement learning — specifically the intersection of model-based online planning and off-policy learning in high-dimensional continuous control.
• Technical level: Advanced. The paper assumes familiarity with Markov Decision Processes, off-policy actor-critic methods, KL divergences, model predictive control (MPPI), and latent world models.
• One-sentence scope: The paper diagnoses "actor divergence" — the mismatch between planner-collected data and the policy's own behavior — and proposes a bootstrap alignment framework that closes this gap with a likelihood-free, value-weighted imitation objective layered on top of a jointly learned world model.

What This Paper Is About

When a model-based RL agent uses online planning (such as MPPI) to pick the actions it actually executes in the environment, the data it collects comes from a different actor than the neural policy it is trying to train. This mismatch — called actor divergence — corrupts value estimates and makes policy updates unreliable. The authors' goal is to keep planning's benefits (high-quality, foresighted actions) while forcing the policy and the planner to converge, so that off-policy learning stays distributionally consistent.

Key Contributions

1. BOOM framework with a bootstrap loop. A design in which the policy initializes the planner, the planner refines actions via model predictive optimization, and the planner in turn bootstraps the policy through behavior alignment — supported by a jointly learned world model that both simulates futures for the planner and supplies value targets for policy improvement.
2. A likelihood-free alignment loss. Because MPPI's reweighting and resampling make the planner's action likelihood intractable, the authors minimize the forward KL divergence KL(β‖π) and discard the term depending only on β, reducing the objective to −log π(a|s) evaluated at planner actions — requiring no access to β(a|s).
3. A soft value-weighted (Q-weighted) mechanism. Alignment targets are weighted by w_i = exp(Q_i/τ) / Σ_j exp(Q_j/τ) with τ set to 1 by default, so that the policy prioritizes high-return planner actions and tolerates quality variability among historical actions stored in the replay buffer.
4. Theoretical guarantees plus state-of-the-art empirical results. Two theorems bound the return gap and the Q-value gap under alignment, and BOOM achieves the best Total Average Return across all 14 evaluated tasks on the DeepMind Control Suite and Humanoid-Bench. Code is released at https://github.com/molumitu/BOOM_MBRL.

Main Findings

• Actor divergence is structural, not incidental. The paper formalizes two consequences: (1) the value function minimizes Bellman error only over the behavior distribution d^β(s,a), producing biased and overconfident estimates in out-of-distribution regions; and (2) the policy is updated using these biased estimates, misguiding improvement.
• Forward KL beats reverse KL in practice. On the Dog-run ablation task (223/38 state/action dims), the likelihood-free forward KL objective consistently reached higher returns than reverse KL, which had to be approximated with a Gaussian surrogate built from value-weighted means and variances of planner actions.
• Q-weighting accelerates learning. Replacing soft Q-weights with uniform weights degraded both training speed and final performance, consistent with the claim that weighting handles variability in the replay buffer's planner actions.
• The alignment coefficient is not sensitive. Testing 0.1× and 10× the default (λ_align = dim(A)/1000 for DMC and dim(A)/50 for Humanoid-Bench) produced stable performance.
• State-of-the-art on the DeepMind Control Suite. BOOM reached an average TAR of 877.7, versus 835.8 for BMPC (+5.0%) and 745.6 for TD-MPC2 (+17.7%). It set new best results on every humanoid and dog task, leading Humanoid-run by +9.7% and Dog-run by +21.8% over the second-best method. SAC and DreamerV3 performed near zero on many of these tasks.
• State-of-the-art on Humanoid-Bench. BOOM reached an average TAR of 820.6, a +47.7% improvement over DreamerV3 with 10M iterations (555.6) and +60.5% over BMPC (511.7). It set records on all seven tasks, including +110.5% on H1hand-slide, +25.8% on H1hand-pole, and +121.0% on H1hand-hurdle.
• Theory supports the empirical claim. Theorem 1 bounds the return gap between behavior and policy: |J(β) − J(π)| ≤ R_max/(1−γ)·√(2ε) whenever KL(β‖π) ≤ ε. Theorem 2 bounds the Q-value gap by L_Q · D(ε) under a Lipschitz-continuity assumption, where D(ε) has an explicit form when MPPI's distribution is approximated as a Gaussian mixture and π is Gaussian.
• Two mechanisms explain improved world-model learning. Better distributional matching yields more policy-relevant value targets, and the resulting more accurate value gradients flow into the encoder, dynamics, and reward predictor because they are optimized jointly in the TD-style objective.
• Ablation scope. All three ablation studies were conducted on the single highest-dimensional DMC task, Dog-run; results are means over three runs with shaded confidence intervals.

Methodology in Plain English

BOOM keeps the standard TD-MPC2-style backbone — an encoder that maps observations to latent states, a latent dynamics model, a reward predictor, and a Q-value function — and trains all four jointly with a temporal-difference objective that also uses cross-entropy losses for rewards and values.

The novelty sits in how the policy is trained. During data collection, actions come from the planner (MPPI), which samples many action trajectories from a factorized Gaussian, scores them with predicted rewards plus a terminal value estimate, and reweights them by a softmax over returns. This produces good actions but an action distribution that is not a clean parametric form — so its likelihood cannot be computed.

Rather than trying to fit that intractable distribution directly, BOOM treats planner actions already stored in the replay buffer as imitation targets and simply maximizes their log-probability under the policy. This is equivalent to minimizing forward KL up to a constant the policy cannot influence. On top of that, each target is weighted by a softmax over its Q-value, so high-return actions dominate the gradient. The final policy loss adds this weighted alignment term to the ordinary Q-maximization term, scaled by λ_align. Conceptually, the policy and planner form a feedback loop: better alignment improves value estimation, and better values let the planner produce even better actions.

Why This Matters

Impact on research. The paper reframes a known but loosely characterized failure mode of planning-driven model-based RL as a testable distributional problem and offers both bounds and a simple fix. It also sharpens the distinction from offline RL: in BOOM, Q-maximization and imitation are complementary rather than conflicting, because planner actions are generally higher quality than policy actions — the opposite of the conservative tension that dominates offline RL. The authors also distinguish their objective from the concurrent TDM(PC)2 work, which uses a TD3+BC-style reverse KL, whereas BOOM is closer to AWAC with critic-guided weights and forward KL.

Real-world applications (as framed by the paper's motivating domains):

• Robotic locomotion control, including humanoid and quadruped platforms such as the Unitree H1hand used in the experiments.
• Autonomous driving, where look-ahead planning over learned dynamics is directly relevant.
• Industrial automation, where sample-efficient control matters when interaction is costly.
• Embodied intelligence more broadly, where agents must act in high-dimensional, long-horizon environments.

Industry relevance. The method is described as easy to implement and builds on an existing, widely used pipeline, so it can be added as a policy-loss modification rather than a new architecture. The demonstrated stability gains matter in practice: the paper contrasts BOOM's consistent curves with the oscillatory training of methods that ignore the Q-function during policy learning.

Future Directions

• Does the alignment idea extend beyond MPPI? The theoretical analysis approximates the planner's distribution as a Gaussian mixture; whether the guarantees and the empirical gains hold for other sampling-based planners is left open in the provided content.
• How should the weighting target be chosen? The authors note that Q can be replaced by other critics such as the state value V or advantage A and describe Q as merely the most convenient to access — an unexplored design axis.
• Can the alignment coefficient be set adaptively? The ablation shows robustness across 0.1× and 10× the default, but the paper does not report an automatic schedule, and the defaults differ between benchmark families (dim(A)/1000 for DMC, dim(A)/50 for Humanoid-Bench).
• Scaling and long-horizon error accumulation. BOOM shows strong gains on the 14 evaluated locomotion tasks, but the paper's conclusion is truncated in the provided content, so explicitly stated future work is not reported.

Target Audience

Reinforcement learning researchers and graduate students working on model-based RL, off-policy learning, or planning-based control will get the most from this paper, along with practitioners applying controller learning to high-dimensional robotics. Readers need prior exposure to actor-critic objectives and MPPI to follow the derivations, though the core idea — align the policy with the planner using value-weighted, likelihood-free imitation — is graspable without the proofs.

Authors’ abstract

Online planning has proven effective in reinforcement learning (RL) for improving sample efficiency and final performance. However, using planning for environment interaction inevitably introduces a divergence between the collected data and the policy's actual behaviors, degrading both model learning and policy improvement. To address this, we propose BOOM (Bootstrap Off-policy with WOrld Model), a framework that tightly integrates planning and off-policy learning through a bootstrap loop: the policy initializes the planner, and the planner refines actions to bootstrap the policy through behavior alignment. This loop is supported by a jointly learned world model, which enables the planner to simulate future trajectories and provides value targets to facilitate policy improvement. The core of BOOM is a likelihood-free alignment loss that bootstraps the policy using the planner's non-parametric action distribution, combined with a soft value-weighted mechanism that prioritizes high-return behaviors and mitigates variability in the planner's action quality within the replay buffer. Experiments on the high-dimensional DeepMind Control Suite and Humanoid-Bench show that BOOM achieves state-of-the-art results in both training stability and final performance. The code is accessible at https://github.com/molumitu/BOOM_MBRL.