For SALE: State-Action Representation Learning for Deep Reinforcement Learning — part 1/4

Cleaned-text code points [0, 15812). Bibliography omitted here; source.md includes it.

# 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 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.

## 1 Introduction

Reinforcement learning (RL) is notoriously sample inefficient, particularly when compared to more straightforward paradigms in machine learning, such as supervised learning. † † Corresponding author: scott.fujimoto@mail.mcgill.ca One possible explanation is the usage of the Bellman equation in most off-policy RL algorithms ( DQN ; DDPG ) , which provides a weak learning signal due to an approximate and non-stationary learning target ( fujimoto2022should ) .

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 ( jaderberg2017reinforcement ; oord2018representation ; anand2019unsupervised ; laskin2020curl ; stooke2021decoupling ; yarats2022mastering ) .

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 ( ota2020can ) , we introduce three important design considerations when learning a state-action representation online. Most importantly, we observe the surprising effect of extrapolation error ( fujimoto2019off ) 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 ( watter2015embed ; ha2018world ; hafner2019learning ; gelada2019deepmdp ; schwarzer2020data ) , 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 ( fujimoto2020equivalence ) , and a behavior cloning term (used only for offline RL) ( fujimoto2021minimalist ) 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 1 1 https://github.com/sfujim/TD7 .

## 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 ( li2006towards ; ferns2011bisimulation ; zhang2020learning ) . For higher-dimensional 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 ( watter2015embed ; finn2016deep ; gelada2019deepmdp ) . 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 ( sutton2011horde ; jaderberg2017reinforcement ; riedmiller2018learning ; lin2019adaptive ) . In recent years, representation learning in RL often refers to both compression and feature learning, and is commonly employed in image-based tasks ( kostrikov2020image ; yarats2021improving ; liu2021return ; cetin2022stabilizing ) 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 ( dayan1993improving ; littman2001predictive ) , spanning many approaches in both model-free RL ( munk2016learning ; van2016stable ; zhang2018decoupling ; gelada2019deepmdp ; schwarzer2020data ; fujimoto2021srdice ; ota2020can ; ota2021training ) and model-based RL in latent space ( watter2015embed ; finn2016deep ; karl2017deep ; ha2018world ; hansen2022temporal ; hafner2019learning ; hafner2023mastering ) . Another related approach is representation learning over actions ( tennenholtz2019natural ; chandak2019learning ; whitney2020dynamics ) . 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 ( liu2021return ) , or MDP homomorphisms ( ravindran2004algebraic ; van2020plannable ; van2020mdp ; rezaei2022continuous ) emphasize abstraction more than feature learning. Our approach can be viewed as an extension of OFENet ( ota2020can ) , 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 ( henderson2017deep ; engstrom2019implementation ) . Our use of checkpoints is most closely related to stabilizing policy performance via safe policy improvement ( trpo ; PPO ; laroche2019safe ) , as well as the combination of evolutionary algorithms ( salimans2017evolution ; mania2018simple ) with RL ( khadka2018evolution ; pourchot2018cem ) , 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 S , A A , R R , p p , γ \gamma ) with state space S S , action space A A , reward function R R , dynamics model p p , and discount factor γ \gamma , where the objective is to find a policy π : S → A \pi:S\rightarrow A , a mapping from state s ∈ S s\in S to action a ∈ A a\in A , which maximizes the return ∑ t = 1 ∞ γ t − 1 ​ r t \sum_{t=1}^{\infty}\gamma^{t-1}r_{t} , the discounted sum of rewards r r obtained when following the policy. RL algorithms commonly use a value function Q π ( s , a ) := 𝔼 [ ∑ t = 1 ∞ γ t − 1 r t | s 0 = s , a 0 = a ] Q^{\pi}(s,a):=\mathbb{E}\left[\sum_{t=1}^{\infty}\gamma^{t-1}r_{t}|s_{0}=s,a_{0}=a\right] , which models the expected return, starting from an initial state s s and action a 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 s ​ a , z s ) (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 ) (f,g) where f ⁡ ( s ) f(s) encodes the state s s into the state embedding z s z^{s} and g ⁡ ( z s , a ) g(z^{s},a) jointly encodes both state s s and action a a into the state-action embedding z s ​ a z^{sa} : z s := f ⁡ ( s ) , z s ​ a := g ⁡ ( z s , a ) . \displaystyle z^{s}:=f(s),\qquad z^{sa}:=g(z^{s},a). (1) 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 ′ s^{\prime} , 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 s ​ a z^{sa} and the embedding of the next state z s ′ z^{s^{\prime}} : ℒ ⁡ ( f , g ) := ( g ⁡ ( f ⁡ ( s ) , a ) − | f ⁡ ( s ′ ) | × ) 2 = ( z s ​ a − | z s ′ | × ) 2 , \displaystyle\mathcal{L}(f,g):=\Bigl(g(f(s),a)-|f(s^{\prime})|_{\times}\Bigr)^{2}=\left(z^{sa}-|z^{s^{\prime}}|_{\times}\right)^{2}, (2) where | ⋅ | × |\cdot|_{\times} 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: Q ⁡ ( s , a ) → Q ⁡ ( z s ​ a , z s , s , a ) , π ⁡ ( s ) → π ⁡ ( z s , s ) . \displaystyle Q(s,a)\rightarrow Q(z^{sa},z^{s},s,a),\qquad\pi(s)\rightarrow\pi(z^{s},s). (3) The encoders ( f , g ) (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 ) (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 ( gelada2019deepmdp ) . 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 x_{i} be the i i -th dimension of an N N -dimensional vector x x , then AvgL1Norm ​ ( x ) := x 1 N ​ ∑ i | x i | . \text{AvgL1Norm}(x):=\frac{x}{\frac{1}{N}\sum_{i}|x_{i}|}. (4) AvgL1Norm is applied to the state embedding z s z^{s} . Similar to the normalized loss functions used by SPR ( schwarzer2020data ) and BYOL ( grill2020bootstrap ) , AvgL1Norm protects from monotonic growth, but also keeps the scale of the downstream input constant without relying on updating statistics (e.g. BatchNorm ( ioffe2015batch ) ). 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 s ​ a z^{sa} , as it is trained to match the normalized next state embedding z s ′ z^{s^{\prime}} .

We also apply AvgL1Norm to the state and action inputs (following a linear layer) to the value function Q Q and policy π \pi , to keep them at a similar scale to the learned embeddings. The input to the value function and policy then becomes: Q ⁡ ( z s ​ a , z s , AvgL1Norm ​ ( Linear ​ ( s , a ) ) ) , π ⁡ ( z s , AvgL1Norm ​ ( Linear ​ ( s ) ) ) . Q(z^{sa},z^{s},\text{AvgL1Norm}(\text{Linear}(s,a))),\qquad\pi(z^{s},\text{AvgL1Norm}(\text{Linear}(s))). (5) Unlike the embeddings ( z s , z s ​ a ) (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 t+1 , the input to the current networks ( Q t + 1 , π t + 1 ) (Q_{t+1},\pi_{t+1}) uses embeddings ( z t s ​ a , z t s ) (z^{sa}_{t},z^{s}_{t}) from the encoders ( f t , g t ) (f_{t},g_{t}) at the previous iteration t t . The value function and policy are thus updated by: Q t + 1 ​ ( z t s ​ a , z t s , s , a ) \displaystyle Q_{t+1}(z_{t}^{sa},z_{t}^{s},s,a) ≈ r + γ ​ Q t ​ ( z t − 1 s ′ ​ a ′ , z t − 1 s ′ , s ′ , a ′ ) , where ​ a ′ ∼ π t ​ ( z t − 1 s ′ , s ′ ) , \displaystyle\approx r+\gamma Q_{t}(z_{t-1}^{s^{\prime}a^{\prime}},z_{t-1}^{s^{\prime}},s^{\prime},a^{\prime}),\qquad\text{ where }a^{\prime}\sim\pi_{t}(z^{s^{\prime}}_{t-1},s^{\prime}), (6) π t + 1 ​ ( z t s , s ) \displaystyle\pi_{t+1}(z_{t}^{s},s) ≈ argmax π Q t + 1 ​ ( z t s ​ a , z t s , s , a ) , where ​ a ∼ π ⁡ ( z t s , s ) . \displaystyle\approx\argmax_{\pi}Q_{t+1}(z_{t}^{sa},z_{t}^{s},s,a),\qquad\text{ where }a\sim\pi(z_{t}^{s},s). (7) The current value function Q t + 1 Q_{t+1} is also trained with respect to the previous value function Q t Q_{t} , known as a target network ( DQN ) . The current embeddings z t + 1 s z^{s}_{t+1} and z t + 1 s ​ a z^{sa}_{t+1} are trained with Equation 2 , using a target z t + 1 s ′ z^{s^{\prime}}_{t+1} (hence, without a target network). Every n n steps the iteration is incremented and all target networks are updated simultaneously: Q t ← Q t + 1 , π t ← π t + 1 , ( f t − 1 , g t − 1 ) ← ( f t , g t ) , ( f t , g t ) ← ( f t + 1 , g t + 1 ) . Q_{t}\leftarrow Q_{t+1},\qquad\pi_{t}\leftarrow\pi_{t+1},\qquad(f_{t-1},g_{t-1})\leftarrow(f_{t},g_{t}),\qquad(f_{t},g_{t})\leftarrow(f_{t+1},g_{t+1}). (8)

Seeds (1-10) Buffer: 1M

Buffer: 1M 100k Buffer: 1M 100k Buffer: 1M 100k Buffer: 1M 100k