SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ruijie Zheng<sup>1</sup> Xiyao Wang<sup>1</sup>
Yanchao Sun  $^{1}$  Shuang Ma  $^{4}$  Jieyu Zhao  $^{1,2}$
Huazhe  $\mathbf{X}\mathbf{u}^{3,4,\S}$  Hal Daumé III1,5,8 Furong Huang1,8
<sup>1</sup> University of Maryland, College Park <sup>2</sup> University of Southern California
$^{3}$  Tsinghua University  $^{4}$  Shanghai Qi Zhi Institute
5 Microsoft Research
rzheng12@umd.edu
Despite recent progress in reinforcement learning (RL) from raw pixel data, sample inefficiency continues to present a substantial obstacle. Prior works have attempted to address this challenge by creating self-supervised auxiliary tasks, aiming to enrich the agent's learned representations with control-relevant information for future state prediction. However, these objectives are often insufficient to learn representations that can represent the optimal policy or value function, and they often consider tasks with small, abstract discrete action spaces and thus overlook the importance of action representation learning in continuous control. In this paper, we introduce TACO: Temporal Action-driven COntrastive Learning, a simple yet powerful temporal contrastive learning approach that facilitates the concurrent acquisition of latent state and action representations for agents. TACO simultaneously learns a state and an action representation by optimizing the mutual information between representations of current states paired with action sequences and representations of the corresponding future states. Theoretically, TACO can be shown to learn state and action representations that encompass sufficient information for control, thereby improving sample efficiency. For online RL, TACO achieves  $40\%$  performance boost after one million environment interaction steps on average across nine challenging visual continuous control tasks from Deepmind Control Suite. In addition, we show that TACO can also serve as a plug-and-play module adding to existing offline visual RL methods to establish the new state-of-the-art performance for offline visual RL across offline datasets with varying quality.
Developing reinforcement learning (RL) agents that can perform complex continuous control tasks from high-dimensional observations, such as pixels, has been a longstanding challenge. A central aspect of this challenge lies in addressing the sample inefficiency problem in visual RL. Despite significant progress in recent years [33, 34, 53, 52, 16, 15, 17, 18, 21, 48], there remains a considerable gap in sample efficiency between RL with physical state-based features and those with pixel inputs. This disparity is particularly pronounced in complex tasks; thus tackling it is crucial for advancing visual RL algorithms and enabling their practical application in real-world scenarios.
$\S$  Equal advising
*Our code is released at https://github.com/FrankZheng2022/TACO.
State representation learning has become an essential aspect of RL research, aiming to improve sample efficiency and enhance agent performance. Initial advancements like CURL [33] utilized a self-supervised contrastive InfoNCE objective [47] for state representation, yet it overlooks the temporal dynamics of the environment. Subsequent works, including CPC [23], ST-DIM [2], and ATC [44], made progress to rectify this by integrating temporal elements into the contrastive loss, linking pairs of observations with short temporal intervals. The objective here was to develop a state representation capable of effectively predicting future observations. A more comprehensive approach was taken by DRIML [37], which incorporated the first action of the action sequence into the temporal contrastive learning framework. However, these methods, while innovative, have their shortcomings. The positive relations in their
contrastive loss designs are often policy-dependent, potentially leading to instability during policy updates throughout the training process. Consequently, they lack the theoretical foundation needed to capture all information representing the optimal policy. Furthermore, these methods except for CURL and ATC typically focus on environments such as Atari games with well-represented, abstract discrete action spaces, thereby overlooking the importance of action representation in continuous control tasks [1, 28]. However, by learning an action representation that groups semantically similar actions together in the latent action space, the agent can better generalize its knowledge across various state-action pairs, enhancing the sample efficiency of RL algorithms. Therefore, learning both state and action representations is crucial for enabling the agent to more effectively reason about its actions' long-term outcomes for continuous control tasks.
In this paper, we introduce Temporal Action-driven COntrastive Learning (TACO) as a promising approach to visual continuous control tasks. TACO simultaneously learns a state and action representation by optimizing the mutual information between representations of current states paired with action sequences and representations of the corresponding future states. By optimizing the mutual information between state and action representations, TACO can be theoretically shown to capture the essential information to represent the optimal value function. In contrast to approaches such as DeepMDP [12] and SPR [42], which directly model the latent environment dynamics, our method transforms the representation learning objective into a self-supervised InfoNCE objective. This leads to more stable optimization and requires minimal hyperparameter tuning efforts. Consequently, TACO yields expressive and concise state-action representations that are better suited for high-dimensional continuous control tasks.
We demonstrate the effectiveness of representation learning by TACO through extensive experiments on the DeepMind Control Suite (DMC) in both online and offline RL settings. TACO is a flexible plug-and-play module that could be combined with any existing RL algorithm. In the online RL setting, combined with the strong baseline DrQ-v2[52], TACO significantly outperforms the SOTA model-free visual RL algorithms, and it even surpasses the strongest model-based visual RL baselines such as Dreamer-v3 [18]. As shown in Figure 1, across nine challenging visual continuous control tasks from DMC, TACO achieves a  $40\%$  performance boost after one million environment interaction steps on average. For offline RL, TACO can be combined with existing strong offline RL methods to further improve performance. When combined with TD3+BC [11] and CQL [31], TACO outperforms the strongest baselines across offline datasets with varying quality.
We list our contributions as follows:
1. We present TACO, a simple yet effective temporal contrastive learning framework that simultaneously 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.4\mathrm{x}$  on nine challenging tasks in Deepmind Control Suite. Applying TACO to offline RL with SOTA
Figure 1: Comparison of average episode reward across nine challenging tasks in Deepmind Control Suite after one million environment steps.
algorithms also achieves significant performance gain in 4 selected challenging tasks with pre-collected offline datasets of various quality.
2.1 Visual reinforcement learning
Let  $\mathcal{M} = \langle \mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},\gamma \rangle$  be a Markov Decision Process (MDP). Here,  $\mathcal{S}$  is the state space, and  $\mathcal{A}$  is the action space. The state transition kernel is denoted by  $\mathcal{P}:S\times \mathcal{A}\to \Delta (\mathcal{S})$ , where  $\Delta (\mathcal{S})$  is a distribution over the state space.  $\mathcal{R}:S\times \mathcal{A}\rightarrow \mathbb{R}$  is the reward function. The objective of the Reinforcement Learning (RL) algorithm is the identification of an optimal policy  $\pi^{*}:S\to \Delta (\mathcal{A})$  that maximizes the expected value  $\mathbb{E}_{\pi}[\sum_{t = 0}^{\infty}\gamma^{t}r_{t}]$ . Additionally, we can define the optimal Q function as follows:  $Q^{*}(s,a) = E_{\pi^{*}}\left[\sum_{t = 0}^{\infty}\gamma^{t}r(s_{t},a_{t})|s_{0} = s,a_{0} = a\right]$ , such that the relationship between the optimal Q function and optimal policy is  $\pi (s) = \arg \max_{a}Q^{*}(s,a)$ . In the domain of visual RL, high-dimensional image data are given as state observations, so the simultaneous learning of both representation and control policy becomes the main challenge. This challenge is exacerbated when the environment interactions are limited and the reward signal is sparse.
2.2 Contrastive learning and the InfoNCE objective
Contrastive Learning, a representation learning approach, imposes similarity constraints on representations, grouping similar/positive pairs and distancing dissimilar/negative ones within the representation space. Contrastive learning objective is often formulated through InfoNCE loss [47] to maximize the mutual information between representations of positive pairs by training a classifier. In particular, let  $X, Y$  be two random variables. Given an instance  $x \sim p(x)$ , and a corresponding positive sample  $y^{+} \sim p(y|x)$  as well as a collection of  $Y = \{y_{1}, \dots, y_{N - 1}\}$  of  $N - 1$  random samples from the marginal distribution  $p(y)$ , the InfoNCE loss is defined as
$$
\mathcal {L} _ {N} = \mathbb {E} _ {x} \left[ \log \frac {f (x , y ^ {+})}{\sum_ {y \in Y \cup \{y ^ {+} \}} f (x , y)} \right] \tag {1}
$$
Optimizing this loss will result in  $f(x,y) \propto \frac{p(y|x)}{p(y)}$  and one can show that InfoNCE loss upper bounds the mutual information,  $\mathcal{L}_N \geq \log (N) - \mathcal{I}(X,Y)$ .
3 TACO: temporal action-driven contrastive Loss
TACO is a flexible temporal contrastive framework that could be easily combined with any existing RL algorithms by interleaving RL updates with its temporal contrastive loss update. In this section, we first present the overall learning objective and theoretical analysis of TACO. Then we provide the architectural design of TACO in detail.
3.1 Temporal contrastive learning objectives and analysis
In the following, we present the learning objectives of TACO. The guiding principle of our method is to learn state and action representations that capture the essential information about the environment's dynamics sufficient for learning the optimal policy. This allows the agent to develop a concise and expressive understanding of both its current state and the potential effects of its actions, thereby enhancing sample efficiency and generalization capabilities.
Let  $S_{t}$ ,  $A_{t}$  be the state and action variables at timestep  $t$ ,  $Z_{t} = \phi(S_{t})$ ,  $U_{t} = \psi(A_{t})$  be their corresponding representations. Then, our method aims to maximize the mutual information between representations of current states paired with action sequences and representations of the corresponding future states:
$$
\mathbb {J} _ {\mathrm {T A C O}} = \mathcal {I} \left(Z _ {t + K}; [ Z _ {t}, U _ {t}, \dots , U _ {t + K - 1} ]\right) \tag {2}
$$
Here,  $K \geq 1$  is a fixed hyperparameter for the prediction horizon. In practice, we estimate the lower bound of the mutual information by the InfoNCE loss, with details of our practical implementation described in §3.2.
We introduce the following theorem extending from Rakely et al. [41] to demonstrate the sufficiency of TACO objective:
Theorem 3.1. Let  $K \in \mathbb{N}^{+}$ , and  $\mathbb{J}_{TACO} = \mathcal{I}(Z_{t + K};[Z_t,U_t,\dots ,U_{t + K - 1}])$ . If for a given state and action representation  $\phi_Z, \psi_U, \mathbb{J}_{TACO}$  is maximized, then for arbitrary state-action pairs  $(s_1,a_1),(s_2,a_2)$  such that  $\phi (s_1) = \phi (s_2),\psi (a_1) = \psi (a_2)$ , it holds that  $Q^{*}(s_{1},a_{1}) = Q^{*}(s_{2},a_{2})$ .
This theorem guarantees that if our mutual information objective Equation (2) is maximized, then for any two state-action pairs  $(s_1, a_1)$  and  $(s_2, a_2)$  with equivalent state and action representations, their optimal action-value functions,  $Q^*(s_1, a_1)$  and  $Q^*(s_2, a_2)$ , will be equal. In other words, maximizing this mutual information objective ensures that the learned representations are sufficient for making optimal decisions.
3.2 TACO implementation
Figure 2: A demonstration of our temporal contrastive loss: Given a batch of state-action transition triples  $\{(s_t^{(i)},[a_t^{(i)},\dots,a_{t + K - 1}^{(i)}],s_{t + K}^{(i)})\}_{i = 1}^N$ , we first apply the state encoder and action encoder to get latent state-action encodings:  $\{(z_{t}^{(i)},[u_{t}^{(i)},\dots,u_{t + K - 1}^{(i)}],z_{t + K}^{(i)})\}_{i = 1}^N$ . Then we apply two different projection layers to map  $(z_{t}^{(i)},[u_{t}^{(i)},\dots,u_{t + K - 1}^{(i)}])$  and  $z_{t + K}^{(i)}$  into the shared contrastive embedding space. Finally, we learn to predict the correct pairings between  $(z_{t},[u_{t},\dots,u_{t + K - 1}])$  and  $z_{t + K}$  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 a base RL algorithm by incorporating the temporal contrastive loss as an auxiliary loss during the batch update process. Specifically, given a batch of state and action sequence transitions  $\{(s_t^{(i)},[a_t^{(i)},\dots,a_{t' - 1}^{(i)}],s_{t'}^{(i)})\}_{i = 1}^N$ $(t^{\prime} = t + K)$ , we optimize:
$$
\mathcal {J} _ {\mathrm {T A C O}} (\phi , \psi , W, G _ {\theta}, H _ {\theta}) = - \frac {1}{N} \sum_ {i = 1} ^ {N} \log \frac {g _ {t} ^ {(i) ^ {\top}} W h _ {t ^ {\prime}} ^ {(i)}}{\sum_ {j = 1} ^ {N} g _ {t} ^ {(i) ^ {\top}} W h _ {t ^ {\prime}} ^ {(j)}} \tag {3}
$$
Here let  $z_{t}^{(i)} = \phi(s_{t}^{(i)})$ ,  $u_{t}^{(i)} = \psi(a_{t}^{(i)})$  be state and action embeddings respectively.  $g_{t}^{(i)} = G_{\theta}(z_{t}^{(i)}, u_{t}^{(i)}, \dots, u_{t' - 1}^{(i)})$ , and  $h_{t}^{(i)} = H_{\theta}(z_{t'}^{(i)})$ , where  $G_{\theta}$  and  $H_{\theta}$  denote two learnable projection layers that map the latent state  $z_{t}^{(i)}$  as well as latent state and action sequence  $(z_{t}^{(i)}, u_{t}^{(i)}, \dots, u_{t' - 1}^{(i)})$  to a common contrastive embedding space.  $W$  is a learnable parameter providing a similarity measure between  $g_{i}$  and  $h_{j}$  in the shared contrastive embedding space. Subsequently, both state and action representations are fed into the agent's  $Q$  network, allowing the agent to effectively reason about the 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 is the CURL [33] loss:
$$
\mathcal {J} _ {\mathrm {C U R L}} (\phi , \psi , W, H _ {\theta}) = - \frac {1}{N} \sum_ {i = 1} ^ {N} \log \frac {h _ {t} ^ {(i)} ^ {\top} W h _ {t} ^ {(i) ^ {+}}}{h _ {t} ^ {(i) ^ {\top}} W h _ {t} ^ {(i) ^ {+}} + \sum_ {j \neq i} h _ {t} ^ {(i) ^ {\top}} W h _ {t} ^ {(j)}} \tag {4}
$$
Here,  $h_t^{(i)^+} = H_\theta(\phi(s_i^{(t)^+}))$ , where  $s_i^{(t)^+}$  is the augmented view of  $s_i^{(t)}$  by applying the same random shift augmentation as DrQ-v2 [52].  $W$  and  $H_\theta$  share the same weight as the ones in TAC0 objective. The third objective is reward prediction:
$$
\mathcal {J} _ {\text {R e w a r d}} (\phi , \psi , \hat {R} _ {\theta}) = \sum_ {i = 1} ^ {N} \left(\hat {R} _ {\theta} \left(z _ {t} ^ {(i)}, u _ {t} ^ {(i)}, \dots , u _ {t ^ {\prime} - 1} ^ {(i)} \right]\right) - r ^ {(i)}) ^ {2} \tag {5}
$$
Here  $r^{(i)} = \sum_{j = t}^{t' - 1}r_{j}^{(i)}$  is the sum of reward from timestep  $t$  to  $t^{\prime} - 1$ , and  $\hat{R}_{\theta}$  is a reward prediction layer. For our final objective, we combine the three losses together with equal weight. As verified in Section 4.1, although TACO serves as the central objective that drives notable performance improvements, the inclusion of both CURL and reward prediction loss can further improve the algorithm's performance.
We have opted to use DrQ-v2 [52] for the backbone algorithm of TACO, although in principle, TACO could be incorporated into any visual RL algorithms. TACO extends DrQ-v2 with minimal additional hyperparameter tuning. The only additional hyperparameter is the selection of the prediction horizon  $K$ . Throughout our experiments, we have limited our choice of  $K$  to either 1 or 3, depending on the nature of the environment. We refer the readers to Appendix A for a discussion on the choice of  $K$ .
This section provides an overview of our empirical evaluation, conducted in both online and offline RL settings. To evaluate our approach under online RL, we apply TACO to a set of nine challenging visual continuous control tasks from Deepmind Control Suite (DMC) [46]. Meanwhile, for offline RL, we combine TACO with existing offline RL methods and test the performance on four DMC tasks, using three pre-collected datasets that differ in the quality of their data collection policies.
4.1 Comparison between TACO and strong baselines in online RL tasks
Environment Settings: In our online RL experiment, we first evaluate the performance of TACO on nine challenging visual continuous control tasks from Deepmind Control Suite [46]: Quadruped Run, Quadruped Walk, Hopper Hop, Reacher Hard, Walker Run, Acrobot Swingup, Cheetah Run, Finger Turn Hard, and Reach Duplo. These tasks demand the agent to acquire and exhibit complex motor skills and present challenges such as delayed and sparse reward. As a result, these tasks have not been fully mastered by previous visual RL algorithms, and require the agent to learn an effective policy that balances exploration and exploitation while coping with the challenges presented by the tasks.
In addition to Deepmind Control Suite, we also present the results of TACO on additional six challenging robotic manipulation tasks from Meta-world [57]: Hammer, Assembly, Disassemble, Stick Pull, Pick Place Wall, Hand Insert. Unlike the DeepMind Control Suite, which primarily concentrates on locomotion tasks, Meta-world domain provides tasks that involve complex manipulation and interaction tasks. This sets it apart by representing a different set of challenges, emphasizing precision and control in fine motor tasks rather than broader locomotion skills. In Appendix G, we provide a visualization for each Meta-world task.
Baselines: We compare TACO with four model-free visual RL algorithms CURL [33], DrQ [53], DrQ-v2 [52], and A-LIX [5]. A-LIX builds on DrQ-v2 by adding adaptive regularization to the encoder's gradients. While TACO could also extend A-LIX, our reproduction of results from its open-source implementation does not consistently surpass DrQ-v2. As such, we do not choose A-LIX as the backbone algorithm for TACO. Additionally, we compare with two state-of-the-art model-based RL algorithms for visual continuous control, Dreamer-v3 [18] and TDMPC [21], which learn world models in latent space and select actions using either model-predictive control or a learned policy.
TACO achieves a significantly better sample efficiency and performance compared with SOTA visual RL algorithm. The efficacy of TACO is evident from the findings presented in Figure 4 (DMC), Table 1 (DMC), and Figure 5 (Meta-world). In contrast to preceding model-free visual RL algorithms, TACO exhibits considerably improved sample efficiency. For example, on the challenging Reacher Hard task, TACO achieves optimal performance in just 0.75 million environment steps, whereas DrQ-v2 requires approximately 1.5 million steps. When trained with only 1 million environment steps, TACO on average achieves  $40\%$  better performance, and it is even better than the model-based visual RL algorithms Dreamer-v3 and TDMPC on 6 out of 9 tasks. In addition, on more demanding tasks such as Quadruped Run, Hopper Hop, Walker Run, and Cheetah Run, TACO continues to outshine competitors, exhibiting superior overall performance after two or three million steps, as illustrated in Figure 4. For robotic manipulation tasks, as shown in Figure 5, TACO also significantly outperform the baseline model-free visual RL algorithms, highlighting the broad applicability of TACO.
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 1M benchmark. Results are averaged over 6 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.
Environment (1M Steps)	Model-free	Model-based
TACO	DrQv2	A-LIX	DrQ	CURL	Dreamer-v3	TDMPC
Quadruped Run	541 ± 38	407 ± 21	454 ± 42	179 ± 18	181 ± 14	331 ± 42	397 ± 37
Hopper Hop	261 ± 52	189 ± 35	225 ± 13	192 ± 41	152 ± 34	369 ± 21	195 ± 18
Walker Run	637 ± 11	517 ± 43	617 ± 12	451 ± 73	387 ± 24	765 ± 32	600 ± 28
Quadruped Walk	793 ± 8	680 ± 52	560 ± 175	120 ± 17	123 ± 11	353 ± 27	435 ± 16
Cheetah Run	821 ± 48	691 ± 42	676 ± 41	474 ± 32	657 ± 35	728 ± 32	565 ± 61
Finger Turn Hard	632 ± 75	220 ± 21	62 ± 54	91 ± 9	215 ± 17	810 ± 58	400 ± 113
Acrobot Swingup	241 ± 21	128 ± 8	112 ± 23	24 ± 8	5 ± 1	210 ± 12	224 ± 20
Reacher Hard	883 ± 63	572 ± 51	510 ± 16	471 ± 45	400 ± 29	499 ± 51	485 ± 31
Reach Duplo	234 ± 21	206 ± 32	199 ± 14	36 ± 7	8 ± 1	119 ± 30	117 ± 12
Concurrently learning state and action representation is crucial for the success of TACO. To demonstrate the effectiveness of action representation learning in TACO, we evaluate its performance on a subset of 4 difficult benchmark tasks and compare it with a baseline method without action representation, as shown in Figure 3. The empirical results underscore the efficacy of the temporal contrastive learning objectives, even in the absence of action representation. For instance, TACO records an enhancement of  $18\%$  on Quadruped Run and a substantial  $51\%$  on Reacher Hard, while the remaining tasks showcase a performance comparable to DrQ-v2. Furthermore, when comparing against TACO without action representation, TACO achieves a consistent performance gain, ranging from  $12.2\%$  on Quadruped Hopper Hop. These results not only emphasize the inherent value objective in TACO, but also underscore the instrumental role of bolstering the performance of the underlying RL algorithms.
Figure 3: 1M 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
Figure 5: (Meta-world) Performance of TACO against DrQ-v2 and A-LIX. All results are averaged over 6 random seeds.