Title: For SALE: State-Action Representation Learning for Deep Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Deep reinforcement learning

representation learning
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

Abstract

In the field of reinforcement learning (RL), representation learning is a proven tool for complex image-based tasks, but is often overlooked for environments with lowlevel 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.

1 Introduction

Reinforcement learning (RL) is notoriously sample inefficient, particularly when compared to more straightforward paradigms in machine learning, such as supervised learning. One possible explanation is the usage of the Bellman equation in most 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].

A near-universal solution to sample inefficiency in 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 as the learning horizon and dynamic programming. While feature learning of this type has found some success in the 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 of representation learning to low-level states is much less common. At first glance, it may seem unnecessary to learn a representation over an already-compact state vector. However, we argue that the difficulty of a 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), a 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

when learning a state-action representation online. Most importantly, we observe the surprising effect of extrapolation error [Fujimoto et al., 2019] when significantly expanding the action-dependent input and introduce a simple clipping technique to mitigate it.

Design study. Learning to model environment dynamics in latent space is a common approach for feature learning which has been widely considered [Watter et al., 2015, Ha and Schmidhuber, 2018, Hafner et 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, we 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 RL 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 aforementioned 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 1 .

2 Related Work

Representation learning. Representation learning has several related interpretations in RL. Historically, 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 higherdimensional spaces, the notion of true abstraction has been replaced with compression , where the intent is to embed the observation space (such as images) into a smaller manageable latent vector [Watter 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, 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 a rich history [Dayan, 1993, Littman and Sutton, 2001], 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, 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, Rezaei-Shoshtari 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 a longstanding challenge, indicated by numerous empirical studies that highlight practical concerns associated with deep RL methods [Henderson 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 et al., 2018] with RL [Khadka and Tumer, 2018, Pourchot and Sigaud, 2018], where the checkpoint resembles the fittest individual and the mutation is defined exclusively by the underlying RL algorithm.

3 Background

In Reinforcement learning (RL) problems are framed as a Markov decision process (MDP). An MDP is a 5-tuple ( S , A , R , p , γ ) with state space S , action space A , reward function R , dynamics model p , and discount factor γ , where the objective is to find a policy π : S → A , a mapping from state s ∈ S to action a ∈ A , which maximizes the return ∑ ∞ t =1 γ t -1 r t , the discounted sum of rewards r obtained when following the policy. RL algorithms commonly use a value function Q π ( s, a ) := E [∑ ∞ t =1 γ t -1 r t | s 0 = s, a 0 = a ] , which models the expected return, starting from an initial state s 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 ( z sa , z s ) which capture relevant structure in the observation space, as well as the transition dynamics of the environment. To do so, SALE utilizes a pair of encoders ( f, g ) where f ( s ) encodes the state s into the state embedding z s and g ( z s , a ) jointly encodes both state s and action a into the state-action embedding z sa :

The embeddings are split into state and state-action components so that the encoders can be trained with a dynamics prediction loss that solely relies on the next state s ′ , independent of the next action or current policy. As a result, the encoders are jointly trained using the mean squared error (MSE) between the state-action embedding z sa and the embedding of the next state z s ′ :

where |·| × denotes the stop-gradient operation. The embeddings are designed to model the underlying structure of the environment. However, they may not 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:

The encoders ( f, g ) are trained online and concurrently with the RL agent (updated at the same frequency as the value function and policy), but are decoupled (gradients from the value function and policy are not propagated to ( 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 not to serve as a 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 a redundant representation [Gelada et al., 2019]. To combat this risk, we introduce AvgL1Norm, a normalization layer that divides the input vector by its average absolute value in each dimension, thus keeping the relative scale of the embedding constant throughout learning. Let x i be the i -th dimension of an N -dimensional vector x , then

g

t

(

z

s

t

, a

)

Linear

(

s, a

)

f

t

(

s

)

g

t

(

z

s

t

, a

)

AvgL1Norm

(

ϕ

sa

)

AvgL1Norm

(

z

s

t

)

z

sa

t

s

a

s

a

Q

t

+1

(

z

sa

t

,z

s

t

,ϕ

sa

)

End-to-end

Decoupled

(a) Value function

Q

with SALE

Linear

(

s

)

π

t

+1

(

z

s

t

, ϕ

s

)

f

t

(

s

)

AvgL1Norm

(

ϕ

s

)

AvgL1Norm

(

z

s

t

)

s

s

End-to-end

Decoupled

(b) Policy

π

with SALE

AvgL1Norm is applied to the state embedding z s . Similar to the normalized loss functions used by SPR [Schwarzer et al., 2020] and BYOL [Grill et al., 2020], AvgL1Norm protects from monotonic growth, but also keeps the scale of the downstream input constant without relying on updating statistics (e.g. BatchNorm [Ioffe and Szegedy, 2015]). This is important for our approach as the embeddings are trained independently from the value function and policy. AvgL1Norm is not applied to the state-action embedding z sa , as it is trained to match the normalized next state embedding z s ′ .

We also apply AvgL1Norm to the state and action inputs (following a 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:

Unlike the embeddings ( z s , z sa ) , these linear layers are learned end-to-end, and can consequently be viewed as an addition to the architecture of the value function or 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 t +1 , π t +1 ) uses embeddings ( z sa t , z s t ) from the encoders ( f t , g t ) at the previous iteration t . The value function and policy are thus updated by:

The current value function Q t +1 is also trained with respect to the previous value function Q t , known as a target network [Mnih et al., 2015]. The current embeddings z s t +1 and z sa t +1 are trained with Equation 2, using a target z s ′ t +1 (hence, without a target network). Every n steps the iteration is incremented and all target networks are updated simultaneously:

Clipped Values. Extrapolation error is the tendency for deep value functions to extrapolate to unrealistic values on 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, as the lack of feedback on overestimated values can result in divergence.

Surprisingly, we observe a similar phenomenon in online RL, when increasing the number of dimensions in the state-action input to the value function, as illustrated in Figure 2. Our hypothesis is that the state-action embedding z sa expands the action input and makes the value function more

0

1

2

3

4

5

-3

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Un-clipped

performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Un-clipped

value estimate by seed

1

2

No

z

sa

0

1

2

3

4

5

1

2

Time steps (1M)

Small

ϕ

sa

1

2

No

z

sa

and small

ϕ

sa

0

1

2

3

4

5

1

2

Time steps (1M)

Clipped

Seeds (1-10)

Buffer: 1M

Buffer:

1M

100k

Buffer:

1M

100k

Buffer:

1M

100k

Buffer:

1M

100k

likely to over-extrapolate on unknown actions. We show in Figure 2 that the dimension size of z sa 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:

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 is 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 an 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 3 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 embedding z sa and a learning target of the next state embedding z s ′ (Equation 2). We test several alternate learning targets. OFENet uses the next state s ′ as the target [Ota et al., 2020] while SPR [Schwarzer et al., 2020] uses the embedding z s ′ target from a target network obtained with an exponential moving average with weight 0 . 01 . Drawing inspiration from Bisimulation metrics [Ferns et al., 2011], DeepMDP [Gelada et al., 2019] use an objective that considers both the next state embedding z s ′ and the

s

′

z

s

′

target

z

s

′

and

r

z

s

′

a

′

20

15

10

5

0

5.0

(2.1)

2.6

(3.6)

0.9

(4.0)

21.0

(3.1)

Mean Percent Loss

Learning Target

Default:

z

s

′

Q

(

z

s

,s,a

)

Q

(

z

sa

,s,a

)

Q

(

z

sa

,z

s

)

Q

(

z

sa

)

Q

(

s,a

)

π

(

s

)

π

(

z

s

)

No fixed

8.0

(4.4)

4.8

(4.3)

10.1

(6.5)

13.8

(4.6)

20.3

(5.8)

5.7

(6.5)

6.2

(4.7)

7.6

(2.8)

Default:

Q

(

z

sa

, z

s

, s, a

)

and

π

(

z

s

, s

)

z

s

only

None

Norm

z

sa

BatchNorm

LayerNorm

Cosine loss

9.5

(7.9)

8.3

(5.1)

4.7

(4.7)

12.4

(5.4)

16.9

(10.8)

10.9

(4.9)

Normalization

Default: AvgL1Norm

0

.

1

1

10

23.6

(7.6)

22.6

(7.3)

24.8

(3.7)

End-to-End

Default: Decoupled

reward r . We test including a prediction loss on the reward by having the encoder g output both z sa and r pred where r pred is trained with the MSE to the reward r . Finally, we test the next state-action embedding z s ′ a ′ as the target, where the action a ′ is sampled from the target policy.

⇒ All learning targets based on the next state s ′ perform similarly, although using the embedding z s ′ further improves the performance. On the contrary, the next state-action embedding z s ′ a ′ performs much worse as a target, highlighting that signal based on the non-stationary policy can harm learning. Including the reward as a signal has little impact on performance.

Network input. In our approach, the learned embeddings ( z sa , z s ) are appended to the state and action input to the value function Q ( z sa , z s , s, a ) and policy π ( z s , s ) (Equation 3). We attempt different combinations of input to both networks. We also evaluate replacing the fixed embeddings (Equations 6 & 7), with the non-static current embeddings ( z sa t +1 , z s t +1 ) .

⇒ The added features have a greater impact on the value function than the policy, but are beneficial for both networks. All components of the value function input ( z sa , z s , s, a ) , are necessary to achieve the highest performance. While the state-action embedding z sa is a 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 z sa as input leads to poor results, but combining it with the original input ( s, a ) significantly improves performance.

Normalization. TD7 uses AvgL1Norm (Equation 4) to normalize the scale of the state embedding z s , as well as on the state-action input ( s, a ) , following a linear layer (Equation 5). We attempt removing AvgL1Norm on ( s, a ) , removing it entirely, and adding it to the state-action embedding z sa . We additionally test swapping AvgL1Norm for BatchNorm [Ioffe 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 2) by the cosine loss from Schwarzer et al. [2020].

⇒ The usage of AvgL1Norm is beneficial and related alternate approaches do not achieve the same performance.

End-to-end. Embeddings can be trained independently or 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 ( z sa , z s ) , 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.

5 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 a 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 a 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 use of checkpoints, to preserve evaluation performance, irrespective of the quality of the current learned policy.

Acheckpoint is a snapshot of the parameters of a model, captured at a specific time during training. In supervised learning, checkpoints are often used to recall a previous set of high-performing parameters based on validation error, and maintain a consistent performance across evaluations [Vaswani et al., 2017, Kenton and Toutanova, 2019]. Yet this technique is 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 to train after each time step (typically at a one-to-one ratio: one gradient step for one data point). However, this means that the policy changes throughout each episode, making it hard to evaluate the performance. Similar to 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 a 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 over 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, it is only sensible to use a 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, we 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], a behavior cloning term for offline RL [Fujimoto and Gu, 2021], SALE (Section 4.1), and policy checkpoints (Section 5.1).

LAP. Gathered experience is stored in a replay buffer [Lin, 1992] and sampled according to LAP [Fujimoto et al., 2020], a prioritized replay buffer D [Schaul et al., 2016] where a transi-

Algorithm 1 Online TD7

tion tuple i := ( s, a, r, s ′ ) is sampled with probability

where y is the learning target. The amount of prioritization used is controlled by a hyperparameter α . Furthermore, the value function loss uses the Huber loss [Huber et al., 1964], rather than the MSE.

Offline RL. To make TD7 amenable to the offline RL setting, we add a behavior cloning loss to the policy update [Silver et al., 2014], inspired by TD3+BC [Fujimoto and Gu, 2021]:

The same loss function is used for both offline and online RL, where λ = 0 for the online setting. | · | × denotes the stop-gradient operation. Unlike TD3+BC, we do not normalize the state vectors. Checkpoints are not used in the offline setting, as there is no interaction with the environment.

Both the value function and policy use the SALE embeddings as input, which we omit from the equations above for simplicity. Pseudocode for TD7 is described in Algorithm 1.

6 Results

In this section, we evaluate the performance of TD7 in both the online and offline regimes. A 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 OpenAI gym [Brockman et al., 2016], we benchmark TD7 against TD3 [Fujimoto et al., 2018], SAC [Haarnoja et al., 2018], TQC [Kuznetsov et al., 2020], and TD3+OFE [Ota et al., 2020] on the MuJoCo environments [Todorov et al., 2012]. SAC and TD3+OFE results are from re-implementations based on author descriptions [Haarnoja et al., 2018, Ota et al., 2020]. TD3 and TQC results use author-provided code [Fujimoto et al., 2018, Kuznetsov et al., 2020], with a consistent evaluation protocol for all methods. Learning curves are displayed in Figure 4 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 et 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

0

1

2

3

4

5

0

4

8

12

16

20

Total Reward (1k)

HalfCheetah

0

1

2

3

4

5

0

1

2

3

4

Hopper

0

1

2

3

4

5

0

2

4

6

8

Walker2d

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Ant

0

1

2

3

4

5

0

2

4

6

8

10

Time steps (1M)

Humanoid

TD7

TD3

SAC

TQC

TD3+OFE

Figure

4:

Learning

curves

on

the MuJoCo benchmark. Results

are averaged over 10 seeds. The

shaded area captures a 95% confi-

dence interval around the average

performance.

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 a consistent evaluation protocol. Final performance is reported in Table 2.

TD7 outperforms all baselines. Since TD7 and TD3+BC employ the same approach to offline RL, the significant performance gap highlights the effectiveness of SALE in the offline setting.

Ablation study . In Figure 5 we report the results of an ablation study over the components of TD7 (SALE, checkpoints, LAP). The interaction between components is explored further in Appendix G.

Run time. To understand the computational cost of using SALE and the TD7 algorithm, we 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.

7 Conclusion

Representation learning has been typically reserved for image-based tasks, where the observations are large and unstructured. However, by learning embeddings which consider the interaction between state and action, we make representation learning more broadly applicable to low-level states. We introduce SALE, a method for learning state-action embeddings by considering a latent space dynamics model. Through an extensive empirical evaluation, we investigate various design choices in SALE.

We highlight the risk of extrapolation error [Fujimoto et al., 2019] due to 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.

While both SALE and policy checkpoints are generalpurpose techniques that can be included with most RL methods, we combine them with TD3 and several other recent improvements [Fujimoto et al., 2020, Fujimoto and Gu, 2021] to introduce the TD7 algorithm. We find our TD7 algorithm is able to match the performance of expensive offline algorithms and significantly outperform the state-of-the-art continuous control methods in both final performance and early learning.

Acknowledgments and Disclosure of Funding

This research was enabled in part by support provided by Calcul Québec and the Digital Research Alliance of Canada.

0

1

2

3

4

5

0

0.2

0.4

0.6

0.8

1.0

Time steps (1M)

Normalized Performance

Ablation

TD7

TD3

No SALE

No checkpoints

No LAP

TD3

SAC

TQC

TD3+OFE

TD7

1

2

3

4

47m

1h 28m

3h 50m

3h 14m

1h 50m

Hours

Run Time

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 preprint arXiv:1607.06450 , 2016.

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

Edoardo Cetin, Philip J Ball, Stephen Roberts, and Oya Celiktutan. Stabilizing 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.

Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv:1511.07289 , 2015.

Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Computation , 5(4):613-624, 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 on Learning Representations , 2019.

Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing , 40(6):1662-1714, 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. IEEE, 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-policy evaluation. In International Conference on Learning Representations , 2021.

Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. In Thirty-Fifth Conference on Neural Information Processing Systems , 2021.

Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actorcritic 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 nonuniform 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 i trust you, bellman? The Bellman error is a poor replacement for value error. In International Conference on 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 , 33:21271-21284, 2020.

David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122 , 2018.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning , 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 machine learning , pages 2555-2565. PMLR, 2019.

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 J Huber et al. Robust estimation of a location parameter. The annals of mathematical statistics , 35(1):73-101, 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-456. pmlr, 2015.

Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on 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 arXiv: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-parameterization inhibits data-efficient deep reinforcement learning. In International Conference on Learning Representations , 2020a.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:1179-1191, 2020b.

Arsenii Kuznetsov, Pavel Shvechikov, Alexander Grishin, and Dmitry Vetrov. Controlling overestimation 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 P Vetrov. Automating control of overestimation bias for continuous reinforcement learning. ArXiv abs/2110.13523 , 2021.

Romain Laroche, Paul Trichelair, and Remi Tachet Des Combes. Safe policy improvement with baseline bootstrapping. In 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 J Walsh, and Michael L Littman. Towards a unified theory of state abstraction for mdps. In AI&M , 2006.

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

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 to reinforcement learning. arXiv preprint arXiv:1803.07055 , 2018.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature , 518(7540):529-533, 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 arXiv:1807.03748 , 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.

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 8024-8035, 2019.

Aloïs Pourchot and Olivier Sigaud. Cem-rl: Combining evolutionary and gradient-based methods for policy search. arXiv preprint arXiv:1810.01222 , 2018.

Balaraman Ravindran. An algebraic approach to abstraction in reinforcement learning . University of Massachusetts Amherst, 2004.

Sahand Rezaei-Shoshtari, Rosie Zhao, Prakash Panangaden, David Meger, and Doina Precup. Continuous 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 preprint arXiv: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.06347 , 2017.

Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive 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: a scalable real-time architecture for learning knowledge from unsupervised sensorimotor interaction. In The 10th International Conference on Autonomous Agents and Multiagent Systems-Volume 2 , pages 761-768, 2011.

Guy Tennenholtz and Shie Mannor. The natural language of actions. In International Conference on Machine Learning , pages 6196-6205. PMLR, 2019.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS) , pages 5026-5033. IEEE, 2012.

Elise van der Pol, Thomas Kipf, Frans A Oliehoek, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. In Proceedings of 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 in reinforcement learning. Advances in Neural Information Processing Systems , 33:4199-4210, 2020b.

Herke Van Hoof, Nutan Chen, Maximilian Karl, Patrick van der Smagt, and Jan Peters. Stable reinforcement learning with autoencoders for tactile and visual data. In 2016 IEEE/RSJ international conference on intelligent robots and systems (IROS) , pages 3928-3934. IEEE, 2016.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is 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 gradient-following algorithms for connectionist reinforcement learning. Machine learning , 8(3-4):229-256, 1992.

Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, 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 on Learning Representations , 2022.

Amy Zhang, Harsh Satija, and Joelle Pineau. Decoupling dynamics and reward for transfer learning. arXiv preprint arXiv:1804.10689 , 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.

Appendix

Table of Contents

A TD7 Additional Details

A.1 Algorithm

TD7 (TD3 + 4 additions) has several networks and sub-components:

Two value functions ( Q t +1 , 1 , Q t +1 , 2 ) .

Two target value functions ( Q t, 1 , Q t, 2 ) .

A policy network π t +1 .

A target policy network π t .

An encoder, with sub-components ( f t +1 , g t +1 ) .

A fixed encoder, with sub-components ( f t , g t ) .

A target fixed encoder with sub-components ( f t -1 , g t -1 ) .

A checkpoint policy π c and checkpoint encoder f c ( g is not needed).

Encoder: The encoder is composed of two sub-networks ( f t +1 ( s ) , g t +1 )( z s , a ) , where each network outputs an embedding:

At each training step, the encoder is updated with the following loss:

where | · | × is the stop-gradient operation.

Value function: TD7 uses a pair of value functions (as motivated by TD3 [Fujimoto et al., 2018]) ( Q t +1 , 1 , Q t +1 , 2 ) , each taking input [ z s ′ a ′ t -1 , z s ′ t -1 , s ′ , a ′ ] . At each training step, both value functions are updated with the following loss:

Taking the minimum of the value functions is from TD3's Clipped Double Q-learning (CDQ) [Fujimoto et al., 2018]). The use of Huber loss [Huber et al., 1964] is in accordance to TD3 with the Loss-Adjusted Prioritized (LAP) experience replay [Fujimoto et al., 2020]. a ′ is sampled and clipped in the same manner as TD3 [Fujimoto et al., 2018]. The same embeddings ( z sa t , z s t ) are used for each value function. Q min and Q max are updated at each time step:

where target is defined by Equation 16.

Policy: TD7 uses a single policy network which takes input [ z s , s ] . On every second training step (according to TD3's delayed policy updates [Fujimoto et al., 2018]) the policy π t +1 is updated with the following loss:

The policy loss is the deterministic policy gradient (DPG) [Silver et al., 2014] with a behavior cloning term to regularize [Fujimoto and Gu, 2021]. λ is set to 0 for online RL.

After every target_update_frequency (250) training steps, the iteration is updated and each target (and fixed) network copies the network of the higher iteration:

The checkpoint policy and checkpoint encoder are only used at test time (see Appendix F).

LAP: Gathered experience is stored in a replay buffer [Lin, 1992] and sampled according to LAP [Fujimoto et al., 2020], a prioritized replay buffer D [Schaul et al., 2016] where a transition tuple i := ( s, a, r, s ′ ) is sampled with probability

where target is defined by Equation 16. As suggested by Fujimoto et al. [2020], | δ ( i ) | is defined by the maximum absolute error of both value functions. The amount of prioritization used is controlled by a hyperparameter α . New transitions are assigned the maximum priority of any sample in the replay buffer.

We outline the train function of TD7 in Algorithm 2. There is no difference between the train function of online TD7 and offline TD7 other than the value of λ , which is 0 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).

4: Update ( Q min , Q max ) (Equations 20 & 21).

5: if i mod policy_update_frequency = 0 then

6:

Train policy (Equation 22).

7: if i mod target_update_frequency = 0 then

8: Update target networks (Equation 26).

A.2 Hyperparameters

The action space is assumed to be in the range [ -1 , 1] (and is normalized if otherwise). Besides a few exceptions mentioned below, hyperparameters (and architecture) are taken directly from TD3 https://github.com/sfujim/TD3 .

Besides algorithmic difference from TD3, there are three implementation-level changes:

Rather than only using Q 1 , both value functions are used when updating the policy (Equation 22).

The value function uses ELU activation functions [Clevert et al., 2015] rather than ReLU activation functions.

The target network is updated periodically (every 250 time steps) rather than using an exponential moving average at every time step. This change is necessary due to the use of fixed encoders.

Weevaluate the importance of these changes in our ablation study (Appendix G). Network architecture details are described in Pseudocode 1.

Pseudocode 1. TD7 Network Details

Variables:

zs_dim = 256

Value Q Network:

▷ TD7 uses two value networks each with the same network and forward pass.

Value Q

Policy π Network:

256) Policy π Forward Pass: AvgL1Norm(l0(input)) x])

State Encoder f Network:

State Encoder f Forward Pass:

State-Action Encoder g Network:

State-Action Encoder g Forward Pass:

B Experimental Details

Environment. Our experimental evaluation is based on the MuJoCo simulator [Todorov et al., 2012] with tasks defined by OpenAI gym [Brockman et al., 2016] using the v4 environments. No modification are made to the state, action, or reward space.

Terminal transitions. All methods use a discount factor γ = 0 . 99 for non-terminal transitions and γ = 0 for terminal transitions. The final transition from an episode which ends due to a time limit is not considered terminal.

Exploration. To fill the replay buffer before training, all methods initially collect data by following a uniformly random policy for the first 25 k 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 a 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. Evaluations 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 to 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 in figures and the ± term in tables refers to a 97 . 5% confidence interval. Given that methods are evaluated over 10 seeds, this confidence interval CI is computed by

where σ is the sample standard deviation with Bessel's correction at the corresponding evaluation. Unless stated otherwise, all curves are smoothed uniformly over a window of 10 evaluations.

Offline Experiments. Offline results are based on 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 on 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 et al., 2016]

MuJoCo 2.3.3 [Todorov et al., 2012]

C Baselines

C.1 TD3 & TD3+OFE Hyperparameters

Our TD3 baseline [Fujimoto et al., 2018] uses the author implementation https://github.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 , 1] .

Environment-specific hyperparameters. OFE uses a variable number of encoder layers depending on the environment. However, the authors also provide an offline approach for hyperparameter tuning based on the representation loss. Our choice of layers per environment is the outcome of 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 are listed in Table 4. Network architecture details of TD3 are described in Pseudocode 2. Network architecture details of TD3+OFE are described in Pseudocode 3 and Pseudocode 4.

Pseudocode 2. TD3 Network Details

Value Q Network:

▷ TD3 uses two value networks each with the same network and forward pass.

Policy π Network:

Policy π Forward Pass:

Pseudocode 3. TD3+OFE Network Details (TD3)

Variables:

Value Q Network:

▷ TD3+OFE uses two value networks each with the same network and forward pass.

Value Q Forward Pass: ReLU(l1(input))

Policy π Network:

l1 = Linear(state_dim + zs_dim, 256) l2 = Linear(256, 256) l3 = Linear(256, action_dim) Policy π Forward Pass: input = zs x = ReLU(l1(input)) x = ReLU(l2(x)) action = tanh(l3(x))

Pseudocode 4. TD3+OFE Network Details (OFE)

C.2 SAC & TQC Hyperparameters

Our SAC baseline [Haarnoja et al., 2018] is based on our TD3 implementation, keeping hyperparameters constant when possible. Remaining details are based on the author implementation https://github.com/haarnoja/sac . Our TQC baseline [Kuznetsov et al., 2020] uses the author PyTorch implementation https://github.com/SamsungLabs/tqc_pytorch (with evaluation code kept consistent for all methods). For hyperparameters, the action space is assumed to be in the range [ -1 , 1] .

TQC is based on SAC and uses similar hyperparameters. One exception is the use of an ϵ offset. As suggested in the appendix of Haarnoja et al. [2018], the log probability of the Tanh Normal distribution is calculated as follows:

where u is the pre-activation value of the action, sampled from N , where N is distributed according to N ( µ, σ 2 ) , from the outputs µ and log σ of the actor network. Kuznetsov et al. [2020] use an alternate calculation of the log of the Tanh Normal distribution which eliminates the need for an ϵ offset:

Environment-specific hyperparameters. TQC varies the number of dropped quantiles depending on the environment. While it is possible to select this quantity based on 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.

Pseudocode 5. SAC Network Details

Value Q Network:

▷ SAC uses two value networks each with the same network and forward pass.

l1 = Linear(state_dim + action_dim, 256) l2 = Linear(256, 256) l3 = Linear(256, 1) Value Q Forward Pass: concatenate([state, action])

Policy π Network:

Policy π Forward Pass:

Pseudocode 6. TQC Network Details

Value Q Network: ▷ TQC uses five value networks each with the same network and forward pass. l1 = Linear(state_dim + action_dim, 512) l2 = Linear(512, 512) l3 = Linear(512, 512) l4 = Linear(512, 25) Value Q Forward Pass: input = concatenate([state, action]) x = ReLU(l1(input)) x = ReLU(l2(x)) x = ReLU(l3(x)) value = l4(x) Policy π Network: l1 = Linear(state_dim, 256) l2 = Linear(256, 256) l3 = Linear(256, action_dim * 2) Policy π Forward Pass: input = state x = ReLU(l1(input)) x = ReLU(l2(x)) mean, log_std = l3(x) x = Normal(mean, exp(log_std.clip(-20, 2))).sample() action = tanh(x)

C.3 Offline RL Baselines

We use 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://github.com/aviralkumar2907/CQL commit d67dbe9

Setting the Lagrange threshold below 0 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 to match the evaluation protocol used by TD7.

TD3+BC [Fujimoto and Gu, 2021]. https://github.com/sfujim/TD3_BC commit 8791ad7

python main.py --env=ENV

IQL [Kostrikov et al., 2021]. https://github.com/ikostrikov/implicit_q_learning commit 09d7002

python train_offline.py --env_name=ENV --config=configs/mujoco_config.py

X -QL [Garg et al., 2023]. https://github.com/Div99/XQL commit dff09af python train_offline.py --env_name=ENV --config=configs/mujoco_config.py \ --max_clip=7 --double=True --temp=2 --batch_size=256

This setting corresponds to the X -QL-C version of the method where hyperparameters are fixed across datasets.

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.

D.1 Learning Target

The encoder ( f t +1 , g t +1 ) in TD7 uses the following update:

In this section, we vary the learning target (originally z s ′ t +1 ).

(Figure 7) Target 1: s ′ . Sets the learning target to the next state s ′ , making the encoder loss:

This target is inspired by OFENet [Ota et al., 2020].

(Figure 7) Target 2: z s ′ target . Sets the learning target to the embedding of the next state z s ′ target taken from a slow-moving target network. This makes the encoder loss:

z s ′ target is taken from a target network f target where f t +1 has parameters θ and f target has parameters θ ′ , and θ ′ is updated at each time step by:

This target is inspired by SPR [Schwarzer et al., 2020].

(Figure 7) Target 3: z s and r . Sets the learning target to the embedding of the next state z s t +1 (which is the same as in 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 z sa t +1 and the second r pred t +1 , making the encoder loss:

This target is inspired by DeepMDP [Gelada et al., 2019].

(Figure 8) Target 4: z s ′ a ′ . Sets the learning target to the next state-action embedding z s ′ a ′ t +1 , where the action a ′ 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:

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

s

′

z

s

′

target

z

s

′

and

r

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

z

s

′

a

′

D.2 Network Input

SALE modifies the value function and policy to take the embeddings as additional input:

In this section, we vary the input of the value function and policy. In all cases, changes are made independently (i.e. we fix the policy to 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 z sa . We remove z sa from the value function input. This makes the input: Q t +1 ( z s t , s, a ) .

(Figure 9) Input 2: Q , remove z s . We remove z s from the value function input. This makes the input: Q t +1 ( z sa t , s, a ) .

(Figure 9) Input 3: Q , remove s, a . We remove s, a from the value function input. This makes the input: Q t +1 ( z sa t , z s ) .

(Figure 10) Input 4: Q , z sa only. We only use z sa in the value function input. This makes the input: Q t +1 ( z sa t ) .

(Figure 10) Input 5: Q , s, a only. We only use s, a in the value function input. This makes the input: Q t +1 ( s, a ) . As for all cases, the policy still takes in the state embedding z s .

(Figure 11) Input 6: π , s only. We remove z s from the policy input, using only the state s as input. This makes the input: π t +1 ( s ) . Note that the value function still uses the embeddings as input.

(Figure 11) Input 7: π , z s only. We remove s from the policy input, using only the embedding z s as input. This makes the input: π t +1 ( z s t ) .

(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: Q t +1 ( z sa t +1 , z s t +1 , s, a ) and π t +1 ( z s t +1 , s ) .

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

Q

, remove

z

sa

Q

, remove

z

s

Q

, remove

s, a

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

Q

,

z

sa

only

Q

,

s, a

only

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

π

,

s

only

π

,

z

s

only

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

No fixed embeddings

D.3 Normalization

SALE normalizes embeddings and input through our proposed AvgL1Norm:

AvgL1Norm is embedded into the architecture used by SALE (see Pseudocode 1), applied to the output of the state encoder f to give the state embedding z s = AvgL1Norm ( 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:

where ϕ sa is the output from a linear layer on the state-action input, Linear ( s, a ) , and ϕ s 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 is used.

(Figure 13) Normalization 1: No normalization on ϕ . Normalization is not applied to either ϕ . This makes the input to the value function and policy:

Q

(

z

sa

,

AvgL1Norm

(

f

(

s

))

, ϕ

sa

)

,

(49)

(Figure 13) Normalization 2: No normalization. Normalization is not applied anywhere. This makes the input to the value function and policy:

(Figure 14) Normalization 3: Normalization used on z sa . Since z sa is trained to predict the next state embedding z s ′ , which is normalized, normalization may not be needed for the output of the state-action encoder g . We try also applying normalization to the output of g . This makes the input to the value function and policy:

(Figure 15) Normalization 4: BatchNorm. BatchNorm is used instead of AvgL1Norm. This makes the input to the value function and policy:

(Figure 15) Normalization 5: LayerNorm. LayerNorm is used instead of AvgL1Norm. This makes the input to the value function and policy:

(Figure 16) Normalization 3: Cosine similarity loss. This is inspired by SPR [Schwarzer et al., 2020], where normalization is 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:

now becomes

No normalization is used elsewhere.

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

No normalization on

ϕ

No normalization

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

Normalization applied to

z

sa

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

BatchNorm

LayerNorm

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

Cosine similarity loss

D.4 End-to-end

The embeddings used in SALE are decoupled, meaning that the encoders used to output the embeddings are trained independently from the value function or 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:

In Table 11 and Figure 17 we display the performance when varying the hyperparameter β .

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

β

= 0

.

1

β

= 1

β

= 10

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 et al., 2019]. If we assume that ( s, a ) is not contained in the dataset, then Q ( s, a ) is effectively a guess generated by the value function Q . During training, s is always sampled from the dataset, but a is often generated by the policy ( a ∼ π ( s )) , resulting in a potentially poor estimate.

Extrapolation error is typically considered a problem for offline RL, where the agent is given a 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 a problem in online RL since the policy interacts with the environment, collecting data for the corresponding actions it generates.

In our empirical analysis (Section E.2) we 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 more prone to erroneous extrapolations (for example, for the Ant environment, the original action dimension size is 8 , but the state-action embedding z sa has a dimension size of 256 ). This is because unseen actions can appear significantly more distinct from seen actions, due to the dramatic increase in the number of dimensions used to represent the action-dependent input.

E.2 Empirical Analysis

In this section we vary the input to 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:

Throughout all experiments, we assume no value clipping (as this is introduced in response to the analysis from this section). This makes the value function loss (originally Equation 15):

Policy checkpoints are not used and the policy is not modified. Since extrapolation error is closely linked to 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 as discussed above.

(Figure 19) Extrapolation 2: No z sa . We remove z sa from the value function input. This makes the input: Q t +1 ( z s t , ϕ ) .

(Figure 20) Extrapolation 3: Small ϕ . We reduce the number of dimensions of ϕ from 256 to 16 .

(Figure 21) Extrapolation 4: No z sa and small ϕ . We remove z sa 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 t +1 ( z s 0 , ϕ ) .

(Figure 23) Extrapolation 6: TD7. The full TD7 method (without checkpoints).

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

Seeds (1-10), buffer size 1M

Seeds (1-10), buffer size 100k

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

Seeds (1-10), buffer size 1M

Seeds (1-10), buffer size 100k

HW

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

Seeds (1-10), buffer size 1M

Seeds (1-10), buffer size 100k

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

Seeds (1-10), buffer size 1M

Seeds (1-10), buffer size 100k

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

Seeds (1-10), buffer size 1M

Seeds (1-10), buffer size 100k

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Total Reward (1k)

Performance by seed

0

1

2

3

4

5

0

1

2

Time steps (1M)

Value Estimate (1k)

Value estimate by seed

Seeds (1-10), buffer size 1M

Seeds (1-10), buffer size 100k

F Policy Checkpoints

F.1 Motivation

Deep RL algorithms are widely known for their inherent instability, which often results in substantial variance in performance during training. Instability in deep RL can occur on a 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).

0

1

2

3

4

5

0

1

2

3

4

Time steps (1M)

Total Reward (1k)

Hopper

0

20

40

60

80

100

0

20

40

60

80

100

> % Decrease

%of Evaluations

Hopper

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 to use training episodes to evaluate the checkpoint policy, as any evaluation or test 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.

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 a 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, if the minimum performance of the current policy drops below the minimum performance of the checkpoint policy. This allows us to use a 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.

Figure 24: Standard presentation Single seed The training curve as commonly presented in RL papers and the training curve of a single seed . Both curves are from the TD3 algorithm, trained for 5M time steps. The standard presentation is to evaluate every N freq steps, average scores over N episodes evaluation episodes and N seeds seeds, then to smooth the curve by averaging over a window of N window evaluations. (In our case this corresponds to N freq = 5000 , N episodes = 10 , N seeds = 10 , N window = 10 ). The learning curve of a single seed has no smoothing over seeds or evaluations ( N seeds = 1 , N window = 1 ). By averaging over many seeds and evaluations, the training curves in RL can appear deceptively smooth and stable.

Early learning. In our empirical analysis we observe that long assessment phases had a negative impact on early learning (Table 13). A final adjustment to policy checkpoints is to keep the number of assessment episodes fixed to 1 during the early stages of learning, and then increase it to a much larger value later in learning. Since the policy or environment may be stochastic, when increasing the number of assessment episodes, we reduce the performance of the checkpoint policy (by a factor of 0 . 9 ), since the checkpoint performance may be overfit to the single episode used to assess its performance.

In Algorithm 1 in the main text, several variables are referenced. These are as follows:

checkpoint_condition : The checkpoint condition refers to either (a) the maximum number 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 is higher than the minimum performance of the checkpoint policy.

Variable descriptions. In Table 3 (in Section A.2), which describes the hyperparameters of TD7, several variables are referenced. These are as follows:

Checkpoint criteria: the measure used to 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 ).

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 an 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, we 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 to 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 use the mean performance instead of the minimum (default) to determine if the current policy outperforms the checkpoint policy. In 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, we terminate the assessment phase if the current mean performance is below the checkpoint performance. This means the number of assessment episodes is variable, up to a maximum of 20.

(Figure 28) Checkpoint 4: 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 1 and increasing to a higher number after the early learning state completes at 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 we present the performance of using immediate checkpoints at 300k time steps. To observe stability benefits, in Figure 30 we display the performance of five individual seeds, with and without checkpoints.

0

1

2

3

4

5

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

1

2

3

4

5

0

1

2

3

4

Time steps (1M)

Hopper

0

1

2

3

4

5

0

2

4

6

8

Time steps (1M)

Walker2d

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Ant

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Humanoid

TD7

TD7 (no checkpoints)

10

50

0

1

2

3

4

5

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

1

2

3

4

5

0

1

2

3

4

Time steps (1M)

Hopper

0

1

2

3

4

5

0

2

4

6

8

Time steps (1M)

Walker2d

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Ant

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Humanoid

TD7

TD7 (no checkpoints)

Mean, fixed assessment length

Mean, variable assessment length

0

1

2

3

4

5

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

1

2

3

4

5

0

1

2

3

4

Time steps (1M)

Hopper

0

1

2

3

4

5

0

2

4

6

8

Time steps (1M)

Walker2d

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Ant

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Humanoid

TD7

TD7 (no checkpoints)

10

20

50

www//m

0

1

2

3

4

5

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

1

2

3

4

5

0

1

2

3

4

Time steps (1M)

Hopper

0

1

2

3

4

5

0

2

4

6

8

Time steps (1M)

Walker2d

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Ant

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Humanoid

TD7

TD7 (no checkpoints)

10

20

50

0

1

2

3

4

5

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

1

2

3

4

5

0

1

2

3

4

Time steps (1M)

Hopper

0

1

2

3

4

5

0

2

4

6

8

Time steps (1M)

Walker2d

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Ant

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Humanoid

TD7

TD7 (no checkpoints)

G Ablation Study

In this section we perform an ablation study on the components of TD7.

G.1 Main Components

(Figure 31) Ablation 1: No SALE. SALE is entirely removed from TD7. This means the input to the value function and policy return to 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 3: 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.

0

1

2

3

4

5

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

1

2

3

4

5

0

1

2

3

4

Time steps (1M)

Hopper

0

1

2

3

4

5

0

2

4

6

8

Time steps (1M)

Walker2d

0

1

2

3

4

5

0

3

6

9

12

Time steps (1M)

Ant

0

1

2

3

4

5

0

2

4

6

8

10

Time steps (1M)

Humanoid

TD7

TD3

No SALE

No checkpoints

No LAP

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 its original input π ( s ) .

(Figure 31) Ablation 5: TD3 with encoder. We add the encoder from SALE to TD3 but train it end-to-end with the value function. The encoder is not applied to the policy.

(Figure 33) Ablation 6: Current policy. TD7 is trained in an identical fashion, but uses the current policy at test time, rather than the checkpoint policy.

(Figure 33) Ablation 7: TD3 with checkpoints. We add policy checkpoints to TD3, using the checkpoint policy at test time.

(Figure 34) Ablation 8: TD3 with LAP. We add LAP to 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):

(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 of AvgL1Norm from TD7. 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 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 t +1 ( z sa t +1 , z s t +1 , s, a ) and π t +1 ( z s t +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 37) Ablation 13: No implementation. We remove the implementation details differences discussed in Section A.1, namely using both value functions when updating the policy, ELU activation functions in the value function network. The target network update is not changed as this may influence the use of fixed encoders.

(Figure 37) Ablation 14: Our TD3. TD3 with the implementation detail differences discussed in Section A.1, 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.

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

No SALE

No SALE, with encoder

TD3 with encoder

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

No checkpoints

Current policy

TD3 with checkpoints

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

No LAP

TD3 with LAP

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

No clipping

TD3 with clipping

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

No normalization

No fixed encoder

0

0.5

1.0

0

4

8

12

16

20

Time steps (1M)

Total Reward (1k)

HalfCheetah

0

0.5

1.0

0

1

2

3

4

Time steps (1M)

Hopper

0

0.5

1.0

0

2

4

6

Time steps (1M)

Walker2d

0

0.5

1.0

0

3

6

9

Time steps (1M)

Ant

0

0.5

1.0

0

2

4

6

8

Time steps (1M)

Humanoid

TD7

TD3

No implementation

Our TD3

H Offline RL Learning Curves

0

0.5

1.0

0

20

40

60

80

100

120

Time steps (1M)

Normalized Score

HalfCheetah

Medium

0

0.5

1.0

0

20

40

60

80

100

120

Time steps (1M)

HalfCheetah

Medium-Replay

0

0.5

1.0

0

20

40

60

80

100

120

Time steps (1M)

HalfCheetah

Medium-Expert

0

0.5

1.0

0

20

40

60

80

100

120

Time steps (1M)

Normalized Score

Hopper

Medium

0

0.5

1.0

0

20

40

60

80

100

120

Time steps (1M)

Hopper

Medium-Replay

0

0.5

1.0

0

20

40

60

80

100

120

Time steps (1M)

Hopper

Medium-Expert

0

0.5

1.0

0

20

40

60

80

100

120

Time steps (1M)

Normalized Score

Walker2d

Medium

0

0.5

1.0

0

20

40

60

80

100

120

Time steps (1M)

Walker2d

Medium-Replay

0

0.5

1.0

0

20

40

60

80

100

120

Time steps (1M)

Walker2d

Medium-Expert

TD7

CQL

TD3+BC

IQL

X

-QL

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 et al., 2019]. All experiments are run on a single Nvidia Titan X GPU and Intel Core i7-7700k CPU. Software is detailed in Appendix B. We display run time adjusted learning curves in Figure 39.

0

4

8

12

16

0

4

8

12

16

20

Hours

Total Reward (1k)

HalfCheetah

0

4

8

12

16

0

1

2

3

4

Hours

Hopper

0

4

8

12

16

0

2

4

6

8

Hours

Walker2d

0

4

8

12

16

0

3

6

9

12

Hours

Ant

0

4

8

12

16

20

0

2

4

6

8

10

Hours

Humanoid

TD7

TD3

SAC

TQC

TD3+OFE

J 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 a setting with image-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 it would be surprising if a method, originally designed for a 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 run time over TD3 (the base algorithm) is more than double.

Theoretical results. We perform extensive empirical analysis to uncover which factors of dynamicsbased 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 a good representation.

Peer reviews:
comments: 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.

STRENGTHS:
1. The paper is well written and easy to follow.
2. Extensive experiments and ablation studies to evaluate the proposed methods.
3. 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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n**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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\nSee Weaknesses Section.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n**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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nSee Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nTD7 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?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-25"}