SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Note we used the notation  $\delta$  instead of  $\delta_{\theta}$  to indicate we are treating  $\delta (\pmb{\mu}_{\mathrm{s}},\pmb{\mu}_{\mathrm{a}})$  as a fixed variable independent of  $\theta$ . This ensures when we take the gradient of this loss relative to  $\theta$ , we do not differentiate through any  $\delta$  terms (following the standard implementation of TD-updates in autodiff frameworks such as PyTorch, see Appendix D). In practice, Taylor TD requires a differentiable model of the environment transitions as well as reward function in order to compute the terms  $\nabla_{\mathbf{a}}\delta_{\theta}(\pmb{\mu}_{\mathrm{s}},\pmb{\mu}_{\mathrm{a}})$  and  $\nabla_{\mathbf{s}}\delta_{\theta}(\pmb{\mu}_{\mathrm{s}},\pmb{\mu}_{\mathrm{a}})$  (see Appendix C for further details).
In Sections 3.1 and 3.2, we have been careful to setup Taylor TD over actions and states in the most general possible setting. Nevertheless, a set-up that may work well with Taylor TD is following,
$$
\mu_ {\mathrm {a}} = \pi^ {\text {d e t - t a r g e t}} (\mathrm {s}) \quad \mu_ {\mathrm {s}} \sim \mathcal {B} \tag {19}
$$
where values for  $\mu_{\mathrm{a}}$  (i.e., the mean-policy) are the outputs of a deterministic (target) policy  $\pi^{\mathrm{det - target}}$  and  $\pmb{\mu}_{\mathrm{s}}$  (i.e., the state-means) are set to states sampled from a reply buffer  $\mathcal{B}$ . In this way, Taylor TD performs TD updates that analytically integrate the expected Q-value of the target policy ( $Q^{\pi^{\mathrm{det - target}}}$ ) over a distribution of actions centered at  $\pi^{\mathrm{det - target}}$  and controlled by  $\lambda_{\mathrm{a}}$  as well as over distributions of states centered at previously visited states (i.e., from the reply buffer) and controlled by  $\lambda_{\mathrm{s}}$ . We provide a Taylor TD algorithm implementation with this set-up in the Algorithm box 1. Note, in the Algorithm box 1, we left the model loss unspecified (i.e.,  $\mathcal{L}^{\mathrm{model}}$ ) to denote that any valid method can be used to learn the model of the transitions. However, in the experiments below, we always learned the transition model based on maximum likelihood on the observed environment transitions (see Appendix E).
3.4 Variance analysis
Here, we show that the expected TD update in Eq. (7) (which Taylor TD estimates) has lower variance than standard (sample-based) TD-updates over the same distribution of actions. We only provide this variance analysis for the distribution over actions, because an analogous theorem can be derived for the distribution over states (i.e. Eq. 16). To begin, we apply the law of total variance [26] to standard TD-updates,
$$
\begin{array}{l} \operatorname {V a r} _ {\mathbf {s}, \mathbf {a}} [ \Delta \theta (\mathbf {s}, \mathbf {a}) ] = \operatorname {V a r} _ {\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}, \boldsymbol {\xi} _ {\mathrm {a}}} [ \Delta \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} + \boldsymbol {\xi} _ {\mathrm {a}}) ] \tag {20} \\ = \operatorname {E} _ {\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}} \left[ \operatorname {V a r} _ {\boldsymbol {\xi} _ {\mathrm {a}}} [ \Delta \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} + \boldsymbol {\xi} _ {\mathrm {a}}) | \mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} ] \right] + \operatorname {V a r} _ {\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}} \left[ \operatorname {E} _ {\boldsymbol {\xi} _ {\mathrm {a}}} [ \Delta \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} + \boldsymbol {\xi} _ {\mathrm {a}}) | \mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} ] \right] \\ \end{array}
$$
The inner expectation and inner variance samples action-noise,  $\xi_{\mathrm{a}}$ , while the outer expectation and outer variance samples initial states, s, and action-means,  $\mu_{\mathrm{a}}$ . To relate this expression to Taylor TD, recall that Taylor TD updates are motivated as performing analytic integration over action-noise,  $\xi_{\mathrm{a}}$  - i.e.  $\Delta_{\mathrm{Exp}}\theta (\mathbf{s},\pmb{\mu}_{\mathrm{a}}) = \mathrm{E}\xi_{\mathrm{a}}[\Delta \theta |\mathbf{s},\pmb{\mu}_{\mathrm{a}}]$ . Thus, the variance of standard (sample-based) TD-updates (Eq. 20) is exactly the variance of  $\Delta_{\mathrm{Exp}}\theta$ , plus an additional term to account for variance induced by sampling action-noise,
$$
\operatorname {V a r} _ {\mathbf {s}, \mathbf {a}} [ \Delta \theta (\mathbf {s}, \mathbf {a}) ] = \mathrm {E} _ {\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}} \left[ \operatorname {V a r} _ {\boldsymbol {\xi} _ {\mathrm {a}}} [ \Delta \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} + \boldsymbol {\xi} _ {\mathrm {a}}) | \mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} ] \right] + \operatorname {V a r} _ {\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}} [ \Delta_ {\text {E x p}} \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) ] \tag {21}
$$
The fact that  $\mathrm{E}_{\mathbf{s},\pmb{\mu}_{\mathrm{a}}}\left[\mathrm{Var}_{\pmb{\xi}_{\mathrm{s}}}\left[\Delta \theta \mid \mathbf{s},\pmb{\mu}_{\mathrm{a}}\right]\right]$  is non-negative directly gives a theorem.
Theorem 3.1. The variance for standard (sample-based) TD-updates,  $\mathrm{Var}_{\mathbf{s},\pmb{\mu}_a,\pmb{\xi}_a}[\Delta \theta]$ , is larger than (or equal to) the variance if we exactly integrate action-noise,  $\mathrm{Var}_{\mathbf{s},\pmb{\mu}_a}[\Delta_{Exp}\hat{\theta}]$ ,
$$
\operatorname {V a r} _ {\mathbf {s}, \boldsymbol {\mu} _ {a}, \boldsymbol {\xi} _ {a}} [ \Delta \theta (\mathbf {s}, \mathbf {a}) ] \geq \operatorname {V a r} _ {\mathbf {s}, \boldsymbol {\mu} _ {a}} [ \Delta_ {E x p} \theta (\mathbf {s}, \boldsymbol {\mu} _ {a}) ] \tag {22}
$$
While this theorem only formally applies to the exact expectation,  $\Delta_{\mathrm{Exp}}\theta$ , we would expect it to become increasingly accurate as  $\lambda_{\mathrm{a}}$  and  $\lambda_{\mathrm{s}}$  become smaller, and hence the Taylor series approximations become more accurate. Furthermore, in Appendix B we prove that the resulting updates do not affect the stability of (standard) TD-learning even when the Taylor series approximation errors do not vanish (under linear function approximation). Additionally, we show empirical reductions in the update variance for Taylor TD in Sec. 4.1.
4.1 Variance reduction
In this section we empirically test the claim that Taylor TD updates are lower variance than standard (sample-based) TD-learning updates. To do so, we take the Taylor TD set-up in Algorithm 1 and perform "batch updates" [12], where given an approximate value function  $Q_{\theta}$  and a policy  $\pi$ , several  $Q_{\theta}$  updates are computed across several sampled states and actions, updating  $Q_{\theta}$  only once, based on the sum of all updates. Batch updates ensure the variance of the updates is estimated based on the same underlying value function. We compute batch updates for both Taylor TD and standard (sample-based) TD updates, comparing the variance of the updates between the two approaches (see Appendix F for more details). Formally, we aim to (empirically) show that the inequality in Theorem 3.1 (for both states and actions) holds even with the Taylor approximation across a range of standard control tasks,
$$
\operatorname {V a r} _ {\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\xi} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}, \boldsymbol {\xi} _ {\mathrm {a}}} [ \Delta \theta (\mathbf {s}, \mathbf {a}) ] \geq \operatorname {V a r} _ {\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}} \left[ \Delta_ {\text {E x p}} \theta \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) \right] \approx \Delta_ {\text {T a}} \theta \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) \tag {23}
$$
Fig. (1) shows Taylor TD provides reliably lower variance updates compared to standard TD-learning across all tested tasks, but the Pendulum environment. This finding is consistent with the idea that Taylor TD updates may be most beneficial in higher dimensional state-action spaces, while the Pendulum environment represents the lowest dimensional environment with only four dimensions. This is because sampling estimates (i.e., standard TD-learning) may become less efficient in higher dimensional spaces. We further explore this proposal in a toy example comparing sample-based estimates and Taylor expansions of expected updates. We perform this comparison across data points of different dimensions, and find the benefits of the Taylor expansion (over purely sample-based estimate) to increase with the dimension of the data points (see Appendix H). Additionally, in the Appendix G, we investigate how the variance of Taylor TD and standard TD-learning updates changes during training.
Figure 1: Mean update variance (and standard error) between Taylor TD and standard (sample-based) TD-learning (batch) updates, based on several sampled states and the distribution of actions for those states (i.e. the policy). All results are based on 10 runs with corresponding error bars.
4.2 Benchmark performance
4.2.1 Algorithm
We combine Taylor TD (i.e. Algorithm 1) with the TD3 algorithm [1] in a model-based off-policy algorithm we call Taylor TD3 (TaTD3) (code available at Appendix I). TaTD3 aims to provide a state-of-the-art implementation of Taylor TD for comparison with baseline algorithms. At each iteration, TaTD3 uses a learned model of the transitions and learned reward function to generate several differentiable (imaginary) 1-step transitions (i.e., Dyna style), starting from real states sampled from a reply buffer. These differentiable 1-step transitions are then used to train two critics (i.e. TD3) using Taylor TD updates. The model of the transitions consists of an ensemble of 8 models trained by maximum likelihood on the observed environment transitions (see Appendix E). These models return a Gaussian distribution over the next state, with zero correlations, and with mean and variance given by applying a neural network to the previous state-action pair. We also learn a model the rewards using a neural network trained with mean-square error on the observed rewards. Hence, TaTD3 does not require any a priori knowledge of the true environment transitions or true reward function. Finally, the actor is trained with the deterministic policy gradient [7] on real states as in standard TD3 [1].
4.2.2 Environments
We employ 6 standard environments for continuous control. The first environment consists of a classic problem in control theory used to evaluate RL algorithms [i.e. Pendulum, 27]. The other 5 environments are stanard MuJoCo continuous control tasks [i.e. Hopper, HalfCheetah, Walker2d, Ant and Humanoid, 28]. All results are reported in terms of average performance across 5 runs, each with a different random seed (shade represents  $95\%$  CI).
4.2.3 Comparison with baselines
Here, we report the comparison of TaTD3 with some state-of-the-art model -free and -based baselines on the six benchmark environments. These baselines include 3 model-based algorithms and one 1 model-free algorithm (see Appendix J). The first model-based algorithm is Model-based Policy Optimization (MBPO) [18], which employs the soft actor-critic algorithm (SAC) [2] within a model-based Dyna setting. The second model-based algorithm is Model-based Action-Gradient-Estimator Policy Optimization (MAGE) [19], which uses a differentiable model of the environment transitions to train the critic by minimising the norm of the action-gradient of the TD-error. The third model-based algorithm is TD3 combined with a model-based Dyna approach (i.e. Dyna-TD3). This algorithm was proposed by [19] and was shown to outperform its model-free counterpart, TD3 [1] on most benchmark tasks. Dyna-TD3 is conceptually similar to MBPO, with the main difference of MBPO relying on SAC instead of TD3. Finally, we included SAC [2] as a model-free baseline. Fig. (2) shows TaTD3 performs at least as well, if not better, than the baseline algorithms in all six benchmark tasks: note the much poorer performance of MAGE on Hopper-v2, Walker2d-v2 and Ant-v2, as well as of MBPO on Humanoid-v2 relative to TaTD3.
Figure 2: Performance in terms of average returns for TaTD3 and four state-of-the-art baseline algorithms on six benchmark continuous control tasks. TaTD3 performs as well, if not better, than the four baseline algorithms on all four tasks. All performance are based on 5 runs, with shade representing  $95\%$  c.i.
Figure 3: Performance comparison of TaTD3 with its sample-based equivalent, Sample-based Expected-TD3. All performance are based on 5 runs, with shade representing  $95\%$  c.i.
4.2.4 Taylor vs sample-based (standard) TD-learning
Next, we ask whether Taylor TD provides any performance benefit in computing the expected TD updates in Eq. (7) and (16) over standard sample-based estimates. To do so, we implement a model-based TD3 algorithm analogous to TaTD3, but where the expected TD updates, Eq. (7) and (16), are estimated using sample-based estimates by drawing multiple samples for the state  $(\pmb{\xi}_{\mathrm{s}})$  and action  $(\pmb{\xi}_{\mathrm{a}})$  noise at each step, instead of using an analytic Taylor series approximation. We call this algorithm Sample-based Expected-TD3 (code available at Appendix I). In practice, at each time step, Sample-based Expected-TD3 uses a (learned) model of the transitions to compute multiple TD-updates by sampling ten state perturbations  $(\pmb{\xi}_{\mathrm{s}})$  of visited states  $(\pmb{\mu}_{\mathrm{s}})$  and ten action perturbations  $(\pmb{\xi}_{\mathrm{a}})$  of the deterministic target policy  $(\pmb{\mu}_{\mathrm{a}})$  (i.e. estimating Eq. (7) and (16) through a sample-based estimate). Crucially, we ensure the variance of the state and action perturbations
(i.e.  $\lambda_{\mathrm{a}}$  and  $\lambda_{\mathrm{s}}$ ) is matched between TaTD3 and Sample-based Expected-TD3. We perform this comparison across 4 standard continuous control tasks, spanning a low dimensional (Pendulum), two moderate dimensional (Half-Cheetah and Walker2d) and one high dimensional environment (Humanoid).
In Fig. (3), we can see TaTD3 provides performance benefits over Sample-based Expected-TD3 across the three most high dimensional environments. This finding is in line with the claim we put forth in the variance reduction section (i.e., 4.1) that the benefits of Taylor TD (i.e. TaTD3) over sampling may be most marked in high dimensional state-action spaces. Indeed, the largest performance advantage of TaTD3 is seen in Humanoid-v2, which has the highest dimensional state-action space by a large margin. Conversely, the least difference between TaTD3 and Sample-based Expected-TD3 is seen in Pendulum-v1, which is the task with smallest dimensional state-action space. The relation between the variance of TD-updates and the dimension of the state-action space should be better explored in the future (see also Appendix H). Finally, we should stress Sample-based Expected-TD3 is different from Dyna-TD3, as the latter does not implement any action or state perturbation in the TD-updates. Hence, unlike Sample-based Expected-TD3, Dyna-TD3 does not directly estimate the expected updates in Eq. (7) and (16), but relies on (standard) TD-learning updates (this is also evident in the massive performance difference between Dyna-TD3 and Sample-based Expected-TD3 - by comparing corresponding performance between Fig. 2 and 3).
Model-based strategies provide a promising solution to improving the sample complexity of RL algorithms [14]. In Dyna methods, a model of the environment transitions is learned through interactions with the environment and then employed to generate additional imaginary transitions [e.g. in the form of model roll-outs, 15]. These imaginary transitions, can be used to enhance existing model-free algorithms, leading to improved sample complexity. For example, within TD-learning, imaginary transitions can be used to train a Q-function by providing additional training examples [e.g. 15, 16, 19]. Alternatively, imaginary transitions can be used to provide better TD targets for existing data points [e.g. 17] or to train the actor and/or critic by generating short-horizon trajectories starting at existing state-action pairs [e.g. 18, 29, 20]. These (Dyna) approaches have a clear relation to our approach (Taylor TD), as they attempt to estimate a similar expected TD-update as in Eq. (7). However, Dyna approaches only use potentially high-variance sample-based estimates, while Taylor TD exploits analytic results to reduce that variance.
Conceptually, our approach may resemble previous methods that also rely on analytical computations of expected updates to achieve lower-variance critic or policy updates [e.g. 24, 30, 31, 12, 32, 33, 34] [see also 35, 36, 37, 38, for a different set of approaches relating to update variance in RL]. The most well-known example of this is Expected-SARSA. Expected-SARSA achieves a lower variance TD-update (relative to SARSA), by analytically computing the expectation over the distribution of target actions in the TD-update (i.e. assuming a stochastic target policy) [31, 12];
$$
\delta_ {\theta} (\mathbf {s}, \mathbf {a}) = r (\mathbf {s}, \mathbf {a}) + \gamma \mathrm {E} _ {a ^ {\prime} \sim \pi} \left[ Q _ {\theta} \left(\mathbf {s} ^ {\prime}, \mathbf {a} ^ {\prime}\right) \right] - Q _ {\theta} (\mathbf {s}, \mathbf {a}) \tag {24}
$$
This approach can only reduce variance of TD-updates at the level of the target actions,  $\mathbf{a}'$ , induced by a stochastic target policy. In the case of a deterministic target policy, Expected-SARSA does not provide any benefit. Conversely, our approach attempts to reduce the variance at the level of the initial state-action pairs, (s, a) at which TD-updates are performed. That is; we take the expectation over (s, a) instead of  $\mathbf{a}'$  (see Eq. 7 and 16) which should yield benefits with both stochastic and deterministic target policies. Other well-known RL approaches exploiting analytical computations of expected updates are Expected Policy Gradients [24, 30, 33], Mean Actor Critic [32] and "all-action" policy gradient [34]. These methods attempt to reduce the variance of the stochastic policy gradient update by integrating over the action distribution. Although similar in principle to our approach, these methods focus on the policy update instead of the critic update and, similarly to Expected-SARSA only apply to stochastic target policies.
In practice, our approach may relate to value gradient methods, as it explicitly incorporates the gradient of the value function into the update [e.g. 39, 29, 40, 19, 41, 42]. To our knowledge, the value gradient approach that most relates to our work is MAGE [19], which, nonetheless, has a radically different motivation from Taylor TD. MAGE is motivated by noting that the action-gradients of Q drive deterministic policy updates [7], so getting the action-gradients right is critical for policy
learning. In order to encourage the action-gradients of Q to be correct, MAGE explicitly adds a term to the objective that consists of the norm of the action-gradient of the TD-error, which takes it outside of the standard TD-framework. In contrast, our motivation is to reduce the gradient variance across standard TD updates by performing some analytic integration. We do this through a first-order Taylor expansion of the TD update. This difference in motivation leads to numerous differences in the method and analysis, the least of which is that MAGE uses only the action-gradients, while Taylor TD additionally suggests using the state-gradients, as both the state and action gradients can be used to reduce the variance in the (TD) updates.
The idea of applying a Taylor expansion to learning updates is not new to RL. For instance, [43] uses a (model-free) Taylor expansion of the policy update to provide a formalism connecting trust-region policy search with off-policy correction. This is very different from our approach (Taylor TD), not only because Taylor TD is model-based, but also because we apply the Taylor expansion to the critic instead of the actor update for very different reasons (i.e., reduce the update variance across TD updates). Finally, Taylor TD may be applicable to risk-sensitive RL [e.g., 44, 45, 46, 47]. In the presence of a good model of the transitions, Taylor TD may be able to approximate the values of risky actions (or states) around safe (e.g. deterministic) actions (or visited states), without actually needing to take those actions (or visit those states), thanks to the Taylor expansion of the TD objective.
In this article, we introduce a model-based RL framework, Taylor TD, to help reduce the variance of standard TD-learning updates and, speed-up learning of critics. We theoretically and empirically show Taylor TD updates are lower variance than standard (sample-based) TD-learning updates. We show the extra gradient terms used by Taylor TD do not affect the stable learning guarantees of TD-learning with linear function approximation under a reasonable assumption. Next, we combine Taylor-TD with the TD3 algorithm [1] into a model-based off-policy algorithm we denote as TaTD3. We show TaTD3 performs as well, if not better, than several state-of-the-art model-free and model-based baseline algorithms on a set of standard benchmark tasks.
Taylor TD has the limitation that it requires continuous state-action spaces as well as a differentiable model of transitions to calculate the additional (TD) gradient terms, i.e. it must be in the model-based rather than model-free setting (see Appendix C). Additionally, the gradient terms in Taylor TD imply additional computational cost; in the Appendix L, we show this cost is not that large in terms of training times and we expect it to reduce as faster automatic differentiation tools are developed [48].
We would like to thank the Wellcome Trust for funding this work as well as Dr. Steward for supporting the purchase of GPU nodes. This work made use of the HPC system Blue Pebble at the University of Bristol, UK.
[1] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning, pages 1587-1596. PMLR, 2018.
[2] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[3] Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971, 2015.
[4] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438, 2015.