Analytic Planning under Uncertainty with Moment Closure

2026-08-03

AI Summary

This work addresses the challenge of analytically modeling and propagating predictive uncertainty in model-based reinforcement learning within stochastic environments, without relying on restrictive policy or reward assumptions. The authors introduce a compatibility principle that reformulates the Bellman backup—parameterized by a quadratic action-value function—as an expectation solely over the state-value function. By leveraging the moment-closure properties of Gaussian transition dynamics and radial basis function value representations, they derive a closed-form expression for this expectation. This approach enables, for the first time, exact analytical propagation of both predictive mean and covariance without resorting to sampling or deterministic approximations. Empirical results in continuous control tasks demonstrate significantly reduced target variance and well-calibrated uncertainty estimates, thereby validating the efficacy of analytical distribution-aware planning.

Abstract

Effective model-based reinforcement learning in stochastic environments requires planning that accounts for predictive uncertainty. Propagating full state distributions analytically offers a principled way to do this, but has traditionally required restrictive policy or reward structures to remain tractable. Consequently, modern deep reinforcement learning has largely retreated to either stochastic sampling, which introduces significant target variance, or deterministic point estimates that ignore predictive covariance entirely. We investigate whether distribution-aware planning is possible without these constraints. Using a quadratic action-value parameterization, we first reduce the Bellman backup to an expectation over the state-value function alone; the key idea is then a compatibility principle between the predictive transition distribution and the value function class, under which this expectation is analytic in the distribution's moments. We instantiate this principle with a Gaussian transition model paired with a radial-basis value function, yielding a closed-form backup that propagates both predictive mean and covariance. Empirically, our approach reduces target variance and yields well-calibrated predictive uncertainty under stochastic observations in continuous control, providing a principled framework for planning with learned distribution models.

Problem

• planning under uncertainty
• distribution-aware planning
• predictive uncertainty
• model-based reinforcement learning
• moment propagation

Innovation

• Moment Closure
• Distribution-Aware Planning
• Bellman Backup
• Gaussian Transition Model
• Radial-Basis Value Function