Bootstrap Off-policy with World Model

TL;DR

BOOM introduces a novel framework that combines planning and off-policy reinforcement learning using a world model, leading to improved stability and performance in high-dimensional control tasks.

Contribution

It proposes a bootstrap loop integrating planning and off-policy learning with a jointly learned world model, enhancing sample efficiency and stability.

Findings

Achieves state-of-the-art results on DeepMind Control Suite.

Improves training stability and final performance.

Effectively handles high-dimensional control tasks.

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.