Taylor TD-learning
Authors: Michele Garibbo
Maxime Robeyns
Laurence Aitchison
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement learning
TD-learning
model-based
variance reduction
TL;DR: Providing lower variance TD updates through a first-order Taylor expansion of expected TD updates with a model-based approach.

Abstract: 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.

Taylor TD-learning
Michele Garibbo
Maxime Robeyns
Department of Engineering Mathematics University of Bristol Bristol, United Kingdom
Department of Engineering Mathematics
University of Bristol
Bristol, United Kingdom
michele. m waw.:e/eggrbbboorrttoo...co ac uk
maxime wtw.oeynns...8@rriiiii
Laurence Aitchison
Department Engineering Mathematics
University of Bristol
Bristol, United Kingdom
1 laurence aitchi ac.uk
Abstract
Many reinforcement learning approaches rely on temporal-differenc (TD) learning to learn critic. However, TD-leaminn updates can be high variance. Here, we introduce mosel-based RL framework, Taylor TD, which reduces this variance continuous state-action settings. Taylor TD uses first-order Taylor series expan- sion 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 dard TD updates. Additionally, we show Taylor TD has the same stable learning guarantees standard TD-Jearning with linear function approximation under a reasonable assumption. Next, we combine Taylor TD with the TD3 algorithm, forming TaTD3. We show TaTD3 performs well, if not better, than several state- of-the art model-free and model-basee baseline algorithms on set of standard benchmark tasks.
indeed lower variance than stan-
Introduction
Actor-critic algorithms underlie many of the recent successes of deep RL 123456 8 In these algorithms, the actor provides the control policy while the critic estimates the policy's expected long-term returns fi.e. the value function; 9 [0] The critic is typically trained using some form of cemporal-difference (TD) update [e.g.
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 []] [2] 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-bacc -bseed 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 [3]
For in- stance, 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
Model-based strategies provide a promising candidate to tackle this high variance 14
can be used as extra training samples for TD methods [e.g. [5] [6 additional (imaginary) transitions help in reducing the variance of performing TD-updates over
19 20). Although the
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

large distribution of state-action pairs, Dyna methods still rely on the same, potentially (sample-basee 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 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- ID-learning under linear function approximation [for a fixed policy as shown by 21]. Next, we propose model- based off- off--ooliy algorithm, Taylor TD3 (TaTD3), which uses Taylor TD in combination with the TD3 algorithm M. 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 "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.
2 Background
Reinforcement learning aims to learn reward-maximising behaviour by interacting with the surround- ing environment. At each discrete time step, t, the agent in state s E S, chooses an action a E AA based a policy S + A. and observes scalar reward, and new state s' E 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 distribution of initial states (or state-action pairs). As such, it is usually necessary to compute the expected return for a state-action pair (s, a) and a policy π; which we can do with value function. Given policy π and initial state-action pair as, we define the value function Q" (s. a) E [R S, = s, A, = a). where R. E "r""r; is the discounted sum of future rewards from the current time step & until termination T, with discount factor y E [0, 1). The value function or critic, Q", quantifies how good the policy, स, is in terms of its expected return when taking action a in state and following the policy π thereafter.
In off- off-policy learning, a function approximation. Qe. is typically used to approximate the true value function, Q waye of a target policy, x zaage based on experience collected from behavioural policies, n behave The function approximation is typically updated with some form of TD-leaming [1] 22),
(ss.))=BB [A0(s, a)] E(s.)) (80 a) VooQoss a)].
60 a) r(s.a)++ + 7000+ (s.
(2)
where B denotes reply buffer containing mixture of states and actions sampled according to behavioural policies, hehave Note that we have not included learning rate these equations, so A0(s, a) only represents an update direction, and not the update magnitude. Additionally, in the tuple (s, a, s', a'), we are gonna refer to as the initial state, a as the initial action, as the next state and a' as the target action. In standard (model-free well as the reward and the next state (s, a, s') are drawn from the replay buffer [c.g.. while the next action comes from the target policy, a' ~ ,. (s'). Critically, this off-policy setting allows to learn Qx aager even when the initial states and actions, (s, a) from the replay buffer, come from different policy (or even an unknown policy).
off-policy learning, the initial state and action, as
In the model-based setting, we can exploit this flexibility in the distribution over the initial state, and action, a. Specifically, we sample the initial state, s from an initial-state distribution, dimit (if we are taking initial states from the replay buffer, then dim would be the distribution over states in the buffer). Next, we sample the initial action, a, from an arbitrary initial policy, π =r" (s) (i.e., the initial action. a, does not have to come from the buffer as for the model-free case). Then, sample the reward and next state (r, s') from the (learned) model, and finally we sample the next action from the target policy,
a' In principle, the off- ooflcolicy setting allows for complete freedom in the distribution over the initial state, dini and the initial policy, "ii That freedom is necessary in the model-free off off-po 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 π init to
a n miit (s)
r,s' learned-model(ss a)
largel

---

equal either the target π -e or behavioural policy hhehave (i.e., the policy used to interact with the environment) [e.g.. [8), this does have be the case. For instance, - could be a broader distribution than n behavv enabling us to update and improve the action value function for a broader range of actions.
3 Taylor TD
The expected update over actions Eq. (T) usually intractable, due to the complexity of the function, aVos.a) a). Standard TD-leanning methods therefore employ sample-based approach, by sampling initial actions from some policy (distribution) xiini 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 expectation in Eq. ① 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 off-policy settings (see Sec. Taylor TD can use any target and behavioural policy, but assumes some structure in the initial policy, iniit Specifically, Taylor TD splits the process of sampling the initial action into two steps. It first samples the mean, H4, from a mean- mean-policy, n a.actio-mean (s) (usually, we would use the target policy, π zlarge,, as the mean-policy, but we do not have to). Then Taylor TD samples "noise" from second distribution, x action-noise and the actual (initial) action is the sum the mean and the noise,
aa n action mean (s)
E n action- moise
F4 Here, we require that the noise, E.. drawn from nr action-noise has zero expectation, and finite, known covariance, Σ.
(5)
Eg, [ = 0
Egg [..5 E,,
The overall initial policy, combining these two steps, can be written, a inint (s).
3.1 Action expansion
Here, we analytically approximate the expectation over the action-noise, Eg. using first-order Taylor expansion. Specifically, we define the expected TD update, averaging over the action noise, E. as,
(s, u.) == Ee, [A pa + E.) is, p.l = Ee, (s ua + Vo 14, + Ea) s, Ma
(7)
The (overall) expected TD update in Eq. can now be written in terms of AExp My)
Ex,a [s0 a)] Ex, As pa + E.) Ex,yu, [Asx pa)
(8)
The exact expectation for A ma) (Eq. is also not tractable typically, due to the complexity of the function, dg (s. H, + C,) Q(. p. T. Instead, Taylor TD uses As nt) to approximate As (.,)) (Eq. 7}
As, m.) ~ Axx pw))
Specifically. (s, (,), arises from a first-order Taylor expansion around &. = 0,
ATa (s. (b)) = Ee, (6o (b)) + e V''vo (ado(s. V. +
14g
(10)
This reduces to (see Appendix A.1)
ATa (b)) = 60(s, Vo + V.
(11)
3

---

For simplicity, we further assume that the covariance of initial actions is isotropic, Σg = X₂I (although this is not requirement), resulting the following (Ist-order) Taylor TD-updtte:
An (4,) 60 (4)) VoQo(s, (4,) Aa V. Qo,, Va5
(12)
This Taylor TD- TD-update includes the standard TD update at state and (mean) action H4 plus 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 provides way to analytically integrate over some of the stochasticity induced by xilit. When doing first order Taylor series approximation, we may worry that errors in the Taylor expansion might affect convergence. In the Appendix [ 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. ① for the distribution of initial actions, but we can extend this approach to a distribution initial states. Namely, instead of performing a TD-update at the single initial state, 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 d' Specifically, sampling initial states from the initial state distribution involves two steps. We first sample the mean of d!"it μ.,. from a mean distribution daaae 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, estate and the actual (initial) state is the sum of the state-mean and the state-noise,
P.
gdate-men
E,
dtate noise
(13)
8 = 11g + Es-
(14)
Again, we require that the state-noise, Es. drawn from molle has zero expectation, and finite, known covariance,
Ee, [E.] -0
Ee, ( E.
(15)
Thus, we can again formulate an expected TD-update, averaging over the state-noise,
Fgg. [A0 + Es. a) a] ( + Es a) + a),a),,,
(16)
Again, we can approximate this expected update with a first-order Taylor approximation, but this time, we expand the states around &, = 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 assumed the state covariance over TD updates to be isotropic, Σ3 A&I (although again this is not a requirement) (see Appendix the full derivation):
A.2 A.2 for
AT, a) = de (1,, a) VoQe(M, A V., Qo(u, V)V
(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 reply buffer), although we expect this only to work for states close-by to the visited ones (i.e. for small values of Ag).
3.3 State-Action expansion implementation
Finally, we can combine the two Taylor expansions into single TD-update involving both state and action expansions. Nevertheless, computing the dot products between Vo, and VQ₂ 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
Initialise reply buffer B
Initialise model, critic and target policy parameters {wc, 0, () Initialise Sa. S. 0
for each training step do
Collect transition (s, r') according to o hchavee B B U B (s, a. r,'s) for each model update step do
fegg. hehaae - r²I)}
<- T/u V Vw s. a,,,),,
(s,,,
end for for each policy update step do
(Us: .) B
det-targot (1g)
Hn = 8, NP (. Ms
8 = F 7Qo((' T (etttareesss) Qo( Hm)
if Action expansion then
ceniiic = CosineSimilartyyVaaooooo Ma), Pao(Ms:Ha))
end if
if State expansion then
Coritic = CosineSimilarity hit.sSmmimaaiitt...0.0.... Mw), Vz Hu))
end if
0 + ne VoQo H1) Vo coriti + noAs ceritic
④ t Φ + T, Qa (ps
end for
end for
this). The cosine distance has the advantage of being bounded. By putting everything together, we propose novel TD update, which we express below in terms of loss:
Coftic =6(,,, (M) H.)) C
(18)
Vao(H,: +メ, CosineSimilartyy/ooo444,. Veo(H,
+Aa CosineSiiilariyy(Vaoollaa
Note used the notation 8 instead of dø indicate independent of 8. This ensures when we take the gradient of this loss relative to 8, we do not differentiate through any 5 terms (following the standard implementation of TD-updates in autodiff frameworks such as PyTorch, see Appendix D). In practice, Taylor TD requires differentiable model of the environment transitions well as reward function in order to compute the terms Vaõo (Ms (a)) and (5 (4)) (see Appendix C for further details).
treating 8(1 M4) as fixed variable
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,
14_ = π det-target 8))
14 B
199)
where values for p4 (i.e., the mean-policy) are the outputs of a deterministic (target) policy π det-tagget and M4 (i.e., the state-means) are set to states sampled from a reply buffer B. In this way, Taylor TD performs TD updates that analytically integrate the expected Q-value of the target policy (Q (Q*") over distribution of actions centered at det-taaget and controlled by X as well as over distributions of states centered at previously visited states (i.e., from the reply buffer) and controlled by A We provide Taylor TD algorithm implementation with this set-up in the Algorithm box Note, in the Algorithm box we left the model loss unspecified (i.e., Luodel) 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 Eb.
5

---

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 TDdupaates,
aar,, [A0(s, a)] = Va [A0(s, M, + +.]]
(20)
Es., Varre [A0(s, M, + E.)ls, Mal + Vars., [Eg, [A0(s, M4 + +G,)ss s, [al]
The inner expectation and inner variance samples action-noise, É., while the outer expectation and outer variance samples initial states, and action-means, M4, To relate this expression to Taylor TD, recall that Taylor TD updates are motivated as performing analytic integration over action-noise, S. -.ce. A (s, M₂) Eg, |A0 Mal. Thus, the variance of standard (sample-base) baseee TD- TD-updates 20) is exactly the variance of A 0. plus an additional term to account for variance induced by sampling action-noise,
Var,,a [AB(s, a)] = Es,1, [Varg, [At(s. p4, + +.))ss (.). [4.]] + Vars,u, [Aspp(s, M4)l
(21)
The fact that Es [Varg, |AB|s, s'll] is non neg-nveative directly gives a theorem.
Theorem 3.1. The variance for standard (sample-base basee- TD-updates, Vara,u, [A0] is larger than (or equal to) the variance if exactly integrate action-noise, Varz, (AEs
Varssu,, [A0(s, a)] 2 Vars., (s.
(22)
While this theorem only formally applies to the exact expectation, A we would expect it to become increasingly accurate as X, and X become smaller, and hence the Taylor series approxima- tions 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 the update variance for Taylor TD in Sec. 4.1
4 Experiments
4.1 Variance reduction
In this section we empirically test the claim that Taylor TD updates are lower variance than standard cample-basede TD-learnnng updates. To do so, we take the Taylor TD set-up in Algorithm and perform "batch updates" [12], where given an approximate value function Qa and policy æ, several Qo updates are computed across several sampled states and actions, updating Q9 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 campleebaseed 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 range of standard control tasks,
Varp,,e..v,. (sss Z Varp..s, (Ass Ha)] AT5 (4)
(23)
Fig. (T) 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 toy example comparing sample-bas 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 we investigate how the variance of Taylor TD and standard TD-learning updates changes during training.

---

models trained by Pardilum-r1 Hopper-122
Tilf-Cheesah-22
Wilke2dd-22
Ant-2
Hummrod-22
20
40
15
10
(x10000
20
20
4
Taylor TDHeaming
Tayler TD-eaming
Taylor TDHeaming
Taylor TDHeaminn
Taylor TD TD-leaming
TD-teanning
Figure I: Mean update variance (and standard error) between Taylor TD and standard (sample-based) TD-learnng (batch) updates, based 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 T with the TD3 algorithm ( model-based off-policy algorithm we call Taylor TD3 (TaTD3) (code available Appendix DI. TaTD3 aims to provide state-of-the-art implementation of Taylor TD for comparison with baseline algorithms. At each iteration, TaTD3 uses learned model of the transitions and Iearned reward function to generate several differentiable (imaginary) 1-step transitions (i.e., Dyna style), starting from real states sampled from reply buffer. These differentiable 1-step transitions then used train critics (i.e. TD3) using Taylor TD updates. The model of the transitions consists of an ensemble maximum likelihood on the observed environment transitions (see Appendix D These models return a Gaussian distribution over the next state, with zero correlations, and with mean and variance given by applying neural network to the previous state-action pair. We also learn 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, on real states as in standard TD3 [T]
the actor is trained with the deterministic policy gradient 7
4.2.2 Environments
We employ standard environments for continuous control. The first environment consists of classic problem in control theory used to evaluate RI. algorithms [i.e. Pendulum, (27). The other environments are stanard MuJoCo continous 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 model and -based baselines on the six benchmark environments. These baselines include model-based algorithms and one 1 model-free algorithm (see Appendix [T]. The first model-based algorithm is Model-based Policy Optimization (MBPO) [8]. which employs the soft actor-critic algorithm (SAC) 2 within a model- Seee--rr-ddddiiiiinenne Policy Optimization (MAGE) [19), which uses 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 model-based Dyna approach (i.e. Dyna-TD3). This algorithm was proposed by [99 and was shown to outperform benchmark tasks. Dyna-TD: TD3 is conceptually similar to MBPO, with the main difference of MBPO relying on SAC instead of TD3. Finally, we included SAC 2 as model-free baseline. Fig. (2) shows TaTD3 performs at least 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-v v2 relative to Ta33
based Dyna setting. The second model-based algorithm is Model-based
model-free counterpart, TD3 [I] on most
7

---

Pendulum-v1
Hopper-v2
HaliCheetah-v2
4000
12000
10000
Average return
-500
3000
8000
-1000
2000
8000
4000
-1500
1000
2000
0
3k
6k
9k
25k
50k
75k
50k
100k
150k
Walker2d-v2
Ant-v2
Humanoid-v2
4000
8000
Average retur
3000
6000
2000
4000
2000
1000
2000
50K
100K
150K
50k
100k
150k
50k
100K
150k
TaTD3 (ours)
MBPO
MAGE
Dyna- TD3
SAC
Figure 2: Performance in terms of average returns for TaTD3 and four state-of-the baseline algorithms on six benchmark continuous control tasks. TaTD3 performs as well, if not better, than the four baseline algorithms all four tasks. All performance are based on 5 runs, with shade representing 95% c.i.
Pandulum-vi
HallCheetah-v2
Walker2d-22
Humanokd-v2
-250
10000
5000
6000
-500
8000
4000
Average
-750
6000
4000
4000
2000
-1250
2000
1000
2000
-1500
Steps
Steps
Steps
Stegs
TsTD3
Sample-based Expecled-TD3
Figure Performance comparison of TaTD3 with its sample-based equivalent, Sample-based Expected-TDD. 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 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 (E.) and action (E,) noise each step, instead of using an analytic Taylor series approximation. -D3 We call this algorithm Sample- base Expected- TD3 (code available at Appendix II. In practice, at each time step, Sample-based Expected-TID TD3 uses (learned) model of transitions to compute multiple TD-updates by sampling ten state perturbations (E,) of visited states (11,) and ten action perturbations (E.) of the deterministic target policy (i.)..... estimating Eq. (7) and (16) through sample-based based estimate). Crucially, ensure the variance of the state and action perturbations

---

(i.e. X and X₂) is matched between TaTD3 and Sample-based Expected-TDS: 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- 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.I that the benefits of Taylor TD (i.e TaTD3) over sampling may be most marked in high dimensional aate-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-base based Expected-TDD is seen in Pendulum-vI, I, 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- 3DD3 different from Dyna-TID 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-TDS and Sample-based Expected-TD3 by comparing corresponding performance between Fig. and
Related work
Model-based strategies provide a promising solution 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-Icarnngg imaginary transitions can be used to train Q-function by providing additional training examples [e.g. 13 16 19]. Alternatively, imaginary transitions can be used provide better TD targets for 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 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.
existing data points [e.g. 17 or to train the actor and/or critic by generating
Conceptually, our approach may resemble previous methods that also rely on analytical computations of expected updates to achieve lower-variance critic or policy updates [c.g. 24 30 31 12 32 33 34 [see also 35 36 37 38 for different set of approaches relating to update variance in RLJ. The
Expected-SA -SRSSA 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) 12]
most well-known example of this is Expected-SARSS
ós a) = r(s,a) + y Egr~x [Qo(s', a²)] (s,
(24)
This approach can only reduce variance of TD "TD-uddaes at the level of the target actions, a', induced by stochastic target policy. In the case of a deterministic target policy, Expected-SARS. does not provide any benefit. Conversely, our approach attempts to reduce the variance 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 a' (see Eq. 7 and [6) which should yield benefits with both stochastic and deterministic target policies. Other well- RL. approaches exploiting analytical computations of expected updates are Expected Policy Gradients [24. 30 (31, 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 Sapected-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
42 To our knowledge, the value gradient approach that most relates to our work is MAGE 19 which, nonetheless, has
40
19
radically different motivation from Taylor TD. MAGE motivated by noting that the -rrdienns of Q drive deterministic policy updates [7]. so getting the action- -rdiiingsa right is critical for policy

---

learning. In order to encourage the action-gras gradients of Q to be correct, MAGE explicitly adds term the objective that consists of the norm of the action-gradiens of the TD-error, which takes outside of the standard TD- fI-ffrawwork. In contrast, our motivation is to reduce the gradient variance across standard TD updates by performing some analytic integration. We do this through 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 iitiiongradients, 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 (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 In the presence of 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.
updates). Finally, Taylor TD may be applicable to risk-sensitive RL [e.g.. 44 46
47
6 Conclusion and Limitations
In this article, we introduce a model-based RL framework, Taylor TD, to help reduce the variance of standard TD-leaming updates and, speed-up learning of critics. We theoretically and empirically show Taylor TD updates are lower variance than standard (sample-based) TD-leaning updates. We show the extra gradient terms used by Taylor TD do not affect the stable learning guarantees of TD-lcarning with linear function approximation under reasonable assumption. Next, combine Taylor-TD with the TD3 algorithm [1] into a model- based off-policy algorithm we denote as TaTD3. We show TaTD3 performs 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 continous state-action spaces as well as differentiable model transitions to calculate the additional (TD) gradient terms, i.e. it must be the model- rather than model-free setting (see Appendix C). Additionally, the gradient terms in Taylor TD imply additional computational cost; in the Appendix L training times and expect it to reduce faster automatic differentiation tools are developed [48]
we show this cost is not that large in terms of
7
Acknowledgement
We would like to thank the Wellcome Trust for funding this work well as Dr. Steward for supporting the purchase GPU nodes. This work made use of the HPC system Blue Pebble at the University of Bristol, UK.
References
[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 Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv: (509.029777 2015.
[4] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. arXiv preprint wrXv::556.00238.. 2015.
10

---

[5] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pages 1889-1897. PMLR, 2015.
[6] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1oXiv:/707.0000 063407 2017.
[7] David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference machine learning. pages 387-395. PMLR, 2014.
[8] Claas Voelcker, Victor Liao, Animesh Garg, and weighted model-based reinforcement learning. arXiv preprint arXiv:2204.01464, 2022.
Farahmand. Value gradient
[9] Andrew G Barto, Richard S Sutton, and Charles W Anderson. Neuronlike adaptive elements that can solve difficult learning control problems. IEEE transactions on systems, man, and cybernetics, pages 834 846, 1983.
[10] Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems, 12, 1999.
[11] Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3(1):9-44, 1988.
[12] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[13] Michael Fairbank and Eduardo Alonso. The local optimality of reinforcement learning by value gradients, and its relationship to policy gradient learning. arXiv preprint arXiv: vcXiv:/101.0422 101.0428, 2011.
[14] Leslie Pack Kaelbling, Michael Littman, and Andrew W Moore. Reinforcement learning: survey, Journal of artificial intelligence research, 4:237-285, 1996.
[15] Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In Machine learning proceedings 1990, pages 216-224. Elsevier, 1990.
[16] Shixiang Gu, Timothy Lillicrap, Ilya Sutskever, and Sergey Levine. Continuous deep q-learning with model-based acceleration. In International conference on machine learning, pages 2829- 2838. PMLR, 2016.
[17] Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I Jordan, Joseph E Gonzalez, and Sergey Levine. Model-based based value estimation for efficient model-free reinforcement learning, arXiv preprint 1 1800000101, 2018.
[18] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model-based policy optimization. Advances in Neural Information Processing Systems, 32, 2019.
[19] Pierluca Oro and Wojciech Jaškowski. How to learn a useful critic? model-based action- eradient-estimator policy optimization. Advances in Neural Information Processing Systems, 33:313-324, 2020.
[20] Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. Advances in neural information processing systems, 31, 2018
[21] John Tsitsiklis and Benjamin Van Roy. Analysis of temporal-diffferen learning with function approximation. Advances neural information processing systems, 9, 1996.
[22] Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8:279-292, 1992.
11

---

[23] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv: 1312.5622 2013.
[24] Kamil Ciosek and Shimon Whiteson. Expected policy gradients. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
[25] Chunjie Luo, Jianfeng Zhan, Xiaohe Xue, Lei Wang. Rui Ren, and Qiang Yang. Cosine nor- malization: Using cosine similarity instead of dot product in neural networks. In International Conference Artificial Neural Networks, pages 382-391. Springer, 2018.
[26] Paul Bratley, Bennet L Fox, and Linus E Schrage. A guide to simulation. Springer Science Business Media, 2011. 1.
[27] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv: 160..011400 2016.
[28] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 JEEE/RSJ international conference on intelligent robots and systems, pages 5026-5033. IEEE, 2012.
[29] Ignasi Clavera, Violet Fu, and Pieter Abbeel. Model-augmented actor-critic: Backpropagating through paths. arXiv preprint arXiv:2005.0800 2020.
[30] Shimon Whiteson, Expected policy gradients for reinforcement learning. Journal of Machine Learning Research, 21, 2020.
[31] Harm Van Seijen, Hado Van Hasselt, Shimon Whiteson, and Marco Wiering. A theoretical and empirical analysis of expected sarsa. In 2009 ieee symposium on adaptive dynamic programming and reinforcement learning, pages 177-184. IEEE, 2009.
[32] Kavosh Asadi, Cameron Allen, Melrose Roderick, Abdel-rahman Mohamed, George Konidaris, Michael Littman, and Brown University Amazon. Mean actor critic. stat, 1050:1, 2017.
[33] Matthew Fellows, Kamil Ciosek, and Shimon Whiteson. Fourier policy gradients. In Interna- tional Conference on Machine Learning, pages 1486-1495. PMLR, 2018.
[34] Benjamin Petit, Loren Amdahl-Culleton, Yao Liu, Jimmy Smith, and Pierre-Luc Bacon. All-action policy gradient methods: A numerical integration approach. arXiv preprint arXiv:191000909,, 2019.
[35] Maxim Kaledin, Alexander Golubev, and Denis Belomestny. Variance reduction for policy- gradient methods via empirical variance minimization. arXiv preprint arXiv:2206.06827 2022.
[36] Arushi Jain, Gandharv Patil, Ayush Jain, Khimya Khetarpal, and Doina Precup. Variance penalized on-policy and off-policy actor-critic, In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 7899-7907, 2021.
[37] Aviv Tamar, Dotan Di Castro, and Shie Mannor. Temporal difference methods for the variance of the reward to go. In International Conference on Machine Learning, pages 495-503 PMIR, 2013.
[38] Matthew Sobel. The variance of discounted markov decision processes. Journal of Applied Probability, 1944),794880, 1982.
[39] Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. Advances in neural informa- tion processing systems, 28, 2015.
[40] Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model-based stochastic value gradient for continuous reinforcement learning. In Learning for Dynamics and Control, pages 6-20. PMLR, 2021.
12

---

[41] David Balduzzi and Muhammad Ghifary. Compatible value gradients for reinforcement learning of continuous deep policies. arXiv preprint 0rXiv:1509003005, 2015.
[42] Michael Fairbank and Eduardo Alonso. Value-gradient learning. The 2012 international joint conference neural networks (ijenn), pages 1-8. IEEE, 2012.
[43] Yunhao Tang, Michal Valko, and Rémi Munos. Taylor expansion policy optimization. In International Conference on Machine Learning, pages 9397-9406. PMLR, 2020.
[44] Shiau Hong Lim and llyas Malik. Distributional reinforcement learning for risk-sensitive policies. Advances in Neural Information Processing Systems, 35:30977-30989, 2022.
[45] LA Prashanth, Michael C Fu, et al. Risk-sensitive reinforcement learning via policy gradient search. Foundations and Trends@ in Machine Learning. 15(5):537-693, 2022.
[46] Ralph Neuneier and Oliver Mihatsch. Risk sensitive reinforcement learning. Advances in Neural Information Processing Systems, 11, 1998.
[47] Yun Shen, Michael Tobia, Tobias Sommer, and Klaus Obermayer. Risk-sensitive reinforcement learning. Neural computation, 26(7):1298 1328, 2014.
[48] Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey, Journal of Marchine Learning Research, 18:1-43 2018
[49] Wojeiech M Czarnecki, Simon Osindero, Max Jaderberg, Grzegorz Swirszez, and Razvan Pascanu. Sobolev training for neural networks. Advances in Neural Information Processing Systems, 30, 2017.
[50] Harris Drucker and Yann Le Cun, Improving generalization performance using double back- propagation. IEEE transactions neural networks, 3(6):991-997) 1992.
13

---

A First-order Taylor expansion
A.1 Action expansion update proof
Transposing the scalar inner product, E/ VaQoss Ma), ) in the first-order Taylor expansion in the main text (Eq. T0p gives,
(s H4) = Es, (5o(s) m.) + E, va)o(s. Ve (Qo(s, + Es)
Ma
(25)
As Vo is linear operator, and the action noise, E. does not depend on the policy parameters, 8,
A,,(s,, U.) = E (oo(s, (.) V (.o0(ss.. (voQo(s. + Vo. M.)..) s, Ms
(26)
where V 22.0Qo(s (2g) is a matrix of second derivatives. The expectation of &. zero, so the terms linear in E, are zero, leading to,
A (s. (4,) (e + Egg, vadg(s (os.
s. Moy (27)
Swapping the order of the terms in the expectation (which is valid, as (5. Vadols (4,) is scalar),
An. (a)) = (s (2) V VoQu p.) + Egg, [ Qa(s, (.)
(aooso, |s, 1s.11a (28)
We then move the terms independent of &, out of the expectation
A (2) = Vo + V2, Qo (.) Ee, [5.5% (s,
(29)
Finally, we know [ = Σ..
(n, m.) = Vo u.) V2 p (t.)
(30)
If assume the action covariance is isotropic, Σ. A,I, get the (Ist-order) Taylor TD-upe for the action expansion,
( m4) = 50 n. Ve Vo
(31)
A.2 State expansion update proof
Applying the first-order Taylor expansion to Eq. (16) then following the approach from the previous section,
(V An, a) = Ee. [(00(H).)) + (,0 al) Vo (Q(H,,))) (V=Q(H,..)) a))² E.) 1 a
=Ee.
(32)
= Ee. [ + + (odo0,,,, a)) (VoQo(H, + 22 a)E.) 14,
(33)
= S0(1,, a) a) a) + Ee, (V A a)E,) (5. Vod a) 1K,
(34)
= S a) Vo + Vi a) Eg. [ (edo(H4,
(35)
Finally, we know E.55" = E,,
Anb = 60(,,, a) (o + Vo, a)E, (eo
(36)
If we assume the state covariance is isotropic, Σ. A.I, we get the (Ist-order) Taylor TD-update for the state expansion,
a) 1 (1, a) 60( a) Vo a) a)+ + No a)
(37)
14

---

B Proof of stable learning for Taylor TD with linear function approximation (with fixed policy)
The Taylor TD update for the action expansion can be written as (an equivalent proof can be derived for the state expansion):
(s, ma)) = (s (₂ VoQoss, m.) + V. Qo(s, Vad m4,))
(38)
This update is composed of two quantities, the standard TD update (i.e., first term in the sum) plus the extra term induced by the Taylor expansion (i.e., second term in the sum). For the purposes of this proof, consider linear function approximation,
Qo(s, 0²
Q0(s',a')= =0'x
(39)
where,
x = (.. ER
x a')
(40)
and where and x' are feature-vecter vectors of length N. We re-write each of the terms in the Taylor TD update in terms of this linear function approximation. The first term, corresponding to standard TD-update can be written,
Sa a) Vo 00'' 0" xxx)x
- - 7x')TT
(41)
The second term, which is the term introduced by Taylor methods, be written,
V a))
Vo.a (x"0) (V.r +Va [Va(x vx(x00)
= (Vax)" (Var + (V (Vax)0)
(Vax)" V + (Vxx) Vax'd (Vax) 7 Vax0
= =(Vax) V - vvxx) (Vax >
(42)
Here, Vax E AAN,, is matrix, while Var E RA is vector.
Putting the two terms together Eq. expected next weight vector as:
42 and factorising the terms multiplying 9.
write the
E [02+1
0,] = Og + 120 = (I 11-1-rl 7 + A))O, + +
(43)
where:
u rx + (Vax) Var
(44)
A = E x(xx "x')
E NxN
(45)
Á = (Vax)* T V e RNN
(46)
Since only A and À multiplies these two quantities exclusively are important for guaranteeing
stable learning. If A is positive definite and positive-semi defiinef, then (A + A₂A) is positive- definite, and for sufficiently small y7, the magnitude of the eigenvalues of ( + A, are all smaller than in which case the system is stable. Crucially, the term A is the same traditional TD-learning so 11] provides a proof that A is always positive definite [see [2] Thus, all we have to prove is that Å is positive semi-definite. In order to prove that Á is positive semi-definite, we require an (very reasonable) assumption, that the timestep At small. Specifically, if we take,
= At f(s, a) +
(47)
Then the Vax' terms must be small. Specifically, applying the chain rule, the full derivative is (expressed here for l case, Vax' da for simplicity).
Ox' = dal Os' da
Ox' da' 0s' + as Da da
Ox' 0s' Ox' (Oxx
(dx' da'
Or') 0s' + 0s' da
(48)
Da
dal as
15

---

This is proportional to At,
(s( Substituting the definition of Á from Eq. 46)
as' - At
(49)
aa
da
The key test for positive semi-definiteness is that for all vectors b
0 0 Ab..
(50)
0 s <b"E[(V.x' V Vaxib 1.xbb-bb' " Vax'b.
(51)
Putting the b's inside the expectation,
0 S (Vxb)
(52)
Now, xVb RIxA is length-A row-vector, and the terms inside the expectations are inner products of two length vectors. We can write these vector inner products as sums,
S
Xb) - (V
(53)
t=1
where a; a particular element of the action- a, so V Ve, RN is a length-N vector, and
Vaix
V x
VaX-
(54)
Critically, the overall inequality in Eq. 52 holds if similar inequality holds for every term in the sum in Eq. 553
0 S E[(V a. xb)² (V a. xb)] - ,,, (V,,'))
(55)
As Va, xb and Vo, x'b b are scalars, we can write,
0 0 B a. xb) - [VV(V
(56)
There are now two cases. 1
we must have that 0 = xb) always (except at
set of measure zero). Thus, the second term must also be zero (except at set of measure zero), ie. 0 = E[(V,, (b)( it must be positive, in that case, the second term also be non-zero, but it scales with At, can always choose At, small enough to ensure that the Eq. (56) holds, in which case A is positive semi-definite.
(V, x b)] and the inequality holds. Alternatively,
( xb) is non-zero,
C Using the chain rule to expand the gradient of Taylor TD
(learned) differentiable model of the transitions and rewards is needed compute the gradient terms V,ōa( (s, a) and Vz00(s, (s, a) in Taylor TD. This is because the TD target, dg r(ss a) + "a(sss a'), comprises the reward and the Q-value at the next time step, both of which depend on and a. The full gradients for the action expansion can be written:
0r a) da
aQo(s. ax)
0s' da
(0Qa a') 0s!
Oal (Qa(s', a') OS
Va)
(57)
aa
Oa'
where f denotes a (differentiable) reward function, & denotes the predicted next state, while 0g' 8
denotes the gradient term computed by differentiating through differentiable model the transitions. An analogous gradient can be written for V.do(s, 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 ōp a) wrt and directly, by re-writing the Taylor TD updates in terms of a simple loss (see Appendix D
16

---

D Computing the updates as the gradient of a loss
Here report how to easily implement the Taylor TD-update (i.e. Eq. 18 an optimizer (e.g. PyTorch optimizer).
a loss be passed to
Co stopgrad ( m.)) Qo( p4a)
stopgrad, Vaōg(1,, m..) VaQo Mb))
X stopgrad, [1l V Vadu(14, (2)) III VaQu(1 [a)l stopgrad, (V (c,))3 VoQu 14)) stopgrad, (Il
"stopgrad, {}" denotes the optimizer should not differentiate the quantity inside the curly brackets relative to the parameters (i.e. equivalent to ".detach()" PyTorch).
E
Model learning
Taylor TD-learning requires a differentiable model of the transitions. We learned this model according maximum-likelime based on observed transitions sampled from a reply buffer,
Cmodel (s' (s' s,);;
(s, a, s') ~ B
(58)
where B denotes a reply buffer from which (observed) transitions can be sampled. In TaTD3, the model of the transitions consists an ensemble of 8 models. These models return 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 neural network trained with -uaare enror 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 Based on policy, value function and 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 a gradient evaluation of the value function relative the TD objective. Next, computed the variance of these gradient terms across all states and summed them up since the variance 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 Tb.
G Variance reduction across training
In this section we test how the variance of Taylor TD and standard TD-learnng 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 training: the start of training (i.e.. for 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. we can see that across all training points, Taylor TD updates provide lower variance updates compared standard TD updates.
H Toy example
We provide toy example to investigate the settings which using Taylor expansion an expected update may provide the largest benefits relative to sample-based estimates of the same expected
17

---

No Traniing
Earl Traning
Late Trsining
14
40
Moan updato vartanos
1.0
(x10000)
0.8
20
0.4
10
02
Taylor TD-leaning
Taylor TD
TD-kamming
Taylor TD
TDilaminng
Figure 4: Mean update variance (and standard error) between Taylor TD and standard (sample-based) TD-learning (batch) updates, computed across time points in training: No Training (i.e., for untrained critic), Early training (i.e., after 5000 training steps) and finally late in training (i.e., after 50000 training steps) All results based on 10 runs with corresponding error bars the Half- Cheetah-v2 environment.
update. To do so, we train a function approximation parameterised by 0 (e.g. a critic) to approximate the expected outputs of underlying target function (e.g. TD targets). Starting from single expected target value, we can formulate the following objective (for multiple targets, we can just sum the objectives):
J(0) 2 Egg. [[v(x + Ex) Jo(x [x))
(59)
where Ex is sampled form Gaussian distribution with [ = O, F. ( X.I. Hence, the task
requires the function approximation (e() to approximate the expected outputs of the target function y(, based on set of randomly perturbed inputs (i.e. x E.). We do this by comparing two different approaches. The first approach involves sampling different outputs of y (based on different input perturbations Ex) and training de () match those outputs. We denote this approach as 'Sample-based targets". This approach aims to mimic standard TD-leaning, where TD-targets are sampled at each time step train the critic (i.e. providing sample- based estimates). The second approach aims to mimic Taylor TD, applying first-order Taylor expansion to the objective in Eq 59
1 Fie, JUa..
(y E
V
go
< Vxjo)²
(60)
2
For clarity used the notation úe (1( and y y(x). The terms that linear in & cancels and after summing up equal terms, we get:
1 00)² (V V V.9o T Vzūe) 1
June, (0)
(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:
V VoJ ~ 3o Voyo
(62)
A VoloV (V V.U V.00))
Vege
A V TV.TTT... VV.0)
Note, this update is reminiscent of of double backpropagation settings the supervised learning literature [49, (0]. The toy example allows to compare the Taylor expansion with the sample-bass estimation under different conditions (i.e. dimension size of the data points as well number of
data points). In particular, we compare the two approaches for inputs of different dimensions (from to 100 dimensional inputs) and across two data regimes, low data regime (i.e. only 15 training samples used to train the function approximation) and high data regime (i.c. over 100 samples
18

---

rogimo
High regimo
0.15
Prodiction aror
0.10
0.10
0.05
0.05
0.00
0.00
20
40
60
80
100
20
10
60
80
100
Inpult
Input
Taylor expansion
Sample-based targets
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 the training data. Fig. show that the benefits of the Taylor approach over sample-bases estimates increase as the dimension of the data points grows in size (e.g. RL tasks involving high dimensional action and state spaces). In partieular, 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.
I Code
All the code is available at httppp///tthhbbb.oo//aaaeeeeree/
J Baseline algorithms
Plotted performance of MBPO I8 was directly taken from the official algorithm repository on GitHub at /tttps://i tained taken from the official algorithm repository on GitHub at htlhtpps://gg thub. com/haarnojal sac Plotted performances of both MAGE and Dyyn-TDD3 [19] were obtained by re-running these algorithms on the benchmark environments, taking the implementations from the official algorithms' repository available at htttps://githuu com/ http://twww...ooo/aaaaeeeee////e mnaiseseeMMMGG
com/ janermm/mbpo Plotted performance of SAC 2 was also ob-
K
Ablations
K.1
State expansion ablation
Here, ask whether the Taylor state expansion brings any benefit performance, 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 A, 0 in the update Eq. 18). Fig. (6) shows that including the state expansion 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.
so, we compare the standard TaTD3 algorithms (trained with the loss in Eq. [8] with version of TaTD3 trained with loss without cosine similarity.
Lg = H, (4) + Xa VaQe(Ht.. Vad(H, + A VeQo(H4,.)) V M.)
(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. Fig. 7 we can see the cosine similarity does improve performance both tasks.
19

---

HatCheetah-w
Walker2-122
12000
10000
Averago rotum
4000
3000
4000
2000
2000
1000
50k
100K
150k
50k
100k
150k
Steps
TaTD3
ActonOniy-TaTOD
Figure 6: Performance in terms of average returns for TaTD3 (with both state and action expansions) compared a version of TaTD3 that 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.).
HollDhistah-22
Welker2d-22
10000
Avarage retum
8000
2000
100k
100k
Steps
Steps
NoCosine-TTTDS
Figure 7: Performance in terms of average returns for TaTD3 compared 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 for a low dimensional (Pendulum), 'medium' dimensional (Walker2d) and high dimensional (Humanoid) environment to span broad range settings. For MBPO, we report the average rumning 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 method (TaTD3), while MBPO lot slower than our method.
M Hyperparameters settings
Below, reported the hyperparameter settings for TaTD3 (and sample-based Expected-TDD.
20

---

Pendulum-vi
HalfCheetah-v2
Walker2d-v2
Steps
10000
150000
150000
Model ensemble size
8
Model architecture (MLP)
h-layers size 512
h-layers of size 512
h-layers size 512
Reward model architecture (MLP)
3 h-layers of size 256
3 h-layers of size 256
h-layers of size 256
Actor-critic architecture (MLP)
2 h-layers of size 400
2 h-layers of size 400
h-layers of size 400
Dyna steps per environment step
10
Model horizon
Aa
0.25
0.25
0.25
le-5
le-5
1e-5
Hopper-v2
Ant-v2
Humanoid-v2
Steps
10000
150000
150000
Model ensemble size
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)
h-layers size 256
3 h-layers of size 512
h-layers size 512
Actor-critic architecture (MLP)
h-layers size 400 4 h-layers of size 400
h-layers of size 400
Dyna steps environment step
10
10
10
Model horizon
0.06
0.06
0.25
lc-5
le-5
le-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 X and A to each environment (i.e., X = 0.06 and X = 0.00005). These parameters were founded by running 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 HalfChectah, Walker2d and Humanoid by using larger A (i.e., A = 0.25). This finding suggests these environments benefit from learning a broader distribution of Q-values notice the much larger values for X compared to Ay. However, we should stress that any direct comparison between the state covariance A and the action covariance X, 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 X, and A hard to assess.
the actions, believe for better exploration. Additionally, you may
21


Reviews:
[{"review_id": "6xwKN7YOBS", "reviewer": "Reviewer_cy77", "summary": "The authors are proposing a method for reducing the variance of TD-learning updates for model-based RL applied to problems with continual state-action spaces. The method relies on Taylor expansion of noise terms in action and initial state distribution, reducing the contribution of those to the variance of TD-update. The authors demonstrate variance reduction both theoretically and empirically, as well as performance on par with SOTA model-based methods such as MBPO (Janner et al, 2019) and MAGE (D'Oro et al, 2020). For empirical evaluation, the authors apply the proposed update adjustment to TD3 with Dyna. ", "questions": "1. In Figure 1, does the variance of each method change during training?\n\n2. In Figure 2, how do the results look like if the models are trained for another 2x steps? What about comparison in terms of training time?\n\n", "limitations": "No concerns regarding potential negative societal impact, and the limitations were discussed to a reasonable amount. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The empirical results seem pretty strong, though I would have found results with longer training times to be more convincing (about 2x longer, i.e. as long as was used in other work such as Janner et al, 2019).\n\nI appreciated the strong theoretical backing, i.e. demonstrating lower variance and stability guarantees (though it should be noted that I did not check the proofs as this work falls outside of my expertise).  \n\n", "weaknesses": "Minor: \n\n- the authors should be clearer in the main body of the paper about the additional computational demands of the proposed method\n\n- some of the particularly interesting results can be only found in the Appendix. For example, the claim that the method performs better on large state-action spaces seems like an important one, and it was good to see it explored in a more controlled setting in Appendix F. Similarly, we can only find ablations (importance of state expansion and cosine similarity) in Appendix I (also it would be good to see these results on all 6 environments). I hope the authors can find a way to move these results to the main body of the paper (personally I found those to be more insightful than the comparison in Fig 3). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments!\n\n> the authors should be clearer in the main body of the paper about the additional computational demands of the proposed method.\n\nWe raised this point in the Limitations section of the original manuscript, and provided a reference to the Appendix for the actual computational costs in terms of training time, highlighting the additional computational cost is not that large (i.e., standard TD learning is only 20\\% faster on average). We will bring this table to the main text in the camera-ready.  However, standard TD learning is perhaps the fastest baseline.  Additional analysis showed that while Taylor-TD was a bit slower than TD, it was much closer to MAGE (for Walker 2d, Taylor-TD took 68 s while MAGE took 63 s).\n\n> In Figure 2, how do the results look like if the models are trained for another 2x steps? What about comparison in terms of training time?\n\nThese runs require quite a bit of additional compute time, so we have prioritised Ant, as that's the only one that seems not to have saturated yet. In line with the rest of the results for Ant, we found that MBPO (5260 at 250k steps) seemed to be doing a bit better than TaTD3 (4860 at 250k steps).  For the camera-ready, we will run Ant out to 300k steps, and increase the number of iterations in the rest of the environments.\n\n> some of the particularly interesting results can be only found in the Appendix. For example, the claim that the method performs better on large state-action spaces seems like an important one, and it was good to see it explored in a more controlled setting in Appendix F. Similarly, we can only find ablations (importance of state expansion and cosine similarity) in Appendix I (also it would be good to see these results on all 6 environments). I hope the authors can find a way to move these results to the main body of the paper (personally I found those to be more insightful than the comparison in Fig 3).\n\nThanks!  We agree that the Appendix contains a number of interesting results, and to the extent that is possible within space constraints, we will move some of these results into the main text.\n\n\n> In Figure 1, does the variance of each method change during training?\n\nWe have run preliminary experiments assessing the variance reduction from Taylor TD at different stages of training.  We found that for completely untrained networks, there was little benefit, likely because the untrained model provides poor gradient estimates.  However, a beneficial variance reduction, similar in magnitude to those in the main-text, emerges early on in training, and remains for the rest of the training run.  For the camera-ready paper, we will modify Fig. 1 to show the full time-course of the variance reduction through training.\n"}, {"review_id": "v5YLMCVt3H", "reviewer": "Reviewer_9c2H", "summary": "The author presents a variance reduction trick in TD learning by taking the analytical expectation of the gradient using first order taylor expansion. Standard RL, replaces the expectation over the gradient of critic (Q-value) with a sampled value from the replay buffer or online learning samples. The paper proposed to replace this expected gradient of Q value function with first order taylor series expansion of the same. They proposed taylor expansion over both state and action.\n\n\n** I have read the rebuttal of the authors.", "questions": "1. It is not clear in the paper how the model (transition, reward) is learnt using the $L_\\theta$ objective in Eq 18?  Algo 1 shows the update $w$ (model parameter ) with the gradient of $L_\\theta$ without showing how $w$ influences it.\n2. Why is $\\mu_a$ in Eq 19 sampled from deterministic target policy? Why do you enforce this determinism on the policy being learnt?\n3. In Fig 1, what is the variance computer over? Is it same $(s,a)$ pair and then compute multiple gradients over that? Is it batch data, comprising of same data for both MC TD and Taylor-TD and then the variance over updates is computed? Because if the variance is computed over different state-action, then the high variance could be a factor of that (s,a) pair having high stochasticity in the dynamics. It wouldn’t mean necessarily that updates have less variance because of taylor approximation. Please explain what this variance update is over.\n4. What is the effect of reducing the variance in critic update vs to that actor update(Ref[1])? Have you tried ablation study for the same - which one is better, or what effects they have on learning separately vs together?\n5. Provide some connections with Risk-sensitive RL and robust RL - inherent stochasticity vs imperfect knowledge? Will be good to connect to the literature in that area (add references for same)!\n6. Comparison with expected sarsa kind of update in tabular environments? That would help to understand how expected update over the gradients help over Expected sarsa. (Would be good to know more insights on this, just something additional)\n\nReferences --\n\n[1] Variance Reduction for Policy-Gradient Methods via Empirical Variance Minimization", "limitations": "Yes the authors have included the limitation section. The negative societal impact section doesn't apply to this line of research.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Originality\n1. The work is proposing novel research in the direction of stabilizing the critic updates by reducing the variance using first order taylor approximation of the expected gradients.\n\nQuality\n1. The paper could be improved on how the experiments are conducted for Fig 1. Furthermore provide better clarity on the questions section later. \n\nSignificance\n1. The work is definitely targeting a significant problem that could have impact on RL algorithms in general. This work is stabilizing the critic updates by reducing the variance.", "weaknesses": "The weakness/questions are mentioned in the next section. If the authors provide justification for the below, I am willing to change the score.\n\nAdd references for the other variance reduction work in direction of - variance reduction of policy updates, inherent stochasticity (risk-sensitive RL). They are very related to your proposed research.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments!\n\n> Add references for the other variance reduction work ...  Risk-sensitive RL and robust RL\n\n\nThanks for these suggestions! We will add your reference [1], and we did a literature review on risk-sensitive RL approaches (e.g., [2],[3],[4],[5]), which we will add to the camera ready. Please do suggest other papers!\n\nWe agree Taylor TD may be applicable to risk-sensitive RL. In the presence of a good model of the transitions, Taylor TD may can approximate the values of risky actions around safe (e.g. deterministic) actions, without actually needing to take those actions, thanks to the Taylor expansion of the TD objective.\nThat said, Taylor TD does not directly estimate the uncertainty induced by imperfect knowledge (i.e., epistemic uncertainty), so it may be less applicable to robust RL.\nWe would also like to stress that in the related work section of the original manuscript, we discussed Expected Policy Gradients (Ciosek and Whiteson, 2018), Mean Actor Critic (Asadi et al., 2017) and \"all-action\" policy gradient (Petit et al., 2019) as related methods. These methods tackle the variance at the level of the policy rather than critic updates by also integrating over the stochasticity induced by the action distribution.\n \n[1] Variance Reduction for Policy-Gradient Methods via Empirical Variance Minimization\n\n[2] Tamar et al. Policy gradients with variance related risk criteria. ICML (2012)\n\n[3] Bellemare et al. A distributional perspective on reinforcement learning. ICML (2017)\n\n[4] Lim & Malik. Distributional Reinforcement Learning for Risk-Sensitive Policies. NeurIPS (2022)\n\n[5] Fu.  Risk-Sensitive Reinforcement Learning via Policy Gradient Search. arXiv (2018).\n\n> It is not clear in the paper how the model (transition, reward) is learnt\n\nWe thank the reviewer for pointing this out. We neglected to include a section in the appendix where we described how the model was learned. To clarify, the reward and transition model were learned by maximum likelihood based on the observed transition, as in most model-based RL approaches (e.g., as in Janner et al., 2019 - MBPO). We will add an Appendix describing the model-learning to the camera-ready. Note, Eq 18 refers to the critic loss (derived from the Taylor state and action expansion).\n\n> Why is $\\mu_\\text{a}$ in Eq 19 sampled from deterministic target policy?\n\nThere are really two separate questions here.  \n\nFirst, there is the question of whether $a$, the initial action in the TD update, is stochastic/deterministic. In our setting, $a$ is always stochastic, which helps learn a broader Q function with performance benefits over deterministic initial actions (e.g., see the poor performance of Dyna-TD3 baseline algorithm).\nHowever, stochastic initial actions also increase the TD-update variance relative to deterministic initial actions. \nTaylor-TD mitigates this increased variance by analytically approximating the expected TD update arising from stochastic initial actions.\n\nSecond, there is the question of whether the target policy (which we use to generate $a'$ in the TD update) is stochastic/deterministic.  It turns out the Taylor TD can be used with deterministic or stochastic target policies. We used deterministic target policies, because in a deterministic environment, the optimal policy is deterministic (or at least, _an_ optimal policy is deterministic).\n\n> In Fig 1, what is the variance computer over?\n\nWe agree that this wasn't quite clear enough in the original manuscript, and will clarify in the camera-ready. We used exactly the same initial state-action pairs when comparing standard (MC) TD and Taylor-TD updates. The additional variance in standard (MC) TD updates comes from the usual approach of adding a small amount of Gaussian noise to the initial action, $a$, in the TD update; while in Taylor-TD, we analytically integrate over a Gaussian distribution over actions (and states).\n\n> What is the effect of reducing the variance in critic update vs to that actor update(Ref[1])?\n\nWe thank the reviewer for bringing this paper to our attention, we will add it to the related work.\nWe would expect that the improvements from each method would compound, as one is reducing variance in the actor update, while the other is reducing variance in the critic update: we will attempt to perform this experiment for the camera-ready deadline.\n\nThat said it is important to note an additional complication.  In particular, TaTD can be used with deterministic or stochastic target policies, and our current experiments use deterministic target policies, as we are mostly working in deterministic environments.  However, [1] can offer no benefits with deterministic target policies, as [1] reduces variance from stochastic target policies.\n\n> Comparison with expected sarsa kind of update in tabular environments?\n\nFirst, as noted in the Abstract, Taylor-TD requires continuous state and action spaces, for the required derivatives, so the connection to tabular settings is unclear.\n\nMore fundamentally, expected SARSA computes the expected $Q(s', a')$ under the distribution over _next_ actions, $a'$, given the current policy.\nIn contrast, Taylor-TD computes the expected TD update under a distribution over _initial_ action, $a$.\nThus, it is possible to combine expected SARSA and Taylor-TD, and as they reduce different components of the variance, we would expect them to be complementary.\nHowever, it is important to note that expected SARSA only makes sense when we have a stochastic policy, which implies a need to compute the expectation of $Q(s', a')$ under a distribution over _next_ actions, $a'$.\nIn contrast, Taylor-TD allows the use of either stochastic or deterministic. \nAs noted above, we use deterministic target policies (as an optimal policy in a deterministic environment is deterministic). And if we use a deterministic target policy, there is no need for an expectation over next actions: we can just compute $Q(s', a')$.\n"}, {"review_id": "pueEmxVPtb", "reviewer": "Reviewer_wJvf", "summary": "This paper, \n1. Proposes Taylor TD, a model-based RL algorithm that uses a Taylor series expansion to analytically estimate the expected TD update over a distribution of nearby state-action pairs. This reduces variance compared to standard MC TD updates.\n2. Provides theoretical analysis showing the variance of Taylor TD updates is lower than standard TD updates. Also proves stability of Taylor TD with linear function approximation.\n3. Empirically demonstrates lower variance of Taylor TD on several RL tasks.\n4. Combines Taylor TD with TD3 in an algorithm called TaTD3. Shows strong performance of TaTD3 relative to model-free and model-based baselines on MuJoCo benchmarks.\n", "questions": "1. In the Taylor TD update derivations (Equations 12, 17 to 18), the Hessian terms present after the Taylor expansions disappear in the final update equation. Can you clarify the assumptions or steps that lead to the Hessian terms dropping out? As these terms arise directly from the Taylor approximations but then vanish, an explicit explanation should be provided about how and why they drop out of the final update form.\n2. It is also not clear how the loss function (equation 18) is derived. Could you please provide the derivation?\nAlgorithm 1 includes an update equation for the model parameters w using some loss L, but does not define what this loss function L is (unless I am mistaken). Please specify what objective or loss function is used to optimize the model parameters w based on the observed environment transitions.\n3. The current evaluations of Taylor TD are limited to a small set of MuJoCo continuous control tasks. Did the authors consider testing on additional environments that evaluate aspects like partial observability, sparse rewards, and higher task dimensionality? A more diverse test suite could provide greater insight into the strengths and weaknesses of Taylor TD across different conditions. What other test beds do you think could be useful for thorough analysis? What are the challenges that you foresee? \n4. The hyperparameters indicate high values for \\lambda_a and very low values for \\lambda_s. Why is that the case?\n", "limitations": "The authors mentioned about the limitations in the last section and I do not forsee any negative societal impacts. ", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The core idea of using Taylor expansions to estimate expected TD updates is novel and well-motivated from a theoretical perspective. This analytic approach to reducing variance is interesting.\n2. The method preserves the convergence guarantees of TD-learning under linear function approximation, as shown formally. This is an important theoretical contribution.\n3. The variance reduction analysis provides evidence Taylor TD reduces variance versus standard TD methods.\n4. The algorithm is straightforward to implement on top of existing model-based RL frameworks.\n", "weaknesses": "1. In Figure 2, Taylor TD3 seems to provide only very small performance improvements over the baselines in most environments (4 out of 6). The gains are noticeable in only 2 environments. Given that Taylor TD3 is much more complex, is the minimal improvement in the majority of environments concerning? The results seem to imply the benefits may be limited to certain environments. Some discussion of why the gains are so marginal in certain tasks would be useful.\n2. Some derivations are not clear. It is not clear what loss function are the authors using to learn the model. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments!\n\n> In Figure 2, Taylor TD3 seems to provide only very small performance improvements\n\nWhile we agree with the reviewer's assessment, it is important to note that there is no single method that performs competitively with TaTD3.\nParticularly, MBPO is competitive with TaTD3 on HalfCheetah-v2, Walker2d-v2 and Ant-v2, but learns extensively slower on the most complex environment, Humanoid-v2.\nIn contrast, MAGE performs well on HalfCheetah-v2 and the most complex environment, Humanoid-v2, but fails in other environments: Walker2d-v2, Ant-v2 and Hopper-v2.\nIn contrast, TaTD3 performs consistently (close to) the best in all these environments.\nIt should also be noted that Taylor TD3 is not necessarily more complex than the two strongest baseline algorithms, MBPO and MAGE.\nFor instance, although MBPO does not require any extra gradient term, it uses additional compute to iterate through the model predictions over multiple steps for each update (up to 25x).\n\n\n> Some derivations are not clear. It is not clear what loss function are the authors using to learn the model.\n\nWe thank the reviewer for pointing this out. We neglected to include a section in the appendix where we described how the model was learned. To clarify, the reward and transition model were learned by maximum likelihood based on the observed transition, as for most model-based RL approaches (e.g. Janner et al., 2019 - MBPO). A section describing this process will be added in the appendix of the camera-ready version of the paper. \n\n\n\n> In the Taylor TD update derivations (Equations 12, 17 to 18), the Hessian terms present after the Taylor expansions disappear\n\nWe have a choice about the order of the Taylor expansion.  For instance, we could choose to do a first-order expansion, which would include gradient but not Hessian terms, or we could choose to do a second-order Taylor expansion, which would include Hessian terms.  We choose the first-order approach.  We did play with the second-order terms, but we found they added considerable additional complexity for little obvious benefit.  Additionally, we were able to get theoretical guarantees on the first-order approach (Appendix B), and it wasn't clear we'd be able to get the same guarantees on the more complex second-order methods.\n\nBy the Hessian, are you referring to $\\Sigma_\\text{a}$, which is present e.g. in Eq 11, but not in Eq 12, 17 to 18?  $\\Sigma_\\text{a}$ isn't the Hessian, it is the covariance of the Gaussian distribution over initial actions, $a$, in the Bellman update (see Eqs 4-6).  We get to choose this covariance, and in going from 11 to 12, we choose $\\Sigma_\\text{a} = \\lambda_\\text{a} I$. We will re-emphasise this point in the camera-ready text.\n\n\n> It is also not clear how the loss function (equation 18) is derived.\n\nThe key derivation of the critic updates is in Sec. 3.1 (Eq. 12) for the actions and Sec. 3.2 (Eq. 17) for the states (see Appendix A for more details).  To get to Eq. 18 requires two steps.\nFirst, we apply the usual RL implementation trick, of writing the updates in terms of a gradient of a loss, with stop-gradients operations to avoid taking the gradient of $\\delta$ (Appendix D).\nSecond, we replace the dot product in Eq 12 with a cosine similarity.\nThis gives a useful normalization, which does seem to improve performance (see Appendix I.2).\nWe agree that these steps were not described clearly enough in the main text, so we will update the camera-ready to spell this out.\n\nThe reward and transition model were learned by maximum likelihood based on the observed transition, as in most model-based RL approaches (e.g. as in Janner et al., 2019 - MBPO). We will add an Appendix describing the model-learning to the camera-ready.\n\n\n\n> The current evaluations of Taylor TD are limited to a small set of MuJoCo continuous control tasks.\n\nWe agree with the reviewer that a more diverse test suite could provide greater insight into the strengths and weaknesses of Taylor TD across different conditions. \nPartial observability is complex, as it requires latent states in the model.  We suspect it should be possible, but anticipate that it is a \"research-project\" level exercise that is out of scope for this paper.\nAdditionally, Humanoid is typically considered a high dimensional control task with the RL literature, given the 376 dimensional state space and the 17 dimensional action space. \nIt should also be noted the set of MuJoCo continuous control tasks that we chose represent the standard benchmark on which RL algorithms are tested for continous control. For instance, the baseline algorithms we tested Taylor TD against were themselves tested on 6 similar MuJoCo continous control environments in the original papers and this is the case for most continuous control RL algorithms (e.g., SAC, TD3). \nNevertheless, we are happy to run further experiments for camera-ready if you can suggest specific benchmarks.\n\n\n\n> The hyperparameters indicate high values for \\lambda_a and very low values for \\lambda_s. \n\nIn general, we found the action-based expansion to bring the largest benefits to performance, although the state-based expansion is still useful (i.e., see Appendix I.1). We think one reason for this may be because, we know that Gaussian distributions over $a$ work well when evaluating the Bellman updates.  However, we don't know whether the same applies for the distribution of state; a Gaussian distribution may only work well for very nearby states. \nWe should also stress that any direct comparison between the magnitude of the state covariance (\\lambda_s) and the action covariance (\\lamba_a) may not be very meaningful.\nThis is because the scale of the actions and states tend be fairly different in the tested environments, making any comparison between $\\lambda_s$ and $\\lambda_a$ harder to assess. We will add this key discussion in section A.6. of the camera-ready.\n\n"}, {"review_id": "a3BxqA97b8", "reviewer": "Reviewer_9ew5", "summary": "This paper proposes an algorithm for improving model-based actor-critic methods in continuous state and action spaces. When using Dyna-style updates, the algorithm can compute an expected update over a small noise distribution using a linearization to reduce the variance of the performed update. This method is shown to be theoretically-sound in a simplified setting and a variety of experiments demonstrate the effectiveness of the proposed approach on standard continuous-control benchmark tasks. \n", "questions": "- Why does aligning the gradients w.r.t the action of the Q-value and the TD error make sense? Is there an interpretation for this? It seems like the inner product between those gradients could be a meaningful quantity more broadly e.g. as an evaluation metric for the quality of a critic. \n- What is $\\eta$ in equation (11)? It wasn't described in the text or the appendix.\n- This method could potentially also be used in cases when the simulator is differentiable to start with. Have you experimented with this at all? \n- Minor point: I would consider modifying the title of the paper to a more descriptive name which could reference model-based RL or Dyna-style updates. \n", "limitations": "The limitations are discussed.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The main idea is interesting, computing an explicit expected update over a small region by using a linearization. \n- The experiments are done well with meaningful baseline algorithms. In particular, the supporting experiments other than standard learning curves are a welcome addition: The variance analysis comparing sampled updates to the proposed one and the ablation study on cosine similarity vs. the inner product (in the appendix).\n- The theoretical result in the linear setting is nice even if it is in a simplified setting and the proofs are correct as far as I can tell.\n\n", "weaknesses": "There are some aspects of the presentation that I think could be improved. These are more minor points overall though.\n- For example, I would avoid using the term \"Monte-Carlo\" estimates when discussing sampling of TD updates since it can lead to some confusion due to \"Monte-Carlo\" often being used to refer to MC estimates of the return in contrast to the bootstrap estimates that TD uses. I would consider using the phrase \"sampled TD update\" instead of \"MC TD update\". E.g. line 31, line 42, ...\n- I would also suggest including a line about continuous state and actions in the introduction to clarify the problem setting since the proposed algorithm would mainly be applicable there. \n- Section 3.1 describing the updates was a bit difficult to follow at first. Some of the notation such as using $\\Delta_{Exp}$ was slightly confusing since it suggests that it is the overall expected update even though it's only the expected update over $\\xi_a$. An alternative could be $\\Delta_{\\mu_a}$. The expectation $E_{s,a$}[\\Delta \\theta(s,a)]$ is slightly unclear as to what distribution is used for $(s,a)$ here. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments!\n\n> For example, I would avoid using the term \"Monte-Carlo\" estimates when discussing sampling of TD updates since it can lead to some confusion due to \"Monte-Carlo\" often being used to refer to MC estimates of the return in contrast to the bootstrap estimates that TD uses. I would consider using the phrase \"sampled TD update\" instead of \"MC TD update\". E.g. line 31, line 42, ...\n\nThanks for spotting this clash with pre-existing RL terminology.  We will switch to your suggested terminology.\n\n\n> I would also suggest including a line about continuous state and actions in the introduction to clarify the problem setting since the proposed algorithm would mainly be applicable there.\n\nThank you for the suggestion. This is included in the original abstract, and we will include a note about it in the Introduction too.\n\n\n> Section 3.1 describing the updates was a bit difficult to follow at first. Some of the notation such as using $\\Delta_\\text{Exp}$ was slightly confusing since it suggests that it is the overall expected update even though it's only the expected update over $\\Xi_\\text{a}$. An alternative could be $\\Delta_{\\mu_a}$. \n\nWe'll consider these notational changes.  We hoped that the $(s, \\mu_\\text{a})$ function arguments in $\\Delta_\\text{Exp}(s, \\mu_\\text{a})$ would emphasise that this is taking an expectation only over the noise in the action, for fixed state, $s$ and mean action, $\\mu_\\text{a}$.\nThe issue with $\\Delta_{\\mu_a}$ is that we contrast the exact expectation, $\\Delta_\\text{Exp}$, with the Taylor-series approximation, $\\Delta_\\text{Ta}$. But they're both exact/approximate expected updates over $\\Xi_\\text{a}$ with fixed $\\mu_a$, so $\\Delta_{\\mu_a}$ wouldn't make the required distinction.\n\nWe could change $\\Delta_\\text{Exp} \\rightarrow \\Delta_\\text{Exact}$?  Alternatively, we could further clarify the definition of $\\Delta_\\text{Exp}$ in Eq. 7?  Let us know if you have a preferred option.\n\n\n\n> The expectation $E_{s,a}[\\Delta \\theta(s,a)]$  is slightly unclear as to what distribution is used for here.\n\nWe will add a sentence clarifying the distribution here. To clarify, this distribution refers to the overall visited state distribution, $s \\sim d^\\pi$, and the initial policy distribution, $a \\sim \\pi^\\text{init}$.  \n\n\n\n\n> Why does aligning the gradients w.r.t the action of the Q-value and the TD error make sense? Is there an interpretation for this? It seems like the inner product between those gradients could be a meaningful quantity more broadly e.g. as an evaluation metric for the quality of a critic.\n\nFirst, it is important to emphasise that our updates should just be understood as the expectation of standard TD updates (Eq. 7) over a distribution over initial actions, $a$, centred at $\\mu_\\text{a}$. In that context, it is possible to understand how the \"alignment\" emerges.  Specifically, lets consider a random action in the direction of $\\nabla_a \\delta_\\theta$, e.g. $a = \\mu_\\text{a} + \\epsilon \\nabla_a \\delta_\\theta$.  That action has a bigger, positive prediction error (as we've moved along the gradient of $\\delta_\\theta$). And under standard TD updates, a bigger positive prediction error should lead to a bigger positive update to $Q_\\theta$.  This interaction of bigger positive prediction errors leading to bigger positive updates to $Q_\\theta$ implies an \"alignment\" intuition.\n\nIts a super-interesting idea to use this quantity as an evaluation metric!  It certainly could work!  But we worry that doing a sufficiently thorough investigation of the properties of this metric is a significant piece of work requiring its own paper.\n\n\n> What is $\\eta$ in equation (11)? It wasn't described in the text or the appendix.\n\nGood catch!  We have deleted the $\\eta$'s.  (They were intended to represent a learning rate in a previous version of the derivation, but we subsequently removed it to simplify the derivations).\n\n\n> This method could potentially also be used in cases when the simulator is differentiable to start with. Have you experimented with this at all?\n\nAgreed, Taylor TD has broad applications to differentiable simulator settings. For instance, [1] highlights the importance of learning accurate value functions with differentiable simulator in order to perform policy updates beyond a short horizon of returns. Taylor TD would greatly aid the process of learning accurate value functions with little-to-none added computational cost to the differentiable simulator. However, in the current manuscript, we decided to focus on the usual (and potentially harder) RL setting where the model must be learned from the transitions (e.g. as in Janner et al., 2019 - MBPO).\nWe'd be happy to run any additional experiments if the reviewer can suggest a specific setting.\n     \n[1] Xu, J., Makoviychuk, V., Narang, Y., Ramos, F., Matusik, W., Garg, A., \\& Macklin, M. (2022). Accelerated policy learning with parallel differentiable simulation. arXiv preprint arXiv:2204.07137.\n\n\n> Minor point: I would consider modifying the title of the paper to a more descriptive name which could reference model-based RL or Dyna-style updates.\n\nSomething like \"Taylor TD-learning in model-based reinforcement learning\"?  We'd be happy to make that change."}]

Author rebuttal:

Meta-review: This paper studies the problem of variance in the temporal difference (TD) update that is widely used in RL. The paper proposes a method for reducing variance based upon an analysis using the Taylor expansion of the TD update. The method requires the use of a learned state transition model and hence can be seen as a hybrid model-free / model-based RL method. The paper primarily shows that the method reduces variance and leads to more data efficient RL through empirical analysis in standard continuous control RL tasks.

Main Strengths: the method is novel and the evaluation of it is sound and shown to provide empirical benefit for an important operator in RL. The proposed method is motivated from solid theoretical analysis.

Main Concerns: Main concerns from the reviewers were lack of clarity in some places and derivations. These can be fixed in the camera ready revision. One reviewer pointed out a number of references that the authors should decide whether they want to discuss or not in the revised version. I would add "Reducing Sampling Error in Batch Temporal Difference Learning" [Pavse et al. 2020] as another paper that seems related (though not conflicting with this paper's novelty) and the authors should review it and consider its relevance. Missing citations and need for some clarifications were not major concerns.

Overall, this paper presents a theoretially motivated method for improving the TD update that leads to empirical improvement in deep RL experiments. Given the importance of the TD update in many RL algorithms I believe this paper makes a valuable contribution to the community.
Decision: Accept (poster)