iQRL – Implicitly Quantized Representations for Sample-efficient Reinforcement Learning

Abstract

Learning representations for reinforcement learning (RL) has shown much promise for continuous control. We propose an efficient representation learning method using only a self-supervised latent-state consistency loss. Our approach employs an encoder and a dynamics model to map observations to latent states and predict future latent states, respectively. We achieve high performance and prevent representation collapse by quantizing the latent representation such that the rank of the representation is empirically preserved. Our method, named iQRL: implicitly Quantized Reinforcement Learning, is straightforward, compatible with any model-free RL algorithm, and demonstrates excellent performance by outperforming other recently proposed representation learning methods in continuous control benchmarks from DeepMind Control Suite.

1 Introduction

Reinforcement learning (RL, e.g., [1]) has shown much promise for solving complex continuous control tasks. However, applying RL in real-world environments is challenging as it typically requires millions of data points which can be unpractical—i.e. RL is sample inefficient. On the other hand, representation learning has become a widely adopted solution for improving sample efficiency in deep learning. The core idea is to learn features which capture the underlying structure and patterns of the data. In the context of RL, such features can be learned independently from the downstream task. Whilst representation learning has had successes in RL, these have mainly been restricted to image-based observations (e.g., CURL [2], DrQ [3], DrQ-v2 [4], and TACO [5]).

The investigation of representation learning for state-based RL is much less common. This is likely due to the fact that learning a compact representation of an already compact state vector seems unnecessary. However, recent work by Fujimoto et al. [6], Zhao et al. [7] suggests that the difficulty of a task is due to the complexity of the underlying transition dynamics, as opposed to the size of the observation space. As such, investigating representation learning for state-based RL is a promising research direction.

Recently, TCRL [7] and SPR [8] have obtained state-of-the-art performance on continuous control benchmarks by learning representations with self-supervised losses. Self-supervised learning (SSL) approaches (which do not reconstruct observations) attempt to learn good features without labels [9]. Whilst they can learn robust representations, self-supervised losses are susceptible to a problem known as representation collapse (see Definition 3.1), where the encoder learns to map all observations to a constant latent representation [10]. As such, when leveraging SSL approaches to learn representations for RL, it is common to combine the self-supervised latent-state consistency loss with other loss terms, such as minimizing the reward prediction error in the latent space [11, 7, 12, 13, 14]. This helps to prevent representation collapse at the cost of learning a task-specific representation.

Figure 1: Overview. iQRL is a stand-alone representation learning technique that is compatible with any model-free RL algorithm (we use TD3 [15]). Importantly, iQRL quantizes the latent representation with Finite Scalar Quantization (FSQ,

), using only a self-supervised latent-state consistency loss, i.e. no decoder (see Eq. 5). Making the latent representation discrete with an implicit codebook ( ) contributes to the very high sample efficiency of iQRL and empirically prevents representation collapse. Thanks to the FSQ-based quantization, iQRL does not need a reward prediction head to prevent representation collapse, a well-known issue with self-supervised learning, making the representation task-agnostic.

In this paper, we propose a simple representation learning technique which learns a task-agnostic representation using only a self-supervised loss. It is based solely on the latent-state consistency loss, i.e. a commonly used self-supervised loss for continuous RL. Importantly, our method empirically prevents representation collapse as it preserves the rank of the representation. We accomplish this by quantizing our latent representation with Finite Scalar Quantization [16], without using any reconstruction loss. As a result, our latent space is bounded and associated with an implicit codebook, whose size we can control. Our method can be combined with any model-free RL method (we use TD3, [15]). See Fig. 1 for an overview of our representation learning method. Importantly, our method (i) alleviates representation collapse, (ii) demonstrates excellent sample efficiency outperforming TCRL and TD7 on a wide range of different continuous control tasks, (iii) is simple to implement, and (iv) learns a task-agnostic representation that could be helpful in downstream tasks.

2 Related Work

In this section, we recap methods for representation learning in RL. In particular, we motivate why researchers are moving towards learning representations using self-supervised learning. Then, as our method builds upon self-supervised representation learning, which is susceptible to representation and dimensional collapse (see Definitions 3.1 and 3.2), we review contrastive self-supervised representation learning approaches; an alternative approach to preventing representation collapse.

Representation learning

Learning representations for RL has been investigated for decades [17, 18, 19, 20, 21, 22, 23, 24, 25, 26]. However, these approaches are usually limited to simple environments. More recently, Fujimoto et al. [6] proposed TD7, an extension of TD3 which learns state and action embeddings and then performs TD3 with this representation, making it highly similar to our method, which also uses TD3 as the base algorithm. However, their method uses a self-supervised loss with no explicit mechanism to prevent representation collapse. In contrast to TD7 and motivated by representation collapse, we quantize our latent space, which we show empirically prevents representation collapse.

Observation reconstruction

A prominent idea in both model-based and model-free RL has been to learn latent representations with reconstruction objectives (e.g. VAE, [27]) [28, 29, 30, 25, 31, 32]. However, as these approaches use observation reconstruction to learn a representation their latent representation contains information about the observation which cannot be controlled by the agent and is not relevant for solving the task, which distorts the optimization landscape [33, 34]. In our experiments, we show that learning representations with observation reconstruction for model-free RL not only harms sample efficiency, but in the complex DMC Dog Run task, it can prevent the agent from solving the task.

Latent-state consistency

How can effective representations be learned efficiently without resorting to reconstruction? A common solution has been to attach an auxiliary loss to the RL objective and perform representation learning [35, 36]. A promising approach for learning suitable representations is the use of self-predictive abstractions, where the model is trained to predict future latent states through an auxiliary loss [37]. Ye et al. [38] introduce a self-supervised consistency loss on the learned latent representation. Instead of relying on a reconstruction-based loss function, Schwarzer et al. [39] propose a cosine similarity loss between predicted future latent states and the true future latent states and then perform Q-learning in the learned latent space. Our approach shares similarities with SPR [39], however, we focus on state-based observations, which leads us to quantizing our representation to prevent representation collapse, instead of using a projection head [40].

Latent-state consistency for model-based RL

Similarly to model-free RL, using the reconstruction loss for learning representation is also unreliable in model-based RL [41] and can have a detrimental effect on the performance of model-based methods in various benchmarks [42, 43]. Therefore, in the context of model-based RL, TD-MPC/TD-MPC2 [12, 44] use a consistency loss to learn representations for planning with Model Predictive Path Integral control together with reward and value functions learned through temporal difference methods [45]. Zhao et al. [7] show that the planning component of TD-MPC is not strictly necessary for high performance and applying model-free RL on top of the self-consistent representations is sufficient for performance competitive with state-of-the-art. We build on top of TCRL to show that we can combat representation collapse through latent-space quantization. As a result, we can drop the reward prediction head to learn a task-agnostic representation.

Contrastive learning

An alternative approach to preventing representation collapse in self-supervised learning is to use contrastive losses. In the context of RL, this was done by CURL [2] and TACO [5]. Whilst CURL and TACO are designed for image-based observations, their contrastive learning approaches could still hold value in state-based RL. The main idea in contrastive learning is to prevent representation collapse (see Definition 3.1) by pushing the latent vectors associated with different observations away from each other. Nevertheless, contrastive methods still experience dimensional collapse [10]. In contrast to TACO, we do not use a contrastive loss and instead leverage quantization to help prevent both representation and dimensional collapse. To offer a fair comparison between contrastive learning and our quantization scheme, we compare our method against a version of TACO tuned for state-based RL.

3 Preliminaries

In this section, we introduce and formally define representation collapse.

Representation collapse

Self-supervised learning methods learn representations by minimizing distances between two embedding vectors. As such, there is a trivial solution where the encoder outputs a constant for all inputs. Formally, this can be defined as follows:

Definition 3.1 (Complete representation collapse).

Given an encoder $e_{\theta}:\mathcal{O}\rightarrow\mathcal{Z}$ which maps observations ${\bm{o}}\in\mathcal{O}$ to latent states ${\bm{z}}\in\mathcal{Z}$, the representation is said to be completely collapsed when the latent representation is constant for all observations, i.e., $e_{\theta}({\bm{o}})=c,\forall{\bm{o}}\in\mathcal{O}$.

In the context of SSL, **g et al. [10] investigated another type of representation collapse known as dimensional collapse, which is defined as follows:

Definition 3.2 (Dimensional collapse).

Given an encoder $e_{\theta}:\mathcal{O}\rightarrow\mathcal{Z}$ which maps observations ${\bm{o}}\in\mathcal{O}$ to latent states ${\bm{z}}_{t}\in\mathcal{Z}=\mathbb{R}^{d}$, of dimension $d$, the representation is said to be dimensionally collapsed when the latent vectors only span a lower dimensional space.

Whilst complete representation collapse is a clear issue when learning representations for RL, it is not immediately obvious if dimensional collapse is an issue because the goal of representation learning is often considered to be learning a lower-dimensional representation of the observations. In other words, there is a trade-off: We want to learn a lower-dimensional representation, but at the same time, we want to ensure that the representation contains all the information required for predicting future states and, thus, state values [36]. Our experiments show that whilst dimensional collapse is not always an issue, in some more complex environments, it can prevent agents from learning to solve a task (see Fig. 3).

Algorithm 1 iQRL

Input: Encoder $e_{\theta}$, dynamics $d_{\phi}$, critics $\{q_{\psi_{1}},q_{\psi_{2}}\}$, policy $\pi_{\eta}$, learning rate $\alpha$, target network update rate $\tau$

for $i$ to $N_{\text{episodes}}$ do

   $\mathcal{D}\leftarrow\mathcal{D}\cup\{{\bm{o}}_{t},{\bm{a}}_{t},{\bm{o}}_{t+1},r_{t+1}\}^{T}_{t=0}$ $\triangleright$ Collect data in environment

   for $i=1$ to $T$ do

      $[\theta,\phi]\leftarrow[\theta,\phi]+\alpha\nabla\left(\mathcal{L}_{\text{rep}}(\theta,\phi;\mathcal{D})\right)$ $\triangleright$ pdate representation, Eq. 5

      $\psi\leftarrow\psi+\alpha\nabla\left(\mathcal{L}_{q}(\psi;\mathcal{D})\right)$ $\triangleright$ Update critic, Eq. 7

      if $i$ % 2 == 0 then

         $\eta\leftarrow\eta+\alpha\nabla\left(\mathcal{L}_{\pi}(\eta;\mathcal{D})\right)$ $\triangleright$ Update actor less frequently than critic, Eq. 8

      end if

      $[\bar{\theta},\bar{\psi},\bar{\eta}]\leftarrow(1-\tau)[\bar{\theta},\bar{\psi},\bar{\eta}]+\tau[{\theta},{\psi},\eta]$ $\triangleright$ Update target networks

   end for

end for

4 Method

In this section, we detail our method, named implicitly Quantized Reinforcement Learning (iQRL). iQRL is conceptually simple, it (i) learns a representation of the observation space and then, (ii) performs model-free RL (e.g., TD3) on this representation. See Fig. 1 and Algorithm 1.

We consider Markov Decision Processes (MDPs, [46]) $\mathcal{M}=(\mathcal{O},\mathcal{A},\mathcal{P},\mathcal{R},\gamma)$, where an agent receives an observation ${\bm{o}}_{t}\in\mathcal{O}$ at time step $t$, performs an action ${\bm{a}}_{t}\in\mathcal{A}$, and then obtains the next observation ${\bm{o}}_{t+1}=\mathcal{P}(\cdot\mid{\bm{o}}_{t},{\bm{a}}_{t})$ and reward $r_{t}=\mathcal{R}({\bm{o}}_{t},{\bm{a}}_{t})$. The discount factor is denoted $\gamma\in[0,1)$.

Method components

iQRL has four main components which we wish to learn:

Encoder: $${\bm{z}}_{t} =f(e_{\theta}({\bm{o}}_{t}))$$ (1)

Dynamics: $$\hat{{\bm{z}}}_{t+1} ={f}({{\bm{z}}}_{t}+d_{\phi}({\bm{z}}_{t},{\bm{a}}_{t}))$$ (2)

Value: $${\bm{q}}_{t} =\mathbf{q}_{\psi}({\bm{z}}_{t},{\bm{a}}_{t})$$ (3)

Policy: $${\bm{a}}_{t} \sim\pi_{\eta}({\bm{z}}_{t})$$ (4)

The encoder $e_{\theta}$ and latent-space dynamics model $d_{\phi}$ are responsible for representation learning. $f(\cdot)$ denotes our quantization scheme, which implicitly quantizes our latent representation (more details to follow). The encoder (with quantization) $f\circ e_{\theta}(\cdot)$ maps observations ${\bm{o}}_{t}$ to latent states ${\bm{z}}_{t}$ and is responsible for learning a representation which can aid RL. The latent-space dynamics model (with quantization) $d_{\phi}(\cdot)$ predicts the next latent states $\hat{\bm{z}}_{t+1}$ given a latent state ${\bm{z}}_{t}$ and an action ${\bm{a}}_{t}$. Its sole purpose is to aid representation learning by making the latent states temporally consistent. Note that we do not use it for model-based RL. Once we have the representation learned by our encoder, we map all observations to the latent space and perform model-free RL in this latent space. Throughout this paper, we use Twin Delayed Deep Deterministic Policy Gradient (TD3, [15]) as the base algorithm. It consists of two state-action value functions $\{q_{\psi_{1}},q_{\psi_{2}}\}$, known as critics, and a deterministic actor $\pi_{\eta}$. Following prior works [4], we use a linear exploration noise schedule which decays from $1$ to $0.1$ during training.

Representation learning

Our representation learning uses the latent-state consistency loss,

$$\mathcal{L}_{\text{rep}}(\theta,\phi;\tau) =\sum_{h=0}^{H-1}\gamma_{\text{rep}}^{h}\left(\frac{{f}(\hat{\bm{z}}_{t+h}+d_{\phi}(\hat{{\bm{z}}}_{t+h},{\bm{a}}_{t+h}))}{\|{f}(\hat{\bm{z}}_{t+h}+d_{\phi}(\hat{{\bm{z}}}_{t+h},{\bm{a}}_{t+h}))\|_{2}}\right)^{\top}\left(\frac{{f}(e_{\bar{\theta}}({\bm{o}}_{t+h+1}))}{\|{f}(e_{\bar{\theta}}({\bm{o}}_{t+h+1}))\|_{2}}\right),$$ (5)

which minimizes the cosine similarity between the next state predicted by the dynamics model $\hat{{\bm{z}}}_{t+1}=f(\hat{\bm{z}}_{t}+d_{\phi}(\hat{{\bm{z}}}_{t},{\bm{a}}_{t}))$ and the next state predicted by the momentum encoder $\bar{{\bm{z}}}_{t+1}=f(e_{\bar{\theta}}({\bm{o}}_{t+1}))$. The latent states are obtained with multi-step predictions in the latent space $\hat{{\bm{z}}}_{t+1}=f(\hat{\bm{z}}_{t}+d_{\phi}(\hat{{\bm{z}}}_{t},{\bm{a}}_{t}))$. The initial map** to the latent space $\hat{{\bm{z}}}_{0}=f(e_{\theta}({\bm{o}}_{0}))$ uses the online encoder which is being trained jointly with the dynamics model $d_{\phi}(\hat{{\bm{z}}_{t}},{\bm{a}}_{t})$. The target $e_{\bar{\theta}}({\bm{o}}_{t+1})$ is calculated with the momentum encoder which uses an exponential moving average (EMA) of the encoder’s weights $\bar{\theta}\leftarrow(1-\tau)\bar{\theta}+\tau\theta$. The target network update rate is denoted $\tau$. Note that we do not use reward or value prediction for learning our representation and as a result, our representation is task-agnostic.

Quantization

Motivated by preventing dimensional collapse we quantize our latent space following the approach from Finite Scalar Quantization (FSQ, [16]). Their important observation is that carefully bounding each dimension gives rise to an implicit codebook $\mathcal{C}$ of a chosen size $|\mathcal{C}|$. Having requested a $d\text{-dimensional}$ latent space, iQRL configures the encoder to output $c$ channels per dimension such that the representation from the encoder ${\bm{x}}=e_{\theta}({\bm{o}})\in\mathbb{R}^{d\times c}$ and the dynamics model $\hat{{\bm{x}}}={{\bm{z}}}+d_{\phi}({\bm{z}},{\bm{a}})\in\mathbb{R}^{d\times c}$ are in $\mathbb{R}^{d\times c}$. To quantize ${\bm{x}}$ (and $\hat{{\bm{x}}}$) into a finite set of codewords, we first apply a bounding function $f(\cdot)$ and then we round to integers. Let us consider a single dimension $j$ of the encoder’s output ${\bm{v}}=[{\bm{x}}]_{j,:}\in\mathbb{R}^{c}$ which consists of $c\text{-channels}$, and demonstrate how it is quantized. We follow FSQ and choose $f(\cdot)$ such that each entry in $\tilde{{\bm{v}}}=\mathrm{round}(f({\bm{v}}))$ takes one of $L_{i}$ unique values,

$$f:{\bm{v}}\rightarrow\lfloor L_{i}/2\rfloor\text{tanh}({\bm{v}}),$$ (6)

where $L_{i}$ is a hyperparameter for channel $i$, specified as FSQ levels $\mathcal{L}=\{L_{1},\ldots,L_{c}\}$. This gives an entry in our codebook $\tilde{\bm{v}}\in\mathcal{C}$, where the implied codebook is given by the product of these per-channel codebook sets. The vectors in $\mathcal{C}$ can be enumerated giving a bijection from any $\tilde{{\bm{v}}}$ to an integer in $\{1,2,\ldots,L^{c}\}$. As an example, in some of our experiments, we used $d=512$ latent dimensions each with $c=2$ channels consisting of 8 levels, i.e. we used FSQ levels $\mathcal{L}=\{L_{1}=8,L_{2}=8\}$. This corresponds to a codebook of size $|\mathcal{C}|=\prod_{i=1}^{c}L_{i}=8\times 8=64=2^{6}$ for each dimension.

Note that this quantization requires a round operation. As such, to propagate gradients through the round operation we use straight-through gradient estimation (STE). This is easily accomplished in deep learning libraries using stop gradient $\mathrm{sg}$ as $\mathrm{round\_ste}(x):x\rightarrow x+\mathrm{sg}(\mathrm{round}(x)-x)$. FSQ has the following hyperparameters: we must specify the number of channels $c$ and the number of levels per channel $\mathcal{L}=\{L_{1},\ldots,L_{c}\}$. Table 1 shows the recommended number of channels and number of levels per channel to obtain codebooks of different sizes [16].

Table 1: FSQ levels $\mathcal{L}$ to approximate different codebook sizes $|\mathcal{C}|$.

Target size $|\mathcal{C}|$	$2^{4}$	$2^{6}$	$2^{8}$	$2^{9}$	$2^{10}$
Proposed $\mathcal{L}$	$\{5,3\}$	$\{8,8\}$	$\{8,6,5\}$	$\{8,8,8\}$	$\{8,5,5,5\}$

In practice, we found codebooks of size $|\mathcal{C}|=2^{6}$ sufficient for all environments in the DeepMind Control suite. However, for more complex environments we hypothesize that larger codebooks will be required.

Figure 2: DeepMind Control Suite results. iQRL (red) is significantly more sample efficient than other model-free baselines TCRL (green), TD7 (purple), TACO (blue) and TD3 (orange). iQRL performs particularly well in the high-dimensional locomotion tasks and outperforms TCRL, which is the most similar baseline. Results are for 20 DMC tasks with UTD=1. We plot the mean (solid line) and the $95\%$ confidence intervals (shaded) across 5 random seeds, where each seed averages over 10 evaluation episodes. See Fig. 11 for results in other DMC tasks.

Model-free reinforcement learning

We learn the policy (actor) and action-value function (critic) using TD3 [15]. However, we follow Yarats et al. [4], Zhao et al. [7] and augment the loss with $n\text{-step}$ returns. The only difference to TD3 is that instead of using the original observations ${\bm{o}}_{t}$, we map them through the online encoder ${\bm{z}}_{t}=f(e_{{\theta}}({\bm{o}}_{t}))$ and learn the actor/critic in the quantized latent space ${\bm{z}}_{t}$. The critic is then updated by minimizing the following objective:

$$\mathcal{L}_{q}(\psi;\tau) =\mathbb{E}_{\tau\sim\mathcal{D}}\left[\textstyle\sum_{k=1}^{2}(q_{\psi_{k}}(f(e_{{\theta}}({\bm{o}}_{t})),{\bm{a}}_{t})-y)^{2}\right],\quad\forall k\in 1,2$$ (7)

$$y =\sum_{n=0}^{N-1}r_{t+n}+\gamma^{n}\min_{k\in\{1,2\}}q_{\bar{\psi}_{k}}(e_{{\theta}}({\bm{o}}_{t+n+1}),{\bm{a}}_{t+n+1}),\quad\text{with}\ {\bm{a}}_{t+n}=\pi_{\bar{\eta}}({\bm{z}}_{t+n})+\epsilon_{t+n},$$

where we use policy smoothing by adding clipped Gaussian noise $\epsilon_{t+n}\sim\text{clip}\left(\mathcal{N}(0,\sigma^{2}),-c,c\right)$ to the action ${\bm{a}}_{t+n}=\pi_{\bar{\eta}}({\bm{z}}_{t+n})+\epsilon_{t+n}$. Note that we use the online encoder to get the latent states in both the prediction and the target. We then use the target action-value functions $\mathbf{q}_{\bar{\psi}}$ and the target policy $\pi_{\bar{\eta}}$ to calculate the TD target. Following TD3, we learn the actor’s parameters by minimizing

$$\mathcal{L}_{\pi}(\eta;\tau)=-\mathbb{E}_{{\bm{o}}_{t}\sim\mathcal{D}}\bigg{[}\min_{k\in\{1,2\}}q_{\psi_{k}}(\underbrace{f(e_{{\theta}}({\bm{o}}_{t}))}_{{\bm{z}}_{t}},\pi_{\eta}(f(e_{{\theta}}({\bm{o}}_{t}))))\bigg{]}.$$ (8)

That is, we maximize the Q-value using the clipped double Q-learning trick to combat overestimation in Q-learning. Note that we do not use the momentum encoder in the actor/critic objectives. In our experiments, using the momentum encoder resulted in worse performance.

Whilst our method shares similarities with TCRL [7], it is important to note that our transition model does not predict the reward. Instead, iQRL leverages quantization to help alleviate representation collapse, and, as a result, learns a task-agnostic representation.

Figure 3: Ablation of quantization. We show how our quantization scheme prevents dimensional collapse. In all tasks, our FSQ scheme prevents dimensional collapse (red) as the rank of the representation remains high. In contrast, when our quantization is not used (blue) the representation undergoes dimensional collapse, indicated by the rank reducing. In the Dog Run task, this results in the agent not learning to solve the task. We plot the mean (solid line) and the $95\%$ confidence intervals (shaded) across 5 random seeds, where each seed averages over 10 evaluation episodes.

5 Experiments

In this section, we evaluate iQRL in a variety of continuous control tasks from the DeepMind Control (DMC) Suite [47]. We aim to answer the following questions:

1. How does iQRL compare to state-of-the-art model-free RL algorithms, especially in the hard DMC tasks?

2. Does our FSQ-based quantization help combat representation and dimensional collapse?

3. Is learning a representation with only latent-state consistency really better than including reward predictions?

4. What impact does reconstruction loss have on the performance of iQRL?

iQRL is simple, fast, and performant

We compare iQRL to the model-free baseline Twin Delayed DDPG (TD3, [15]), and the representation learning-based RL methods Temporal Consistency Reinforcement Learning (TCRL, [7]), TD7 ([6]), and Temporal Action-driven Contrastive Learning (TACO, [5]). In Fig. 2, we evaluate sample efficiency by plotting the average performance of the algorithms across 20 DeepMind Control Suite tasks as a function of environment steps. We see that, on average, iQRL outperforms the baselines and shows significant advantages in many environments. We outperform TCRL, which is the most similar baseline to our work. Furthermore, TD3 is noncompetitive with iQRL, highlighting the importance of representation learning in state-based reinforcement learning. For complete results on all 20 tasks, see Fig. 11 in Appendix G. For more details of the tasks on which the algorithms were evaluated, see Appendix C. For more details of the baselines used in our work and how we implemented them, see Appendix B.

High-dimensional control

Many tasks in DeepMind Control Suite are particularly high-dimensional. For instance, the observation space of the Dog tasks is $\mathcal{O}\in\mathbb{R}^{223}$ and the action space is $\mathcal{A}\in\mathbb{R}^{38}$, and for Humanoid, the observation space is $\mathcal{O}\in\mathbb{R}^{67}$ and the action space $\mathcal{A}\in\mathbb{R}^{24}$. Fig. 2 and Fig. 11 show that iQRL excels in the high dimensional Dog and Humanoid environments when compared to the baselines. We hypothesize that our discretized representations are particularly beneficial for simplifying learning the transition dynamics in high-dimensional spaces, making iQRL highly sample efficient in these tasks.

iQRL does not suffer from rank collapse

We examine the behaviour of adding quantization to our MLP encoder during training. Following Ni et al. [36], we estimate the rank of the linear operator associated with the MLP encoder by calculating the matrix rank of the latent states for a batch of inputs. We ensure full rank at the start of training by orthogonally initializing the MLP encoders. Fig. 3 shows the orthogonality-preserving effect of our quantization scheme as the matrix rank stays close to the maximum. Without quantization, a dimensional collapse occurs, which can have significant harmful effects as the representational power of the latent state diminishes [10]. Correspondingly, in three of the four environments, removing the quantization has a deteriorating impact on the sample efficiency of iQRL, and in Dog Run, the algorithm completely fails to learn to solve the task without the quantization.

Figure 4: Reward prediction is not necessary for representation learning. We compare iQRL to a variant of our method with a reward prediction head trained to predict the reward from the current latent state. Adding a reward prediction head to iQRL leads into a slight increase in performance in Dog Run, but has a slightly harmful impact on sample efficiency in Humanoid Walk and Quadruped Run. We plot the mean (solid line) and the $95\%$ confidence intervals (shaded) across 5 random seeds, where each seed averages over 10 evaluation episodes.

Reward prediction is not necessary for representation learning

Unlike prior methods such as TCRL [7], TD-MPC [12], Dreamer-V2 [48] and TACO [49], our representation learning loss (Eq. 5) does not include a term for learning to predict the reward or the value of the latent state. Instead, we rely solely on the self-predictive temporal consistency loss. To analyze the impact of not including the reward prediction term, we compare our method to a variant of our method, where we have included a reward prediction head similar to that of Zhao et al. [7]. Formally, we define a reward head as $\hat{r}_{t}=g_{\xi}({\bm{z}}_{t},{\bm{a}}_{t})$ (see also Eq. 2), and include a reward prediction term (discounted MSE loss) in the representation loss:

$$\mathcal{L}_{\text{rew}}=\sum_{h=0}^{H-1}\gamma_{\text{rep}}^{h}\|\hat{r}_{t+h}-r_{t+h}\|_{2}^{2},$$

where $r_{t+h}$ is the ground-truth H-step reward and $\hat{r}_{t+h}$ is the predicted H-step reward.

The results for this ablation study are shown in Fig. 4. The plots show that our method, iQRL, without a reward prediction term in the loss, has equal or superior performance to the variant with a reward prediction term except in Dog Run. Our results imply that learning to predict the reward is not necessary for learning a suitable latent representation. The upside of not including the reward prediction head is that it makes the representation task-agnostic, which we believe to be important for downstream applications such as speeding up learning in an incremental multi-task setting in the same domain.

We also evaluated whether including a reward head alone without our FSQ-based normalization scheme is sufficient for preserving the rank of the latent representation and found that iQRL with a reward prediction head but without FSQ suffers from poor performance and dimensional collapse. Therefore, the reward prediction head is not a substitute for our quantization. For more details of this experiment, see Appendix F.

Reconstruction loss has a detrimental impact

Learning to minimize the observation reconstruction error has been widely applied in model-based RL [1, 50, 51], and an observation decoder has been a component of many of the most successful RL algorithms to date [52]. However, recent work in representation learning for RL [7] and model-based RL [12] has shown that incorporating a reconstruction term into the representation loss can hurt the performance, as learning to reconstruct the observations is inefficient due to the observations containing irrelevant details and visuals like shading that are uncontrollable by the agent and do not affect the tasks.

To provide a thorough analysis of iQRL, we include results where we add a reconstruction term to our representation loss in Eq. 5:

$$\mathcal{L}_{{\bm{o}}}=\mathbb{E}_{{\bm{o}}_{t}\sim\mathcal{D}}[\|\hat{{\bm{o}}}_{t}-{\bm{o}}_{t}\|_{2}^{2}],\quad\hat{{\bm{o}}}_{t}=h_{\kappa}({\bm{z}}_{t}),$$ (9)

where $h_{\kappa}$ is a learned observation decoder that takes the latent state as the input and outputs the reconstructed observation. The decoder $h_{\kappa}$ is a standard MLP. We perform reconstruction at each time step in the horizon. The results in Fig. 5 show that in no environments does reconstruction aid learning, and in some tasks, such as the difficult Dog Run and Humanoid Walk tasks, including the reconstruction term has a significant detrimental effect on the performance, and can even prevent learning completely. Our results support the observations of Zhao et al. [7] and Hansen et al. [12] about the lack of need for a reconstruction target in continuous control tasks.

Figure 5: Reconstruction loss has detrimental impact. Unlike many methods, such as SAC-AE [53], iQRL neither has an observation decoder nor a reconstruction term in the loss function. We show that adding a reconstruction loss harms the performance of iQRL across a mixture of easy and hard evaluation environments. We plot the mean (solid line) and the $95\%$ confidence intervals (shaded) across 5 random seeds, where each seed averages over 10 evaluation episodes.

Projection head

Wen and Li [40] and Schwarzer et al. [8] investigated the role of a learnable projection head in non-contrastive self-supervised learning and found that it helps RL algorithms learn more diversified and therefore, superior representations. Whilst iQRL shares similarities with SPR [8], in particular, a temporal consistency loss using cosine similarity, it differs in that it does not use a learnable projection head and quantizes the representation instead. In Fig. 10, we show the impact of adding a projection head to iQRL. It shows that the projection head decreases the sample efficiency of iQRL. Whilst projection heads are effective for learning representations from images, our results suggest that they have a significant negative impact on sample efficiency when learning representations of state-based observations, reaffirming that state-based RL has a different set of challenges to image-based RL and techniques designed to combat representation collapse are not always transferable between the settings.

Stop gradient

[36] proved that using stop gradients should suffice for preventing representation collapse. However, their experiments suggested that using an EMA encoder improves performance over simply using stop gradients. In Fig. 9, we show how replacing iQRL’s EMA encoder with a stop gradient operation can have a negative impact on performace. For example, using stop gradient in the Acrobot Swingup task results in the agent struggling to solve the task.

Codebook size

In Appendix D we evaluate how the size of the codebook $|\mathcal{C}|$ influences the performance of the agent. It shows that size of the codebook and the activeness is intuitive: the smaller the codebook the larger the active proportion. The best codebook size varies between environments but the rank of the representation appears to be preserved for all codebook sizes.

Latent dimension

As a final experiment, we evaluate how the dimension of the latent space $d$ impacts iQRL’s performance. We find that iQRL is fairly robust to different latent dimensions. We find that a latent dimension of $d=1024$ with FSQ levels $\mathcal{L}=[8,8]$, which corresponds to a codebook size $|\mathcal{C}|=2^{6}$, performs best in the harder DMC tasks. See Appendix E for more details.

6 Conclusion

We have presented iQRL, a technique for learning representations using only a self-supervised temporal consistency loss, which demonstrates strong performance in continuous control tasks, including the complex DMC Humanoid and Dog tasks. Our quantization of the latent space empirically preserves the representation’s matrix rank, indicating that it alleviates representation and dimensional collapse. Our experiments further demonstrate that iQRL is extremely sample efficient whilst being fast to train, which we believe is a strong selling point. Importantly, our method is (i) straightforward, (ii) compatible with any model-free RL algorithm, and (iii) learns a task-agnostic representation.

Limitations and future work

Given that iQRL learns a task-agnostic representation, exploring its use for multi-task RL is an exciting direction for future work. Can iQRL learn a single representation which is shared across a wide variety of tasks? In this paper, we have only evaluated iQRL in deterministic environments so extending iQRL to stochastic environments is another important direction for future work.

Appendices

Appendix A Implementation Details

Architecture

We implemented iQRL with PyTorch [54] and used the AdamW optimizer [55] for training the models. All components (encoder, dynamics, actor and critic) are implemented as MLPs. Following Hansen et al. [44] we let all intermediate layers be linear layers followed by LayerNorm [56]. Using LayerNorm is what led to our base TD3 implementation performing so well. We use Mish activation functions throughout. Below we summarize the iQRL architecture for our base model.

⬇

iQRL(

(fsq): FSQ(

(project_in): Identity()

(project_out): Identity()

)

(encoder): ModuleDict(

(state): Sequential(

(0): NormedLinear(in_features=O, out_features=256, act=Mish)

(1): Linear(in_features=256, out_features=512)

)

)

(encoder_tar): ModuleDict(

(state): Sequential(

(0): NormedLinear(in_features=O, out_features=256, act=Mish)

(1): Linear(in_features=256, out_features=512)

)

)

(dynamics): Sequential(

(0): NormedLinear(in_features=512+A, out_features=512, act=Mish)

(1): NormedLinear(in_features=512, out_features=512, act=Mish)

(2): Linear(in_features=512, out_features=512)

)

(pi): Actor(

(_pi): Sequential(

(0): NormedLinear(in_features=512, out_features=512, act=Mish)

(1): NormedLinear(in_features=512, out_features=512, act=Mish)

(2): Linear(in_features=512, out_features=A)

)

)

(pi_tar): Actor(

(_pi): Sequential(

(0): NormedLinear(in_features=512, out_features=512, act=Mish)

(1): NormedLinear(in_features=512, out_features=512, act=Mish)

(2): Linear(in_features=512, out_features=A)

)

)

(critic): Critic(

(_q1): Sequential(

(0): NormedLinear(in_features=512+A, out_features=512, act=Mish)

(1): NormedLinear(in_features=512, out_features=512, act=Mish)

(2): Linear(in_features=512, out_features=1)

)

(_q2): Sequential(

(0): NormedLinear(in_features=512+A, out_features=512, act=Mish)

(1): NormedLinear(in_features=512, out_features=512, act=Mish)

(2): Linear(in_features=512, out_features=1)

)

)

(critic_tar): Critic(

(_q1): Sequential(

(0): NormedLinear(in_features=512+A, out_features=512, act=Mish)

(1): NormedLinear(in_features=512, out_features=512, act=Mish)

(2): Linear(in_features=512, out_features=1, bias=True)

)

(_q2): Sequential(

(0): NormedLinear(in_features=512+A, out_features=512, act=Mish)

(1): NormedLinear(in_features=512, out_features=512, act=Mish)

(2): Linear(in_features=512, out_features=1)

)

)

)

where $O$ is the dimensionality of the observation space and $A$ is the dimensionality of the action spaces.

Hyperparameters

Table 2 lists all of the hyperparameters for training iQRL which were used for the main experiments and the ablations.

Table 2: iQRL hyperparameters. We kept most hyperparameters fixed across all tasks.

Hyperparameter	Value	Description
Training
Action repeat	2
Max episode length	500	Action repeat makes this 1000
Num. eval episodes	$10$
Random episodes	$10$	Num. random episodes at start
TD3
Actor update freq.	$2$	Update actor less than critic
Batch size	$256$
Buffer size	$10^{6}$
Discount factor $\gamma$	$0.99$
Exploration noise	$\mathrm{Linear}(1.0,0.1,50)$ (easy)
	$\mathrm{Linear}(1.0,0.1,150)$ (medium)
	$\mathrm{Linear}(1.0,0.1,500)$ (hard)
Learning rate	$3\times 10^{-4}$
MLP dims	$[512,512]$	For actor/critic/dynamics
Momentum coef. ($\tau$)	$0.005$
Noise clip	$0.3$
N-step TD	$1$ or $3$
Policy noise	$0.2$
Update-to-data (UTD) ratio	$1$
Encoder
Discount factor $\gamma_{\text{rep}}$	$0.9$
Encoder learning rate	$10^{-4}$
Encoder MLP dims	$[256]$
Encoder momentum coef. ($\tau$)	0.005
FSQ levels	$[8,8]$
Horizon $(H)$	$5$	For representation learning
Latent dimension ($d$)	$512$
	$1024$ (Humanoid/Dog)

Statistical significance

We used five seeds for the main figures, at least three seeds for all ablations, and plotted the 95 % confidence intervals as the shaded area, which corresponds to approximately two standard errors of the mean.

Hardware

We used Nvidia A100s and AMD Instinct MI250X GPUs to run our experiments. All our experiments have been run on a single GPU with a single-digit number of CPU workers.

Open-source code

For full details of the implementation, model architectures, and training, please check the code, which is available in the submitted supplementary material and will be made public upon acceptance to guarantee seamless reproducibility.

–appendices continue on next page–

Appendix B Baselines

In this section, we provide further details of the baselines we compare against. In particular, we provide details of how we modified the original codebases and tuned the hyperparameters in an effort to offer a fair comparison.

• Temporal Consistency Reinforcement Learning (TCRL, [7]) is a reinforcement learning algorithm consisting of four components, an encoder and transition, policy and value functions, similarly to iQRL. TCRL uses a temporal consistency loss similar to model-based reinforcement learning to learn a representation used for model-free policy and value function training. The most crucial difference between TCRL and iQRL is that we replace the reward prediction head in the transition function with the FSQ-based normalization scheme. We used the official TCRL implementation on GitHub to run the TCRL experiments in our paper. For the DeepMind Control Suite (DMC) tasks, we used the tuned hyperparameters from the original paper. We used the official PyTorch implementation.

• Temporal Action-driven Contrastive Learning (TACO, [5]) is a temporal contrastive learning framework that learns a latent representation of states and actions with a contrastive loss that optimizes the mutual information between the representations of current states and the following action sequences, and those of the corresponding future states. TACO was primarily designed for vision-based tasks, whereas our benchmarks are state-based. We adapted TACO to the state-based setting by increasing the learning rate and update-to-data ratios to match those of iQRL. We also replaced their CNN-based encoder with the MLP-based encoder of iQRL. Then, we performed a grid search over feature dimensions of 50 and 128, hidden dimensions of 256, 512, and 1024, and frame stacking and no frame stacking. We found the combination of a feature dimension of 50 and a hidden dimension of 1024 without frame stacking to perform the best.

• Twin Delayed DDPG (TD3, [15]) is a model-free RL algorithm for continuous control, extending deep deterministic policy gradient (DDPG) to deal with value overestimation bias. Compared to DDPG, this algorithm uses two critics and takes the minimum over the two for training, adds clipped noise to the actions selected for bootstrap** (policy smoothing), and updates the actor less frequently compared to the critics. iQRL is based on TD3 and we simply replace the observations with their corresponding latent representation by map** them through our encoder. This baseline uses our TD3 implementation which obtains very strong results. Comparing to this baseline allows us to investigate the impact of representation learning on sample-efficiency.

• TD7 [6] is a model-free reinforcement learning algorithm for continuous control that builds on TD3. TD7 builds on a representation learning method, state-action learned embeddings (SALE). The embeddings are learned using a temporal consistency term in the latent state. Other improvements that TD7 has over TD3 are prioritized experience replay and checkpointing. TD7 was initially evaluated on MuJoCo. To adapt it for DeepMind Control Suite, we added action repeats, essential for good performance on DMC. Then, we compared the original hyperparameters of TD7 to those of iQRL and found iQRL to perform the best, so we used those for the final evaluation. In particular, the exploration noise decay of iQRL was crucial for high performance in the DMC environments, and without it, TD7 struggled. Note that both TD7 and iQRL use TD3 as the underlying algorithm, allowing us to reliably compare the impact of SALE and our FSQ-based representations. We used the official PyTorch implementation of TD7.

–appendices continue on next page–

Appendix C Tasks

We evaluate our method in 20 tasks from the DeepMind Control suite [47]. Table 3 provides details of the environments we used, including the dimensionality of the observation and action spaces.

Table 3: DMControl. We consider a total of 20 continuous control tasks from the DeepMind Control suite.

Task	Observation dim	Action dim	Sparse?
Acrobot Swingup	6	1	N
Cheetah Run	17	6	N
Cup Catch	8	2	Y
Dog Run	223	38	N
Dog Trot	223	38	N
Dog Stand	223	38	N
Dog Walk	223	38	N
Fish Swim	24	5	N
Hopper Hop	15	4	N
Hopper Stand	15	4	N
Humanoid Run	67	24	N
Humanoid Stand	67	24	N
Humanoid Walk	67	24	N
Quadruped Run	78	12	N
Quadruped Walk	78	12	N
Reacher Easy	6	2	Y
Reacher Hard	6	2	Y
Walker Run	24	6	N
Walker Stand	24	6	N
Walker Walk	24	6	N

–appendices continue on next page–

Appendix D Ablation of Codebook Size

In this section, we evaluate how the size of the codebook $|\mathcal{C}|$ influences training. We indirectly configure different codebook sizes via the FSQ levels $\mathcal{L}=\{L_{1},\ldots,L_{c}\}$ hyperparameter. This is because the codebook size is given by $|\mathcal{C}|=\prod_{i=1}^{c}L_{i}$. The top row of Fig. 6 compares the training curves for different codebook sizes. The algorithm’s performance is not particularly sensitive to the codebook size. A codebook that is too large can result in slower learning. The best codebook size varies between environments. The most difficult environment, Humanoid Run, benefits from the largest codebook.

Given that a codebook has a particular size, we can gain insights into how quickly iQRL’s encoder starts to activate all of the codebook. The connection between the codebook size and the activeness of the codebook is intuitive: the middle row of Fig. 6 shows that the smaller the codebook, the larger the active proportion.

In the bottom row of Fig. 6, we evaluate how different codebook sizes affect the encoder’s ability to preserve the rank of the representation. We see that the rank of the representation is maintained no matter the codebook size.

Figure 6: Codebook size ablation. We compare how the codebook size affects the performance of iQRL (top), the percentage of the codebook that is active during training (middle), and how the different codebook sizes affect the encoder’s ability to preserve the rank of the representation (bottom). In general, smaller codebooks become fully active faster than larger codebooks, and the rank of the representation is maintained for all codebook sizes. We plot the mean and the $95\%$ confidence intervals (shaded) across 3 random seeds for all environments.

–appendices continue on next page–

Appendix E Ablation of Latent Dimension $d$

This section investigates how the latent dimension $d$ affects the behavior and performance of iQRL in four different environments. The latent dimension $d$ corresponds to the dimension of the representation corresponding to each FSQ level before and after quantization is applied. In the top row of Fig. 7, we see that the performance of our algorithm is robust to the latent dimension $d$, although a latent dimension too small can result in inferior performance, especially in the more difficult environments. The bottom row of Fig. 7 demonstrates that iQRL learns to use the complete codebook irrespective of the latent dimension. However, a larger $d$ can also correspond to the codebook becoming fully active slightly slower.

Figure 7: Latent dim $d$ ablation. We compare how the latent dimension $d$ affects the performance of iQRL (top) and the percentage of the codebook that is active during training (bottom). In general, our algorithm is robust to the latent dimension of the representation, although in more difficult environments, such as Humanoid Walk, a $d$ too small can harm the agent’s performance. We plot the mean and the $95\%$ confidence intervals (shaded) across 3 random seeds for all environments.

–appendices continue on next page–

Appendix F Further Ablations

Figure 8: Adding a reward head is not enough to prevent loss of rank. We show how removing our quantization scheme leads to dimensional collapse measured in terms of the rank of the representation, and in addition, how adding a reward prediction head to iQRL without quantization is insufficient to counteract this and maintain full rank. We plot the mean and the $95\%$ confidence intervals (shaded) across 3 random seeds for all environments.

Figure 9: Replacing EMA encoder with stop gradient. We show that removing iQRL’s EMA encoder and replacing it with only stop gradient hurts performance in DMC tasks. This is particularly apparent in the Acrobot Swingup task. We plot the mean and the $95\%$ confidence intervals (shaded) across 3 random seeds for all environments.

Figure 10: Adding a projection head decreases sample efficiency. We show that adding a projection head to iQRL, similar to what is done in SPR [8], decreases iQRL’s sample efficiency. We plot the mean and the $95\%$ confidence intervals (shaded) across 3 random seeds for all environments.

–appendices continue on next page–

Appendix G Further DMC Results

Fig. 11 compares iQRL to the baselines in the 20 DMC tasks. iQRL’s representation learning significantly improves sample efficiency when compared to TD3. Note that iQRL uses the same TD3 implementation with the same hyperparameters, so the only difference is our representation learning. iQRL also outperforms TCRL in terms of sample efficiency, even without the reward prediction head. Our experiments indicate that this improvement is due to our quantization and the inclusion of LayerNorm in our encoder. We compare iQRL to TACO (which uses a contrastive loss) and observe that iQRL outperforms TACO in most environments. TACO seems to particularly struggle in the Dog tasks. Finally, iQRL outperforms TD7, a state-of-the-art representation learning method for state-based RL.

Figure 11: DeepMind Control results. iQRL performs well across a variety of DMC tasks. We plot the mean (solid line) and the $95\%$ confidence intervals (shaded) across 5 random seeds, where each seed averages over 10 evaluation episodes.