SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Scott Fujimoto  
a, McGill University
Wei-Di Chang McGill University
Edward J. Smith McGill University
Shixiang Shane Gu Google DeepMind
Doina Precup  
Mila, McGill University
David Meger  
Mila, McGill University
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.
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
Corresponding author: scott.fujimoto@mail.mcgill.ca
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<sup>1</sup>.
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 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 [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
1https://github.com/sfujim/TD7
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.
In Reinforcement learning (RL) problems are framed as a Markov decision process (MDP). An MDP is a 5-tuple  $(S, A, R, p, \gamma)$  with state space  $S$ , action space  $A$ , reward function  $R$ , dynamics model  $p$ , and discount factor  $\gamma$ , where the objective is to find a policy  $\pi: S \to A$ , a mapping from state  $s \in S$  to action  $a \in A$ , which maximizes the return  $\sum_{t=1}^{\infty} \gamma^{t-1} r_t$ , the discounted sum of rewards  $r$  obtained when following the policy. RL algorithms commonly use a value function  $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$  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}$ :
$$
z ^ {s} := f (s), \quad z ^ {s a} := g \left(z ^ {s}, a\right). \tag {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'$ , 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'}$ :
$$
\mathcal {L} (f, g) := \left(g (f (s), a) - | f \left(s ^ {\prime}\right) | _ {\times}\right) ^ {2} = \left(z ^ {s a} - \left| z ^ {s ^ {\prime}} \right| _ {\times}\right) ^ {2}, \tag {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) \rightarrow Q \left(z ^ {s a}, z ^ {s}, s, a\right), \quad \pi (s) \rightarrow \pi \left(z ^ {s}, s\right). \tag {3}
$$
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
$$
\operatorname {A v g L 1 N o r m} (x) := \frac {x}{\frac {1}{N} \sum_ {i} | x _ {i} |}. \tag {4}
$$
Figure 1: Diagram of State-Action Learned Embeddings (SALE). SALE uses encoders  $(f, g)$  to output embeddings  $(z^s, z^{sa})$  to enhance the input of the value function  $Q$  and policy  $\pi$ .  $\phi$  denotes the output of the corresponding linear layer. The encoders  $(f, g)$  are jointly trained to predict the next state embedding (where  $|\cdot|_{\times}$  denotes the stop-gradient operation), decoupled from the training of the value function and policy (Equation 2). The end-to-end linear layers are trained with gradients from the corresponding network. AvgL1Norm is used to keep the scale of each of the inputs to the value function and policy constant.
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  $\pi$ , to keep them at a similar scale to the learned embeddings. The input to the value function and policy then becomes:
$$
Q \left(z ^ {s a}, z ^ {s}, \operatorname {A v g L 1 N o r m} (\operatorname {L i n e a r} (s, a))\right), \quad \pi \left(z ^ {s}, \operatorname {A v g L 1 N o r m} (\operatorname {L i n e a r} (s))\right). \tag {5}
$$
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},\pi_{t + 1})$  uses embeddings  $(z_t^{sa},z_t^s)$  from the encoders  $(f_{t},g_{t})$  at the previous iteration  $t$ . The value function and policy are thus updated by:
$$
Q _ {t + 1} \left(z _ {t} ^ {s a}, z _ {t} ^ {s}, s, a\right) \approx r + \gamma Q _ {t} \left(z _ {t - 1} ^ {s ^ {\prime} a ^ {\prime}}, z _ {t - 1} ^ {s ^ {\prime}}, s ^ {\prime}, a ^ {\prime}\right), \quad \text {w h e r e} a ^ {\prime} \sim \pi_ {t} \left(z _ {t - 1} ^ {s ^ {\prime}}, s ^ {\prime}\right), \tag {6}
$$
$$
\pi_ {t + 1} \left(z _ {t} ^ {s}, s\right) \approx \underset {\pi} {\operatorname {a r g m a x}} Q _ {t + 1} \left(z _ {t} ^ {s a}, z _ {t} ^ {s}, s, a\right), \quad \text {w h e r e} a \sim \pi \left(z _ {t} ^ {s}, s\right). \tag {7}
$$
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_{t + 1}^{s}$  and  $z_{t + 1}^{sa}$  are trained with Equation 2, using a target  $z_{t + 1}^{s'}$  (hence, without a target network). Every  $n$  steps the iteration is incremented and all target networks are updated simultaneously:
$$
Q _ {t} \leftarrow Q _ {t + 1}, \quad \pi_ {t} \leftarrow \pi_ {t + 1}, \quad \left(f _ {t - 1}, g _ {t - 1}\right) \leftarrow \left(f _ {t}, g _ {t}\right), \quad \left(f _ {t}, g _ {t}\right) \leftarrow \left(f _ {t + 1}, g _ {t + 1}\right). \tag {8}
$$
**Clipped Values.** Extrapolation error is the tendency for deep value functions to extrapolate to unrealistic values on state-action 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
Figure 2: Extrapolation error can occur in online RL when using state-action representation learning. All figures use the Ant environment.  $\phi^{sa}$  corresponds to the output of the linear layer (Linear(s, a) =  $\phi^{sa}$ ) (Equation 5). Both embeddings and  $\phi^{sa}$  have a default dimension size of 256. Small  $\phi^{sa}$  means that  $\mathrm{Dim}(\phi^{sa})$  is set to 16. No  $z^{sa}$  means the value function input is  $Q(z^s, s, a)$ . The default performance and value estimate of 10 individual seeds without value clipping. While the performance trends upwards there are large dips in reward, which correspond with jumps in the estimated value. Varying the input dimension can improve or harm stability of the value estimate. The severity is impacted by the replay buffer size (1M or 100k). The state embedding  $z^s$  is left unchanged in all settings, showing that the state-action embedding  $z^{sa}$  and the linear layer over the state-action input  $\phi^{sa}$  are the primary contributors to the extrapolation error. This shows the potential negative impact from increasing the dimension size of an input which relies on a potentially unseen action. Clipping stabilizes the value estimate, without modifying the input dimension size (Equation 9).
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:
$$
Q _ {t + 1} (s, a) \approx r + \gamma \operatorname {c l i p} \left(Q _ {t} \left(s ^ {\prime}, a ^ {\prime}\right), \min  _ {(s, a) \in D} Q _ {t} (s, a), \max  _ {(s, a) \in D} Q _ {t} (s, a)\right). \tag {9}
$$
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_{\text{target}}^{s'}$  from a target network obtained with an exponential moving average with weight 0.01. Drawing inspiration from Bisimulation metrics [Ferns et al., 2011], Deep-MDP [Gelada et al., 2019] use an objective that considers both the next state embedding  $z^{s'}$  and the
Figure 3: The mean percent loss from using alternate design choices in TD7 at 1M time steps, over 10 seeds and the five benchmark MuJoCo environments. Bracketed values describe the range of the  $95\%$  confidence interval around the mean. Percent loss is computed against TD7 where the default choices correspond to a percent loss of 0. See Section 4.2 for a description of each design choice and key observations. See the Appendix for further implementation-level details.
reward  $r$ . We test including a prediction loss on the reward by having the encoder  $g$  output both  $z^{sa}$  and  $r^{\mathrm{pred}}$  where  $r^{\mathrm{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.
$\Rightarrow$  All learning targets based on the next state  $s^{\prime}$  perform similarly, although using the embedding  $z^{s^{\prime}}$  further improves the performance. On the contrary, the next state-action embedding  $z^{s^{\prime}a^{\prime}}$  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  $\pi(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_{t+1}^{sa}, z_{t+1}^s)$ .
$\Rightarrow$  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].
$\Rightarrow$  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.
$\Rightarrow$  Learning the embeddings end-to-end with the value function performs significantly worse than decoupled representation learning.
5 Stabilizing RL with Decoupled Representation Learning