Reproduce “A KL-regularization Framework for Learning to Plan with Adaptive Priors” on OpenResearch

About this paper

Álvaro Serra-Gomez, Daniel Jarne Ornia, Dhruva Tirumala, Thomas Moerland

This research introduces Policy Optimization–Model Predictive Control (PO-MPC), a framework that unifies and extends MPPI-based reinforcement learning by reformulating policy learning as a KL-regularized problem with an adaptive prior. The approach demonstrates improved performance and sample efficiency in high-dimensional continuous control tasks by mitigating policy mismatch and reducing gradient variance through a learned intermediate prior.

Key insights

• PO-MPC provides a unifying framework that shows existing state-of-the-art MPPI-based RL methods, such as TD-MPC2 and BMPC, as specific limiting cases of the proposed approach.
• Intermediate KL-regularization strengths (λ), rather than extreme values, yield more robust performance and enable learning in challenging high-dimensional environments where baseline methods struggle.
• Utilizing a learned intermediate policy prior significantly reduces the gradient variance in policy updates, which typically arises from using potentially outdated and multi-modal planner samples directly from the replay buffer.

Summary by alphaXiv.