For SALE: State-Action Representation Learning for Deep Reinforcement Learning
Authors: Scott Fujimoto
Wei-Di Chang
Edward J. Smith
Shixiang Shane Gu
Doina Precup
David Meger
Venue: NeurIPS 2023 poster (2023)
Keywords: Deep reinforcement learning
representation learning
TL;DR: We add state-action representation learning to TD3 and significantly outperform existing methods.

Abstract: In reinforcement learning (RL), representation learning is a proven tool for complex image-based tasks, 
but is often overlooked for environments with low-level states, such as physical control problems. This paper introduces SALE, a novel approach for learning embeddings that model the nuanced interaction between state and action, enabling effective representation learning from low-level states. We extensively study the design space of these embeddings and highlight important design considerations. We integrate SALE and an adaptation of checkpoints for RL into TD3 to form the TD7 algorithm, which significantly outperforms existing continuous control algorithms. On OpenAI gym benchmark tasks, TD7 has an average performance gain of 276.7% and 50.7% over TD3 at 300k and 5M time steps, respectively, and works in both the online and offline settings.

For SALE: State-Action Representation Learning for Deep Reinforcement Learning
Scott Fujimoto Mila, McGill University
Wei-Di Chang McGill University
Edward Smith McGill University
Shixiang Shane Gu Google DeepMind
Doina Precup Mila, McGill University
David Meger Mila, McGill University
Abstract
field of reinforcement learning (RL), representation learning is a proven tool for complex image-based tasks, but is often overlooked for environments with low- level states, such as physical control problems. This paper introduces SALE, novel approach for learning embeddings that model the nuanced interaction between state and action, enabling effective representation learning from low-level states. We extensively study the design space of these embeddings and highlight important design considerations. We integrate SALE and an adaptation of checkpoints for RL into TD3 form the TD7 algorithm, which significantly outperforms existing continuous control algorithms. On OpenAl gym benchmark tasks, TD7 has an average performance gain of 276.7% and 50.7% over TD3 at 300k and 5M time steps, respectively, and works in both the online and offline settings.
Introduction
Reinforcement learning (RL) is notoriously sample inefficient, particularly when compared to more straightforward paradigms in machine learning, such as supervised learning. One possible expla- nation is the usage of the Bellman equation in most off- off-policy RL algorithms [Mnih et al., 2015, Lillicrap et al., 2015], which provides a weak learning signal due to an approximate and non-stationary learning target [Fujimoto et al., 2022].
near-universal solution to sample inefficiency deep learning is representation learning, whereby intermediate features are learned to capture the underlying structure and patterns of the data. These features can be found independently from the downstream task and considerations such the learning horizon and dynamic programming. While feature learning of this type has found some success in RL setting, it has been mainly limited to vision-based environments [Jaderberg et al., 2017, Oord et al., 2018, Anand et al., 2019, Laskin et al., 2020, Stooke et al., 2021, Yarats et al., 2022].
On the other hand, the application representation learning low-level states is much less common. At first glance, it may seem unnecessary to learn representation over an already-compen state vector. However, we argue that the difficulty of task is often defined by the complexity of the underlying dynamical system, rather than the size of the observation space. This means that regardless of the original observation space, there exists an opportunity to learn meaningful features by capturing the interaction between state and action.
SALE. In this paper, we devise state-action learned embeddings (SALE), method that learns embeddings jointly over both state and action by modeling the dynamics of the environment in latent space. Extending prior work [Ota et al., 2020], we introduce three important design considerations
Corresponding author: scott. fujimoto@mail mcgiill.c
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

when learning a state-action representation online. Most importantly, we observe the surprising effect of extrapolation error [Fujimoto al.. 2019] when significantly expanding the action-dependent imput and introduce simple clipping technique to mitigate it.
Design study. Learning to model environment dynamics in latent space is common approach for feature learning which has been widely considered [Watter et al., 2015, Ha and Schmidhuber, 2018, Hafner al., 2019, Gelada et al., 2019, Schwarzer et al., 2020], with many possible variations in design. Consequently, the optimal design decision is often unclear without considering empirical performance. To this end, we perform an extensive empirical evaluation over the design space, with the aim of discovering which choices are the most significant contributors to final performance.
Checkpoints. Next, explore the usage of checkpoints in RL. Similar to representation learning, early stopping and checkpoints are standard techniques used to enhance the performance of deep learning models. A similar effect can be achieved in RI. by fixing each policy for multiple training episodes, and then at test time, using the highest-performing policy observed during training.
TD7. We combine TD3 with our state-action representation learning method SALE, the aforemen- tioned checkpoints, prioritized experience replay [Fujimoto et al., 2020], and a behavior cloning term (used only for offline RL) [Fujimoto and Gu, 2021] to form the TD7 (TD3+4 additions) algorithm. We benchmark the TD7 algorithm in both the online and offline RL setting. TD7 significantly outperforms existing methods without the additional complexity from competing methods such as large ensembles, additional updates per time step, or per-environment hyperparameters. Our key improvement, SALE, works in tandem with most RL methods and can be used to enhance existing approaches in both the online and offline setting. Our code is open-sourced!
2 Related Work
Representation learning. Representation learning has several related interpretations in RL. His- torically, representation learning referred to abstraction, mapping an MDP to a smaller one via bisimulation or other means [Li et al., 2006, Ferns et al., 2011, Zhang et al., 2020]. For higher- dimensional spaces, the notion of true abstraction has been replaced with compression, where the intent is embed the observation space (such as images) into smaller manageable latent vector [Wat- ter et al., 2015, Finn et al., 2016, Gelada et al., 2019]. Representation learning can also refer to feature learning, where the objective is to learn features that capture relevant aspects of the environment or task, via auxiliary rewards or alternate training signals [Sutton et al., 2011 1, Jaderberg et al., 2017, Riedmiller et al., 2018, Lin et al., 2019]. In recent years, representation learning in RL often refers to both compression and feature learning, and is commonly employed in image-based tasks [Kostrikov et al., 2020, Yarats et al., 2021, Liu et al., 2021, Cetin et al., 2022] where the observation space is characterized by its high dimensionality and the presence of redundant information.
Representation learning by predicting future states draws inspiration from rich history [Dayan, 1993, Lituman and Sutton, 2001], J, spanning many approaches in both model-free RL [Munk et al., 2016, Van Hoof et al., 2016, Zhang et al., 2018, Gelada et al., 2019, Schwarzer et al., 2020, Fujimoto et al., 2021, Ota et al., 2020, 2021] and model-based RL in latent space [ [Watter et al., 2015, Finn et al., 2016, Karl et al., 2017, Ha and Schmidhuber, 2018, 8, Hansen et al., 2022, Hafner et al., 2019, 2023]. Another related approach is representation learning over actions [Tennenholtz and Mannor, 2019, Chandak et al., 2019, Whitney et al., 2020]. Our key distinction from many previous approaches is the emphasis on learning joint representations of both state and action.
Methods which do learn state-action representations, by auxiliary rewards to the value function [Liu et al., 2021], or MDP homomorphisms! [Ravindran, 2004, van der Pol et al., 2020a,b, Shossarrr et al., 2022] emphasize abstraction more than feature learning. Our approach can be viewed as an extension of OFENet [Ota et al., 2020]. which also learns a state-action embedding. We build off of OFENet and other representation learning methods by highlighting crucial design considerations and addressing the difficulties that arise when using decoupled state-action embeddings. Our resulting improvements are reflected by significant performance gains in benchmark tasks.
Stability in RL. Stabilizing deep RL algorithms has been longstanding challenge, indicated by numerous empirical studies that highlight practical concerns associated with deep RL methods [Hen- derson et al., 2017, Engstrom et al., 2019]. Our use of checkpoints is most closely related to stabilizing

---

policy performance via safe policy improvement [Schulman et al., 2015, 2017, Laroche et al., 2019], as well as the combination of evolutionary algorithms [Salimans et al., 2017, Mania al.. 2018] with RI. [Khadka and Tumer, 2018, Pourchot and Sigaud, 2018), where the checkpoint resembles the fittest individual and the mutation is defined exclusively by the underlying RI. algorithm.
3 Background
In Reinforcement learning (RL) problems are framed as a Markov decision process (MDP). An MDP a 5-tuple (S, A, R. p. ?) with state space S, action space A, reward functiin R. dynamics model p. and discount factor where the objective is to find a policy æ S - A, a mapping from state 8
to action E A, which maximizes the return Σ Ii=1f=1rr r the discounted sum of rewards obtained when following the policy. RL algorithms commonly use value function Q' a) == E [_x= >r-1re|s0 = 8, ag = =a. which models the expected return, starting from an initial state and action a.
4 State-Action Representation Learning
In this section, we introduce state-action learned embeddings (SALE) (Figure 1). We begin with the basic outline of SALE and then discuss three important considerations in how SALE is implemented. We then perform an extensive empirical evaluation on the design space to highlight the critical choices when learning embeddings from the dynamics of the environment.
4.1 State-Action Learned Embeddings
The objective of SALE is to discover learned embeddings (2*a, 2²)) which capture relevant structure in the observation space, as well as the transition dynamics of the environment. To do so, SALE utilizes pair of encoders (f.g) where f(s) encodes the state into the state embedding =" and g(z", a) jointly encodes both state s and action into the state-action embedding ""
(1) The embeddings split into state and state-action components that the encoders be trained with dynamics prediction loss that solely relies on the next state 8', independent of the next action current policy. As a result, the encoders are jointly trained using the mean squared error (MSE)
(s))
g(=", a).
between the state-action embedding * and the embedding of the next state z²
C(S,g) = (g(S(s). -1(ss))x) -
where |-|x denotes the stop-gradient operation. The embeddings are designed to model the underlying structure of the environment. However, they may encompass all relevant information needed by the value function and policy, such as features related to the reward, current policy, or task horizon. Accordingly, we concatenate the embeddings with the original state and action, allowing the value and policy networks to learn relevant internal representations for their respective tasks:
(s.a) + 0(="0,"** ,",,),))
=(s) - 8))
The encoders (f)g) trained online and concurrently with the RL agent (updated the same frequency as the value function and policy), but are decoupled (gradients from the value function and policy are not propagated (f. g)). Although the embeddings are learned by considering the dynamics of the environment, their purpose is solely to improve the input to the value function and policy, and
serve world model for planning or estimating rollouts.
There are three additional considerations in how SALE is implemented in practice.
Normalized embeddings. The minimization of distances in embedding space can result in instability due to either monotonic growth or collapse to redundant representation [Gelada et al., 2019]. To combat this risk, we introduce AvgLINomm a normalization layer that divides the input vector by its average absolute value each dimension, thus keeping the relative scale of the embedding constant throughout learning. Let z; be the i-th dimension of an N-dimensional vector x, then
AvgLINorm (x)
(4)
Σ; x
N
3

---

Decoupled
Decoupled
fc())
gt [2i
(()
AvgL (Norm(N)))))))
2/
A (=))
Linear(s, a)
Qt+il2i sarppng
Linear ()
A (
AvgLi SgLLNNor(e')
End-to-end
End-to-end (b) Policy x with SALE
(a) Value function Q with SALE
Figure 1: Diagram of State-Action Learned Embeddings (SALE). SALE uses encoders (f.g) to output embeddings (=" to enhance the input of the value function Q and policy corresponding linear layer. The encoders (S.g) are jointly trained to predict the state embedding (where ix denotes the stop-gradient operation), decoupled from training of the value function and policy (Equation 2). The end-to-end linear layers are trained with gradients from the corresponding network. AvgL is used to keep the scale of cach of the inputs to the value function and policy constant.
denotes the output of the
AvgL is applied to the state embedding 2 Similar the normalized loss functions used by SPR [Schwarzer et al., 2020] and BYO [Grill et al., 2020], AvgLINorm protects from monotonic growth, but also keeps the scale of the downstream input constant without relying updating statistics (e.g. BatchNorm [loffe and Szegedy, 2015]). This is important for our approach as the embeddings trained independently from the value function and policy. AvgL. Norm not applied the state-action embedding =0
trained match the normalized next state embedding z"'.
We also apply AvgL to the state and action inputs (following linear layer) to the value function Q and policy to keep them at a similar scale to the learned embeddings. The input to the value function and policy then becomes:
(x=a0 2,, WogLLNNommLLneet a))),
T (== AvgL [Nomm(inea(())))))
(5)
Unlike the embeddings (a*, ".)), these linear layers are learned andtto-end and can consequently be viewed addition to the architecture of the value function policy.
Fixed embeddings. Since an inconsistent input can cause instability, we freeze the embeddings used to train the current value and policy networks. This means at the iteration t + 1, the input to the current networks (Q++1, Te+1) uses embeddings (z/* ,2") from the encoders (fe- ) at the previous iteration t. The value function and policy are thus updated by:
Qut1 s,
= +r+ (1,1,,1,27,8s0)
where ~ =e(2i
(6)
(e41 8) = argmax Q (2 s,a). where s))
(7)
The current value function Q++1 is also trained with respect the previous value function Qs, known as a target network [Mnih et al., 2015]. The current embeddings zi+1 and 2 are trained with Equation using a target 2 (hence, without target network). Every steps the iteration is incremented and all target networks are updated simultancously:
Q - Qe+1,
(f-1119-11) + (fe,9)),
(fe. + (f4t1)9411.
(8)
Clipped Values. Extrapolation error is the tendency for deep value functions to extrapolate to unrealistic values state-actions pairs which are rarely seen in the dataset [Fujimoto et al., 2019]. Extrapolation error has a significant impact in offline RL, where the RL agent learns from a given dataset rather than collecting its own experience, the lack of feedback overestimated values can result in divergence.
Surprisingly, we observe a similar phenomenon in online RI, when increasing the number of dimensions in the state-action input the value function, as illustrated in Figure Our hypothesis is that the state action embedding =0 expands the action input and makes the value function more

---

Un-clipped
Un-elipped
performance by seed
value estimate by seed
No 2
n and small d
Total Reward (Ik)
12
Value Estimate (Ik)
9
Small de
Clipped
3
2
0
1
2 Time steps (1M)
Time steps (IM)
Time steps (IM)
Time steps (IM)
Seeds (1-10) Buffer: 1M
Buffer: 1M
100k
Buffer:
IM
100k
Buffer:
1M
100k
Buffer:
IM
100k
Figure 2: Extrapolation error can occur in online RL when using state-action representation learning. All figures use the Ant environment. 6t corresponds to the output of the linear layer (Lincar(s, a) = 8n0 (Equation 5). Both embeddings and 6ḥh have a default dimension size 256. Small @ means that Dim " set to 16, No 2= means the value function input is Q(z", s, a). estimate of 10 individual seeds without value clipping. While the performance trends upwards there are large Varying the input dimension
The default performance and value
dips in reward, which correspond with jumps in the estimated value. can improve or harm stability of the value estimate. The severity is impacted by the replay buffer size (1M 100k). The state embedding =" is left unchanged all settings, showing that the state-action embedding 2 * and the linear layer over the state-action input 6 are the primary contributors to the extrapolation error. This shows the potential negative impact from increasing the dimension size of an input which relies potentially unseen action. Clipping stabilizes the value estimate, without modifying the input dimension size (Equation 9).
/
likely to over-extrapolate unknown actions. We show in Figure 2 that the dimension size of =a as well as the state-action input plays an important role in the stability of value estimates.
Fortunately, extrapolation error can be combated in a straightforward manner in online RL, where poor estimates are corrected by feedback from interacting with the environment. Consequently, we only need to stabilize the value estimate until the correction occurs. This can be achieved in SALE by tracking the range of values in the dataset D (estimated over sampled mini-batches during training), and then bounding the target used in Equation 6 by the range:
a)+ = r r + clip (Qi(s) d'), min Qe(s,a), max Pc(s.a))
(9)
(s,a)ED
(e.0)ED
Additional discussion of extrapolation error, experimental details, and ablation of the proposed value clipping in SALE can be found in Appendix E & D.
4.2 Evaluating Design Choices
The effectiveness of learning embeddings by modeling the dynamics of the environment a natural consequence of the relationship between the value function and future states. However, there are many design considerations for which all alternatives are potentially valid and the approach adopted differs among related methods in the literature. In this section, we perform extensive study over the design space to (1) show SALE uses the correct and highest performing set of choices, and (2) better understand which choices are the biggest contributors to performance when using SALE.
In Figure we display the mean percent loss when modifying SALE in the TD7 algorithm (to be fully introduced in Section 5.2). The percent loss is determined from the average performance at 1M time steps, over 10 seeds and five benchmark environments (HalfCheetah, Hopper, Walker2d, Ant, Humanoid) [Brockman et al., 2016]. A more detailed description of each variation and complete learning curves can be found in Appendix D.
Learning target. TD7 trains the encoders by minimizing the MSE between the state-action embed- ding =0 and a learning target of the next state embedding 2" (Equation 2). We test several alternate learning targets. OFENet uses the next state s' as the target [Ota al., 2020] while SPR [Schwarzer et al., 2020] uses the embedding zülgge from a target network obtained with an exponential moving average with weight 0.01. Drawing inspiration from Bisimulation metrics [Ferns et al., 2011]. I]. Deep- MDP [Gelada et al., 2019] use an objective that considers both the next state embedding 2² and the
5

---

Learning Target Detault: s*
Network Input
Normalization
End-to-End
Default: Qis st,
and x(s"
Default: A
Default: Decoupled
Qie aa))
Q ssa)
Q(
Qlenn
Q(aa))
No final
only
Naana ait
BatdkNoom
LayerNoom
Crsine bosa
x(a)
Nogst
Mean Percent Loss
0.9
2.6
(4.0)
10
15
(2.1)
(3.6)
8.0
(4.3)
10.1 (6.5)
13.8
(6.5)
5.7
6.2
(2.8)
(7.9)
9.5
(5.1)
(4.7)
4.7
12.4
(5:4)
(4.9)
10.9
21
(4.6)
23.6
22.6
24.8
(3.1)
(58)
(10.8)
(7.6)
(7.3)
(3.7)
Figure 3: The mean percent loss from using alternate design choices in TD7 at 1M time steps, over 10 seeds and the five benchmark MuJoCo environments. Bracketed values describe the range of the 95% confidence interval around the mean. Percent loss is computed against TD7 where the default choices correspond to percent loss of See Section for description of cach design choice and key observations. See the Appendix for further implementation-Ier level details.
reward r. We test including prediction loss on the reward by having the encoder output both 2e0 and prod where ppred is trained with the MSE to the reward r. Finally, we test the next state-action embedding zz the target, where the action a' is sampled from the target policy.
All learning targets based the next state s' perform similarly, although using the embed- ding z further improves the performance. On the contrary, the next state-action embedding za performs much worse as a target, highlighting that signal based on the non-stationary policy can harm learning. Including the reward signal has little impact performance.
Network input. In our approach, the learned embeddings (2 2n) are appended to the state and action input the value function Q(=00 2", 8, a) and policy (2>",s) (Equation 3). We attempt different combinations of input to both networks. We also evaluate replacing the fixed embed- dings (Equations 7), with the non-static current embeddings (zff1:
The added features have greater impact the value function than the policy, but are beneficial for both networks. All components of the value function input (2*0, ,²,,,, a). are necessary to achieve the highest performance. While the state-action embedding = is useful representation for value learning, it is only trained to predict the next state and may overlook other relevant aspects of the original state-action input (s, a). Solely using the state-action embedding =a input leads to poor results, but combining it with the original input (s, a) significantly improves performance.
Normalization. TD7 uses AvgL Norm (Equation 4) to normalize the scale of the state embedding z", as well on the state-action input (s, a). following a linear layer (Equation 5). We attempt removing AvgLINorm on (s, a), removing it entirely, and adding to the state-action embedding =* We additionally test swapping AvgL for BatchNorm [loffe and Szegedy, 2015] and LayerNorm [Ba et al., 2016]. Finally, instead of directly applying normalization to the embeddings, we replace the MSE in the encoder loss (Equation by the cosine loss from Schwarzer et al. [2020].
The usage of AvgLINorm is beneficial and related alternate approaches do not achieve the same performance.
End-to-end. Embeddings be trained independently end-to-end with the downstream task. We test our approach as an auxiliary loss to the value function. The encoders and the value function are trained end-to-end, thus allowing the value loss to affect the embeddings (2*2 z²), where the encoder loss (Equation 2) is multiplied by a constant to weigh its importance versus the value loss.
Learning the embeddings end-to-end with the value function performs signficantly worse than decoupled representation learning.

---

Stabilizing RL with Decoupled Representation Learning
In this section, we present the TD7 algorithm (TD3+4 additions). We begin by introducing the use of checkpoints in RL to improve the stability of RL agents. We then combine SALE with checkpoints and various previous algorithmic modifications to TD3 [Fujimoto et al., 2018] to create single RL algorithm for both the online and offline setting.
5.1 Policy Checkpoints
Deep RL algorithms are notoriously unstable [Henderson et al., 2017]. The unreliable nature of deep RL algorithms suggest need for stabilizing techniques. While we can often directly address the source of instability, some amount of instability is inherent to the combination of function approximation and RL. In this section, we propose the of checkpoints, to preserve evaluation performance, irrespective of the quality of the current learned policy.
A checkpoint is snapshot of the parameters of a model, captured specific time during training. In supervised learning, checkpoints are often used to recall a previous set of serepperforiing parameters based validation error, and maintain consistent performance across evaluations [Vaswani et al., 2017, Kenton and Toutanova, 2019]. Yet this technique surprisingly absent from the deep RL toolkit for stabilizing policy performance.
In RL, using the checkpoint of a policy that obtained a high reward during training, instead of the current policy, could improve the stability of the performance at test time.
For off-policy deep RL algorithms, the standard training paradigm is train after each time step (typically at one-to-one ratio: one gradient step for one data point). However, this means that the policy changes throughout each episode, making hard to evaluate the performance. Similar many on-policy algorithms [Williams, 1992, Schulman et al., 2017], we propose to keep the policy fixed for several assessment episodes, then batch the training that would have occurred.
Standard off-policy RL: Collect data point - train once.
Proposed: Collect N data points over several assessment episodes train N times.
In a similar manner to evolutionary approaches [Salimans et al., 2017], we can use these assessment episodes to judge if the current policy outperforms the previous best policy and checkpoint accordingly. At evaluation time, the checkpoint policy is used, rather than the current policy.
We make two additional modifications to this basic strategy.
Minimum mean. Setting aside practical considerations, the optimal approach would be to evaluate the average performance of each policy using as many trials as possible. However, to preserve learning speed and sample efficiency, is only sensible to use handful of trials. As such, to penalize unstable policies using a finite number of assessment episodes, we use the minimum performance, rather than the mean performance. This approach also means that extra assessment episodes do not need to be wasted on poorly performing policies, since training can resume early if the performance of any episode falls below the checkpoint performance.
Variable assessment length. In Appendix F, we examine the caliber of policies trained with a varied number of assessment episodes and observe that a surprisingly high number of episodes (20+) can be used without compromising the performance of the final policy. However, the use of many assessment episodes negatively impacts the early performance of the agent. Freezing training for many episodes means that the environment is explored by a stale policy, reducing data diversity, and delaying feedback from policy updates. To counteract this effect, restrict the number of assessment episodes used during the initial phase of training before increasing it.
Additional details of our approach to policy checkpoints can be found in Appendix F.
5.2 TD7
TD7 is based on TD3 [Fujimoto et al., 2018] with LAP [Fujimoto et al., 2020], behavior cloning term for offline RL [Fujimoto and Gu, 2021], SALE (Section 4.1), and policy checkpoints (Section 5.
LAP. Gathered experience is stored in replay buffer [Lin, 1992] and sampled according to LAP [Fujimoto et al., 2020], prioritized replay buffer D [Schaul al., 2016] where transi-
7

---

Algorithm Online TD7
I:
Initialize:
Before training
Policy R44-1, value function Q++ encoders
Target policy target value function Qt. fixed encoders (ʃt, 91), target fixed encoders 9t-))
Checkpoint policy æa. checkpoint encoder fe.
2: for spde to final_epi sode do
b Data collection
3:
Using current policy m--1. collect transitions and store in the LAP replay buffer.
4:
if checkpoint_.onditiie then if actor # outper checkpoint policy T. then
Checkpointing
Update checkpoint networks ne <- Te+1, Sc + fc-
for to Limesteps__st sicrssssseee___ieeeeeee stt_s_iiiiittnnnnnnggnggge training do
D> Training
8:
Sample transitions from LAP replay buffer (Equation 10).
9:
Train encoder (Equation 2), value function (Equations & 9), and policy (Equation 11). if parget_update_.. _ww.uuteefeeeueeeeeeeeeeeeeeee frequuncyyyyyy steps have passed then Update target networks (Equation 8).
10:
11:
Detailed hyperparameter explanations found in the Appendix.
tion tuple i = (s, a, r, sampled with probability max (16(1)1911
p(1) EJED max (16(5)),1)
where 5(1) = Q(, a) y.
(10)
where y the learning target. The amount of prioritization used controlled by a hyperparameter a. Furthermore, the value function loss uses the Huber loss [Huber et al., 1964], rather than the MSE.
Offline RL. To make TD> amenable the offline RL setting, add behavior cloning loss the policy update [Silver et al., 2014], inspired by TD3+BC [Fujimoto and Gu, 2021]:
argmax (e.) [Q(s, (s))) A|E (Q(s, (s))]lxx (T(() - a)²
(11)
The same loss function is used for both offline and online RL, where À = 0 for the online setting. | denotes the stop-grade graddeta operation. Unlike TD3+BC, we not normalize the state vectors. Checkpoints are not used in the offline setting, there is no interaction with the environment.
Both the value function and policy the SALE embeddings input, which omit from the equations above for simplicity, Pseudocode for TD7 is described in Algorithm 1.
Results
In this section, we evaluate the performance of TD7 in both the online and offline regimes. detailed description of the experimental setup, baselines, and hyperparameters can be found in the Appendix, along with additional learning curves and ablation studies.
Online. Using OpenAl gym [Brockman et al., 2016), we benchmark TD7 against TD3 [Fujimoto et al., 2018], SAC [Haarnoja al., 2018], TQC [Kuznetsov et al., 2020], and TD3+OFE [Ota et al., 2020] on the MuloCo environments [Todorov et al., 2012]. SAC and TD3+OFE results are from re-implementations based author descriptions [Haarnoja et al., 2018, Ota al., 2020]. TD3 and TQC results use author-provides code [Fujimoto et al., 201 8, Kuznetsov et al., 2020], with a consistent evaluation protocol for all methods. Learning curves are displayed in Figure and final and intermediate results are listed in Table 1.
Although TQC and TD3+OFE use per-environment hyperparameters along with larger and more computationally expensive architectures, TD7 outperforms these baselines significantly in terms of both early (300k time steps) and final performance (5M time steps). At 300k time steps, TD7 often surpasses the performance of TD3 at 5M time steps, highlighting the considerable performance gains.
Offline. We benchmark TD7 against CQL [Kumar et al., 2020b], TD3+BC [Fujimoto and Gu, 2021], IQL [Kostrikov al., 2021] and X-QL. [Garg et al., 2023] using the MuJoCo datasets in D4RL [Todorov et al., 2012, Fu et al., 2021]. While there are methods that use per-dataset

---

HalfCheetah
Hopper
Walker2d
Total Reward (1k)
12
8
0
2
4
0
2
3
5
Ant
Humanoid
10
TD7
TD3 TQC
SAC
TD3+OFE
Total Reward
Figure Learning
the MuJoCo benchmark. Results are averaged over seeds. The shaded area captures a 95% confi- dence interval around the average performance.
3
0
4 Time steps (IM)
Time steps (1M)
Table 1: Average performance on the MuloCo benchmark at 300k, 1M. and SM time steps. over 10 trials, where + captures 95% confidence interval. The highest performance highlighted. Any performance which is not statistically significantly worse than the highest performance (according to a Welch's 's (-test with significance level 0.05) is highlighted.
Environment
Time step
TD3
SAC
TQC
TD3+FFF
TD7
300k
7715
633
8052 515
7006 + 891
11294 + 247
15031
401
HallCheetah
1M
10574 + 897
10484 + 659
12349 + 878
13758 + 544
17434
155
5M
14337 + 1491
15526 + 697
17459 + 258
16596 + 164
18165
255
300k
1289 1. 768
2370 1 626
3251 1: 461
1581 1: 682
2948
464
Hopper
IM
3226
315
2785 634
3526 1. 244
3121 1. 506
3512
315
5M
3682
1.
3167
3462 1. 818
3423 584
4075
300k
386 3946
1989 500
2812 1. 838
4018 570
5379
Walker2d
IM
292
4314 1 256
5321 1 322
5195 1: 512
6097
570
5M
5078
343
5681 1: 329
6137 1: 1194
6379 1: 332
7397 454
300k
1704 + 655
1478 + 354
1830 + 572
6348 + 441
6171 = 831
Ant
1M
3942 + 1030
3681 + 506
3582 + 1093
7398 + 118
8509 422
5M
758 5589
4615 2022
6329 + 1510
8547
10133 966
300k
1344 365
1997 483
3117 + 910
3181 771
5332 714
Humanoid
IM
5165
145
4909 + 364
6029 + 531
6032 - 334
7429 153
5M
5433 + 245
6555 + 279
8361 + 1364
8951 + 246
10281 + 588
hyperparameters to attain higher total results, we omit these methods because it makes it difficult to directly compare. Baseline results are obtained by re-running author-provided code with a single set of hyperparameters and consistent evaluation protocol. Final performance is reported in Table 2.
TD7 outperforms all baselines. Since TD7 and TD3+BC employ the same approach offline RL. the significant performance gap highlights the effectiveness of SALE in the offline setting.
Ablation study. In Figure 5 report the results of an ablation study (SALE, checkpoints, LAP). The interaction between components is explored further in Appendix G.
the components of TD7
Run time. To understand the computational cost of using SALE and the TD7 algorithm, benchmark the run time of each of the online baselines with identical computational resources and deep learning framework. The results are reported in Figure 6.

---

Table 2: Average final performance the D4RI henchmark after training for 1M time steps. over 10 trials, where + captures 95% confidence interval. The highest performance is highlighted. Any performance which not statistically significantly worse than the highest performance (according to a Welch's t-test with significance level 0.05) is highlighted.
Environment Dataset
CQL
TD3+BC
IQL
X'QQ
TD7
Medium
46.7 + 0.3
48.1 + 0.1
47.4 + 0.2
47.4
1:4:0.1 4
58.0 = 0.4
HalfCheetah
Bediuum-Replan
45.5 +: 0.3
44.6 + 0.4
43.9 + 1.3
44.2 0.7
53.8 = 0.8
Medium-ENie Expeert
76.8 7.4
93.7 0.9
89.6 3.5
90.2
104.6 = 1.6
Medium
59.3 3.3
59.1 3.0
63.9 4.9
67.7 3.6
76.1 + 5.1
Hopper
Medium-Rephe Replay
78.8 + 10.9
52.0 + 10.6
93.4 7.8
82.0 14.9
91.1 + 8.0
Medium-Exper
79.9 + 19.8
98.1 + 10.7
64.2 32.0
92.0 - 10.0
108.2 + 4.8
Medium
81.4 + 1.7
84.3 + 0.8
84.2 1.6
79.2 + 4.0
91.1 + 7.8
Walker2d
Medium-Replie
79.9 + 3.6
81.0 + 3.4
71.2 + 8.3
61.8 + 7.7
89.7 + 4.7
Medium- Expert
108.5 + 1.2
110.5 + 0.4
108.9 + 1.4
110.3 + 0.2
111.8 + 0.6
Total
656.7 1: 24.3
671.3 1: 15.7
666.7 1 34.6
674.9 1 20.4
784.4 L 14.1
7 Conclusion
Ablation
Normalized Porformance
1.0
Representation learning has been typically reserved for image-based tasks, where the observations are large and unstructured. However, by learning embed- dings which consider the interaction between state and action, we make representation learning more broadly applicable to low-level states. We introduce SALE, method for learning state-action embed- dings by considering a latent space dynamics model. Through extensive empirical evaluation, inves- tigate various design choices SALE.
0.8
0.6
0.4
0.2
2
Time steps (1M)
We highlight the risk of extrapolation error [Fujimoto al., 2019] due the increase in input dimensions from using state-action embeddings, but show this instability can be corrected by clipping the target with seen values. We further introduce stability by including policy checkpoints.
TD7
TD3
No SALE
No checkpoints
No LAP
Figure 5: Ablation study over the components of TD7. The y-axis corresponds to the average per- formance over all five MuloCo tasks, normalized with respect the performance TD7 at 5M time steps. The shaded area captures 95% confidence interval. The impact can be ranked (1) SALE, (2) LAP, (3) policy checkpoints.
While both SALE and policy checkpoints are general- purpose techniques that can be included with most RL methods, we combine them with TD3 and sev- eral other recent improvements [Fujimoto al., 2020, Fujimoto and Gu, 2021] to introduce the TD7 algo- rithm. We find our TD7 algorithm is able match the performance of expensive offline algorithms and significantly outperform the state-of-the-art continu- ous control methods in both final performance and early learning.
Run Time
3h 50m
14m
3
Hours
2
1h 50m
1h 28m
47m
Acknowledgments
TD3
SAC
TQC
TD3+FF
and Disclosure of Funding
Figure Run time of each method for 1M time steps on the HalfCheetah environment, using same hardware and deep learning framework (Py- Torch [Paszke et al., 2019]).
This research was enabled in part by support provided by Calcul Québec and the Digital Research Alliance of Canada
10

---

References
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. Advances in neural information processing systems, 32, 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv:1607.064500, 2016
arXiv preprint
Greg Brockman, Vicki Cheung. Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym, 2016.
Edoardo Cetin, Philip J Ball, Stephen Roberts, and Oya Celiktutan. Stabilizing off off-policy deep reinforcement learning from pixels. In International Conference on Machine Learning, pages 2784-2810. PMLR, 2022.
Yash Chandak, Georgios Theocharous, James Kostas, Scott Jordan, and Philip Thomas. Learning action representations for reinforcement learning. In International conference on machine learning, pages 941-950. PMLR, 2019.
Amé Clevert Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv: 2015.
Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Computation, 5(4):613 6244, 1993.
Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. Implementation matters in deep rl: A case study on ppo and trpo. In International Conference Learning Representations, 2019.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metries for continuous markov decision processes. SIAM Journal on Computing, 1066):16221114, 2011.
Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pages 512-519. IEFE, 2016.
Justin Fu, Mohammad Norouzi, Ofir Nachum, George Tucker, Ziyu Wang, Alexander Novikov, Mengjiao Yang, Michael R Zhang, Yutian Chen, Aviral Kumar, Cosmin Paduraru, Sergey Levine, and Thomas Paine. Benchmarks for deep off- off-policy evaluation. In International Conference Learning Representations, 2021.
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. In Thirty-Fifih Conference on Neural Information Processing Systems, 2021.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International Conference on Machine Learning, volume 80, pages 1587-1596. PMLR, 2018.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, pages 2052-2062, 2019.
Scott Fujimoto. David Meger, and Doina Precup. An equivalence between loss functions and non- uniform sampling in experience replay. Advances in Neural Information Processing Systems, 33, 2020.
Scott Fujimoto. David Meger, and Doina Precup. A deep reinforcement learning approach to marginalized importance sampling with the successor representation. In Proceedings of the 38th International Conference on Machine Learning, volume 139, pages 3518-3529. PMLR, 2021.
Scott Fujimoto, David Meger, Doina Precup, Ofir Nachum, and Shixiang Shane Gu. Why should trust you, bellman? The Bellman error a poor replacement for value error. In International Conference Machine Learning, volume 162, pages 6918-6943 PMLR, 2022.

---

Divyansh Garg, Joey Hejna, Matthieu Geist, and Stefano Ermon. Extreme q-learning: Maxent RL without entropy. In The Eleventh International Conference on Learning Representations. 2023.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pages 2170-2179 PMLR, 2019.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning, Advances in neural information processing systems, 1-211211-2124. 2020.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint 1rXiv:1803.10111 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International Conference on Machine Learning, volume 80, pages 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on PMLR, 2019.
machine learning, pages 2555-2566
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 2023.
Nicklas A Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predictive control. In International Conference on Machine Learning, pages 8387-8406. PMLR, 2022.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In AAAI Conference on Artificial Intelligence, 2017.
Peter Huber et al. Robust estimation of location parameter. The annals of mathematical statistics, 1011, 1964.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448 458. pmlr, 2015.
Max Jaderberg, Volodymyr Mnih, Wojeiech Marian Czarnecki, Tom Schaul, Joel Z Leibo. David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference Learning Representations, 2017.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. stat, 1050:3, 2017.
Jacob Devlin Ming-Wei. Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT,, pages 4171-4186, 2019.
Shauharda Khadka and Kagan Tumer. Evolution-guided policy gradient in reinforcement learning. In Advances in Neural Information Processing Systems 31, pages 1188-1200, 2018.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 44X.v:1412.6980, 2014.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations, 2021.
Aviral Kumar, Rishabh Agarwal, Dibya Ghosh, and Sergey Levine. Implicit under-parammterrzaati
inhibits data-efficient deep reinforcement learning. In International Conference on Learning Representations, 2020a.
12

---

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179-111 1191, 2020b.
Arsenii Kuznetsov, Pavel Shvechikov, Alexander Grishin, and Dmitry Vetrov. Controlling overesti- mation bias with truncated mixture of continuous distributional quantile critics, In International Conference on Machine Learning, pages 5556-5566. PMLR, 2020.
Arsenii Kuznetsov, Alexander Grishin, Artem Tsypin, Arsenii Ashukha, and Dmitry Vetrov. Automating control of overestimation bias for continuous reinforcement learning. ArXiv abs/2/ 13523,, 2021.
Romain Laroche, Paul Trichelair, and Remi Tachet Des Combes. Safe policy improvement with baseline bootstrapping. International Conference on Machine Learning, pages 3652-3661. PMLR, 2019.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, pages 5639-5650. PMLR, 2020.
Lihong Li, Thomas Walsh, and Michael Littman. Towards a unified theory of state abstraction for mdps. In AI&M, 2006.
Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint w:Xiv:159.009777 2015.
Long-Ji Lin. Self-improving reactive agents based on reinforcement learning, planning and teaching. Machine learning, 8(3-4):293-321, 1992.
Xingyu Lin, Harjatin Baweja, George Kantor, and David Held. Adaptive auxiliary task weighting for reinforcement learning. Advances in neural information processing systems, 32, 2019.
Michael Littman and Richard S Sutton. Predictive representations of state. Advances in neural information processing systems, 14, 2001.
Guoqing Liu, Chuheng Zhang, Li Zhao, Tao Qin, Jinhua Zhu, Li Jian, Nenghai Yu, and Tie- Yan Liu. Return-based contrastive representation learning for reinforcement learning. In International Conference on Learning Representations, 2021.
Clare Lyle, Mark Rowland, and Will Dabney. Understanding and preventing capacity loss in reinforcement learning. In International Conference on Learning Representations, 2021.
Horia Mania, Aurelia Guy, and Benjamin Recht. Simple random search provides a competitive approach reinforcement learning. arXiv preprint arXiv:/8022 07055, 2018
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, (ituue,,18170, 518(7540):59-5332 2015.
Jelle Munk, Jens Kober, and Robert Babuška. Learning state representation for deep actor-critic control. In 2016 IEEE 55th Conference on Decision and Control (CDC), pages 4667-4673. IEEE, 2016.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 17X.008877,03488 2018.
Kei Ota. Tomoaki Oiki, Devesh Jha, Toshisada Mariyama, and Daniel Nikovski. Can increasing input dimensionality improve deep reinforcement learning? In International Conference on Machine Learning, pages 7424 7433. PMLR, 2020.
Kei Ota, Devesh K Jha, and Asako Kanezaki. Training larger networks for deep reinforcement learning, arXiv preprint arXiv:2102.07920, 2021.
13

---

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style. high-performance deep learning library. In Advances in Neural Information Processing Systems, pages
2019.
Aloïs Pourchot and Olivier Sigaud. Cem-rl: Combining evolutionary and gradient-based methods for policy search, arXiv preprint arXiv: arXiv: 1810.01222, 2018
Balaraman Ravindran, An algebraic approach to abstraction in reinforcement learning. University of Massachusetts Amherst, 2004.
Sahand -rraei-Shoshtari, Rosie Zhao, Prakash Panangaden, David Meger, and Doina Precup. Contin- uous mdp homomorphisms and homomorphic policy gradient. In Advances in Neural Information Processing Systems, 2022.
Martin Riedmiller, Roland Hafner, Thomas Lampe, Michael Neunert, Jonas Degrave, Tom Wiele, Vlad Mnih, Nicolas Heess, and Jost Tobias Springenberg. Learning by playing solving sparse reward tasks from scratch. In International conference on machine learning, pages 4344-4353 PMLR, 2018.
Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. Evolution strategies as a scalable alternative to reinforcement learning. arXiv preprini arXiv: xrX1v:1703.03864, 2017.
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. In International Conference on Learning Representations, Puerto Rico, 2016.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization, In International Conference on Machine Learning, pages 1889-1897, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.063777 2017.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self-predictio representations. In International Conference on Learning Representations, 2020.
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, 2014.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning, pages 9870-9879, PMLR, 2021.
Richard S Sutton, Joseph Modayil, Michael Delp, Thomas Degris, Patrick M Pilarski, Adam White, and Doina Precup. Horde: scalable real-time architecture for learning knowledge from unsuper- vised sensorimotor interaction. In The 10th International Conference on Autonomous Agents and Multiagent Systems-Voluned Volume 2, pages 761-768, 2011.
Guy Tennenholtz and Shie Mannor. The natural language of actions. In International Conference on Machine Learning, pages 1995-2055. PMLR, 2019.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In IEEE/RS/ International Conference on Intelligent Robots and Systems (IROS), pages 5026-5033. IEEE, 2012.
Elise van der Pol, Thomas Kipf, Frans A Olichock, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. In Proceedings the 19th International Conference on Autonomous Agents and MultiAgent Systems, pages 1431-1439, 2020a.
Elise van der Pol, Daniel Worrall, Herke van Hoof, Frans Oliehoek, and Max Welling. Mdp homomorphic networks: Group symmetries reinforcement learning. Advances in Neural Information Processing Systems, 33:4199-4210, 2020b.
14

---

Herke Van Hoof, Nutan Chen, Maximilian Karl, Patrick van der Smagt, and Jan Peters. Stable rein- forcement learning with autoencoders for tactile and visual data. 2016 JEEE/RSJ international conference on intelligent robots and systems (IROS), pages 3928-3934 IEER, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances in neural information processing systems, 28, 2015.
William Whitney, Rajat Agarwal, Kyunghyun Cho, and Abhinav Gupta. Dynamics-aware embeddings. In International Conference on Learning Representations, 2020.
Ronald J Williams. Simple statistical follownn algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229-22 256, 1992.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pincau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 10674-10681, 2021.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. In International Conference Learning Representations, 2022.
Amy Zhang, Harsh Satija, and Joelle Pineau. Decoupling dynamics and reward for transfer learning. arXiv preprint arXiv: 10...8884...,,,.. 2018.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2020.
15

---

Appendix
Table of Contents
A TD7 Additional Details
17
A.1 Algorithm
17
A.2 Hyperparameters.
19
B Experimental Details
21
C Baselines
22
C.I TD3 & TD3+OFE Hyperparameters.
22
C.2 SAC & TQC Hyperparameters
25
C,3 Offline RL Basclines.
28
Design Study
29
Learning Target
D.2 Network Input
................................
32
D.3 Normalization
34
D.4
End-to-end.
37
E Extrapolation Error Study
38
E.1 Extrapolation Error
38
Empirical Analysis
38
F
Policy Checkpoints
41
F.1
Motivation
41
F2
Method
41
F3
Empirical Analysis
43
Ablation Study
46
Main Components
46
G.2 Subcomponents
47
H Offline RL Learning Curves
50
1 Run time
51
J Limitations
52
16

---

A
TD7 Additional Details
A.1 Algorithm
TD7 (TD3+4 additions) has several networks and sub-components:
Two value functions (Q++1,1, Q++1,2).
Two target value functions (Qt,1: Q1,2).
A policy network T++1-
A target policy network ne.
An encoder, with sub-components (1+1199411
A fixed encoder, with sub-components (Jt, St).
A target fixed encoder with sub-components (fr-1, 9t-1)
A checkpoint policy Te and checkpoint encoder Je (g is not needed).
The encoder is composed of two sub-networks outputs an embedding:
Encoder:
9(+1)) a), where each network
f(s).
2 = g==""""00 a).
(12)
At each training step, the encoder is updated with the following loss:
(( = [e+t (Je+1 - (b)+)
(13)
== [ lett1lx)
(14)
where ||x is the stop-gradient operation.
Value function: TD7 uses pair of value functions (as motivated by TD3 [Fujimoto et al., 2018])
(Q++1,1, 12+1,2) each taking input [24-4 24 s',"",, At each training step, both value functions
are updated with the following loss:
((Q++1) Huber (= s,)
(15)
target == r r + clip(min (Q,, Qx,((x)) Qmin: Q
(16)
= 1 zád s a' = . ( 8') te, ~ clip(AV o
(17)
(18)
(19)
Taking the minimum of the value functions is from TD3's Clipped Double Q-learning (CDQ) [Fu- jimoto et al., 2018J). The use of Huber loss [Huber et al., 1964] is in accordance to TD3 with the Loss- Adjustdde Prioritized (LAP) experience replay [Fujimoto et al., 2020J. a' sampled and clipped in the same manner as TD3 [Fujimoto et al., 2018]. The same embeddings (=fa 2f) are used for each value function. Qmin and Qmax are updated at each time step:
Qmin min (Qmins target), Q + max (Qmax, target),
(20)
(21)
where target defined by Equation 16.
Policy: TD7 uses a single policy network which takes input [2", On every second training step (according to TD3's delayed policy updates [Fujimoto et al., 2018]) the policy π is updated with the following loss:
[(T(+1) == -Q+ ABE [Q]lx (a, - a)² Q = 0.5 (Q(+1.1 + (++1,((x))
(22)
(23)
[=²"= 25; an),
(24)
= u s).
(25)
The policy loss the deterministic policy gradient (DPG) [Silver et al., 2014] with behavior cloning term to regularize [Fujimoto and Gu, 2021]. X is set to 0 for online RL.
17

---

After every target_update_ _rreuuue (250) training steps, the iteration is updated and each target (and fixed) network copies the network of the higher iteration:
(Qt,1, Q1,2) Q1112)
(26)
Tt - 7t+1,
(27)
(fr-1.9t-1) (fr1gt)
(28)
(ft,9t) t (fe+1.9411).
(29)
The checkpoint policy and checkpoint encoder only used test time (see Appendix F).
LAP: Gathered experience is stored in replay buffer [Lin, 1992] and sampled according to LAP [Fu-
jimoto al., 2020], prioritized replay buffer D [Schaul et al., 2016] where transition tuple i = (s, a, is sampled with probability
max (18(t))
p(t) ΣjED max (10(i)1,,1)
(30)
[8()) == max (101+1, 2f, s. targett, 2f s,a) targetl).
(31)
where target is defined by Equation 16. As suggested by Fujimoto al. [2020], 18(i) is defined by the maximum absolute error of both value functions. The amount of prioritization used is controlled by hyperparameter a. New transitions assigned the maximum priority of any sample in the replay buffer.
We outline the train function of TD7 in Algorithm 2. There is difference between the train function of online TD7 and offline TD7 other than the value of X, which for online and 0.1 for offline.
Algorithm 2 TD7 Train Function
1: Sample transition from LAP replay buffer with probability (Equation 30).
2: Train encoder (Equation 13).
3: Train value function (Equation 15).
Update (Qmin, Q (Equations 20 21).
mod policy_update suy_uuaate__rrquuncyy
then
6: Train policy (Equation 22).
7: if mod sateet_udate__eeeeee frequency then
8: Update target networks (Equation 26).
18

---

A.2 Hyperparameters
The action space is assumed be in the range [-1, (and is normalized if otherwise). Besides few exceptions mentioned below, hyperparameters (and architecture) are taken directly from TD3 htttps://githbb com/ wtt.coooo/sssaaaee sfujim/TD3.
Table 3: TD7 Hyperparameters.
Hyperparameter
Value
Target policy noise Target policy noise clipping c
N(0,
TD3
(-0.5,
[Fujimoto et al., 2018]
Policy update frequency
LAP [Fujimoto et al., 2020]
Probability smoothing a
0.4
Minimum priority
TD3+BC [Fujimoto and Gu, 2021]
Behavior cloning weight À (Online)
0.0
Behavior cloning weight À (Offline)
0.1
Checkpoint criteria
minimum
Early assessment episodes
Policy Checkpoints (Appendix F)
Late assessment episodes
20
Early time steps
750k
Criteria reset weight
0.9
Initial random exploration time steps Exploration noise
25k
Exploration
N(0,0.1%)
Discount factor 7
0.99
Replay buffer capacity
1M
Common
Mini-batch size
256
Target update frequency
250
(Shared) Optimizer
Adam [Kingma and Ba, 2014]
Optimizer
(Shared) Learning rate
3c
Besides algorithmic difference from TD3, there three implemennattion-leell changes:
Rather than only using QI, both value functions are used when updating policy (Equa-
tion 22).
2. The value function uses ELU activation functions [Clevert et al., 2015] rather than ReLU activation functions.
3. The target network updated periodically (every 250 time steps) rather than using an exponential moving average every time step. This change is necessary due to the use of fixed encoders.
We evaluate the importance of these changes in our ablation study (Appendix G). Network architecture details are described in Pseudocode
19

---

Pseudocode 1. TD>7 Network Details
Variables:
zs_dim 256
Value Q Network:
b TD7 uses two value networks each with the same network and forward pass.
10 Linear istate__tt
action_dinn 256)
11 Linear (zs dim * 2 256, 256)
12 Linear (256, 256)
13 = Linear L256 1)
Value Q Forward Pass:
input concatenate((statte, action])
X concatenate(Izsa, zs, x])
ELU(11(x))
x ELU(12(x))
value 13(x)
Policy Network:
10 = Linear (state_sstt__dngg 256)
11 Linear (zs_dim 256, 256)
12 = Linear (256, 256)
13 Linear (256, action_dim)
Policy Forward Pass:
input state
x (vgLINNor(((((nnutt)
concatenate([2s, x])
x ReLU(11(x))
ReLU(12(x)) action t1n)(13(x))
State Encoder / Network:
11 Linear (state_ddim 256)
12 = Linear (256, 256)
13 Linear (256, zs_dim)
State Encoder Forward Pass:
input = state
ELU(11 (input))
ELU(12(x))
zs AvgL (N))
State-Action Encoder g Network:
11 Linear faction_dim zs_dim, 256)
12 Linear (256, 256)
13 = Linear (256, zs_dim)
State-Action Encoder Forward Pass:
input = Concatenate(Taction, zs])
ELU(11(input) ELU(12(x))
zsa 13(x)
20

---

B Experimental Details
Environment. Our experimental evaluation is based on the MuJoCo simulator [Todorov al., 2012] with tasks defined by OpenAl gym [Brockman et al., 2016] using the v4 environments. No modification are made the state, action, or reward space.
Terminal transitions. All methods use a discount factor 0.99 for non-terminal transitions and y for terminal transitions. The final transition from episode which ends due time limit is not considered terminal.
Exploration. To fill the replay buffer before training, all methods initially collect data by following uniformly random policy for the first 25k time steps (256 for TQC). This initial data collection is accounted for in all graphs and tables, meaning "time steps" refers to the number of environment interactions, rather than the number of training steps. Methods based on the deterministic TD3 add Gaussian noise to the policy. Methods based on the stochastic SAC do not add noise.
Evaluation. Agents are evaluated every 5000 time steps, taking the average undiscounted sum of rewards over 10 episodes. Each experiment is repeated over 10 seeds. Evaluations use deterministic policy, meaning there is no noise added to the policy and stochastic methods use the mean action. Methods with our proposed policy checkpoints use the checkpoint policy during evaluations. Evalua- tions are considered entirely independent from training, meaning no data is saved, nor is any network updated.
Visualization and tables. As aforementioned, time steps in figures and tables refers the number of environment interactions (all methods are trained once per environment interaction, other than the initial random data collection phase). The value reported by the table corresponds to the average evaluation over the 10 seeds, at the given time step (300k, 1M or 5M for online results and 1M for offline results). The shaded area figures and the term tables refers to a 97.5% confidence interval. Given that methods are evaluated over 10 seeds, this confidence interval CI is computed by
1.96
v10 a.
(32)
where is the sample standard deviation with Bessel's correction the corresponding evaluation. Unless stated otherwise, all curves are smoothed uniformly over window of 10 evaluations.
Offlnne Experiments. Offline results are based the D4RL datasets [Fu et al., 2021], using the v2 version. No modification are made to the state, action, or reward space. The reported performance is from a final evaluation occurring after 1M training steps, which uses the average D4RL score over 10 episodes. TD7 results are repeated over 10 seeds. Baseline results are taken from other papers, and may not be based an identical evaluation protocol in terms of number of episodes and seeds used.
Software. We use the following software versions:
Python 3.9.13
Pytorch 2.0.0 [Paszke et al., 2019]
CUDA version 11.8
Gym 0.25.0 [Brockman al., 2016] MuJoCo 2.3.3 [Todooov et al., 2012]
21

---

C Baselines
C.1 TD3 TD3+OFE Hyperparameeerss
Our TD3 baseline [Fujimoto et al., 2018] uses the author implementation https://gi thub. com/ sfujim/TD3. Our TD3+OFE baseline [Ota et al., 2020] uses the aforementioned TD3 code alongside a re-implementation of OFENet. For hyperparameters, the action space is assumed to be in the range 1].
Environment-speeffic hyperparameters. OFE uses a variable number of encoder layers depending on the environment. However, the authors also provide offline approach for hyperparameter tuning based on the representation loss. Our choice of layers per environment is the outcome their hyperparameter tuning. Additionally, when training the state-action encoder, OFE drops certain state dimensions which correspond to external forces which are hard to predict. In practice this only affects Humanoid, which uses 292 out of a possible state 376 dimensions on Humanoid. All other tasks use the full state space.
Hyperparameters of both methods listed in Table 4. Network architecture details of TD3 described Pseudocode Network architecture details of TD3+OFE are described in Pseudocode and Pseudocode 4.
Table 4; TD3 & TD3+OFE Hyperparameters.
Hyperparameter
Value
Encoder layers
OFE [Ota et al., 2020]
HalfCheetah
Hopper
6
Walker2d
6
Ant
6
Humanoid
8
Target policy noise
0.0.
TD3 [Fujimoto et al., 2018]
Target policy noise clipping c
(-5.
Policy update frequency
2
Initial random exploration time steps
25k
Exploration
Exploration noise
00,00.13)
Discount factor y
0.99
Replay buffer capacity
1M
Common
Mini-batch size
256
Target update rate
0.005
(Shared) Optimizer
Adam [Kingma and Ba, 2014]
Optimizer
(Shared) Learning rate
22

---

Pseudocode 2. TD33 Network Details
Value Q Network:
>> TD3 uses two value networks each with the same network and forward pass.
11 = Linear Linear(stat__diii) __mate_dimmi_e + action_dim, 256)
12 Linear (256, 256)
13 = Linear (256, 1)
Value Q Forward Pass:
input (Cecatttnate(stateee action])
ReLU (input))
x ReLU(12(x))
value 13(x)
Policy Network:
11 Linear (state_diin _dme 256)
12 Linear (256, 256)
13 = Linear (256, action _im)
Policy ı Forward Pass:
input state ReLU(11 (input))
x ReLU(12(x)) action tanh
Pseudocode 3. TD3+OFE Network Details (TD3)
Variables:
zs_dim
zsa_dim = 240
num num_layers = 6 (or 8)
zs_hdim ss_dim/numl haiin/nummllyerrs 40 (or 30)
zsa_hdim aesa_ddmmnummmaaaaeeee
40 (or 30)
target_dim = state_dim if not Humanoid else 292
Value Q Network:
b TD3+OFE uses two value networks each with the same network and forward pass.
11 Linear (state_di dimm action_dim dim + zsa_dim, 256)
12 Linear (256, 256)
13 Linear (256, 1)
Value Q Forward Pass:
input = zsa
ReLU(11 (input)) ReLU(12(x))
value 13(x)
Policy Network:
11 Linear -im zs_dim, 256)
12 Linear (256, 256) Linear Linear(256, action_dim)
13 =
Policy Forward Pass:
input zs
ReLU(11 (input))
x ReLU(12(x)) action (13((13(x))
23

---

Pseudocode 4. TD3+OFF Network Details (OFE)
State Encoder f Network:
11 Linear Linear(Sttae__ingg zs_hdim)
12 Linear (state_dim zs_hdim, zs_hdim)
13 Linear(state_dinn zs_hdim * 2, zs_ zs_hdim)
Linear(state_die + zs_hdim * (num_layers-1), zs_ ss_hdim)
State Encoder f Forward Pass:
input = state
(input))
input = concatenate(Cinput, x])
swii:hr(cacchhommmm__mmmaaaaaar output concatenatee(iinput, x])
State-Action Encoder far Network:
11 Linear action_dim zs_dim, zsa_hdim)
12 = Linear(state__iin action_dim zs_dim + zsa_hdim, zsa_hdim)
13 Linear(state_ddis action_dim zs_dim ssa_ 2, zsa_ esa_hdim)
...
Lnye__ isnearsstate,,.. action_dim
+ =s ds_ zsa_hdim * (num__ayers--), esa_hdim)
final_layer Linear
action_dim zs_dim zsa_dim, target_dim)
State-Action Encoder fa Forward Pass:
input = concateeatte([action zs])
input = ponncatenatte(innuus x])
swissh(fsacchmmmm(((((ittt))))))
swisshhbattchmmmm___nL_eaeeeeeee (input)))
output concatenatee x])
Final Layer t Forward Pass:
input zsa
output final_layer (zsa)
24

---

C.2 SAC & TQC Hyperparameters
Our SAC baseline [Haamoja al., 2018] based on our TD3 implementation, keeping hyper- parameters constant when possible. Remaining details are based on the author implementation https: https://github http:///gihhub..o/aaaaaaaaaaaa com/haarnoja/saaaaaaaaaaaaaoaaaa Our TQC baseline [Kuznetsov al., 2020] uses the author PyTorch implementation https: //gi cwtccoom//aaaae comm/SamuuggLabs. http//tccpppootteeeeeeeeoomm (with evaluation code kept consistent for all methods). For hyperparameters, the action space assumed to be in the range [-1,1].
TQC is based on SAC and uses similar hyperparameters. One exception the use of an e offset. As suggested in the appendix of Haarnoja et al. [2018], the log probability of the Tanh Normal distribution is calculated as follows:
log (a]s) log (uls) log(1 tanh(u)² +),
(33)
where u is the pre-activation value of the action, sampled from N, where N is distributed according M(u, )²), from the outputs j and log a of the actor network. Kuznetsov [2020] use an alternate calculation of the log of the Tanh Normal distribution which eliminates the need for an offset:
log (als) log (uls)) log(2) log(sigmoid(2u)) log(sigmoid(-2u))
(34)
Environment-speeing hyperparameters. TQC varies the number of dropped quantiles depending the environment. While it possible to select this quantity based heuristics [Kuznetsov et al., 2021], we use the author suggested values for each environment.
Hyperparameters of SAC are listed in Table 5. Network architecture details of SAC are described in Pseudocode 5.
Hyperparameters of TQC are listed in Table 6. Network architecture details of TQC are described in Pseudocode 6.
25

---

Table 5: SAC Hyperparameters.
Hyperparameter
Value
Target Entropy
action_dim
SAC [Haarnoja al., 2018]
Policy log standard deviation clamp
--2, 2]
Numerical stability offset
le 6
Exploration
Initial random exploration time steps
25k
Discount factor 7
0.99
Replay buffer capacity
1M
Common
Mini-batche size
256
Target update rate
0.005
(Shared) Optimizer
Adam [Kingma and Ba, 2014]
Optimizer
(Shared) Learning rate
3c -4
Pseudocode 5. SAC Network Details
Value Q Network:
D SAC uses two value networks each with the same network and forward pass.
11 = Linear(state, _emee_dim-- action_dim, 256)
12 Linear (256, 256)
13 = Linear (256, 1)
Value Q Forward Pass:
input concatenate action])
ReLU (inputt)
x ReLU(12(x))
value = 13(x)
Policy Network:
11 Linear (state_dis d_mee_dimm__ 256)
12 Linear (256, 256)
13 Linear (256, action_ding 2)
Policy π Forward Pass:
input state
x ReLU(11 (input))
ReLU(12(x))
mean, log_std = 13(x)
Normal (mean, (allo__std.clip(-22 2))) 2))) sample()
action (x)
26

---

Table 6: TQC Hyperparameeerss.
Hyperparameter
Value
Number of networks
Quantiles per network
25
TQC
Quantiles dropped per network
[Kuznetsov al., 2020]
HalfCheetah
Hopper Walker2d
Ant
Humanoid
SAC [Haarnoja et al., 2018]
Target Entropy
-action_dim
Policy log standard deviation clamp
[-20, 2]
Exploration
Initial random exploration time steps
256
Discount factor
0.99
Replay buffer capacity
1M
Common
Mini-bated size
256
Target update rate
0.005
(Shared) Optimizer
Adam [Kingma and Ba, 2014]
Optimizer
(Shared) Learning rate
3
Pseudocode 6. TQC Network Details
Value Q Network:
TQC uses five value networks each with the same network and forward pass.
11 Linear(state_die -_itee___m-m-_ action_dim, 512)
12 Linear (512, 512)
13 = Linear(512, 512)
14 Linear (512, 25)
Value Q Forward Pass:
input = concatenate([state, action])
x ReLU(11 (input))
ReLU(12(x))
x ReLU(13(x))
value 14(x)
Policy Network:
11 = Linear Linear (staae_dim, 256)
12 Linear (256, 256)
13 = Linear (256, action_dim * 2)
Policy ı Forward Pass:
input state
ReLU(11 (input))
x ReLU(12(x))
mean, log_std = 13(x)
x Noormal(mean, exp(log_std.cl) clip 2))). sample(
action txnh(x)
27

---

C.3
Offline RL Baselines
We four offline RL baseline methods. The results of each method are obtained by re-running the author-provided code with the following commands:
CQL [Kumar et al., 2020b]. https: h/tp:::///gthbbb httpt//...oo/aaaaaaaauua//////// 2tt000/tttttto.ot.too.oo/o///o/r commit d67dbe9
python examples/cal_mujocan - --Langrange_thresti--1 wiii___wwight=1111111
-~env=ENV
Setting the Lagrange threshold below means the Lagrange version of the code is not used (as dictated by the settings defined in the CQL paper). Default hyperparameters in the GitHub performed substantially worse (total normalized score: 221.3 + 34.4). Some modifications were made to the evaluation code match the evaluation protocol used by TD7.
TD3+BC [Fujimoto and Gu, 2021]. https::/githhbb.ooo com wwteeuujimmuuoo httpitmmmmTDaaaS TD3_BC commit 8791ad7
python main. py --env=ENy
IQL [Kostrikov et al., 2021]. http:///ggtthi...///rrrrr/////lt commi 09d7002
python train_offline.pu - otwpa//www.oooo.oooooooooooaoooo
X--Q [Garg et al., 2023]. http:://ggthubb.cm///////// commit dff09af
python train_offline. py >>eny_name-EENy oww.a///ooooooooooori//////ooooo
--max_clip=7 --double=True --temp=2 --batch_size=255
This setting corresponds to the X'-QL-C version of the method where hyperparameters are fixed datasets.
28

---

D Design Study
In this section we analyze the space of design choices for SALE. All results are based on modifying TD7 without policy checkpoints. A summary of all results is contained in Table 7. Each choice is described in detail in the subsections that follow.
Table 7: Average performance on the MuloCo benchmark at 1M time steps. + captures a 95% confidence interval around the average performance. Results are over 10 seeds.
Algorithm
HalfChectah
Hopper
Walker2d
Ant
Humanoid
TD7 (no checkpoints) TD3
17123 + 296
3361 + 429
5718 308
1008
7381
10574 1 897
3226 + 315
3946 1: 292
3942 1030
5165 + 145
Learning Target (Section D.1)
16438 1: 219
3507 1: 313
5882 262
7415
6337 387
ingget and r
17317 +
3034 647
6227
218
7939
578
17089 + 219
3264
+ 670
6255 + 220
8300
675 6288
6841 + 252
12445 +: 417
3009 +: 528
4757
352
5657
Network Input (Section D.2)
Q. remove a
16948 +: 250
2914
756
5746
254
6911 552
6907 + 443
Q, remove z²
16856 + 241
2783
+ 710
5987
407
7954
7196 363
Q, remove 8,
17042 + 208
2268
+ 723
5486 + 466
8546
655
6449 + 478
Q. z" only
15410 + 691
2781
+ 585
5205 + 409
7916 + 976
5551 + 1144
Q. a only
12776 + 774
2929
631
4938 1 422
5695 875
6209 1 384
T. only
16997 + 265
3206
5744
7710 593
1521
T, 2 only
17283 + 388
2429 + 705
6319 + 331
8378 732
6507 + 496
No fixed embeddings
17116 + 149
2262 + 590
5835 + 481
8027 + 551
7334 + 153
Normalization (Section D.3)
normalization on &
17231 + 246
2647 + 466
5639 + 1248
8191 846
5862 1471
No normalization
17275 + 288
3359 + 479
6168 + 164
7274 + 662
4803 + 1706
Normalization zaa
16947 + 284
3383 + 262
5502 + 1142
8049 + 538
6418 - 302
BatchNorm
17299 +: 218
2318 + 473
4839 + 987
7636 769
6979 325
LayerNorm
17132 + 360
2451
+ 902
4470 1096
1024
6712
1297
Cosine similarity loss
16897 + 372
3324
249
5566 + 353
7873 511
4370 + 1573
End-to-end (Section D.4)
End-to-end, 0.1
16186 + 360
1820 + 674
5013 729
6601
1251
5076 897
End-to-end, 1
15775 + 658
1779 + 537
4882 + 710
6604
1135
5880 + 703
End-to-end, 10
16472 1: 365
1534 +: 341
4900 +: 760
6626 1: 1253
5279 1 1095
29

---

D.1 Learning Target
The encoder (fr+1. 90+1) in TD7 uses the following update:
[(fe+1,9441 (g++1(fc+1 (s)x((')lx)
(35)
(2 -
(36)
In this section, we vary the learning target (originally =f+1). (Figure 7) Target 1: s'. Sets the learning target to the next state s', making the encoder loss:
( == (=t+1 - 8))²
(37)
This target is inspired by OFENet [Ota al., 2020].
(Figure 7) Target 2: "target Sets the learning target to the embedding of the next state Fürget taken from a low-moving target network. This makes the encoder loss:
[(J++1, Ic+1) = == |=(41|x)
(38)
Faaget is taken from target network furger where fe+1 has parameters 0 and furget has parameters 0', and 0' is updated at each time step by:
0' - ([0.9))0'+00001
(39)
This target is inspired by SPR [Schwarzer et al., 2020].
(Figure 7) Target 3: = and r. Sets the learning target to the embedding of the next state 2,41 (which the same as the original loss) but adds a second loss towards the target r. This is achieved by having the output of g be two dimensional, with the first dimension being a and the second pred making the encoder loss:
[(ft+1.9+11) = (942((fe11()) a) 111-2((()))'))) + (0-100=133)))011
(40)
22+1
let4ulx) + (1-td (()1--)) -r)
(41)
This target is inspired by DeepMDP [Gelada et al., 2019].
(Figure 8) Target 4: z'aa, Sets the learning target to the next state-action embedding 2+11 where the action is sampled according to the target policy with added noise and clipping (as done in TD3 [Fujimoto et al., 2018]). This makes the encoder loss:
s'a'
[(f+11 = ( 12
(42)
a' == و s')+E, +
(43)
E <~clip(N -0,
(44)
30

---

Table 8: Average performance on the MuJaCo henchmark at 1M time steps. + captures 95% confidence interval around the average performance. Results are over 10 seeds.
Algorithm
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
TD7 (no checkpoints)
17123 + 296
3361 + 429
5718 + 308
8605 + 1008
7381 + 172
TD3
10574 + 897
3226 + 315
3946 + 292
3942 + 1030
5165 + 145
16438 +: 219
3507 +: 313
5882 1: 262
7415 1: 753
6337 +: 387
Angget
17317 +: 65
3034 + 647
6227 +: 218
7939 + 334
6984 + 578
x" and
17089 + 219
3264 + 670
6255 + 220
8300
+ 675
6841 + 252
12445 1 417
3009 1 528
4757 1: 352
6288 1: 570
5657 1: 110
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
20
16
12
6
3
Time steps (1M)
Time steps (IM)
Time steps (IM)
Time steps (1M)
Time steps (1M)
TD7
TD3
Piarget
and
Figure Learning targets from baseline methods. Learning curves on the MuJoCo benchmark, using learning targets inspired by baseline methods, the next state from OFFNet [Ota al., 2020), the next state embedding zügger from slow-moving target network from SPR [Schwarzer et al., 2020], and including reward term (z" and r) as in DeepMDP [Gelada et al., 2019]. Note that this study only accounts for changes to the learning target and does not encompass all aspects of the baseline algorithms. Results are averaged over 10 seeds. The shaded arca captures 95% confidence interval around the average performance.
HalfChectah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
20
16
12
3
0
0.5
1.0
0.5
1.0
0.5
0.5
1.0
0
0.5
1.0
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TDD7
TD3
Figure 8: Alternate learning target. Learning curves on the MuJoCo benchmark, using a learning target of the next state-action embedding z²'a', Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
31

---

D.2 Network Input
SALE modifies the value function and policy take the embeddings as additional input:
Qs+1(s, + Q++1 s,a,,
T(+1(s) TT+1 (2",
(45)
In this section, we vary the input of the value function and policy. In all cases, changes are made independently (i.e. fix the policy the above (using embeddings) and change the policy, or fix the value function to the above (using embeddings) and change the policy).
(Figure 9) Input 1: Q, remove = We remove 00 from the value function input. This makes the input: Q(+1 (2f; ; 8, a).
(Figure 9) Input 2: Q, remove z². We remove 2 from the value function input. This makes the
input: (2 s,a).
(Figure 9) Input 3: Q, remove We remove a from the value function input. This makes the input: (2 ,2*) 2²).
=00 (Figure 10) Input 4: 2, only. We only use 0 in the value function input. This makes the input: (=f").
(Figure 10) Input 5: Q, 8, only. We only use 8, the value function input. This makes the input: Q1+1(s,)) a). As for all cases, the policy still takes in the state embedding 2".
(Figure 11) Input 6: only. We remove from the policy input, using only the state as input. This makes the input: (s). Note that the value function still uses the embeddings as input.
(Figure 11) Input 7: n, 2 only. We remove from the policy input, using only the embedding 2 as input. This makes the input: (2f).
(Figure 12) Input 8: No fixed embeddings. We remove the fixed embeddings from the value function and the policy. This means that the networks use embeddings from the current encoder, rather than the fixed encoder from the previous iteration. This makes the input: (2f4 (2f41 2841; a) and Tt-1 s).
Table 9: Average performance on the MuloCo benehmark at 1M time steps. + captures 95% confidence interval around the average performance. Results are over 10 seeds.
Algorithm
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
TD7 (no checkpoints) TD3
17123 + 296
3361 + 429
5718 + 308
8605 + 1008
7381 + 172
10574 + 897
3226 + 315
3946 + 292
3942 + 1030
5165 + 145
Q, remove nd
16948 + 250
2914 + 756
5746 + 254
6911 + 552
6907 + 443
Q. remove x*
16856 +: 241
2783 +: 710
5987 1: 407
7954 +: 248
7196 + 363
Q. remove a,a
17042 + 208
2268 +: 723
5486 + 466
8546 + 655
6449 + 478
Q. s0 only
15410 691
2781 585
5205 409
7916 976
5551 1144
Q, 8, a only
12776 + 774
2929
+ 631
4938 + 422
5695 + 875
6209 + 384
TT, only
16997 + 265
3206
+ 419
5744 + 430
7710 + 593
6413 + 1521
T, x only
17283 + 388
2429
+ 705
6319
+ 331
8378 + 732
6507 + 496
No fixed embeddings
17116 1 149
2262
1 590
5835 1 481
8027 1 551
7334 1 153
32

---

HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
20
16
12
3
0
0
0.5
1.0
0.5
1.0
0.5
1.0
0
0.5
1.0
0
0.5
1.0
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
TD3
Q. remove nd
Q, remove
Qu remove 8, a
Figure 9: Removing terms from the value function input. Learning curves on the MuJoCo benchmark, when removing components from the input to the value function. Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
20
16
12
3
0
0.5
1.0
0.5
1.0
0.5
1.0
0.5
1.0
0
0.5
1.0
Time steps (IM)
Time steps (IM)
Time steps (IM)
Time steps (IM)
Time steps (1M)
TD7
TD3
=" only
Q. s, a only
Figure 10: Exclusive value function input. Learning curves on the MuloCo benchmark, changing value function input to a single component. Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
20
16
12
3
0
0.5
1.0
1.0
0.5
0.5
1.0
0
0.5
1.0
Time steps (IM)
Time steps (IM)
Time steps (IM)
Time steps (1M)
Time steps (1M)
TD7
TD3
n, only
प, z" only
Figure 11: Modified policy input. Learning curves on the MuJoCo benchmark, changing the policy input to a single component. Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
T Reward (Ik)
20
16
12
3
0
0.5
1.0
1.0
0.5
1.0
0.5
1.0
0
0.5
1.0
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
TD3
No fixed embeddings
Figure 12: Fixed embeddings. Learning curves on the MuJoCo benchmark, without fixed embeddings for the value function and policy, so that the current encoder is used rather than the fixed encoder. Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
33

---

D.3 Normalization
SALE normalizes embeddings and input through our proposed AvgL Norm:
Avgl. A ==
I
(46)
Σg
AvgL is embedded into the architecture used by SALE (see Pseudocode 1), applied to the output of the state encoder f give the state embedding 2* AvgLINorm(f(s)) It is also used following a linear layer in the input of both the value function and policy. This makes the input to the value function and policy:
Q(2th (frgL AvgL
(47)
πI (AvgI.. (s))) I (o)))
(48)
where daa is the output from a linear layer the state-action input, Linear(s, a), and d= is the output from a linear layer on the state input. Linear(s).
In this section, we vary how normalization is used, swapping which components it is used on and what normalization function used.
(Figure 13) Normalization 1; No normalization on ø. Normalization is not applied to either ø. This makes the input to the value function and policy:
(= (frm 6="),
(49)
x (AvgLINorm(f(s)) o²).
(50)
(Figure 13) Normalization 2: No normalization. Normalization is not applied anywhere. This makes the input to the value function and policy:
Q(=" ((),
(51)
xf{(s)
(52)
(Figure 14) Normalization Normalization used on "0 Since 2T is trained to predict the next state embedding 2²', which normalized, normalization may not be needed for the output of the state-action encoder g. We try also applying normalization the output of g. This makes the input to the value function and policy:
(orgLLNN a,),,
(53)
ग (AvgLINorm Avom
(54)
(Figure 15) Normalization 4: BatchNorm. BatchNorm is used instead of A This makes the input to the value function and policy:
(( (==00 BatchNNom(m/(s). BatchNorm(dd""))
(55)
न (BatchNorm(f(ss) BatchNormmo"))
(56)
(Figure 15) Normalization 5: LayerNorm. LayerNorm is used instead of AvgL This makes the input to the value function and policy:
Q(= LayerNNrmm((ss) LayerNorm(( Layer, (d*)),
(57)
क (1 (LayerNorm([((s) LayerNorm(p"))
(58)
(Figure 16) Normalization 3: Cosine similarity loss. This is inspired by SPR [Schwarzer et al., 2020], where normalization not used and a cosine similarity loss is used instead of MSE for updating the encoder. This means the encoder loss function, which is originally:
C(f.g) = (o(f(s), -(f(s)))) = (200 1006 1x)
(59)
now becomes
C(fg)) = Cosine sa |222
6,00|2201 1,22111x) 1xx,
(60)
20
1200
Cosine (x00,120|1 = || ||2
(61)
No normalization is used elsewhere.
34

---

Table 10: Average performance on the MuJoCo benchmark at 1M time steps. + captures 95% confidence interval around the average performance. Results are over 10 seeds.
Algorithm
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
TD7 (no checkpoints) TD3
17123 + 296
3361 + 429
5718 + 308
8605 1008
7381 + 172
10574 + 897
3226 + 315
3946 + 292
3942 + 1030
5165 + 145
No normalization on &
17231 +: 246
2647 1: 466
5639 + 1248
8191 +: 846
5862 + 1471
No normalization
17275 1: 288
3359 1 479
6168 1 164
7274 662
4803 +: 1706
Normalization gda
16947 + 284
3383 + 262
5502 1142
8049 538
6418 302
BatchNorm
17299 + 218
2318 + 473
4839 + 987
7636 769
6979 + 325
LayerNorm
17132 + 360
2451 + 902
4470 + 1096
6331 + 1024
6712 + 1297
Cosine similarity loss
16897 1 372
3324
1 249
5566 1 353
7873 + 511
4370 + 1573
HalfChcetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
16
12
3
Time steps (1M)
Time steps (1M)
Time steps (IM)
Time steps (1M)
Time steps (1M)
TD7
TD3
No normalization on &
No normalization
the MuloCo benchmark, where normalization is removed on certain components the input the value function and policy. Results are averaged over 10 seeds. The shaded area captures a 95% confidence interval around the average performance.
Figure 13: Removing normalization. Learning
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (Ik)
20
9
16
6
0.5
1.0
0.5
0.5
1.0
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
TD3
Normalization applicd to "
Figure 14: Normalization
Learning curves the MuJoCo benehmark, where normalization is also
applied the state-action embedding =d Results averaged over 10 seeds. The shaded captures 95%
confidence interval around the average performance.
35

---

HalfChcctah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
16
12
3
Time steps (1M)
Time steps (IM)
Time steps (IM)
Time steps (1M)
Time steps (1M)
TD7
TD3
BatchNorm
LayerNorm
Figure 15:
Alternate normalization. Learning curves the MuJoCo benchmark, where BatchNorm
LayerNorm is used instead AvoLINorm. Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Roward (Ik)
20
16
6
3
0.5
1.0
0.5
0.5
1.0
Time steps (1M)
Time steps (1M)
Time steps (IM)
Time steps (1M)
Time steps (1M)
TD7
TD3
Cosine similarity loss
Figure 16: Cosine similarity instead of normalization. Learning curves the MuJoCo benchmark, where normalization is removed and replaced with cosine similarity loss, inspired by SPR [Schwarzer al., 2020].
Results are averaged over 10 seeds. The shaded area captures a 95% confidence interval around the average performance.
36

---

D.4 End-to-end
The embeddings used SALE are decoupled, meaning that the encoders used output the embed- dings are trained independently from the value function policy, Instead, we might consider training the encoders end-to-end with the value function. Doing so requires training the current embedding with the value function loss (rather than using the fixed embedding). This means that the encoders are trained with the following loss:
C(1+119941) := Huber target Q0+1 8,)) +
(= + l+ttaxx)
(62)
target clip(min (z), Qx,2(1))) Qmin Qmax),
(63)
x = (z| = ,,'), d' ne (z ,s)) te, <~ clipp(N o²), -c.
(64)
(65)
(66)
In Table 11 and Figure 17 display the performance when varying the hyperparameter B.
Table 11: Average performance the MuJoCo benchmark at IM time steps. + captures 95% confidence interval around average performance, Results are over 10 seeds.
Algorithm
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
TD7 (no checkpoints) TD3
17123 + 296
3361 + 429
5718 +: 308
8605 + 1008
7381 + 172
10574 + 897
3226 + 315
3946 + 292
3942 + 1030
5165 + 145
End-to-end, 0.1
16186 + 360
1820 +: 674
5013 + 729
6601 + 1251
5076 + 897
End-to-end, 1
15775 + 658
1779
+ 537
4882 + 710
6604 + 1135
5880
+ 703
End-to-end, 10
16472 + 365
1534 + 341
4900 760
6626 + 1253
5279
1095
HaifCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
20
9
16
12
6
0
0.5
1.0
0.5
1.0
0.5
1.0
0
0.5
1.0
0
0.5
10
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
=
3 10
Figure 17: End-to-end learning. Learning curves on the MuloCo benchmark, when training the embeddings end-to-end with the value function, where the encoder loss is weighted with respect to B. Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
37

---

E Extrapolation Error Study
E.1 Extrapolation Error
Extrapolation error is an issue with deep RL where the value function poorly estimates the value of unseen or rare actions [Fujimoto al., 2019]. we assume that (s, a) not contained in the dataset, then Q(s, a) is effectively a guess generated by the value function Q. During training, is always sampled from the dataset, but is often generated by the policy (a (s)). resulting potentially poor estimate.
Extrapolation error is typically considered a problem for offline RL, where the agent is given fixed dataset and cannot interact further with the environment, as actions sampled from the policy may not be contained in the dataset. Extrapolation error is not considered problem in online RL since the policy interacts with the environment, collecting data for the corresponding actions it generates.
In empirical analysis (Section E.2) observe the presence of extrapolation error when using SALE. Our hypothesis is that by significantly expanding the input dimension dependent on the action, the network becomes the original action dimension size 8, but the state-action embedding 2 has a dimension size of 256). This is because unseen actions can appear significantly more distinet from seen actions, due to the dramatic increase in the number of dimensions used to represent the action-dependent input.
prone to erroneous extrapolations (for example, for the Ant environment,
E.2 Empirical Analysis
In this section we vary the input the value function in TD7 to understand the role of the input dimension size and extrapolation error. The default input to the value function in TD7 is as follows:
Q( 2000,,2"00 22
(67)
Φ = AvgLINorm(Lineal a)).
(68)
Throughout all experiments, we assume no value clipping (as this is introduced response to the analysis from this section). This makes the value function loss (originally Equation 15):
[(Q++1) = Huber (target - Q2t1 12, 2,),
(69)
target == + 7 min (2), (z,)
(70)
= |- 24' ,
(71)
a' = n. (=(==_11))) 5) + te,
(72)
<~>lip(AV o
c).
(73)
Policy checkpoints are not used and the policy is not modified. Since extrapolation error is closely linked available data, we also vary the maximum size of the replay buffer (default: 1M).
(Figure 18) Extrapolation 1: No clipping. This is TD7 without value clipping in the value function loss discussed above.
(Figure 19) Extrapolation 2: No aa We remove t0 from the value function input. This makes the input: Q4+1(2f o).
(Figure 20) Extrapolation 3: Small Ø. We reduce the number of dimensions of from 256 16.
(Figure 21) Extrapolation 4: No s0 and small di We remove =a from the value function input and reduce the number of dimensions of from 256 to 16.
(Figure 22) Extrapolation 5: Frozen embeddings. The encoders are left unchanged throughout training, by leaving them untrained. This means the input of the value function is Q++1 e).
(Figure 23) Extrapolation 6: TD7, The full TD7 method (without checkpoints).
38

---

Performance by seed
Value estimate by seed
Performance by seed
Value estimate by seed
Total Reward (Ik)
Value Estimate (Ik)
Total Reward (Ik)
Value Estimate (1k)
9
3
4
2
5
2 3
4 5 Time steps (1M)
Time steps (IM)
Time steps (IM)
Time steps (IM)
Seeds (1-10), buffer size 1M
Seeds (1-10), buffer size 100k
Figure 18: No clipping. The performance curve and corresponding value estimate made by the value function, where no value clipping is used. The results for each individual seed are presented for replay buffer size of either IM (left) 100k (right). The environment used is Ant.
Total Reward (Ik)
Performance by seed
Value Estimate (Ik)
Value estimate by seed
Total Reward (Ik)
Performance by seed
Value Estimate (Ik)
Value estimate by seed
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Seeds (1-10), buffer size 1M
Seeds (1-10), buffer size 100k
Figure 19: No 50 The performance and corresponding value estimate made by the value function, where 9a not included input. The results each individual seed are presented for a replay buffer size of
either IM (left) or 100k (right). The environment used is Ant.
Total Reward (Ik)
12
Performance by seed
Value Estimate (Ik)
Value estimate by seed
Total Reward (Ik)
Performance by seed
Value Estimate (Ik)
Value estimate by seed
3
2
2
3
0
5
Time steps (IM)
Time steps (IM)
Time steps (IM)
Time steps (IM)
Seeds (1-10), buffer size IM
Seeds (1- 10), buffer 100k
Figure 20: Small d. The performance curve and corresponding value estimate made by the value function, where the dimension size of is reduced. The results for each individual seed are presented for a replay buffer size of either IM (left) 100k (right). The environment used is Ant.
39

---

Performance by seed
Value estimate by seed
Performance by seed
Valuc estimate by seed
Total Reward (Ik)
Value Estimate (Ik)
Total Reward (Ik)
Value Estimate (1k)
9
3
4 Time steps (IM)
3
1
2
5
2 3
0
4 Time steps (IM)
5
Time steps (IM)
Time steps (IM)
Seeds (1-10), buffer size 1M
Seeds (1-10), buffer size 100k
Figure 21: No =" small Φ.. The performance curve and corresponding value estimate made by the value function, where 2 is not included in the input and the dimension size of é is reduced. The results for each individual seed presented for replay buffer size of either 1M (left) 100k (right). The environment used Ant.
Total Reward (Ik)
Performance by seed
Value Estimate (Ik)
Value estimate by seed
Total Reward (lk)
12
Performance by seed
Value Estimate (Ik)
Value estimate by seed
6
6
4
1
2 3
5
Time steps (1M)
Time steps (IM)
Time steps (1M)
Time steps (1M)
Seeds (1-10), buffer size 1M
Seeds (1-10), buffer size 100k
Figure 22:
Frozen embeddings. The performance curve corresponding value estimate made by value
function, where the encoders are not trained. The results for cach individual seed are presented for replay buffer size of cither 1M (left) or 100k (right). The environment used is Ant.
Performance by seed
Value estimate by seed
Performance by seed
Value estimate by seed
Total Reward (Ik)
Value Estimate (Ik)
Total Reward (Ik)
Value Estimate (1k)
3
3
4
1
2
5
2
3
4
0
4
5
Time steps (IM)
Time steps (IM)
Time steps (IM)
Time steps (IM)
Seeds (1-10), buffer size 1M
Seeds (1-10), buffer size 100k
Figure 23: TD7. The performance curve and corresponding value estimate made by the value function of TD7 (using value clipping). The results for each individual seed are presented for replay buffer size of either 1M (left) 100k (right). The environment used is Ant.
40

---

F
Policy Checkpoints
E1
Motivation
Deep RL algorithms are widely known for their inherent instability, which often results in substantial variance in performance during training. Instability deep RI. can oceur on micro timescale (performance can shift dramatically between episodes) [Henderson et al., 2017, Fujimoto and Gu, 2021] and a macro timescale (the algorithm can diverge or collapse with too much training) [Kumar et al., 2020a, Lyle et al., 2021]. In Figure 24 & 25 we show the instability of RL methods by (a) showing how presenting the average performance can hide the instability in a single trial of RL (Figure 24), and (b) measuring how much the performance can drop between nearby evaluations (Figure 25).
Hopper
Figure 24:
Standard presentation Single seed
The training curve as commonly presented in RL papers and the training curve of single seed M. Both curves are from the TD3 algorithm, trained for 5M time steps. The standard presentation is to evaluate every Nieq steps, average scores over N Napisodss evaluation episodes and Needs seeds, then smooth the curve by averaging over a window Nwinoww evaluations. (In our case this corresponds to Nieq 5000, Nepisades = 10, Ness = 10, Nwindow = 10). The learning curve of a single seed has no smoothing over seeds or evaluations Loass winkow 1). By averaging over many seeds and evaluations, the training curves RL can appear deceptively smooth and stable.
Total Reward (1k)
2
1
Time steps (1M)
Hopper
% of Evaluations
100
Figure 25:
window = 1 window = 5 window = 10
80
The % of evaluations which suffer from a performance drop of at least 2% in the next window evaluations for the TD3 algorithm, trained for 5M time steps. For example, for any given evaluation, the likelihood of one of the next 10 evaluations performing at least 50% worse is around 30% Only evaluations from the last 1M training steps are used, so that early training instability is not accounted for. This shows that large drops in performance are common.
60
40
20
80
> %
Decrease
F.2
Method
The overview for our approach for using checkpoints in RL:
Assess the current policy (using training episodes).
Train the current policy (with a number of time steps equal (or proportional) to the number of time steps viewed during assessment).
If the current policy outperforms the checkpoint policy, then update the checkpoint policy,
Note that we aim time episodes are considered entirely separate from the training process. For clarity, we will use the term assessment (rather than evaluation) to denote any measure of performance which occurs during training. In Algorithm 3 we outline the basic policy checkpoint strategy.
training episodes evaluate the checkpoint policy, as any evaluation test
Minimum performance criteria. An interesting choice that remains is how to determine whether the current policy outperforms the checkpoint policy. Instead of using the average performance, we consider the minimum performance over small number of assessment episodes. Using the minimum provides a criteria which is more sensitive to instability, favoring policies which achieve a consistent performance.
Additionally, using the minimum allows us to end the assessment phase early, the minimum performance of the current policy drops below the minimum performance of the checkpoint policy. This allows us to use higher maximum number of assessment episodes, while not wasting valuable training episodes assessing suboptimal policies. In Algorithm 4 we outline policy checkpoints where the minimum performance is used, and the assessment phase is variable length due to early termination.
41

---

Algorithm 3 Basic Policy Checkpoints
1:
for episode = to assessment_epi do
Assessment
2: Follow the current policy ボ4+1 and determine epi soweereewarddd...eeeeo
3: Update the performance measure of the current policy.
4:
Increment
seapsttes_ssinneetraiiniee by the length of the episode.
5:
if actor ke+1 furperforms checkpoint policy then
D Checkpointing
6:
Update checkpoint networks iee Re+1, ft.
Update checkpoint performance.
8:
for to timesteps_ since_training do
Training
9: Train RL agent.
10: Reset the performance measure of the current policy.
Algorithm Policy Checkpoints with Minimum Performance and Early Termination
I:
for episode assessment_ eet sodes do
Assessment
2: Follow the current policy ke+11 and determine epi sode_
3: fin__perfrrance + waanmmi__eefrrmmaaa
episode_ reward)..
4:
Increment timesteps_sincet ttnprttpss_sineee_taininnnn by the length of the episode.
5:
min_ mmii_pprrfrrmncee <. cceckpooint..erformmnn then
> Early termination
6:
End current assessment.
min_performancs
checkpoint_performmant Update checkpoint networks лa + ke+1, fa +- Jt- checkpoint_perf per
then
b Checkpointing
8:
9:
porrtttaeeroorreeeeeee formance + pein_performane
10:
for i- to
siice_traaningg do
b Training
11: Train RL agent,
12:
Reset min_per min_performance. formance.
Early learning. In our empirical analysis we observe that long assessment phases had a negative impact early learning (Table 13). A final adjustment to policy checkpoints is to keep the number of assessment episodes fixed during the early stages learning, and then increase it to much larger value later in learning. Since the policy environment may be stochastic, when increasing the number of assessment episodes, we reduce the performance the checkpoint policy (by factor of 0.9), since the checkpoint performance may be overfit to the single episode used to assess its performance.
In Algorithm I in the main text, several variables are referenced. These are as follows:
checkpoint ieciooiit_ccodiiii tion: The checkpoint condition refers to either (a) the maximum num- ber of assessment episodes being reached (20) or the current minimum performance dropping below the minimum performance of the checkpoint policy.
outperforms: The maximum number of assessment episodes were reached and the current minimum performance higher than the minimum performance of the checkpoint policy.
Variable descriptions. In Table (in Section A.2), which describes the hyperparameters of TD7, several variables referenced. These are follows:
Checkpoint criteria: the measure used evaluate the performance of the current and checkpoint policy (default: minimum performance over the assessment episodes).
Early assessment episodes; the maximum number of assessment episodes during the early stage of learning (default: 1).
Late assessment episodes: the maximum number of assessment episodes during the late stage of learning (default: 20).
Early time steps: the duration of the early learning stage in time steps. We consider early learning to be the first 750k time steps.
Criteria reset weight: The multiplier on the current performance of the checkpoint policy that is applied once, after the early learning stage ends (default: 0.9).
42

---

F.3 Empirical Analysis
TD7 uses a checkpointing scheme which:
Measures the performance of the current and checkpoint policy by the minimum performance over 20 assessment episodes.
Uses a variable length assessment phase, where the assessment phase terminates early if the minimum performance of the current policy is less than the minimum performance of the checkpoint policy.
Uses early learning stage where the maximum number of assessment episodes starts at 1 and is increased to 20 after 750k time steps.
In this section, vary some of the choices made when using policy checkpoints and analyze the stability benefits of using checkpoints.
(Figure 26) Checkpoint 1: Maximum number of assessment episodes. We vary the maximum number episodes used assess the performance of the current policy (default 20). Early termination of the assessment phase is still used.
(Figure 27) Checkpoint 2:
Mean, fixed assessment length. We the mean performance instead of the minimum (default) to determine if the current policy outperforms the checkpoint policy. this case, early termination of the assessment phase is not used. The number of assessment episodes is fixed at 20.
(Figure 27) Checkpoint 3: Mean, variable assessment length. The mean performance is used to determine if the current policy outperforms the checkpoint policy. Early termination of the assessment phase is used. After each episode, terminate the assessment phase if the current mean performance below the checkpoint performance. This means the number assessment episodes variable, up maximum of 20.
(Figure 28) Checkpoint Fixed assessment length. The same number of assessment episodes is used during each assessment phase, without early termination.
(Figure 29) Checkpoint 5: Immediate. The number of assessment episodes starts at the corresponding number (as opposed to starting at and increasing to a higher number after the early learning state completes 750k time steps).
In Table 12 we present the main set of results at 5M time steps. To measure early learning performance, in Table 13 present the performance of using immediate checkpoints 300k time steps. To observe stability benefits, in Figure 30 we display the performance five individual seeds, with and without checkpoints.
Table 12: Average performance the MuJoCo benchmark at 5M time steps. + captures 95% confidence interval around the average performance. Results over 10 seeds.
Algorithm
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
TD7
18165 + 255
4075 + 225
7397 + 454
10133 = 966
10281 + 588
TD7 (no checkpoints)
18328 = 331
3851 372
6519 + 209
10388 = 1024
9521 + 820
Max 10 episodes
18257 + 338
4208 + 52
6856 +: 273
9890 + 661
10689 + 576
Max 50 episodes
17875 + 192
4110 83
7104 + 559
8226 = 743
10239 + 273
Mean, fixed Mean, variahle
18137 + 417
4061 + 160
7347 + 459
9866 + 804
9850 + 497
17943 + 291
4217
6845 + 476
9614 849
10784 + 188
Immediate, 10 episodes
17767 + 374
4218 17
7416 + 451
10090 = 961
9573 + 255
Immediate, 20 episodes
17900 + 127
4078
53
7696 616
9576 935
10219 551
Immediate, 50 episodes
17619 + 275
4033 + 154
6712 1 313
9716 + 1091
9687 1 584
Fixed, 10 episodes Fixed, 20 episodes Fixed, 50 episodes
18209 + 332
4130 56
6890 + 410
10357 - 893
9682 1408
17925
1 334
3777
1 337
6882
1 452
9716 1 772
10471 1. 560
17807
1 183
3759 1 142
6800 1 422
8881 1 408
8830 1. 655
43

---

Tahle 13: Average performance the MuJoCo benchmark at 300k time steps. + captures a 95% confidence interval around the average performance. Results are over 10 seeds.
Algorithm
HalsCheetah
Hopper
Walker2d
Ant
Humanoid
TD7
15031 + 401
2948 + 464
5379 + 328
6171 + 831
5332 + 714
Immediate, 10 episodes
13778 + 856
3306 + 34
4595 + 343
5977 + 764
4660 + 735
Immediate, 20 episodes
12289 1 507
3052 + 195
4518 1 321
5504
1 889
4531 + 738
Immediate, 50 episodes
6658 1 1309
2668
247
2897 1 816
3709
1848
2060 1 620
HallCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (Ik)
20
16
91
12
8
4
3
0
2
3
2
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
TD7 (no checkpoints)
10
50
Figure 26: Maximum number of assessment episodes. Learning curves on the MuJoCo benchmark, varying the maximum number episodes that the policy is fixed for. Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (lk)
20
16
12
6
8
2
3
4
0
2
3
4
0
3
4
5
0
2
3
4
2
3
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (IM)
TD7
TD7 (no checkpoints)
Mean, fixed assessment length
Mcan, variable assessment length
Figure 27: Checkpoint condition. Learning curves on the MuJoCo benehmark, using the mean performance of assessment episodes rather than the minimum. The number of assessment episodes is kept at 20. Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (lk)
20
16
12
6
8
3
2
2
3
4
0
4
2
3
2
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
TD7 (no checkpoints)
10
20
50
Figure 28: Fixed assessment length. Learning curves on MuJoCo benchmark, where there is no early termination the assessment phase (i.c the number of assessment episodes is fixed), and the number of assessment episodes is varied. Results are averaged over 10 sceds. The shaded area captures a 95% confidence interval around the average performance.
44

---

HalfChectah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
16
13
36
3
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
TD7 (no checkpoints)
10
20
50
Figure 29: Immediate. Learning curves the MuJoCo benchmark, where the maximum number of assessment episodes does not start at (i.e. there early learning phase) varied. Results are averaged over 10 seeds. The shaded area captures a 95% confidence interval around the average performance.
HallCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (IK)
20
16
3
2
3
2
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
TD7 (no checkpoints)
Figure 30: Performance of individual seeds with and without checkpoints, Learinng curves of five individual seeds, with and without checkpoints, the MuJoCo benchmark. The shaded area captures 95% confidence interval around average performance.
45

---

G
Ablation Study
In this section we perform ablation study on the components of TD7.
G.1 Main Components
(Figure 31) Ablation I:
No SALE. SALE entirely removed from TD7. This means the input to the value function and policy return the original Q(s, a) and (s), respectively. TD7 keeps LAP, policy checkpoints, and any implementation differences over TD3 (discussed in Section A.1).
(Figure 31) Ablation 2: No checkpoints. Similar to TD3 and other off-policy deep RL algorithms, TD7 is trained at every time step, rather than after a batch of episodes. The current policy is used at test time.
(Figure 31) Ablation No LAP. TD7 uses the standard replay buffer where transitions are sampled with uniform probability. The value function loss (Equation 15) uses the mean-squared error (MSE) rather than the Huber loss.
Table 14:
Average performance on the MuloCo benchmark at 5M time steps. + captures a 95% confidence interval around the average performance. Results are over 10 seeds.
Algorithm
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
TD7
18163
4075
225
7397
10133 +: 966
10281 588
TD3
14337 + 1491
3682 = 83
5078
343
5589 + 758
5433 245
SALE
17099
4018
170
6418
7861 + 253
7275
No checkpoints
18328 + 331
3851
372
6519
209
10388 + 1024
9521 + 820
No LAP
18104 + 315
4188
22
7233 + 251
6940 + 1044
10155 + 522
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
16
13
4 5
88
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1 (1M)
Time steps (1M)
TD7
TD3
D No SALE
No checkpoints
No LAP
Figure 31: Main ablation. Learning curves the MuloCo benchmark, removing the main components of TD7.
Results are averaged over 10 seeds. The shaded area captures a 95% confidence interval around the average performance.
46

---

G.2 Subcomponents
(Figure 31) Ablation 4: No SALE, with encoder. The encoder is kept but trained end-to-end with the value function, rather than via the SALE loss function. This means the encoder can be simply treated as an extension of the architecture of the value function. Since the embeddings trained with the value function may be conflicting with the policy, the policy does not use the state embedding, and only takes original input iT (s)-
(Figure 31) Ablation 5: TD3 with encoder. We add the encoder from SALE to TD3 but train end-to-end with the value function. The encoder applied to the policy.
(Figure 33) Ablation 6: Current policy, TD7 is trained in an identical fashion, but uses the current policy test time, rather than the checkpoint policy.
(Figure 33) Ablation 7: TD3 with checkpoints. We add policy checkpoints TD3, using the checkpoint policy at test time.
(Figure 34) Ablation 8: TD3 with LAP. We add LAP TD3 (identical to the TD3 LAP [Fujimoto et al., 2020)).
(Figure 35) Ablation 9: No clipping. We remove our proposed value clipping for mitigating extrapolation error. This makes the value function loss (originally Equation 15):
[(Q++1) == Huber (target - Qr+1(=i" zi, (()
(74)
target := + y min (2),
(75)
= | 2 8 ,a), a' = n. ( 5)) te,
(76)
(77)
clVi o -c, c).
(78)
(Figure 35) Ablation 10: TD3 with clipping. We add value clipping to the loss function of TD3.
(Figure 36) Ablation 11: No normalization. We remove the use AvgL from TD7. This identical to no normalization in the design study in Section D.3, except this version includes policy checkpoints (the design study does not use checkpoints).
No fixed encoder. We remove the fixed embeddings from the value function and the policy, this means the networks use embeddings from the current encoder, rather than the fixed encoder from the previous iteration. This makes the input: Q++1 (=04 2f+1, a) and TT41 (2f 1. s). This is identical to no normalization in the design study in Section D.3, except this version includes policy checkpoints (the design study does not use checkpoints).
(Figure 36) Ablation 12:
No implementation. We remove the implementation details differences discussed in Section 1, namely using both value functions when updating the policy, ELU activation functions in the value function network. The target network update not changed this may influence the use of fixed encoders.
(Figure 37) Ablation 13:
(Figure 37) Ablation 14: Our TD3. TD3 with the implementation detail differences discussed in Section A. namely using both value functions when updating the policy, ELU activation functions in the value function network, and updating the target network every 250 time steps, rather than using an exponential moving average.
47

---

Table 15: Average performance on the MuJoCo benchmark at 1M time steps. + captures 95% confidence interval around the average performance. Results are over 10 seeds.
Algorithm
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
TD7
17434 155
3512 + 315
6097 570
8509 + 422
7429 + 153
TD3
10574 897
3226 + 315
3946 + 292
3942 + 1030
5165 + 145
No SALE
12981 + 261
3536 +: 65
5237 1: 376
5296 +: 1336
6263 1: 289
No SALE, with encoder
15639 548
3544 + 371
5350 1: 427
6003 1 920
5480 1 83
TD3 with encoder
12495 813
1750 + 302
4226 491
5255 + 579
5082
317
No checkpoints
17123 296
3361 -. 429
5718
308
8605 1 1008
7381 1 172
Current policy
17420
2940 + 636
5763
8748 1 397
7162
TD3 with checkpoints
10255 656
3414 1 77
3266 474
3843 + 749
5349 72
No LAP
17347 207
3697 - 144
6382
339
6571 + 1504
8082 260
TD3 with LAP
10324 1159
3117 + 554
4127 330
4310 + 1150
5090 190
No clipping
17378 100
3762 + 118
6198
289
7695 + 497
7251 +: 274
TD3 with clipping
10283 = 422
2969 = 682
3990 + 258
3711 + 799
5254 + 203
No normalization
17391 275
3640 + 95
6256
317
7807 + 266
4829 + 1809
No fixed encoder
17145 + 138
3710 + 120
5869 + 531
8287 + 379
7412 + 337
No implementation
17334 + 99
3750 + 136
5819 + 71
7756 + 704
7042 + 259
Our TD3
11068
1399
2791 + 632
4179 + 297
5489 + 448
5186 + 108
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
20
16
12
3
0
0.5
1.0
0.5
0.5
1.0
0
0.5
1.0
Time steps (1M)
Time steps (IM)
Time steps (IM)
Time steps (IM)
Time steps (1M)
I TD7
TD3
No SALE
No SALE, with encoder
TD3 with encoder
Figure 32: SALE. Learning
the MuloCo benchmark, varying the usage of SALE. Results averaged
over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
HalfChectah
Hopper
Walker2d
Ant
Humanoid
Total Reward (Ik)
20
16
12
6
0
0.5
1.0
0.5
1.0
0.5
1.0
0
0.5
1.0
0
0.5
1.0
Time steps (IM)
Time steps (1M)
Time steps (IM)
Time steps (IM)
Time steps (1M)
TD7
TD3
No checkpoints
Current policy
TD3 with checkpoints
Figure 33: Checkpoints. Learning curves on the MuloCo benchmark, varying the usage of policy checkpoints.
Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
48

---

HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (Ik)
20
16
12
2
0
0.5
1.0
0.5
1,0
0.5
1.0
0
0.5
1,0
0
0.5
1.0
Time steps (IM)
Time steps (IM)
Time steps (IM)
Time steps (1M)
Time steps (1M)
TD7
TD3
No LAP
TD3 with LAP
Figure 34: LAP. Learning curves the MuloCo henchmark, varying the usage of LAP. Results averaged
over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
HalfChectah
Hopper
Walker2d
Ant
Humanoid
Total Reward (1k)
16
12
3
Time steps (1M)
Time steps (IM)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
TD3
No clipping
TD3 with clipping
Figure 35: Clipping. Learning curves the MuloCo benchmark, varying the usage of proposed value clipping for mitigating extrapolation error. Results are averaged over seeds. The shaded area captures 95% confidence interval around the average performance.
HalfCheetah
Hopper
Walker2d
Ant
Humanood
Total Reward (1k)
20
16
12
3
0
0.5
1.0
1.0
0.5
0.5
1.0
0
0.5
1.0
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
TD3
No normalization
No fixed encoder
Figure 36: SALE components. Learning curves on the MuloCo benchmark, removing components of SALE.
Results averaged over 10 seeds. The shaded area captures 95% confidence interval around the average
performance.
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Reward (Ik)
20
16
6
3
0.5
1.0
0.5
0.5
1.0
Time steps (1M)
Time steps (1M)
Time steps (IM)
Time steps (1M)
Time steps (1M)
TD7
TD3
No implementation Our TD3
Figure 37: Implementation differences. Learning curves the MuJoCo benchmark, varying the minor implementation details hetween TD3 and TD7. Results averaged 10 seeds. The shaded captures 95% confidence interval around the average performance.
49

---

H
Offline RL Learning Curves
HalfChcetah
HalfChectah Medium-Reye Reelaa
HalfChcetah
Medium
Medium- Expert
Normalized Score
120
120
120
100
100
100
80
80
80
60
60
60
40
40
40
20
20
20
0.5
1.0
0.5
1.0
0
0.5
1.0
Time steps (1M)
Time steps (IM)
Time steps (1M)
Hopper
Hopper
Hopper
Medium
Medium-Rephe
Medium Medium-Expert Expert
Normalized Score
100
80
80
60
60
40
40
26
20
20
0
0.5
1.0
0.5
0.5
1.0
Time steps (1M)
Time steps (IM)
Time steps (1M)
Walker2d
Walker2d
Walker2d
Medium
Reeium-Reppaay
Medium-Exxe -xpmetttttttt-t
Normalized Score
120
120
120
100
100
100
80
80
80
60
60
40
40
40
20
20
20
0.5
1.0
0.5
1.0
0
0.5
1.0
Time steps (1M)
Time steps (1M)
Time steps (1M)
TD7
M
CQL
TD3+BC
IQL X-QL
Figure 38: Learning curves on the offline D4RL benchmark. Results are averaged over 10 seeds. The shaded area captures a 95% confidence interval around the average performance.
50

---

I Run time
We benchmark the run time of TD7 and the baseline algorithms. Evaluation episodes were not included in this analysis. Each algorithm is trained using the same deep learning framework, PyTorch [Paszke al., 2019]. All experiments are run on a single Nvidia Titan GPU and Intel Core 17-7700k CPU. Software is detailed in Appendix B. We display run time adjusted learning curves in Figure 39.
HalfCheetah
Hopper
Walker2d
Ant
Humanoid
Total Rewand (Ik)
20
16
13
3
4
12
16
4
16
0
12
16
12
8
16
20
Hours
Hours
Hours
Hours
Hours
TDD7
TD3
E
SAC
TQC
TD3+OFE
Figure 39: Run time adjusted curves. Learning curves on the MuJoCo benchmark over 5M time steps, where the x-axis is run time. Results are averaged over 10 seeds. The shaded area captures 95% confidence interval around the average performance.
51

---

Limitations
Depth vs. Breadth. Although our empirical study is extensive, the scope of our experimentation emphasizes depth over breadth. There remains many other interesting benchmarks for continuous control, such as examining setting with image-basee based observations.
Baselines. We benchmark TD7 against the highest performing agents in both the online and offline setting. However, due to the fact that most representation learning methods do not cover low level states, we only benchmark against a single other method focused on representation learning. While would be surprising if method, originally designed for different setting, outperformed the strongest existing baselines, exploring additional methods could provide new insight into representation learning for low level states.
Computational cost. While TD7 has a lower computational cost than other competing methods, the time over TD3 (the base algorithm) is more than double.
Theoretical results. We perform extensive empirical analysis to uncover which factors of dynamics- based representation learning are important for performance. However, we do not address the theoretical side of representation learning. There are many important questions regarding why dynamics-based representation learning is effective, or what makes good representation.
52


Reviews:
[{"review_id": "l2mqZqFUjm", "reviewer": "Reviewer_Casd", "summary": "This paper introduces a couple of ideas to improve the empirical performance of the TD3 algorithm on continuous-action RL problems.\n\nThe core contribution is to show that learning state and action embeddings that are designed to predict themselves in successive timestep can help achieve more reward. This is pretty nice because the representation learning process is done in conjunction with reward maximization and is shown to be useful and stable enough to enable faster learning in the same task.\n\nA few other contributions are presented which I list below:\n\n1- a normalization approach that scales the output of the embeddings, and which is compared against some of the existing normalization approaches such as batch and layer norm.\n\n2- checkpointing the policy during online RL, akin to supervised learning, so rather than using the latest learned policy one can continue learning by using the best-performing policy.\n\n3- clipping the value estimates during training to ensure that the estimates remain within a meaningful range.\n\n\nVery nice empirical results are then provided, most notably on the Mujoco benchmark in online RL, where the proposed approach named TD7 is capable of beating competitive baselines such as TD3 and SAC.", "questions": "- In what way is the new embedding-learning approach different than the previous work? What are the advantages and disadvantages, and what can we learn about designing embedding-learning algorithms in light of your experiments?\n\n- How is the normalization hedging against the 0 mode collapse?\n\n- Do we not need to do off-policy learning during check-pointing to account for exploration?", "limitations": "NA", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The highlight of the paper is the impressive empirical results provided on the online RL experiments with the Mujoco baseline. I did my best to cross-check the performance of the baseline agents and it does seem like that TD7 is capable of beating TD3 and SAC (these are the two baselines I checked) on these benchmarks. I do see some discrepancy between SAC results reported here and those reported (for Hopper) in the original paper, but otherwise the results are consistent with published papers.", "weaknesses": "I did enjoy reading this paper, and I applaud authors for their successful implementation and empirical results. That said, I am not sure what the core message of the paper is. I think the contributions need to be motivated, framed, and highlighted better. For instance, is this paper primarily about the advantage of learning state and action embeddings in RL? If yes, existing work has already demonstrated that, so other than the fact that the empirical results are superior, what kind of statements can we make about this new approach to learning embeddings that we did not know previously? Is it better for transfer learning? Is it faster in terms of run time? Is it better motivated theoretically?\n\nJust to name a few papers, with the danger of missing other related and interesting work:\n\n- An approach very similar to the one presented here was presented earlier by Zhang et al: \"Learning Action Representations for Reinforcement Learning\"\n\n- Gelada et al also propose a very similar approach in \"Deepmdp: Learning continuous latent space models for representation learning\", and I am not clear in what sense SALE is doing something different than them.\n\n- Chandak et al show in \"Learning Action Representations for Reinforcement Learning\" that one can learn a low-dimensional representation for actions.\n\nThese are just a few examples to show that the idea of learning state and action representations is well-explored and to me this limits the novelty of the work.\n\n\nMoving to the other contributions, and in terms of the normalization approach, I am not quite positive in what way this normalization is going to hedge against the issue of collapsing all states and actions to 0. To me it seems like that if the reward is not part of the process of learning the embeddings, and we use the kind of successive predictive loss in the embedding space, then trivially the 0 solution would be optimal with or without normalization.\n\n\n\nI found the checkpointing idea creative and interesting. That said, I have a few issues with it: 1- it seems like based on Figure 7 that checkpointing is rarely effective in the online case, and as the paper states, not applicable in the offline case. 2- it seems to me that when computing the goodness of a checkpointed policy, a potential issue is that during training we add some exploration noise to the action suggested by the network, so unless we do zero exploration, we need to account for the fact that we have not exactly executed the policy and therefore some off-policy learning is needed to compute the true goodness of the checkpointed policies. Do I get this right?\n\nThe clipping idea also makes sense, and it is interesting, but it is still a limited contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the very detailed review, and we appreciate the highlighted positives in our work! We address your key points below. \n\n**Contribution:** We will aim to tighten up the writing in the introduction and provide a stronger and clearer message for the paper. We believe our paper adds to the current landscape of representation learning in RL by answering (or adding to) the following questions:\n\n- **What design choices should I make if I want to use representation learning for low-level states?** Almost every representation learning paper in RL is applied to image-based tasks (this includes both Zhang et al., and Gelada et al. that the reviewer mentioned). This raises some important design questions when transferring ideas from an image-based task to low-level states. What should the target in the loss function be? What should the embedding be over? Is normalization needed and if yes, how should it be used? Should learning be end-to-end or not? For each of these questions we can point to several papers which take different approaches. Our paper tests each of these choices, shows what choices matter the most, and suggests the best option.\n\n- **What challenges should I be aware of when using state-action representation learning?** While state-only representation learning is fairly common, state-action representation learning is not. Interestingly enough, we find that using this combination can introduce new challenges (i.e. extrapolation error in online RL, section 5.1) that have not been discussed or adequately addressed in prior work.  \n\n- **I just want the highest performing continuous control algorithm for online or offline RL.** The MuJoCo benchmark has been the cornerstone benchmark for continuous control for many years, and many important algorithms such as PPO, TD3 and SAC have been designed around it. This space is highly saturated and improvements over TD3/SAC are either incremental or come at high computational costs. TD7 suffers from neither of these problems, as it significantly outperforms existing methods (including the aforementioned expensive ones!) and comes at a much lower run time cost than the expensive methods (for run time see Appendix I). \n\n**Related work:** There is a lot of important related work in this area and we were careful to cite all three of the papers mentioned. Our method is similar to many existing methods in that they all do representation learning from a dynamics-modelling based objective. We cite at least a dozen papers which also use this same idea. This idea appears often in the literature because it is a fundamental concept relating to Bisimulation metrics. The important question is how do we best implement this idea with complexity of deep RL where there is function approximation, continuous states and actions, finite data, etc. We believe our work adds to this important story on the empirical side. \n\nSpecifically, here’s how our method differs from the mentioned approaches:\n\n- **Zhang et al.** The title listed is presumably a copy error. Our best guess is “Learning Invariant Representations for Reinforcement Learning without Reconstruction” by Zhang et al., but let us know if otherwise. The primary goal of this paper is abstraction in the state space for visual tasks. In SALE we are interested in feature learning in a low-level state space, rather than abstraction in a complex visual space. Instead of trying to summarize the input, we are trying to expand the input with useful features. We also learn a representation over state and action, rather than just the state. This distinction is important– In section 5.1 we find that doing so introduces new challenges which need to be addressed.\n- **Gelada et al.** is closely related. A few obvious differences: (setting) they work in the image space with discrete actions, (conceptual)  they only learn a representation over the state, (design) they include the reward in the target, they learn the representation end-to-end as an auxiliary loss, they don’t use a target network for the encoder. In section 4.2 in our paper, we find that the nuances in design choices can have a significant impact on performance. Similar to Zhang et al, they work in a visual space and learn state-only features. \n- **Chandak et al.** is similar in that they consider actions in their representation. However, while we consider both state and action, they don’t consider the state at all. Furthermore, their approach to representation learning is very different from ours, based on a policy gradient method, rather than learned from a disentangled signal.  \n\n**Normalization:** With AvgL1Norm it is impossible for the embedding to be 0. Even for a very small $\\epsilon$ we have $\\frac{\\epsilon}{ \\frac{1}{N} \\sum \\epsilon} \\approx 1$. Prior work does something similar by normalizing only in the loss function. It turns out that normalizing the embedding is important experimentally, since we use the embeddings concatenated with the original state-action input. \n\n**Checkpoints:** Not sure which figure you are referring to (7 is a typo presumably). Regardless, the value of checkpointing is more in stability than raw performance, since instability can get averaged out. In the general response (and the attached PDF), we show how checkpoints improve the stability of RL algorithms (Figure 4). \n\nThere is indeed some bias introduced by the exploration noise used during checkpoints. However, this isn’t necessarily problematic since the checkpoint is more likely to prefer policies which are robust to noise in the action space. One issue that could arise is if the random exploration noise results in some kind of positive benefit. To combat this, we actually use the minimum performance over N episodes. This choice is discussed and evaluated further in Appendix F.  "}, {"review_id": "DZzeaz76lp", "reviewer": "Reviewer_sfVv", "summary": "This paper proposes TD7, an improved version of the popular TD3 algorithm with 4 additional techniques: state-action learned embeddings (SALE, the major one among the four), using policy checkpoints for stable evaluation, an existing prioritized experience replay method called LAP, an existing offline RL algorithm called TD3+BC. The four techniques strengthen vanilla TD3 algorithm from different aspects. TD7 is evaluated on MuJoCo and D4RL benchmarks for online and offline RL settings, significantly outperforming existing continuous control algorithms. Comprehensive ablation studies are also included.", "questions": "TD7 needs to keep current policy fixed for several episodes:\n\n1. What if the episode horizon is long?\n2. Will this be inconvenient to combine high update frequency (like used in REDQ)?\n3. How to consider the gap between using the checkpoint policy for evaluation (exploitation) while the comparison and selection of checkpoint policies are based on the performance of exploration policies?", "limitations": "- The proposed improvements seems to be general and orthogonal to different continuous control RL algorithms. However, SALE is not combined with SAC, TQC and REDQ to evaluate the generality.\n- TD7 is evaluated on MuJoCo and D4RL whose backends are both OpenAI MuJoCo suite. I would be interested in seeing TD7 in other continuous control benchmarks like DMC.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- This paper is well written and the content is clear.\n- The proposed improvements (almost) make sense to me and well motivated or explained.\n- The part of SALE (Section 4) is interesting to me, although SALE can be viewed as an improvement over OFENet. I think the detailed studies on useful techniques for stable representation learning could be inspiring to other similar problems in RL.\n- This paper conduct ablation studies for most design choices to rule out the alternative choices empirically.\n- The code implementation is neat and I assume it is easy to reproduce the experiments.", "weaknesses": "- The proposed improvements are piecemeal, each of which is incremental or existing.\n- A few additional hyperparameters (e.g., the dimensionality of $z^{s,a}$, the episode number to maintain the policy unchanged) are introduced, although the authors give recommended values.\n- Although the paper has comprehensive empirical studies, intuitive explanations are lacked.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and considerate questions.\n\n**Intuitive explanations:** Space was obviously a bit of a concern with the paper (and fortunately the camera-ready allows for an additional page). We would be happy to expand on some of the intuition/reasoning behind our improvements with this additional space. We are happy to discuss details during the rebuttal phase if there are any specific improvements you have questions about. \n\n**Checkpoint questions:** It’s worth highlighting that checkpoints are not a fundamental aspect of TD7, as they are easy to remove. We ablate over the usage of checkpoints in Figure 5 and more extensively in Appendix F. Notably, we still outperform existing state-of-the-art methods without the use of checkpointing. In situations where checkpoints are problematic, there does exist the option to not use them. However to answer your questions: \n- **Long episode horizon:** This is a good thought. Ultimately any method which relies on any kind of MC estimate could have problems for long horizons. There are some possible modifications to checkpoints such as only considering a finite/shorter horizon when evaluating the checkpoint, or similarly, accounting for the discount factor, or simply evaluating with fewer episodes. Another possibility is to use an off-policy evaluation method rather than using raw values. Ultimately there’s a lot of room for new research for unique settings.\n- **Combining with high frequency updates:** We don’t believe this is an issue, although there is some potential for conflict if the number of updates is very significant, say 1000x, then training approaches the offline setting. For smaller increases (say 2-5x), our empirical analysis (in Appendix F) did not suggest that there were any stability or training issues with training for a large number of updates before collecting new data. In practice, since our method trains after an episode finishes (rather than during the episode), it could be argued that our method is better suited for expensive training methods. This is because training the networks during an episode could cause practical issues with latency. \n- **The gap between the checkpoint policy and the exploration policy:** Great question! We don’t really cover this in the main paper due to a lack of space, but discuss it more in Appendix F. In practice, we use the minimum performance over N episodes, rather than the mean (quantiles could also be used for particularly stochastic environments). Since our goal is stability, the idea is that by using the minimum we put a much higher preference on robust policies. Following that logic, the performance gap between the exploration policy (which includes noise) and the true deterministic policy just helps us choose more robust policies. Again, there are some other possible research questions which arise depending on the ultimate goal from using checkpoints but we leave that to future work. \n\nWe agree with the reviewer that these are all interesting thoughts. We believe that checkpoints are a promising avenue for research and we hope to explore some of these ideas in future work. \n\n**Combining SALE with other algorithms:** Good point! We present results of SAC with SALE compared against vanilla SAC. SALE significantly improves the results of SAC in the majority of the environments. This change only represents the addition of SALE, and with no underlying hyperparameter changes to either SAC or SALE. Since TQC and REDQ are both based on SAC, this shows that SALE could offer performance benefits to those methods as well. Learning curves are included in the PDF of the general response. Results are for 5M time steps and 10 seeds. $\\pm$ captures a 95\\% confidence interval.  \n\n|             | SAC + SALE          | SAC                |\n|-------------|---------------------|--------------------|\n| HalfCheetah | **17316 $\\pm$ 236** | 15526 $\\pm$ 697    |\n| Hopper      | 2842 $\\pm$ 951      | **3167 $\\pm$ 485** |\n| Walker2d    | **6887 $\\pm$ 306**  | 5681 $\\pm$ 329     |\n| Ant         | **9288 $\\pm$ 462**  | 4615 $\\pm$ 2022    |\n| Humanoid    | **9062 $\\pm$ 335**  | 6555 $\\pm$ 279     |\n\n\n**TD7 on other benchmarks:** We have gathered results on the DMC suite. As shown in the table below, TD7 consistently maintains its advantage over TD3 in DMC as well. In the attached PDF of the general response, we include the full learning curves for these environments. Results are for 1M time steps and 10 seeds. $\\pm$ captures a 95\\% confidence interval.  \n\n|                         | TD7                 | TD3                |\n|-------------------------|---------------------|--------------------|\n| Ball in Cup Catch       | **982 $\\pm$ 1.8**   | 981 $\\pm$ 3.3      |\n| Cartpole Balance Sparse | **1000 $\\pm$ 0.0**  | **1000 $\\pm$ 0.0** |\n| Cartpole Swingup        | **875 $\\pm$ 5.8**   | 871 $\\pm$ 6.6      |\n| Cheetah Run             | **914 $\\pm$ 2.9**   | 866 $\\pm$ 23.5     |\n| Finger Spin             | **989 $\\pm$ 1.4**   | 975 $\\pm$ 14.1     |\n| Finger Turn Hard        | **742 $\\pm$ 136.8** | 552 $\\pm$ 80.7     |\n| Hopper Stand            | **774 $\\pm$ 177.4** | 480 $\\pm$ 180.8    |\n| Hopper Hop              | **175 $\\pm$ 98.7**  | 58 $\\pm$ 42.3      |\n| Humanoid Stand          | **504 $\\pm$ 336.5** | 7 $\\pm$ 0.4        |\n| Humanoid Walk           | **389 $\\pm$ 157.6** | 36 $\\pm$ 72.4      |\n| Walker Stand            | **991 $\\pm$ 1.5**   | 985 $\\pm$ 3.1      |\n| Walker Walk             | **981 $\\pm$ 2.7**   | 965 $\\pm$ 5.2      |\n| Walker Run              | **807 $\\pm$ 62.5**  | 631 $\\pm$ 74.7     |\n| Quadruped Walk          | **950 $\\pm$ 4.9**   | 750 $\\pm$ 168.9    |\n| Quadruped Run           | **684 $\\pm$ 87.5**  | 503 $\\pm$ 193.7    |\n| Mean                    | **784**             | 644                |\n| Median                  | **875**             | 750                |\n| IQM                     | **856**             | 733                |"}, {"review_id": "wYtBD5uSjs", "reviewer": "Reviewer_YfJA", "summary": "This work introduces a novel state-action representation learning framework SALE and two other techniques (e.g. checkpointing, a new type of Q value clipping) that substantially improve the data efficiency and final performance of TD3 in online and offline RL.", "questions": "1. In principle, any latent state/action RL framework can be modified such that  $Q$ and $\\pi$ take both the original state/action and the latent state/action as inputs. Would it be fair to say the ablations in Fig. 2 suggests that any state/action representation learning method should have $Q$ and $\\pi$ depend on both the original and latent state/actions?\n1. Fig. 2: Can the authors clarify what is meant by percent loss here? Are these quantities all with respect to default values? Where do the default values fall on these figures?\n1. The authors mention that checkpointing harms performance early in training – is this because the policy changes rapidly early on? Can the authors provide more detail? \n1. Line 217: How does checkpointing change the relevance of data?", "limitations": "See Weaknesses.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The work studies joint state-action representation learning, which is under-studied with respect to state representation learning.\n2. The method is very clear, and the authors provide extensive empirical analysis and ablation explaining their design choices.\n3. Some techniques introduced in this work (checkpointing, Q clipping in Eq. 6) can be applied to most RL methods and are thus influential beyond representation learning.\n", "weaknesses": "**Core comments:**\n1. I believe the paper would read more easily after reorganization.\n    * Section 4.2 feels like an ablation study that should follow the core empirical results. Also, Section 4.2 references TD7 results, though TD7 is not introduced until section 5.\n    * Section 5 begins by mentioning stability and extrapolation error jumped out at me; the intro and related work emphasize the representation learning side of this work much more, and extrapolation error is not mentioned in sections 1-4. I believe the paper would read better if the authors painted a clearer picture of what the stability side of the work entails in the introduction, and included transitions between section 4 and 5 that clearly indicate “we talked about one issue, and now we’re going to talk about another.”\n    * Transition sentences before bolded subheadings (e.g. Normalizing embeddings, Fixed embeddings) would help in Section 4 – something to the effect of “We now discuss two important aspect of SALE…” \n1. The checkpoint policy $\\pi_c$ denotes the policy is the largest minimum return during the assessment phase. During agent evaluation, you compute the return over N rollouts of $\\pi_c$ rather than $\\pi_t$, correct? If so, I’m a bit skeptical of the results; is it possible that the perceived benefits of checkpointing are simply due to the batched nature of training it requires? Suppose you run vanilla TD3 within the checkpointing framework but perform separate evaluations using the checkpoint policy *and* current policy. This experiment would be equivalent to vanilla TD3 with variable length batched updates. Does the batched updates alone improve performance?\n\n\n**Minor comments:**\n1. Figure 2 would be more readable with error bars rather than text to indicate the 95% confidence interval.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and helpful comments.\n\n**Reorganization and writing:** Thank you for bringing this up and providing concrete suggestions. We will add transition statements and expand the introduction to clearly establish the challenges associated with our section on stability. We will also look at moving section 4.2 to a later part of the paper, although it is harder to promise this change without first doing so and making sure the flow of the paper still makes sense. \n\n**Does the batched updates alone improve performance?** Great question. We show this result for TD7 in Figure 25 and Table 13 in Appendix G, where “current policy” refers to TD7 trained in an identical fashion but we evaluate with the current policy rather than the checkpoint policy. The result only goes to 1M time steps however, so in the table below we present results for the full 5M time steps.\n\n|             | TD7 (checkpoints, checkpoint policy) | TD7 (checkpoints, current policy) | TD7 (no checkpoints) |\n|-------------|--------------------------------------|-----------------------------------|----------------------|\n| HalfCheetah | 18165 $\\pm$ 255                      | 18214 $\\pm$ 237                   | **18328 $\\pm$ 331**  |\n| Hopper      | **4075 $\\pm$ 225**                   | 3690 $\\pm$ 562                    | 3851 $\\pm$ 372       |\n| Walker2d    | **7397 $\\pm$ 454**                   | 7258 $\\pm$ 469                    | 6519 $\\pm$ 209       |\n| Ant         | 10133 $\\pm$ 966                      | 9807 $\\pm$ 730                    | **10388 $\\pm$ 1024** |\n| Humanoid    | **10281 $\\pm$ 588**                  | 10157 $\\pm$ 560                   | 9521 $\\pm$ 820       |\n\nAdditionally, we ran the suggested experiment for TD3. The results are presented in the table below. Full learning curves are included in the PDF of the general response (Figure 3). For both tables, the results are for 5M time steps and 10 seeds. $\\pm$ captures a 95\\% confidence interval. While there are some performance benefits for some environments, the results do not suggest that the batched updates are the main reason for the performance benefit. \n\n|             | TD3 (checkpoints, checkpoint policy) | TD3 (checkpoints, current policy) | TD3 (no checkpoints) |\n|-------------|--------------------------------------|-----------------------------------|----------------------|\n| HalfCheetah | 14075 $\\pm$ 1064                     | 13946 $\\pm$ 1022                  | **14337 $\\pm$ 1491** |\n| Hopper      | 3366 $\\pm$ 362                       | 3285 $\\pm$ 578                    | **3682 $\\pm$ 83**    |\n| Walker2d    | **5189 $\\pm$ 405**                   | 5126 $\\pm$ 386                    | 5078 $\\pm$ 343       |\n| Ant         | **6168 $\\pm$ 677**                   | 6149 $\\pm$ 713                    | 5589 $\\pm$ 758       |\n| Humanoid    | **5785 $\\pm$ 154**                   | 5534 $\\pm$ 136                    | 5433 $\\pm$ 245       |\n\n**Questions:**\n1. **Including the original input:** Correct! This is what we found during our testing (and shown in Figure 2). The input to TD7’s critic includes both the embeddings $z^{sa}$ and $z^s$ (latent state action) as well as the original state and action. The input to the actor includes the embedding $z^s$ and the original state.\n2. **Percent loss:** Percent loss refers to $\\frac{\\text{performance of modification} - \\text{performance of TD7}}{\\text{performance of TD7}} \\times 100$. TD7 uses all the default choices. The default choices would be 0 in all cases. We will include these details in the next version of the paper. \n3. **Early training with checkpoints:** Experimentally, we looked at this in Table 11 in Appendix F.  Your intuition is correct, in that the policy (and corresponding data) improves very rapidly early in learning so batching the training means the data will be more distant from the behavior of the current policy (distribution shift), and corresponds to a lower reward. \n4. **Relevance of data:** By relevance we mean how much the data corresponds to the current policy (how much distribution shift there is). In the standard training paradigm, after every training step, a new data point is collected. However, using checkpoints means that the training gets batched, so as the policy trains, before collecting the next batch, it may select actions which are not yet contained in the dataset. We will clarify this in the next version of the paper. Empirically we did not find this problematic for 20-50 episodes (after early learning stages), but presumably if this number was set to very large values (say, 1000), then the training paradigm would resemble offline RL and would likely negatively affect performance. "}, {"review_id": "ZA82LRVxB3", "reviewer": "Reviewer_xwBa", "summary": "The paper introduces an approach dubbed SALE for learning embeddings that model the interaction between state and action in low-level state environments. The authors extensively study the design space of these embeddings and integrate SALE into the TD3 algorithm along with 3 other components to form a new algorithm, TD7. \nThey perform an extensive empirical evaluation over the design space to discover the most significant contributors to final performance. The paper shows that TD7 outperforms existing continuous control algorithms on MuJoCo OpenAI gym tasks.", "questions": "See Weaknesses Section.", "limitations": "The authors adequately addressed the limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Originality**\nThe paper's originality lies in its comprehensive study of learning embeddings that model the interaction between state and action in low-level state environments. The experiments conducted provide evidence of the effectiveness of their approach in improving the performance of TD3.\n\n**Quality**\nSee Weaknesses Section.\n\n**Clarity**\nThe paper is well-written and organized, making it easy for readers to follow the authors' thought process and understand the methodology and results. \n\n**Significance**\nThe authors demonstrate that their proposed method, when integrated into TD3 to form what they name TD7, performs significantly better than TD3 for the MuJoCo environments. This work brings to light the importance of learning the interaction between state and action information to improve performance in RL.", "weaknesses": "While the paper presents a novel approach to learning embeddings for state and action information for RL, there are areas where it could be improved.\n\nFirstly, the evaluation of TD7 could be expanded to include more benchmarks. While the MuJoCo environments provide valuable data, including additional benchmarks such as Procgen or Brax could offer a more comprehensive understanding of the method's performance and any potential limitations.\n\nThese considerations, while not detracting from the originality and potential significance of the method, do highlight areas where further work could be beneficial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Is part of this review missing?**\n\nWe just wanted to double check that this review is complete. It currently reads as if some of the weaknesses were mistakenly deleted. If some other weaknesses did indeed get removed, let us know so we can improve the paper accordingly. Regardless, thank you for your time and the review. We address your comments below. \n\n**TD7 on other benchmarks:** To address this concern, we have collected results for TD7 on the DMC suite. As shown in the table below, TD7 consistently maintains its advantage over TD3 in DMC as well. In the attached PDF of the general response, we include the full learning curves for these environments (Figure 1). We hope this addresses your primary concern!\n\n|                         | TD7                 | TD3                |\n|-------------------------|---------------------|--------------------|\n| Ball in Cup Catch       | **982 $\\pm$ 1.8**   | 981 $\\pm$ 3.3      |\n| Cartpole Balance Sparse | **1000 $\\pm$ 0.0**  | **1000 $\\pm$ 0.0** |\n| Cartpole Swingup        | **875 $\\pm$ 5.8**   | 871 $\\pm$ 6.6      |\n| Cheetah Run             | **914 $\\pm$ 2.9**   | 866 $\\pm$ 23.5     |\n| Finger Spin             | **989 $\\pm$ 1.4**   | 975 $\\pm$ 14.1     |\n| Finger Turn Hard        | **742 $\\pm$ 136.8** | 552 $\\pm$ 80.7     |\n| Hopper Stand            | **774 $\\pm$ 177.4** | 480 $\\pm$ 180.8    |\n| Hopper Hop              | **175 $\\pm$ 98.7**  | 58 $\\pm$ 42.3      |\n| Humanoid Stand          | **504 $\\pm$ 336.5** | 7 $\\pm$ 0.4        |\n| Humanoid Walk           | **389 $\\pm$ 157.6** | 36 $\\pm$ 72.4      |\n| Walker Stand            | **991 $\\pm$ 1.5**   | 985 $\\pm$ 3.1      |\n| Walker Walk             | **981 $\\pm$ 2.7**   | 965 $\\pm$ 5.2      |\n| Walker Run              | **807 $\\pm$ 62.5**  | 631 $\\pm$ 74.7     |\n| Quadruped Walk          | **950 $\\pm$ 4.9**   | 750 $\\pm$ 168.9    |\n| Quadruped Run           | **684 $\\pm$ 87.5**  | 503 $\\pm$ 193.7    |\n| Mean                    | **784**             | 644                |\n| Median                  | **875**             | 750                |\n| IQM                     | **856**             | 733                |\n\nResults are for 1M time steps and 10 seeds. $\\pm$ captures a 95% confidence interval.  "}, {"review_id": "AWFlMuI5LL", "reviewer": "Reviewer_FeJr", "summary": "This paper proposes a new RL algorithm called TD7, which is based on TD3, and adopts additional techniques including (1) learning state-action representations, (2) LAP prioritized replay, (3) a behavior cloning term in the learning objective, and (4) checkpoints, where (3) and (4) are applied exclusively to the offline RL setting. They demonstrate the superior empirical performance of TD7 on the D4RL benchmark.", "questions": "1. Could you provide more evidence that shows SALE's applicability to other existing RL methods?", "limitations": "The limitations are not explicitly discussed in the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well written and easy to follow.\n2. Extensive experiments and ablation studies to evaluate the proposed methods.\n3. The proposed method is shown to have strong empirical performance on established benchmarks.", "weaknesses": "1. **Limited contribution**: While I appreciate TD7 as a competent algorithm for both online and offline RL, I'm not so sure about the contribution of SALE especially regarding its applicability to other existing baselines. It doesn't seem conclusive to me that the proposed embedding is generally superior for RL agents. From the ablation studies it seems many independent design choices could significantly affect the final performance, which may also be algorithm dependent. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and the feedback. \n\n**Combining SALE with other algorithms:** Thank you for bringing up this concern. In the table below we have included results with SALE applied to SAC.  Learning curves are included in the PDF of the general response (Figure 2). No modifications or hyperparameter changes were made to either SALE or SAC, we simply applied SALE directly. We hope this addresses your primary concern!\n  \n|             | SAC + SALE          | SAC                |\n|-------------|---------------------|--------------------|\n| HalfCheetah | **17316 $\\pm$ 236** | 15526 $\\pm$ 697    |\n| Hopper      | 2842 $\\pm$ 951      | **3167 $\\pm$ 485** |\n| Walker2d    | **6887 $\\pm$ 306**  | 5681 $\\pm$ 329     |\n| Ant         | **9288 $\\pm$ 462**  | 4615 $\\pm$ 2022    |\n| Humanoid    | **9062 $\\pm$ 335**  | 6555 $\\pm$ 279     |\n\nResults are for 5M time steps, and 10 seeds. $\\pm$ captures a 95% confidence interval. \n\nThis is further evidence that while there are many design choices to be made, we are presenting a set of design choices which are robust to algorithm and task. Regardless, it is our hope that by presenting an in-depth empirical study, readers can determine which choices are more important than others, and can start with a good default set. \n\n**Limitations:** As stated in the checklist, these are included in Appendix J. Sorry for the misunderstanding, space was a concern in the paper. "}]

Author rebuttal:
We would like to thank all the reviewers for their comments, suggestions for improvement, and interest in the paper.

Overall, there were two main comments that were repeated among the reviewers which were:

**Does TD7 work on other benchmarks?** To answer this, we ran TD7 on 15 new environments from the DeepMind Control suite and benchmarked it against TD3 (**Figure 1**). The final results at 1M time steps are shown below. 

|                         | TD7                 | TD3                |
|-------------------------|---------------------|--------------------|
| Ball in Cup Catch       | **982 $\pm$ 1.8**   | 981 $\pm$ 3.3      |
| Cartpole Balance Sparse | **1000 $\pm$ 0.0**  | **1000 $\pm$ 0.0** |
| Cartpole Swingup        | **875 $\pm$ 5.8**   | 871 $\pm$ 6.6      |
| Cheetah Run             | **914 $\pm$ 2.9**   | 866 $\pm$ 23.5     |
| Finger Spin             | **989 $\pm$ 1.4**   | 975 $\pm$ 14.1     |
| Finger Turn Hard        | **742 $\pm$ 136.8** | 552 $\pm$ 80.7     |
| Hopper Stand            | **774 $\pm$ 177.4** | 480 $\pm$ 180.8    |
| Hopper Hop              | **175 $\pm$ 98.7**  | 58 $\pm$ 42.3      |
| Humanoid Stand          | **504 $\pm$ 336.5** | 7 $\pm$ 0.4        |
| Humanoid Walk           | **389 $\pm$ 157.6** | 36 $\pm$ 72.4      |
| Walker Stand            | **991 $\pm$ 1.5**   | 985 $\pm$ 3.1      |
| Walker Walk             | **981 $\pm$ 2.7**   | 965 $\pm$ 5.2      |
| Walker Run              | **807 $\pm$ 62.5**  | 631 $\pm$ 74.7     |
| Quadruped Walk          | **950 $\pm$ 4.9**   | 750 $\pm$ 168.9    |
| Quadruped Run           | **684 $\pm$ 87.5**  | 503 $\pm$ 193.7    |
| Mean                    | **784**             | 644                |
| Median                  | **875**             | 750                |
| IQM                     | **856**             | 733                |

Results are for 1M time steps and 10 seeds. $\pm$ captures a 95% confidence interval. Our results conclusively demonstrate that the improvements we have proposed to TD3 allow TD7 to consistently outperform it on this new benchmark of varied tasks. Between these new results, our online MuJoCo results, and our offline D4RL results, we hope that this convinces the reviewers that TD7 (and SALE) provides a consistent and meaningful improvement over TD3 in a wide range of tasks and settings. 

**Does SALE work with other RL algorithms?** To answer this, we applied SALE to SAC, making no hyperparameter changes or additional modifications to SALE or SAC (**Figure 2**). The final results at 5M time steps are shown below. 

|             | SAC + SALE          | SAC                |
|-------------|---------------------|--------------------|
| HalfCheetah | **17316 $\pm$ 236** | 15526 $\pm$ 697    |
| Hopper      | 2842 $\pm$ 951      | **3167 $\pm$ 485** |
| Walker2d    | **6887 $\pm$ 306**  | 5681 $\pm$ 329     |
| Ant         | **9288 $\pm$ 462**  | 4615 $\pm$ 2022    |
| Humanoid    | **9062 $\pm$ 335**  | 6555 $\pm$ 279     |

Results are for 5M time steps and 10 seeds. $\pm$ captures a 95% confidence interval. On 4 of the 5 tasks, SAC+SALE significantly outperforms vanilla SAC. Since most off-policy RL algorithms are based on either TD3 or SAC, we hope that this demonstrates that SALE can be applied more widely than just TD3. 

**Additional experiments on checkpoints:** Some of the reviewers had some additional thoughts on checkpoints that we thought could be best answered with additional visualization. 

In **Figure 3** we looked to answer a question posed by reviewer YfJA, which was: how does the change to batched training affect performance? This figure shows the performance of the checkpoint policy taken from TD3 trained with checkpoints, the current policy obtained simultaneously from TD3 trained with checkpoints, and vanilla TD3. The result shows that batching training can slightly improve performance in some cases (Ant), but largely has limited impact on performance. 

This raises another question (which was asked by reviewer Casd), which is: how does using checkpoints contribute? In **Figure 4** we show the performance of 5 individual trials (10 trials is too visually cluttered), with checkpoints and without checkpoints. What we see is that while the average performance of the current policy is similar to the checkpoint policy, the stability of the checkpoint policy is much higher. For those who have not examined individual trials before, the instability of the current policy of TD7 may be surprising, however this is consistent with other algorithms as well. In Appendix F, Figure 18, we plot the learning curve of a single seed of TD3. 

This stability benefit is perhaps better understood numerically. In the following table we present the average standard deviation over the last 20 evaluations (corresponding to the final 100k time steps of training). In brackets we include the [min performance, max performance] over those evaluations. 

|             | TD7    (checkpoints) | TD7 (no checkpoints) | TD3 (checkpoints)  | TD3 (no checkpoints) |
|-------------|----------------------|----------------------|--------------------|----------------------|
| HalfCheetah | 14, [18149, 18202]   | 227, [17700, 18671]  | 98, [13745, 14118] | 319, [13404, 14665]  |
| Hopper      | 28, [4025, 4132]     | 573, [2200, 4224]    | 121, [3109, 3571]  | 344, [2412, 3727]    |
| Walker2d    | 83, [7108, 7408]     | 607, [4315, 6685]    | 26, [5124, 5213]   | 538, [3159, 5270]    |
| Ant         | 440, [8763, 10223]   | 609, [8748, 10676]   | 196, [5568, 6280]  | 716, [3287, 5709]    |
| Humanoid    | 422, [9414, 10808]   | 1251, [5213, 10501]  | 149, [5335, 5882]  | 158, [5045, 5658]    |

What we see is that checkpoints significantly reduce the variability of performance between evaluations, and improve the worst case performance. 

We hope this addresses the concerns of the reviewers and are happy to discuss further. 


Meta-review: The paper proposes a method for learning state-action representations for continuous control off-policy deep RL and integrates it into the TD3 algorithm to instantiate the TD7 algorithm. 

Main Strengths: Reviewers appreciated novelty in the under-explored area of state-action representation learning. Representatio learning methods aside from state representation learning are novel and interesting to the RL community. The paper has strong empirical evaluation including ablation studies that detail the utility of these techniques across a variety of benchmarks and algorithm settings. 

Main Concerns: The main weakness noted by the reviewers is that a variety of techniques are introduced in the method and intuition and theory for these is sometimes limited.

Overall, the study of the interplay between state and action in an agents representation is interesting and valuable to the RL community and so I recommend acceptance.
Decision: Accept (poster)