Title: $\texttt{TACO}$: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Deep Reinforcement Learning

Visual Reinforcement Learning

Online Visual RL

Offline Visual RL

Action Representation
Abstract: Despite recent progress in reinforcement learning (RL) from raw pixel data, sample inefficiency continues to present a substantial obstacle. 
Prior works have attempted to address this challenge by creating self-supervised auxiliary tasks, aiming to enrich the agent's learned representations with control-relevant information for future state prediction.
However, these objectives are often insufficient to learn representations that can represent the optimal policy or value function, and they often consider tasks with small, abstract discrete action spaces and thus overlook the importance of action representation learning in continuous control.
In this paper, we introduce $\texttt{TACO}$: $\textbf{T}$emporal $\textbf{A}$ction-driven $\textbf{CO}$ntrastive Learning, a simple yet powerful temporal contrastive learning approach that facilitates the concurrent acquisition of latent state and action representations for agents. 
$\texttt{TACO}$ simultaneously learns a state and an action representation by optimizing the mutual information between representations of current states paired with action sequences and representations of the corresponding future states. 
Theoretically, $\texttt{TACO}$ can be shown to learn state and action representations that encompass sufficient information for control, thereby improving sample efficiency.
For online RL, $\texttt{TACO}$ achieves 40% performance boost after one million environment interaction steps on average across nine challenging visual continuous control tasks from Deepmind Control Suite. 
In addition, we show that $\texttt{TACO}$ can also serve as a plug-and-play module adding to existing offline visual RL methods to establish the new state-of-the-art performance for offline visual RL across offline datasets with varying quality.

TACO : Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning

Abstract

Despite recent progress in reinforcement learning (RL) from raw pixel data, sample inefficiency continues to present a substantial obstacle. Prior works have attempted to address this challenge by creating self-supervised auxiliary tasks, aiming to enrich the agent's learned representations with control-relevant information for future state prediction. However, these objectives are often insufficient to learn representations that can represent the optimal policy or value function, and they often consider tasks with small, abstract discrete action spaces and thus overlook the importance of action representation learning in continuous control. In this paper, we introduce TACO : T emporal A ction-driven CO ntrastive Learning, a simple yet powerful temporal contrastive learning approach that facilitates the concurrent acquisition of latent state and action representations for agents. TACO simultaneously learns a state and an action representation by optimizing the mutual information between representations of current states paired with action sequences and representations of the corresponding future states. Theoretically, TACO can be shown to learn state and action representations that encompass sufficient information for control, thereby improving sample efficiency. For online RL, TACO achieves 40% performance boost after one million environment interaction steps on average across nine challenging visual continuous control tasks from Deepmind Control Suite. In addition, we show that TACO can also serve as a plug-and-play module adding to existing offline visual RL methods to establish the new state-of-the-art performance for offline visual RL across offline datasets with varying quality.

1 Introduction

Developing reinforcement learning (RL) agents that can perform complex continuous control tasks from high-dimensional observations, such as pixels, has been a longstanding challenge. A central aspect of this challenge lies in addressing the sample inefficiency problem in visual RL. Despite significant progress in recent years [33, 34, 53, 52, 16, 15, 17, 18, 21, 48], there remains a considerable gap in sample efficiency between RL with physical state-based features and those with pixel inputs. This disparity is particularly pronounced in complex tasks; thus tackling it is crucial for advancing visual RL algorithms and enabling their practical application in real-world scenarios.

State representation learning has become an essential aspect of RL research, aiming to improve sample efficiency and enhance agent performance. Initial advancements like CURL [33] utilized a self-supervised contrastive InfoNCE objective [47] for state representation, yet it overlooks the temporal dynamics of the environment. Subsequent works, including CPC [23], ST-DIM [2], and ATC [44], made progress to rectify this by integrating temporal elements into the contrastive loss, linking pairs of observations with short temporal intervals. The objective here was to develop a state representation capable of effectively predicting future observations. A more comprehensive approach was taken by DRIML [37], which incorporated the first action of the action sequence into the temporal contrastive learning framework. However, these methods, while innovative, have their shortcomings. The positive relations in their

CURL

DrQ

Deepmind Control Suite 1M

(9 Challenging Tasks)

CURL

DrQ

DrQ-v2

TACO

Model-free

Model-based

contrastive loss designs are often policy-dependent, potentially leading to instability during policy updates throughout the training process. Consequently, they lack the theoretical foundation needed to capture all information representing the optimal policy. Furthermore, these methods except for CURL and ATC typically focus on environments such as Atari games with well-represented, abstract discrete action spaces, thereby overlooking the importance of action representation in continuous control tasks [1, 28]. However, by learning an action representation that groups semantically similar actions together in the latent action space, the agent can better generalize its knowledge across various state-action pairs, enhancing the sample efficiency of RL algorithms. Therefore, learning both state and action representations is crucial for enabling the agent to more effectively reason about its actions' long-term outcomes for continuous control tasks.

In this paper, we introduce T emporal A ction-driven CO ntrastive Learning ( TACO ) as a promising approach to visual continuous control tasks. TACO simultaneously learns a state and action representation by optimizing the mutual information between representations of current states paired with action sequences and representations of the corresponding future states. By optimizing the mutual information between state and action representations, TACO can be theoretically shown to capture the essential information to represent the optimal value function. In contrast to approaches such as DeepMDP [12] and SPR [42], which directly model the latent environment dynamics, our method transforms the representation learning objective into a self-supervised InfoNCE objective. This leads to more stable optimization and requires minimal hyperparameter tuning efforts. Consequently, TACO yields expressive and concise state-action representations that are better suited for high-dimensional continuous control tasks.

We demonstrate the effectiveness of representation learning by TACO through extensive experiments on the DeepMind Control Suite (DMC) in both online and offline RL settings. TACO is a flexible plug-and-play module that could be combined with any existing RL algorithm. In the online RL setting, combined with the strong baseline DrQ-v2[52], TACO significantly outperforms the SOTA model-free visual RL algorithms, and it even surpasses the strongest model-based visual RL baselines such as Dreamer-v3 [18]. As shown in Figure 1, across nine challenging visual continuous control tasks from DMC, TACO achieves a 40% performance boost after one million environment interaction steps on average. For offline RL, TACO can be combined with existing strong offline RL methods to further improve performance. When combined with TD3+BC [11] and CQL [31], TACO outperforms the strongest baselines across offline datasets with varying quality.

We list our contributions as follows:

We present TACO , a simple yet effective temporal contrastive learning framework that simultaneously learns state and action representations.

The framework of TACO is flexible and could be integrated into both online and offline visual RL algorithms with minimal changes to the architecture and hyperparameter tuning efforts.

Wetheoretically show that the objectives of TACO is sufficient to capture the essential information in state and action representation for control.

Empirically, we show that TACO outperforms prior state-of-the-art model-free RL by 1.4x on nine challenging tasks in Deepmind Control Suite. Applying TACO to offline RL with SOTA

algorithms also achieves significant performance gain in 4 selected challenging tasks with pre-collected offline datasets of various quality.

2 Preliminaries

2.1 Visual reinforcement learning

Let M = ⟨S , A , P , R , γ ⟩ be a Markov Decision Process (MDP). Here, S is the state space, and A is the action space. The state transition kernel is denoted by P : S × A → ∆( S ) , where ∆( S ) is a distribution over the state space. R : S × A → R is the reward function. The objective of the Reinforcement Learning (RL) algorithm is the identification of an optimal policy π ∗ : S → ∆( A ) that maximizes the expected value E π [ ∑ ∞ t =0 γ t r t ] . Additionally, we can define the optimal Q function as follows: Q ∗ ( s, a ) = E π ∗ [ ∑ ∞ t =0 γ t r ( s t , a t ) | s 0 = s, a 0 = a ] , such that the relationship between the optimal Q function and optimal policy is π ( s ) = arg max a Q ∗ ( s, a ) . In in the domain of visual RL, high-dimensional image data are given as state observations, so the simultaneous learning of both representation and control policy becomes the main challenge. This challenge is exacerbated when the environment interactions are limited and the reward signal is sparse.

2.2 Contrastive learning and the InfoNCE objective

Contrastive Learning, a representation learning approach, imposes similarity constraints on representations, grouping similar/positive pairs and distancing dissimilar/negative ones within the representation space. Contrastive learning objective is often formulated through InfoNCE loss [47] to maximize the mutual information between representations of positive pairs by training a classifier. In particular, let X,Y be two random variables. Given an instance x ∼ p ( x ) , and a corresponding positive sample y + ∼ p ( y | x ) as well as a collection of Y = { y 1 , ..., y N -1 } of N -1 random samples from the marginal distribution p ( y ) , the InfoNCE loss is defined as

Optimizing this loss will result in f ( x, y ) ∝ p ( y | x ) p ( y ) and one can show that InfoNCE loss upper bounds the mutual information, L N ≥ log( N ) -I ( X,Y ) .

3 TACO : temporal action-driven contrastive Loss

TACO is a flexible temporal contrastive framework that could be easily combined with any existing RL algorithms by interleaving RL updates with its temporal contrastive loss update. In this section, we first presnt the overall learning objective and theoretical analysis of TACO . Then we provide the architectural design of TACO in detail.

3.1 Temporal contrastive learning objectives and analysis

In the following, we present the learning objectives of TACO . The guiding principle of our method is to learn state and action representations that capture the essential information about the environment's dynamics sufficient for learning the optimal policy. This allows the agent to develop a concise and expressive understanding of both its current state and the potential effects of its actions, thereby enhancing sample efficiency and generalization capabilities.

Let S t , A t be the state and action variables at timestep t , Z t = ϕ ( S t ) , U t = ψ ( A t ) be their corresponding representations. Then, our method aims to maximize the mutual information between representations of current states paired with action sequences and representations of the corresponding future states:

Here, K ≥ 1 is a fixed hyperparameter for the prediction horizon. In practice, we estimate the lower bound of the mutual information by the InfoNCE loss, with details of our practical implementation described in §3.2.

We introduce the following theorem extending from Rakely et al. [41] to demonstrate the sufficiency of TACO objective:

Theorem 3.1. Let K ∈ N + , and J TACO = I ( Z t + K ; [ Z t , U t , ..., U t + K -1 ]) . If for a given state and action representation ϕ Z , ψ U , J TACO is maximized, then for arbitrary state-action pairs ( s 1 , a 1 ) , ( s 2 , a 2 ) such that ϕ ( s 1 ) = ϕ ( s 2 ) , ψ ( a 1 ) = ψ ( a 2 ) , it holds that Q ∗ ( s 1 , a 1 ) = Q ∗ ( s 2 , a 2 ) .

This theorem guarantees that if our mutual information objective Equation (2) is maximized, then for any two state-action pairs ( s 1 , a 1 ) and ( s 2 , a 2 ) with equivalent state and action representations, their optimal action-value functions, Q ∗ ( s 1 , a 1 ) and Q ∗ ( s 2 , a 2 ) , will be equal. In other words, maximizing this mutual information objective ensures that the learned representations are sufficient for making optimal decisions.

3.2 TACO implementation

𝒔

𝒕

𝒂

𝒕

𝒕"𝑲

𝒛

𝒕

𝒖

𝒕

Action

Enc.

𝒛

𝒕"𝑲

…

Q Net

Proj. Layer

𝐺

Proj. Layer

𝐻

Contrastive

Loss

𝑄(𝑠

!

, 𝑎

!

)

…

𝒖

𝒕"𝑲$𝟏

Action

Enc.

State

Enc.

𝒂

𝒕"𝑲$𝟏

Action

Enc.

State

Enc.

𝒔

Here we provide a detailed description of the practical implementation of TACO . In Figure 2, we illustrate the architecture design of TACO . Our approach minimally adapts a base RL algorithm by incorporating the temporal contrastive loss as an auxiliary loss during the batch update process. Specifically, given a batch of state and action sequence transitions { ( s ( i ) t , [ a ( i ) t , ..., a ( i ) t ′ -1 ] , s ( i ) t ′ ) } N i =1 ( t ′ = t + K ), we optimize:

Here let z ( i ) t = ϕ ( s ( i ) t ) , u ( i ) t = ψ ( a ( i ) t ) be state and action embeddings respectively. g ( i ) t = G θ ( z ( i ) t , u ( i ) t , ..., u ( i ) t ′ -1 ) , and h ( i ) t = H θ ( z ( i ) t ′ ) , where G θ and H θ denote two learnable projection layers that map the latent state z ( i ) t as well as latent state and action sequence ( z ( i ) t , u ( i ) t , ..., u ( i ) t ′ -1 ) to a common contrastive embedding space. W is a learnable parameter providing a similarity measure between g i and h j in the shared contrastive embedding space. Subsequently, both state and action representations are fed into the agent's Q network, allowing the agent to effectively reason about the longterm effects of their actions and better leverage their past experience through state-action abstractions.

In addition to the main TACO objective, in our practical implementation, we find that the inclusion of two auxiliary objectives yields further enhancement in the algorithm's overall performance. The first is the CURL [33] loss:

̸

Here, h ( i ) t + = H θ ( ϕ ( s ( t ) i + )) , where s ( t ) i + is the augmented view of s ( t ) i by applying the same random shift augmentation as DrQ-v2 [52]. W and H θ share the same weight as the ones in TACO objective. The third objective is reward prediction:

Here r ( i ) = ∑ t ′ -1 j = t r ( i ) j is the sum of reward from timestep t to t ′ -1 , and ˆ R θ is a reward prediction layer. For our final objective, we combine the three losses together with equal weight. As verified in Section 4.1, although TACO serves as the central objective that drives notable performance improvements, the inclusion of both CURL and reward prediction loss can further improve the algorithm's performance.

We have opted to use DrQ-v2 [52] for the backbone algorithm of TACO , although in principle, TACO could be incorporated into any visual RL algorithms. TACO extends DrQ-v2 with minimal additional hyperparameter tuning. The only additional hyperparameter is the selection of the prediction horizon K . Throughout our experiments, we have limited our choice of K to either 1 or 3, depending on the nature of the environment. We refer the readers to Appendix A for a discussion on the choice of K .

4 Experiments and results

This section provides an overview of our empirical evaluation, conducted in both online and offline RL settings. To evaluate our approach under online RL, we apply TACO to a set of nine challenging visual continuous control tasks from Deepmind Control Suite (DMC) [46]. Meanwhile, for offline RL, we combine TACO with existing offline RL methods and test the performance on four DMC tasks, using three pre-collected datasets that differ in the quality of their data collection policies.

4.1 Comparison between TACO and strong baselines in online RL tasks

Environment Settings : In our online RL experiment, we first evaluate the performance of TACO on nine challenging visual continuous control tasks from Deepmind Control Suite [46]: Quadruped Run, Quadruped Walk, Hopper Hop, Reacher Hard, Walker Run, Acrobot Swingup, Cheetah Run, Finger Turn Hard, and Reach Duplo. These tasks demand the agent to acquire and exhibit complex motor skills and present challenges such as delayed and sparse reward. As a result, these tasks have not been fully mastered by previous visual RL algorithms, and require the agent to learn an effective policy that balances exploration and exploitation while coping with the challenges presented by the tasks.

In addition to Deepmind Control Suite, we also present the results of TACO on additional six challenging robotic manipulation tasks from Meta-world [57]: Hammer, Assembly, Disassemble, Stick Pull, Pick Place Wall, Hand Insert. Unlike the DeepMind Control Suite, which primarily concentrates on locomotion tasks, Meta-world domain provides tasks that involve complex manipulation and interaction tasks. This sets it apart by representing a different set of challenges, emphasizing precision and control in fine motor tasks rather than broader locomotion skills. In Appendix G, we provide a visualization for each Meta-world task.

Baselines : We compare TACO with four model-free visual RL algorithms CURL [33], DrQ [53], DrQ-v2 [52], and A-LIX [5]. A-LIX builds on DrQ-v2 by adding adaptive regularization to the encoder's gradients. While TACO could also extend A-LIX , our reproduction of results from its opensource implementation does not consistently surpass DrQ-v2 . As such, we do not choose A-LIX as the backbone algorithm for TACO . Additionally, we compare with two state-of-the-art model-based RL algorithms for visual continuous control, Dreamer-v3 [18] and TDMPC [21], which learn world models in latent space and select actions using either model-predictive control or a learned policy.

TACO achieves a significantly better sample efficiency and performance compared with SOTA visual RL algorithm. The efficacy of TACO is evident from the findings presented in Figure 4 (DMC) , Table 1 (DMC), and Figure 5 (Meta-world). In contrast to preceding model-free visual RL algorithms, TACO exhibits considerably improved sample efficiency. For example, on the challenging Reacher Hard task, TACO achieves optimal performance in just 0.75 million environment steps, whereas DrQv2 requires approximately 1.5 million steps. When trained with only 1 million environment steps, TACO on average achieves 40% better performance, and it is even better than the model-based visual RL algorithms Dreamer-v3 and TDMPC on 6 out of 9 tasks. In addition, on more demanding tasks such as Quadruped Run, Hopper Hop, Walker Run, and Cheetah Run, TACO continues to outshine competitors, exhibiting superior overall performance after two or three million steps, as illustrated in Figure 4. For robotic manipulation tasks, as shown in Figure 5, TACO also significantly outperform the baseline model-free visual RL algorithms, highlighting the broad applicability of TACO .

0m

0.6m

1.2m

1.8m

2.4m

3m

100

200

300

400

500

600

700

800

0m

0.6m

1.2m

1.8m

2.4m

3m

0

50

100

150

200

250

300

350

400

0m

0.6m

1.2m

1.8m

2.4m

3m

Number of Environment Steps

0

100

200

300

400

500

600

700

800

0m

0.4m

0.8m

1.2m

1.6m

2m

100

200

300

400

500

600

700

800

900

Episode Reward

0m

0.4m

0.8m

1.2m

1.6m

2m

0

200

400

600

800

Episode Reward

0m

0.4m

0.8m

1.2m

1.6m

2m

Number of Environment Steps

0

200

400

600

800

1000

Episode Reward

Finger Turn Hard

0m

0.4m

0.8m

1.2m

1.6m

2m

0

50

100

150

200

250

300

350

400

0m

0.4m

0.8m

1.2m

1.6m

2m

0

200

400

600

800

1000

Reacher Hard

0m

0.4m

0.8m

1.2m

1.6m

2m

Number of Environment Steps

0

50

100

150

200

Reach Duplo

TACO

A-LIX

DrQ-v2

Concurrently learning state and action representation is crucial for the success of TACO . To demonstrate the effectiveness of action representation learning in TACO , we evaluate its performance on a subset of 4 difficult benchmark tasks and compare it with a baseline method without action representation, as shown in Figure 3. The empirical results underscore the efficacy of the temporal contrastive learning objectives, even in the absence of action representation. For instance, TACO records an enhancement of 18% on Quadruped Run and a substantial 51% on Reacher Hard, while the remaining tasks showcase a performance comparable to DrQ-v2. Furthermore, when comparing against TACO without action representation, TACO achieves a

Reacher Hard

200

400

600

800

Episode Reward

TACO

TACO w/o action representation

DrQv2

consistent performance gain, ranging from 12.2% on Quadruped Run to a significant 70.4% on Hopper Hop. These results not only emphasize the inherent value of the temporal contrastive learning objective in TACO , but also underscore the instrumental role of high-quality action representation in bolstering the performance of the underlying RL algorithms.

TACO learns action representations that group semantically similar actions together. To verify that indeed our learned action representation has grouped semantically similar actions together, we conduct an experiment within the Cheetah Run task. We artificially add 20 dimensions to the action

Hammer

Assembly

Disassemble

100%

100%

80%

Task Sucess Rate

80%

80%

60%

60%

60%

40%

40%

40%

20%

20%

20%

0%

0%

0%

0m

0.4m

0.8m

1.2m

1.6m

2m

0m

0.4m

0.8m

1.2m

1.6m

2m

0m

0.4m

0.8m

1.2m

1.6m

2m

StickPull

PickPlaceWall

Hand Insert

100%

80%

Task Sucess Rate

80%

60%

80%

60%

60%

40%

40%

40%

20%

20%

20%

0%

0%

0%

0m

0.4m

0.8m

1.2m

1.6m

2m

0m

0.4m

0.8m

1.2m

1.6m

2m

0m

0.4m

0.8m

1.2m

1.6m

2m

Number of Environment Steps

TACO

A-LIX

DrQ-v2

space of task Cheetah Run, although only the first six were utilized in environmental interactions. We first train an agent with TACO online to obtain the action representation. Then we select four actions within the original action space a 1 , a 2 , a 3 , a 4 to act as centroids. For each of the four centroids, we generate 1000 augmented actions by adding standard Gaussian noises to the last 20 dimensions. We aim to determine if our action representation could disregard these 'noisy' dimensions while retaining the information of the first six. Using t-SNE for visualization, we embed the 4000 actions before and after applying action representation. As shown in Figure 6, indeed our learned action representation could group the four clusters, demonstrating the ability of our action representation to extract control relevant information from the raw action space.

The effectiveness of our temporal contrastive loss is enhanced with a larger batch size. As is widely acknowledged in contrastive learning research [8, 22, 40, 13], our contrastive loss sees significant benefits from utilizing a larger batch size. In Figure 7a, we illustrate the performance of our algorithms alongside DrQv2 after one million environment steps on the Quadruped Run task. As evident from the plot, batch size greatly influences the performance of our algorithm, while DrQ-v2's baseline performance remains fairly consistent throughout

t-SNE Embedding of actions

t-SNE Embedding of action representations

training. In order to strike a balance between time efficiency and performance, we opt for a batch size of 1024, which is 4 times larger than the 256 batch size employed in DrQ-v2, but 4 times smaller than the 4096 which is commonly used in the contrastive learning literature [8, 22, 13]. For an analysis on how batch size affects the algorithm's runtime, we direct the reader to Appendix B.

256

512

1024

2048

200

300

400

500

600

Episode Reward

TACO

DrQ-v2

300

350

400

450

500

550

600

Episode Reward

TACO

w/o

Reward

w/o

CURL

w/o

TACO

CURL

Reward prediction and CURL loss serves an auxiliary role to further improve the performance of TACO , while the temporal contrastive loss of TACO is the most crucial component. In the practical deployment of TACO , two additional objectives, namely reward prediction and CURL loss,

are incorporated to enhance the algorithm's performance. In Figure 7b, we remove one objective at a time on the Quadruped Run task to assess its individual impact on the performance after one million environment steps. As illustrated in the figure, the omission of TACO ' temporal contrastive objective results in the most significant performance drop, emphasizing its critical role in the algorithm's operation. Meanwhile, the auxiliary reward prediction and CURL objectives, although secondary, contribute to performance improvement to some degree.

InfoNCE-based temporal action-driven contrastive objective in TACO outperforms other representation learning objectives including SPR [43], ATC [44], and DRIML [37] . In Table 2, we have showcased a comparison between our approach and other visual RL representation learning objectives such as SPR , ATC , and DRIML . Given that SPR and DRIML were not initially designed for continuous control tasks, we have re-implemented their learning objectives using the identical backbone algorithm, DrQ-v2. A similar approach was taken for ATC, with their learning objectives also being reimplemented on DrQ-v2 to ensure a fair comparison. (Without the DrQ-v2 backbone algorithm, the performance reproduced by their original implementation is significantly worse.) Furthermore, recognizing the significance of learning action encoding, as discussed earlier, we have integrated action representation learning into all these baselines. Therefore, the model architecture remains consistent across different representation learning objectives, with the sole difference being the design of the temporal contrastive loss. For DRIML , given that only the first action of the action sequence is considered in the temporal contrastive loss, TACO and DRIML differ when the number of steps K is greater than one. Thus, we indicate N/A for tasks where we choose K = 1 for TACO .

Table 2 showcases that while previous representation learning objectives have proven benefit in assisting the agent to surpass the DrQ-v2 baseline by learning a superior representation, our approach exhibits consistent superiority over other representation learning objectives in all five evaluated environments. These results reinforce our claim that TACO is a more effective method for learning state-action representations, allowing agents to reason more efficiently about the long-term outcomes of their actions in the environment.

4.2 Combining TACO with offline RL algorithms

In this part, we discuss the experimental results of TACO within the context of offline reinforcement learning, emphasizing the benefits our temporal contrastive state/action representation learning objective brings to visual offline RL. Offline visual reinforcement learning poses unique challenges, as algorithms must learn an optimal policy solely from a fixed dataset without further interaction with the environment. This necessitates that the agent effectively generalizes from limited data while handling high-dimensional visual inputs. The state/action representation learning objective of TACO plays a vital role in addressing these challenges by capturing essential information about the environment's dynamics, thereby enabling more efficient generalization and improved performance. TACO can be easily integrated as a plug-and-play module on top of existing strong offline RL methods, such as TD3+BC [11] and CQL [31].

For evaluation, we select four challenging visual control tasks from DMC: Hopper Hop, Cheetah Run, Walker Run, and Quadruped Run. For each task, we generate three types of datasets. The medium dataset consists of trajectories collected by a single policy of medium performance. The precise definition of 'medium performance' is task-dependent but generally represents an intermediate level of mastery, which is neither too poor nor too proficient. The medium-replay dataset contains trajectories randomly sampled from the online learning agent's replay buffer before it reaches a medium performance level. The full-replay dataset includes trajectories randomly sampled throughout the online learning phase, from the beginning until convergence. The dataset size for Walker, Hopper, and Cheetah is 100K, while for the more challenging Quadruped Run task, a larger dataset size of 500K is used to account for the increased difficulty. We compute the normalized reward by diving the offline RL reward by the best reward we get during online TACO training.

We compare the performance of TD3+BC and CQL with and without TACO on our benchmark. Additionally, we also compare with the decision transformer ( DT ) [7], a strong model-free offline RL baseline that casts the problem of RL as conditional sequence modeling, IQL [30], another commonly used offline RL algorithm, and the behavior cloning ( BC ) baseline. For TD3+BC , CQL and IQL , which were originally proposed to solve offline RL with vector inputs, we add the their learning objective on top of DrQ-v2 to handle image inputs.

Table 3 provides the normalized reward for each dataset. The results underscore that when combined with the strongest baselines TD3+BC and CQL , TACO achieves consistent performance improvements across all tasks and datasets, setting new state-of-the-art results for offline visual reinforcement learning. This is true for both the medium dataset collected with a single policy and narrow data distribution, as well as the medium-replay and replay datasets with a diverse distribution.

5 Related work

5.1 Contrastive learning in visual reinforcement learning

Contrastive learning has emerged as a powerful technique for learning effective representations across various domains, particularly in computer vision [47, 8, 22, 23, 49]. This success is attributed to its ability to learn meaningful embeddings by contrasting similar and dissimilar data samples. In visual reinforcement learning, it's used as a self-supervised auxiliary task to improve state representation learning, with InfoNCE [47] being a popular learning objective. In CURL [33], it treats augmented states as positive pairs, but it neglects the temporal dependency of MDP. CPC [47], ST-DIM [2], and ATC [44] integrate temporal relationships into the contrastive loss by maximizing mutual information between current state representations (or state histories encoded by LSTM in CPC) and future state representations. However, they do not consider actions, making positive relationships in the learning objective policy-dependent. DRIML [37] addresses this by maximizing mutual information between state-action pairs at the current time step and the resulting future state, but its objective remains policy-dependent as it only provides the first action of the action sequence. Besides, ADAT [29] and ACO [59] incorporate actions into contrastive loss by labeling observations with similar policy action outputs as positive samples, but these methods do not naturally extend to tasks with nontrivial continuous action spaces. A common downside of these approaches is the potential for unstable encoder updates due to policy-dependent positive relations. In contrast, TACO is theoretically sufficient, and it tackles the additional challenge of continuous control tasks by simultaneously learning state and action representations.

In addition to the InfoNCE objective, other self-supervised learning objective is also proposed. Approahces such as DeepMDP [12], SPR [42], SGI [43], and EfficientZero [56] direct learn a latentspace transition model. Notably, these methods predominantly target Atari games characterized by their small, well-represented, and abstract discrete action spaces. When dealing with continuous control tasks, which often involve a continuous and potentially high-dimensional action space, the relationships between actions and states become increasingly intricate. This complexity poses a significant challenge in effectively capturing the underlying dynamics. In contrast, by framing the latent dynamics model predictions as a self-supervised InfoNCE objective, the mutual information

guided approach used by TACO is better suited for continuous control task, resulting in more stable optimization and thus better state and action representations.

5.2 Action representation in reinforcement learning

Although state or observation representations are the main focus of prior research, there also exists discussion on the benefits and effects of learning action representations. Chandak et al. [6] propose to learn a policy over latent action space and transform the latent actions into actual actions, which enables generalization over large action sets. Allshire et al. [1] introduce a variational encoderdecoder model to learn disentangled action representation, improving the sample efficiency of policy learning. In model-based RL, strategies to achieve more precise and stable model-based planning or roll-out are essential. To this end, Park et al. [39] propose an approach to train an environment model in the learned latent action space. In addition, action representation also has the potential to improve multi-task learning [25], where latent actions can be shared and enhance generalization.

6 Conclusion

In this paper, we have introduced a conceptually simple temporal action-driven contrastive learning objective that simultaneously learns the state and action representations for image-based continuous control TACO . Theoretically sound, TACO has demonstrated significant practical superiority by outperforming SOTA online visual RL algorithms. Additionally, it can be seamlessly integrated as a plugin module to enhance the performance of existing offline RL algorithms. Despite the promising results, TACO does present limitations, particularly its need for large batch sizes due to the inherent nature of the contrastive InfoNCE objective, which impacts computational efficiency. Moving forward, we envisage two primary directions for future research. Firstly, the creation of more advanced temporal contrastive InfoNCE objectives that can function effectively with smaller data batches may mitigate the concerns related to computational efficiency. Secondly, the implementation of a distributed version of TACO , akin to the strategies employed for DDPG in previous works [3, 24], could significantly enhance training speed. These approaches offer promising avenues for further advancements in visual RL.

References

Arthur Allshire, Roberto Martín-Martín, Charles Lin, Shawn Manuel, Silvio Savarese, and Animesh Garg. Laser: Learning a latent action space for efficient reinforcement learning. In 2021 IEEE International Conference on Robotics and Automation (ICRA) , pages 6650-6656. IEEE, 2021. 2, 10

Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. 2, 9

Gabriel Barth-Maron, Matthew W. Hoffman, David Budden, Will Dabney, Dan Horgan, Dhruva TB, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributional policy gradients. In International Conference on Learning Representations , 2018. 10

Marc Bellemare, Will Dabney, Robert Dadashi, Adrien Ali Taiga, Pablo Samuel Castro, Nicolas Le Roux, Dale Schuurmans, Tor Lattimore, and Clare Lyle. A geometric perspective on optimal representations for reinforcement learning. Advances in neural information processing systems , 32:4358-4369, 2019. 22

Edoardo Cetin, Philip J Ball, Stephen Roberts, and Oya Celiktutan. Stabilizing off-policy deep reinforcement learning from pixels. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 2784-2810. PMLR, 17-23 Jul 2022. 5

Yash Chandak, Georgios Theocharous, James Kostas, Scott Jordan, and Philip Thomas. Learning action representations for reinforcement learning. In International conference on machine learning , pages 941-950. PMLR, 2019. 10

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 15084-15097. Curran Associates, Inc., 2021. 9, 17

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 1597-1607. PMLR, 13-18 Jul 2020. 7, 9

Will Dabney, André Barreto, Mark Rowland, Robert Dadashi, John Quan, Marc G. Bellemare, and David Silver. The value-improvement path: Towards better representations for reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence , 35(8):7160-7168, May 2021. 22

Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Ruslan Salakhutdinov. Contrastive learning as goal-conditioned reinforcement learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. 23

Scott Fujimoto and Shixiang (Shane) Gu. A minimalist approach to offline reinforcement learning. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 20132-20145. Curran Associates, Inc., 2021. 2, 8, 17

Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning , pages 2170-2179. PMLR, 2019. 2, 9, 22

Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, koray kavukcuoglu, Remi Munos, and Michal Valko. Bootstrap your own latent a new approach to self-supervised learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 21271-21284. Curran Associates, Inc., 2020. 7, 23

Zhaohan Daniel Guo, Shantanu Thakoor, Miruna Pislar, Bernardo Avila Pires, Florent Altché, Corentin Tallec, Alaa Saade, Daniele Calandriello, Jean-Bastien Grill, Yunhao Tang, Michal Valko, Remi Munos, Mohammad Gheshlaghi Azar, and Bilal Piot. BYOL-explore: Exploration by bootstrapped prediction. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. 23

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations , 2020. 1, 22

Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 2555-2565. PMLR, 09-15 Jun 2019. 1, 22

Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations , 2021. 1

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2023. 1, 2, 5

Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. In Conference on Neural Information Processing Systems , 2021. 22

Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation. In International Conference on Robotics and Automation , 2021. 22

Nicklas A Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predictive control. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 8387-8406. PMLR, 17-23 Jul 2022. 1, 5

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 9726-9735, 2020. 7, 9, 23

Olivier Henaff. Data-efficient image recognition with contrastive predictive coding. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 4182-4192. PMLR, 13-18 Jul 2020. 2, 9

Matthew W. Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Nikola Momchev, Danila Sinopalnikov, Piotr Sta´ nczyk, Sabela Ramos, Anton Raichuk, Damien Vincent, Léonard Hussenot, Robert Dadashi, Gabriel Dulac-Arnold, Manu Orsini, Alexis Jacq, Johan Ferret, Nino Vieillard, Seyed Kamyar Seyed Ghasemipour, Sertan Girgin, Olivier Pietquin, Feryal Behbahani, Tamara Norman, Abbas Abdolmaleki, Albin Cassirer, Fan Yang, Kate Baumli, Sarah Henderson, Abe Friesen, Ruba Haroun, Alex Novikov, Sergio Gómez Colmenarejo, Serkan Cabi, Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Andrew Cowie, Ziyu Wang, Bilal Piot, and Nando de Freitas. Acme: A research framework for distributed reinforcement learning. arXiv preprint arXiv:2006.00979 , 2020. 10

Pu Hua, Yubei Chen, and Huazhe Xu. Simple emergent action representations from multi-task policy training. arXiv preprint arXiv:2210.09566 , 2022. 10

Riashat Islam, Manan Tomar, Alex Lamb, Yonathan Efroni, Hongyu Zang, Aniket Didolkar, Dipendra Misra, Xin Li, Harm van Seijen, Remi Tachet des Combes, and John Langford. Agent-controller representations: Principled offline rl with rich exogenous information, 2022. 22

Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations , 2017. 22

Hong Jun Jeon, Dylan Losey, and Dorsa Sadigh. Shared Autonomy with Learned Latent Actions. In Proceedings of Robotics: Science and Systems , Corvalis, Oregon, USA, July 2020. 2

Minbeom Kim, Kyeongha Rho, Yong-duk Kim, and Kyomin Jung. Action-driven contrastive representation for reinforcement learning. PLOS ONE , 17(3):1-14, 03 2022. 9

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations , 2022. 9

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 1179-1191. Curran Associates, Inc., 2020. 2, 8

Alex Lamb, Riashat Islam, Yonathan Efroni, Aniket Didolkar, Dipendra Misra, Dylan Foster, Lekan Molu, Rajan Chari, Akshay Krishnamurthy, and John Langford. Guaranteed discovery of control-endogenous latent states with multi-step inverse models, 2022. 22

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representations for reinforcement learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 5639-5650. PMLR, 13-18 Jul 2020. 1, 2, 4, 5, 9

Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 19884-19895. Curran Associates, Inc., 2020. 1

Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. Advances in Neural Information Processing Systems , 33:741-752, 2020. 22

Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. VIP: Towards universal visual reward and representation via value-implicit pre-training. In The Eleventh International Conference on Learning Representations , 2023. 23

Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 3686-3698. Curran Associates, Inc., 2020. 2, 8, 9

Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsurprising effectiveness of pre-trained vision models for control. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 17359-17371. PMLR, 17-23 Jul 2022. 22

Seohong Park and Sergey Levine. Predictable mdp abstraction for unsupervised model-based rl, 2023. 10

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 8748-8763. PMLR, 18-24 Jul 2021. 7

Kate Rakelly, Abhishek Gupta, Carlos Florensa, and Sergey Levine. Which mutual-information representation learning objectives are sufficient for control? In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. 4, 21, 22

Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In International Conference on Learning Representations , 2021. 2, 8, 9, 22

Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R Devon Hjelm, Philip Bachman, and Aaron Courville. Pretraining representations for dataefficient reinforcement learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. 8, 9

Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 9870-9879. PMLR, 18-24 Jul 2021. 2, 8, 9

Yanchao Sun, Ruijie Zheng, Xiyao Wang, Andrew E Cohen, and Furong Huang. Transfer RL across observation feature spaces via model-based regularization. In International Conference on Learning Representations , 2022. 22

Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy Lillicrap, and Martin Riedmiller. Deepmind control suite, 2018. 5

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, 2019. 2, 3, 9

Xiyao Wang, Ruijie Zheng, Yanchao Sun, Ruonan Jia, Wichayaporn Wongkamjan, Huazhe Xu, and Furong Huang. Coplanner: Plan to roll out conservatively but to explore optimistically for model-based rl, 2023. 1

Zhirong Wu, Yuanjun Xiong, Stella X. Yu, and Dahua Lin. Unsupervised feature learning via non-parametric instance discrimination. 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 3733-3742, 2018. 9

Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control, 2022. 23

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 11920-11931. PMLR, 18-24 Jul 2021. 23

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. In International Conference on Learning Representations , 2022. 1, 2, 4, 5

Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations , 2021. 1, 5

Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations , 2021. 22

Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 10674-10681, 2021. 22

Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. 9, 22

Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Avnish Narayan, Hayden Shively, Adithya Bellathur, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning, 2021. 5

Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations , 2021. 22

Qihang Zhang, Zhenghao Peng, and Bolei Zhou. Learning to drive by watching youtube videos: Action-conditioned contrastive policy pretraining. European Conference on Computer Vision (ECCV) , 2022. 9

Appendix

A Effects of the choice of prediction horizon K

TACO functions as a versatile add-on to current online and offline RL algorithms, requiring only the prediction horizon K as an additional hyperparameter. In our investigations, we select K as either 1 or 3 across all tasks. The performance of TACO with K values of 1 and 3 is compared against the DrQ-v2 baselines in Figure 8. As illustrated in the figure, TACO outperforms the baseline DrQ-v2 for both K = 1 and 3 , indicating the effectiveness of the temporal contrastive loss in concurrent state and action representation learning. In comparing K = 1 to K = 3 , we observe that a longer prediction horizon ( K = 3 ) yields superior results for four out of nine tasks, specifically Hopper Hop, Quadruped Walk, Acrobot Swingup, and Reacher Hard. Conversely, for the Quadruped Run and Walker Run tasks, a shorter temporal contrastive loss interval ( K = 1 ) proves more beneficial. For the remaining tasks, the choice between K = 1 and K = 3 appears to have no discernable impact.

Quadruped Walk Finger Turn Hard Acrobot Swingup

Reacher Hard

Reach Duplo

0

200

400

600

800

Episode Reward

DrQ-v2

TACO (K=1)

TACO (K=3)

Our speculation hinges on the rate of changes in the agent's observations. While Theorem 3.1 applies regardless of the prediction horizon K , a shorter horizon (such as K = 1 ) can offer a somewhat shortsighted perspective. This becomes less beneficial in continuous control tasks where state transitions within brief time intervals are negligible. Therefore, in environments exhibiting abrupt state transitions, a larger K would be advantageous, whereas a smaller K would suffice for environments with gradual state transitions.

To substantiate this conjecture, we conduct a simple experiment on the task Hopper Hop where we aim to simulate the different rates of changes in the agent's observations. We verify this by varying the size of action repeats, which correspond to how many times a chosen action repeats per environmental step. Consequently, a larger action repeat size induces more pronounced observational changes. In our prior experiments, following the settings of DrQ-v2, we fix the size of the action repeat to be 2. In this experiment, we alter the size of the action repeat for Hopper Hop to be 1, 2, and 4. For each action repeat size, we then compare the 1M performance of TACO with the prediction horizon K selected from the set { 1 , 3 , 5 } . Interestingly, as demonstrated in Figure 9, the optimal K values for different action repeats are reversed: 5 for an action repeat of 1, 3 for an action repeat of 2, and 1 for an action repeat of 4. This observation further substantiates our assertion that the optimal choice of prediction horizon correlates with the rate of change in environmental dynamics.

B Time-efficiency of TACO

In this section, we compare the time efficiency of different visual RL algorithms. In Table 4, we present a comprehensive comparison of the speed of these algorithms, measured in frames per second

Action Repeat=1

Action Repeat=2

Action Repeat=4

50

100

150

200

250

300

Episode Reward

K=1

K=5

( FPS ). Additionally, we also provide FPS for TACO with different batch sizes in Table 5 as a reference. To ensure a fair comparison, all the algorithms are tested on Nvidia A100 GPUs.

As mentioned in §4.1, the InfoNCE objective in TACO requires a large batch size of 1024, which is 4 times larger than the batch size used in DrQ-v2. Consequently, this increases the processing time, making our method to run about 3.6 times slower than DrQ-v2, with similar time efficiency as DrQ. Therefore, the primary limitation of our method is this time inefficiency caused by the usage of large batch sizes. Potential solutions to improve the time efficiency could involve the implementation of a distributed version of our method to speed up training. We intend to explore these improvements in future work related to our paper.

C Experiment details

C.1 Online RL

In this section, we describe the implementation details of TACO for the online RL experiments. We implement TACO on top of the released open-source implementation of DrQ-v2, where we interleave the update of TACO objective with the original actor and critic update of DrQ-v2. Below we demonstrate the pseudo-code of the new update function.

Listing 1: Pytorch-like pseudo-code how TACO is incorporated into the update function of existing visual RL algorithms.

Next, we demonstrate the pseudo-code of how TACO objective is computed with pytorch-like pseudocode.

Then when we compute the Q values for both actor and critic updates, we use the trained state and action encoder. Same as what is used in DrQ-v2, we use 1024 for the hidden dimension of all the encoder layers and 50 for the feature dimension of state representation. For action representation, we choose the dimensionality of the action encoding, which corresponds to the output size of the action encoding layer, as ⌈ 1 . 25 ×|A|⌉ . In practice, we find this works well as it effectively extracts relevant control information from the raw action space while minimizing the inclusion of irrelevant control information in the representation. See Appendix D for an additional experiment on testing the robustness of TACO against the dimensionality of action representation.

C.2 Offline RL

TD3+BC, CQL, IQL all are originally proposed for vector-input. We modify these algorithms on top of DrQ-v2 so that they are able to deal with image observations. For TD3+BC, the behavior cloning regularizer is incorporated into the actor function update, with the regularizer weight α TD3+BC = 2 . 5 as defined and used in Scott et al. [11]. In our experiments, no significant performance difference was found for α within the [2 , 3] range. In the case of CQL, we augment the original critic loss with a Q-value regularizer and choose the Q-regularizer weight, α CQL, from { 0 . 5 , 1 , 2 , 4 , 8 } . Table 6 presents the chosen α CQL for each dataset.

For IQL, we adopt the update functions of both policy and value function from its open-source JAX implementation into DrQ-v2, setting the inverse temperature β to 0.1, and τ = 0 . 7 for the expectile. Lastly, for the Decision Transformer (DT), we adapt from the original open-source implementation by Chen et al. [7] and use a context length of 20.

D Sensitivity of TACO to action representation dimensionality: an additional experiment

In all of our previous experiments, we use ⌈ 1 . 25 ×|A|⌉ as the latent dimensions of the action space for TACO . In practice, we find this works well so that it retains the rich information of the raw actions while being able to group semantically similar actions together. To test the algorithm's sensitivity to this hyperparameter, we conduct an experiment on the Quadruped Run task, which has a 12-dimensional action space. In Figure 10, we show the 1M performance of TACO with different choices of latent action dimensions. Notably, we observe that as long as the dimensionality of the action space is neither too small (6-dimensional), which could limit the ability of latent actions to capture sufficient information from the raw actions, nor too large (24-dimensional), which might introduce excessive controlirrelevant information, the performance of TACO remains robust to the choice of action representation dimensionality.

6

9

12

15

18

21

24

Dimensionality of Action Representation

300

350

400

450

500

550

Performance

E Sensitivity of TACO vs. DrQ-v2 to random seeds

700

750

800

850

1M Episode Reward

DrQ-v2

TACO

(a)

50

100

150

200

250

300

350

1M Episode Reward

DrQ-v2

TACO

(b)

150

200

250

300

350

400

450

500

550

1M Episode Reward

DrQ-v2

TACO

(c)

450

500

550

600

650

700

750

800

850

1M Episode Reward

DrQ-v2

TACO

(d)

300

400

500

600

700

800

900

1000

1M Episode Reward

DrQ-v2

TACO

Reacher Reach

(e)

Reacher Hard

0

50

100

150

200

250

1M Episode Reward

DrQ-v2

TACO

(f)

0

100

200

300

400

500

600

1M Episode Reward

DrQ-v2

TACO

Finger Turn Hard

(g)

Finger Turn Hard

160

170

180

190

200

210

220

1M Episode Reward

DrQ-v2

TACO

Reach Duplo

(h)

Reach Duplo

350

400

450

500

550

600

650

700

1M Episode Reward

DrQ-v2

TACO

(i)

In this section, we conduct a comprehensive analysis of the algorithm's robustness under random seed, comparing the performance of our proposed method, TACO , with the baseline algorithm, DrQ-v2, across multiple seeds on nine distinct tasks shown in Figure 4. Our aim is to gain insights into the stability and robustness of both algorithms under varying random seeds, providing valuable information to assess their reliability.

Remarkably, Figure 11 demonstrates the impressive robustness of TACO compared to DrQ-v2. TACO consistently demonstrates smaller performance variations across different seeds, outperforming the baseline on almost every task. In contrast, runs of DrQ-v2 frequently encounter broken seeds. These results provide strong evidence that TACO not only improves average performance but also significantly enhances the robustness of the training process, making it more resilient against failure cases. This enhanced robustness is crucial for real-world applications, where stability and consistency are essential for successful deployment.

F Tackling the complex 'Manipulator Bring Ball" task: an additional online RL experiment within the DeepMind Control Suite

In addition to the nine tasks depicted in Figure Figure 4, we conduct an additional experiment on an even more complex task "Manipulator Bring Ball". This task necessitates the robotic arm to precisely locate, grab, and transport the ball to a specified goal location. A visualization of a successful trajectory is illustrated in Figure 12a.

The challenge of this domain lies in the sparse reward structure, the demand for accurate arm control, and the need for temporally extended skills. As a result, despite training for 30 million frames, as demonstrated in Figure 12b, DrQ-v2 fails to successfully complete this task. In comparison, TACO manages to successfully solve this task for half of the random initializations, i.e., 3 out of 6 seeds.

T=0

T=40

T=60

T=80

T=100

T=120

T=140

T=160

T=180

1m

5m

10m

15m

20m

25m

30m

Number of Environment Steps

0

100

200

300

400

500

600

Episode Reward

Manipulator Bring Ball

DrQ-v2

TACO

(b) Online learning curves on manipulator bring ball aggregated over 6 random seeds.

G Visualization of Meta-world Tasks

Hammer

T=0

T=40

T=60

T=80

T=100

Assembly

Disassemble

Stick Pull

Pick Place Wall

Hand Insert

H Proof of Theorem 3.1

In this section, we prove Theorem 3.1, extending the results of Rakely et al. [41]. We will use notation A t : t + K to denote the action sequence A t , ..., A t + K from timestep t to t + K , and we will use U t : t + K to denote the sequence of latent actions U t , ..., U t + K from timestep t to t + K .

𝑆

!

𝑆

!"#

𝐴

!

𝑈

!

𝑍

!

𝑍

!"#

𝐴

!"#

𝑈

!"#

Proposition H.1. Let X be the return to go H -t -K ∑ i =0 γ i R t + K + i , with the conditional indepen- dence assumptions implied by the graphical model in Figure 14. If I ( Z t + K ; Z t , U t : t + K -1 ) = I ( S t + K ; S t , A t : t + K -1 ) , then I ( X ; Z t , U t : t + K -1 ) = I ( X ; S t , A t : t + K -1 )

Proof. We prove by contradictions. Suppose there exists a pair of state action representation ϕ Z , ψ U and a reward function r such that I ( Z t + K ; Z t , U t : t + K -1 ) = I ( S t + K ; S t , A t : t + K -1 ) , but I ( X ; Z t , U t : t + K -1 ) < I ( X ; S t , U t : t + K -1 ) < I ( X ; S t , A t : t + K -1 ) . Then it suffices to show that I ( S t + k ; Z t , U t : t + K -1 ) < I ( S t + K ; S t , U t : t + K -1 ) , which would give us the desired contradiction since I ( Z t + k ; Z t , U t : t + K -1 ) ≤ I ( S t + k ; Z t , U t : t + K -1 ) , and I ( S t + K ; S t , U t : t + K -1 ) ≤ I ( S t + K ; S t , A t : t + K -1 ) .

Now we look at I ( X ; Z t , S t , U t : t + K -1 ) . Apply chain rule of mutual information:

Applying the chain rule in another way, we get

Therefore, we get

By our assumption that I ( X ; Z t , U t : t + K -1 ) < I ( X ; S t , U t : t + K -1 ) , we must have

Next, we expand I ( S t + K ; Z t , S t , U t : t + K -1 ) :

On the other hand, we have

(14)

Thus we have

But then because I ( S t + K ; S t | Z t , U t : t + K -1 ) > I ( X ; S t | Z t , U t : t + K -1 ) as S t → S t + K → X forms a Markov chain, it is greater than zero by the Inequality (10). As a result, I ( S t + K ; Z t , U t : t + K -1 ) < I ( S t + K ; S t , U t : t + K -1 ) < I ( S t + K ; S t , A t : t + K -1 ) . This is exactly the contradiction that we would like to show.

Before proving Theorem 3.1, we need to cite another proposition, which is proved as Lemma 2 in Rakely et al. [41]

Proposition H.2. Let X,Y.Z be random variables. Suppose I ( Y, Z ) = I ( Y, X ) and Y ⊥ Z | X , then ∃ p ( Z | X ) s.t. ∀ x, p ( Y | X = x ) = ∫ p ( Y | Z ) p ( Z | X = x ) dz

Proof of Theorem 3.1 : Based on the graphical model, it is clear that

Now define the random variable of return-to-go R t such that

Based on Proposition H.1, because

Now applying Proposition H.2, we get

As a result, when K = 1 , for any reward function r , given a state-action pair ( s 1 , a 1 ) , ( s 2 , a 2 ) such that ϕ ( s 1 ) = ϕ ( s 2 ) , ψ ( a 1 ) = ψ ( a 2 ) , we have Q r ( s 1 , a 1 ) = E p ( R t | S t = s 1 ,A t = a 1 ) [ R t ] = E p ( R t | S t = s 2 ,A t = a 2 ) [ R t ] . This is because p ( R t | S t = s 1 , A t = a 1 ) = p ( R t | S t = s 2 , A t = a 2 ) by Equation (18) as p ( z t | S t = s 1 ) = p ( z t | S t = s 2 ) , p ( u t | A t = a 1 ) = p ( u t | A t = a 2 ) . In case when K > 1 , because if E [ Z t + K , [ Z t , U t , ..., U t + K -1 ]] = E [ S t + K , [ S t , A t , ..., A t + K -1 ]] , then for any 1 ≤ k ≤ K , E [ Z t + k , [ Z t , U t , ..., U t + k -1 ]] = E [ S t + k , [ S t , A t , ..., A t + k -1 ]] , including K = 1 , by Data processing Inequality. (Intuitively, this implies that if the information about the transition dynamics at a specific step is lost, the mutual information decreases as the timestep progresses, making it impossible to reach its maximum value at horizon K .) Then the same argument should also apply here.

I Additional related work

I.1 Visual reinforcement learning

In this paper, we focus primarily on visual-control tasks, and this section reviews relevant prior work in visual RL. For visual-control environments, representation learning has shown to be a key. Many prior works show that learning auxiliary tasks can encourage the representation to be better aligned with the task and thus enhance the performance, such as reconstructing pixel observations [55], minimizing bisimulation distances [58], fitting extra value functions [4, 9], learning latent dynamics models [12, 42, 45], multi-step inverse dynamics model [32, 26] or various control-relevant objectives [27]. Model-based methods are also shown to be successful in efficient visual RL, which learn the transition dynamics based on the encoded observation [16, 35, 15, 56]. Data augmentation can also be used to smooth out the learned representation or value functions to improve the learning performance [54, 20, 19].

I.2 Additional works on self-supervised/contrastive learning in reinforcement learning

In §5.1, we summarize the works that apply self-supervised/contrastive learning objectives to improve the sample efficiency of visual reinforcement learning. In this section, we discuss the additional works that apply self-supervised/contrastive learning objectives to a broader set of topics in reinforcement learning.

Several recent works have investigated the use of self-supervised/contrastive learning objectives to pre-train representations for reinforcement learning (RL) agents. Parisi et al. [38] propose PVR, we could conclude that

which leverages a pre-trained visual representation from MoCo [22] as the perception module for downstream policy learning. Xiao et al. [50] introduce VIP, a method that pre-trains visual representations using masked autoencoders. Additionally, Ma et al. [36] propose VIP, which formulates the representation learning problem as offline goal-conditioned RL and derives a self-supervised dual goal-conditioned value-function objective.

Besides pretraining state representations, in goal-conditioned RL, the work by Eysenbach et al. [10] establishes a connection between learning representations with a contrastive loss and learning a value function. Moreover, self-supervised learning has also been employed for effective exploration in RL. Guo et al. [14] introduce BYOL-Explore, a method that leverages the self-supervised BYOL objective [13] to acquire a latent forward dynamics model and state representation. The disagreement in the forward dynamics model is then utilized as an intrinsic reward for exploration. Another approach, ProtoRL by Yarats et al. [51], presents a self-supervised framework for learning a state representation through a clustering-based self-supervised learning objective in the reward-free exploration setting.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a simple yet effective temporal contrastive learning approach for encoding the high-dimensional observations and inputs for reinforcement learning. The authors propose a loss function (TACO) related to the mutual information between representations of current states paired with action sequences and the future states. By jointly optimizing this TACO loss with CURL loss and reward prediction loss, the proposed method outperforms the SOTA representation learning algorithms in on-policy and off-policy RL frameworks. They also validate their effectiveness in model-based frameworks.\n\nSTRENGTHS:\n- The proposed method outperforms the SOTA algorithms both in model-free and model-based frameworks\n- The experiments were conducted in various environments with multiple baselines (including on-policy, off-policy, and model-based).\n- An appropriate ablation study was performed for each loss term.\n\nWEAKNESSES:\n- In terms of TACO loss, the difference between DRIML and the proposed method lies in the fact that the proposed method considers the entire action sequence up to t+K, rather than just observing single action. Looking at Figure 7 in the supplementary materials, it can be seen that in 4 out of 9 environments, TACO performs better when K=1 compared to K=3. This could suggest that extending DRIML appropriately to a continuous action space yields performance similar to TACO.\n- The authors mention that they implemented some baselines based on DrQ-v2 and additionally considered action as an input (lines 258-260). However, there is no detailed explanation of this implementation, making it somewhat difficult to ensure a fair comparison.\n\nQUESTIONS:\n- Regarding “recognizing the significance of learning action encoding, as discussed earlier, we’ve integrated action representation learning into all these baselines.” could you provide a more detailed explanation of how action inputs were considered in ATC and SPR in the existing methodologies?\n- In lines 257-258, it is stated, \"Without the DrQ-v2 backbone algorithm, the performance reproduced by their original implementation is significantly worse.\" What is the difference between extending DRIML to a continuous action space and DRIML with the DrQ-v2 backbone you used as a baseline? Which of these differences do you think has the most significant impact on performance?\n- Does the RL loss flow into the encoder, or does the encoder only update through the loss in Equation 3-5? If it is the latter, how does the RL loss affect performance?\n- Does the performance decrease if the action encoder is not used and only projection is performed?\n- Line 318: Please check the name of the algorithm (ACO → DBC for [50]).\n\nLIMITATIONS:\nThe authors mentioned the limitation in the conclusion section. Since they utilize InfoNCE objective and the performance of the proposed method depends on the batch size, the main objective of TACO impacts computational efficiency.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces an auxiliary objective based on contrastive learning to learn action and state representation for continuous control benchmarks. The auxiliary objective is called TACO. The main idea behind the objective is to maximize the mutual information between the current state s_t, current and future actions {a_t, a_{t + 1}, … , a_{t + k}} and the future state s_{t + k}. They show that they can outperform various model-free, model-based methods on various environment in the deepmind control suite.\n\nSTRENGTHS:\nThe main thing I like about the paper is the breadth of applicability of the approach. The authors use the objective for both model-free and offline RL methods. The method is compared against relevant baselines in both these cases. The authors also compare against other relevant methods that use auxiliary pretaining objectives for learning state representations. This thorough comparison shows the effectiveness of the approach.\n\nWEAKNESSES:\nThe main concern I have is the use of reward prediction and curl objectives. The main novelty of the paper seems to be the contrastive objective in equation 3. But it seems that there are no experiments that evaluates and studies this objective in isolation. Figure 6b removes one objective at a time but it would be nice to study the effective of removing both reward prediction and curl objectives. Without this it is hard to say whether the empirical gains are justified by the motivation and theoretical claims of the paper. I would be happy to increase the score if this point is addressed.\n\nQUESTIONS:\nIt seems that the TACO objective is used in conjunction with the underlying RL algorithm and not as a pertaining step. I wonder if the authors have tried using it for pretraining? If not, is there any reason why the authors have not used it for pretraining?\n\nI could not find the architectural details of H_\\theta and G_\\theta, could the authors specify the details of these?\n\nLIMITATIONS:\nThe authors have discussed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces TACO, a framework that learns state and action representations simultaneously in visual reinforcement learning for continuous control tasks. TACO optimizes mutual information between current state-action pairs and future state representations. It additionally optimizes 2 auxiliary losses. Experimental results demonstrate TACO can achieve great performance gains in online and offline RL settings. TACO offers a flexible and stable approach for capturing essential information in high-dimensional continuous control tasks.\n\nSTRENGTHS:\n- TACO is a simple yet effective framework that learns state and action representations by an auxiliary contrastive learning task. TACO could be integrated into both online and offline visual RL algorithms flexibly.\n- Extensive experiments on the DeepMind Control Suite demonstrate that TACO has outstanding performance.\n- The paper provides theoretical analysis of TACO's objectives. \n- In general, the paper is well-written.\n\nWEAKNESSES:\n- The proposed contrastive learning objective is very similar to DRIML. Both methods improve the performance of the model-free agent by enhancing the predictability of the latent representation through contrastive learning. The major difference is that in this paper, the whole action sequence is given instead of the first action only. Therefore, I think the novelty is limited.\n- Extra parameter K is required to tune on different tasks.\n- It would be better if the proposed method can be evaluated in more environments besides DMC.\n\nQUESTIONS:\n- Line 134-136, \"maximizing this mutual information objective ensures that the learned representations are sufficient for making optimal decisions.\" However, Thereom 3.1 only shows that the optimal Q function is the same with equivalent state and action representations. How do you ensure the optimal action will be the same?\n- In Eq. 4, Why do you use a learnable parameter $W$ as a similarity measure, as the features have already been projected to the latent space by network G and H？\n- Is the action encoder only used in the temporal contrastive loss or is it also incorporated in policy learning?\n\nLIMITATIONS:\nThe computational limitation is considered in the paper. In addition to the time complexity, I am also curious about the GPU memory cost caused by the large batch size.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces TACO, a novel state-action representation learning technique based on contrastive learning. Empirically, TACO outperforms both model-free and model-based visual RL baselines in both online and offline settings.\n\nSTRENGTHS:\n1. The work studies joint state-action representation learning, which is less studied than state representation learning.\n2. The empirical results for the online setting seem promising and highlights the importance of both state *and* action representation learning.\n\nWEAKNESSES:\n1. The core novelty of this work is the contrastive temporal objective for learning state-action representations and section Fig. 6b shows that it is most responsible for the observed performance improvement. However, it's unclear how important TACO is in the offline RL experiments; TACO makes use of data augmentation (for the CURL loss) while the remaining baselines don't have access to augmented data; thus TACO has an unfair advantage. I would like to see ablations similar to Fig. 6b for the offline experiments. \n2. None of the baselines considered learn a latent action representation. While Fig. 3 shows that TACO's action representation is important, it's unclear if it's any better than other methods that learn an action representation via e.g. a state-conditioned variational autoencoder [1].\n\nI am willing to raise my score if the authors address these concerns.\n\n**Minor comments:**\n1. Fig. 6 a and b should have the same vertical scale to make the figures more easily comparable.\n\nQUESTIONS:\n1. Is there a specific reason why the authors use CURL as an additional auxiliary tasks (as opposed to e.g. SODA [2])? If CURL can be swapped with another method, it would be worth mentioning this. \n2. Could the authors clarify the experiment in lines 215-231? You add 20 dummy dimensions to the action vector when learning the action representation, though I'm not sure what is meant by \"1000 noise dimensions\", and I am not sure where \"the 4000 actions\" come from. \n3. Line 120: Is there a reason why the authors chose to let the action representation be independent of the state?\n4. Line 115-120: The goal stated here is to compress the state-action representation while retaining information needed to solve the task. Is the action representation being compressed? The sensitivity analysis in Appendix D considers increasing the action dimensionality in Quadruped Run by a factor of 2.\n\n[1] Laser: Learning a latent action space for efficient reinforcement learning. Allshire et al. ICRA 2021.\n[2] Generalization in reinforcement learning by soft data augmentation. Hansen & Wang. ICRA 2021.\n\nLIMITATIONS:\nSee Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}