SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\alpha_{\mathrm{TD3 + BC}} = 2.5$  as defined and used in Scott et al. [11]. In our experiments, no significant performance difference was found for  $\alpha$  within the [2, 3] range. In the case of CQL, we augment the original critic loss with a Q-value regularizer and choose the Q-regularizer weight,  $\alpha_{\mathrm{CQL}}$ , from  $\{0.5, 1, 2, 4, 8\}$ . Table 6 presents the chosen  $\alpha_{\mathrm{CQL}}$  for each dataset.
Table 6: Hyperparameter of  ${\alpha }_{\mathrm{{CQL}}}$  used in different tasks/datasets.
Task/ Dataset	αCQL
Hopper Hop	Medium	0.5
Medium-replay	0.5
Replay	2
Cheetah Run	Medium	0.5
Medium-replay	2
Replay	4
Walker Run	Medium	0.5
Medium-replay	1
Replay	4
Quadruped Run	Medium	0.5
Medium-replay	2
Replay	4
For IQL, we adopt the update functions of both policy and value function from its open-source JAX implementation into DrQ-v2, setting the inverse temperature  $\beta$  to 0.1, and  $\tau = 0.7$  for the expectile. Lastly, for the Decision Transformer (DT), we adapt from the original open-source implementation by Chen et al. [7] and use a context length of 20.
20   
21   
22   
23   
24   
25   
26   
27   
28   
29   
30   
31   
32   
33   
34   
35   
36   
37   
38
D Sensitivity of TACO to action representation dimensionality: an additional experiment
In all of our previous experiments, we use  $\lceil 1.25 \times |\mathcal{A}| \rceil$  as the latent dimensions of the action space for TACO. In practice, we find this works well so that it retains the rich information of the raw actions while being able to group semantically similar actions together. To test the algorithm's sensitivity to this hyperparameter, we conduct an experiment on the Quadruped Run task, which has a 12-dimensional action space. In Figure 10, we show the 1M performance of TACO with different choices of latent action dimensions. Notably, we observe that as long as the dimensionality of the action space is neither too small (6-dimensional), which could limit the ability of latent actions to capture sufficient information from the raw actions, nor too large (24-dimensional), which might introduce excessive control-irrelevant information, the performance of TACO remains robust to the choice of action representation dimensionality.
Figure 10: 1M Performance of TACO with different dimensionality of latent action representations on Quadruped Run  $(|\mathcal{A}| = 12)$ . Error bar represents standard deviation across 8 random seeds.
E Sensitivity of TACO vs. DrQ-v2 to random seeds
(a) Cheetah Run
(d) Quadruped Walk
(g) Finger Turn Hard
Figure 11: Robustness analysis of TACO vs. DrQ-v2 across multiple (6) random seeds
(b) Acrobot Swingup
(e) Reacher Hard
(h) Reach Duplo
(c) Quadruped Run
(f) Hopper Hop
(i) Walker Run
In this section, we conduct a comprehensive analysis of the
algorithm's robustness under random seed, comparing the performance of our proposed method, TACO, with the baseline algorithm, DrQ-v2, across multiple seeds on nine distinct tasks shown in Figure 4. Our aim is to gain insights into the stability and robustness of both algorithms under varying random seeds, providing valuable information to assess their reliability.
Remarkably, Figure 11 demonstrates the impressive robustness of TACO compared to DrQ-v2. TACO consistently demonstrates smaller performance variations across different seeds, outperforming the baseline on almost every task. In contrast, runs of DrQ-v2 frequently encounter broken seeds. These results provide strong evidence that TACO not only improves average performance but also significantly enhances the robustness of the training process, making it more resilient against failure cases. This enhanced robustness is crucial for real-world applications, where stability and consistency are essential for successful deployment.
F Tackling the complex "Manipulator Bring Ball" task: an additional online RL experiment within the DeepMind Control Suite
In addition to the nine tasks depicted in Figure Figure 4, we conduct an additional experiment on an even more complex task "Manipulator Bring Ball". This task necessitates the robotic arm to precisely locate, grab, and transport the ball to a specified goal location. A visualization of a successful trajectory is illustrated in Figure 12a.
The challenge of this domain lies in the sparse reward structure, the demand for accurate arm control, and the need for temporally extended skills. As a result, despite training for 30 million frames, as demonstrated in Figure 12b, DrQ-v2 fails to successfully complete this task. In comparison, TACO manages to successfully solve this task for half of the random initializations, i.e., 3 out of 6 seeds.
(a) Visualization of an expert trajectory on manipulator bring ball.
(b) Online learning curves on manipulator bring ball aggregated over 6 random seeds.
Figure 12: Manipulator Bring Ball Task
G Visualization of Meta-world Tasks
Figure 13: Visualization of expert trajectories for each Meta-world task
H Proof of Theorem 3.1
In this section, we prove Theorem 3.1, extending the results of Rakely et al. [41]. We will use notation  $A_{t:t + K}$  to denote the action sequence  $A_{t},\ldots ,A_{t + K}$  from timestep  $t$  to  $t + K$ , and we will use  $U_{t:t + K}$  to denote the sequence of latent actions  $U_{t},\ldots ,U_{t + K}$  from timestep  $t$  to  $t + K$ .
Figure 14: The graphical diagram
Proposition H.1. Let  $X$  be the return to go  $\sum_{i=0}^{H-t-K} \gamma^i R_{t+K+i}$ , with the conditional independence assumptions implied by the graphical model in Figure 14. If  $I(Z_{t+K}; Z_t, U_{t:t+K-1}) = I(S_{t+K}; S_t, A_{t:t+K-1})$ , then  $I(X; Z_t, U_{t:t+K-1}) = I(X; S_t, A_{t:t+K-1})$ .
Proof. We prove by contradictions. Suppose there exists a pair of state action representation  $\phi_Z, \psi_U$  and a reward function  $r$  such that  $I(Z_{t + K}; Z_t, U_{t:t + K - 1}) = I(S_{t + K}; S_t, A_{t:t + K - 1})$ , but  $I(X; Z_t, U_{t:t + K - 1}) < I(X; S_t, U_{t:t + K - 1}) < I(X; S_t, A_{t:t + K - 1})$ . Then it suffices to show that  $I(S_{t + k}; Z_t, U_{t:t + K - 1}) < I(S_{t + K}; S_t, U_{t:t + K - 1})$ , which would give us the desired contradiction since  $I(Z_{t + k}; Z_t, U_{t:t + K - 1}) \leq I(S_{t + k}; Z_t, U_{t:t + K - 1})$ , and  $I(S_{t + K}; S_t, U_{t:t + K - 1}) \leq I(S_{t + K}; S_t, A_{t:t + K - 1})$ .
Now we look at  $I(X;Z_t,S_t,U_{t:t + K - 1})$ . Apply chain rule of mutual information:
$$
\begin{array}{l} I \left(X; Z _ {t}, S _ {t}, U _ {t: t + K - 1}\right) = I (= X; Z _ {t} | S _ {t}, U _ {t: t + K - 1}) + I (X; S _ {t}, U _ {t: t + K - 1}) (6) \\ = 0 + I \left(X; S _ {t}, U _ {t: t + K - 1}\right) (7) \\ \end{array}
$$
Applying the chain rule in another way, we get
$$
I \left(X; Z _ {t}, S _ {t}, U _ {t: t + K - 1}\right) = I \left(X; S _ {t} \mid Z _ {t}, U _ {t: t + K - 1}\right) + I \left(X; Z _ {t}, U _ {t: t + K - 1}\right) \tag {8}
$$
$$
I \left(X; S _ {t}, U _ {t: t + K - 1}\right) = I \left(X; S _ {t} \mid Z _ {t}, U _ {t: t + K - 1}\right) + I \left(X; Z _ {t}, U _ {t: t + K - 1}\right) \tag {9}
$$
By our assumption that  $I(X;Z_{t},U_{t:t + K - 1}) < I(X;S_{t},U_{t:t + K - 1})$ , we must have
$$
I \left(X; S _ {t} \mid Z _ {t}, U _ {t: t + K - 1}\right) > 0 \tag {10}
$$
Next, we expand  $I(S_{t + K};Z_t,S_t,U_{t:t + K - 1})$
$$
\begin{array}{l} I \left(S _ {t + K}; Z _ {t}, S _ {t}, U _ {t: t + K - 1}\right) = I \left(S _ {t + K}; Z _ {t} \mid S _ {t}, U _ {t: t + K - 1}\right) + I \left(S _ {t + K}; S _ {t}, U _ {t: t + K - 1}\right) (11) \\ = 0 + I \left(S _ {t + K}; S _ {t}, U _ {t: t + K - 1}\right) (12) \\ \end{array}
$$
$$
I \left(S _ {t + K}; Z _ {t}, S _ {t}, U _ {t: t + K - 1}\right) = I \left(S _ {t + K}; S _ {t} \mid Z _ {t}, U _ {t: t + K - 1}\right) + I \left(S _ {t + K}; Z _ {t}, U _ {t: t + K - 1}\right) \tag {13}
$$
$$
I \left(S _ {t + K}; S _ {t} \mid Z _ {t}, U _ {t: t + K - 1}\right) + I \left(S _ {t + K}; Z _ {t}, U _ {t: t + K - 1}\right) = I \left(S _ {t + K}; S _ {t}, U _ {t: t + K - 1}\right) \tag {15}
$$
But then because  $I(S_{t + K}; S_t | Z_t, U_{t:t + K - 1}) > I(X; S_t | Z_t, U_{t:t + K - 1})$  as  $S_t \to S_{t + K} \to X$  forms a Markov chain, it is greater than zero by the Inequality (10). As a result,  $I(S_{t + K}; Z_t, U_{t:t + K - 1}) < I(S_{t + K}; S_t, U_{t:t + K - 1}) < I(S_{t + K}; S_t, A_{t:t + K - 1})$ . This is exactly the contradiction that we would like to show.
Before proving Theorem 3.1, we need to cite another proposition, which is proved as Lemma 2 in Rakely et al. [41]
Proposition H.2. Let  $X, Y, Z$  be random variables. Suppose  $I(Y, Z) = I(Y, X)$  and  $Y \perp Z|X$ , then  $\exists p(Z|X)$  s.t.  $\forall x, p(Y|X = x) = \int p(Y|Z)p(Z|X = x)dz$
Proof of Theorem 3.1: Based on the graphical model, it is clear that
$$
\max  _ {\phi , \psi} I \left(Z _ {t + K}, \left[ Z _ {t}, U _ {t}, \dots , U _ {t + K - 1} \right]\right) = I \left(S _ {t + K}; \left[ S _ {t}, A _ {t}, \dots , A _ {t + K - 1} \right]\right) \tag {16}
$$
Now define the random variable of return-to-go  $\overline{R}_t$  such that
$$
\bar {R} _ {t} = \sum_ {k = 0} ^ {H - t} \gamma^ {k} R _ {t + k} \tag {17}
$$
Based on Proposition H.1, because
$$
I \left(Z _ {t + K}; Z _ {t}, U _ {t: t + K - 1}\right) = I \left(S _ {t + K}; S _ {t}, A _ {t: t + K - 1}\right)
$$
we could conclude that
$$
I \left(\bar {R} _ {t + K}; Z _ {t}, U _ {t: t + K - 1}\right) = I \left(\bar {R} _ {t + K}; S _ {t}, A _ {t: t + K - 1}\right) \tag {18}
$$
Now applying Proposition H.2, we get
$$
\mathbb {E} _ {p \left(z _ {t}, u _ {t: t + K - 1} \mid S _ {t} = s, A _ {t: t + K - 1} = a _ {t: t + K - 1}\right)} \left[ p \left(\bar {R} _ {t} \mid Z _ {t}, U _ {t: t + K - 1}\right) \right] = p \left(\bar {R} _ {t} \mid S _ {t} = s, A _ {t: t + K - 1}\right) \tag {19}
$$
As a result, when  $K = 1$ , for any reward function  $r$ , given a state-action pair  $(s_1, a_1)$ ,  $(s_2, a_2)$  such that  $\phi(s_1) = \phi(s_2)$ ,  $\psi(a_1) = \psi(a_2)$ , we have  $Q_r(s_1, a_1) = \mathbb{E}_{p(\overline{R}_t | S_t = s_1, A_t = a_1)}[\overline{R}_t] = \mathbb{E}_{p(\overline{R}_t | S_t = s_2, A_t = a_2)}[\overline{R}_t]$ . This is because  $p(\overline{R}_t | S_t = s_1, A_t = a_1) = p(\overline{R}_t | S_t = s_2, A_t = a_2)$  by Equation (18) as  $p(z_t | S_t = s_1) = p(z_t | S_t = s_2)$ ,  $p(u_t | A_t = a_1) = p(u_t | A_t = a_2)$ . In case when  $K > 1$ , because if  $\mathbb{E}[Z_{t + K}, [Z_t, U_t, ..., U_{t + K - 1}]] = \mathbb{E}[S_{t + K}, [S_t, A_t, ..., A_{t + K - 1}]]$ , then for any  $1 \leq k \leq K$ ,  $\mathbb{E}[Z_{t + k}, [Z_t, U_t, ..., U_{t + k - 1}]] = \mathbb{E}[S_{t + k}, [S_t, A_t, ..., A_{t + k - 1}]]$ , including  $K = 1$ , by Data processing Inequality. (Intuitively, this implies that if the information about the transition dynamics at a specific step is lost, the mutual information decreases as the timestep progresses, making it impossible to reach its maximum value at horizon  $K$ .) Then the same argument should also apply here.
I Additional related work
I.1 Visual reinforcement learning
In this paper, we focus primarily on visual-control tasks, and this section reviews relevant prior work in visual RL. For visual-control environments, representation learning has shown to be a key. Many prior works show that learning auxiliary tasks can encourage the representation to be better aligned with the task and thus enhance the performance, such as reconstructing pixel observations [55], minimizing bisimulation distances [58], fitting extra value functions [4, 9], learning latent dynamics models [12, 42, 45], multi-step inverse dynamics model [32, 26] or various control-relevant objectives [27]. Model-based methods are also shown to be successful in efficient visual RL, which learn the transition dynamics based on the encoded observation [16, 35, 15, 56]. Data augmentation can also be used to smooth out the learned representation or value functions to improve the learning performance [54, 20, 19].
I.2 Additional works on self-supervised/contrastive learning in reinforcement learning
In §5.1, we summarize the works that apply self-supervised/contrastive learning objectives to improve the sample efficiency of visual reinforcement learning. In this section, we discuss the additional works that apply self-supervised/contrastive learning objectives to a broader set of topics in reinforcement learning.
Several recent works have investigated the use of self-supervised/contrastive learning objectives to pre-train representations for reinforcement learning (RL) agents. Parisi et al. [38] propose PVR,
which leverages a pre-trained visual representation from MoCo [22] as the perception module for downstream policy learning. Xiao et al. [50] introduce VIP, a method that pre-trains visual representations using masked autoencoders. Additionally, Ma et al. [36] propose VIP, which formulates the representation learning problem as offline goal-conditioned RL and derives a self-supervised dual goal-conditioned value-function objective.
Besides pretraining state representations, in goal-conditioned RL, the work by Eysenbach et al. [10] establishes a connection between learning representations with a contrastive loss and learning a value function. Moreover, self-supervised learning has also been employed for effective exploration in RL. Guo et al. [14] introduce BYOL-Explore, a method that leverages the self-supervised BYOL objective [13] to acquire a latent forward dynamics model and state representation. The disagreement in the forward dynamics model is then utilized as an intrinsic reward for exploration. Another approach, ProtoRL by Yarats et al. [51], presents a self-supervised framework for learning a state representation through a clustering-based self-supervised learning objective in the reward-free exploration setting.