SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Boyi Liu
Zhaoran Wang†
Tuo Zhao $^{2\dagger}$
$^{1}$ Northwestern University  $^{2}$ Georgia Tech
shenao@u.northwestern.edu
ReParameterization (RP) Policy Gradient Methods (PGMs) have been widely adopted for continuous control tasks in robotics and computer graphics. However, recent studies have revealed that, when applied to long-term reinforcement learning problems, model-based RP PGMs may experience chaotic and non-smooth optimization landscapes with exploding gradient variance, which leads to slow convergence. This is in contrast to the conventional belief that reparameterization methods have low gradient estimation variance in problems such as training deep generative models. To comprehend this phenomenon, we conduct a theoretical examination of model-based RP PGMs and search for solutions to the optimization difficulties. Specifically, we analyze the convergence of the model-based RP PGMs and pinpoint the smoothness of function approximators as a major factor that affects the quality of gradient estimation. Based on our analysis, we propose a spectral normalization method to mitigate the exploding variance issue caused by long model unrolls. Our experimental results demonstrate that proper normalization significantly reduces the gradient variance of model-based RP PGMs. As a result, the performance of the proposed method is comparable or superior to other gradient estimators, such as the Likelihood Ratio (LR) gradient estimator. Our code is available at https://github.com/agentification/RP_PGM.
Reinforcement Learning (RL) has seen tremendous success in a variety of sequential decision-making applications, such as strategy games [51, 59] and robotics [15, 61], by identifying actions that maximize long-term accumulated rewards. As one of the most popular methodologies, the policy gradient methods (PGM) [56, 31, 52] seek to search for the optimal policy by iteratively computing and following a stochastic gradient direction with respect to the policy parameters. Therefore, the quality of the stochastic gradient estimation is essential for the effectiveness of PGMs.
Two main categories have emerged in the realm of stochastic gradient estimation: (1) Likelihood Ratio (LR) estimators, which perform zeroth-order estimation through the sampling of function evaluations [64, 32, 31], and (2) ReParameterization (RP) gradient estimators, which harness the differentiability of the function approximation [17, 48, 12, 53]. Despite the wide adoption of both LR and RP PGMs in practice, the majority of the literature on the theoretical properties of PGMs focuses on LR PGMs. The optimality and approximation error of LR PGMs have been heavily investigated under various settings [3, 60, 7]. Conversely, the theoretical underpinnings of RP PGMs remain to be fully explored, with a dearth of research on the quality of RP gradient estimators and the convergence of RP PGMs.
RP gradient estimators have established themselves as a reliable technique for training deep generative models such as variational autoencoders [17]. From a stochastic optimization perspective, previous
studies [48, 40] have shown that RP gradient methods enjoy small variance, which leads to better convergence and performance. However, recent research [42, 37] has reported an opposite observation: When applied to long-horizon reinforcement learning problems, model-based RP PGMs tend to encounter chaotic optimization procedures and highly non-smooth optimization landscapes with exploding gradient variance, causing slow convergence.
Such an intriguing phenomenon inspires us to delve deeper into the theoretical properties of RP gradient estimators in search of a remedy for the issue of exploding gradient variance in model-based RP PGMs. To this end, we present a unified theoretical framework for the examination of model-based RP PGMs and establish their convergence results. Our analysis implies that the smoothness and accuracy of the learned model are crucial determinants of the exploding variance of RP gradients: (1) both the gradient variance and bias exhibit a polynomial dependence on the Lipschitz continuity of the learned model and policy w.r.t. the input state, with degrees that increase linearly with the steps of model value expansion, and (2) the bias also depends on the error of the estimated model and value.
Our findings suggest that imposing smoothness on the model and policy can greatly decrease the variance of RP gradient estimators. To put this discovery into practice, we propose a spectral normalization method to enforce the smoothness of the learned model and policy. It's worth noting that this method can enhance the algorithm's efficiency without substantially compromising accuracy when the underlying transition kernel is smooth. However, if the transition kernel is not smooth, enforcing smoothness may lead to increased error in the learned model and introduce bias. In such cases, a balance should be struck between model bias and gradient variance. Nonetheless, our empirical study demonstrates that the reduced gradient variance when applying spectral normalization leads to a significant performance boost, even with the cost of a higher bias. Furthermore, our results highlight the potential of investigating model-based RP PGMs, as they demonstrate superiority over other model-based and Likelihood Ratio (LR) gradient estimator alternatives.
Reinforcement Learning. We consider learning to optimize an infinite-horizon  $\gamma$ -discounted Markov Decision Process (MDP) over repeated episodes of interaction. We denote by  $S \subseteq \mathbb{R}^{d_s}$  and  $\mathcal{A} \subseteq \mathbb{R}^{d_a}$  the state and action space, respectively. When taking an action  $a \in \mathcal{A}$  at a state  $s \in S$ , the agent receives a reward  $r(s, a)$  and the MDP transits to a new state  $s'$  according to  $s' \sim f(\cdot | s, a)$ .
We aim to find a policy  $\pi$  that maps a state to an action distribution to maximize the expected cumulative reward. We denote by  $V^{\pi}:S\to \mathbb{R}$  and  $Q^{\pi}:S\times \mathcal{A}\rightarrow \mathbb{R}$  the state value function and the state-action value function associated with  $\pi$ , respectively, which are defined as follows,
$$
Q ^ {\pi} (s, a) = (1 - \gamma) \cdot \mathbb {E} _ {\pi , f} \left[ \sum_ {i = 0} ^ {\infty} \gamma^ {i} \cdot r (s _ {i}, a _ {i}) \Bigg | s _ {0} = s, a _ {0} = a \right], \quad V ^ {\pi} (s) = \mathbb {E} _ {a \sim \pi} \left[ Q ^ {\pi} (s, a) \right].
$$
Here  $s \in S, a \in \mathcal{A}$ , and the expectation  $\mathbb{E}_{\pi,f}[\cdot]$  is taken with respect to the dynamic induced by the policy  $\pi$  and the transition probability  $f$ . We denote by  $\zeta$  the initial state distribution. Under policy  $\pi$ , the state and state-action visitation measure  $\nu_{\pi}(s)$  over  $S$  and  $\sigma_{\pi}(s,a)$  over  $S \times \mathcal{A}$  are defined as
$$
\nu_ {\pi} (s) = (1 - \gamma) \cdot \sum_ {i = 0} ^ {\infty} \gamma^ {i} \cdot \mathbb {P} (s _ {i} = s), \qquad \sigma_ {\pi} (s, a) = (1 - \gamma) \cdot \sum_ {i = 0} ^ {\infty} \gamma^ {i} \cdot \mathbb {P} (s _ {i} = s, a _ {i} = a),
$$
where the summations are taken with respect to the trajectory induced by  $s_0 \sim \zeta$ ,  $a_i \sim \pi(\cdot \mid s_i)$ , and  $s_{i+1} \sim f(\cdot \mid s_i, a_i)$ . The objective  $J(\pi)$  of RL is defined as the expected policy value as follows,
$$
J (\pi) = \mathbb {E} _ {s _ {0} \sim \zeta} \left[ V ^ {\pi} (s _ {0}) \right] = \mathbb {E} _ {(s, a) \sim \sigma_ {\pi}} \left[ r (s, a) \right]. \tag {2.1}
$$
Stochastic Gradient Estimation. The underlying problem of policy gradient, i.e., computing the gradient of an expectation with respect to the parameters of the sampling distribution, takes the form  $\nabla_{\theta}\mathbb{E}_{p(x;\theta)}[y(x)]$ . To restore the RL objective, we can set  $p(x;\theta)$  as the trajectory distribution conditioned on the policy parameter  $\theta$  and  $y(x)$  as the cumulative reward. In the sequel, we introduce two commonly used gradient estimators.
Likelihood Ratio (LR) Gradient (Zeroth-Order): By leveraging the score function, LR gradients only require samples of the function values. Since  $\nabla_{\theta}\log p(x;\theta) = \nabla_{\theta}p(x;\theta) / p(x;\theta)$ , the LR gradient is
$$
\nabla_ {\theta} \mathbb {E} _ {p (x; \theta)} [ y (x) ] = \mathbb {E} _ {p (x; \theta)} [ y (x) \nabla_ {\theta} \log p (x; \theta) ]. \tag {2.2}
$$
ReParameterization (RP) Gradient (First-Order): RP gradient benefits from the structural characteristics of the objective, i.e., how the overall objective is affected by the operations applied to the sources of randomness as they pass through the measure and into the cost function [40]. From the simulation property of continuous distribution, we have the following equivalent sampling processes:
$$
\widehat {x} \sim p (x; \theta) \Longleftrightarrow \widehat {x} = g (\widehat {\epsilon}, \theta), \widehat {\epsilon} \sim p (\epsilon), \tag {2.3}
$$
which states that an alternative way to generate a sample  $\widehat{x}$  from the distribution  $p(x;\theta)$  is to sample from a simpler base distribution  $p(\epsilon)$  and transform this sample using a deterministic function  $g(\epsilon ,\theta)$ . Derived from the law of the unconscious statistician (LOTUS) [21], i.e.,  $\mathbb{E}_{p(x;\theta)}[y(x)] = \mathbb{E}_{p(\epsilon)}[y(g(\epsilon ;\theta))]$ , the RP gradient can be formulated as  $\nabla_{\theta}\mathbb{E}_{p(x;\theta)}[y(x)] = \mathbb{E}_{p(\epsilon)}[\nabla_{\theta}y(g(\epsilon ;\theta))]$ .
3 Analytic Reparameterization Gradient in Reinforcement Learning
In this section, we present two fundamental analytic forms of the RP gradient in RL. We first consider the Policy-Value Gradient (PVG) method, which is model-free and can be expanded sequentially to obtain the Analytic Policy Gradient (APG) method. Then we discuss potential obstacles that may arise when developing practical algorithms.
We consider a policy  $\pi_{\theta}(s,\varsigma)$  with noise  $\varsigma$  in continuous action spaces. To ensure that the first-order gradient through the value function is well-defined, we make the following continuity assumption.
Assumption 3.1 (Continuous MDP). We assume that  $f(s' \mid s, a)$ ,  $\pi_{\theta}(s, \varsigma)^{1}$ ,  $r(s, a)$ , and  $\nabla_{a}r(s, a)$  are continuous in all parameters and variables  $s, a, s'$ .
Policy-Value Gradient. The reparameterization PVG takes the following general form,
$$
\nabla_ {\theta} J (\pi_ {\theta}) = \mathbb {E} _ {s \sim \zeta , \varsigma \sim p} \left[ \nabla_ {\theta} Q ^ {\pi_ {\theta}} (s, \pi_ {\theta} (s, \varsigma)) \right]. \tag {3.1}
$$
In sequential decision-making, any immediate action could lead to changes in all future states and rewards. Therefore, the value gradient  $\nabla_{\theta}Q^{\pi_{\theta}}$  possesses a recursive structure. Adapted from the deterministic policy gradient theorem [52, 34] by considering stochasticity, we rewrite (3.1) as
$$
\nabla_ {\theta} J (\pi_ {\theta}) = \mathbb {E} _ {s \sim \nu_ {\pi , \varsigma}} \left[ \nabla_ {\theta} \pi_ {\theta} (s, \varsigma) \cdot \nabla_ {a} Q ^ {\pi_ {\theta}} (s, a) \big | _ {a = \pi_ {\theta} (s, \varsigma)} \right].
$$
Here,  $\nabla_{a}Q^{\pi_{\theta}}$  can be estimated using a critic, which leads to model-free frameworks [25, 4]. Notably, as a result of the recursive structure of  $\nabla_{\theta}Q^{\pi_{\theta}}$ , the expectation is taken over the state visitation  $\nu_{\pi}$  instead of the initial distribution  $\zeta$ .
By sequentially expanding PVG, we obtain the analytic representation of the policy gradient.
Analytic Policy Gradient. From the Bellman equation  $V^{\pi_{\theta}}(s) = \mathbb{E}_{\varsigma}[(1 - \gamma)r(s,\pi_{\theta}(s,\varsigma)) + \gamma \mathbb{E}_{\xi^{*}}[V^{\pi_{\theta}}(f(s,\pi_{\theta}(s,\varsigma),\xi^{*}))]]$ , we obtain the following backward recursions:
$$
\nabla_ {\theta} V ^ {\pi_ {\theta}} (s) = \mathbb {E} _ {\varsigma} \left[ (1 - \gamma) \nabla_ {a} r \nabla_ {\theta} \pi_ {\theta} + \gamma \mathbb {E} _ {\xi^ {*}} \left[ \nabla_ {s ^ {\prime}} V ^ {\pi_ {\theta}} \left(s ^ {\prime}\right) \nabla_ {a} f \nabla_ {\theta} \pi_ {\theta} + \nabla_ {\theta} V ^ {\pi_ {\theta}} \left(s ^ {\prime}\right) \right] \right], \tag {3.2}
$$
$$
\nabla_ {s} V ^ {\pi_ {\theta}} (s) = \mathbb {E} _ {\varsigma} \left[ (1 - \gamma) \left(\nabla_ {s} r + \nabla_ {a} r \nabla_ {s} \pi_ {\theta}\right) + \gamma \mathbb {E} _ {\xi^ {*}} \left[ \nabla_ {s ^ {\prime}} V ^ {\pi_ {\theta}} \left(s ^ {\prime}\right) \left(\nabla_ {s} f + \nabla_ {a} f \nabla_ {s} \pi_ {\theta}\right) \right] \right]. \tag {3.3}
$$
See §A for detailed derivations of (3.2) and (3.3). Now we have the RP gradient backpropagated through the transition path starting at  $s$ . By taking an expectation over the initial state distribution, we obtain the Analytic Policy Gradient (APG)  $\nabla_{\theta}J(\pi_{\theta}) = \mathbb{E}_{s\sim \zeta}[\nabla_{\theta}V^{\pi_{\theta}}(s)]$ .
There remain challenges when developing practical algorithms: (1) the above formulas require the gradient information of the transition function  $f$ . In this work, however, we consider a common RL setting where  $f$  is unknown and needs to be fitted by a model. It is thus natural to ask how the properties of the model (e.g., prediction accuracy and model smoothness) affect the gradient estimation and the convergence of the resulting algorithms, and (2) even if we have access to an accurate model, unrolling it over full sequences faces practical difficulties. The memory and computational cost scale linearly with the unroll length. Long chains of nonlinear mappings can also lead to exploding or vanishing gradients and even worse, chaotic phenomena [9] and difficulty in optimization [43, 36, 58, 38]. These difficulties demand some form of truncation when performing RP PGMs.
<sup>1</sup>Due to the simulation property of continuous distributions in (2.3), we interchangeably write  $a \sim \pi_{\theta}(\cdot \mid s)$  with  $a = \pi_{\theta}(s, \varsigma)$  and  $s' \sim f(\cdot \mid s, a)$  with  $s' = f(s, a, \xi^*)$ , where  $\xi^*$  is sampled from an unknown distribution.
4 Model-Based RP Policy Gradient Methods
Through the application of Model Value Expansion (MVE) for model truncation, this section unveils two RP policy gradient frameworks constructed upon MVE.
4.1  $h$ -Step Model Value Expansion
To handle the difficulties inherent in full unrolls, many algorithms employ direct truncation, where the long sequence is broken down into short sub-sequences and backpropagation is applied accordingly, e.g., Truncated BPTT [63]. However, such an approach over-prioritizes short-term dependencies, which leads to biased gradient estimates.
In model-based RL (MBRL), one viable solution is to adopt the  $h$ -step Model Value Expansion [16], which decomposes the value estimation  $\widehat{V}^{\pi}(s)$  into the rewards gleaned from the learned model and a residual estimated by a critic function  $\widehat{Q}_{\omega}$ , that is,
$$
\widehat {V} ^ {\pi_ {\theta}} (s) = (1 - \gamma) \cdot \left(\sum_ {i = 0} ^ {h - 1} \gamma^ {i} \cdot r (\widehat {s} _ {i}, \widehat {a} _ {i}) + \gamma^ {h} \cdot \widehat {Q} _ {\omega} (\widehat {s} _ {h}, \widehat {a} _ {h})\right),
$$
where  $\widehat{s}_0 = s$ ,  $\widehat{a}_i = \pi_\theta (\widehat{s}_i,\varsigma)$ , and  $\widehat{s}_{i + 1} = \widehat{f}_{\psi}(\widehat{s}_i,\widehat{a}_i,\xi)$ . Here, the noise variables  $\varsigma$  and  $\xi$  can be sampled from the fixed distributions or inferred from the real samples, which we now discuss.
4.2 Model-Based RP Gradient Estimation
Utilizing the pathwise gradient with respect to  $\theta$ , we present the following two frameworks.
Model Derivatives on Predictions (DP). A straightforward way to compute the first-order gradient is to link the reward, model, policy, and critic together and backpropagate through them. Specifically, the differentiation is carried out on the trajectories simulated by the model  $\widehat{f}_{\psi}$ , which serves as a tool for both the prediction of states and the evaluation of derivatives. The corresponding RP-DP estimator of gradient  $\nabla_{\theta}J(\pi_{\theta})$  is denoted as  $\widehat{\nabla}_{\theta}^{\mathrm{DP}}J(\pi_{\theta})$ , which takes the form of
$$
\widehat {\nabla} _ {\theta} ^ {\mathrm {D P}} J \left(\pi_ {\theta}\right) = \frac {1}{N} \sum_ {n = 1} ^ {N} \nabla_ {\theta} \left(\sum_ {i = 0} ^ {h - 1} \gamma^ {i} \cdot r \left(\widehat {s} _ {i, n}, \widehat {a} _ {i, n}\right) + \gamma^ {h} \cdot \widehat {Q} _ {\omega} \left(\widehat {s} _ {h, n}, \widehat {a} _ {h, n}\right)\right), \tag {4.1}
$$
where  $\widehat{s}_{0,n} \sim \mu_{\pi_\theta}$ ,  $\widehat{a}_{i,n} = \pi_\theta(\widehat{s}_{i,n}, \varsigma_n)$ , and  $\widehat{s}_{i+1,n} = \widehat{f}_\psi(\widehat{s}_{i,n}, \widehat{a}_{i,n}, \xi_n)$  with noises  $\varsigma_n \sim p(\varsigma)$  and  $\xi_n \sim p(\xi)$ . Here,  $\mu_{\pi_\theta}$  is the distribution where the initial states of the simulated trajectories are sampled. In Section 5, we study a general form of  $\mu_{\pi_\theta}$  that is a mixture of the initial state distribution  $\zeta$  and the state visitation  $\nu_{\pi_\theta}$ .
Various algorithms can be instantiated from (4.1) with different choices of  $h$ . When  $h = 0$ , the framework reduces to model-free policy gradients, such as  $\mathrm{RP}(0)$  [4] and the variants of DDPG [34], e.g., SAC [23]. When  $h \to \infty$ , the resulting algorithm is BPTT [22, 13, 6] where only the model is learned. Recent model-based approaches, such as MAAC [12] and related algorithms [42, 4, 33], require a carefully selected  $h$ .
Model Derivatives on Real Samples (DR). An alternative approach is to use the learned differentiable model solely for the calculation of derivatives, with the aid of Monte-Carlo estimates obtained from real samples. By replacing  $\nabla_{a}f$ ,  $\nabla_{s}f$  in (3.2)-(3.3) with  $\nabla_{a}\widehat{f}_{\psi}$ ,  $\nabla_{a}\widehat{f}_{\psi}$  and setting the termination of backpropagation at the  $h$ -th step as  $\widehat{\nabla} V^{\pi_{\theta}}(\widehat{s}_{h,n}) = \nabla \widehat{V}_{\omega}(\widehat{s}_{h,n})$ , we are able to derive a dynamic representation of  $\widehat{\nabla}_{\theta}V^{\pi_{\theta}}$ , which we defer to §A. The corresponding RP-DR gradient estimator is
$$
\widehat {\nabla} _ {\theta} ^ {\mathrm {D R}} J \left(\pi_ {\theta}\right) = \frac {1}{N} \sum_ {n = 1} ^ {N} \widehat {\nabla} _ {\theta} V ^ {\pi_ {\theta}} \left(\widehat {s} _ {0, n}\right), \tag {4.2}
$$
where  $\widehat{s}_{0,n} \sim \mu_{\pi_\theta}$ . Equation (4.2) can be specified as (A.9), which is in the same format as (4.1), but with the noise variables  $\varsigma_n, \xi_n$  inferred from the real data sample  $(s_i, a_i, s_{i+1})$  via the relation  $a_i = \pi_\theta(s_i, \varsigma_n)$  and  $s_{i+1} = \widehat{f_\psi}(s_i, a_i, \xi_n)$  (see §A for details). Algorithms such as SVG [25] and its variants [1, 5] are examples of this RP-DR method.
4.3 Algorithmic Framework
The pseudocode of model-based RP PGMs is presented in Algorithm 1, where three update procedures are performed iteratively. In other words, the policy, model, and critic are updated at each iteration  $t \in [T]$ , generating sequences of  $\{\pi_{\theta_t}\}_{t \in [T + 1]}$ ,  $\{\widehat{f}_{\psi_t}\}_{t \in [T]}$ , and  $\{\widehat{Q}_{\omega_t}\}_{t \in [T]}$ , respectively.
Algorithm 1 Model-Based Reparameterization Policy Gradient
Input: Number of iterations  $T$ , learning rate  $\eta$ , batch size  $N$ , empty dataset  $\mathcal{D}$   
1: for iteration  $t \in [T]$  do  
2: Update the model parameter  $\psi_t$  by MSE or MLE  
3: Update the critic parameter  $\omega_t$  by performing Temporal Difference  
4: Sample states from  $\mu_{\pi_t}$  and estimate  $\widehat{\nabla}_{\theta} J(\pi_{\theta_t}) = \widehat{\nabla}_{\theta}^{\mathrm{DP}} J(\pi_{\theta_t})$  (4.1) or  $\widehat{\nabla}_{\theta}^{\mathrm{DR}} J(\pi_{\theta_t})$  (4.2)  
5: Update the policy parameter  $\theta_t$  by  $\theta_{t+1} \gets \theta_t + \eta \cdot \widehat{\nabla}_{\theta} J(\pi_{\theta_t})$   
6: Execute  $\pi_{\theta_{t+1}}$  and save data to  $\mathcal{D}$  to obtain  $\mathcal{D}_{t+1}$   
7: end for
Policy Update. The update rule for the policy parameter  $\theta \in \Theta$  with learning rate  $\eta$  is as follows,
$$
\theta_ {t + 1} \leftarrow \theta_ {t} + \eta \cdot \widehat {\nabla} _ {\theta} J \left(\pi_ {\theta_ {t}}\right), \tag {4.3}
$$
where  $\widehat{\nabla}_{\theta}J(\pi_{\theta_t})$  can be specified as  $\widehat{\nabla}_{\theta}^{\mathrm{DP}}J(\pi_{\theta_t})$  or  $\widehat{\nabla}_{\theta}^{\mathrm{DR}}J(\pi_{\theta_t})$ .
Model Update. By predicting the mean of transition with minimized mean squared error (MSE) or fitting a probabilistic model with maximum likelihood estimation (MLE), e.g.,  $\psi_t = \operatorname{argmax}_{\psi \in \Psi} \mathbb{E}_{\mathcal{D}_t}[\log \widehat{f}_{\psi}(s_{i+1}|s_i, a_i)]$ , canonical MBRL methods learn forward models that predict how the system evolves when an action is taken at a state.
However, accurate state predictions do not imply accurate RP gradient estimation. Thus, we define  $\epsilon_f(t)$  to denote the model (gradient) error at iteration  $t$ :
$$
\epsilon_ {f} (t) = \max  _ {i \in [ h ]} \mathbb {E} _ {\mathbb {P} \left(s _ {i}, a _ {i}\right), \mathbb {P} \left(\widehat {s} _ {i}, \widehat {a} _ {i}\right)} \left[ \left\| \frac {\partial s _ {i}}{\partial s _ {i - 1}} - \frac {\partial \widehat {s} _ {i}}{\partial \widehat {s} _ {i - 1}} \right\| _ {2} + \left\| \frac {\partial s _ {i}}{\partial a _ {i - 1}} - \frac {\partial \widehat {s} _ {i}}{\partial \widehat {a} _ {i - 1}} \right\| _ {2} \right], \tag {4.4}
$$
where  $\mathbb{P}(s_i,a_i)$  is the true state-action distribution at the  $i$ -th timestep by following  $s_0\sim \nu_{\pi_{\theta_t}}$ ,  $a_j\sim \pi_{\theta_t}(\cdot |s_j)$ ,  $s_{j + 1}\sim f(\cdot |s_j,a_j)$ , with policy and transition noise sampled from a fixed distribution. Similarly,  $\mathbb{P}(\widehat{s}_i,\widehat{a}_i)$  is the model rollout distribution at the  $i$ -th timestep by following  $\widehat{s}_0\sim \nu_{\pi_{\theta_t}}$ ,  $\widehat{a}_j\sim \pi_{\theta_t}(\cdot |\widehat{s}_j)$ ,  $\widehat{s}_{j + 1}\sim \widehat{f}_{\psi_t}(\cdot |\widehat{s}_j,\widehat{a}_j)$ , where the noise is sampled when we use RP-DP gradient estimator and is inferred from real samples when we use RP-DR gradient estimator (in this case  $\mathbb{P}(\widehat{s}_i,\widehat{a}_i) = \mathbb{P}(s_i,a_i)$ ).
In MBRL, it is common to learn a state-predictive model that can make multi-step predictions. However, this presents a challenge in reconciling the discrepancy between minimizing state prediction error and the gradient error of the model. Although it is natural to consider regularizing the models' directional derivatives to be consistent with the samples [33], we contend that the use of state-predictive models does not cripple our analysis of gradient bias based on  $\epsilon_f$ : For learned models that extrapolate beyond the visited regions, the gradient error can still be bounded via finite difference. In other words,  $\epsilon_f$  can be expressed as the mean squared training error with an additional measure of the model class complexity to capture its generalizability. This same argument can also be applied to the case of learning a critic through temporal difference.
Critic Update. For any policy  $\pi$ , its value function  $Q^{\pi}$  satisfies the Bellman equation, which has a unique solution. In other words,  $Q = \mathcal{T}^{\pi}Q$  if and only if  $Q = Q^{\pi}$ . The Bellman operator  $\mathcal{T}^{\pi}$  is defined for any  $(s,a)\in S\times \mathcal{A}$  as
$$
\mathcal {T} ^ {\pi} Q (s, a) = \mathbb {E} _ {\pi , f} \left[ (1 - \gamma) \cdot r (s, a) + \gamma \cdot Q \left(s ^ {\prime}, a ^ {\prime}\right) \right].
$$