Paper deep dive
Analytic Planning under Uncertainty with Moment Closure
Shishir Sharma, Doina Precup
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2608.02519v1
- Canonical: https://arxiv.org/abs/2608.02519v1
Trouble viewing inline? Open PDF directly →
Full Text
51,263 characters extracted from source content.
Expand or collapse full text
Analytic Planning under Uncertainty with Moment Closure Shishir Sharma McGill University, Montreal, Canada Mila – Quebec Artificial Intelligence Institute, Montreal, Canada Doina Precup McGill University, Montreal, Canada Mila – Quebec Artificial Intelligence Institute, Montreal, Canada 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. Figure 1: Analytic versus Monte Carlo estimation of the Bellman expectation [V(s′)]E[V(s )], on a one-dimensional toy example with a Gaussian predictive distribution p(s′)p(s ) and a radial-basis value function V(s′)V(s ). (A) The expectation is the integral of the density-weighted value p(s′)V(s′)p(s )\,V(s ) (shaded); Here equal to [V(s′)]=0.057E[V(s )]=0.057. (B) A Monte Carlo estimate from K samples is unbiased but is high-variance, especially when V is peaked. (C) For a choice between two actions whose predictive means differ ([A]=0.057E[A]=0.057, [B]=0.025E[B]=0.025, so A is better), sampling noise leads to the wrong action being chosen with higher probability at smaller K, whereas the analytic expectation always selects correctly. 1 Introduction Recent advances in model-based reinforcement learning (MBRL) increasingly revolve around world models that enable imagination-based value learning and planning in either state space or learned latent spaces [gu2016continuous, chua2018deep, janner2019mbpo, hafner2019learning, Hafner2020Dream, hansen2022temporal, Hafner2025]. A common pattern in these systems is that Bellman targets (or policy updates) are computed from sampled imagined rollouts or trajectory optimization under the learned dynamics, which implicitly propagates predictive uncertainty through Monte Carlo variation. While this approach scales to expressive models, it makes value estimation inherently stochastic and often sensitive to sampling budgets and target variance. Non-sampling approaches avoid this variance but have their own limitations. Methods that collapse a stochastic model to its mean, or learn a deterministic model directly, discard predictive covariance and can be systematically biased when value functions are nonlinear. On the other hand, methods that propagate the full predictive distribution analytically can often require strong structural constraints, for instance on the policy or reward, to remain tractable. Our focus is distinct from both mean-substitution and rollout-style uncertainty propagation. Rather than propagating distributions over trajectories, we integrate stochastic one-step transitions directly inside the Bellman backup used for value-based control. The question we address is whether the Bellman expectation itself can be computed analytically under learned stochastic dynamics. We show that, under a quadratic action-value parameterization, closed-form Bellman expectations are possible when the transition model and value representation are moment-compatible. We instantiate this principle with a Gaussian dynamics model that predicts mean and covariance, paired with a value function whose expectation under a Gaussian distribution is analytic, resulting in a closed-form expression for the Bellman target. This removes Monte Carlo sampling from target construction while preserving propagation of predictive variance. Concretely, we learn a heteroscedastic Gaussian transition model via maximum likelihood to capture aleatoric uncertainty and use an ensemble, combined via moment matching, to represent epistemic uncertainty. Greedy action selection is made tractable separately, via a quadratic action-value parameterization that removes the inner maximization from the Bellman backup; uncertainty then enters exclusively through the analytic Bellman expectation. The value function is parameterized as a mixture of Gaussian radial basis functions, whose expectation under the predictive Gaussian is available in closed form. Empirically, analytic greedy planning remains stable as observation noise increases, whereas sampling-based and mean-only baselines degrade. Ablations show that our moment-matched predictive distribution stably captures the combined uncertainty, and that this robustness stems from planning efficiently under that uncertainty rather than from increased Monte Carlo budgets. 2 Related Work MBRL under uncertainty has been approached from multiple directions, differing primarily in how stochastic dynamics influence value estimation and planning. Most contemporary deep MBRL systems estimate value targets by sampling imagined trajectories from learned dynamics models, using particle rollouts or model predictive control-style shooting methods [chua2018deep, janner2019mbpo, hafner2019learning, Hafner2020Dream]. These approaches scale well and accommodate expressive neural models, but Monte Carlo sampling introduces variance directly into Bellman targets. Stability often depends on rollout depth, particle count, or model bias. At the other end, expectation-style planning constructs targets through a deterministic dynamics representation — either by collapsing a stochastic model to its predicted mean, or by learning a deterministic model directly [nagab2017neural, wan2019planning]. Such approaches simplify target computation and preserve determinism, but discard predictive covariance. Probabilistic dynamics models and deep ensembles are widely used to represent aleatoric and epistemic uncertainty [lakshminarayanan2017simple]. In online MBRL, ensembles mitigate model bias and improve exploration [chua2018deep]. However, even when models are probabilistic, value targets are typically constructed either by sampling ensemble members or by collapsing ensemble predictions to point estimates. Analytic uncertainty propagation offers a third perspective. In control and Gaussian process-based RL, moment matching has been used to propagate state distributions through nonlinear dynamics for policy search [deisenroth2011pilco]. More recently, deterministic approximations have been proposed to reduce Monte Carlo noise in deep value updates by propagating Gaussian approximations through critics [akgul2024deterministic]. These approaches typically approximate intermediate activations or rely on multi-step rollout evaluation. Our work instead targets the Bellman backup itself in value-based deep MBRL, decoupling the tractability of greedy action selection from the tractability of the predictive expectation. Quadratic Q-function parameterizations enable closed-form greedy action selection in continuous control [gu2016continuous]; we leverage this structure so that the maximization disappears from the backup entirely, leaving an expectation over the state-value function alone. This decoupling requires tractability only of the value representation, allowing analytic, uncertainty-aware planning without restricting the policy class or reward structure. 3 Analytic Greedy Planning via Moment Closure 3.1 Motivation To illustrate the impact of sampling noise in Bellman backups, we construct a one-dimensional toy example. Let the predictive distribution p(s′)p(s ) be Gaussian and the value function V(s′)V(s ) be a single Gaussian radial basis function. The expectation is then the integral of their pointwise product p(s′)V(s′)p(s )\,V(s ), shown shaded in Figure 1(A). When both p and V are Gaussian-shaped, this integral admits a simple closed form, so the exact expectation is available without any sampling. A sampling-based planner instead draws K next states s1′,…,sK′∼p(s′)s _1,…,s _K p(s ) and averages the value at those points, 1K∑iV(si′) 1K _iV(s _i). This estimator is unbiased, but its variance can be severe, especially when V is sharply peaked relative to p, i.e. when only a narrow band of next states is highly valued (imagine the edge of a cliff, where a small region of states is safe and the rest are not). In this regime most sampled states fall where V≈0V≈ 0 and contribute nothing, while the estimate is carried by the rare samples that land near the peak. Figure 1(B) shows the consequence: at small K the Monte Carlo estimate scatters widely around the true expectation (dashed line), tightening only as K grows. This variance can easily corrupt control decisions. Greedy action selection does not require the expected value itself, only the ranking of actions by expected value, and sampling noise can flip that ranking. To show this, we consider two actions whose predictive distributions differ only in their means, so that one action is genuinely better (higher true expected value) than the other. A Monte Carlo planner selects whichever action has the higher sampled average; Figure 1(C) reports how often this disagrees with the correct choice. At small K the wrong action is selected a substantial fraction of the time, and the error rate falls toward zero only as K increases. The analytic expectation, computed exactly, never makes this error and removes this source of instability entirely. 3.2 Preliminaries We consider the standard reinforcement learning setting formalized as a Markov decision process (MDP) [sutton2018reinforcement], defined by the tuple (,,p,r,γ),(S,A,p,r,γ), where ⊂ℝdsS ^d_s is a continuous state space, ⊂ℝdaA ^d_a is a continuous action space, p(s′|s,a)p(s |s,a) is the transition function, r(s,a)r(s,a) is a bounded reward function, and γ∈[0,1)γ∈[0,1) is a discount factor. For a policy π, the action-value function is Qπ(s,a)=π[∑k=0∞γkr(St+k,At+k)|St=s,At=a]Q^π(s,a)=E_π [ _k=0^∞γ^kr(S_t+k,A_t+k)\; |\;S_t=s,A_t=a ] The optimal action-value function satisfies the Bellman optimality equation Q∗(s,a)=s′∼p(⋅∣s,a)[r(s,a)+γmaxa′∈Q∗(s′,a′)]Q^*(s,a)=E_s p(· s,a) [r(s,a)+γ _a Q^*(s ,a ) ] (1) Learned stochastic dynamics In MBRL, the true transition density p(s′|s,a)p(s |s,a) is replaced by a learned model p^θ(s′|s,a) p_θ(s |s,a). When p^θ p_θ is stochastic, the Bellman target in (1) requires computing an expectation over next states under the learned predictive distribution. 3.3 Decoupling Greedy Action Selection from Expectation Representation challenge A sampling-free Bellman backup must compute s′∼p^θ(⋅∣s,a)[maxa′∈Q(s′,a′)]E_s p_θ(· s,a) [ _a Q(s ,a ) ] The difficulty arises because the Bellman expectation involves the greedy value maxa′Q(s′,a′) _a Q(s ,a ) evaluated at each possible next state s′s . Even if Q(s′,a′)Q(s ,a ) is smooth, the maximization over actions produces a highly nonlinear function of s′s . As a result, even under Gaussian predictive dynamics, the Bellman expectation is generally intractable for expressive deep Q-functions. One classical route to tractability is to impose strong structural assumptions. In linear–quadratic settings (LQR/LQG), the Q-function is quadratic jointly in the state and action, yielding closed-form greedy actions and analytic expectations under Gaussian disturbances [anderson2007optimal]. Similarly, trajectory-optimization methods such as iterative LQG or differential dynamic programming maintain local quadratic approximations to the cost-to-go to obtain analytic feedback updates [todorov2005generalized, tassa2012synthesis]. However, enforcing this joint quadratic structure in state and action directly on a deep Q-function would significantly restrict representational capacity. Key observation: decoupling maximization from expectation The intractability of the Bellman expectation stems from integrating the greedy value maxa′Q(s′,a′) _a Q(s ,a ) under a stochastic transition. If the maximization can be removed from inside the expectation, the backup reduces to evaluating s′∼p^θ(⋅∣s,a)[V(s′)],E_s p_θ(· s,a)[V(s )], which is substantially simpler. We achieve this decoupling using a quadratic action-value parameterization (NAF-style) [gu2016continuous]. Specifically, we decompose Q(s,a)=V(s)−12(a−μ~(s))⊤P(s)(a−μ~(s)),Q(s,a)=V(s)- 12(a- μ(s)) P(s)(a- μ(s)), (2) where P(s)≻0P(s) 0. Crucially, unlike the classical quadratic-in-state-and-action routes above, none of V(s)V(s), μ~(s) μ(s), or P(s)P(s) are restrained to be quadratic: only the action-slice at each fixed s is quadratic, so Q retains its full flexibility in how it varies with s. In the unconstrained case =ℝdaA=R^d_a, greedy action selection is available in closed form, argmaxaQ(s,a)=μ~(s),maxaQ(s,a)=V(s). _aQ(s,a)= μ(s), _aQ(s,a)=V(s). Substituting this structure into the Bellman equation removes the inner maximization entirely, so that the backup value depends only on the state-value term: y(s,a)=r(s,a)+γs′∼p(⋅∣s,a)[V(s′)]y(s,a)=r(s,a)+ _s p(· s,a)[V(s )] This decomposition has a consequence central to our approach. The value part V(s)V(s) and the policy part μ~(s) μ(s) play separate roles: uncertainty enters the backup only through the expectation of V(s′)V(s ), so it is the value representation alone that must be made compatible with the predictive distribution, while the greedy policy μ~(s) μ(s) carries no such compatibility requirement. In analytic planning with distribution models, the two are typically tied together, with the policy class restricted so that the propagated state distribution stays tractable. Here the tractability requirement falls entirely on the value representation, freeing the policy from any such constraint: we obtain analytic uncertainty propagation and an unrestricted, network-parameterized policy at the same time. Bounded actions In practical continuous-control settings, actions are typically bounded (e.g., =[−1,1]daA=[-1,1]^d_a), and the unconstrained maximizer μ~(s) μ(s) need not lie in A. Following the bounded NAF (BNAF) construction of plaksin2022continuous, we define the greedy action as the projection μ(s)=Π(μ~(s))μ(s)= _A( μ(s)) and shift the quadratic by a state-dependent constant, Q(s,a)=V(s) Q(s,a)=V(s) −12(a−μ~(s))⊤P(s)(a−μ~(s)) - 12(a- μ(s)) P(s)(a- μ(s)) +12(μ(s)−μ~(s))⊤P(s)(μ(s)−μ~(s)) + 12(μ(s)- μ(s)) P(s)(μ(s)- μ(s)) (3) This preserves quadratic structure while ensuring argmaxa∈Q(s,a)=μ(s),maxa∈Q(s,a)=V(s) _a Q(s,a)=μ(s), _a Q(s,a)=V(s) Consequently, even under bounded actions, the Bellman backup reduces to computing s′∼p(⋅∣s,a)[V(s′)]E_s p(· s,a)[V(s )] such that uncertainty influences planning exclusively through the expectation of the state-value function. For the class of RL problems arising from time-discretized optimal control with control-affine dynamics and control-quadratic costs, plaksin2022continuous show that this bounded family is expressive enough to solve the Bellman optimality equation to arbitrary accuracy as the discretization is refined, and that any such approximate solution yields a greedy policy with correspondingly bounded suboptimality. They further show that the original vertex-constrained NAF family does not admit this guarantee, which motivates our use of the bounded construction. 3.4 Stochastic Dynamics Model We learn a stochastic dynamics model as a conditional Gaussian, pθ(s′∣s,a)=(s′;μθ(s,a),Σθ(s,a)),p_θ(s s,a)=N\! (s ; _θ(s,a), _θ(s,a) ), where μθ:×→ℝds _θ:S×A ^d_s predicts the conditional mean and Σθ:×→ℝds×ds _θ:S×A ^d_s× d_s predicts the conditional covariance. The covariance Σθ(s,a) _θ(s,a) captures aleatoric uncertainty, corresponding to irreducible stochasticity in the environment. Ensemble-based epistemic uncertainty To capture epistemic uncertainty arising from limited data, instead of a single model, we train an ensemble of K independently initialized copies (μθk,Σθk)k=1K\( _ _k, _ _k)\_k=1^K. The ensemble induces a uniformly weighted mixture of Gaussian transition models, p(s′∣s,a)=1K∑k=1K(s′;μθk(s,a),Σθk(s,a))p(s s,a)= 1K _k=1^KN\! (s ; _ _k(s,a), _ _k(s,a) ) Rather than sampling ensemble members during planning, we approximate this mixture by a single Gaussian whose mean and covariance match the first two moments of the mixture, following lakshminarayanan2017simple. Using the law of total expectation and total variance, the moment-matched mean is μ¯(s,a)=1K∑k=1Kμθk(s,a), μ(s,a)= 1K _k=1^K _ _k(s,a), and the predictive covariance decomposes as Σ¯(s,a)= (s,a)= 1K∑k=1KΣθk(s,a) 1K _k=1^K _ _k(s,a) + + 1K∑k=1KΔμk(s,a)Δμk(s,a)⊤, 1K _k=1^K _k(s,a)\, _k(s,a) , where Δμk(s,a)= _k(s,a)= μθk(s,a)−μ¯(s,a) _ _k(s,a)- μ(s,a) (4) The first term in Σ¯(s,a) (s,a) is the ensemble-averaged aleatoric covariance from each member. The second term captures epistemic uncertainty, reflecting divergence across ensemble members [kendall2017uncertainties]. Because members are trained on independent bootstrap resamples, this divergence is a genuine disagreement signal, inflating the predictive covariance in regions where ensemble members disagree. We therefore plan under the approximation p(s′∣s,a)≈(s′;μ¯(s,a),Σ¯(s,a)),p(s s,a)\;≈\;N\! (s ; μ(s,a), (s,a) ), without sampling ensemble members. Under this model, the Bellman target becomes y(s,a)=r(s,a)+γs′∼(μ¯(s,a),Σ¯(s,a))[V(s′)]y(s,a)=r(s,a)+ _s ( μ(s,a), (s,a))\! [V(s ) ] (5) The remaining challenge is therefore to evaluate [V(s′)]E[V(s )] in closed form. 3.5 Moment-Compatible Value Representation The tractability of the analytic backup rests on a compatibility condition between the predictive transition distribution and the value representation. Let p(s′|s,a)p(s |s,a) be the predictive distribution with parameters θp _p. We say the value class is moment-compatible with p if the expectation s′∼p[V(s′)]E_s p[V(s )] admits a closed-form expression in θp _p. Whenever this holds, the state-value expectation in the Bellman backup can be evaluated analytically, with no sampling. Moment compatibility is a design principle rather than a prescription of any single functional form. Several pairings of predictive distribution and value class could satisfy it, such as polynomial value functions with elliptical distributions (e.g. Gaussian, Laplace, Student-t) or Random Fourier Features with a Gaussian predictive distribution. For concreteness and simplicity, we parameterize the state-value function as a mixture of Gaussian radial basis functions, V(s)=∑i=1Mwiexp(−12(s−ci)⊤Λi(s−ci)),V(s)= _i=1^Mw_i \! (- 12(s-c_i) _i(s-c_i) ), (6) where Λi≻0 _i 0. The centers cic_i and precision matrices Λi _i are fixed throughout training (see Appendix B.4-B.5 for initialization details); only the weights wi\w_i\ are learned, updated via temporal-difference learning as part of the critic optimization (Algorithm 1, line 12). Under a Gaussian next-state distribution s′∼(μ,Σ)s (μ, ), the expectation of each basis function can be computed exactly. Writing ‖x−a‖B2:=(x−a)⊤B(x−a)\|x-a\|_B^2:=(x-a) B\,(x-a) and Ai:=Λi(I+ΣΛi)−1A_i:= _i(I+ _i)^-1, [V(s′)]=∑i=1Mwi|I+ΣΛi|−1/2exp(−12‖μ−ci‖Ai2)E[V(s )]= _i=1^Mw_i\,|I+ _i|^-1/2 \! (- 12\|μ-c_i\|_A_i^2 ) (7) The full derivation is given in Appendix B.2. Algorithm 1 Moment-Compatible Analytic Planning 1:Initialize twin critics ϕ1,ϕ2 _1, _2 (each with Vϕ,μϕV_φ, _φ) and target parameters ϕ1−,ϕ2− _1^-, _2^-, ensemble dynamics models pθkk=1K\p_ _k\_k=1^K, replay buffer D, update counter n←0n← 0 2:for each environment step do 3: With probability p, a∼Uniform()a (A); otherwise a=Π(μϕ1(s)+ϵ)a= _A( _ _1(s)+ε), ϵ∼(0,σ2I)ε (0,σ^2I) ⊳ exploration 4: Execute a, observe (r,s′,d)(r,s ,d), and store (s,a,r,s′,d)(s,a,r,s ,d) in D 5: for each gradient step do 6: Sample minibatch ℬ⊂B 7: for k=1k=1 to K do 8: Update dynamics parameters θk _k via NLL on bootstrap resample of ℬB 9: end for 10: Compute moment-matched predictive Gaussian (μ¯,Σ¯)( μ, ) from (μθk,Σθk)k=1K\( _ _k, _ _k)\_k=1^K using Eq. (4) 11: Compute target y←r+γ(1−d)mini=1,2s′∼(μ¯,Σ¯)[Vϕi−(s′)]y← r+γ(1-d) _i=1,2E_s ( μ, )[V_ _i^-(s )] 12: Update critics ϕ1,ϕ2 _1, _2 by minimizing (Qϕi(s,a)−y)2(Q_ _i(s,a)-y)^2 on ℬB 13: n←n+1n← n+1; every NtargetN_target updates, ϕi−←ϕi _i^-← _i 14: end for 15:end for 4 Learning Procedure We introduce Moment-Compatible Analytic Planning (MoCA) and detail the overall training procedure (Algorithm 1). We integrate the analytic Bellman backup into a standard off-policy learning loop with a learned stochastic dynamics ensemble. The agent alternates between collecting transitions, updating the dynamics model, and updating value critics using analytic targets. Exploration Actions are selected from the greedy policy implied by the quadratic critic, with injected exploration noise. Concretely, we perturb the greedy action with additive Gaussian noise and clip to the feasible action set. In addition, with a fixed probability we instead take a uniformly random action (also clipped). This mirrors the exploration mechanism used in Spinning Up’s TD3 implementation [SpinningUp2018], combining local exploration around the greedy action with occasional global exploration. Dynamics model Given a replay buffer =(st,at,st+1,rt,dt)D=\(s_t,a_t,s_t+1,r_t,d_t)\, we train each member of a Gaussian dynamics ensemble s′∼(μθk(s,a),Σθk(s,a))s ( _ _k(s,a), _ _k(s,a)) by minimizing the negative log-likelihood. Models are trained independently on bootstrap resamples of the sampled minibatch ℬB. At target-construction time, ensemble predictions are aggregated via moment matching according to (4), yielding a single predictive Gaussian (μ¯(s,a),Σ¯(s,a))N( μ(s,a), (s,a)). Twin value critics and analytic targets To mitigate overestimation bias, we adopt clipped double estimation (as in TD3), adapted to our setting where greedy control reduces the backup to a state-value term. We maintain two critics, each inducing a state-value function Vϕ1V_ _1 and Vϕ2V_ _2 (with corresponding target parameters ϕ1−,ϕ2− _1^-, _2^-). For a transition (st,at,rt,dt)(s_t,a_t,r_t,d_t), we compute two analytic expected next-values under the moment-matched predictive Gaussian: vi=s′∼(μ¯(st,at),Σ¯(st,at))[Vϕi−(s′)],i∈1,2,v_i=E_s ( μ(s_t,a_t), (s_t,a_t))\! [V_ _i^-(s ) ], i∈\1,2\, where the expectation is evaluated in closed form via V’s moment-compatible representation (Section 3.5). The Bellman target then uses clipped double estimation: yt=rt+γ(1−dt)min(v1,v2)y_t=r_t+γ(1-d_t)\, (v_1,v_2) Each critic is updated by minimizing a squared-error regression to this shared target. Target parameters are refreshed periodically by copying the online critic parameters. Action bounds and model evaluation A related and well-known failure mode in model-based RL is model exploitation (or planner overfitting), where planning discovers actions that exploit inaccuracies of the learned dynamics in poorly supported regions of the state–action space, leading to overly optimistic value estimates. In our setting, this risk is particularly transparent because analytic expectation removes Monte Carlo variability from target construction: any optimism induced by off-support model queries is no longer masked by sampling noise and can appear as systematic bias. For this reason, we restrict model evaluation to observed (s,a)(s,a) pairs when constructing Bellman targets and avoid Dyna-style synthetic rollouts; greedy actions are projected onto the feasible action set via the BNAF construction, and uncertainty enters only through the closed-form expectation of V(s′)V(s ). If one were to extend the method to Dyna-style synthetic experience generation, the BNAF correction term in (3) would need to be incorporated inside the expectation. Analytic treatment under bounded actions would likely require truncated Gaussian moments, which we leave for future work. Figure 2: Learning curves on Cartpole under Gaussian observation noise. Top row: homoscedastic noise with standard deviation σ∈0.025,0.05σ∈0.025,0.05. Bottom: heteroscedastic noise with σmin=0.0 _ =0.0, σmax=0.05 _ =0.05. Curves show mean return over 10 seeds; shaded regions denote ± 1 standard deviation. Figure 3: Return and uncertainty diagnostics on Pendulum under Gaussian observation noise (σ∈0.025,0.05σ∈\0.025,0.05\, left and right columns). Top: evaluation return during training. Middle: for MoCA, empirical 1σ1σ coverage remains near the nominal 68% level throughout training. Bottom: epistemic standard deviation for MoCA, which stabilizes as the ensemble converges. 5 Empirical Study 5.1 Implementation Details Our model-free baseline follows the continuous Normalized Advantage Function (NAF) formulation of gu2016continuous and its bounded-action analysis in plaksin2022continuous. We adopt the faithful bounded variant (BNAF), in which the quadratic vertex μ~(s)∈ℝA μ(s) ^A is learned without squashing and only clipped to the environment action limits at evaluation time. The Q-function is parameterized so that Q(s,π(s))=V(s)Q(s,π(s))=V(s) for the clipped greedy action π(s)=clip(μ~(s))π(s)=clip( μ(s)), eliminating the saturation pathologies that arise from naïve tanh-constrained quadratic parameterizations. The critic uses a shared MLP trunk with two heads: (i) an unconstrained vertex head μ~(s) μ(s) and (i) a Cholesky head L(s)L(s) defining a positive-definite curvature matrix P(s)=L(s)L(s)⊤P(s)=L(s)L(s) . The value function V(s)V(s) is parameterized linearly in fixed radial basis features. To ensure strict positive definiteness and prevent curvature collapse during training, we enforce a minimum-curvature constraint on P(s)P(s) via the Cholesky parameterization.111We also apply a curvature keep-alive penalty that discourages near-flat quadratic advantages by adding λ[ReLU(pmin−‖L(s)‖F2)2]λ\,E[ReLU(p_ -\|L(s)\|_F^2)^2] to the TD loss, where ‖L‖F2=tr(P)\|L\|_F^2=tr(P). See Appendix C. Training uses twin critics with clipped double targets, i.e., Bellman targets are constructed using minV1−(s′),V2−(s′) \V_1^-(s ),V_2^-(s )\ from periodically updated target networks. Optimization is performed with Adam on minibatches sampled from a replay buffer. 5.2 Noise Models Our objective is to evaluate whether explicitly modeling predictive uncertainty improves planning performance under stochastic observations. To isolate the effect of uncertainty propagation in value estimation, we introduce controlled observation noise processes while keeping the underlying system dynamics unchanged. A key consideration is avoiding transition-model misspecification. Since our learned dynamics model assumes Gaussian structure, we restrict injected stochasticity to Gaussian observation noise. This ensures that performance differences reflect how uncertainty is represented and propagated in planning, rather than mismatches in the transition distribution. We consider two observation-noise regimes. Homoscedastic Gaussian Observation Noise We first inject independent Gaussian noise directly into the observations: s~t=st+ϵt,ϵt∼(0,σ2I) s_t=s_t+ _t, _t (0,σ^2I) (8) The underlying dynamics remain deterministic; only the agent’s perception is corrupted. The variance is constant across the state space, providing a controlled and analytically simple stochastic regime. By varying σ, we obtain degradation curves that quantify robustness as observation uncertainty increases. This homoscedastic setting serves as our primary evaluation regime, as it cleanly isolates the effect of uncertainty-aware planning under well-specified Gaussian noise. State-Norm–Scaled (Heteroscedastic) Observation Noise Many practical systems exhibit observation uncertainty that increases with state magnitude. To reflect this structure, we consider state-dependent Gaussian observation noise of the form s~t=st+ϵt,ϵt∼(0,σ(st)2I), s_t=s_t+ _t, _t \! (0,σ(s_t)^2I ), (9) where the standard deviation scales with a normalized state norm: σ(st)=σmin+(σmax−σmin)clip(‖stsmax‖2mmax,0,1)σ(s_t)= _ +( _ - _ )\,clip\! ( \| s_ts_ \|_2m_ ,0,1 ) (10) Here smax∈ℝds_ ^d is a fixed per-dimension scaling vector, and the division is understood elementwise. The quantity mmaxm_ bounds the normalized norm, ensuring that σ(st)∈[σmin,σmax]σ(s_t)∈[ _ , _ ]. Under this model, observation variance increases smoothly as the state moves away from nominal regions, capturing a common phenomenon in control systems where sensor uncertainty grows in high-magnitude or dynamically extreme regimes. To maintain experimental clarity and controlled comparisons, our primary ablations focus on the homoscedastic regime, while the state-norm noise serves as a practically motivated extension that stresses the planner under structured, state-dependent uncertainty. 5.3 Evaluation Protocol We evaluate all methods on continuous-control tasks from the MuJoCo Playground zakka2025mujocoplayground, which provides JAX-based implementations of DeepMind Control Suite environments tassa2018deepmindcontrolsuite. We report results on Cartpole Balance and Pendulum Swingup. These environments do not contain terminal states; we therefore use fixed-length episodes of 1000 timesteps for both training and evaluation. After each training episode, the current policy is evaluated over 10 episodes without exploration noise, and the average return is recorded. All methods use identical interaction budgets, episode lengths, and update-to-data ratios. Results are averaged over ten random seeds. On Cartpole Balance, we study robustness to observation noise. We consider homoscedastic Gaussian noise with σ∈0.025,0.05σ∈\0.025,0.05\ and state-norm–scaled heteroscedastic noise with σmin=0.0 _ =0.0 and σmax=0.05 _ =0.05 (Section 5.2). We compare Soft-Actor Critic (SAC) [haarnoja2018soft] and BNAF against our analytic uncertainty-propagation method (MoCA) as well as a certainty-equivalent version (MoCA-CE), which computes the Bellman target using only the predictive mean μ¯ μ (i.e., Σ¯→0 → 0 in Eq. (7)). On Pendulum Swingup, we isolate analytic versus Monte Carlo uncertainty propagation under homoscedastic Gaussian observation noise (σ∈0.025,0.05σ∈\0.025,0.05\). In addition to SAC, BNAF, and MoCA-CE, we include Monte Carlo variants MoCA-MC with K=4K=4 and K=16K=16 samples, and compare them against MoCA, which computes the same expectation analytically. We also report empirical 1σ1σ coverage and epistemic standard deviation for MoCA to assess calibration during training. 6 Results We conduct two complementary experiments. First, we evaluate robustness to stochastic observations under controlled Gaussian noise regimes. Second, we isolate the effect of analytic versus Monte Carlo uncertainty propagation in the Bellman expectation through a targeted ablation study. 6.1 Robustness to Stochastic Observations Figure 2 reports learning curves on Cartpole under increasing observation noise. Homoscedastic noise Under moderate noise (σ=0.025σ=0.025), all methods learn, but clear performance differences emerge. SAC and BNAF achieve reduced asymptotic return relative to low-noise settings, and MoCA-CE performs similarly to BNAF. In contrast, MoCA consistently attains higher final return and exhibits more stable learning dynamics. As noise increases (σ=0.05σ=0.05), the gap widens substantially. SAC, BNAF, and MoCA-CE degrade further in both convergence speed and asymptotic performance. MoCA remains comparatively stable and maintains a clear performance advantage. State-dependent noise Under state-norm–scaled observation noise, the same qualitative ordering persists. All baselines deteriorate noticeably, while MoCA retains superior stability and final return. Overall, the performance gains are attributable to analytic propagation of predictive uncertainty, rather than to model-based learning alone. 6.2 Analytic vs. Monte Carlo Uncertainty Propagation Figure 3 compares analytic and sampling-based uncertainty propagation on Pendulum Swingup. Under moderate noise, MoCA consistently achieves the highest return and fastest convergence. MoCA-MC improves as the number of samples increases from K=4K=4 to K=16K=16, but remains significantly below the analytic method in both learning speed and asymptotic performance. MoCA-CE and model-free baselines perform substantially worse, indicating that ignoring predictive variance is detrimental under stochastic observations. As noise increases, these differences become more pronounced. Monte Carlo variants exhibit higher variance and slower improvement, even with K=16K=16 samples. In contrast, MoCA maintains stable learning dynamics and a clear performance advantage. The lower panels report uncertainty diagnostics for MoCA. Empirical 1σ1σ coverage remains close to the nominal 68% level throughout training, while the epistemic standard deviation stabilizes as learning progresses. This indicates that the analytic moment propagation remains well-calibrated while enabling superior control performance. Overall, analytic uncertainty propagation provides both improved sample efficiency and greater robustness compared to stochastic Monte Carlo estimation under identical model assumptions. 7 Discussion The role of each structural assumption Our method combines three structural choices that differ in both origin and function. Two are inherited from prior work: modeling the transition kernel as a Gaussian (mixture) is standard in MBRL [chua2018deep, janner2019mbpo], and condensing the ensemble mixture to a single moment-matched Gaussian follows lakshminarayanan2017simple. The third, the radial-basis value class, is the element we introduce. It is worth separating what each assumption is for. The quadratic action-value (BNAF) parameterization is a structural necessity of the framework: without a closed-form maxaQ _aQ, the Bellman backup cannot be posed as a state-value expectation in the first place. The dynamics-value pairing plays a different role. It is what makes the resulting expectation [V(s′)]E[V(s )] tractable, and the specific Gaussian–RBF pair is one of several pairings that could satisfy this requirement. Limitations of the value representation The radial-basis instantiation carries two distinct limitations in higher dimensions. The first is representational: RBF mixtures are universal approximators on compact state domains, but the number of bases M needed to maintain approximation quality grows poorly with state dimension, a manifestation of the curse of dimensionality. The second is computational: for each basis, the closed-form expectation requires the determinant |I+ΣΛi|−1/2|I+ _i|^-1/2 and the inverse (I+ΣΛi)−1(I+ _i)^-1 of a d×d× d matrix, each O(d3)O(d^3) in general, giving a per-backup cost of O(Md3)O(Md^3). When Λi _i and the predictive covariance are diagonal, as in our experiments, I+ΣΛiI+ _i is diagonal and both reduce to elementwise operations, lowering the cost to O(Md)O(Md). One way to scale to higher dimensions could be to apply the analytic backup in a learned lower-dimensional latent space, where a compact value representation suffices. A full exploration of this direction is left for future work. 8 Conclusion We introduced a moment-closure formulation for model-based reinforcement learning that enables analytic Bellman expectations under learned stochastic dynamics. Two ingredients combine to make this possible. First, a quadratic action-value parameterization removes the inner maximization from the Bellman backup, reducing it to an expectation of the state-value function alone. Second, a moment-compatibility principle: whenever the predictive transition distribution and the value representation are chosen so that this remaining state-value expectation is closed-form in the distribution’s moments, the backup can be evaluated analytically, without sampling. This principle is not tied to any single functional form. We realize it concretely with a Gaussian transition model paired with a radial-basis value function, for which the expectation has a closed form that propagates both predictive mean and covariance while remaining deterministic and computationally efficient. Empirically, analytic uncertainty propagation improves robustness to stochastic observations and consistently outperforms both mean-only (certainty-equivalent) and Monte Carlo baselines. In particular, it yields higher final performance, greater stability under increasing noise, and well-calibrated predictive uncertainty. These results demonstrate that analytic moment propagation can serve as a practical and effective alternative to sampling-based planning in continuous-control settings. References Appendix A Observation Noise Construction This appendix provides full implementation details of the observation noise processes used in Section 5.2. A.1 Homoscedastic Gaussian Noise In the homoscedastic setting, independent Gaussian noise is added directly to the observed state: s~t=st+ϵt,ϵt∼(0,σ2I). s_t=s_t+ _t, _t (0,σ^2I). The underlying dynamics remain deterministic; stochasticity enters only through the agent’s observations. The variance σ2σ^2 is fixed globally and identical across all state dimensions. We evaluate σ∈0.025,0.05σ∈\0.025,0.05\. A.2 State-Norm–Scaled (Heteroscedastic) Noise To model state-dependent sensing uncertainty, we inject Gaussian noise with variance scaled by a normalized state magnitude. Noise is defined as: s~t=st+ϵt,ϵt∼(0,σ(st)2I), s_t=s_t+ _t, _t \! (0,σ(s_t)^2I ), where the standard deviation is σ(st)=σmin+(σmax−σmin)⋅clip(‖st⊘smax‖2mmax, 0,1)σ(s_t)= _ +( _ - _ )·clip ( \|s_t s_ \|_2m_ ,\,0,1 ) (11) Here: • ⊘ denotes elementwise division. • smax∈ℝks_ ^k is a fixed per-dimension scaling vector over the k normalized state dimensions (defined below). • mmaxm_ bounds the normalized norm. • clip(⋅,0,1)clip(·,0,1) truncates the scaling factor to the interval [0,1][0,1]. Per-Dimension Scaling We select k=3k=3 state dimensions to include in the norm — cart position, cart velocity, and pole angular velocity — and exclude cosθ,sinθ θ, θ, since these are already bounded by construction. Each selected dimension is first normalized by a fixed constant reflecting its nominal magnitude: s^t=st⊘smax. s_t=s_t s_ . (12) For Cartpole, we use: • cart position scale: 2.42.4 • cart velocity scale: 2.02.0 • pole angular velocity scale: 3.03.0 These constants ensure all normalized state components are approximately (1)O(1) within the typical operating region. Normalized Magnitude The scaling term uses the (unsquared) ℓ2 _2 norm: m(st)=‖s^t‖2mmax.m(s_t)= \| s_t\|_2m_ . (13) We set mmax=km_ = k, where k is the number of normalized dimensions, such that m(st)≈1m(s_t)≈ 1 when each normalized component has magnitude 1. For Cartpole, k=3k=3, so mmax=3≈1.732m_ = 3≈ 1.732. This produces smooth growth of observation variance as the system moves toward dynamically extreme regimes. Hyperparameters For the heteroscedastic regime we use σmin=0.0,σmax=0.05 _ =0.0, _ =0.05. Appendix B Moment-Compatible RBF Value Representation To enable analytic Bellman expectations under Gaussian predictive dynamics, the state-value function must admit a closed-form expectation under a Gaussian distribution. We therefore parameterize the value function as a finite mixture of Gaussian radial basis functions (RBFs). B.1 Functional Form The state-value function is defined as: V(s)=∑i=1Mwiexp(−12(s−ci)⊤Λi(s−ci)),V(s)= _i=1^Mw_i (- 12(s-c_i) _i(s-c_i) ), (14) where: • M is the number of RBF components, • ci∈ℝdc_i ^d are fixed centers, • Λi≻0 _i 0 are positive-definite precision matrices, • wi∈ℝw_i are learned weights. The weights wi\w_i\ are optimized via temporal-difference learning. The centers and precision matrices are fixed throughout training. B.2 Gaussian Expectation If the predictive next-state distribution is Gaussian, s′∼(μ,Σ)s (μ, ), the expectation of each basis function admits a closed-form expression. For brevity, write the quadratic form ‖x−a‖B2:=(x−a)⊤B(x−a)\|x-a\|_B^2:=(x-a) B\,(x-a) Consider a single component with center cic_i and precision Λi _i: [Vi(s′)]=wi∫exp(−12‖s′−ci‖Λi2)(s′;μ,Σ)s′E[V_i(s )]=w_i \! (- 12\|s -c_i\|_ _i^2 )\,N(s ;μ, )\,ds Writing the Gaussian density explicitly, the integrand is the exponential of a quadratic form in s′s , Zi∫exp(−12[‖s′−ci‖Λi2+‖s′−μ‖Σ−12])s′,where Zi:=wi(2π)−d/2|Σ|−1/2 split&Z_i \! (- 12 [\,\|s -c_i\|_ _i^2+\|s -μ\|_ ^-1^2\, ] )\,ds ,\\ &where Z_i:=w_i(2π)^-d/2| |^-1/2 split Completing the square in s′s combines the two quadratic forms into a single Gaussian with precision Λi+Σ−1 _i+ ^-1, leaving a constant factor that depends on μ and cic_i. Carrying out the Gaussian integral and simplifying with the identity Σ−1(Λi+Σ−1)−1Λi=Λi(I+ΣΛi)−1 ^-1( _i+ ^-1)^-1 _i= _i(I+ _i)^-1 yields [Vi(s′)]= [V_i(s )]= wi|I+ΣΛi|−1/2exp(−12‖μ−ci‖Ai2), \;w_i\,|I+ _i|^-1/2 \! (- 12\|μ-c_i\|_A_i^2 ), Ai:= A_i:= Λi(I+ΣΛi)−1 \; _i(I+ _i)^-1 Summing over the M components gives the full expectation, [V(s′)]=∑i=1Mwi|I+ΣΛi|−1/2exp(−12‖μ−ci‖Ai2)E[V(s )]= _i=1^Mw_i\,|I+ _i|^-1/2 \! (- 12\|μ-c_i\|_A_i^2 ) which is Eq. (7). This closed form enables deterministic Bellman targets without Monte Carlo sampling. B.3 Number of Basis Functions We set the number of RBFs based on the effective complexity of the value landscape in each environment: • Cartpole: M=256M=256 bases. Cartpole uses a 5D observation with two unbounded-velocity dimensions (practically bounded for initialization), producing a broader range of relevant regimes; we therefore allocate more bases to avoid underfitting. • Pendulum: M=128M=128 bases. Pendulum observations are 3D with (cosθ,sinθ)( θ, θ) bounded by construction, making the value function smoother over a lower-dimensional bounded domain, requiring fewer bases. These choices follow the principle that M should scale with (i) state dimension and (i) the extent of the region in which accurate value curvature is needed for stable greedy control. B.4 Center Initialization and Support Bounds RBF centers cii=1M\c_i\_i=1^M are sampled uniformly over a fixed hyper-rectangle defined by environment-specific practical bounds: ci∼(ℓ,h),c_i ( ,h), where (ℓ,h)( ,h) are per-dimension lower/upper bounds. Cartpole bounds We use bounds ℓ=[−2.5,−1.0,−1.0,−10.0,−15.0], =[-2.5,\,-1.0,\,-1.0,\,-10.0,\,-15.0], h=[2.5, 1.0, 1.0, 10.0, 15.0], h=[2.5,\,1.0,\,1.0,\,10.0,\,15.0], (15) corresponding to practical operating ranges for stable initialization, while cosθ θ and sinθ θ remain bounded by construction. Pendulum bounds We use bounds ℓ=[−1.0,−1.0,−8.0],h=[1.0, 1.0, 8.0]. =[-1.0,\,-1.0,\,-8.0], h=[1.0,\,1.0,\,8.0]. B.5 Anisotropic Lengthscales Each basis uses a diagonal precision (anisotropic bandwidth) to reflect heterogeneous units and curvature across state dimensions: Λ=diag(λ1−2,…,λd−2). =diag ( _1^-2,…, _d^-2 ). Cartpole lengthscales We set λ=[0.6, 0.5, 0.5, 3.0, 4.0],λ=[0.6,\,0.5,\,0.5,\,3.0,\,4.0], using tighter scales for (cosθ,sinθ)( θ, θ) and looser scales for velocity-related dimensions. Pendulum lengthscales We set λ=[0.5, 0.5, 2.0],λ=[0.5,\,0.5,\,2.0], again using tighter scales on bounded trigonometric coordinates and a looser scale for angular velocity. Appendix C Curvature Regularization (Keep-Alive) For NAF-based critics, the advantage term uses a positive-definite curvature matrix P(s)=L(s)L(s)⊤P(s)=L(s)L(s) , where L(s)L(s) is produced by a Cholesky head. In addition to the Cholesky parameterization, we optionally apply a curvature keep-alive penalty that prevents the quadratic form from becoming degenerate (near-flat) early in training. Keep-alive penalty We add the following term to the critic TD loss on each replay minibatch: ℒkeepalive=λs∼ℬ[ReLU(pmin−‖L(s)‖F2)2],L_keepalive=λ\;E_s [ReLU\! (p_ -\|L(s)\|_F^2 )^2 ], (16) where ℬB denotes the replay batch and ‖L(s)‖F2=∑i,jLij(s)2\|L(s)\|_F^2= _i,jL_ij(s)^2. Interpretation Since tr(P(s))=‖L(s)‖F2tr(P(s))=\|L(s)\|_F^2, Eq. (16) enforces a lower bound on the overall curvature magnitude (trace of P), rather than directly constraining individual eigenvalues. This reduces training pathologies where the learned quadratic advantage becomes nearly flat, which can destabilize greedy control and TD learning. Scope This keep-alive penalty is applied uniformly to all BNAF-based variants that use the shared-trunk critic. Appendix D Hyperparameters Hyperparameter Cartpole Pendulum Critic LR (init) 1×10−41× 10^-4 1×10−41× 10^-4 Critic LR (end) 1×10−41× 10^-4 1×10−41× 10^-4 μ~ μ LR (init) 1×10−51× 10^-5 5×10−65× 10^-6 μ~ μ LR (end) 1×10−51× 10^-5 5×10−65× 10^-6 Batch size 128 128 σexplore _explore (init) 1.0 0.3 σexplore _explore (end) 0.1 0.3 Random-action probability prandp_rand 0.2 0.2 Target update period NtargetN_target 100 100 Curvature keep-alive λ 1×10−41× 10^-4 1×10−41× 10^-4 Curvature keep-alive pminp_ 0.5 2.0 Table 1: MoCA hyperparameters used for Cartpole and Pendulum. Learning rates are specified separately for the critic and the quadratic vertex head μ~(s) μ(s). For Cartpole, exploration noise σexplore _explore is linearly annealed from its initial to final value over the first 50% of training steps; for Pendulum, σexplore _explore is held constant throughout. Curvature keep-alive uses the hinge-squared penalty described in Appendix C.