Technical Summary: A KL-Regularization Framework for Learning to Plan with Adaptive Priors

1. Problem Statement

In Model-Based Reinforcement Learning (MBRL), effective exploration remains a critical challenge, particularly in high-dimensional continuous control tasks where sample efficiency is paramount. Recent approaches, such as TD-MPC and BMPC, utilize learned policies as proposal distributions for Model-Predictive Path Integral (MPPI) planning. However, a fundamental limitation arises from policy mismatch: the sampling policy ($\pi_{\theta_s}$) is often updated independently of the planner's action distribution ($\pi_P$).

This decoupling creates a distribution mismatch where the bootstrap action-value function is trained on trajectories generated by the MPPI planner (which explores a specific distribution), while the sampling policy is updated to maximize local action values, often diverging from the planner's distribution. This misalignment leads to unreliable value estimates, particularly for short horizons, and degrades long-term performance. Previous attempts to align these distributions, such as minimizing reverse KL divergence (BMPC), suffer from high variance because they rely on "stale" planner samples stored in the replay buffer, which represent a mixture of outdated planning policies rather than a coherent current distribution.

2. Methodology: PO-MPC Framework

The authors propose Policy Optimization–Model Predictive Control (PO-MPC), a unified framework that treats the planner's action distribution as a prior in policy optimization. PO-MPC casts the sampling policy learning step as an instance of KL-regularized Reinforcement Learning.

Core Formulation

The objective is to learn a sampling policy $\pi_{\theta_s}$ that maximizes expected return while remaining close to an adaptive prior $\pi_p$ (derived from the MPPI planner). The objective function is:

$$J(\pi_{\theta_s}) = \mathbb{E} \left[ \sum_{t=0}^{T-1} \gamma^t \left( r(z_t, a_t) - \lambda \text{KL}[\pi_{\theta_s}(\cdot | z_t) \| \pi_p(\cdot | z_t)] \right) \right]$$

Where:

• $z_t$ is the latent state representation.
• $\lambda$ is a hyperparameter controlling the strength of the KL regularization.
• $\pi_p$ is the policy prior.

Key Components

1. KL-Regularized RL: The framework unifies existing methods as special cases of this formulation:

   • $\lambda = 0$: Recovers TD-MPC2, where the policy is updated solely via action-value maximization and entropy regularization.
   • $\lambda \to \infty$: Recovers BMPC (with reverse KL), where the policy is constrained to mimic the planner's distribution.
   • Intermediate $\lambda$: The paper explores intermediate values (e.g., $\lambda \in [0.1, 9]$) to balance return maximization and distribution alignment.

2. Adaptive Intermediate Prior ($\pi_{\theta_p}$):
   To address the high variance caused by using stale planner samples directly from the replay buffer, PO-MPC introduces a learned intermediate prior. Instead of cloning the raw distribution of stored planner samples (which may be a time-varying Gaussian mixture), the framework trains a neural network $\pi_{\theta_p}$ to approximate the current MPPI-induced distribution $\pi_P$.

   • This prior is distilled from the planner's output, shielding the sampling policy updates from the variance of outdated samples.
   • The prior can be trained using either Reverse KL (minimizing $\text{KL}[\pi_{\theta_p} \| \pi_P]$) or Forward KL (minimizing $\text{KL}[\pi_P \| \pi_{\theta_p}]$).
     • Reverse KL: Biases the policy toward matching specific modes of the planner (faster convergence, potentially less exploration).
     • Forward KL: Biases the policy toward covering the support of all sampled distributions (enhanced exploration, potentially slower convergence).

3. Algorithm Loop:
   The training follows a Plan $\to$ Infer $\to$ Regularize cycle:

   • Plan: MPPI generates a trajectory distribution $\pi_P$ using the current sampling policy and value function.
   • Infer: The adaptive prior $\pi_{\theta_p}$ is updated to approximate $\pi_P$.
   • Regularize: The sampling policy $\pi_{\theta_s}$ is updated using the KL-regularized objective, balancing the Q-value maximization and the KL penalty toward $\pi_{\theta_p}$.

3. Key Contributions

• Unification: The paper unifies disparate MPPI-based RL methods (TD-MPC2, BMPC) under a single theoretical framework (PO-MPC), clarifying their design choices and relationships.
• Adaptive Priors: It introduces a learned intermediate prior to replace direct cloning of replay buffer samples, theoretically and empirically demonstrating that this reduces gradient variance and stabilizes training.
• Flexible Objectives: The framework enables the exploration of previously unstudied variants, specifically tuning the KL strength $\lambda$ and choosing between Forward/Reverse KL losses for the prior.
• State-of-the-Art Performance: The method achieves superior sample efficiency and final performance in high-dimensional continuous control tasks compared to existing baselines.

4. Experimental Results

The authors evaluated PO-MPC on 21 challenging continuous control tasks from the DeepMind Control Suite (Humanoid and Dog) and HumanoidBench.

• Performance Gains: PO-MPC with intermediate $\lambda$ values consistently outperformed both TD-MPC2 ($\lambda=0$) and BMPC ($\lambda \to \infty$).
  • In high-dimensional tasks (e.g., HumanoidBench), intermediate $\lambda$ values (particularly $\lambda=1$ or $\lambda=9$) yielded significant improvements in final return and sample efficiency.
  • In some environments where baselines failed to learn (e.g., HumanoidBench "Balance Simple," "Slide," "Stair"), PO-MPC successfully converged.
• Intermediate Prior Efficacy: Using the learned intermediate prior ($\pi_{\theta_p}$) matched or surpassed the performance of using raw planner samples from the replay buffer, confirming the reduction in update variance.
• Loss Function Sensitivity: The choice of KL loss for training the prior mattered. Forward KL was beneficial in exploration-heavy tasks (e.g., "Stair"), while Reverse KL accelerated convergence in precision-critical tasks (e.g., "Balance Simple").
• Ablation: The results showed that performance degrades if $\lambda$ is too low (policy mismatch returns) or too high (premature mode collapse), highlighting the importance of tuning this trade-off.

5. Significance and Claims

The paper claims that PO-MPC represents a principled unification of MPPI-based approaches. By explicitly aligning the sampling policy with the planner via a KL-regularized objective and an adaptive prior, the framework:

1. Closes the loop between planning and policy learning, ensuring the value function is evaluated on distributions relevant to the current planner.
2. Improves stability by mitigating the variance introduced by outdated planning samples in the replay buffer.
3. Advances the state of the art in model-based RL for continuous control, offering a flexible framework that can be tuned to specific task requirements (exploration vs. exploitation).

The authors note that while the framework is powerful, it assumes Gaussian distributions for policies and does not fully exploit all simulated trajectories for value learning. Future work is suggested to extend policy distributions beyond Gaussians and automatically tune the $\lambda$ hyperparameter.