Year: 2021

Abstract:
Deep reinforcement learning (RL) agents that exist in high-dimensional state spaces, such as those composed of images, have interconnected learning burdens. Agents must learn an action-selection policy that completes their given task, which requires them to learn a representation of the state space that discerns between useful and useless information. The reward function is the only supervised feedback that RL agents receive, which causes a representation learning bottleneck that can manifest in poor sample efficiency. We present k-Step Latent (KSL), a new representation learning method that enforces temporal consistency of representations via a self-supervised auxiliary task wherein agents learn to recurrently predict action-conditioned representations of the state space. The state encoder learned by KSL produces low-dimensional representations that make optimization of the RL task more sample efficient. Altogether, KSL produces state-of-the-art results in both data efficiency and asymptotic performance in the popular PlaNet benchmark suite. Our analyses show that KSL produces encoders that generalize better to new tasks unseen during training, and its representations are more strongly tied to reward, are more invariant to perturbations in the state space, and move more smoothly through the temporal axis of the RL problem than other methods such as DrQ, RAD, CURL, and SAC-AE.

heading: 1 INTRODUCTION
text: Learning sample-efficient continuous-control directly from pixels is a challenging task for reinforcement learning (RL) agents. Despite recent advances, modern approaches require hundreds of thousands, if not millions, of agent-environment interactions to attain acceptable performance (Kalashnikov et al., 2018; Jaderberg et al., 2019). In real-world settings, this magnitude of data collection may take a prohibitive amount of time. Therefore, advances in learning efficiency for continuouscontrol in pixel-based state spaces would significantly accelerate the success of real-world RL applications such as industrial robotics and autonomous vehicles (Dulac-Arnold et al., 2021).
The standard RL loop ties policy optimization and representation learning together with a single supervisory signal. This dual-objective, single-signal learning scheme results in a representation learning bottleneck that ultimately produces agents that understand their environment slowly throughout training (Shelhamer et al., 2017), thereby reducing performance in the early stages of learning. Approaches that have seen the most success in alleviating the bottleneck rely on methods that explicitly encourage the distillation of useful information from the state space.
Recently, data augmentation (Chen et al., 2020a;b; Hénaff et al., 2020) has been studied in the context of RL as a way to improve sample efficiency (Yarats et al., 2021b; Laskin et al., 2020b). Other approaches to improve sample efficiency in RL add auxiliary tasks to the learning process. For example, image reconstruction (Yarats et al., 2020), latent variable models (Lee et al., 2020a; Gelada et al., 2019), variational objectives (Lee et al., 2020b), and contrastive learning (Laskin et al., 2020a; Stooke et al., 2021). While the aforementioned approaches improve performance over base RL algorithms, they do not fully leverage the long-term temporal connection between actions, future states, and rewards. For example, some approaches do not interact with the temporal axis of the RL problem’s underlying Markov decision process (MDP) (Yarats et al., 2020; 2021b; Laskin et al.,
2020b), or are only concerned with next-step dynamics (Laskin et al., 2020a; Stooke et al., 2021; Gelada et al., 2019). As we will show with our work, ignoring the long-term temporal axis of the MDP withholds information that can help agents to understand consequences of their environment interactions more quickly. Approaches that do attempt to capture longer-term information either require a significant amount of random exploration and pre-training (Lee et al., 2020a;b), which counteracts the goal of minimizing agent-environment interactions, or condition their predictions of the future on the full-dimensional state space (Lee et al., 2020b), thereby not encouraging latent representations themselves to be predictive of latent representations multiple steps into the future.
In this paper, we revisit the problem of sample-efficient continuous-control in image-based state spaces. We introduce k-Step Latent (KSL)1, a representation learning module that aids RL agents with continuous control on pixels. In contrast to aforementioned prior methods, KSL’s auxiliary task directly exploits the RL problem’s underlying MDP, requires no pre-training, and encourages latent representations to be predictive of latent representations that are multiple steps into the future. KSL agents learn to create temporally-coherent representations by projecting states into the latent space such that latent representations, when conditioned upon a series of actions, are predictive of latent representations that are multiple steps into the future. KSL’s auxiliary task propels agents to new state-of-the-art results on both the data efficiency (100k steps) and asymptotic performance (500k steps) checkpoints in the popular PlaNet benchmark suite (Hafner et al., 2019).
This paper makes the following contributions. First, KSL augments the learning process of the popular Soft Actor-Critic (Haarnoja et al., 2018a) algorithm to produce new state-of-the-art results in the PlaNet benchmark suite. KSL significantly outperforms current methods such as DrQ (Yarats et al., 2021b) and RAD (Laskin et al., 2020b). Across all six tasks in PlaNet, KSL produces an average improvement over DrQ of 46.9% and 14.5% at the 100k and 500k steps mark, respectively, and an average improvement over RAD of 58.1% and 21.6% at the 100k and 500k steps mark, respectively. Second, we analyze the learned encoders and latent projections that are produced by several methods. Our analyses show that KSL creates representations that are more strongly tied to reward, are less sensitive to perturbation in the state space, and move more smoothly over time than other baseline methods. Also, our analyses show that KSL creates encoders that generalize to unseen tasks better than other baseline methods.

heading: 2 BACKGROUND
text: We consider the usual formulation of RL agents that act within a Markov decision process (MDP), defined by the tuple (S,A, P,R, γ), where S is the state space,A is the action space, P (st+1|st, at) is the transition probability distribution, R : S × A → R is a reward function that maps states and actions to a scalar feedback signal, and γ ∈ [0, 1) is a discount factor. The agent’s goal is to learn an action-selection policy π(a|s) that maximizes the sum of discounted rewards over the time horizon T of the given task: arg maxπ Ea∼π,s∼P [ ∑T t=1 γ
tR(st, at)]. Specifically, we examine learning agents in state spaces comprised of images o ∈ O. Image-based states are high-dimensional and may give an incomplete description of the environment. For example, incomplete information may arise when the camera’s point-of-view shows occluded objects. As such, it is common to stack a series of n consecutive images to help alleviate the problems associated with incomplete information: st = {ot, ..., ot−n+1} (Mnih et al., 2015).

heading: 2.1 SOFT ACTOR-CRITIC
text: Soft Actor-Critic (SAC) (Haarnoja et al., 2018a;b) is a current state-of-the-art off-policy, model-free RL algorithm for continuous control. SAC learns a state-action value-function criticQθ, a stochastic actor πφ, and a temperature α that weighs between maximizing reward and entropy in an augmented RL objective: Est,at∼π[ ∑ tR(st, at) + αH(π(·|st))].
SAC’s critic is learned by minimizing the squared Bellman error over trajectories τ = (st, at, rt, st+1) sampled from a replay memory D:
Lcritic(θ) = Eτ∼D[(Qθ(st, at)− (rt + γy))2], (1) 1Code available at: https://github.com/anon-researcher-repo/ksl
where the temporal-difference target, y, is computed by sampling an action from the current policy:
y = Ea′∼π[Qθ̄(st+1, a′)− α log πφ(a′|st+1)]. (2)
The target critic Qθ̄ is updated as an exponential moving average (EMA): θ̄ ← ζθ̄+ (1− ζ)θ where ζ controls the “speed” at which θ̄ tracks θ (e.g., He et al. (2020)). SAC’s actor parameterizes a multivariate GaussianN (µ, σ) where µ is a vector of means and σ is the diagonal of the covariance matrix. The actor is updated via minimizing :
Lactor(φ) = −Ea∼π,τ∼D[Qθ(st, a)− α log πφ(a|st)], (3) and α is simply learned against a target entropy value.
3 k-STEP LATENT

heading: 3.1 MOTIVATION
text: We design KSL to address the representation learning bottleneck of base RL algorithms by providing an auxiliary task that encourages representational consistency of states that are nearby in time. States that are nearby in time are likely to share high levels of mutual information (Mobahi et al., 2009), which implies that the states’ representations should as well. KSL contains a recurrent mechanism that encourages representations of states, when conditioned upon actions, to be predictive of states that are multiple steps into the future. Intuitively, we hypothesize that agents who understand the long-term consequences of their actions would perform better than agents who do not.

heading: 3.2 TRAINING DETAILS
text: A visual depiction of KSL can be seen in Figure 1 and pseudocode can be found in Appendix B. At each learning step, we uniformly sample M trajectories of length k, τ = {(s1, a1, ..., ak−1, sk)i}Mi=1, from the replay memory, where s1 can be any state in the episode. For each state in each trajectory, we apply a random translation augmentation (Laskin et al., 2020b). Using augmentations within KSL encourages representations to be invariant to perturbations over the temporal axis. KSL contains an online and momentum pathway, which are highlighted in Figure 1 as orange and green, respectively. At step j = 1 of trajectory2 i, sj is fed through an online encoder zo = Φo(sj) and sj+1 is fed through a momentum encoder zm = Φm(sj+1). Next, zo is concatenated with aj and fed through a learned transition model z′o = T ([zo|aj ]). This transition model enables KSL to learn temporal projections from latent representations of state sj and applied action aj to a representation of the following state sj+1 in the form of z′o. Then, z ′ o and zm are fed through their respective projection modules qo = Ψo(z′o), ym = Ψm(zm). Such modules after an encoder have been shown to improve the learned representations of the encoder (Chen et al., 2020a). Finally, qo is fed through a prediction head yo = P(qo).
2Note that representations through both pathways are computed for every step and trajectory. We omit batch notation for steps and trajectories in the paragraph for clarity.
KSL’s loss for step j is computed as the distance between the `2 normalized versions of the final two vectors from both paths: ‖(yo/‖yo‖) − (ym/‖ym‖)‖22. This operation is performed in a loop over all k steps per each trajectory in τ , with one notable difference at each step after j = 1. Instead of encoding every s along the online path, we recursively feed z′o from the previous step j − 1 into the learned transition model T . This recurrent operation is shown with the red arrow in Figure 1. Therefore, Φo is only used to compute a direct representation of the first state within each trajectory. The minimization of KSL’s loss requires Φo to create latent projections such that multi-step consequences can be inferred from any arbitrary starting state using the transition model T . We formulate the loss function as an average over the computed distance for every state and trajectory in a given batch:
LKSL = 1
Mk M∑ i=1 k∑ j=1 ∣∣∣∣∣ ∣∣∣∣∣ y(i,j)o‖y(i,j)o ‖ − y (i,j) m ‖y(i,j)m ‖ ∣∣∣∣∣ ∣∣∣∣∣ 2
2
, (4)
with y(i,j)p denoting the final vector of pathway p for trajectory i at step j.
We implement a stop gradient to disallow backpropagation from Equation 4 through the momentum pathway. Instead, we update the momentum encoder and momentum projection module as an EMA of their online counterparts. We use the state encoding of the online encoder Φo as input into the critic Qθ and actor πφ. While we allow the gradient of Qθ to update weights within Φo, we implement a stop gradient that disallows signal from the actor’s computation graph (Equation 2 and Equation 3) to propagate into Φo. Gradient flow is depicted in Figure 1 as purple, dotted lines. Φm is an EMA of Φo and Qθ̄ is an EMA Qθ, so it is natural for Φm to transform states for input into Qθ̄. See Appendix F for an expanded description of how KSL fits within SAC. We also provide ablations across architecture choices and values of k in Appendix E and Appendix G. We note that KSL is not restricted to SAC and can be applied alongside any RL algorithm for representation learning in tasks with continuous or discrete action spaces.

heading: 3.3 OPTIMIZATION STRATEGY
text: The weights of Φo are affected by two task-specific losses (Equation 1 and Equation 4). Usually, learning algorithms that attempt to optimize parameters across multiple losses use a weighted linear combination of each loss. However, doing so causes the sharing of inductive biases across each task, which may be an issue if the two tasks compete (Yu et al., 2020; Sener & Koltun, 2018). The representation learning bottleneck of the RL signal may cause disruptions in Φo during optimization of KSL’s loss. Additionally, this problem may be exacerbated with optimizers that accumulate gradient statistics (Jean et al., 2018), such as Adam (Kingma & Ba, 2015). Much of the current research in multi-objective learning has focused on methods that actively learn to balance between multiple losses in the context of a single optimizer (Kendall et al., 2018; Groenendijk et al., 2021). Instead, we explicitly separate the optimization of the RL and KSL tasks by sequestering the weights and losses of each task into separate optimizers. This division is not meant to solve the conflicting gradient problem. The division avoids conflation of the Adam optimizer’s momentum estimates between the two tasks. As a consequence of our multi-optimizer choice, Φo is granted a higher diversity in gradient-update directions. In our experiments, we see a performance boost due to this separation. See Appendix I for more details.

heading: 4 EXPERIMENTAL EVALUATION
text: We investigate several facets of KSL and other baseline methods. Most importantly, does KSL result in agents that converge to asymptotic performance in fewer agent-environment interactions than other current methods? We examine KSL’s learning progress in environments from the DeepMind Control Suite (DMControl) (Tassa et al., 2018; 2020) (Section 4.1) against several relevant baselines (Section 4.2) and show these results in Section 4.3. Also, we examine the characteristics of encoders and latent representations that KSL and other methods produce. We inspect the relationship between representations and reward, the robustness of encoders to perturbations in the state space, the ability of encoders to generalize to other tasks in the DMControl suite, and the temporal coherence of representations.

heading: 4.1 DMCONTROL
text: We examine the learning speed of agents by measuring performance within the popular PlaNet benchmark suite (Hafner et al., 2019). This suite contains environments from DMControl (Tassa et al., 2018) and has been widely used to measure the sample efficiency of RL agents on imagebased states. DMControl presents a wide variety of continuous-control tasks, including rigid-body control, control of loose appendages with generated momentum, and bipedal traversal. There are six tasks that range from one-dimensional to six-dimensional action spaces and agents in DMControl are exposed to both dense and sparse reward signals. See Appendix J for more details.
Two checkpoints of interest have emerged in recent literature (e.g., Laskin et al. (2020a)): the 100k steps mark and the 500k steps mark. The former is considered to be the data efficiency mark and the latter the asymptotic performance mark. The term “steps” refers to the number of environment transitions and not the number of learning updates. Each task has an episode length of 1k steps and has a specific number of action repeats. The agent’s chosen action is repeated several times, each of which results in an environment transition. For example, Cheetah, Run has an action repeat of four, which results in 250 actions selections per episode.

heading: 4.2 BASELINES
text: We compare KSL against six baseline methods in the data-efficiency experiment: RAD (Laskin et al., 2020b), DrQ (Yarats et al., 2021b), CURL (Laskin et al., 2020a), SAC-AE (Yarats et al., 2020), SAC State, and SAC Pixel. We specifically choose the first four methods because they were (i) state-of-the-art in the PlaNet benchmark suite at one point in time and (ii) all use the exact same encoder and SAC architecture as well as experimental design in PlaNet. Likewise, KSL uses the same encoder and SAC architecture as these four methods to ensure fair comparison. SAC State and SAC Pixel are commonly chosen as experimental control agents. SAC State learns on proprioceptive versions of the state space (e.g., vectors that contain “ground truth” information such as joint angles and velocity). SAC Pixel learns on the pixel version of the state space but with no image augmentation and no representation learning routine. We produce the results for baseline methods using code provided by the original authors. See Appendix D for a comparison between KSL and other methods that do not meet criteria (i) and (ii).

heading: 4.3 DATA EFFICIENCY EXPERIMENT
text: During training, all agents perform random actions for the first 1k steps to collect experience tuples for the replay buffer. Each agent type performs one gradient update per action selection by using randomly sampled experience tuples. We perform an agent evaluation every 20k steps. Agent evaluation uses a current snapshot of the learning agent and deterministically samples actions. The agent’s final score for each evaluation is the mean of total rewards across 10 episodes. For all methods, we use a batch size of 128. A full table of hyperparameters can be found in Appendix A. We also provide a wall clock time comparison between KSL, RAD, DrQ, CURL, and SAC-AE in Appendix C.
Figure 2 presents the evaluation curves for KSL, RAD, DrQ, CURL, and SAC-AE over 10 seeds. Table 1 highlights the 100k and 500k steps mark as averages ± one standard deviation, with best mean evaluation returns between image-based methods bolded. We use Welch’s t-test to probe for statistical significance between the highest average result and all other results. We highlight that KSL achieves the best average result in all six tasks at both checkpoints of interest and especially outperforms all other methods at the data efficiency checkpoint.

heading: 4.4 ANALYZING LATENT REPRESENTATIONS AND LEARNED ENCODERS
text: We analyze the latent representations and the learned encoders of all image-based methods from Section 4.2. We collect 10 episodes of state-reward pairs in the Walker, Walk task with a fullytrained agent, which we refer to as the evaluation states and rewards. We choose to focus on Walker, Walk as it is one of the more difficult tasks in PlaNet, as judged by the dimension of its action space, the number of action selections per episode, and the difficulty of learning bipedal movement. For the following analyses, we train encoders for each method in Walker, Walk for 100k steps across five seeds, saving the encoders’ weights every 5k steps. Specifically, we use the encoders that produce
representations for SAC’s critic Qθ and actor πφ, which is Φo for KSL. In this section, we refer to these encoders as the evaluation encoders.
Organization Around Reward: The reward function is the only supervised signal in RL, so it follows that learned representations of the state space should relate to reward. We first observe this relationship qualitatively. Figure 3 shows the first two components of PCA projections of latent representations produced by KSL (top row) and DrQ (bottom row). Here, we focus on the early stages of training to observe emergent behavior in learned representations. The left, middle, and right columns of plots use encoders after 5k, 10k, and 15k environment steps, respectively. Point colors denote the reward. We highlight that KSL’s plots show a clear separation between low-reward and high-reward states more quickly than DrQ’s plots. By 10k steps, KSL’s encoder has grouped the highest-reward states together and by 15k steps, high-reward and low-reward states seem to be linearly separable within the PCA projections. In contrast, DrQ’s projections show little sign of reward-based organization by 15k steps. All other baseline methods show the same pattern as DrQ. See Appendix L for PCA plots of other baseline methods. Despite the concrete connection between states and rewards in the context of RL, KSL is the only method that we observe discovering this relationship in the early stages of training.
We also quantify the relationship between learned representations and reward by measuring the mean squared error (MSE) of a linear regression that is fitted to estimate reward from latent projections. Figure 4 displays the mean MSE for all evaluation encoders over all evaluation states. We note
that KSL consistently produces representations that are better linear predictors of reward while also displaying lower variance from one environment step checkpoint to the next than other baseline methods. Also, we note that SAC Pixel produces results that are significantly worse than other baseline methods. This suggests that the RL signal, alone, is not enough to produce representations that are coherently tied to reward.
Perturbation Invariance: We measure the invariance of encoders to perturbations in the state space. Specifically, we measure the `2 distance between latent representations of several augmented versions of states. We produce three versions of each evaluation state per augmentation using random translation, random erasing, and Gaussian noise (N (0, .15)). Then, we project representations with all evaluation encoders for all augmented states. Figure 5 displays the average `2 distance for all evaluation encoders. We note that y-axes are scaled to highlight the curves of baseline methods, which results in SAC Pixel being off-screen. See Appendix M for un-scaled y-axes. These results show the importance of combining image augmentation and auxiliary tasks. Methods that use neither (SAC Pixel) show high sensitivity to all three augmentations. Methods that only use auxiliary tasks (SAC-AE) show a high sensitivity to geometric translations. Methods that use only augmentation (RAD and DrQ) sometimes show instability to the visual perturbance of erasing and Gaussian noise. Methods that use both augmentation and auxiliary tasks (CURL and KSL) show stability across all three augmentations. We highlight that KSL’s encoders are generally the most consistent. For example, the change in sensitivity from one evaluation checkpoint to the next is generally much lower for KSL than other baseline methods. See Appendix K for more details.
Generalization of Encoders: We test the ability of learned encoders to generalize their projections from one task to another. To this end, we first train each agent in each task for 100k steps five times and save the encoders’ weights. Then, we freeze each encoder and train a base SAC agent for 100k steps in every task other than the original training task. Table 2 presents the mean ± one
standard deviation of evaluation returns at 100k steps of training. We highlight that KSL’s encoders outperform other methods in five of the six tasks. See Appendix H for full evaluation curves.
Temporal Coherence: Control of robotic components, such as those similar to the objects in DMControl, involves important elements that move smoothly over time, such as joint angles. Ideally, our learned representations should capture this important information from the state space. Therefore, our representations should also move slowly and smoothly across time. We quantify this characteristic by measuring the `2 distance between the latent representations of consecutive state pairs in our evaluation set. Figure 6 displays the mean across all evaluation encoders. We highlight that KSL learns representations that move more slowly and consistently than the representations learned by other methods.

heading: 5 RELATED WORK

heading: 5.1 REPRESENTATION LEARNING AND REINFORCEMENT LEARNING
text: KSL builds off of the success of Bootstrap Your Own Latent (BYOL) (Grill et al., 2020) by adapting its architecture to work over the temporal axis of the RL problem’s underlying MDP. BYOL is a self-supervised representation learning method for images that does not require negative examples to avoid collapsed representations.
Adding auxiliary tasks to base RL algorithms is a common method to encourage the learning of useful representations. Yarats et al. (2020) introduce SAC-AE, which adds an auxiliary imagereconstruction task with a deterministic auto-encoder (Ghosh et al., 2020) and show that gradients from the actor harm the learning process of encoders. Laskin et al. (2020a) introduce CURL, which augments the RL loop with a projection head akin to contrastive predictive coding (van den Oord et al., 2018) with a momentum encoding procedure (e.g., Chen et al. (2020c)). Alternative tasks have been explored such as aligning the latent- and state-space with bisimulation metrics (Zhang et al., 2021), maximizing mutual temporal-information via learning on hidden activations (Anand et al., 2019), using contrastive learning for task generalization (Agarwal et al., 2021), and using latent dynamics models (Lee et al., 2020a; Hafner et al., 2020; 2021; Gelada et al., 2019; Guo et al., 2020).
Other approaches have relied on using tricks from computer vision, such as image augmentation. RAD (Laskin et al., 2020b) explores how simply augmenting incoming images can help SAC improve without any other alterations. DrQ (Yarats et al., 2021b), published concurrently with RAD,
extends RAD to include regularization via averaging Q-function values over several image augmentations during training.
KSL is similar to SPR (Schwarzer et al., 2021) in that both use an architecture similar to BYOL. However, there are several key differences. Unlike SPR, KSL explicitly separates the weight-update process of the RL task and the representation learning task. This separation results in independent tracking of accumulated gradient information within optimizers between tasks. Due to the overlap in parameters between tasks, this separation leads to differences in weight updates. We show in Appendix I that the update direction of the KSL and RL task often conflict throughout training. A high diversity of gradient directions across multiple tasks may be beneficial in the context of noni.i.d data streams (Aljundi et al., 2019), which is a common scenario in RL. Also, KSL and SPR have significant architectural differences. For one, SPR was designed for discrete-action settings and its architecture relies on one-hot encoded actions. In contrast, KSL’s architecture is general enough to be applied in both discrete- and continuous-action domains. Additionally, KSL’s transition module operates on vectors, and therefore it contains no convolutional layers. Unlike SPR, this allows KSL to fully share the entire spatial learning process between both the RL and representation learning task. Also, SPR shares more layers than just an encoder between the representation learning and RL task (see note in Figure 2 in Schwarzer et al. (2021)). We refer to this choice as “knowledge sharing”. We find that both knowledge sharing (Appendix E) and a convolutional transition model (Appendix G) are detrimental for KSL.

heading: 5.2 ANALYSIS OF LEARNED REPRESENTATIONS
text: Most of the focus in RL research is on developing algorithms that provide state-of-the-art performance. Comparatively, less attention has been spent on analyzing the learned representations themselves. Laskin et al. (2020b) qualitatively analyze spatial attention maps (Zagoruyko & Komodakis, 2017) of the learned encoder on various image augmentations. Zhang et al. (2021) provide plots of t-SNE-projected representations to show organization around learned state values. Hafner et al. (2019), and Yarats et al. (2020) observe the relationship between learned representations and the proprioceptive versions of states. We refer the interested reader to the survey work of Bengio et al. (2013), which provides a thorough overview of hypotheses around the question, “what makes for a good representation, in general?”.

heading: 6 CONCLUSION
text: In this work, we addressed the problem of improving the sample efficiency of reinforcement learning (RL) agents for continuous control in image-based state spaces. To this end, we introduced k-Step Latent (KSL), an auxiliary learning task for RL agents. KSL leverages advancements in representation learning and data augmentation to produce state-of-the-art results in the PlaNet benchmark suite. We also analyzed the encoders and latent representations produced by KSL and other baseline methods. These analyses showed that KSL produces encoders that are more robust to perturbations in the state space and generalize to other tasks better than other baseline methods. Also, the analyses showed that KSL’s representations are more strongly tied to reward and move more steadily across the temporal axis of the RL problem’s underlying MDP. Despite the sample-efficiency gains of KSL, it requires increased wall clock time compared to previous methods. It may be possible to reduce this burden via parallel processing or by reducing the size of KSL’s neural networks.

heading: 7 REPRODUCIBILITY STATEMENT
text: Please see the footnote on page 2 for a link to an anonymized code repository. Also, see Appendix A for implementation details such as hyperparameters and architecture choices. Finally, see Appendix B for KSL’s pseudocode.

heading: B KSL PSEUDOCODE
text: Algorithm 1 presents a PyTorch-like pseudocode for KSL. For clarity, we use the same mathematical symbols here as in the main body of the paper. Also, we include PyTorch-like tensor-indexing for clarity. For example, line 3 shows s[:, 1, :, :, :], which means all batches, 1 out of k, all color
channels across the stack of images, all pixels in the width, all pixels in the height. There is one minor difference between the notation in Algorithm 1 and the notation in the main body of the paper. To allow for the for loop, z′o is kept as zo in the pseudocode.
Algorithm 1: PyTorch-like pseudocode for KSL’s loss. 1 Input : Trajectories τ = {(s1, a1, ..., ak−1, sk)i}Mi=1 sampled from replay buffer with augmented states 2 loss← 0 3 z0 ← Φo(s[:, 1, :, :, :]) 4 for j in 1...k − 1 do 5 zm ← Φm(s[:, j + 1, :, :, :]) 6 zo ← T ([zo|a[:, j]]) 7 qo ← Ψo(zo) 8 ym ← Ψm(zm).detach() 9 yo ← P(qo)
10 loss← loss + ‖ yo‖yo‖ − ym ‖ym‖‖ 2 2.mean() 11 end 12 ksl optimizer.zero grad() 13 loss.backward() 14 ksl optimizer.step()

heading: C WALL CLOCK TIME
text: To compare wall clock time, we average the total number of seconds to complete 15k episodes of training of Cartpole, Swingup across three seeds. Table 4 shows this average as a value that is normalized to SAC Pixel. All metrics were collected on a machine with a 1080Ti GPU, 8-core AMD Threadripper 1900x @ 3.8GHz, and 128GB of RAM. Codebases for all methods are written in PyTorch.

heading: D COMPARISON OF KSL TO PI-SAC AND PROTO-RL
text: Here, we include a performance comparison between KSL, PI-SAC (Lee et al., 2020b) and ProtoRL (PRL) (Yarats et al., 2021a). PI-SAC uses both a forwards and backwards encoder to learn a compressed representation of image-based states that maximizes the predictive information between representations and future states while minimizing useless information in the state space. We choose to not include this method in Section 4.2 due to significant differences in architecture, hyperparameters, and experiment settings in the PI-SAC paper and papers of methods presented in Section 4.2. For one, PI-SAC’s encoder and SAC network architectures are different than those in Section 4.2. Even small differences in network architecture can explain significant differences in performance (Islam et al., 2017; Henderson et al., 2018). Second, PI-SAC uses a significant amount of pre-training to achieve its reported results, while all methods presented in Section 4.2 do not. For example, in Cheetah, Run, Lee et al. (2020b) use 10k pretraining steps, which results in PISAC agents receiving the same number of gradient updates that the agents described in Section 4.2 would only receive by the 40k step mark. Thirdly, the hyperparameters used for tasks (e.g., the action repeat) in the PI-SAC paper are not the same as the ones used in previous PlaNet studies. For example, Lee et al. (2020b) use an action repeat of 4 for Cartpole, Swingup, while all methods in Section 4.2 use 8. This results in PI-SAC agents receiving 2x the number of gradient updates for any given environment-steps checkpoint as our baseline methods.
PRL introduces an auxiliary task that encourages intrinsically-motivated agent exploration. PRL agents are trained in two stages. First, the state encoder is pre-trained via a task-agnostic method that attempts to maximize the entropy of the distribution of agent-visited states. Second, the state encoder is frozen and a standard SAC agent is trained on top of the learned representations. We choose to not include this method in Section 4.2 as there are differences in architecture and evaluation procedure between the PRL paper and the papers of methods shown in Section 4.2. For one, PRL’s
encoder produces larger latent vectors than the encoders’ of methods in Section 4.2 (50 vs. 128). Second, Yarats et al. (2021a) use an action repeat of two for all tasks and evaluate agents over 1M steps of training. The evaluations in PRL’s paper are done in a two-part 500k/500k split. First, 500k steps of auxiliary-task pre-training. Second, 500k steps of RL-task training. We keep this one-toone ratio in our evaluations. For the 100k steps checkpoint, we train PRL agents in a 50k/50k split. Similarly, we train PRL agents in a 250k/250k split for the 500k steps checkpoint.
We adjust the encoder and SAC networks of PI-SAC and PRL to be congruent with methods in Section 4.2, all hyperparameters to be congruent with Table 3 in Appendix A, and all environment settings to be congruent with Table 7 in Appendix J. Otherwise, we keep PI-SAC and PRL in their original state. We make these adjustments within code provided by the original authors. Table 5 shows the evaluation returns across 10 seeds using the same methodology as described in Section 4.3. Figure 7 shows the evaluation curves. We highlight that KSL agents outperform both PI-SAC and PRL agents in all six tasks for both the 100k and 500k steps mark.

heading: E KNOWLEDGE SHARING AND CHOICE OF K
text: Figure 8 shows potential options for the knowledge sharing (KS) mechanism within KSL. On the far left, we show the Q-function of the critic Qθ. Although the KS mechanism uses layers from both Q-functions, we only show one in the “Q-function layers” diagram for simplicity. When l = 1, only the first layer of the Q-functions are used for KS, and when l = 2, the first two layers are used. When l = 0, the KS section of KSL’s transition module T is made of a single linear layer that is not part of the Critic’s Q-functions.
The choice of k creates a trade-off in terms of algorithm performance and compute expense. A larger k gives KSL access to deeper levels of temporal information. However, a larger k results in a larger amount of compute, both in terms of data sampling and the recurrent loop within KSL. With every incremental k, KSL needs to sample and perform the translation augmentation on batch
size more states from the replay memory. Then, each additional state needs to perform several additional passes through neural network modules as well as compute a loss, therefore leading to an exponential growth in compute requirements.
We perform a full grid-sweep across all combinations of k ∈ {1, 3, 5} and l ∈ {0, 1, 2} in the task of Cheetah, Run. Figure 9, below, shows the mean (bold line) and one standard deviation (shaded area) across five random seeds. Table 6 summarizes these results for the 100k and 500k checkpoints.
These results show that knowledge sharing provides no benefit and may actually harm the RL learning process. In general, when k > 1, l > 0 leads to lower levels of performance with higher variance between runs. Also, the results show that there is little reason to go beyond k = 3. As such, the main results in the paper use l = 0 and k = 3.

heading: F HOW KSL FITS WITHIN SAC
text: Figure 10 shows how KSL fits within the greater SAC algorithm. Φo takes in current states st and creates latent projections zo for both the critic Qθ (blue) and the actor πφ (red). Φm takes in future states st+1 and creates latent projections zm for the target critic Qθ̄ (green).

heading: G CONVOLUTIONAL TRANSITION MODEL
text: We test the performance of a convolutional transition model T within KSL. To make a convolutional T fit within KSL, we adjust its architecture. First, we remove the final dense layer from the encoders Φ. T is then made of two convolutional layers with 32 and 16 filters, respectively, and both have a 3×3 kernel with a stride of two. Following the convolutional layers is a single dense layer with 1024 +dim(A) hidden units that outputs a vector of length 50, followed by a layer norm operation. Action vectors are concatenated to the flattened feature maps from the transition model’s convolutional layers before being fed through its dense layer. Both the critics and the actor have two convolutional layers added to their networks, all with the same architecture as described here for T . Additionally, they contain one dense layer with 1024 hidden units that outputs a vector of length 50. The remaining architectures of Φm, Φo, P , Qθ, Qθ̄, and πφ remain unchanged. The additional layers in the critic target are updated as an EMA of their counterparts in the critic.
Figure 11 shows the evaluation results for both KSL with (five seeds) and without (ten seeds) a convolutional T . We highlight that the version of KSL with a convolutional T performs significantly worse than the version of KSL with a fully-dense T . Also, we notice a significantly higher probability of representational collapse in the version of KSL with a convolutional T . As a result, we use the version of KSL with a fully-dense T for all other results in this paper.

heading: H GENERALIZATION LEARNING CURVES
text: Figure 12 shows the learning curves for the generalization experiment.

heading: I OPTIMIZERS VS. TASKS
text: The Adam optimizer (Kingma & Ba, 2015) performs optimization updates by using exponentiallydecayed first moment m̂ and second moment v̂ estimates of historical gradients. Adam updates parameters θ at timestep t as
θt+1 ← θt − α Update direction︷ ︸︸ ︷ m̂t√ v̂t + , (5)
where is a small constant.
In KSL’s formulation, the RL task and the representation learning task both send feedback from their loss functions into Φo. Two unique situations can arise with the optimization process in KSL: (i) a single optimizer handling updates from both tasks or (ii) separate optimizers handling updates for either task. The differences in optimizer-held gradient statistics results in different optimization steps for Φo in either scenario. Using the same optimizer for both tasks may result in updates that conflate the accumulated gradient statistics of Φo between tasks.
To confirm this, we perform two evaluations. First, we measure the learning performance of KSL agents for scenario (i) and (ii) in both Cheetah, Run and Walker, Walk. Figure 13, below, shows the average and one standard deviation across five seeds. We highlight that KSL agents that use separate optimizers for either task result, on average, in better learning performance. Also, we note that these results are statistically significant at p = 0.05.
Second, we query m̂ and v̂ from either optimizer in case (ii) throughout training for all tasks in the PlaNet benchmark suite. We extract the accumulated statistics for Φo and measure the angle φij = 〈gi, gj〉/‖gi‖‖gj‖ where gi and gj are the “update direction”, as denoted in Equation 5, for the RL- and KSL-task optimizers, respectively. When φij < 0, task-specific updates are considered to be conflicting. When gradients are conflicting, updates move weights in “opposing” directions. Figure 14 depicts φij for all tasks in PlaNet over three runs of training. We highlight that, on average, φij < 0 for all tasks. This implies that maintaining a combined m̂ and v̂ for both the RL and KSL task conflates update information from both tasks.
The RL task is unlike other optimization problems in that its sequential nature and non-i.i.d data create a moving target. Given that KSL’s auxiliary task improves the performance on the RL task, we hypothesize that the pull from KSL’s conflicting update helps to regularize learning of the RL task. Also, the independent tracking of gradient statistics helps to increase the diversity of gradients, which has been shown to improve learning in streams of non-i.i.d data (Aljundi et al., 2019).
Alternatives to our multi-optimizer choice may include altering the gradients of the tasks before they are combined and passed to the optimizer (e.g., Yu et al. (2020)), changing the architecture of Φo to include task-specific attention (Liu et al., 2019), or applying additional learning procedures that control the learning schedule between tasks (e.g., Xu et al. (2019)).

heading: J ENVIRONMENT DETAILS
text: Table 7 describes the dimensions of the action space, the number of action repeats, and the type of reward (dense vs. sparse) for each task in the PlaNet benchmark suite. Figure 15 depicts the six tasks within the PlaNet benchmark suite.

heading: K EXTENDED ENCODER SENSITIVITY
text: We measure the change in sensitivity to perturbations as the absolute difference between one encoder checkpoint to the next. A lower value is desireable, as it indicates smooth changes in latent representations throughout the training process. Smooth changes in latent representations results in less variability in input for the downstream RL training task. Best results are bolded. Table 8 shows this metric.

heading: L EXTENDED PCA PLOTS
text: Here we provide PCA plots for all methods not shown in the main body of the paper. Figure 16, Figure 17, Figure 18, Figure 19 show plots for SAC Pixel, RAD, CURL, and SAC-AE, respectively.

heading: M EXTENDED AUGMENTATION PLOTS
text: Figure 20 shows plots for the random erasing augmentation and Figure 21 shows plots for the Gaussian noise augmentation.

Review summaries:
The paper tackles the problem of sample inefficiency in continuous control, by noting that standard RL methods deal with both policy optimisation and representation learning jointly with a single supervisory signal, namely the reward. Consequently, authors propose to leverage long-term temporal connections between actions in the representation learning, and introduce k-Step Latent (KSL), a representation learning module for learning temporally consistent representations of the state space. Authors show that KSL improves previous state-of-the-art methods in PlaNet benchmark suite and provide some analysis of the representations learned by KSL.

This paper proposes a representation-learning method (k-step latent, or KSL) that uses a self-supervised auxiliary loss between recurrently-predicted action-conditioned representations of the state space and non-recurrently predicted target representations, in the style of BYOL. The method is trained using two separate optimizers on different parts of the models, to avoid interference in the statistics maintained by the optimizers. Results at 100k and 500k steps on 6 tasks from the DM control suite (from pixels) compared to methods using alternative self-supervised auxiliary losses show that the proposed method improves data efficiency. Analysis of the learned latent representations shows that those from KSL produce more robust encoders and are more consistent with the underlying MDP.

Summary: The authors introduced k-Step Latent (KSL), a representation learning method for visual-based continuous control tasks. KSL utilises multi-step latent action-dependent predictive supervision for training the representation. The empirical evaluations are based on the dm-control suite benchmarks, KSL demonstrates improved sample efficiency (100k evaluation) and asymptotic performance (500k evaluation) across the six tasks presented, comparing to the baseline algorithms (mainly based on image augmentation). The authors further empirically examined the properties of the learned representations, and showed that the trained encoder quickly learns to be representative of the reward structure. The authors also argue that KSL supports more robust representation learning and stronger generalisability.