Enhancing Diffuser for Decision-Making with Velocity Prediction, MPPI, and Replanning

The VanishingGrads team presents an enhanced approach to the Diffuser framework for offline decision-making tasks by incorporating velocity prediction, Model Predictive Path Integral (MPPI) sampling, and replanning strategies. Our goal is to adapt techniques that have proven effective in generative vision models, particularly velocity-based diffusion parameterization, into the domain of reinforcement learning and trajectory planning. We achieve approximately 50% reduction in planning horizon length without sacrificing trajectory quality, and our full pipeline achieves state-of-the-art performance in Maze2D tasks, particularly in U-Maze and Medium Maze environments.

• ›Velocity prediction integrated into the Diffuser denoising process for improved trajectory dynamics and temporal consistency.
• ›MPPI sampling computing normalised returns over multiple trajectories with softmax-weighted averaging for smoother, goal-directed outputs.
• ›Replanning module dividing the planning horizon into smaller segments to reduce model drift and compounding errors.
• ›~50% reduction in planning horizon length with state-of-the-art Maze2D performance.

We modify the Diffuser framework's denoising parameterization from predicting noise ε to predicting velocity v = αε − σx, inspired by Stable Diffusion 2.0. MPPI sampling generates N trajectory candidates, computes normalised returns, and forms a softmax-weighted average to produce the final plan. The replanning module splits the horizon H into segments of length K, re-running the denoising process at each segment boundary to correct accumulated drift.

Maze2D normalised scores at half the usual diffusion steps. All Diffuser variants outperform the classical baselines (MPPI, CQL, IQL) by a wide margin; our full pipeline improves on the original Diffuser by 21.1% on U-Maze and 14.5% on Medium.

• ›The architectural change alone (v = 0.0, still predicting noise) already beats baseline everywhere, so part of the gain is architectural rather than from the velocity parameterisation itself.
• ›Explicit velocity prediction (v = 0.5) helps in U-Maze and Medium but hurts in Large Maze (1.094 vs 1.230), suggesting it suits environments of moderate complexity rather than very large state spaces.
• ›MPPI contributes most in Medium Maze, a 5.2% improvement over velocity prediction alone.
• ›Replanning recovers from suboptimal initial plans, and matters most where the horizon is long enough for drift to accumulate.