SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Michele Garibbo
Department of Engineering Mathematics
Bristol, United Kingdom
michele.garibbo@bristol.ac.uk
Maxime Robeyns
maxime.robeyns.2018@bristol.ac.uk
Laurence Aitchison
laurence.aitchison@bristol.ac.uk
Many reinforcement learning approaches rely on temporal-difference (TD) learning to learn a critic. However, TD-learning updates can be high variance. Here, we introduce a model-based RL framework, Taylor TD, which reduces this variance in continuous state-action settings. Taylor TD uses a first-order Taylor series expansion of TD updates. This expansion allows Taylor TD to analytically integrate over stochasticity in the action-choice, and some stochasticity in the state distribution for the initial state and action of each TD update. We include theoretical and empirical evidence that Taylor TD updates are indeed lower variance than standard TD updates. Additionally, we show Taylor TD has the same stable learning guarantees as standard TD-learning with linear function approximation under a reasonable assumption. Next, we combine Taylor TD with the TD3 algorithm, forming 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.
Actor-critic algorithms underlie many of the recent successes of deep RL [1, 2, 3, 4, 5, 6, 7, 8]. In these algorithms, the actor provides the control policy while the critic estimates the policy's expected long-term returns [i.e. the value function; 9, 10]. The critic is typically trained using some form of temporal-difference (TD) update [e.g. 3, 7, 1, 2, 6]. These TD updates need to be computed in expectation over a large distribution of visited states and actions, induced by the policy and the environment dynamics [11, 12]. Since this expectation is analytically intractable, TD updates are typically performed based on individually sampled state-action pairs from real environmental transitions (i.e. sample-based estimates). However, the variance of (sample-based) TD updates can be quite large, meaning that we need to average over many TD updates for different initial states and actions to get a good estimate of the expected updates [13].
Model-based strategies provide a promising candidate to tackle this high variance [14]. For instance, Dyna methods, among the most popular model-based strategies, use a learned model of the environment transitions to generate additional imaginary transitions. These imaginary transitions can be used as extra training samples for TD methods [e.g. 15, 16, 17, 18, 19, 20]. Although the additional (imaginary) transitions help in reducing the variance of performing TD-updates over a
large distribution of state-action pairs, Dyna methods still rely on the same, potentially high-variance (sample-based) TD-updates as standard TD-learning.
We address the issue of high-variance TD-updates by formulating an expected TD-update over a small distribution of state-action pairs. We show this expected update can analytically be estimated with a first-order Taylor expansion for continuous state-action spaces, in an approach we call Taylor TD. By analytically estimating this expected update, rather than exclusively relying on sample-based estimates (as in e.g. Dyna), we get lower variance TD updates. Additionally, we show Taylor TD does not affect the stable learning guarantees of TD-learning under linear function approximation [for a fixed policy as shown by 21]. Next, we propose a model-based off-policy algorithm, Taylor TD3 (TaTD3), which uses Taylor TD in combination with the TD3 algorithm [1]. We show TaTD3 performs as well as if not better than several state-of-the-art model-free and model-based baseline algorithms on a set of standard benchmark tasks. Finally, we compare TaTD3 to its "Dyna" equivalent, which exclusively relies on sample-based TD-updates. We found the largest benefits of Taylor TD may appear in high dimensional state-action spaces.
Reinforcement learning aims to learn reward-maximising behaviour by interacting with the surrounding environment. At each discrete time step,  $t$ , the agent in state  $\mathbf{s} \in S$ , chooses an action  $\mathbf{a} \in \mathcal{A}$  based on a policy  $\pi : S \to \mathcal{A}$ , and observes a scalar reward,  $r$  and a new state  $\mathbf{s}' \in S$  from the environment. The agent's goal is to find the policy that maximises the expected sum of rewards (i.e. the expected return), from a distribution of initial states (or state-action pairs). As such, it is usually necessary to compute the expected return for a state-action pair  $(\mathbf{s}, \mathbf{a})$  and a policy  $\pi$ ; which we can do with a value function. Given a policy  $\pi$  and an initial state-action pair  $(\mathbf{s}, \mathbf{a})$ , we define the value function  $Q^{\pi}(\mathbf{s}, \mathbf{a}) = \operatorname{E}[R_t \mid S_t = \mathbf{s}, A_t = \mathbf{a}]$ , where  $R_t = \sum_{i=t}^{T} \gamma^{i-t} r_i$  is the discounted sum of future rewards from the current time step  $t$  until termination  $T$ , with discount factor  $\gamma \in [0,1]$ . The value function or critic,  $Q^{\pi}$ , quantifies how good the policy,  $\pi$ , is in terms of its expected return when taking action  $\mathbf{a}$  in state  $\mathbf{s}$  and following the policy  $\pi$  thereafter.
In off-policy learning, a function approximation,  $Q_{\theta}$ , is typically used to approximate the true value function,  $Q^{\pi^{\mathrm{target}}}$  of a target policy,  $\pi^{\mathrm{target}}$ , based on experience collected from behavioural policies,  $\pi^{\mathrm{behave}}$ . The function approximation is typically updated with some form of TD-learning [11, 22],
$$
\mathrm {E} _ {(s, a) \sim \mathcal {B}} [ \Delta \theta (\mathbf {s}, \mathbf {a}) ] = \mathrm {E} _ {(s, a) \sim \mathcal {B}} [ \delta_ {\theta} (\mathbf {s}, \mathbf {a}) \nabla_ {\theta} Q _ {\theta} (\mathbf {s}, \mathbf {a}) ], \tag {1}
$$
$$
\delta_ {\theta} (\mathbf {s}, \mathbf {a}) = r (\mathbf {s}, \mathbf {a}) + \gamma Q _ {\theta} \left(\mathbf {s} ^ {\prime}, \mathbf {a} ^ {\prime}\right) - Q _ {\theta} (\mathbf {s}, \mathbf {a}). \tag {2}
$$
where  $\mathcal{B}$  denotes a reply buffer containing a mixture of states and actions sampled according to behavioural policies,  $\pi^{\mathrm{behave}}$ . Note that we have not included a learning rate in these equations, so  $\Delta \theta (\mathbf{s},\mathbf{a})$  only represents an update direction, and not the update magnitude. Additionally, in the tuple  $(\mathbf{s},\mathbf{a},\mathbf{s}^{\prime},\mathbf{a}^{\prime})$ , we are gonna refer to  $\mathbf{s}$  as the initial state,  $\mathbf{a}$  as the initial action,  $\mathbf{s}^{\prime}$  as the next state and  $\mathbf{a}^{\prime}$  as the target action. In standard (model-free) off-policy learning, the initial state and action, as well as the reward and the next state  $(\mathbf{s},\mathbf{a},r,\mathbf{s}^{\prime})$  are drawn from the replay buffer [e.g., 23], while the next action comes from the target policy,  $\mathbf{a}^{\prime}\sim \pi^{\mathrm{target}}(\mathbf{s}^{\prime})$ . Critically, this off-policy setting allows to learn  $Q^{\pi^{\mathrm{target}}}$ , even when the initial states and actions,  $(\mathbf{s},\mathbf{a})$  from the replay buffer, come from a different policy (or even an unknown policy).
In the model-based setting, we can exploit this flexibility in the distribution over the initial state,  $\mathbf{s}$ , and action,  $\mathbf{a}$ . Specifically, we sample the initial state,  $\mathbf{s}$  from an initial-state distribution,  $d^{init}$  (if we are taking initial states from the replay buffer, then  $d^{init}$ , would be the distribution over states in the buffer). Next, we sample the initial action,  $\mathbf{a}$ , from an arbitrary initial policy,  $\pi^{init}(\mathbf{s})$  (i.e., the initial action,  $\mathbf{a}$ , does not have to come from the buffer as for the model-free case). Then, we sample the reward and next state  $(r,\mathbf{s}')$  from the (learned) model, and finally we sample the next action from the target policy,
$$
\mathbf {a} \sim \pi^ {\text {i n i t}} (\mathbf {s}) \quad r, \mathbf {s} ^ {\prime} \sim \text {l e a n e r e d - m o d e l} (\mathbf {s}, \mathbf {a}) \quad \mathbf {a} ^ {\prime} \sim \pi^ {\text {t a r g e t}} (\mathbf {s} ^ {\prime}). \tag {3}
$$
In principle, the off-policy setting allows for complete freedom in the distribution over the initial state,  $d^{\mathrm{init}}$ , and the initial policy,  $\pi^{\mathrm{init}}$ . That freedom is necessary in the model-free off-policy setting, because the states and actions in the replay buffer often come from an unknown distribution (e.g., all past behavioural policies). Although most model-based approaches take the initial policy  $\pi^{\mathrm{init}}$  to
equal either the target  $\pi^{\mathrm{target}}$  or behavioural policy  $\pi^{\mathrm{behave}}$  (i.e., the policy used to interact with the environment) [e.g., 19, 18], this does not have to be the case. For instance,  $\pi^{\mathrm{init}}$  could be a broader distribution than  $\pi^{\mathrm{behave}}$ , enabling us to update and improve the action value function for a broader range of actions.
3 Taylor TD
The expected update over actions in Eq. (1) is usually intractable, due to the complexity of the function,  $\delta_{\theta}(\mathbf{s},\mathbf{a})\nabla_{\theta}Q_{\theta}(\mathbf{s},\mathbf{a})$ . Standard TD-learning methods therefore employ a sample-based approach, by sampling initial actions from some policy (distribution)  $\pi^{\mathrm{init}}$  and performing the (TD) updates for those actions. However, sampling can give rise to high-variance estimators, and hence slow learning [24]. Taylor TD reduces this variance by introducing an analytic approximation to the expectation in Eq. (1) over some stochasticity in the distribution over actions (and over the state, i.e., Sec. 3.2). This approach requires a differentiable model of the transitions and rewards as well as continuous state and action space.
Taylor TD takes advantage of the flexibility in choice of the initial, target and behavioural policies arising in model-based off-policy settings (see Sec. 2). Taylor TD can use any target and behavioural policy, but assumes some structure in the initial policy,  $\pi^{\mathrm{init}}$ . Specifically, Taylor TD splits the process of sampling the initial action into two steps. It first samples the mean,  $\mu_{\mathrm{a}}$ , from a mean-policy,  $\pi^{\mathrm{action-mean}}(\mathbf{s})$  (usually, we would use the target policy,  $\pi^{\mathrm{target}}$ , as the mean-policy, but we do not have to). Then Taylor TD samples "noise" from a second distribution,  $\pi^{\mathrm{action-noise}}$ , and the actual (initial) action is the sum of the mean and the noise,
$$
\boldsymbol {\mu} _ {\mathrm {a}} \sim \pi^ {\text {a c t i o n - m e a n}} (\mathbf {s}) \quad \boldsymbol {\xi} _ {\mathrm {a}} \sim \pi^ {\text {a c t i o n - n o i s e}} \tag {4}
$$
$$
\mathbf {a} = \boldsymbol {\mu} _ {\mathrm {a}} + \boldsymbol {\xi} _ {\mathrm {a}}. \tag {5}
$$
Here, we require that the noise,  $\xi_{\mathrm{a}}$ , drawn from  $\pi^{\mathrm{action - noise}}$  has zero expectation, and finite, known covariance,  $\Sigma_{\mathrm{a}}$ ,
$$
\mathrm {E} _ {\boldsymbol {\xi} _ {\mathrm {a}}} \left[ \boldsymbol {\xi} _ {\mathrm {a}} \right] = \mathbf {0} \quad \mathrm {E} _ {\boldsymbol {\xi} _ {\mathrm {a}}} \left[ \boldsymbol {\xi} _ {\mathrm {a}} \boldsymbol {\xi} _ {\mathrm {a}} ^ {T} \right] = \boldsymbol {\Sigma} _ {\mathrm {a}}. \tag {6}
$$
The overall initial policy, combining these two steps, can be written,  $\mathbf{a}\sim \pi^{\mathrm{init}}(\mathbf{s})$
3.1 Action expansion
Here, we analytically approximate the expectation over the action-noise,  $\xi_{\mathrm{a}}$ , using a first-order Taylor expansion. Specifically, we define the expected TD update, averaging over the action noise,  $\xi_{\mathrm{a}}$ , as,
$$
\Delta_ {\mathrm {E x p}} \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) := \mathrm {E} _ {\boldsymbol {\xi} _ {\mathrm {a}}} [ \Delta \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} + \boldsymbol {\xi} _ {\mathrm {a}}) | \mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} ] = \mathrm {E} _ {\boldsymbol {\xi} _ {\mathrm {a}}} \left[ \delta_ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} + \boldsymbol {\xi} _ {\mathrm {a}}) \nabla_ {\theta} Q (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} + \boldsymbol {\xi} _ {\mathrm {a}}) | \mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} \right] \tag {7}
$$
The (overall) expected TD update in Eq. 1 can now be written in terms of  $\Delta_{\mathrm{Exp}}\theta (\mathbf{s},\pmb{\mu}_{\mathrm{a}})$ ,
$$
\mathrm {E} _ {\mathbf {s}, \mathbf {a}} [ \Delta \theta (\mathbf {s}, \mathbf {a}) ] = \mathrm {E} _ {\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}, \boldsymbol {\xi} _ {\mathrm {a}}} [ \Delta \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} + \boldsymbol {\xi} _ {\mathrm {a}}) ] = \mathrm {E} _ {\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}} [ \Delta_ {\mathrm {E x p}} \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) ] \tag {8}
$$
The exact expectation for  $\Delta_{\mathrm{Exp}}\theta (\mathbf{s},\pmb {\mu}_{\mathrm{a}})$  (Eq. 7) is also not tractable typically, due to the complexity of the function,  $\delta_{\theta}(\mathbf{s},\pmb{\mu}_{\mathrm{a}} + \pmb{\xi}_{\mathrm{a}})\nabla_{\theta}Q(\mathbf{s},\pmb{\mu}_{\mathrm{a}} + \pmb{\xi}_{\mathrm{a}})$ . Instead, Taylor TD uses  $\Delta_{\mathrm{Ta}}\theta (\mathbf{s},\pmb{\mu}_{\mathrm{a}})$  to approximate  $\Delta_{\mathrm{Exp}}\theta (\mathbf{s},\pmb{\mu}_{\mathrm{a}})$  (Eq. 7),
$$
\Delta_ {\mathrm {T a}} \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) \approx \Delta_ {\mathrm {E x p}} \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) \tag {9}
$$
Specifically,  $\Delta_{\mathrm{Ta}}\theta (\mathbf{s},\pmb{\mu}_{\mathrm{a}})$ , arises from a first-order Taylor expansion around  $\xi_{\mathrm{a}} = 0$
$$
\Delta_ {\mathrm {T a}} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) = \mathrm {E} _ {\boldsymbol {\xi} _ {\mathrm {a}}} \left[ \left(\delta_ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) + \boldsymbol {\xi} _ {\mathrm {a}} ^ {T} \nabla_ {\mathbf {a}} \delta_ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}})\right) \nabla_ {\theta} \left(Q _ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) + \boldsymbol {\xi} _ {\mathrm {a}} ^ {T} \nabla_ {\mathbf {a}} Q _ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}})\right) \mid \mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}} \right] \tag {10}
$$
This reduces to (see Appendix A.1),
$$
\Delta_ {\mathrm {T a}} \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) = \delta_ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) \nabla_ {\theta} Q _ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) + \nabla_ {\theta , \mathbf {a}} ^ {2} Q _ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) \Sigma_ {\mathrm {a}} \nabla_ {\mathbf {a}} \delta_ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}). \tag {11}
$$
For simplicity, we further assume that the covariance of initial actions is isotropic,  $\Sigma_{\mathrm{a}} = \lambda_{\mathrm{a}}\mathbf{I}$  (although this is not a requirement), resulting in the following (1st-order) Taylor TD-update:
$$
\Delta_ {\mathrm {T a}} \theta (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) = \delta_ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) \nabla_ {\theta} Q _ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) + \lambda_ {\mathrm {a}} \nabla_ {\theta , \mathbf {a}} ^ {2} Q _ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) \nabla_ {\mathbf {a}} \delta_ {\theta} (\mathbf {s}, \boldsymbol {\mu} _ {\mathrm {a}}) \tag {12}
$$
This Taylor TD-update includes the standard TD update at state  $s$  and (mean) action  $\mu_{\mathrm{a}}$  plus a new term, which tries to align the action gradient of the critic (Q-function) with the action gradient of the TD target. Conceptually, this gradient matching should help reduce the variance across TD-updates since it provides a way to analytically integrate over some of the stochasticity induced by  $\pi^{\mathrm{init}}$ . When doing a first-order Taylor series approximation, we may worry that errors in the Taylor expansion might affect convergence. In the Appendix B, we include a proof that for linear function approximation, any errors in the first-order Taylor expansion do not affect the stability of TD-learning (assuming small simulator time steps). Importantly, this result is not a trivial, as there are still approximation errors in the mapping from state and action to the features that we use for linear function approximation. Moreover, we provide empirical evidence that the first-order Taylor expansion reduces the variance of standard TD-updates and support efficient learning, even under non-linear function approximation (see Sections, 3.4, and 4.2.3).
3.2 State expansion
We are not limited to analytically approximate the expectation in Eq. (1) for the distribution of initial actions, but we can extend this approach to a distribution of initial states. Namely, instead of performing a TD-update at the single initial state, s, we perform this update over a distribution of states. Similarly to the action expansion, we assume structure in the distribution over the initial state distribution (denoted here as  $d^{\mathrm{init}}$ ). Specifically, sampling initial states from the initial state distribution involves two steps. We first sample the mean of  $d^{\mathrm{init}}$ ,  $\mu_{\mathrm{s}}$ , from a state-mean distribution  $d^{\mathrm{state-mean}}$  (usually, we would use the distribution of states in a reply buffer as the state-mean distribution, but we do not have to). Then we sample state-noise from a second distribution,  $\xi^{\mathrm{state-noise}}$ , and the actual (initial) state is the sum of the state-mean and the state-noise,
$$
\boldsymbol {\mu} _ {\mathrm {s}} \sim d ^ {\text {s t a t e - m e a n}} \quad \xi_ {\mathrm {s}} \sim d ^ {\text {s t a t e - n o i s e}} \tag {13}
$$
$$
\mathbf {s} = \boldsymbol {\mu} _ {\mathrm {s}} + \boldsymbol {\xi} _ {\mathrm {s}}. \tag {14}
$$
Again, we require that the state-noise,  $\xi_{s}$ , drawn from  $d^{state - noise}$  has zero expectation, and finite, known covariance,
$$
\mathrm {E} _ {\boldsymbol {\xi} _ {\mathrm {s}}} \left[ \boldsymbol {\xi} _ {\mathrm {s}} \right] = \mathbf {0} \quad \mathrm {E} _ {\boldsymbol {\xi} _ {\mathrm {s}}} \left[ \boldsymbol {\xi} _ {\mathrm {s}} \boldsymbol {\xi} _ {\mathrm {s}} ^ {T} \right] = \boldsymbol {\Sigma} _ {\mathrm {s}}. \tag {15}
$$
Thus, we can again formulate an expected TD-update, averaging over the state-noise,
$$
\mathrm {E} _ {\boldsymbol {\xi} _ {\mathrm {s}}} \left[ \Delta \theta \left(\boldsymbol {\mu} _ {\mathrm {s}} + \boldsymbol {\xi} _ {\mathrm {s}}, \mathbf {a}\right) \mid \boldsymbol {\mu} _ {\mathrm {s}}, \mathbf {a} \right] = \mathrm {E} _ {\boldsymbol {\xi} _ {\mathrm {s}}} \left[ \delta_ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}} + \boldsymbol {\xi} _ {\mathrm {s}}, \mathbf {a}\right) \nabla_ {\theta} Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}} + \boldsymbol {\xi} _ {\mathrm {s}}, \mathbf {a}\right) \mid \boldsymbol {\mu} _ {\mathrm {s}}, \mathbf {a} \right] \tag {16}
$$
Again, we can approximate this expected update with a first-order Taylor approximation, but this time, we expand the states around  $\pmb{\xi}_{\mathrm{s}} = 0$  instead of the actions. This requires that the states are continuous, which is usually the case for control tasks. Based on a similar derivation to the action expansion, we get the following Taylor TD-update, where we assumed the state covariance over TD updates to be isotropic,  $\pmb{\Sigma}_{\mathrm{s}} = \lambda_{\mathrm{s}}\mathbf{I}$  (although again this is not a requirement) (see Appendix A.2 for the full derivation):
$$
\Delta_ {\mathrm {T a}} \theta \left(\boldsymbol {\mu} _ {\mathrm {s}}, \mathbf {a}\right) = \delta_ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \mathbf {a}\right) \nabla_ {\theta} Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \mathbf {a}\right) + \lambda_ {s} \nabla_ {\theta , \mathbf {s}} ^ {2} Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \mathbf {a}\right) \nabla_ {\mathbf {s}} \delta_ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \mathbf {a}\right) \tag {17}
$$
The rational behind this update is trying to tackle some of the TD-update variance induced by the initial state distribution (e.g., the distribution of sates in a reply buffer), although we expect this only to work for states close-by to the visited ones (i.e. for small values of  $\lambda_{\mathrm{s}}$ ).
3.3 State-Action expansion implementation
Finally, we can combine the two Taylor expansions into a single TD-update involving both state and action expansions. Nevertheless, computing the dot products between  $\nabla \delta_{\theta}$  and  $\nabla Q_{\theta}$  terms for both state and action terms may not be optimal. One reason for this is dot products are unbounded, increasing the risk of high variance (TD) updates [e.g. 25]. To tackle this issue, we use cosine distances between the gradient terms instead of dot products (see Appendix K.2 for the benefits of
Algorithm 1: Taylor TD
Initialize reply buffer  $\mathcal{B}$    
Initialize model, critic and target policy parameters  $\{w,\theta ,\phi \}$    
Initialize  $\zeta_{\mathrm{a}},\zeta_{\mathrm{s}} = 0$    
for each training step do Collect transition  $(s,a,r,s^{\prime})$  according to  $\pi^{\mathrm{behave}}$  {e.g.,  $\pi^{\mathrm{behave}} = \mathcal{N}(\pi_{\phi}^{\mathrm{det - target}},\sigma^{2}\mathbf{I})\}$ $\mathcal{B}\gets \mathcal{B}\cup \mathcal{B}(\mathbf{s},\mathbf{a},r,\mathbf{s}^{\prime})$  for each model update step do  $w\leftarrow w - \eta_m\nabla_w\mathcal{L}_w^{\mathrm{model}}(\mathbf{s},\mathbf{a},r,\mathbf{s}^{\prime}),\quad (\mathbf{s},\mathbf{a},r,\mathbf{s}^{\prime})\sim \mathcal{B}$    
end for   
for each policy update step do  $(\mu_{\mathrm{s}},\cdot ,\cdot)\sim \mathcal{B}$ $\mu_{\mathrm{a}} = \pi_{\phi}^{\mathrm{det - target}}(\mu_{\mathrm{s}})$ $\hat{r},\hat{s}^{\prime}\sim p_{w}(\cdot |\mu_{s},\mu_{a})$ $\delta = \hat{r} +\gamma Q_{\theta}(\hat{s}^{\prime},\pi_{\phi}^{\mathrm{det - target}}(\hat{s}^{\prime})) - Q_{\theta}(\mu_{s},\mu_{a})$  if Action expansion then  $\zeta_{\mathrm{a}}^{\mathrm{critic}} = \mathrm{CosineSimilarity}(\nabla_{\mathbf{a}}Q_{\theta}(\mu_{\mathrm{s}},\mu_{\mathrm{a}}),\nabla_{\mathbf{a}}\delta (\mu_{\mathrm{s}},\mu_{\mathrm{a}}))$    
end if   
if State expansion then  $\zeta_{\mathrm{s}}^{\mathrm{critic}} = \mathrm{CosineSimilarity}(\nabla_{\mathbf{s}}Q_{\theta}(\mu_{\mathrm{s}},\mu_{\mathrm{a}}),\nabla_{\mathbf{s}}\delta (\mu_{\mathrm{s}},\mu_{\mathrm{a}}))$    
end if   
 $\theta \leftarrow \theta +\eta_c\delta \nabla_\theta Q_\theta (\mu_s,\mu_a) + \eta_c\lambda_a\nabla_\theta \zeta_a^{\mathrm{critic}} + \eta_c\lambda_s\nabla_\theta \zeta_s^{\mathrm{critic}}$ $\phi \leftarrow \phi +\eta_p\nabla_\phi Q_\theta (\mu_s,\mu_a)$    
end for
this). The cosine distance has the advantage of being bounded. By putting everything together, we propose a novel TD update, which we express below in terms of a loss:
$$
\begin{array}{l} \mathcal {L} _ {\theta} ^ {\text {c r i t i c}} = \delta \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right) \tag {18} \\ + \lambda_ {\mathrm {a}} \operatorname {C o s i n e S i m i l a r i t y} \left(\nabla_ {\mathbf {a}} Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right), \nabla_ {\mathbf {a}} \delta \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right)\right) \\ + \lambda_ {s} \operatorname {C o s i n e S i m i l a r i t y} \left(\nabla_ {\mathbf {s}} Q _ {\theta} \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right), \nabla_ {\mathbf {s}} \delta \left(\boldsymbol {\mu} _ {\mathrm {s}}, \boldsymbol {\mu} _ {\mathrm {a}}\right)\right) \\ \end{array}
$$