Analytic Planning under Uncertainty with Moment Closure

2026-08-03Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors study how to improve decision-making in unpredictable environments using model-based reinforcement learning. They propose a new method that considers both the average expected outcome and the uncertainty of possible outcomes when planning. Their approach combines a mathematical model of transitions with a special way to represent value, allowing for exact calculation without relying on random sampling or ignoring uncertainty. Experiments show their method better captures uncertainty and reduces noise in learning signals, helping make more reliable plans.

model-based reinforcement learningstochastic environmentspredictive uncertaintyBellman backupquadratic action-valueGaussian transition modelradial-basis functionvalue functiontarget variancecontinuous control
Authors
Shishir Sharma, Doina Precup
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.