$\texttt{TACO}$: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning
Authors: Ruijie Zheng
Xiyao Wang
Yanchao Sun
Shuang Ma
Jieyu Zhao
Huazhe Xu
Hal Daumé III
Furong Huang
Venue: NeurIPS 2023 poster (2023)
Keywords: Deep Reinforcement Learning
Visual Reinforcement Learning
Online Visual RL
Offline Visual RL
Action Representation
TL;DR: We propose a temporal action-driven contrastive loss that concurrently learn state and action representations for continuous control tasks.

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
Ruijie Zheng'
Xiyao Wang'
Yanchao Sun
Shuang Ma²
Jieyu Zhao
Huazhe X
Hal Daumé 5.58
Furong Huang 1.8
University of Maryland, College Park 2 University of Southern California
3 Tsinghua University Shanghai Qi Zhi Institute Microsoft Research
rzheng eeteng120unn 120umd. edu
Abstract
Despite recent progress in reinforcement learning (RL) from raw pixel data, sample inefficiency continues to present substantial obstacle. Prior works have attempted address this challenge by creating suff-supeeeieed auxiliary tasks, aiming to enrich the agent's '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: Temporal Action-driven COntrastive Learning, a simple yet powerful temporal contrastive learning approach that facilitates the concurrent ac- quisition latent state and action representations for agents. TACO simultaneously learns state and an action representation by optimizing the mutual information between representations current states paired with action sequences and repre- sentations 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 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 longstanding challenge. A central aspect of this challenge lies in addressing the sample inefficiency problem 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.
$ Equal advising
Our code is released at http://ghhaabccmm//aaa///a//2//2
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

State representation learning has become an essential as- pect of RL research, aiming to improve sample efficiency and enhance agent performance. Initial advancements like CURL. [33] utilized self-supervised contrastive InfoNCE objective [47] for state representation, yet it overlooks the temporal dynamics 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 state representation capable of effectively predicting fu- ture observations. A more comprehensive approach was taken by DRIML [37], which incorporated the first action of the action sequence into the temporal contrastive learn- ing framework. However, these methods, while innovative, have their shortcomings. The positive relations in their 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 rell-sepres,nted, 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 more effectively reason about its actions' long-term outcomes for continuous control tasks.
Deepmind Control Suite 1M (9 Challenging Tasks)
Average Episode Reward
300
Dreamer v33
TACO
300-
DrQ-v2
CURL
DrQ
100
Model
Comparison of average episode reward across nine challenging tasks in Deep- mind Control Suite after one million environ- ment steps.
Figure 1:
In this paper, we introduce Temporal Action-driven COntrastive Learning (TACO) promising approach to visual continuous control tasks. TACO simultaneously learns a state and action represen- tation 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 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 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-ple 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 v33 [18]. As shown in Figure 1, across nine challenging visual continuous control tasks from DMC, TACO achieves 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:
1. We present TACO, simple yet effective temporal contrastive learning framework that simultane- ously learns state and action representations.
2. 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.
3. We theoretically show that the objectives of TACO is sufficient to capture the essential information in state and action representation for control.
4. 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 = P, 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 x A + A(S), where (SS) is a distribution over the state space. S + the reward function. The objective of the Reinforcement Learning (RL) algorithm is the identification of an optimal policy n- S - ((A) that maximizes the expected value E,, [-r= y're]. Additionally, we can define the optimal Q function as follows: Q" (s. a) = Ex. Ei=o r'r(St, (s) = 8, ag = a) such that the relationship between the optimal Q function and optimal policy (()) = arg max, Q" (s, a). In the domain of visual RI,, high-dimensional image data are given state observations, so the simultaneous learning of both representation and control policy becomes the main challenge. This challenge 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 representa- tions, grouping similar/positive pairs and distancing dissimilar/neegative 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 two random variables. Given an instance ~ p(r), and a corresponding positive sample yt p(y(æ) well collection of Y {uI1... WN-1) } of N random samples from the marginal distribution p(y), the InfoNCE loss is defined as
f(x,y" yx
LN = Ex log Σ EverU(ut) ((,
(I)
Optimizing this loss will result in fxx 3) x plylx) and one can show that InfoNCE loss upper bounds the mutual information, LN log(N) I(X,Y).
ply)
3 TACD: temporal action-driven contrastive Loss
TACO is 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 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 concise and expressive understanding of both its current state and the potential effects of its actions, thereby enhancing sample efficiency and generalization capabilities.
Let St. A be the state and action variables timestep t, Z, ((S)). U. (bA) 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:
J I[ k; [4,, Ut, -1xx-11))
Here, K > 1 is fixed hyperparameter for the prediction horizon. In practice, 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 sufficieney of TACO objective:
Theorem 3.1. Let K E N*, and TACo = Ki Ki IZ,, tion representation oz, Yu, J raco maximized, then for arbitrary stale-action pairs (51,01), (32, 02) such that ((81) = (())), (((1)) = (a2)). it holds that Q (51, 01) = (82 02).
Ut,...,ee Ku-KU++K--- K-1).. If for given state and ac-
This theorem guarantees that if our mutual information objective Equation (2) is maximized, then for any two state-action pairs ($1, 01) and (82, a2) with equivalent state and action representations, their optimal action-value functions, (51, a1) and Q* (82, 02), 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
$to.
State Enc.
Action
Actioon Enc.
Action Enc.
State Enc.
Z
M
Contrastive Loss
(G.. 0r) +-
Q Net
Proj. Loyer
Proj. Leyer
Figure 2: demonstration of our temporal contrastive loss: Given batch state-action transition triples
((s,") [a( aa d;-x-1l (4 we first apply state encoder and action encoder get latent state-
action encodings: ((=}" [u, ...., x (t) }) Then we apply two different projection layers to map
(=(*) [u[" .... # -1l) and *+* into the shared contrastive embedding space. Finally, we learn predict the
(1)
(4)
correct pairings between (žt, [ut...., Mc4K and 2 using an InfoNCE loss.
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 base RL algorithm by incorporating the temporal contrastive loss as an auxiliary loss during the batch update process.
Specifically, given batch of state and action sequence transitions {(st (t' =t K), optimize:
(2) [a(
(6) (2211 N
N
9i
W
W.W.CGG Go, Ho) = N
log
(3)
Σ
(4)
WhO
2j=1
Here let (2) )(
ui (2) = (a(a.))) (1) be state and action embeddings respectively.
Go(z(") u. .... 1/2_)). and (1) Ho Hz(2) where G9 and Ha denote two learnable projection
layers that map the latent state =1(0) well as latent state and action sequence (=(() un u 1) to a common contrastive embedding space. W' is learnable parameter providing a similarity measure between Sr and h, in the shared contrastive embedding space. Subsequently, both state and action rep- resentations are fed into the agent's 's Q network. allowing the agent to effectively reason about the long- term 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 the CURL [33] loss:
Wh
JCURL v. W. Ho) = N
Σ log
(4)
WA + +
WhO
Here, (0)
the augmented view of s by applying the same random shift augmentation as DrQ-v2 [52]. W and Ho share the same weight the ones in TACO objective. The third objective is reward prediction:
Ho( }). where 8 (t)
(t)
JReward (0, $1 Rio)
(Ro(=))) (0) ()) .... up 1)) (0)
(5)

---

Here r(0) = 13 is the sum of reward from timestep t to t' 1, and Rg is 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 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 or 3, depending on the nature of the environment. We refer the readers to Appendix A for discussion 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 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 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, 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 challeng- ing 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-worke 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 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, reproduction of results from its open- source implementation does not consistently surpass DrQ-v2. As such, we do not choose A-LIX as the backbone algorithm for TACO. Additionally, compare with two state-of-the-art model-basee 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- 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 (DMC), and Figure 5 (Mela-world).
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 DrQ- v2 requires approximately 1.5 million steps. When trained with only I 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 6 out of 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 three million steps, 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.
5

---

Cheetah
Swingup
Quadiruped
Eptodd Reward
800
380
720
420
200
234
Quadruped Walk
Reacher Hard
Hopper
430
310
Epuode Rowars
250
370
34h
L2n
3.5m
12n
Finger Turn Hard
Reach Duplo
We Walket Run
3035
230
200
720
600
530
410
300
200
300
Number Environment Stops
Number Enviroommnnt Stops
38m
N.mber of Brwironment
2/47
TACO
A-LIX
DrQ+2
Figure 4: (Deepmind Control Suite) Performance of TACO against two strongest model-free visual RL baselines. Results of DrQ-v2 and A-LIX are reproduced from their open-source implementations, and all results are averaged over 6 random seeds.
Table 1: Episode reward of TACO and SOTA visual RL algorithms on the image-based DMControl IM benchmark.
Results are averaged over random seeds. Within the table, entries shaded represent the best performance of model-free algorithms, while text in bold signifies the highest performance across all baseline algorithms, including model-based algorithms.
Model-free
Model-based
Enviromment (IM Steps)
TACO
DrQv2
A-LIX
DrQ
CURL
Dreamer-v3
TDMPC
Quadruped
38
454 42
397
Hopper Hop
261 + 52
189
+ 35
225 = 13
192 1:
152 + 34
369
21
195 :: 18
Walker Run
637 + 11
517 43
617 =
451 + 73
387 1: 21
765 + 32
600 + 28
Quadruped
560
Cheetah Run
821 48
691 42
676 =1
474 32
657 +: 35
728 + 32
565 + 61
Finger Turn Hard
632 + 75
220
62 54
911:9
215 17
810 1 58
400 100 113
Acrobot Swingup
21
128
112
24::8
210
224 :
Reacher Hard
883 63
572
+ 51
510 16
4-1+-45
400 +
499 + 51
485 +
Reach Duplo
234 21
206
1: 32
199
199=10
36 47
8
119 1 30
117
Concurrently learning state and action representation is erucial for the success of TACO. To demonstrate the effective- ness of action representation learning in TACO, we evaluate its performance on subset of difficult benchmark tasks and com- pare it with baseline method without action representation, as shown in Figure 3. The empirical results underscore the effi-
IACD
TACO wlo representation
DrQv2
newar
Episode Rewe
cacy of the temporal contrastive learning objectives, even in the absence 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 perfor- mance comparable to DrQ-v2. Furthermore, when comparing against TACO without action representation, TACO achieves a consistent performance gain, ranging from 12.2% 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 TACD, but also underscore the instrumental role of high-quality action representation in bolstering the performance of the underlying RL algorithms.
Quadnieed Mun Wakeer Aus Hopper Rescher
Figure 3: IM Performance of TACO with and without action representation
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
Disessemble
Suivess Helte
60%
40%
40%
Nvv hagh
2015
20%
09
0m
0.4m
0.8m
1.2m
1.5m
0m
0.4m
8.8m
3,2m
1i0m
2m
0.4m
0.5m
1mm
1.mm
Pull
Wall
Haind
1esk Sucess Hate
2095
60%
00%
40%
60%
40%
20%
20%
05/
DMV
Number Environment Steps
Number of Lnvirnnment Steps
Number Environment Steps
TACO
A-LIX
DrQ-v2
Figure 5: (Meta-world) Performance of TACO against DrQ-v2 and A-LIX. All results are averaged over random sceds.
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 d1, 022 ag, 0,4 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 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 learned action representation could group the four clusters, demonstrating the ability of our action representation extract control relevant information from the raw action space.
The effectiveness our temporal contrastive loss is enhanced with a larger batch size. As widely acknowledged in contrastive learning re- search [8, 22, 40, 13], contrastive loss sees sig- nificant benefits from utilizing a larger batch size. In Figure 7a, illustrate the performance of our algo- rithms alongside DrQv2 after one million environ- ment steps the Quadruped Run task. As evident from the plot, batch size greatly influences the per- formance of algorithm, while DrQ-v2's baseline performance remains fairly consistent throughout training. In order to strike balance between time efficiency and performance, we opt for 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 's runtime, we direct the reader to Appendix B.
of actions
action represencations
Figure 6: Left: t-SNE embedding of actions with distracting dimensions. Right: t-SNE embedding of latent representations for actions with distracting dimensions.
TACO
DrQ v2
$50-
Epicode Reward
500-
sco
800.
#50
TACO
"N:0 JEURL
tnia Jraco
400
Jeeaard
CCUR
2014
2013
3:0
Batch Size
(a) TACO and DrQ-v2 across different batch sizes.
(b) TACD with different learning objective removed Figure 7: TACO with different batch sizes and with different component of its learning ohjectives removed
Reward prediction and CURL loss serves 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 's performance. In Figure 7b, we remove one objective at 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 '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 repre- sentation 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 identiaal backbone algorithm, DrQ-v2. A similar approach was taken for ATC, with their learning objectives also being reimplemented on DrQ-v2 to ensure fair comparison. (Without the DrQ-v2 backbone algorithm, the performance reproduced by their original implementation significantly worse.) Furthermore, recognizing the significance of learning action encoding, as discussed carlier, 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: Comparison with other objectives including SPR [42], ATC [44), and DRIMI. [37]
Environment
TACO
SPR
ATC
DRIML
DrQ-v2
Quadruped Run
541 :: 38
448 + 79
432 + 54
N/A
107 + 21
Walker Run
637 + 21
560
+ 71
502 + 171
N/A
517
+ 43
Hopper Hop
261 52
154
+ 10
112 98
216 13
192
Reacher Hard
883 + 63
711
+
863
+ 12
835 + 72
572
+ 51
Acrobot Swingup
241
= 21
198
21
206
+ 61
222 + 39
210
+ 12
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 fixed dataset without further interaction with the environment. This necessitates that the agent effectively generalizes from limited data while handling nigh-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 's dynamics, thereby enabling more efficient generalization and improved performance. TACO can be easily integrated as huug-ndd-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 trajec- tories randomly sampled from the online learning agent's '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.
8

---

Table 3: Offline Performance (Normalized Reward) for different offline RL methods. Results are averaged over
random seeds. +: captures the standard deviation over seeds.
Task/ Dataset
TD3+BC TACO
TD3+BC
CQL TACO
CQL
DT
IQL
BC
Hopper Hop
Medium-replay
Medium
67.2 0.1
52.4 = 0.4
51.2 = 0.8
62.9 0.1
74.6 : 0.4
47.9 0.6
46.7 0.2
68.7 + 0.1
6
3.6
57.6
2.0
=00=0 1.7
25.9 + 3.2
48.2 + 0.8
replar
97.6 144
83.8 = 2.3
101.2 1 1.9
94.2 2.0
92.4
22.40:03 0.3
47.7 1:22
63.7 2.7
Cheetah Run
Medium
66.6 = 0.5
06.1 0.3
70.1 + 0.4
68.7
6.73
6..9
Medium replay
62.6 0.2
61.1 0.1
72.3 1.2
67.3
67.0
1 0.6
28.5 4:3.2
48.0 3.10
Fall-replny
92.5 2.4
91.2= 0.8
86.9 2.4
65.0 3.9
80.0
69.0 1..6
Walker Run
Medium-replay
Medium
6..
0.5
62.3 0.2
18.0 0.2
62.3 + 2.6
19.6 = 1.0
59.9
19.4 0.9
61.7
17.3 0.3
414
4.4
46.2 0.6
18.5 0.8
Fall-repiny
86.8 + 0.6
84.0 = 1.6
68.1 + 0.1
79.8 + 0.6
$1.6 0.8
18.1 + 3.7
30.8 + 1.8
Quadruped Run
Medium
60.6=011
3.7
116
0.8
56,2 1.1
Medium-replay
61.3 0.3
58.1=0.5 0.5
61.9 + 0.2
61.2 0.9
19.5 = 2.2
68.4 4.4
51
3.3
Fell-repiny
92.6 0.7
80.3 0.4
92.1 + 0.1
85.2 2.5
14.5
36.3 5.9
57.6 0.7
Average Normalized Score
71.0
72.1
66.7
48.4
25.8
54.9
We compare the performance of TD3+BC and CQL with and without TACO Additionally, we also compare with the decision transformer (DT) [7], a strong model-free offline RL baseline that casts the problem of RL 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 solve offline RL with vector inputs, add the their learning objective on top of DrQ-v2 to handle image inputs.
benchmark.
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, well the medium-rephe and replay datasets with diverse distribution.
5
Related work
5.1 Contrastive learning in visual reinforcement learning
Contrastive learning has emerged 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 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 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 positive samples, but these methods do not naturally extend to tasks with non- trivial continuous action spaces. A common downside of these approaches the potential for unstable encoder updates due to policy-dependent positive relations. In contrast, TACO is theoretically sufficient, and tackles the additional challenge of continuous control tasks by simultancously learning state and action representations.
In addition 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 latent- space 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 significant challenge effectively capturing the underlying dynamics. In contrast, by framing the latent dynamics model predictions 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 effeets of learning action representations. Chandak et al. [6] propose to learn 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 encoder- decoder 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 essential. To this end, Park et al. [39] propose approach train 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 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 - TACD. Theoretically sound, TACO has demonstrated significant practical superiority by out- performing SOTA online visual RL algorithms. Additionally, it can be seamlessly integrated as a plug- in 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 en- visage two primary directions for future research. Firstly, the creation of more advanced temporal con- trastive 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.
7
Acknowledgement
Zheng, Wang. Sun and Huang are supported by National Science Foundation NSF-IIS-FAI pro- gram, DOD-ONR-Oftice of Naval Research, DOD Air Force Office of Scientific Research, DOD- DARPA-Defense Advanced Research Projects Agency Guaranteeing AI Robustness against Deception (GARD), Adobe, Capital One and JP Morgan faculty fellowships.
References
[1] Arthur Allshire, Roberto Martin-Martin, Charles Lin, Shawn Manuel, Silvio Savarese, and Animesh Garg. Laser: Learning latent action space for efficient reinforcement learning. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 6650-6656. IEEE, 2021. 2, 10
[2] Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Mlee-Ale Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Bue, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. 2,9
[3] 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
[4] 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 optimal representations for reinforcement learning. Advances in neural information processing systems, 0144358-43661 2019. 22
10

---

[5] Edoardo Cetin, Philip J Ball, Stephen Roberts, and Oya Celiktutan. Stabilizing off-policy deep reinforeement learning from pixels. In Kamalika Chaudhuri, Stefanic 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-2811 10. PMLR, 17-23 Jul 2022. 5
[6] 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
[7] 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
[8] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple frame- work 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, Mare G. Bellemare, and David Silver. The galue-improvement path: Towards better representations for reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence, 35(8):7160-7168, May 2021. 22
[10] 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
[11] 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
[12] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deep- mdp: Learning continuous latent space models for representation learning. In International Conference Machine Learning, pages 1770-119,, PMLR, 2019. 2, 9, 22
[13] 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 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 1122122122 2124. Curran Associates, Inc., 2020. 7.23
[14] 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
[15] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representa- lions, 2020. 1, 22
[16] 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
11

---

[17] Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021.
[18] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2023. 2.5
[19] 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
[20] Nicklas Hansen and Xiaolong Wang- Generalization reinforcement learning by soft data augmentation. In International Conference on Robotics and Automation, 2021. 22
[21] Nicklas A Hansen, Hao Su, and Xiaolong Wang- Temporal difference learning for model predictive control. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaha Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference Machine Learming, volume 162 of Proceedings of Machine Learning Research, pages 8387-8406. PMLR,
17-23 Jul 2022. 1,5
[22] 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
[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 PMLR, 13-18 Jul 2020. 2,
[24] Matthew W. Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Nikola Momchev, Danila Sinopalnikov, Piotr Stanczyk, Sabela Ramos, Anton Raichuk, Damien Vincent, Léonard Hussenot, Robert Dadashi, Gabriel Dulac-Arold, 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 20Xi::2006.00979 2020. 10
[25] Pu Hua, Yubei Chen, and Huazhe Xu. Simple emergent action representations from multi-task policy training. arXiv preprint 2rXiv:2210.09566, 2022. 10
[26] 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 with rich exogenous information, 2022. 22
[27] Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czameekki Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations, 2017. 22
[28] 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
[29] Minbeom Kim, Kyeongha Rho, Yong-duk Kim, and Kyomin Jung. Action-driven contrastive representation for reinforcement learning. PLOS ONE, 17(3):1114, 03 2022. 9
[30] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning In International Conference on Learning Representations, 2022.
[31] 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- 191. Curran Associates, Inc., 2020. 2, 8
12

---

[32] 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
[33] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised repre- sentations for reinforcement learning. In Hal Daumé IIII 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. 2, 5.9
[34] 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
[35] Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. Advances Neural Information Processing Systems, 33:741-752, 2020. 22
[36] 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
[37] 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
[38] Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsur- prising 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
[39] Scohong Park and Sergey Levine. Predictable mdp abstraction for unsupervised model-based rl, 2023. 10
[40] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, 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 8763. PMLR. 18-24 Jul 2021.7 7
[41] Kate Rakelly, Abhishek Gupta, Carlos Florensa, and Sergey Levine. Which representation learning objectives are sufficient for control? In A. Beygelzimer, Y. Dauphin, Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. 4, 21, 22
[42] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In Interna- tional Conference on Learning Representations, 2021. 2, 8, 9, 22
[43] Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R Devon Hjelm, Philip Bachman, and Aaron Courville. Pretraining representations for data- efficient reinforcement learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. 8,9
[44] 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
13

---

[45] 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
[46] 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
[47] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, 2019. 2, 3,9
[48] Xiyao Wang, Ruijie Zheng, Yanchao Sun, Ruonan Jia, Wichayaporn Wongkamjan, Huazhe Xu, and Furong Huang. Coplanner: Plan roll out conservatively but to explore optimistically for model-basee ri, 2023.
[49] 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
[50] Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control, 2022. 23
[51] 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
[52] Denis Yarats, Rob Fergus. Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. In International Conference
Learning Representations, 2022. 2, 0,2,4,3 4,5
[53] Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regu- larizing deep reinforcement learning from pixels. In International Conference Learning Representations, 2021. 5,5
[54] Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regu- larizing deep reinforcement learning from pixels. In International Conference on Learning Representations, 2021. 22
[55] 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
(56) 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 Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. 9, 22
[57] Tianhe Yu, Deirdre Quillen, Zhanpeng He. Ryan Julian, Avnish Narayan, Hayden Shively, Adithya Bellathur, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: benchmark and evaluation for multi-task and meta reinforcement learning, 2021. 5
[58] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine Learning invariant representations for reinforcement learning without reconstruction. In Inter- national Conference on Learning Representations, 2021. 22
[59] 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.
14

---

Appendix
A Effects of the choice of prediction horizon K
TACO functions as a versatile add-on to current online and offline RI. algorithms, requiring only the prediction horizon K as additional hyperparameter. In our investigations, we select K as either 1 or 3 across all tasks. The performance of TACO with values DrQ-v2 baselines in Figure 8. As illustrated in the figure, TACO outperforms the baseline DrQ-v2 for both
and is compared against the
and 3, indicating the effectiveness of the temporal contrastive loss in concurrent state and action representation learning. In comparing K 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, shorter temporal contrastive loss interval (K = 1) proves more beneficial. For the remaining tasks, the choice between K and K appears have discernable impact.
K 3, observe that longer
DrQ-22
TACO (K-1)
(K-3)
Episode Reward
Hop
Quadi Walk Hard Aumbot
Reachee Hard
Beach
Quuetruped Rum
Walte
Chewlahh
Figure 8: IM Performance of TACO with step size K = 1 and K 3.
Our speculation hinges on the rate changes in the agent's observations. While Theorem 3.1 applies regardless of the prediction horizon K, shorter horizon (such K 1) can offer somewhat shortsighted perspective. This becomes less beneficial in continuous control tasks where transitions within brief time intervals negligible. Therefore, in environments exhibiting abrupt state transitions, larger K would be advantagcous, whereas smaller K would suffice for environments with gradual state transitions.
To substantiate this conjecture, we conduct a simple experiment the task Hopper Hop where aim 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 chosen action repeats per environmental step. Consequently, 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 For each action repeat size, we then compare the 1M performance of TACO with the prediction horizon K selected from the 3., 3,5.. Interestingly, as demonstrated in Figure 9, the optimal K values for different action repeats are reversed: for an action repeat of 1, 3 for an action repeat of 2, and for action repeat of This observation further substantiates our assertion that the optimal choice of prediction horizon correlates with the rate of change environmental dynamics.
K=1
K-3
250
Episode Reward
K-5
200
Ropoat-
Acton Ropcst-2 Repesti
Performance of TACO with different prediction horizon K under different action repeat sizes. Shaded columns correspond to the hest prediction horizon K' under fixed ac- tion repeat size.
Figure
B Time-efficiency of TACO
In this section, we compare the time efficiency of different visual RI. algorithms. Table 4, we present a comprehensive comparison of the speed of these algorithms, measured in frames per second
15

---

(FPS). Additionally, we also provide FPS for TACO with different batch sizes in Table as a reference. To ensure 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 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 distributed version of our method speed up training. We intend to explore these improvements in future work related to our paper.
Table 4: Frames per second (FPS) for visual RL algorithms. B stands for batch size used in their original implementations.
TACO (B: (B:1024) DrQv2 (B:256)
A-LIX (B:256)
DrQ (B:512)
CURL (B:512)
Dreamer-v3 (B:256)
TDMPC (B:256)
31
22
Table 5: Frames per second (FPS) for TACO with different batch size
TACO (B:1024) TACO (B:512)
TACD (B:256)
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.
### Extract feature representation for state and actions.
def update (batch): ubattebbatch):
obs, action_sequence, reward, next_obs batch ### Update Agent's critic function
action_sequence, reward, next_obs) ### Update the agent's value function update_actor (obs)
### Update TACO loss
update taco (obs,
reward, next_obs)
Listing 1: Pytorch-like pseudo-code how TACD 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 pseu- docode.
state_ State/Observation Encoder (CNN)
action_encoder:
Action Encoder (MLP with 1-hidden layer)
sequence_
Action Sequence encoder (Linear Layer) Revard Prediction Layer (MLP with 1-hidden layer)
revard_predictor:
G:
Projection Layer I (MLP with 1-hidden layer)
Projection Layer II (MLP with 1-hidden layer)
aug: W:
Data Augmentation Function (Radnom Shift)
Matrix for computing similarity score
def compute, sacee_ttcc___e tettttacccceeeeeoe obj.etiveeeobs, action_sequence, reward, next_obs): Compute feature representation for both state and actions.
11
12
2
z_anchor = state_ (aug (obs))
13
(aug (obs), stop_grad=True)
next_z state_encoder (aug(next_obb))) stop-grad=True)
15
u_seq sequence encoderreeeeeeeeeeeeeeee (action encoder (actioniaciion.. (action_
16
Project joint contrastive embedding space
17
x O(torch.cate u_seq], dim=-1)) H(next_z)
18
### Compute bilinear product T"y
16

---

### Diagonal entries of x correspond to positive pairs
logits torch. matmul (x, torch matmul (W, y.t))
22
logits logits torch.max
labels aorch. malgiissss (logits, 1)
23
24
taco_loss cross_ eross_entropy..o loss (logits, labels)
*## Compute CURL loss H(z)
27
28
= H(z_anchor) detach() ()
logits torch. matmul (x, torch matmul (H, y.T))
logits logits torch. (logits, 1)
31
labels torch. arange
32
curl_loss cross_entropp. loss (osgrr (logite, labels)
33
### Reward Prediction Loss reward_pred revard.prediitor __s)q)
36
reward_loss torch. mse_ loss (reward_pred, reward)
37
38
return taco faso_loss curl_loss reward_loss
Listing Pytorch-like pseudo-code how TACD objective computed
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, encoder layers and 50 for the feature dimension 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 x [AA]]. In practice, we find this works well as effectively extracts relevant control information from the 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.
1024 for the hidden dimension of all the
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 Q'D3-BCC' = 2.5 as defined and used in Scott et al. [11]. In experiments. significant performance difference was found for a within the [2, 3] range. In the case CQL, we augment the original critic loss with Q-value regularizer and choose the Q-regularizer weight, acQu, from {0.5, 1, 4, 8}. Table 6
presents the chosen ocQI. for each dataset.
Table 6: Hyperparameter of acor. used in different tasks/datascts.
Task/ Dataset
aCQL
Medium
0.5
Hopper Hop
Medium-replay
0.5
Replay
Medium
0.5
Cheetah Run
Medium-replay
Replay
Medium
0.5
Walker Run
Medium-replay
Replay
Medium
0.5
Quadruped Run
Medium-replay Replay
For IQL, adopt the update functions of both policy and value function from its open-source JAX implementation into DrQ-v2, setting the inverse temperature 3 to 0.1, and T = 0.7 for the expectile. Lastly, for the Decision Transformer (DT), adapt from the original open-source implementation by Chen et al. [7] and use a context length of 20.
17

---

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 ac-
Pefformarce
tions together. To test the algorithm's sensitivity to this hyperpa- rameter, we conduct an experiment on the Quadruped Run task, which has 12-dimensional action space. In Figure 10, show the 1M performance of TACO with different choices of latent ac- tion dimensions. Notably, observe that long the dimen- sionality of the action space is neither too small (6-dimensional), which could limit the ability of latent actions capture sufficient information from the raw actions, nor too large (24-dimensional), which might introduce excessive control- irrelevant information, the performance of TACO remains robust the choice of action representation dimensionality.
12 Oimosleddly Action Represertaiinn
Figure 10: 1M Performance TACO with different dimensionality of latent action representations on Quadruped Run (AA) 12). Error bar repre- standard deviation across ran- dom seeds.
E
Sensitivity of TACO vs. DrQ-v2 to random seeds
Quadruped
1ACD
1A00
D:0-22
100
850
Rewarr
Ephoode Recret
JM Episode Reward
(a) Cheetah Run
(b) Acrobot Swingup
(c) Quadruped Run
Reacher Seach
Uopper
TACO
TACO
D/0-22
Drov2
539
550
2M Ep 50de Naward
19 sodl Naward
1M Eploi hower
(d) Quadruped Walk
(c) Reacher Hard
(1) Hopper Hop
Finger Turn Hand
Reach
TACO
TACO
TACO
02
Dro-2
400
1M Episoda howard
Epuede Naward
1M tgroda Neward
(g) Finger Tumn Hard
(h) Reach Duplo
(i) Walker Run
Figure 11: Robustness analysis TACO vs. DrQ-v2 across multiple (6) random seeds
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 bascline 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 essential for successful deployment.
18

---

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 specified goal location. A visualization of successful trajectory illustrated 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 result, despite training for 30 million frames, as demonstrated in Figure 12b, DrQ-v2 fails to successfully complete this task. In comparison, TACO manages successfully solve this task for half of the random initializations, i.e., 3 out of 6 seeds.
T=0
T=20
T=40
T=60
T=80
T=100
T=120
T=140
T=160
T=180
(a) Visualization of an expert trajectory on manipulator bring ball.
Manipulator Bring Ball
600
DrQ-v2
500
TACO
Episode Reward
400
300
200
100
1m
10m
15m
20m
25m
30m
Number Environment Steps
Online learning manipulator bring ball aggregated 6 random seeds.
Figure 12: Manipulator Bring Ball Task
19

---

G
Visualization of Meta-world Tasks
T=0
T=20
T=40
T=60
T=80
T=100
Hammer
Assembly
Disassemble
Stick Pull
Pick Place Wall
Hand Insert
Figure 13: Visualization of expert trajectories for each Meta-world task
20

---

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 to denote the action sequence At, A4t from timestep t to t + K, and we will use Upt+K denote the sequence of latent actions Ut.... ULtK from timestep to + K.
A
Az+1
Ut+1
St
Sc+1
4
Ze+1
Figure 14: The graphical diagram
H H-1-K Σ R' with the conditional indepen-
Proposition H.1. Let X be the return to go
dence assumptions implied by the graphical model in Figure 14. If I(Z++Ki Ze, UL++K-1) I[S+*K; St, K-1), then [[;; Ze, UL(+K-1) I(X; St, A
Proof. We prove by contradictions. Suppose there exists pair of state action representation øz, thy and a reward function r such that I(Z Ze. U((4K-1) = I(Sr4K; S,, --I+ but [X;; Ze, U24+K-1)
(X;;S St, A Then it suffices show that 1(S44R; Zz, U1a+K-1) < I(S:KK; St,UUKaKKil which would give us the desired contra- diction since [(2442i Zz, U U144K-1) S 1(Su4a5 Zz U-)) and I(S+ St, U2(+K-1) S I[S+K;; St, A1)+K-1).
I(X; -,,KU--KK)))
Now we look at I(X; Zi, Ula+K-1) Apply chain rule of mutual information:
T(x; Ze, St, UnLt K-1) = I[= X; Z, S, U + S( U
(6)
=0+ T(X; Ss U+++K-1)
(7)
Applying the chain rule in another way, we get
I(x; Ze, St, UL+4K11) IX: S, Z,, U + I(X; Ze, U
(8)
Therefore, we get
Sx U = (Xx; U [(X; 2,, U
(9)
By our assumption that 1 1(x; 2,, U < ((; U1sKK11) we must have
(10)
Next, we expand 1(8.4K; Ze, St, 11)+
I(S4k: Ze, St, U = I(S+KK; Zi U+4+K-1 + I(S S.
(11)
=0+ 10+1(S St, U1)+K-1)
(12)
On the other hand, we have
T(S++K; Za St, UL+tK-1) = T(S Un44 T(StKK Ze, Un++K-1)
(13)
(14)
Thus we have
I(S+*K; S. U + Z,, Ut(tK-1)= [(S St. U+(+K-1)
(15)
But then because I(St+ Sz U-(+k=1) > T(X; UI++K-1) as S, + S > forms a Markov chain, it greater than zero by the Inequality (10). As a result, I(S+*K; Ze, UK+
T(S x; St, U+++K-1) < T(S k; S., Am+*x-).. This is exactly the contradiction that we would like to show.
21

---

Before proving Theorem 3.1, we need to cite another proposition, which is proved as Lemma 2 in Rakely al. [41]
Proposition H.2. Let X, Y.Z be random variables. Suppose I(Y,Z) I(Y,)) and Y ZIX, z)d2
then X s.t. Vx, pe,p((Y = z =
Proof of Theorem 3.1: Based on the graphical model, it is clear that
max
T 1 Ze, Ur. U+tK-11 -1) = I( [St, Au,.., ActK-1)) (16)
6,4
Now define the random variable of return-to-go R4 such that
H-t
R4 =
Re+k
(17)
K=
Based on Proposition H.1, because
I(Z+4Ki Ze U144K-1) I(S++K3s St A
we could conclude that
I(R4Kii Ze, U1i-K-1) S4
(18)
Now applying Proposition H.2, we get
Lippeiiiilllleeee (p(RelZ, U1++k p( p s, (19)
As a result, when = 1, for any reward function r, given state-action pair (81, 01); (82, 02) such that (081) (((1),
va(a2), we have (s,, 01) =E 1S,
E PT. p, R.y. This is because p(Re(S, $1, =a)) p(, St = 82, A, a2)
by Equation (18) as p(2 S = 81) S5 82), s(),p(u Ar = a1) p(u 1 02). In case when K > 1, because if E Z Z2, ..., U2+K-11 K, E|Z Ze, Ut.... U+] = E[S [S,, At.... A...,11 including K
EJS Ar, Ar. A then for
any 1 S by Data processing Inequality. (Intuitively, this implies that if the information about the transition dynamics at specific step lost, the mutual information decreases as the timestep progresses, making it impossible to reach its maximum value horizon K.) Then the same argument should also apply here.
I Additional related work
L1 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 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 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 selfssupervised//ontrassvvv learning in reinforcement learning
In 1, we summarize the works that apply seffsuuuprriiee/oootastieee learning objectives to improve the sample efficiency of visual reinforcement learning. In this section, discuss the additional works that apply alffsuuppevised/coottsivee learning objectives to broader set of topics in reinforcement learning.
Several recent works have investigated the use of suressupprriisd/oontrstiiee learning objectives to pre-train representations for reinforcement learning (RL) agents. Parisi et al. [38] propose PVR.
22

---

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 represen- tations using masked autoencoders. Additionally, Ma et al. [36] propose VIP, which formulates the representation learning problem as offline goal-conditioned RI. and derives self-supervised dual goal-conditioned value-function objective.
Besides pretraining state representations, in goal-conditioned RL, the work by Eysenbach et al. [10] establishes connection between learning representations with a contrastive loss and learning a value function. Moreover, sulf--uupr learning has also been employed for effective exploration in RL. Guo et al. [14] introduce BYOL- -YO--Explore, a method that leverages the self-supervised B 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, ProtoRI. by Yarats et al. [51], presents sulf-supervised framework for learning state representation through a clustering baaen self-superviser learning objective in the reward-free -reee exploration setting.
23


Reviews:
[{"review_id": "ITljB7ic2s", "reviewer": "Reviewer_gh9U", "summary": "This 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.", "questions": "1. 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.", "limitations": "See Weaknesses.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback and review! We are encouraged that you recognize the novelty of our approach in joint state-action representation learning and appreciate the promising empirical results of TACO. Below we first address your two concerns.\n\n---\n**Offline RL ablation**:\nWe conduct additional experiments in which we remove the CURL and reward prediction loss from TACO, referring to this modified version as \"TACO w/o R&C.\" In **Figure 16** of the attached PDF, we present the normalized rewards for both CQL and TD3+BC with TACO, TACO w/o R&C, and without TACO. These results underscore that the removal of reward prediction and CURL losses does not significantly diminish TACO's performance, emphasizing that they are not the primary drivers of TACO's superior performance.\n\n**Comparison with LASER**:\nFirst, LASER and TACO differ significantly in their action representation learning. While TACO uses action representation only for critic learning, LASER alternates between policy/critic learning and action representation learning, potentially causing instability in RL learning, as the agent's latent action space continuously evolves throughout the training process.\n\nLASER does not make their code publicly available, and it builds upon the state-input SAC algorithm without providing comprehensive experimental details. Therefore, we re-implement their learning objectives from scratch based on DrQ-v2. On Walker Run, we evaluated the performance of LASER with different dimensionalities of the learned latent action space.\n\n|    | LASER (dim=4)| LASER (dim=6)| LASER (dim=8) | DrQ-v2 | TACO |\n| -----------   | ----------- | -----------|-----------|-----------|-----------|\n| Walker Run  | 251 +- 79   | 492 +- 51  | 511 +- 41| 517 +- 43| 637 +- 11 |\n\n(Note: the original action space of Walker Run is 6 dimensional and TACO's latent action space dimension is 8.)\n\nOur findings suggest that LASER's performance suffers from the iterative updates between RL and action representation learning, limiting its improvement over DrQ-v2, particularly when compressing the action space. However, we must acknowledge that due to the time constraints for the rebuttal phase, our choice of hyperparameters for LASER may not have been optimal, potentially failing to reproduce their best results.\n\n---\n\nNext, here is the response to your specific questions.\n>Question 1: Is there a specific reason why the authors use CURL as an additional auxiliary tasks (as opposed to e.g. SODA)?\n\nWe chose to use CURL, but we could also apply SODA loss in place of CURL loss. We test TACO with SODA loss on the Quadruped Run and Walker Run tasks, As shown in **Figure 15 (Left)** of the attached PDF, the results suggest that for Quadruped Run, CURL as the auxiliary loss outperforms SODA, but yields comparable performance on Walker Run.\n\n>Question 2: Clarification on Line 215-231 (Figure 5)\n\nThe primary objective of this experiment is to assess whether TACO's learned action representation can effectively extract control-relevant action information. To further clarify, here is the Gym-like pseudocode for our modified Cheetah Run environment:\n\n```\nclass CheetahRunNew:\n    def __init__(self, env, action_dim=6, distract_dim=20):\n        self.orginal_env = env\n        self.action_space = Box(-1.0, 1.0, shape=(action_dim + distract_dim))\n        self.action_dim = action_dim\n    def step(self, action):\n        return self.orginal_env.step(action[:self.action_dim])\n```\nIn this modified Cheetah Run environment, while the action space dimensionality has been expanded to 26, only the first 6 dimensions are utilised. We then train TACO in this modified environment, setting the dimensionality of learned latent action embedding to be 6.\nTo evaluate whether the learned action representation indeed captures the information of the first 6 dimensions, we sample four actions, $a_1, a_2, a_3, a_4 \\in R^6$, 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. \nOur hypothesis is that actions with the same centroid should possess similar latent representations, given that they are fundamentally the same action. To validate this, we generated a t-SNE action embedding plot (**Figure 5** of our original manuscript), which revealed successful clustering of semantically similar actions.\n\n>Question 3: Is there a reason why the authors chose to let the action representation be independent of the state?\n\nThe learned action representation does not strictly need to be state-independent. Indeed, we have explored the possibility of allowing the action representation to be dependent on the latent state. Through empirical testing on the Quadruped Run and Walker Run tasks (**Figure 15 (Right)** of the attached PDF), we find that whether the action representation is state-dependent or independent does not create a major difference in the performance of TACO. Thus, we opted for the simpler approach, allowing the action representation to be state-independent. We acknowledge that our empirical results do not rule out the possibility of allowing state dependent action representations, and we will clarify this in our revised manuscript.\n\n>Question 4: Is the action representation being compressed? \n\nThe action representation in our approach is not primarily aimed at compression. Instead, our focus is on shaping this representation to align with the optimal Q function. Drawing an analogy with state representation, we assert that a valuable action representation ought to be able to linearly represent the optimal Q function and should also be predictive of subsequent states. These guiding principles inform our method in TACO, where the action embedding is systematically learned through a thoughtful integration of critic TD loss and temporal contrastive loss updates. We appreciate the reviewer's insight and question, and we will clarify this point in our revised manuscript.\n"}, {"review_id": "hMhDgiS0Jl", "reviewer": "Reviewer_1TK9", "summary": "The 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.", "questions": "- 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?", "limitations": "The 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful feedback and review! We are encouraged that you recognize TACO's simplicity, flexibility, outstanding performance, and theoretical analysis, all of which contribute to the strength of our approach. Below we address the concerns and questions that you have raised.\n\n---\n\n> Novelty of our approach: comparison with DRIML\n\nWe would like to clarify the key distinction between our approach and DRIML, as well as other representation learning baselines. Below we list the comparison with DRIML in bullet points.\n\n1. DRIML & SPR builds on top of C51/DQN algorithm which do not extend to environments with continuous action space.\n2. DRIML & SPR focus on environments with small, well-represented, abstract discrete action spaces, overlooking the importance of action representation learning.\n3. In contrast, we identify the importance of action representaion learning in continuous control, an under-explored topic in previous works. We introduce TACO as a simple yet effecive approach to utilize temporal contrastive loss to learn state and action representations.\n4. For the comparison with DRIML in Table 2, we also incorporate action representation learning in the same way as TACO and only focus on the comparison of the design of temporal contrastive loss.\n5. Non-trivial improvement on Hopepr Hop and Acrobot Swingup: We still observe notable enhancements in tasks such as hopper hop (showing a 20.8% improvement) and acrobot swingup (with an 8.5% increase). This can be attributed to DRIML's contrastive loss positive pairs being policy-dependent, potentially causing stability concerns during policy updates. Contrarily, TACO's design is policy-neutral, offering a more stable solution.\n\nIn summary, the main contribution of our paper lies in both identifying the importance of action representation in continuous control and proposing TACO as a simple yet effective solution to address this problem.\n\n---\n\n> Extra parameter K is required to tune\n\nWe ackowledge that this is one limitation of TACO. However, we would like to point out that in TACO, we only use K=1 or K=3 in all of our experiments. Thus, it requires minimum hyperparameter tuning efforts to find the best K. \n\n\n---\n\n> It would be better if the proposed method can be evaluated in more environments besides DMC.\n\nWe appreciate your suggestion to evaluate our method in diverse environments. In response, we have selected six challenging tasks within the **Meta-world** domain to test the online learning performance of TACO against the DrQ-v2 baseline. The learning curve, presented in **Figure 17** of the attached PDF file, once again demonstrates TACO's significant performance improvement. It successfully solves complex tasks such as hammer, assembly, disassemble, stick pull, and pick place wall, which DrQ-v2 cannot accomplish within 2 million online interaction steps. These results underscore that the insights of TACO extend beyond the locomotion tasks of the DeepMind Control Suite to complex and intricate robotic manipulation tasks.\n\n---\n\n> How do you ensure the optimal action will be the same?\n> \nIn lines 132-134, we're not claiming that the optimal action will be identical. The result indicates that if two state-action pairs possess the same state and action representations, their optimal Q values coincide. Thus, the optimal value function can be factorized as $Q^*(\\phi(s), \\psi(a))$. This means for a given state, if two actions share the same representation, they'll have an equal optimal Q value. Thus, optimal actions do not have to be unique.\n\n---\n\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\nThis is a design choice that we have made, following the same design choice as CURL and CPC. We could instead use the cosine similiarity measure, as done in MoCo, SimCLR, and CLIP. In our empirical evaluation on the quadruped run task, we find that using cosine similarity lead to a 3.5% decrease in performance at the 1M mark (541 ± 38 vs. 522 ± 56). \n\n---\n\n>  Is the action encoder only used in the temporal contrastive loss or is it also incorporated in policy learning?\n\nThe action encoder is utilized both in the temporal contrastive loss and the critic learning, as illustrated in Figure 2. It is updated based on both the TACO loss as well as the Temporal Difference (TD) loss from the critic functions. \n\n---\n\n> GPU memory cost:\n\nTACO's batch size of 1024 is four times that of DrQ-v2 but also only a quarter of the standard batch size in the contrastive learning literature. Using the same CNN architecture as DrQ-v2, TACO's GPU memory requirement for the Quadruped Run task is 4.8 GB, in contrast to DrQ-v2's 2.2 GB. This fits easily under a single RTX 2080 Ti GPU."}, {"review_id": "jM8NH2imRp", "reviewer": "Reviewer_Y2ee", "summary": "This 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", "questions": "It 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", "limitations": "The authors have discussed limitations. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The 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", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback and review! We are encouraged by your recognition of the broad applicability of our approach, as manifested in our application of TACO for both model-free and offline visual RL settings. \n\n---\nTo address your question of CURL and reward prediction loss, below we conduct a comprehensive experiment by running TACO without CURL and reward prediction loss on five tasks (same as the five tasks in **Table 2** of our original manuscript): Quadruped Run, Hopper Hop, Walker Run, Reacher Hard, Acrobot Swingup. Here we show the 1M online RL performance.\n\n|               | TACO     | TACO w/o reward & CURL| DrQv2|\n| -----------    | ----     | ----------- |----------- |\n| Quadruped Run  | 541 +- 38|  501 +- 24  | 407 +- 21|\n| Walker Run     | 637 +- 21|  615 +- 11  | 517 +- 43|\n| Hopper Hop     | 261 +- 38| 242 +- 12   | 192 +- 41|\n| Reacher Hard   | 883 +- 21| 882 +-  67  | 572 +- 51|\n| Acrobot Swingup| 241 +- 21| 301 +-  42  | 210 +- 12|\n\nAs the experimental result suggests, for initial test tasks like Quadruped Run, Walker Run, and Hopper Hop, omitting reward prediction and CURL loss seems to have a slight negative influence on performance, albeit not substantial. This observation led us to incorporate these two auxiliary losses into our final objective. Interestingly, we also found that in the Acrobot Swingup task, the agent's 1M performance actually improved when these two losses were removed, increasing from 241 to 301.\nThese experimental findings reinforce our claim that while CURL and reward prediction loss could further improve the performance in many tasks, the proposed temporal contrastive loss of TACO is indeed the central and most impactful component.\n\n---\n\nNext, we would like to address your two other questions.\n\n> Question1: I wonder if the authors have tried using it for pretraining?\n\nExploring the use of TACO loss as a self-supervised pretraining objective for both state and action representation is indeed an exciting future direction. In the current paper, our primary focus lies in online and offline RL learning from scratch without any prior knowledge. Recognizing the interest in this aspect, and despite the limited rebuttal period, we conduct an additional experiment to highlight the potential of using TACO for pretraining these representations. The detailed results can be found in **Figure 18** of our attached PDF.\n\nSpecifically, we experimented with pretraining on an offline dataset of **Walker Walk Replay** (generated in a manner consistent with our \"Replay\" dataset in offline RL experiments) to test the learned representation's generalization to a new task, **Walker Run**, evaluating with both online RL and few-shot behavior cloning. For online RL, we initialize DrQ-v2 with the pretained state and action encoders. For few-shot behavior cloning, we initialize the policy with the pretrained state encoder. As demonstrated in **Figure 18** of the attached 1-page PDF, by pretraining on the walker walk dataset, the state and action representation indeed captures the essential information for the shared embodiment (walker) across two different tasks. Thus, it facilitates both efficient online RL training and few-shot imitation learning.\n\n---\n\n> Question 2: I could not find the architectural details of H_\\theta and G_\\theta, could the authors specify the details of these?\n\nG is a two-layer Multilayer Perceptron (MLP) where the input size is observation feature dimension plus K (number of timesteps in TACO) times the latent action dimension. The output size of G is the same as the observation feature dimension. H is also a two-layer MLP, with both its input and output sizes being the observation feature dimension. Both G and H utilize a hidden layer size of 1024, and same as DrQ-v2, the observation feature dimension used in TACO is 50."}, {"review_id": "mtuQqh3Y9a", "reviewer": "Reviewer_oD4S", "summary": "This 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.", "questions": "- 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]).", "limitations": "The 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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful feedback and review! Below we address the concerns and questions that you have raised. We are encouraged that you appreciate TACO's significant empirical performance and recognize the comprehensiveness of our experiments in both online and offline RL, alongside the ablation studies for each loss term.\n\n---\n\nWe would like to first explain the question around the implementation of other representation learning baselines.\n> The authors mention that they implemented some baselines based on DrQ-v2 and additionally considered action as an input. However, there is no detailed explanation of this implementation, making it somewhat difficult to ensure a fair comparison. \n> Could you provide a more detailed explanation of how action inputs were considered in ATC and SPR in the existing methodologies?\n\nThe learning objective of ATC loss does not involve action input. For SPR, it is implemented as follows. \n\n```\nz_hat = state_encoding(observation[0])\nspr_loss = 0\nfor k in range(K):\n    u = action_encoding(action[k])\n    z_hat = h(torch.concat([z_hat,u], dim=-1)) ### transition model\n    z_next = state_encoding_target(observation[k+1]) ### target encoder\n    y_hat = q(g(z_hat)) \n    y_next = g_target(z_hat)\n    spr_loss += -cosine_similarity(y_hat, y_next)\n```\nHere, h is a latent transition model, g is a projection layer, and q is the prediction head. Following the insights of TACO, we let the critic and SPR loss share the same acton encoder so that we only focus on the comparison of the temporal contrastive loss. In other words, we update action encoder from both SPR loss as well as the critic's TD loss. (See my comments below for a detailed discussion on this design choice for TACO.)\n\n---\n\n> 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> What is the difference between extending DRIML to a continuous action space and DRIML with the DrQ-v2 backbone you used as a baseline?\n\nWe would like to clarify the key distinction between our approach and DRIML, as well as other representation learning baselines.\n\n1. The original DRIML and SPR papers, build their methods on top of the C51 algorithm (DQN for SPR) and is specifically tailored to environments with discrete action spaces. C51 algorithm itself cannot extend to continuous action spaces. This is why we choose to re-implement DRIML on top of DrQ-v2, a simple yet strong online RL algorithm for visual continuous control.\n2. DRIML focuses on environments with small, well-represented, abstract discrete action spaces, overlooking the importance of action representation learning. \n3. In contrast, we identify the importance of action representaion learning in continuous control, an under-explored topic in previous works. We introduce TACO as a simple yet effecive approach to utilize temporal contrastive loss to learn state and action representations.\n4. For the comparison with DRIML in Table 2, we also incorporate action representation learning in the same way as TACO and only focus on the comparison of the design of temporal contrastive loss.\n5. Non-trivial improvement on Hopepr Hop and Acrobot Swingup: We still observe notable enhancements in tasks such as hopper hop (showing a 20.8% improvement) and acrobot swingup (with an 8.5% increase). This can be attributed to DRIML's contrastive loss positive pairs being policy-dependent, potentially causing stability concerns during policy updates. Contrarily, TACO's design is policy-neutral, offering a more stable solution.\n\nIn summary, the main contribution of our paper lies in both identifying the importance of action representation in continuous control and proposing TACO as a simple yet effective solution to address this problem.\n\n\n---\n\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\nThe action encoder is updated through both the critic's TD loss and the TACO loss, allowing it to learn more informative action representations. Specifically, when we exclude the TD loss from updating the action encoder, there is a noticeable performance drop, with the 1M online performance falling from 541 +- 38 to 492 +- 44 on Quadruped Run, and from 261 +- 52 to 211 +- 86 on Hopper Hop.\n\n---\n\n> Does the performance decrease if the action encoder is not used and only projection is performed?\n\nYes, the performance does indeed decrease if the action encoder is not used and only projection is performed. By excluding the action encoder and taking raw action as input for the critic, the 1M online performance drops to 499 +- 32 on Quadruped Run and to 221 +- 36 on Hopper Hop, down from 541 +- 38 and 261 +- 52, respectively. Compared to DrQ-v2's 1M performance of 407 +- 21 for Quadruped Run and 192 +- 41 for Hopper Hop, we still observe a noticeable improvement as state and action representations benefit from the temporal contrastive loss. Nevertheless, these findings emphasize the importance of the action encoder and suggest that allowing the critic and temporal contrastive loss to share the same action embedding effectively aids in learning action representations. We appreciate the reviewer's insight and question, and we will clarify this point in our revised manuscript."}]

Author rebuttal:
We thank all reviewers for their insightful questions and valuable feedback. We are encouraged that reviewers recognize the importance of our tackled problem in state-action representation learning (gh9U). They also appreciate the flexibility and applicability of our proposed approach, TACO, in both offline and online RL settings (1TK9, Y2ee, od4S), and highlight the outstanding empirical performance of our methods (gh9U, 1TK9, Y2ee, od4S). We have addressed all individual questions of reviewers in separate responses. 

Additionally, we have attached a one-page PDF with more experimental results, including (**Figure 15**) two ablation studies - one substituting CURL with SODA, another using state-dependent over state-independent action representation; (**Figure 16**) an additional ablation study on CURL and reward prediction loss in offline RL; (**Figure 17**) online RL experiments for six Meta-world manipulation tasks; and (**Figure 18**) further experiments on representation pretraining using TACO.

Here we briefly outline the response to the two most common questions raised by the reviewers. The detailed response is included in each individual response.

**CURL and reward prediction objectives** (Reviewer gh9U, Y2ee):
In our initial tests on the Quadruped Run and Walker Run tasks, we discovered that CURL and reward prediction enhanced TACO's performance, leading us to incorporate these two losses into TACO's final objectives. However, we would like to emphasize again that the CURL and reward prediction objective are added only as auxiliary losses to further improve the performance. We conduct additional experiments on both online (Reviewer Y2ee) and offline RL (Reviewer gh9U) settings. Together with **Figure 6(b)** in our original manuscript, they demonstrate that CURL and reward prediction objectives are not the main driver of the superior performance of TACO.

**Novelty of our approach compared with other representation learning methods such as DRIML** (Reviewer 1TK9, oD4S):
We want to emphasize other representation learning objectives, except for CURL and ATC, were studied in environments with well-represented, small discrete action spaces, thus overlooking the importance of action representation learning. In contrast, we identify the importance of action representation learning in continuous control, an under-explored topic in previous works. We introduce TACO as a simple yet effective approach to utilize temporal contrastive loss to learn state and action representation. In our comparison with other representation learning objectives in the **Table 2** of our original manuscript, we fix the action representation learning part and focus only on the design of temporal contrastive loss. Still, we see some non-trivial improvements over other representation learning objectives. In summary, our contribution lies in both underscoring action representation in continuous control tasks and introducing a simple yet effective temporal contrastive TACO loss as a solution to state and action representation learning in visual continuous control problems.

Meta-review: Reviewers are unanimous in recommending acceptance (although to varying degrees). I have gone through the discussions in detail and am glad to see that these discussions have resulted in an improved paper. I do not think there are any significant concerns remaining and recommend acceptance.
Decision: Accept (poster)