SlowGuess/NeurIPS_2023_OCR huggingface dataset:
space of task Cheetah Run, although only the first six were utilized in environmental interactions. We first train an agent with TACO online to obtain the action representation. Then we select four actions within the original action space  $a_1, a_2, a_3, a_4$  to act as centroids. For each of the four centroids, we generate 1000 augmented actions by adding standard Gaussian noises to the last 20 dimensions. We aim to determine if our action representation could disregard these "noisy" dimensions while retaining the information of the first six. Using t-SNE for visualization, we embed the 4000 actions before and after applying action representation. As shown in Figure 6, indeed our learned action representation could group the four clusters, demonstrating the ability of our action representation to extract control relevant information from the raw action space.
The effectiveness of our temporal contrastive loss is enhanced with a larger batch size. As is widely acknowledged in contrastive learning research [8, 22, 40, 13], our contrastive loss sees significant benefits from utilizing a larger batch size. In Figure 7a, we illustrate the performance of our algorithms alongside DrQv2 after one million environment steps on the Quadruped Run task. As evident from the plot, batch size greatly influences the performance of our algorithm, while DrQ-v2's baseline performance remains fairly consistent throughout
training. In order to strike a balance between time efficiency and performance, we opt for a batch size of 1024, which is 4 times larger than the 256 batch size employed in DrQ-v2, but 4 times smaller than the 4096 which is commonly used in the contrastive learning literature [8, 22, 13]. For an analysis on how batch size affects the algorithm's runtime, we direct the reader to Appendix B.
Figure 6: Left: t-SNE embedding of actions with distracting dimensions. Right: t-SNE embedding of latent representations for actions with distracting dimensions.
efficiency and performance, we opt for a batch size of size employed in DrQ-v2, but 4 times smaller than the learning literature [8, 22, 13]. For an analysis on it to direct the reader to Appendix B.
(a) TACO and DrQ-v2 across different batch sizes.
Figure 7: TACO with different batch sizes and with different component of its learning objectives removed
(b) TACO with different learning objective removed
Reward prediction and CURL loss serves an auxiliary role to further improve the performance of TACO, while the temporal contrastive loss of TACO is the most crucial component. In the practical deployment of TACO, two additional objectives, namely reward prediction and CURL loss,
are incorporated to enhance the algorithm's performance. In Figure 7b, we remove one objective at a time on the Quadruped Run task to assess its individual impact on the performance after one million environment steps. As illustrated in the figure, the omission of TAC0' temporal contrastive objective results in the most significant performance drop, emphasizing its critical role in the algorithm's operation. Meanwhile, the auxiliary reward prediction and CURL objectives, although secondary, contribute to performance improvement to some degree.
InfoNCE-based temporal action-driven contrastive objective in TACO outperforms other representation learning objectives including SPR [43], ATC [44], and DRIML [37]. In Table 2, we have showcased a comparison between our approach and other visual RL representation learning objectives such as SPR, ATC, and DRIML. Given that SPR and DRIML were not initially designed for continuous control tasks, we have re-implemented their learning objectives using the identical backbone algorithm, DrQ-v2. A similar approach was taken for ATC, with their learning objectives also being reimplemented on DrQ-v2 to ensure a fair comparison. (Without the DrQ-v2 backbone algorithm, the performance reproduced by their original implementation is significantly worse.) Furthermore, recognizing the significance of learning action encoding, as discussed earlier, we have integrated action representation learning into all these baselines. Therefore, the model architecture remains consistent across different representation learning objectives, with the sole difference being the design of the temporal contrastive loss. For DRIML, given that only the first action of the action sequence is considered in the temporal contrastive loss, TACO and DRIML differ when the number of steps  $K$  is greater than one. Thus, we indicate N/A for tasks where we choose  $K = 1$  for TACO.
Table 2: Comparison with other objectives including SPR [42], ATC [44], and DRIML [37]
Environment	TACO	SPR	ATC	DRIML	DrQ-v2
Quadruped Run	541 ± 38	448 ± 79	432 ± 54	N/A	407 ± 21
Walker Run	637 ± 21	560 ± 71	502 ± 171	N/A	517 ± 43
Hopper Hop	261 ± 52	154 ± 10	112 ± 98	216 ± 13	192 ± 41
Reacher Hard	883 ± 63	711 ± 92	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 a fixed dataset without further interaction with the environment. This necessitates that the agent effectively generalizes from limited data while handling high-dimensional visual inputs. The state/action representation learning objective of TACO plays a vital role in addressing these challenges by capturing essential information about the environment's dynamics, thereby enabling more efficient generalization and improved performance. TACO can be easily integrated as a plug-and-play module on top of existing strong offline RL methods, such as TD3+BC [11] and CQL [31].
For evaluation, we select four challenging visual control tasks from DMC: Hopper Hop, Cheetah Run, Walker Run, and Quadruped Run. For each task, we generate three types of datasets. The medium dataset consists of trajectories collected by a single policy of medium performance. The precise definition of "medium performance" is task-dependent but generally represents an intermediate level of mastery, which is neither too poor nor too proficient. The medium-replay dataset contains trajectories randomly sampled from the online learning agent's replay buffer before it reaches a medium performance level. The full-replay dataset includes trajectories randomly sampled throughout the online learning phase, from the beginning until convergence. The dataset size for Walker, Hopper, and Cheetah is 100K, while for the more challenging Quadruped Run task, a larger dataset size of 500K is used to account for the increased difficulty. We compute the normalized reward by diving the offline RL reward by the best reward we get during online TACO training.
Table 3: Offline Performance (Normalized Reward) for different offline RL methods. Results are averaged over 6 random seeds.  $\pm$  captures the standard deviation over seeds.
Task/ Dataset	TD3+BC w. TACO	TD3+BC	CQL w. TACO	CQL	DT	IQL	BC
Hopper Hop	Medium	52.4 ± 0.4	51.2 ± 0.8	47.9 ± 0.6	46.7 ± 0.2	40.5 ± 3.6	2.0 ± 1.7	48.2 ± 0.8
Medium-replay	67.2 ± 0.1	62.9 ± 0.1	74.6 ± 0.4	68.7 ± 0.1	65.3 ± 1.8	57.6 ± 1.4	25.9 ± 3.2
Full-replay	97.6 ± 1.4	83.8 ± 2.3	101.2 ± 1.9	94.2 ± 2.0	92.4 ± 0.3	47.7 ± 2.2	65.7 ± 2.7
Cheetah Run	Medium	66.6 ± 0.5	66.1 ± 0.3	70.1 ± 0.4	66.7 ± 1.7	64.3 ± 0.7	1.7 ± 1.1	62.9 ± 0.1
Medium-replay	62.6 ± 0.2	61.1 ± 0.1	72.3 ± 1.2	67.3 ± 1.1	67.0 ± 0.6	26.5 ± 3.2	48.0 ± 3.1
Full-replay	92.5 ± 2.4	91.2 ± 0.8	86.9 ± 2.4	65.0 ± 3.9	89.6 ± 1.4	14.6 ± 3.7	69.0 ± 0.3
Walker Run	Medium	49.2 ± 0.5	48.0 ± 0.2	49.6 ± 1.0	49.4 ± 0.9	47.3 ± 0.3	4.4 ± 0.4	46.2 ± 0.6
Medium-replay	63.1 ± 0.6	62.3 ± 0.2	62.3 ± 2.6	59.9 ± 0.9	61.7 ± 1.1	41.4 ± 2.8	18.5 ± 0.8
Full-replay	86.8 ± 0.6	84.0 ± 1.6	88.1 ± 0.1	79.8 ± 0.6	81.6 ± 0.8	18.1 ± 3.7	30.8 ± 1.8
Quadruped Run	Medium	60.6 ± 0.1	60.0 ± 0.2	58.1 ± 3.7	55.9 ± 9.1	14.6 ± 3.8	0.8 ± 0.8	56.2 ± 1.1
Medium-replay	61.3 ± 0.3	58.1 ± 0.5	61.9 ± 0.2	61.2 ± 0.9	19.5 ± 2.2	58.4 ± 4.4	51.6 ± 3.3
Full-replay	92.6 ± 0.7	89.3 ± 0.4	92.1 ± 0.1	85.2 ± 2.5	14.5 ± 1.1	36.3 ± 5.9	57.6 ± 0.7
Average Normalized Score	71.0	68.2	72.1	66.7	48.4	25.8	54.9
We compare the performance of TD3+BC and CQL with and without TACO on our benchmark. Additionally, we also compare with the decision transformer (DT) [7], a strong model-free offline RL baseline that casts the problem of RL as conditional sequence modeling, IQL [30], another commonly used offline RL algorithm, and the behavior cloning (BC) baseline. For TD3+BC, CQL and IQL, which were originally proposed to solve offline RL with vector inputs, we add the their learning objective on top of DrQ-v2 to handle image inputs.
Table 3 provides the normalized reward for each dataset. The results underscore that when combined with the strongest baselines TD3+BC and CQL, TACO achieves consistent performance improvements across all tasks and datasets, setting new state-of-the-art results for offline visual reinforcement learning. This is true for both the medium dataset collected with a single policy and narrow data distribution, as well as the medium-replay and replay datasets with a diverse distribution.
5.1 Contrastive learning in visual reinforcement learning
Contrastive learning has emerged as a powerful technique for learning effective representations across various domains, particularly in computer vision [47, 8, 22, 23, 49]. This success is attributed to its ability to learn meaningful embeddings by contrasting similar and dissimilar data samples. In visual reinforcement learning, it's used as a self-supervised auxiliary task to improve state representation learning, with InfoNCE [47] being a popular learning objective. In CURL [33], it treats augmented states as positive pairs, but it neglects the temporal dependency of MDP. CPC [47], ST-DIM [2], and ATC [44] integrate temporal relationships into the contrastive loss by maximizing mutual information between current state representations (or state histories encoded by LSTM in CPC) and future state representations. However, they do not consider actions, making positive relationships in the learning objective policy-dependent. DRIML [37] addresses this by maximizing mutual information between state-action pairs at the current time step and the resulting future state, but its objective remains policy-dependent as it only provides the first action of the action sequence. Besides, ADAT [29] and ACO [59] incorporate actions into contrastive loss by labeling observations with similar policy action outputs as positive samples, but these methods do not naturally extend to tasks with nontrivial continuous action spaces. A common downside of these approaches is the potential for unstable encoder updates due to policy-dependent positive relations. In contrast, TACO is theoretically sufficient, and it tackles the additional challenge of continuous control tasks by simultaneously learning state and action representations.
In addition to the InfoNCE objective, other self-supervised learning objective is also proposed. Approaches 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 a significant challenge in effectively capturing the underlying dynamics. In contrast, by framing the latent dynamics model predictions as a self-supervised InfoNCE objective, the mutual information
guided approach used by TACO is better suited for continuous control task, resulting in more stable optimization and thus better state and action representations.
5.2 Action representation in reinforcement learning
Although state or observation representations are the main focus of prior research, there also exists discussion on the benefits and effects of learning action representations. Chandak et al. [6] propose to learn a policy over latent action space and transform the latent actions into actual actions, which enables generalization over large action sets. Allshire et al. [1] introduce a variational 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 are essential. To this end, Park et al. [39] propose an approach to train an environment model in the learned latent action space. In addition, action representation also has the potential to improve multi-task learning [25], where latent actions can be shared and enhance generalization.
In this paper, we have introduced a conceptually simple temporal action-driven contrastive learning objective that simultaneously learns the state and action representations for image-based continuous control — TACO. Theoretically sound, TACO has demonstrated significant practical superiority by outperforming SOTA online visual RL algorithms. Additionally, it can be seamlessly integrated as a plugin 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 envisage two primary directions for future research. Firstly, the creation of more advanced temporal contrastive InfoNCE objectives that can function effectively with smaller data batches may mitigate the concerns related to computational efficiency. Secondly, the implementation of a distributed version of TACO, akin to the strategies employed for DDPG in previous works [3, 24], could significantly enhance training speed. These approaches offer promising avenues for further advancements in visual RL.
Zheng, Wang, Sun and Huang are supported by National Science Foundation NSF-IIS-FAI program, DOD-ONR-Office 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.
[1] Arthur Allshire, Roberto Martin-Martin, Charles Lin, Shawn Manuel, Silvio Savarese, and Animesh Garg. Laser: Learning a latent action space for efficient reinforcement learning. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 6650-6656. IEEE, 2021. 2, 10
[2] Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. 2, 9
[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 on optimal representations for reinforcement learning. Advances in neural information processing systems, 32:4358-4369, 2019. 22
A Effects of the choice of prediction horizon  $K$
TACO functions as a versatile add-on to current online and offline RL algorithms, requiring only the prediction horizon  $K$  as an additional hyperparameter. In our investigations, we select  $K$  as either 1 or 3 across all tasks. The performance of TACO with  $K$  values of 1 and 3 is compared against the DrQ-v2 baselines in Figure 8. As illustrated in the figure, TACO outperforms the baseline DrQ-v2 for both  $K = 1$  and 3, indicating the effectiveness of the temporal contrastive loss in concurrent state and action representation learning. In comparing  $K = 1$  to  $K = 3$ , we observe that a longer prediction horizon ( $K = 3$ ) yields superior results for four out of nine tasks, specifically Hopper Hop, Quadruped Walk, Acrobot Swingup, and Reacher Hard. Conversely, for the Quadruped Run and Walker Run tasks, a shorter temporal contrastive loss interval ( $K = 1$ ) proves more beneficial. For the remaining tasks, the choice between  $K = 1$  and  $K = 3$  appears to have no discernable impact.
Figure 8: 1M Performance of TACO with step size  $K = 1$  and  $K = 3$ .
Our speculation hinges on the rate of changes in the agent's observations. While Theorem 3.1 applies regardless of the prediction horizon  $K$ , a shorter horizon (such as  $K = 1$ ) can offer a somewhat shortsighted perspective. This becomes less beneficial in continuous control tasks where state transitions within brief time intervals are negligible. Therefore, in environments exhibiting abrupt state transitions, a larger  $K$  would be advantageous, whereas a smaller  $K$  would suffice for environments with gradual state transitions.
To substantiate this conjecture, we conduct a simple experiment on the task Hopper Hop where we aim to simulate the different rates of changes in the agent's observations. We verify this by varying the size of action repeats, which correspond to how many times a chosen action repeats per environmental step. Consequently, a larger action repeat size induces more pronounced observational changes. In our prior experiments, following the settings of DrQ-v2, we fix the size of the action repeat to be 2. In this experiment, we alter the size of the action repeat for Hopper Hop to be 1, 2, and 4. For each action repeat size, we then compare the 1M performance of TACO with the prediction horizon  $K$  selected from the set  $\{1,3,5\}$ . Interestingly, as demonstrated in Figure 9, the optimal  $K$  values for different action repeats are reversed: 5 for an action repeat of 1, 3 for an action repeat of 2, and 1 for an action repeat of 4. This observation further substantiates our assertion that the optimal choice of prediction horizon correlates with the rate of change in environmental dynamics.
Figure 9: 1M Performance of TACO with different prediction horizon  $K$  under different action repeat sizes. Shaded columns correspond to the best prediction horizon  $K$  under a fixed action repeat size.
B Time-efficiency of TACO
In this section, we compare the time efficiency of different visual RL algorithms. In Table 4, we present a comprehensive comparison of the speed of these algorithms, measured in frames per second
(FPS). Additionally, we also provide FPS for TACO with different batch sizes in Table 5 as a reference. To ensure a fair comparison, all the algorithms are tested on Nvidia A100 GPUs.
As mentioned in §4.1, the InfoNCE objective in TACO requires a large batch size of 1024, which is 4 times larger than the batch size used in DrQ-v2. Consequently, this increases the processing time, making our method to run about 3.6 times slower than DrQ-v2, with similar time efficiency as DrQ. Therefore, the primary limitation of our method is this time inefficiency caused by the usage of large batch sizes. Potential solutions to improve the time efficiency could involve the implementation of a distributed version of our method to speed up training. We intend to explore these improvements in future work related to our paper.
Table 4: Frames per second (FPS) for visual RL algorithms. B stands for batch size used in their original implementations.
TACO (B:1024)	DrQv2 (B:256)	A-LIX (B:256)	DrQ (B:512)	CURL (B:512)	Dreamer-v3 (B:256)	TDMPC (B:256)
35	130	98	33	20	31	22
Table 5: Frames per second (FPS) for TACO with different batch size
TACO (B:1024)	TACO (B:512)	TACO (B:256)
35	65	94
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.
1 ## Extract feature representation for state and actions.   
2 def update(batch): obs, action_sequence, reward, next_obs = batch   
4 # Update Agent's critic function update_critic(obs, action_sequence, reward, next_obs)   
6 # Update the agent's value function update Actor(obs)   
8 # Update TACO loss   
9 update_taco(obs, action_sequence, reward, next_obs)
Listing 1: Pytorch-like pseudo-code how TACO is incorporated into the update function of existing visual RL algorithms.
Next, we demonstrate the pseudo-code of how TACO objective is computed with pytorch-like pseudocode.
1 # stateEncoder: State/Observation Encoder (CNN)  
2 # actionEncoder: Action Encoder (MLP with 1-hidden layer)  
3 # sequenceEncoder: Action Sequence encoder (Linear Layer)  
4 # reward Predictor: Reward Prediction Layer (MLP with 1-hidden layer)  
5 # G: Projection Layer I (MLP with 1-hidden layer)  
6 # H: Projection Layer II (MLP with 1-hidden layer)  
7 # aug: Data Augmentation Function (Radnom Shift)  
8 # W: Matrix for computing similarity score  
9  
10 def compute_taco_objective(obs, action_sequence, reward, next_obs):  
11 ## Compute feature representation for both state and actions.  
12 z = stateEncoder(aug(obs))  
13 zanchor = stateEncoder(aug(obs), stop_grad=True)  
14 next_z = stateEncoder(aug(next_obs), stop_grad=True)  
15 u_seq = sequenceEncoder(actionEncoder(action_sequence))  
16 ## Project to joint contrastive embedding space  
17 x = G(torch.cat([z, u_seq], dim=-1))  
18 y = H(next_z)  
19 ## Compute bilinear product x^TWy
Listing 2: Pytorch-like pseudo-code for how TACO objective is computed
## Diagonal entries of x^T Wy correspond to positive pairs
logits = torch/matmul(x, torch/matmul(W, y.T))
logits = logits - torch.max(logits, 1)
labels = torch.arange(n)
taco_loss = cross_entropy_loss(logits, labels)
## Compute CURL loss
x = H(z)
y = H(zanchor).detach()
logits = torch/matmul(x, torch/matmul(W, y.T))
logits = logits - torch.max(logits, 1)
labels = torch.arange(n)
curl_loss = cross_entropy_loss(logits, labels)
## Reward Prediction Loss
reward_pred = reward_prediction(z, u_seq)
reward_loss = torch.mse_loss(reward_pred, reward)
return taco_loss + curl_loss + reward_loss
Then when we compute the Q values for both actor and critic updates, we use the trained state and action encoder. Same as what is used in DrQ-v2, we use 1024 for the hidden dimension of all the encoder layers and 50 for the feature dimension of state representation. For action representation, we choose the dimensionality of the action encoding, which corresponds to the output size of the action encoding layer, as  $\lceil 1.25 \times |\mathcal{A}| \rceil$ . In practice, we find this works well as it effectively extracts relevant control information from the raw action space while minimizing the inclusion of irrelevant control information in the representation. See Appendix D for an additional experiment on testing the robustness of TACO against the dimensionality of action representation.
C.2 Offline RL