A KL-regularization framework for learning to plan with adaptive priors

Effective exploration remains a key challenge in model-based reinforcement learning (MBRL), especially in high-dimensional continuous control tasks where sample efficiency is critical. Recent work addresses this by using learned policies as proposal distributions for Model-Predictive Path Integral (MPPI) planning. Early approaches update the sampling policy independently of the planner, typically via deterministic policy gradients with entropy regularization. However, since the data distribution is induced by the MPPI planner, misalignment between the policy and planner degrades value estimation and long-term performance. To address this, recent methods explicitly align the policy with the planner by minimizing KL divergence to the planner distribution or by incorporating planner-guided regularization. In this work, we unify these approaches under the Policy Optimization–Model Predictive Control (PO-MPC) framework, a family of KL-regularized MBRL methods that treat the planner’s action distribution as a prior in policy optimization. We show how existing methods emerge as special cases of this family and explore previously unstudied variants. Experiments demonstrate that these variants yield significant performance gains, advancing the state of the art in MPPI-based RL.

Many modern AI systems learn by trial and error, but they often struggle to explore efficiently in complex environments such as robotics or simulated control tasks. One promising approach combines reinforcement learning with online planning, where the AI uses a learned model of the world to test possible future actions generated by the learned agent before acting. However, existing methods can become unstable because the learned agent and the planner gradually drift apart, causing poor long-term predictions and unreliable learning. We introduce PO-MPC, a framework that keeps the trajectories generated with the learned agent aligned with those of the planner using a mathematical regularization technique based on KL divergence. Our method unifies several previously separate approaches under one framework and reveals new variants that had not been explored before. We also propose learning an adaptive “planner prior” that reduces instability caused by outdated planning information. Across challenging continuous-control benchmarks, including high-dimensional humanoid tasks, PO-MPC consistently improves learning speed and final performance over prior state-of-the-art methods.