SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\nabla_ {\mathbf {a}} \delta_ {\theta} (\mathbf {s}, \mathbf {a}) = \frac {\partial \hat {r} (\mathbf {s} , \mathbf {a})}{\partial \mathbf {a}} - \frac {\partial Q _ {\theta} (\mathbf {s} , \mathbf {a} _ {t})}{\partial \mathbf {a}} + \gamma \frac {\partial \hat {\mathbf {s}} ^ {\prime}}{\partial \mathbf {a}} \left(\frac {\partial Q _ {\theta} (\hat {\mathbf {s}} ^ {\prime} , \mathbf {a} ^ {\prime})}{\partial \hat {\mathbf {s}} ^ {\prime}} + \frac {\partial \mathbf {a} ^ {\prime}}{\partial \hat {\mathbf {s}} ^ {\prime}} \frac {\partial Q _ {\theta} (\hat {\mathbf {s}} ^ {\prime} , \mathbf {a} ^ {\prime})}{\partial \mathbf {a} ^ {\prime}}\right) \tag {57}
$$
where  $\hat{r}$  denotes a (differentiable) reward function,  $\hat{s}$  denotes the predicted next state, while  $\frac{\partial\hat{s}'}{\partial\mathbf{a}}$  denotes the gradient term computed by differentiating through a differentiable model of the transitions. An analogous gradient can be written for  $\nabla_{\mathbf{s}}\delta_{\theta}(\mathbf{s},\mathbf{a})$  (i.e., state expansion). However, we do not have to explicitly implement these expressions, as we use autodiff in PyTorch to find gradients of  $\delta_{\theta}(\mathbf{s},\mathbf{a})$  wrt  $\mathbf{a}$  and  $\mathbf{s}$  directly, by re-writing the Taylor TD updates in terms of a simple loss (see Appendix D).
D Computing the updates as the gradient of a loss
Here we report how to easily implement the Taylor TD-update (i.e. Eq. 18) as a loss to be passed to an optimizer (e.g. PyTorch optimizer).
$$
\begin{array}{l} \mathcal {L} _ {\theta} ^ {\text {c r i t i c}} = - \operatorname {s t o p g r a d} _ {\theta} \left\{\delta_ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) \right\} Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) \\ - \lambda_ {\mathrm {a}} \frac {\operatorname {s t o p g r a d} _ {\theta} \left\{\nabla_ {\mathbf {a}} \delta_ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}} , \boldsymbol {\mu} _ {\mathrm {a}}\right) \right\} \cdot \nabla_ {\mathbf {a}} Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}} , \boldsymbol {\mu} _ {\mathrm {a}}\right)}{\operatorname {s t o p g r a d} _ {\theta} \left\{\| \left(\nabla_ {\mathbf {a}} \delta_ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}} , \boldsymbol {\mu} _ {\mathrm {a}}\right)\right) \| \| \nabla_ {\mathbf {a}} Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}} , \boldsymbol {\mu} _ {\mathrm {a}}\right) \| \right\}} \\ - \lambda_ {s} \frac {\operatorname {s t o p g r a d} _ {\theta} \left\{\left(\nabla_ {s} \delta_ {\theta} \left(\boldsymbol {\mu} _ {s} , \boldsymbol {\mu} _ {a}\right)\right) \right\} \cdot \nabla_ {s} Q _ {\theta} \left(\boldsymbol {\mu} _ {s} , \boldsymbol {\mu} _ {a}\right)}{\operatorname {s t o p g r a d} _ {\theta} \left\{\| \left(\nabla_ {s} \delta_ {\theta} \left(\boldsymbol {\mu} _ {s} , \boldsymbol {\mu} _ {a}\right) \right. \right\| \| \nabla_ {s} Q _ {\theta} \left(\boldsymbol {\mu} _ {s} , \boldsymbol {\mu} _ {a}\right) \| \}} \\ \end{array}
$$
"stopgrad $\theta\}$ " denotes the optimizer should not differentiate the quantity inside the curly brackets relative to the parameters  $\theta$  (i.e. equivalent to ".detach") in PyTorch).
E Model learning
Taylor TD-learning requires a differentiable model of the transitions. We learned this model according to maximum-likelihood based on observed transitions sampled from a reply buffer,
$$
\mathcal {L} _ {w} ^ {\text {m o d e l}} = \hat {p} _ {w} \left(s ^ {\prime} \mid s, a\right); \quad (s, a, s ^ {\prime}) \sim \mathcal {B} \tag {58}
$$
where  $\mathcal{B}$  denotes a reply buffer from which (observed) transitions can be sampled. In TaTD3, the model of the transitions consists of an ensemble of 8 models. 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. Conversely, the model of the rewards consists of a neural network trained with mean-square error on the observed rewards sampled from the reply buffer.
F Variance reduction
Here we describe in more details how we computed the variance of Taylor TD and standard TD-learning updates on the 6 continuous control benchmark tasks (i.e. Section 4.1). Based on a policy, a value function and a set of randomly sampled states from a memory buffer, we computed the Taylor and the standard (sample-based) TD batch updates across all sampled states (under the policy). Note, we ensured the exact same (random) states were used to compute the Taylor and the standard (sample-based) TD batch updates. The update for each state takes the form of a gradient evaluation of the value function relative to the TD objective. Next, we computed the variance of these gradient terms across all states and summed them up since the variance of each parameter update contributes to the variance of the value function relative to the TD objective. We repeated this process for 10 different seeds and plotted the mean update variance and standard error for both Taylor and the standard TD updates (i.e. Fig.1).
G Variance reduction across training
In this section we test how the variance of Taylor TD and standard TD-learning updates changes across training. In particular, we want to assess how the variance advantage of Taylor TD updates (over standard TD updates) changes across 3 different points in training: at the start of training (i.e., for an untrained critic and an untrained model of the transitions), early in training (i.e., after 5000 training steps) and finally late in training (i.e., after 50000 training steps).
In Fig. 4, we can see that across all training points, Taylor TD updates provide lower variance updates compared to standard TD updates.
H Toy example
We provide a toy example to investigate the settings in which using a Taylor expansion of an expected update may provide the largest benefits relative to sample-based estimates of the same expected
Figure 4: Mean update variance (and standard error) between Taylor TD and standard (sample-based) TD-learning (batch) updates, computed across 3 time points in training: No Training (i.e., for an untrained critic), Early training (i.e., after 5000 training steps) and finally late in training (i.e., after 50000 training steps). All results are based on 10 runs with corresponding error bars on the Half-Cheetah-v2 environment.
update. To do so, we train a function approximation parameterised by  $\theta$  (e.g. a critic) to approximate the expected outputs of an underlying target function (e.g. TD targets). Starting from a single expected target value, we can formulate the following objective (for multiple targets, we can just sum the objectives):
$$
J (\theta) = \frac {1}{2} \mathrm {E} _ {\boldsymbol {\xi} _ {\mathrm {x}}} \left[ \left(y \left(x + \boldsymbol {\xi} _ {\mathrm {x}}\right) - \hat {y} _ {\theta} \left(x + \boldsymbol {\xi} _ {\mathrm {x}}\right)\right) ^ {2} \right] \tag {59}
$$
where  $\pmb{\xi}_{\mathrm{x}}$  is sampled from a Gaussian distribution with  $\operatorname{E}[\pmb{\xi}_{\mathrm{x}}] = 0$ ,  $\operatorname{E}\left[\pmb{\xi}_{\mathrm{x}}\pmb{\xi}_{\mathrm{x}}^{T}\right] = \lambda_{\mathrm{x}}\mathbf{I}$ . Hence, the task requires the function approximation  $\hat{y}_{\theta}()$  to approximate the expected outputs of the target function  $y()$ , based on a set of randomly perturbed inputs (i.e.  $x + \pmb{\xi}_{\mathrm{x}}$ ). We do this by comparing two different approaches. The first approach involves sampling different outputs of  $y$  (based on different input perturbations  $\pmb{\xi}_{\mathrm{x}}$ ) and training  $\hat{y}_{\theta}()$  to match those outputs. We denote this approach as "Sample-based targets". This approach aims to mimic standard TD-learning, where TD-targets are sampled at each time step to train the critic (i.e. providing sample-based estimates). The second approach aims to mimic Taylor TD, applying a first-order Taylor expansion to the objective in Eq 59,
$$
J _ {\mathrm {T a} \xi_ {\mathrm {a}}} (\theta) = \frac {1}{2} \mathrm {E} _ {\xi_ {\mathrm {x}}} \left[ \left(y + \xi_ {\mathrm {x}} ^ {T} \nabla_ {\mathrm {x}} y - \hat {y} _ {\theta} - \xi_ {\mathrm {x}} ^ {T} \nabla_ {\mathrm {x}} \hat {y} _ {\theta}\right) ^ {2} \right] \tag {60}
$$
For clarity we used the notation  $\hat{y}_{\theta} = \hat{y}_{\theta}(x)$  and  $y = y(x)$ . The terms that are linear in  $\xi_{\mathrm{x}}$  cancel and after summing up equal terms, we get:
$$
J _ {\mathrm {T a} \xi_ {\mathrm {a}}} (\theta) = \frac {1}{2} (y - \hat {y} _ {\theta}) ^ {2} + \frac {1}{2} \lambda_ {\mathrm {x}} \left(\nabla_ {x} y - \nabla_ {x} \hat {y} _ {\theta}\right) ^ {T} \left(\nabla_ {x} y - \nabla_ {x} \hat {y} _ {\theta}\right) \tag {61}
$$
We can then take the gradient of this approximation relative to the function approximation weights to get the weight update for the Taylor approach:
$$
\begin{array}{l} \nabla_ {\theta} J (\theta) \sim \hat {y} _ {\theta} \nabla_ {\theta} \hat {y} _ {\theta} \tag {62} \\ + \lambda_ {x} \nabla_ {\theta} \left(\nabla_ {x} \hat {y} _ {\theta} ^ {T} \nabla_ {x} \hat {y} _ {\theta}\right) \\ - y \nabla_ {\theta} \hat {y} _ {\theta} \\ - \lambda_ {x} \nabla_ {\theta} \left(\nabla_ {x} y ^ {T} \nabla_ {x} \hat {y} _ {\theta} ^ {T}\right) \\ \end{array}
$$
Note, this update is reminiscent of double backpropagation settings in the supervised learning literature [49, 50]. The toy example allows us to compare the Taylor expansion with the sample-based estimation under different conditions (i.e. dimension size of the data points as well as number of data points). In particular, we compare the two approaches for inputs of different dimensions (from 1 to 100 dimensional inputs) and across two data regimes, a low data regime (i.e. only 15 training samples are used to train the function approximation) and a high data regime (i.e. over 100 samples
Figure 5: Average performance of the Taylor expansion approach relative to sample-based estimates on unseen input examples across several input dimensions (x-axis) and two data regimes (5 runs,  $95\%$  c.i.).
are used to train the function approximation). Performance are then assessed based on a novel set of inputs (i.e. 50), sampled from the same underlying distribution as the training data. Fig. 5 show that the benefits of the Taylor approach over sample-based estimates increase as the dimension of the data points grows in size (e.g. RL tasks involving high dimensional action and state spaces). In particular, these benefits are even larger in the presence of a low data regime, such as RL settings in which for any given state-action pair we can only sample a few transitions from the environment.
All the code is available at https://github.com/maximerobeyns/taylortd
J Baseline algorithms
Plotted performance of MBPO [18] was directly taken from the official algorithm repository on GitHub at https://github.com/jannerm/mbpo. Plotted performance of SAC [2] was also obtained taken from the official algorithm repository on GitHub at https://github.com/haarnoja/sac. Plotted performances of both MAGE and Dyna-TD3 [19] were obtained by re-running these algorithms on the benchmark environments, taking the implementations from the official algorithms' repository available at https://github.com/nnaisense/MAGE.
K Ablations
K.1 State expansion ablation
Here, we ask whether the Taylor state expansion brings any benefit to performance, on top of the Taylor action expansion. To do so, we compare the TaTD3 algorithms with and without state expansion on two standard benchmark tasks (i.e. analogous to setting  $\lambda_{\mathrm{s}} = 0$  in the update Eq. 18). Fig. (6) shows that including the state expansion is beneficial to both environments.
K.2 Cosine similarity ablation
Here, we ask whether taking the cosine similarity of state and action gradient terms benefit the performance of TaTD3. To do so, we compare the standard TaTD3 algorithms (trained with the loss in Eq. 18) with a version of TaTD3 trained with a loss without cosine similarity,
$$
\mathcal {L} _ {\theta} = \delta \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) + \lambda_ {\mathrm {a}} \nabla_ {\mathbf {a}} Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) \cdot \nabla_ {\mathbf {a}} \delta \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) + \lambda_ {\mathrm {s}} \nabla_ {\mathbf {s}} Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) \cdot \nabla_ {\mathbf {s}} \delta \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) \tag {63}
$$
Namely, this loss optimizes the dot-product for state and action gradient terms instead of the cosine similarity. We assess this comparison on two standard benchmark tasks (i.e. see Fig. 7). In Fig. 7, we can see the cosine similarity does improve performance on both tasks.
Figure 6: Performance in terms of average returns for TaTD3 (with both state and action expansions) compared to a version of TaTD3 that uses the action expansion only, on two benchmark continuous control tasks. Including the state expansion in TaTD3 seem to improve performance on both tasks (5 runs,  $95\%$  c.i.).
Figure 7: Performance in terms of average returns for TaTD3 compared to a version of TaTD3 that computes the dot product of state and action gradient terms, instead of cosine similarity (5 runs,  $95\%$  c.i.).
L Computing
All experiments were run on a cluster of GPUs, including NVIDIA GeForce RTX 2080, 3090 and NDVIDIA A100. Here, we report the difference in computing time between standard TD-learning, Taylor TD and the two strongest baselines (i.e., MAGE and MBPO) based on the same GPU for each comparison. We do so for a low dimensional (Pendulum), a 'medium' dimensional (Walker2d) and a high dimensional (Humanoid) environment to span a broad range of settings. For MBPO, we report the average running time across all training epochs using the horizon lengthening schedule outlined in the original paper [18].
Environment	TD-learning time	Taylor TD time	MAGE time	MBPO time	n. time steps
Pendulum	24s	38s	36s	52s	200
Walker2d	50s	68s	63s	133s	1000
Humanoid	94s	117s	127s	235s	1000
MAGE has about the same runtime as our method (TaTD3), while MBPO is a lot slower than our method.
M Hyperparameters settings
Below, we reported the hyperparameter settings for TaTD3 (and sample-based Expected-TD3),
Pendulum-v1	HalfCheetah-v2	Walker2d-v2
Steps	10000	150000	150000
Model ensemble size	8	8	8
Model architecture (MLP)	4 h-layers of size 512	4 h-layers of size 512	4 h-layers of size 512
Reward model architecture (MLP)	3 h-layers of size 256	3 h-layers of size 256	3 h-layers of size 256
Actor-critic architecture (MLP)	2 h-layers of size 400	2 h-layers of size 400	2 h-layers of size 400
Dyna steps per environment step	10	10	10
Model horizon	1	1	1
λa	0.25	0.25	0.25
λs	1e-5	1e-5	1e-5
Hopper-v2	Ant-v2	Humanoid-v2
Steps	10000	150000	150000
Model ensemble size	8	8	8
Model architecture (MLP)	4 h-layers of size 512	4 h-layers of size 512	4 h-layers of size 512
Reward model architecture (MLP)	3 h-layers of size 256	3 h-layers of size 512	3 h-layers of size 512
Actor-critic architecture (MLP)	2 h-layers of size 400	4 h-layers of size 400	4 h-layers of size 400
Dyna steps per environment step	10	10	10
Model horizon	1	1	1
λa	0.06	0.06	0.25
λs	1e-5	1e-5	1e-5
Note, "h-layers" stands for hidden layers and the size is in terms of number of units.
We found we could achieve good performance for TaTD3 across all tested environments without needing to fine tune the value of  $\lambda_{\mathrm{a}}$  and  $\lambda_{\mathrm{s}}$  to each environment (i.e.,  $\lambda_{\mathrm{a}} = 0.06$  and  $\lambda_{\mathrm{s}} = 0.00005$ ). These parameters were founded by running a grid search over potential values of these parameters based on a single environment (i.e., Pendulum), then using the best values for all other environments. Nevertheless, we reached top performance in HalfCheetah, Walker2d and Humanoid by using a larger  $\lambda_{\mathrm{a}}$  (i.e.,  $\lambda_{\mathrm{a}} = 0.25$ ). This finding suggests these 3 environments benefit from learning a broader distribution of Q-values over the actions, we believe for better exploration. Additionally, you may notice the much larger values for  $\lambda_{\mathrm{a}}$  compared to  $\lambda_{\mathrm{s}}$ . However, we should stress that any direct comparison between the state covariance  $\lambda_{\mathrm{s}}$  and the action covariance  $\lambda_{\mathrm{a}}$  may not be very meaningful. This is because the distribution of initial states and the distribution of initial actions may be very different, making any direct comparison between  $\lambda_{\mathrm{s}}$  and  $\lambda_{\mathrm{a}}$  hard to assess.