Title: Taylor TD-learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement learning

TD-learning

model-based

variance reduction
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

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 stateof-the art model-free and model-based baseline algorithms on a set of standard benchmark tasks.

1 Introduction

Actor-critic algorithms underlie many of the recent successes of deep RL [1, 2, 3, 4, 5, 6, 7, 8]. In these algorithms, the actor provides the control policy while the critic estimates the policy's expected long-term returns [i.e. the value function; 9, 10]. The critic is typically trained using some form of temporal-difference (TD) update [e.g. 3, 7, 1, 2, 6]. These TD updates need to be computed in expectation over a large distribution of visited states and actions, induced by the policy and the environment dynamics [11, 12]. Since this expectation is analytically intractable, TD updates are typically performed based on individually sampled state-action pairs from real environmental transitions (i.e. sample-based estimates). However, the variance of (sample-based) TD updates can be quite large, meaning that we need to average over many TD updates for different initial states and actions to get a good estimate of the expected updates [13].

Model-based strategies provide a promising candidate to tackle this high variance [14]. For instance, Dyna methods, among the most popular model-based strategies, use a learned model of the environment transitions to generate additional imaginary transitions. These imaginary transitions can be used as extra training samples for TD methods [e.g. 15, 16, 17, 18, 19, 20]. Although the additional (imaginary) transitions help in reducing the variance of performing TD-updates over a

large distribution of state-action pairs, Dyna methods still rely on the same, potentially high-variance (sample-based) TD-updates as standard TD-learning.

We address the issue of high-variance TD-updates by formulating an expected TD-update over a small distribution of state-action pairs. We show this expected update can analytically be estimated with a first-order Taylor expansion for continuous state-action spaces, in an approach we call Taylor TD . By analytically estimating this expected update, rather than exclusively relying on sample-based estimates (as in e.g. Dyna), we get lower variance TD updates. Additionally, we show Taylor TD does not affect the stable learning guarantees of TD-learning under linear function approximation [for a fixed policy as shown by 21]. Next, we propose a model-based off-policy algorithm, Taylor TD3 (TaTD3), which uses Taylor TD in combination with the TD3 algorithm [1]. We show TaTD3 performs as well as if not better than several state-of-the art model-free and model-based baseline algorithms on a set of standard benchmark tasks. Finally, we compare TaTD3 to its 'Dyna' equivalent, which exclusively relies on sample-based TD-updates. We found the largest benefits of Taylor TD may appear in high dimensional state-action spaces.

2 Background

Reinforcement learning aims to learn reward-maximising behaviour by interacting with the surrounding environment. At each discrete time step, t , the agent in state s ∈ S , chooses an action a ∈ A based on a policy π : S → A , and observes a scalar reward, r and a new state s ′ ∈ S from the environment. The agent's goal is to find the policy that maximises the expected sum of rewards (i.e. the expected return), from a distribution of initial states (or state-action pairs). As such, it is usually necessary to compute the expected return for a state-action pair ( s , a ) and a policy π ; which we can do with a value function. Given a policy π and an initial state-action pair ( s , a ) , we define the value function Q π ( s , a ) = E[ R t | S t = s , A t = a ] , where R t = ∑ T i = t γ i -t r i is the discounted sum of future rewards from the current time step t until termination T , with discount factor γ ∈ [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 s and following the policy π thereafter.

In off-policy learning, a function approximation, Q θ , is typically used to approximate the true value function, Q π target of a target policy, π target , based on experience collected from behavioural policies, π behave . The function approximation is typically updated with some form of TD-learning [11, 22],

where B denotes a reply buffer containing a mixture of states and actions sampled according to behavioural policies, π behave . Note that we have not included a learning rate in these equations, so ∆ θ ( 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 s as the initial state, a as the initial action, s ′ as the next state and a ′ as the target action. In standard (model-free) off-policy learning, the initial state and action, as well as the reward and the next state ( s , a , r, s ′ ) are drawn from the replay buffer [e.g., 23], while the next action comes from the target policy, a ′ ∼ π target ( s ′ ) . Critically, this off-policy setting allows to learn Q π target , even when the initial states and actions, ( s , a ) from the replay buffer, come from a different policy (or even an unknown policy).

In the model-based setting, we can exploit this flexibility in the distribution over the initial state, s , and action, a . Specifically, we sample the initial state, s from an initial-state distribution, d init (if we are taking initial states from the replay buffer, then d init , would be the distribution over states in the buffer). Next, we sample the initial action, a , from an arbitrary initial policy, π init ( s ) (i.e., the initial action, a , does not have to come from the buffer as for the model-free case). Then, we sample the reward and next state ( r, s ′ ) from the (learned) model, and finally we sample the next action from the target policy,

In principle, the off-policy setting allows for complete freedom in the distribution over the initial state, d init , and the initial policy, π init . That freedom is necessary in the model-free off-policy setting, because the states and actions in the replay buffer often come from an unknown distribution (e.g., all past behavioural policies). Although most model-based approaches take the initial policy π init to

equal either the target π target or behavioural policy π behave (i.e., the policy used to interact with the environment) [e.g., 19, 18], this does not have to be the case. For instance, π init could be a broader distribution than π behave , enabling us to update and improve the action value function for a broader range of actions.

3 Taylor TD

The expected update over actions in Eq. (1) is usually intractable, due to the complexity of the function, δ θ ( s , a ) ∇ θ Q θ ( s , a ) . Standard TD-learning methods therefore employ a sample-based approach, by sampling initial actions from some policy (distribution) π init and performing the (TD) updates for those actions. However, sampling can give rise to high-variance estimators, and hence slow learning [24]. Taylor TD reduces this variance by introducing an analytic approximation to the expectation in Eq. (1) over some stochasticity in the distribution over actions (and over the state, i.e., Sec. 3.2). This approach requires a differentiable model of the transitions and rewards as well as continuous state and action space.

Taylor TD takes advantage of the flexibility in choice of the initial, target and behavioural policies arising in model-based off-policy settings (see Sec. 2). Taylor TD can use any target and behavioural policy, but assumes some structure in the initial policy, π init . Specifically, Taylor TD splits the process of sampling the initial action into two steps. It first samples the mean, µ a , from a mean-policy, π action-mean ( s ) (usually, we would use the target policy, π target , as the mean-policy, but we do not have to). Then Taylor TD samples 'noise' from a second distribution, π action-noise , and the actual (initial) action is the sum of the mean and the noise,

Here, we require that the noise, ξ a , drawn from π action-noise has zero expectation, and finite, known covariance, Σ a ,

The overall initial policy, combining these two steps, can be written, a ∼ π init ( s ) .

3.1 Action expansion

Here, we analytically approximate the expectation over the action-noise, ξ a , using a first-order Taylor expansion. Specifically, we define the expected TD update, averaging over the action noise, ξ a , as,

The (overall) expected TD update in Eq. 1 can now be written in terms of ∆ Exp θ ( s , µ a ) ,

The exact expectation for ∆ Exp θ ( s , µ a ) (Eq. 7) is also not tractable typically, due to the complexity of the function, δ θ ( s , µ a + ξ a ) ∇ θ Q ( s , µ a + ξ a ) . Instead, Taylor TD uses ∆ Ta θ ( s , µ a ) to approximate ∆ Exp θ ( s , µ a ) (Eq. 7),

Specifically, ∆ Ta θ ( s , µ a ) , arises from a first-order Taylor expansion around ξ a = 0 ,

This reduces to (see Appendix A.1),

For simplicity, we further assume that the covariance of initial actions is isotropic, Σ a = λ a I (although this is not a requirement), resulting in the following (1st-order) Taylor TD-update:

This Taylor TD-update includes the standard TD update at state s and (mean) action µ a plus a new term, which tries to align the action gradient of the critic (Q-function) with the action gradient of the TD target. Conceptually, this gradient matching should help reduce the variance across TDupdates since it provides a way to analytically integrate over some of the stochasticity induced by π init . When doing a first order Taylor series approximation, we may worry that errors in the Taylor expansion might affect convergence. In the Appendix B, we include a proof that for linear function approximation, any errors in the first-order Taylor expansion do not affect the stability of TD-learning (assuming small simulator time steps). Importantly, this result is not a trivial, as there are still approximation errors in the mapping from state and action to the features that we use for linear function approximation. Moreover, we provide empirical evidence that the first-order Taylor expansion reduces the variance of standard TD-updates and support efficient learning, even under non-linear function approximation (see Sections, 3.4, and 4.2.3).

3.2 State expansion

We are not limited to analytically approximate the expectation in Eq. (1) for the distribution of initial actions, but we can extend this approach to a distribution of initial states. Namely, instead of performing a TD-update at the single initial state, s , we perform this update over a distribution of states. Similarly to the action expansion, we assume structure in the distribution over the initial state distribution (denoted here as d init ). Specifically, sampling initial states from the initial state distribution involves two steps. We first sample the mean of d init , µ s , from a state-mean distribution d state-mean (usually, we would use the distribution of states in a reply buffer as the state-mean distribution, but we do not have to). Then we sample state-noise from a second distribution, ξ state-noise , and the actual (initial) state is the sum of the state-mean and the state-noise,

Again, we require that the state-noise, ξ s , drawn from d state-noise has zero expectation, and finite, known covariance,

Thus, we can again formulate an expected TD-update, averaging over the state-noise,

Again, we can approximate this expected update with a first-order Taylor approximation, but this time, we expand the states around ξ s = 0 instead of the actions. This requires that the states are continuous, which is usually the case for control tasks. Based on a similar derivation to the action expansion, we get the following Taylor TD-update, where we assumed the state covariance over TD updates to be isotropic, Σ s = λ s I (although again this is not a requirement) (see Appendix A.2 for the full derivation):

The rational behind this update is trying to tackle some of the TD-update variance induced by the initial state distribution (e.g., the distribution of sates in a reply buffer), although we expect this only to work for states close-by to the visited ones (i.e. for small values of λ s ).

3.3 State-Action expansion implementation

Finally, we can combine the two Taylor expansions into a single TD-update involving both state and action expansions. Nevertheless, computing the dot products between ∇ δ θ and ∇ Q θ 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

this). The cosine distance has the advantage of being bounded. By putting everything together, we propose a novel TD update, which we express below in terms of a loss:

Note we used the notation δ instead of δ θ to indicate we are treating δ ( µ s , µ a ) as a fixed variable independent of θ . This ensures when we take the gradient of this loss relative to θ , we do not differentiate through any δ terms (following the standard implementation of TD-updates in autodiff frameworks such as PyTorch, see Appendix D). In practice, Taylor TD requires a differentiable model of the environment transitions as well as reward function in order to compute the terms ∇ a δ θ ( µ s , µ a ) and ∇ s δ θ ( µ s , µ a ) (see Appendix C for further details).

In Sections 3.1 and 3.2, we have been careful to setup Taylor TD over actions and states in the most general possible setting. Nevertheless, a set-up that may work well with Taylor TD is following,

where values for µ a (i.e., the mean-policy) are the outputs of a deterministic (target) policy π det-target and µ s (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 π det-target ) over a distribution of actions centered at π det-target and controlled by λ a as well as over distributions of states centered at previously visited states (i.e., from the reply buffer) and controlled by λ s. We provide a Taylor TD algorithm implementation with this set-up in the Algorithm box 1. Note, in the Algorithm box 1, we left the model loss unspecified (i.e., L model ) to denote that any valid method can be used to learn the model of the transitions. However, in the experiments below, we always learned the transition model based on maximum likelihood on the observed environment transitions (see Appendix E).

3.4 Variance analysis

Here, we show that the expected TD update in Eq. (7) (which Taylor TD estimates) has lower variance than standard (sample-based) TD-updates over the same distribution of actions. We only provide this variance analysis for the distribution over actions, because an analogous theorem can be derived for the distribution over states (i.e. Eq. 16). To begin, we apply the law of total variance [26] to standard TD-updates,

The inner expectation and inner variance samples action-noise, ξ a , while the outer expectation and outer variance samples initial states, s , and action-means, µ a . To relate this expression to Taylor TD, recall that Taylor TD updates are motivated as performing analytic integration over action-noise, ξ a - i.e. ∆ Exp θ ( s , µ a ) = E ξ a [∆ θ | s , µ a ] . Thus, the variance of standard (sample-based) TD-updates (Eq. 20) is exactly the variance of ∆ Exp θ , plus an additional term to account for variance induced by sampling action-noise,

The fact that E s , µ a [ Var ξ a [∆ θ | s , µ a ] ] is non-negative directly gives a theorem.

Theorem 3.1. The variance for standard (sample-based) TD-updates, Var s , µ a , ξ a [∆ θ ] , is larger than (or equal to) the variance if we exactly integrate action-noise, Var s , µ a [∆ Exp θ ] ,

While this theorem only formally applies to the exact expectation, ∆ Exp θ , we would expect it to become increasingly accurate as λ a and λ s become smaller, and hence the Taylor series approximations become more accurate. Furthermore, in Appendix B we prove that the resulting updates do not affect the stability of (standard) TD-learning even when the Taylor series approximation errors do not vanish (under linear function approximation). Additionally, we show empirical reductions in the update variance for Taylor TD in Sec. 4.1.

4 Experiments

4.1 Variance reduction

In this section we empirically test the claim that Taylor TD updates are lower variance than standard (sample-based) TD-learning updates. To do so, we take the Taylor TD set-up in Algorithm 1 and perform "batch updates" [12], where given an approximate value function Q θ and a policy π , several Q θ updates are computed across several sampled states and actions, updating Q θ only once, based on the sum of all updates. Batch updates ensure the variance of the updates is estimated based on the same underlying value function. We compute batch updates for both Taylor TD and standard (sample-based) TD updates, comparing the variance of the updates between the two approaches (see Appendix F for more details). Formally, we aim to (empirically) show that the inequality in Theorem 3.1 (for both states and actions) holds even with the Taylor approximation across a range of standard control tasks,

Fig. (1) shows Taylor TD provides reliably lower variance updates compared to standard TD-learning across all tested tasks, but the Pendulum environment. This finding is consistent with the idea that Taylor TD updates may be most beneficial in higher dimensional state-action spaces, while the Pendulum environment represents the lowest dimensional environment with only four dimensions. This is because sampling estimates (i.e., standard TD-learning) may become less efficient in higher dimensional spaces. We further explore this proposal in a toy example comparing sample-based estimates and Taylor expansions of expected updates. We perform this comparison across data points of different dimensions, and find the benefits of the Taylor expansion (over purely sample-based estimate) to increase with the dimension of the data points (see Appendix H). Additionally, in the Appendix G, we investigate how the variance of Taylor TD and standard TD-learning updates changes during training.

Pendulum-v1

Hopper-v2

Half-Cheetah-v2

Humanoid-v2

20

14

6

8

40

12

40

Meanupdatevariance

5

15

10

(x1000)

4

6

30

8

30

10

3

4

20

6

20

2

2

10

5

4

10

1

2

0

Q

0

0

0

0

4.2 Benchmark performance

4.2.1 Algorithm

We combine Taylor TD (i.e. Algorithm 1) with the TD3 algorithm [1] in a model-based off-policy algorithm we call Taylor TD3 (TaTD3) (code available at Appendix I). TaTD3 aims to provide a state-of-the-art implementation of Taylor TD for comparison with baseline algorithms. At each iteration, TaTD3 uses a learned model of the transitions and learned reward function to generate several differentiable (imaginary) 1-step transitions (i.e., Dyna style), starting from real states sampled from a reply buffer. These differentiable 1-step transitions are then used to train two critics (i.e. TD3) using Taylor TD updates. The model of the transitions consists of an ensemble of 8 models trained by maximum likelihood on the observed environment transitions (see Appendix E). These models return a Gaussian distribution over the next state, with zero correlations, and with mean and variance given by applying a neural network to the previous state-action pair. We also learn a model the rewards using a neural network trained with mean-square error on the observed rewards. Hence, TaTD3 does not require any a priori knowledge of the true environment transitions or true reward function. Finally, the actor is trained with the deterministic policy gradient [7] on real states as in standard TD3 [1].

4.2.2 Environments

We employ 6 standard environments for continuous control. The first environment consists of a classic problem in control theory used to evaluate RL algorithms [i.e. Pendulum, 27]. The other 5 environments are stanard MuJoCo 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 art model -free and -based baselines on the six benchmark environments. These baselines include 3 model-based algorithms and one 1 model-free algorithm (see Appendix J). The first model-based algorithm is Model-based Policy Optimization (MBPO) [18], which employs the soft actor-critic algorithm (SAC) [2] within a modelbased Dyna setting. The second model-based algorithm is Model-based Action-Gradient-Estimator Policy Optimization (MAGE) [19], which uses a differentiable model of the environment transitions to train the critic by minimising the norm of the action-gradient of the TD-error. The third model-based algorithm is TD3 combined with a model-based Dyna approach (i.e. Dyna-TD3). This algorithm was proposed by [19] and was shown to outperform its model-free counterpart, TD3 [1] on most benchmark tasks. Dyna-TD3 is conceptually similar to MBPO, with the main difference of MBPO relying on SAC instead of TD3. Finally, we included SAC [2] as a model-free baseline. Fig. (2) shows TaTD3 performs at least as well, if not better, than the baseline algorithms in all six benchmark tasks: note the much poorer performance of MAGE on Hopper-v2, Walker2d-v2 and Ant-v2, as well as of MBPO on Humanoid-v2 relative to TaTD3.

Pendulum-v1

Hopper-v2

HalfCheetah-v2

4000

12000

Fo

10000

return

-500

3000

8000

Averag

-1000

2000

6000

4000

-1500

1000

2000

0

0

0

3k

6k

9k

0

25k

50k

75k

0

50k

100k

150k

Humanoid-v2

4000

8000

Average return

4000

3000

6000

2000

4000

2000

1000

2000

1

0

0

50k

100k

150k

0

50k

100k

150k

0

50k

100k

150k

TaTD3 (ours)

MBPO

MAGE

Dyna-TD3

SAC

Pendulum-v1

HalfCheetah-v2

Humanoid-v2

0

12000

8000

-250

10000

5000

6000

return

-500

8000

4000

Average

-750

6000

3000

4000

-1000

4000

2000

-1250

1000

2000

2000

-1500

0

0

T

1

0

1

0

5k

10k

0

50k

100k

150k

0

50k

100k

150k

0

50k

100k

150k

Steps

Steps

Steps

Steps

TaTD3

Sample-basedExpected-TD3

4.2.4 Taylor vs sample-based (standard) TD-learning

Next, we ask whether Taylor TD provides any performance benefit in computing the expected TD updates in Eq. (7) and (16) over standard sample-based estimates. To do so, we implement a model-based TD3 algorithm analogous to TaTD3, but where the expected TD updates, Eq. (7) and (16), are estimated using sample-based estimates by drawing multiple samples for the state ( ξ s ) and action ( ξ a ) noise at each step, instead of using an analytic Taylor series approximation. We call this algorithm Sample-based Expected-TD3 (code available at Appendix I). In practice, at each time step, Sample-based Expected-TD3 uses a (learned) model of the transitions to compute multiple TD-updates by sampling ten state perturbations ( ξ s ) of visited states ( µ s ) and ten action perturbations ( ξ a ) of the deterministic target policy ( µ a )(i.e. estimating Eq. (7) and (16) through a sample-based estimate). Crucially, we ensure the variance of the state and action perturbations

(i.e. λ a and λ s) is matched between TaTD3 and Sample-based Expected-TD3. We perform this comparison across 4 standard continuous control tasks, spanning a low dimensional (Pendulum), two moderate dimensional (Half-Cheetah and Walker2d) and one high dimensional environment (Humanoid).

In Fig. (3), we can see TaTD3 provides performance benefits over Sample-based Expected-TD3 across the three most high dimensional environments. This finding is in line with the claim we put forth in the variance reduction section (i.e., 4.1) that the benefits of Taylor TD (i.e TaTD3) over sampling may be most marked in high dimensional state-action spaces. Indeed, the largest performance advantage of TaTD3 is seen in Humanoid-v2, which has the highest dimensional state-action space by a large margin. Conversely, the least difference between TaTD3 and Sample-based Expected-TD3 is seen in Pendulum-v1, which is the task with smallest dimensional state-action space. The relation between the variance of TD-updates and the dimension of the state-action space should be better explored in the future (see also Appendix H). Finally, we should stress Sample-based Expected-TD3 is different from Dyna-TD3, as the latter does not implement any action or state perturbation in the TD-updates. Hence, unlike Sample-based Expected-TD3, Dyna-TD3 does not directly estimate the expected updates in Eq. (7) and (16), but relies on (standard) TD-learning updates (this is also evident in the massive performance difference between Dyna-TD3 and Sample-based Expected-TD3 - by comparing corresponding performance between Fig. 2 and 3).

5 Related work

Model-based strategies provide a promising solution to improving the sample complexity of RL algorithms [14]. In Dyna methods, a model of the environment transitions is learned through interactions with the environment and then employed to generate additional imaginary transitions [e.g. in the form of model roll-outs, 15]. These imaginary transitions, can be used to enhance existing model-free algorithms, leading to improved sample complexity. For example, within TD-learning, imaginary transitions can be used to train a Q-function by providing additional training examples [e.g. 15, 16, 19]. Alternatively, imaginary transitions can be used to provide better TD targets for existing data points [e.g. 17] or to train the actor and/or critic by generating short-horizon trajectories starting at existing state-action pairs [e.g. 18, 29, 20]. These (Dyna) approaches have a clear relation to our approach (Taylor TD), as they attempt to estimate a similar expected TD-update as in Eq. (7). However, Dyna approaches only use potentially high-variance sample-based estimates, while Taylor TD exploits analytic results to reduce that variance.

Conceptually, our approach may resemble previous methods that also rely on analytical computations of expected updates to achieve lower-variance critic or policy updates [e.g. 24, 30, 31, 12, 32, 33, 34] [see also 35, 36, 37, 38, for a different set of approaches relating to update variance in RL]. The most well-known example of this is Expected-SARSA. Expected-SARSA achieves a lower variance TD-update (relative to SARSA), by analytically computing the expectation over the distribution of target actions in the TD-update (i.e. assuming a stochastic target policy) [31, 12];

This approach can only reduce variance of TD-updates at the level of the target actions, a ′ , induced by a stochastic target policy. In the case of a deterministic target policy, Expected-SARSA does not provide any benefit. Conversely, our approach attempts to reduce the variance at the level of the initial state-action pairs, ( s , a ) at which TD-updates are performed. That is; we take the expectation over ( s , a ) instead of a ′ (see Eq. 7 and 16) which should yield benefits with both stochastic and deterministic target policies. Other well-known RL approaches exploiting analytical computations of expected updates are Expected Policy Gradients [24, 30, 33], Mean Actor Critic [32] and "all-action" policy gradient [34]. These methods attempt to reduce the variance of the stochastic policy gradient update by integrating over the action distribution. Although similar in principle to our approach, these methods focus on the policy update instead of the critic update and, similarly to Expected-SARSA only apply to stochastic target policies.

In practice, our approach may relate to value gradient methods, as it explicitly incorporates the gradient of the value function into the update [e.g. 39, 29, 40, 19, 41, 42]. To our knowledge, the value gradient approach that most relates to our work is MAGE [19], which, nonetheless, has a radically different motivation from Taylor TD. MAGE is motivated by noting that the action-gradients of Q drive deterministic policy updates [7], so getting the action-gradients right is critical for policy

learning. In order to encourage the action-gradients of Q to be correct, MAGE explicitly adds a term to the objective that consists of the norm of the action-gradient of the TD-error, which takes it outside of the standard TD-framework. In contrast, our motivation is to reduce the gradient variance across standard TD updates by performing some analytic integration. We do this through a first-order Taylor expansion of the TD update. This difference in motivation leads to numerous differences in the method and analysis, the least of which is that MAGE uses only the action-gradients, while Taylor TD additionally suggests using the state-gradients, as both the state and action gradients can be used to reduce the variance in the (TD) updates.

The idea of applying a Taylor expansion to learning updates is not new to RL. For instance, [43] uses a (model-free) Taylor expansion of the policy update to provide a formalism connecting trust-region policy search with off-policy correction. This is very different from our approach (Taylor TD), not only because Taylor TD is model-based, but also because we apply the Taylor expansion to the critic instead of the actor update for very different reasons (i.e., reduce the update variance across TD updates). Finally, Taylor TD may be applicable to risk-sensitive RL [e.g., 44, 45, 46, 47]. In the presence of a good model of the transitions, Taylor TD may be able to approximate the values of risky actions (or states) around safe (e.g. deterministic) actions (or visited states), without actually needing to take those actions (or visit those states), thanks to the Taylor expansion of the TD objective.

6 Conclusion and Limitations

In this article, we introduce a model-based RL framework, Taylor TD, to help reduce the variance of standard TD-learning updates and, speed-up learning of critics. We theoretically and empirically show Taylor TD updates are lower variance than standard (sample-based) TD-learning updates. We show the extra gradient terms used by Taylor TD do not affect the stable learning guarantees of TD-learning with linear function approximation under a reasonable assumption. Next, we combine Taylor-TD with the TD3 algorithm [1] into a model-based off-policy algorithm we denote as TaTD3. We show TaTD3 performs as well, if not better, than several state-of-the art model-free and model-based baseline algorithms on a set of standard benchmark tasks.

Taylor TD has the limitation that it requires continous state-action spaces as well as a differentiable model of transitions to calculate the additional (TD) gradient terms, i.e. it must be in the model-based rather than model-free setting (see Appendix C). Additionally, the gradient terms in Taylor TD imply additional computational cost; in the Appendix L, we show this cost is not that large in terms of training times and we expect it to reduce as faster automatic differentiation tools are developed [48].

References

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.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. PMLR, 2018.

Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971 , 2015.

John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. Highdimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438 , 2015.

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.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning , pages 387-395. PMLR, 2014.

Claas Voelcker, Victor Liao, Animesh Garg, and Amir-massoud Farahmand. Value gradient weighted model-based reinforcement learning. arXiv preprint arXiv:2204.01464 , 2022.

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.

Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems , 12, 1999.

Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning , 3(1):9-44, 1988.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Michael Fairbank and Eduardo Alonso. The local optimality of reinforcement learning by value gradients, and its relationship to policy gradient learning. arXiv preprint arXiv:1101.0428 , 2011.

Leslie Pack Kaelbling, Michael L Littman, and Andrew W Moore. Reinforcement learning: A survey. Journal of artificial intelligence research , 4:237-285, 1996.

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.

Shixiang Gu, Timothy Lillicrap, Ilya Sutskever, and Sergey Levine. Continuous deep q-learning with model-based acceleration. In International conference on machine learning , pages 28292838. PMLR, 2016.

Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I Jordan, Joseph E Gonzalez, and Sergey Levine. Model-based value estimation for efficient model-free reinforcement learning. arXiv preprint arXiv:1803.00101 , 2018.

Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in Neural Information Processing Systems , 32, 2019.

Pierluca D'Oro and Wojciech Ja´ skowski. How to learn a useful critic? model-based actiongradient-estimator policy optimization. Advances in Neural Information Processing Systems , 33:313-324, 2020.

Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sampleefficient reinforcement learning with stochastic ensemble value expansion. Advances in neural information processing systems , 31, 2018.

John Tsitsiklis and Benjamin Van Roy. Analysis of temporal-diffference learning with function approximation. Advances in neural information processing systems , 9, 1996.

Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning , 8:279-292, 1992.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602 , 2013.

Kamil Ciosek and Shimon Whiteson. Expected policy gradients. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 32, 2018.

Chunjie Luo, Jianfeng Zhan, Xiaohe Xue, Lei Wang, Rui Ren, and Qiang Yang. Cosine normalization: Using cosine similarity instead of dot product in neural networks. In International Conference on Artificial Neural Networks , pages 382-391. Springer, 2018.

Paul Bratley, Bennet L Fox, and Linus E Schrage. A guide to simulation . Springer Science & Business Media, 2011.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540 , 2016.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems , pages 5026-5033. IEEE, 2012.

Ignasi Clavera, Violet Fu, and Pieter Abbeel. Model-augmented actor-critic: Backpropagating through paths. arXiv preprint arXiv:2005.08068 , 2020.

Shimon Whiteson. Expected policy gradients for reinforcement learning. Journal of Machine Learning Research , 21, 2020.

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.

Kavosh Asadi, Cameron Allen, Melrose Roderick, Abdel-rahman Mohamed, George Konidaris, Michael Littman, and Brown University Amazon. Mean actor critic. stat , 1050:1, 2017.

Matthew Fellows, Kamil Ciosek, and Shimon Whiteson. Fourier policy gradients. In International Conference on Machine Learning , pages 1486-1495. PMLR, 2018.

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:1910.09093 , 2019.

Maxim Kaledin, Alexander Golubev, and Denis Belomestny. Variance reduction for policygradient methods via empirical variance minimization. arXiv preprint arXiv:2206.06827 , 2022.

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.

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. PMLR, 2013.

Matthew J Sobel. The variance of discounted markov decision processes. Journal of Applied Probability , 19(4):794-802, 1982.

Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. Advances in neural information processing systems , 28, 2015.

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.

David Balduzzi and Muhammad Ghifary. Compatible value gradients for reinforcement learning of continuous deep policies. arXiv preprint arXiv:1509.03005 , 2015.

Michael Fairbank and Eduardo Alonso. Value-gradient learning. In The 2012 international joint conference on neural networks (ijcnn) , pages 1-8. IEEE, 2012.

Yunhao Tang, Michal Valko, and Rémi Munos. Taylor expansion policy optimization. In International Conference on Machine Learning , pages 9397-9406. PMLR, 2020.

Shiau Hong Lim and Ilyas Malik. Distributional reinforcement learning for risk-sensitive policies. Advances in Neural Information Processing Systems , 35:30977-30989, 2022.

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.

Ralph Neuneier and Oliver Mihatsch. Risk sensitive reinforcement learning. Advances in Neural Information Processing Systems , 11, 1998.

Yun Shen, Michael J Tobia, Tobias Sommer, and Klaus Obermayer. Risk-sensitive reinforcement learning. Neural computation , 26(7):1298-1328, 2014.

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.

Wojciech M Czarnecki, Simon Osindero, Max Jaderberg, Grzegorz Swirszcz, and Razvan Pascanu. Sobolev training for neural networks. Advances in Neural Information Processing Systems , 30, 2017.

Harris Drucker and Yann Le Cun. Improving generalization performance using double backpropagation. IEEE transactions on neural networks , 3(6):991-997, 1992.

A First-order Taylor expansion

A.1 Action expansion update proof

Transposing the scalar inner product, ξ T a ∇ a Q θ ( s , µ a ) , in the first-order Taylor expansion in the main text (Eq. 10) gives,

As ∇ θ is a linear operator, and the action noise, ξ a , does not depend on the policy parameters, θ ,

where ∇ 2 a ,θ Q θ ( s , µ a ) is a matrix of second derivatives. The expectation of ξ a is zero, so the terms linear in ξ a are zero, leading to,

Swapping the order of the terms in the expectation (which is valid, as ( ξ T a ∇ a δ θ ( s , µ a ) ) is a scalar),

We can then move the terms independent of ξ a out of the expectation

Finally, we know E [ ξ a ξ T a ] = Σ a ,

If we assume the action covariance is isotropic, Σ a = λ a I , we get the (1st-order) Taylor TD-update for the action expansion,

A.2 State expansion update proof

Applying the first-order Taylor expansion to Eq. (16), then following the approach from the previous section,

Finally, we know E [ ξ s ξ T s ] = Σ s ,

If we assume the state covariance is isotropic, Σ s = λ s I , we get the (1st-order) Taylor TD-update for the state expansion,

B Proof of stable learning for Taylor TD with linear function approximation (with a fixed policy)

The Taylor TD update for the action expansion can be written as (an equivalent proof can be derived for the state expansion):

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, we consider linear function approximation,

where,

and where x and x ′ are feature-vectors of length N . We can re-write each of the two terms in the Taylor TD update in terms of this linear function approximation. The first term, corresponding to a standard TD-update can be written,

The second term, which is the new term introduced by Taylor-TD methods, can be written,

Here, ∇ a x ∈ R A × N , is a matrix, while ∇ a r ∈ R A is a vector.

Putting the two terms together Eq. 41 & 42 and factorising the terms multiplying θ t , we can write the expected next weight vector as:

where:

Since only A and ˜ A multiplies θ , these two quantities exclusively are important for guaranteeing stable learning. If A is positive definite and ˜ A is positive-semi-definite, then ( A + λ a ˜ A ) is positivedefinite, and for sufficiently small η , the magnitude of the eigenvalues of I -η ( A + λ a ˜ A ) are all smaller than 1 , in which case the system is stable. Crucially, the term A is the same as traditional TD-learning so [11] provides a proof that A is always positive definite [see also 12]. Thus, all we have to prove is that ˜ A is positive semi-definite. In order to prove that ˜ A is positive semi-definite, we require an (very reasonable) assumption, that the timestep ∆ t is small. Specifically, if we take,

Then the ∇ a x ′ terms must be small. Specifically, applying the chain rule, the full derivative is (expressed here for 1d case, ∇ a x ′ = ∂x ′ ∂a , for simplicity),

This is proportional to ∆ t , as

The key test for positive semi-definiteness is that for all vectors b ,

Substituting the definition of ˜ A from Eq. (46),

Putting the b 's inside the expectation,

Now, ( ∇ a xb ) ∈ R 1 × A is a lengthA row-vector, and the terms inside the expectations are inner products of two length A vectors. We can write these vector inner products as sums,

where a i is a particular element of the action-vector, a , so ∇ a i x ∈ R N is a length-N vector, and

Critically, the overall inequality in Eq. (52) holds if a similar inequality holds for every term in the sum in Eq. (53),

As ∇ a i xb and ∇ a i x ′ b are scalars, we can write,

There are now two cases. If 0 = E [( ∇ a i xb ) 2 ] , we must have that 0 = ( ∇ a i xb ) always (except at a set of measure zero). Thus, the second term must also be zero (except at a set of measure zero), i.e. 0 = E [( ∇ a i xb )( ∇ a i x ′ b )] and the inequality holds. Alternatively, if E [( ∇ a i xb ) 2 ] is non-zero, it must be positive, in that case, the second term can also be non-zero, but it scales with ∆ t , so we can always choose a ∆ t , small enough to ensure that the Eq. (56) holds, in which case ˜ A is positive semi-definite.

C Using the chain rule to expand the gradient of Taylor TD

A(learned) differentiable model of the transitions and rewards is needed to compute the gradient terms ∇ a δ θ ( s , a ) and ∇ s δ θ ( s , a ) in Taylor TD. This is because the TD target, δ θ = r ( s , a ) + γQ θ ( s ′ , a ′ ) , comprises the reward and the Q-value at the next time step, both of which depend on s and a . The full gradients for the action expansion can be written:

where ˆ r denotes a (differentiable) reward function, ˆ s denotes the predicted next state, while ∂ ˆ s ′ ∂ a denotes the gradient term computed by differentiating through a differentiable model of the transitions. An analogous gradient can be written for ∇ s δ θ ( 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 δ θ ( s , a ) wrt a and s directly, by re-writing the Taylor TD updates in terms of a simple loss (see Appendix D).

D Computing the updates as the gradient of a loss

Here we report how to easily implement the Taylor TD-update (i.e. Eq. 18) as a loss to be passed to an optimizer (e.g. PyTorch optimizer).

E Model learning

Taylor TD-learning requires a differentiable model of the transitions. We learned this model according to maximum-likelihood based on observed transitions sampled from a reply buffer,

where B denotes a reply buffer from which (observed) transitions can be sampled. In TaTD3, the model of the transitions consists of an ensemble of 8 models. These models return a Gaussian distribution over the next state, with zero correlations, and with mean and variance given by applying a neural network to the previous state-action pair. Conversely, the model of the rewards consists of a neural network trained with mean-square error on the observed rewards sampled from the reply buffer.

F Variance reduction

Here we describe in more details how we computed the variance of Taylor TD and standard TDlearning updates on the 6 continuous control benchmark tasks (i.e. Section 4.1). Based on a policy, a value function and a set of randomly sampled states from a memory buffer, we computed the Taylor and the standard (sample-based) TD batch updates across all sampled states (under the policy). Note, we ensured the exact same (random) states were used to compute the Taylor and the standard (sample-based) TD batch updates. The update for each state takes the form of a gradient evaluation of the value function relative to the TD objective. Next, we computed the variance of these gradient terms across all states and summed them up since the variance of each parameter update contributes to the variance of the value function relative to the TD objective. We repeated this process for 10 different seeds and plotted the mean update variance and standard error for both Taylor and the standard TD updates (i.e. Fig.1).

G Variance reduction across training

In this section we test how the variance of Taylor TD and standard TD-learning updates changes across training. In particular, we want to assess how the variance advantage of Taylor TD updates (over standard TD updates) changes across 3 different points in training: at the start of training (i.e., for an untrained critic and an untrained model of the transitions), early in training (i.e., after 5000 training steps) and finally late in training (i.e., after 50000 training steps).

In Fig. 4, we can see that across all training points, Taylor TD updates provide lower variance updates compared to standard TD updates.

H Toy example

We provide a toy example to investigate the settings in which using a Taylor expansion of an expected update may provide the largest benefits relative to sample-based estimates of the same expected

No Training

EarlyTraining

LateTraining

0.0006

1.4

0.0005

1.2

40

ice

/ariar

1.0

0.0004

30

Meanupdate

>

000

0.8

0.0003

0.6

20

0.0002

0.4

10

0.0001

0.2

0.0000

0.0

0

TD-learning

TD-learning

TD-learning

update. To do so, we train a function approximation parameterised by θ (e.g. a critic) to approximate the expected outputs of an underlying target function (e.g. TD targets). Starting from a single expected target value, we can formulate the following objective (for multiple targets, we can just sum the objectives):

where ξ x is sampled form a Gaussian distribution with E[ ξ x ] = 0 , E [ ξ x ξ T x ] = λ x I . Hence, the task requires the function approximation ˆ y θ () to approximate the expected outputs of the target function y () , based on a set of randomly perturbed inputs (i.e. x + ξ x ). We do this by comparing two different approaches. The first approach involves sampling different outputs of y (based on different input perturbations ξ x ) and training ˆ y θ () to match those outputs. We denote this approach as "Sample-based targets". This approach aims to mimic standard TD-learning, where TD-targets are sampled at each time step to train the critic (i.e. providing sample-based estimates). The second approach aims to mimic Taylor TD, applying a first-order Taylor expansion to the objective in Eq 59,

For clarity we used the notation ˆ y θ = ˆ y θ ( x ) and y = y ( x ) . The terms that are linear in ξ x cancels and after summing up equal terms, we get:

We can then take the gradient of this approximation relative to the function approximation weights to get the weight update for the Taylor approach:

Note, this update is reminiscent of of double backpropagation settings in the supervised learning literature [49, 50]. The toy example allows us to compare the Taylor expansion with the sample-based estimation under different conditions (i.e. dimension size of the data points as well as number of data points). In particular, we compare the two approaches for inputs of different dimensions (from 1 to 100 dimensional inputs) and across two data regimes, a low data regime (i.e. only 15 training samples are used to train the function approximation) and a high data regime (i.e. over 100 samples

Low dataregime

High dataregime

0.20

0.15

error

0.15

0.10

Prediction

0.10

品

0.05

0.05

0.00

0.00

0

20

40

60

80

100

0

20

40

60

80

100

Input dim

Input dim

Taylorexpansion

Sample-based targets

are used to train the function approximation). Performance are then assessed based on a novel set of inputs (i.e. 50), sampled from the same underlying distribution as the training data. Fig. 5 show that the benefits of the Taylor approach over sample-based estimates increase as the dimension of the data points grows in size (e.g. RL tasks involving high dimensional action and state spaces). In particular, these benefits are even larger in the presence of a low data regime, such as RL settings in which for any given state-action pair we can only sample a few transitions from the environment.

I Code

All the code is available at https://github.com/maximerobeyns/taylortd

J Baseline algorithms

Plotted performance of MBPO [18] was directly taken from the official algorithm repository on GitHub at https://github.com/jannerm/mbpo . Plotted performance of SAC [2] was also obtained taken from the official algorithm repository on GitHub at https://github.com/haarnoja/ sac . Plotted performances of both MAGE and Dyna-TD3 [19] were obtained by re-running these algorithms on the benchmark environments, taking the implementations from the official algorithms' repository available at https://github.com/nnaisense/MAGE .

K Ablations

K.1 State expansion ablation

Here, we ask whether the Taylor state expansion brings any benefit to performance, on top of the Taylor action expansion. To do so, we compare the TaTD3 algorithms with and without state expansion on two standard benchmark tasks (i.e. analogous to setting λ s = 0 in the update Eq. 18). Fig. (6) shows that including the state expansion is beneficial to both environments.

K.2 Cosine similarity ablation

Here, we ask whether taking the cosine similarity of state and action gradient terms benefit the performance of TaTD3. To do so, we compare the standard TaTD3 algorithms (trained with the loss in Eq. 18) with a version of TaTD3 trained with a loss without cosine similarity,

Namely, this loss optimizes the dot-product for state and action gradient terms instead of the cosine similarity. We assess this comparison on two standard benchmark tasks (i.e. see Fig. 7). In Fig. 7, we can see the cosine similarity does improve performance on both tasks.

HalfCheetah-v2

12000

10000

5000

8000

4000

ret

6000

3000

Avera

2000

4000

1000

2000

一

0

0

0

50k

100k

150k

0

50k

100k

150k

Steps

Steps

TaTD3

ActionOnly-TaTD3

HalfCheetah-v2

12000

10000

5000

return

8000

4000

Average

6000

3000

2000

4000

一

1000

2000

0

0

0

50k

100k

150k

0

50k

100k

150k

Steps

Steps

TaTD3

NoCosine-TaTD3

L Computing

All experiments were run on a cluster of GPUs, including NVIDIA GeForce RTX 2080, 3090 and NDVIDIA A100. Here, we report the difference in computing time between standard TD-learning, Taylor TD and the two strongest baselines (i.e., MAGE and MBPO) based on the same GPU for each comparison. We do so for a low dimensional (Pendulum), a 'medium' dimensional (Walker2d) and a high dimensional (Humanoid) environment to span a broad range of settings. For MBPO, we report the average running time across all training epochs using the horizon lengthening schedule outlined in the original paper [18].

MAGE has about the same runtime as our method (TaTD3), while MBPO is a lot slower than our method.

M Hyperparameters settings

Below, we reported the hyperparameter settings for TaTD3 (and sample-based Expected-TD3),

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 λ a and λ s to each environment (i.e., λ a = 0 . 06 and λ s = 0 . 00005 ). These parameters were founded by running a grid search over potential values of these parameters based on a single environment (i.e., Pendulum), then using the best values for all other environments. Nevertheless, we reached top performance in HalfCheetah, Walker2d and Humanoid by using a larger λ a (i.e., λ a = 0 . 25 ). This finding suggests these 3 environments benefit from learning a broader distribution of Q-values over the actions, we believe for better exploration. Additionally, you may notice the much larger values for λ a compared to λ s . However, we should stress that any direct comparison between the state covariance λ s and the action covariance λ a may not be very meaningful. This is because the distribution of initial states and the distribution of initial actions may be very different, making any direct comparison between λ s and λ a hard to assess.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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\nWEAKNESSES:\nThere 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.\n\nQUESTIONS:\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- 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\nLIMITATIONS:\nThe limitations are discussed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe authors mentioned about the limitations in the last section and I do not forsee any negative societal impacts.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nOriginality\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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. 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\n\nLIMITATIONS:\nYes the authors have included the limitation section. The negative societal impact section doesn't apply to this line of research.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nMinor: \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).\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nNo concerns regarding potential negative societal impact, and the limitations were discussed to a reasonable amount.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-30"}