SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Bohan Zhou $^{1}$  Ke Li $^{2}$  Jiechuan Jiang $^{1}$  Zongqing Lu $^{1,2\dagger}$
<sup>1</sup> School of Computer Science, Peking University
$^{2}$  Beijing Academy of Artificial Intelligence
Learning from visual observation (LfVO), aiming at recovering policies from only visual observation data, is promising yet a challenging problem. Existing LfVO approaches either only adopt inefficient online learning schemes or require additional task-specific information like goal states, making them not suited for open-ended tasks. To address these issues, we propose a two-stage framework for learning from visual observation. In the first stage, we introduce and pretrain State-to-Go (STG) Transformer offline to predict and differentiate latent transitions of demonstrations. Subsequently, in the second stage, the STG Transformer provides intrinsic rewards for downstream reinforcement learning tasks where an agent learns merely from intrinsic rewards. Empirical results on Atari and Minecraft show that our proposed method outperforms baselines and in some tasks even achieves performance comparable to the policy learned from environmental rewards. These results shed light on the potential of utilizing video-only data to solve difficult visual reinforcement learning tasks rather than relying on complete offline datasets containing states, actions, and rewards. The project's website and code can be found at https://sites.google.com/view/stgtransformer.
Figure 1: A two-stage framework for learning from visual observation. The first stage involves three concurrently pretrained components. A feature encoder is trained in a self-supervised manner to provide easily predicted and temporally aligned representations for stacked-image states. State-to-Go (STG) Transformer is trained in an adversarial way to accurately predict transitions in latent space. A discriminator is updated simultaneously to distinguish state transitions of prediction from expert demonstrations, which provides high-quality intrinsic rewards for downstream online reinforcement learning in the next stage.
† Corresponding author <zongqing.lu@pku.edu.cn>
Reinforcement learning (RL) from scratch imposes significant challenges due to sample inefficiency and hard exploration in environments with sparse rewards. This has led to increased interest in imitation learning (IL). IL agents learn policies by imitating expert demonstrations in a data-driven manner rather than through trial-and-error processes. It has been proven effective in various domains, including games [1] and robotics [2].
However, acquiring demonstrated actions can be expensive or impractical, e.g., from videos that are largely available though, leading to the development of learning from observation (LfO) [3, 4, 5, 6, 7, 8, 9]. This line of research utilizes observation-only data about agent behaviors and state transitions for policy learning. Humans naturally learn from visual observation without requiring explicit action guidance, such as beginners in video games improving their skills by watching skilled players' recordings. However, LfO agents face challenges in extracting useful features from raw visual observations and using them to train a policy due to the lack of explicit action information. Thus, further study of learning from visual observation (LfVO) has the potential to grant agents human-like learning capabilities.
In this paper, we investigate a reinforcement learning setting in which agents learn from visual observation to play challenging video games, such as Atari and Minecraft. Existing LfO approaches like [4, 5, 6, 7, 8] primarily apply to vector-observation environments such as MuJoCo. Others like [3, 9, 10, 11, 12, 13, 14] explicitly consider or can be applied to high-dimensional visual observations. Among them, representation-learning methods [9, 10, 11] learn effective visual representations and recover an intrinsic reward function from them. Adversarial methods [3, 12, 13] learn an expert-agent observation discriminator online to directly indicate visual differences. However, noises or local changes in visual observations may easily cause misclassification [15]. In [12, 13], additional proprioceptive features (e.g., joint angles) are used to train a discriminator, which are unavailable in environments that only provide visual observations. Moreover, as these methods require online training, sample efficiency is much lower compared to offline learning. Goal-oriented methods, like [14], evaluate the proximity of each visual observation to expert demonstrations or predefined goals. However, defining explicit goals is often impractical in open-ended tasks [16]. Furthermore, the continuity of observation sequences in video games cannot be guaranteed due to respawn settings or unanticipated events.
To address these limitations and hence enable RL agents to effectively learn from visual observation, we propose a general two-stage framework that leverages visual observations of expert demonstrations to guide online RL. In the first stage, unlike existing online adversarial methods, we introduce and pretrain State-to-Go (STG) Transformer, a variant of Decision Transformer (DT) [17], for offline predicting transitions in latent space. In the meanwhile, temporally-aligned and predictable visual representations are learned. Together, a discriminator is trained to differentiate expert transitions, generating intrinsic rewards to guide downstream online RL training in the second stage. That said, in the second stage, agents learn merely from generated intrinsic rewards without environmental reward signals. Our empirical evaluation reveals significant improvements in both sample efficiency and overall performance across various video games, demonstrating the effectiveness of the proposed framework.
- We propose a general two-stage framework, providing a novel way to enable agents to effectively learn from visual observation. We introduce State-to-Go Transformer, which is pretrained offline merely on visual observations and then employed to guide online reinforcement learning without environmental rewards.
- We simultaneously learn a discriminator and a temporal distance regressor for temporally-aligned embeddings while predicting latent transitions. We demonstrate that the jointly learned representations lead to enhanced performance in downstream RL tasks.
- Through extensive experiments in Atari and Minecraft, we demonstrate that the proposed method substantially outperforms baselines and even achieves performance comparable to the policies learned from environmental rewards in some games, underscoring the potential of leveraging offline video-only data for reinforcement learning.
Learning from Observation (LfO) is a more challenging setting than imitation learning (IL), in which an agent learns from a set of demonstrated observations to complete a task without the assistance of action or reward guidance. Many existing works [5, 18, 19, 20] attempt to train an inverse dynamic model (IDM) to label observation-only demonstrations with expert actions, enabling behavior cloning. However, these methods often suffer from compounding error [21]. On the other hand, [4] learns a latent policy and a latent forward model, but the latent actions can sometimes be ambiguous and may not correspond accurately with real actions. GAIfo [3], inspired by [22], is an online adversarial framework that couples the process of learning from expert observations with RL training. GAIfo learns a discriminator to evaluate the similarity between online-collected observations and expert demonstrations. Although helpful in mitigating compounding error [3], it shows limited applicability in environments with high-dimensional observations. Follow-up methods [12, 13] pay more attention to visual-observation environments, but require vector-state in expert observations to either learn a feasible policy or proper visual representations. More importantly, learning a discriminator online is less sample-efficient, compared to LfO via offline pretraining. A recent attempt [23] demonstrates some progress in action-free offline pretraining, but return-to-gos are indispensable in addition to observations because of upside down reinforcement learning (UDRL) framework [24]. Moreover, it only shows satisfactory results in vector-observation environments like MuJoCo. In this work, we focus on learning from offline visual observations (LfVO), a more general and practical setting, which is also considered by the concurrent work VIPER [25]. STG and VIPER both offline pretrain a reward function and then use it for online reinforcement learning. STG learns discriminative intrinsic rewards in WGAN style while VIPER leverages logarithm predictive probability as guidance.
Visual Representation Learning in RL. High-quality visual representations are crucial for LfVO. Many previous works [26, 27, 28, 9, 10, 11] have contributed to this in various ways. For example, [26] employs GANs to learn universal representations from different viewpoints, and [27, 28] learn representations via contrastive learning to associate pairs of observations separated by a short time difference. In terms of LfVO, a wide range of methods such as [9, 10, 11] learn temporally continuous representations in a self-supervised manner and utilize temporal distance to assess the progress of demonstrations. Nevertheless, in games like Atari, adjacent image observations may exhibit abrupt or subtle changes due to respawn settings or unanticipated events, not following a gradual change along the timeline. Moreover, over-reliance on temporal information often results in over-optimistic estimates of task progress [14], potentially misleading RL training.
Transformer in RL. Transformer [29] is widely acknowledged as a kind of powerful structure for sequence modeling, which has led to domination in a variety of offline RL tasks. Decision Transformer (DT) [17] and Trajectory Transformer (TT) [30] redefine the offline RL problem as a context-conditioned sequential problem to learn an offline policy directly, following the UDRL framework [24]. DT takes states, actions, and return-to-gos as inputs and autoregressively predicts actions to learn a policy. TT predicts the complete sequence dimension by dimension and uses beam search for planning. MGDT [31] samples from a learned return distribution to avoid manually selecting expert-level returns as DT. ODT [32] extends DT to bridge the gap between offline pretraining and online fine-tuning. MADT [33] extends DT to a multi-agent reinforcement learning setting.
Reinforcement Learning. The RL problem can be formulated as a Markov decision process (MDP) [34], which can be represented by a tuple  $\mathcal{M} = < S, A, \mathcal{P}, \mathcal{R}, \gamma, \rho_0 >$ .  $\mathcal{S}$  denotes the state space and  $\mathcal{A}$  denotes the action space.  $\mathcal{P}: S \times \mathcal{A} \times S \to [0,1)$  is the state transition function and  $\mathcal{R}: S \times \mathcal{A} \to \mathbb{R}$  is the reward function.  $\gamma \in [0,1]$  is the discount factor and  $\rho_0: S \to [0,1]$  represents the initial state distribution. The objective is to find a policy  $\pi(a|s): S \to \mathcal{A}$ , which maximizes the expected discounted return:
$$
J (\pi) = \mathbb {E} _ {\rho_ {0}, a _ {t} \sim \pi (\cdot | s _ {t}), s _ {t} \sim \mathcal {P}} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r (s _ {t}, a _ {t}) \right]. \tag {1}
$$
Transformer. Stacked self-attention layers with residual connections in Transformer is instrumental in processing long-range dependencies, each of which embeds  $n$  tokens  $\{x_i\}_{i=1}^n$  and outputs  $n$  embeddings  $\{z_i\}_{i=1}^n$  of the same dimensions considering the information of the whole sequence. In this study, we utilize the GPT [35] architecture, an extension of the Transformer model, that incorporates a causal self-attention mask to facilitate autoregressive generation. Specifically, each input token  $x_i$  is mapped to a key  $k_i$ , a query  $q_i$ , and a value  $v_i$  through linear transformations, where  $z_i$  is obtained by computing the weighted sum of history values  $v_{1:i}$ , with attention weights determined by the normalized dot product between the query  $q_i$  and history keys  $k_{1:i}$ :
$$
z _ {i} = \sum_ {j = 1} ^ {i} \operatorname {s o f t m a x} \left(\left\{q _ {i} ^ {\mathbf {T}}, k _ {j ^ {\prime}} \right\} _ {j ^ {\prime} = 1} ^ {i}\right) _ {j} \cdot v _ {j}. \tag {2}
$$
The GPT model only attends to the previous tokens in the sequence during training and inference, thereby avoiding the leakage of future information, which is appropriate in state prediction.
Learning from Observation. The goal is to learn a policy from an expert state sequence dataset  $\mathcal{D}^e = \{\tau^1, \tau^2, \dots, \tau^m\}$ ,  $\tau^i = \{s_1^i, s_2^i, \dots, s_n^i\}$ ,  $s_j^i \in S$ . Denote the transition distribution as  $\mu(s, s')$ . The objective of LfO can be formulated as a distribution matching problem, finding a policy that minimizes the  $f$ -divergence between  $\mu^\pi(s, s')$  induced by the agent and  $\mu^e(s, s')$  induced by the expert [7]:
$$
J _ {\mathrm {L f O}} (\pi) = \mathbb {E} _ {\tau^ {i} \sim \mathcal {D} ^ {e}, (s, s ^ {\prime}) \sim \tau^ {i}} \mathrm {D} _ {f} [ \mu^ {\pi} (s, s ^ {\prime}) \| \mu^ {e} (s, s ^ {\prime}) ]. \tag {3}
$$
It is almost impossible to learn a policy directly from the state-only dataset  $\mathcal{D}^e$ . However, our delicately designed framework (see Figure 1) effectively captures transition features in expert demonstrations to provide informative guidance for RL agents, which will be expounded in the following.
3.2 Offline Pretraining Framework
STG Transformer is built upon GPT [35] similar to DT [17], but with a smaller scale and more structural modifications to better handle state sequence prediction tasks. Unlike DT, in our setting, neither the action nor the reward can be accessible, so the STG Transformer primarily focuses on predicting the next state embedding given a sequence of states.
As depicted in Figure 2, first we concatenate a few consecutive image frames in the expert dataset to approximate a single state  $s_t$ . Then, a sequence of  $n$  states  $\{s_t, \dots, s_{t+n-1}\}$  are encoded into a sequence of  $n$  token embeddings  $\{e_t, \dots, e_{t+n-1}\}$  by the feature encoder  $E_{\xi}$  composed of several CNN layers and a single-layer MLP, where  $e_t = E_{\xi}(s_t)$ . A group of learnable positional embedding parameters is added to the token embedding sequence to remember temporal order. These positional-encoded embeddings are then processed by the causal self-attention module which excels in incorporating information about the previous state sequence to better capture temporal
Figure 2: State-to-Go Transformer
dependencies, followed by layer normalization. The linear decoder outputs the final latent prediction sequence  $\{\hat{e}_{t + 1},\dots ,\hat{e}_{t + n}\}$ . Denote the positional encoding, transition predicting, and linear decoding model together as  $T_{\sigma}$ . It is worth noting that instead of predicting the embeddings of the next state sequence directly, we predict the embedding change and combine it with token embeddings in a residual way, which is commonly applied in transition prediction [4] and trajectory forecasting [36] to improve prediction quality.
For brevity, in further discussion we will refer to  $T_{\sigma}(e_t)$  directly as the predicted  $\hat{e}_{t + 1}$ .
Expert Transition Discrimination. Distinguishing expert transiting patterns is the key to leveraging the power of offline expert datasets to improve sample efficiency in downstream RL tasks. Traditional online adversarial methods [3, 12, 13] employ a discriminator to maximize the logarithm probability of transitions sampled from expert datasets while minimizing that from transitions collected online,
which is often sample-inefficient in practice. Moreover, in the case of visual observation, the traditional discriminator may rapidly and strictly differentiate expert transitions from those collected online within a few updates. As a result, the collected observations will be assigned with substantially low scores, making it arduous for policy improvement.
To overcome these limitations, we draw inspiration from WGAN [37] and adopt a more generalized distance metric, known as the Wasserstein distance, to measure the difference between the distributions of expert and online transitions. Compared to the sigmoid probability limited in  $[0,1]$ , the Wasserstein distance provides a wider range and more meaningful measure of the difference between two transition distributions, as it captures the underlying structure rather than simply computing the probability. More importantly, unlike traditional online adversarial methods like GAIfo [3] that use the Jensen-Shannon divergence or Kullback-Leibler divergence, the Wasserstein distance is more robust to the issues of vanishing gradients and mode collapse, making offline pretraining possible. Specifically, two temporally adjacent states  $s_t$ ,  $s_{t+1}$  are sampled from the expert dataset, then we have  $e_t = E_\xi(s_t)$ ,  $e_{t+1} = E_\xi(s_{t+1})$ , and  $\hat{e}_{t+1} = T_\sigma(E_\xi(s_t))$ . The WGAN discriminator  $D_\omega$  aims to maximize the Wasserstein distance between the distribution of expert transition  $(e_t, e_{t+1})$  and the distribution of predicted transition  $(e_t, \hat{e}_{t+1})$ , while the generator tries to minimize it. The objective can be formulated as:
$$
\min  _ {\xi , \sigma} \max  _ {w \in \mathcal {W}} \mathbb {E} _ {\tau^ {i} \sim \mathcal {D} ^ {e}, (s _ {t}, s _ {t + 1}) \sim \tau^ {i}} \left[ D _ {\omega} \left(E _ {\xi} (s _ {t}), E _ {\xi} (s _ {t + 1})\right) - D _ {\omega} \left(E _ {\xi} (s _ {t}), T _ {\sigma} \left(E _ {\xi} (s _ {t})\right)\right) \right]. \tag {4}
$$
$\{D_{\omega}\}_{\omega \in \mathcal{W}}$  represents a parameterized family of functions that are 1-Lipschitz, limiting the variation of the gradient. We clamp the weights to a fixed box  $(\mathcal{W} = [-0.01, 0.01]^l)$  after each gradient update to have parameters  $w$  lie in a compact space. Besides, to suppress the potential pattern collapse, an additional  $L_2$  norm penalizes errors in the predicted transitions, constraining all  $e_t$  and  $\hat{e}_t$  in a consistent representation space. Thus, the loss functions can be rewritten as follows.
For discriminator:
$$
\min  _ {w \in \mathcal {W}} \mathcal {L} _ {d i s} = \mathbb {E} _ {\tau^ {i} \sim \mathcal {D} ^ {e}, (s _ {t}, s _ {t + 1}) \sim \tau^ {i}} \left[ D _ {\omega} \left(E _ {\xi} (s _ {t}), T _ {\sigma} \left(E _ {\xi} (s _ {t})\right)\right) - D _ {\omega} \left(E _ {\xi} (s _ {t}), E _ {\xi} (s _ {t + 1})\right) \right]. \tag {5}
$$
For STG Transformer (generator):
$$
\begin{array}{l} \min  _ {\xi , \sigma} \mathcal {L} _ {a d v} + \mathcal {L} _ {m s e} = - \mathbb {E} _ {\tau^ {i} \sim \mathcal {D} ^ {e}, s _ {t} \sim \tau^ {i}} D _ {\omega} \left(E _ {\xi} \left(s _ {t}\right), T _ {\sigma} \left(E _ {\xi} \left(s _ {t}\right)\right)\right) \\ + \mathbb {E} _ {\tau^ {i} \sim \mathcal {D} ^ {c}, (s _ {t}, s _ {t + 1}) \sim \tau^ {i}} \| T _ {\sigma} \left(E _ {\xi} (s _ {t})\right) - E _ {\xi} (s _ {t + 1}) \| _ {2}. \tag {6} \\ \end{array}
$$
By such an approach, the discriminator can distinguish between expert and non-expert transitions without collecting online negative samples, providing an offline way to generate intrinsic rewards for downstream reinforcement learning tasks.
Temporally-Aligned Representation Learning. Having a high-quality representation is crucial for latent transition prediction. To ensure the embedding is temporally aligned, we devise a self-supervised auxiliary module, named temporal distance regressor (TDR). Since the time span between any two states  $s_i$  and  $s_j$  in a state sequence may vary significantly, inspired by [38], we define symlog temporal distance between two embeddings  $e_i = E_{\xi}(s_i)$  and  $e_j = E_{\xi}(s_j)$ :
$$
t _ {i j} = \operatorname {s i g n} (j - i) \ln (1 + | j - i |). \tag {7}
$$
This bi-symmetric logarithmic distance helps scale the value and accurately capture the fine-grained temporal variation. The TDR module  $P_{\phi}$  consists of MLPs with 1D self-attention for symlog prediction. The objective of TDR is to simply minimize the MSE loss:
$$
\min  _ {\xi , \phi} \mathcal {L} _ {t d r} = \mathbb {E} _ {\tau^ {i} \sim \mathcal {D} ^ {e}, (s _ {i}, s _ {j}) \sim \tau^ {i}} \| P _ {\phi} \left(E _ {\xi} (s _ {i}), E _ {\xi} (s _ {j})\right) - t _ {i j} \| _ {2}. \tag {8}
$$
Offline Pretraining. In our offline pretraining, the transition predictor  $T_{\sigma}$  and transition discriminator  $D_{\omega}$  share the same feature encoder  $E_{\xi}$  similar to online methods [39], which allows them to both operate in an easily-predictable and temporally-continuous representation space.
At each training step, a batch of transitions is randomly sampled from the expert dataset. The model is trained autoregressively to predict the next state embedding without accessing any future information. When backpropagating,  $\mathcal{L}_{mse}$  and  $\mathcal{L}_{adv}$  concurrently update  $E_{\xi}$  and  $T_{\sigma}$  to provide high-quality visual embeddings as well as accurate embedding prediction.  $\mathcal{L}_{tdr}$  is responsible for updating the  $E_{\xi}$  and  $P_{\phi}$  as an auxiliary component, and  $\mathcal{L}_{dis}$  updates  $D_{\omega}$ . Algorithm 1 in Appendix A details the offline pretraining of the STG Transformer.
3.3 Online Reinforcement Learning
Intrinsic Reward. For downstream RL tasks, our idea is to guide the agent to follow the pretrained STG Transformer to match the expert state transition distribution. Unlike [14], our experimental results show that the pretrained WGAN discriminator is robust enough to clearly distinguish between the sub-optimal online transitions and expert transitions from datasets.
Thus, there is no need for fine-tuning and we just use the discrimination score as the intrinsic reward for online RL. Moreover, we do not use 'progress' like what is done in [9]. This is because, in games with multiple restarts, progress signals can easily be inaccurate and hence mislead policy improvement, while the WGAN discriminator mastering the principle of transitions can often make the correct judgment. The intrinsic reward at timestep  $t$  is consequently defined as follows:
$$
r _ {t} ^ {i} = - \left[ D _ {\omega} \left(E _ {\xi} (s _ {t}), T _ {\sigma} \left(E _ {\xi} (s _ {t})\right)\right) - D _ {\omega} \left(E _ {\xi} (s _ {t}), E _ {\xi} (s _ {t + 1})\right) \right]. \tag {9}
$$
A larger  $r_t^i$  means a smaller gap between the current transition and the expert transition.
Online Learning Procedure. Given an image observation sequence collected by an agent, the feature encoder first generates corresponding visual representations, followed by the STG Transformer predicting the latent one-step expert transition. Then the discriminator compares the difference between real transitions and predicted transitions. Their Wasserstein distances, as intrinsic rewards  $r^i$ , is used to calculate generalized advantage, based on which the agent policy  $\pi_{\theta}$  is updated using PPO [40]. It is worth noting that the agent learns the policy merely from intrinsic rewards and environmental rewards are not used.
In this section, we conduct a comprehensive evaluation of our proposed STG on diverse tasks from two environments: a classical Atari environment and an open-ended Minecraft environment. Among the three mainstream methods mentioned in Section 1, goal-oriented methods are not appropriate for comparison because there is no pre-defined target state. Therefore, we choose IDM-based method BCO [19], online adversarial method GAIfo [3] and its follow-ups AMP [41], IDDM [42], and representation-learning method ELE [9].