SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For each task, we conduct 4 runs with different random seeds and report the mean and standard deviation. The same network architecture is applied for each algorithm to maintain consistency. Similar to [39], the discriminator and policy network share the same visual encoder for all online adversarial methods, and we choose one-step prediction for ELE to align with STG. Through extensive experiments, we answer the following questions:
- Is our proposed framework effective and efficient in visual environments?
- Is our offline pretrained discriminator better than the ones which are trained online?
- Does TDR make a difference to visual representations? And do we need to add 'progress' rewards, as is done in ELE?
4.1 Atari
Atari Datasets. Atari is a well-established benchmark for visual control tasks and also a popular testbed for evaluating the performance of various LfVO algorithms. We conduct experiments on four Atari games: Breakout, Freeway, Qbert, and Space Invaders. To ensure the quality of expert datasets, two approaches are utilized to collect expert observations. For Qbert and SpaceInvaders, we collect the last  $10^{5}$  transitions (around 50 trajectories) from Google Dopamine [43] DQN replay experiences. For Breakout and Freeway, we find the quality of transitions from Dopamine is far from expert. Therefore, we alternatively train a SAC agent [44] from scratch for  $5 \times 10^{6}$  steps and leverage the trained policy to gather approximately 50 trajectories (around  $10^{5}$  transitions) in each game to construct the expert dataset.
Performance in Atari Games. As illustrated in Figure 3, STG is on par with BCO in Freeway but outperforms all baselines in the rest tasks. Table 1 displays the final scores of all methods along with expert datasets and PPO learned from environmental rewards. Remarkably, in Breakout
(a) Breakout
(b) Freeway
(c) Qbert
(d) Space Invaders
Figure 3: STG is compared with BCO, GAIfo, AMP, IDDM, and ELE in four Atari games. The learning curves demonstrate that STG combines the advantage of adversarial learning and the benefit of representation learning, showing substantially better performance in four Atari games.
(a) Pick a flower
Figure 4: Average success rates of STG, ELE, and BCO in Minecraft tasks. STG shows superiority over the two baselines in challenging open-ended tasks.
(b) Milk a cow
(c) Harvest tallgrass
(d) Gather wool
and Qbert, the performance of STG surpasses offline datasets. In Breakout, we observe that STG learns to consecutively bounce the ball up into the top gaps to hit the upper bricks to obtain high intrinsic rewards, while online adversarial methods and ELE fail. In Qbert, STG achieves a prominent breakthrough in the later stages, substantially outperforming other baselines.
Table 1: Mean final scores of last 100 episodes in Atari games. The last two columns display the average episodic scores of expert datasets and PPO with environmental rewards reported in [40].
Environment	GAIfO	AMP	IDDM	ELE	BCO	STG	Expert	PPO
Breakout	1.5	0.6	1.2	22.0	180.4	288.8	212.5	274.8
Freeway	0.6	3.0	10.5	2.7	21.6	21.8	31.9	32.5
Qbert	394.4	874.9	423.3	4698.6	234.1	27234.1	15620.7	14293.3
Space Invaders	260.2	268.1	290.4	384.6	402.2	502.1	1093.9	942.5
Analysis. During the training process, we observe that online adversarial methods tend to easily get stuck in a suboptimal policy and struggle to explore a better policy. This is because the discriminator can easily distinguish between the visual behavior of the expert and the imitator based on relatively insignificant factors within just a few online interactions. In contrast, STG learns better temporally-aligned representations in an offline manner, enabling the discriminator to detect more substantial differences. Besides, instead of relying on probability, STG employs the Wasserstein distance metric to provide more nuanced and extensive reward signals. Consequently, even without fine-tuning during the online RL process, STG can offer valuable guidance to the RL agent. Additionally, in Breakout and Freeway we observe that ELE drops in final performance primarily due to the over-optimistic progress, which will be further investigated in Section 4.3. In comparison, STG ensures precise expert transition prediction and discriminative transition judgment, avoiding over-optimistically driving the agent to transfer to new states.
4.2 Minecraft
Brief Introduction. Built upon Minecraft, Minedojo [45] provides a simulation platform with thousands of diverse open-ended tasks, which propose new challenges for goal completion. Furthermore, the 3D egocentric view and intricate scenes make it extremely hard to extract task-relevant visual signals for LfVO. We evaluate STG on four Minecraft tasks including "pick a flower", "milk a cow",
(e) Pick a flower
(f) Milk a cow
(g) Harvest tallgrass
Figure 5: The removal of the TDR loss from STG, denoted as STG-, induces a decline in performance and sample efficiency in Atari and Minecraft.
(h) Gather wool
"harvest tallgrass", and "gather wool". All tasks are sparse-reward, where only a binary reward is emitted at the end of the episode, thus the performance is measured by success rates.
Minecraft Datasets. Recently, various algorithms, e.g., Plan4MC [16] and CLIP4MC [46] have been proposed for Minecraft tasks. To create expert datasets, for each task, we utilize the learned policies of these two algorithms to collect 100 trajectories (around  $5 \times 10^{4}$  observations).
Performance in Minecraft. The results on Atari show that online adversarial methods can hardly tackle LfVO tasks well. Therefore, in Minecraft, we focus on comparing BCO, ELE, and STG. As depicted in Figure 4, the success rates across four Minecraft tasks reveal a consistent superiority of STG over ELE. In terms of IDM-based BCO, albeit equipped with an additional online-learned inverse dynamic module, it only shows a negligible advantage over STG in "gather wool" and fails to achieve satisfactory performance in the rest three tasks, and even suffers from a severe performance drop in "harvest tallgras".
Notably, we observe that in "milk a cow", STG attains a success rate approaching  $25\%$ , significantly eclipsing the  $5\%$  success rate of ELE or BCO. The reasons for this stark contrast are not yet entirely elucidated. However, a plausible conjecture could be attributed to the task's primary objective, i.e. locating the cow. STG, which is proficient in memorizing high-quality transitions, can effectively accomplish this subgoal, while ELE may lose the intended viewpoint due to its tendency for over-optimistic progress estimations and BCO can easily get stuck in local online sub-optimal transitions.
TDR Ablation. Since STG outperforms other baselines in Atari and Minecraft domain, it raises a question regarding the contribution of the designed temporal distance regressor. To answer this question, we conduct an ablation study, denoted as STG-, where we remove the TDR loss  $\mathcal{L}tdr$  from STG to figure out how the TDR module contributes to enhancing performance and representation. Consequently, the feature encoder  $E_{\xi}$  and the STG Transformer  $T_{\sigma}$  are pretrained under the guidance of a lin
ear combination of  $\mathcal{L}_{mse}$  and  $\mathcal{L}_{adv}$ . We keep the other settings fixed and leverage the pretrained model to tackle the respective downstream tasks. The results are shown in Figure 5, where STG is substantially superior to STG- in most tasks.
Figure 6: T-SNE visualization of a sampled Qbert trajectory embeddings.
In order to figure out the underlying reasons for their discrepancy in performance, we compare the visualization of embeddings encoded by STG and STG-. We randomly select an expert trajectory from Qbert and utilize t-SNE projection to visualize their embedding sequences. As illustrated in Figure 6, the embeddings learned by STG exhibit remarkable continuity in adjacent states, significantly different from the scattered and disjoint embeddings produced by STG-. The superior temporal alignment of the STG representation plays a critical role in capturing latent transition patterns, thereby providing instructive information for downstream RL tasks. For more multi-trajectory visualizations, see Appendix G.
Multi-Task STG. In the previous experiments, the demonstrations used for pretraining come from agents solving the same tasks, it raises a question on how well do STG generalize to a variety of downstream tasks. To this end, we pretrain a new instance of the STG Transformer named STG-Multi, with the same network architecture, on the whole Atari datasets encompassing all four downstream tasks to further assess the efficacy of multi-task adaptation. Considering the four times increase in datasets, we enlarge the model capacity by about four times via stacking more multi-head causal self-attention modules. More details are available in Appendix E. As shown in Figure 7, STG-Multi shows comparable performance across four Atari games. These results strongly reveal the potential of pretraining STG on multi-task datasets for guiding downstream tasks.
(b) Freeway
Figure 7: Multi-task STG (STG-Multi) is pretrained on the whole Atari datasets to guide RL training in downstream tasks.
Pretraining Loss Design. During pre-training, STG concurrently minimizes  $\mathcal{L}_{tot} = \alpha \mathcal{L}_{mse} + \beta \mathcal{L}_{adv} + \kappa \mathcal{L}_{tdr}$ . We set  $\alpha = 0$  to investigate the contribution of  $\mathcal{L}_{mse}$ , denoted as STG(rm MSE), and  $\beta = 0$  to investigate the contribution of WGAN, denoted as STG(rm Adv). For STG(rm Adv), we can only use prediction error as intrinsic rewards like what has been done in [23]. The final performance of STG(rm MSE) and STG(rm Adv) in SpaceInvaders are listed in Figure 8a and 8b. The sample efficiency slightly drops without  $L_{2}$  penalty while the final performance declines heavily without WGAN. Thus, we can conclude each item in  $\mathcal{L}_{tot}$  makes an indispensable contribution to the final performance.
It is worth noting that we arbitrarily choose  $\alpha = \beta = 0.5$  and  $\kappa = 0.1$  in primary experiments according to the value scales of these losses. We discover that these coefficients have a slight influence on training stability and final performance so we do not bother to tune these parameters and leave them consistent across all tasks.
Datasets. Given that the reward guidance thoroughly originates from offline dataset, undoubtedly the capacity and the quality of the pretraining dataset have a non-negligible impact on performance of LfVO tasks. We first conduct extra ablation studies on reduction of dataset size. 50,000 transitions, i.e. half of the original dataset are randomly sampled to train STG(Half-Data) in SpaceInvaders with other settings unchanged. The outcomes of this experiment are presented in Figure 8c. The ablation clearly demonstrates that an increased number of demonstrations contributes to enhanced performance. This observation aligns seamlessly with prior investigations, as illustrated in Figure 4 of [3] and Figure 3 of [19].
Furthermore, we ablate the quality of dataset in Breakout. We manually select 50 trajectories with an episodic return greater than 250 in Breakout to construct expert datasets to train STG(Expert) to solve the Breakout game. Figure 8d shows that higher-quality offline datasets contribute to better sample efficiency. Meanwhile, it also reflects that STG can achieve excellent performance learning from sub-optimal observations.
**Progression Reward.** Considering ELE follows the same scenario of pretraining to enhance downstream tasks as STG, we conduct experiments to figure out whether we can additionally add progress-
(a) Ablate removing  $\mathcal{L}_{mse}$
Figure 8: Learning curves of four pre-training ablations: (a) removing  $\mathcal{L}_{mse}$  in SpaceInvaders; (b) removing  $\mathcal{L}_{adv}$  in SpaceInvaders; (c) using half dataset to train STG in SpaceInvaders; (d) using expert dataset to train STG in Breakout.
(b) Ablate removing  $\mathcal{L}_{adv}$
(c) Ablate dataset size
(d) Ablate dataset quality
sion rewards derived from TDR like ELE to further the performance. To this end, we train the agent under the guidance of both the discriminative and progression rewards from the same pretrained STG Transformer in Atari tasks, denoted as  $\mathbf{STG^{*}}$ . As illustrated in Figure 9, STG outperforms  $\mathrm{STG^{*}}$  in all tasks. We analyze that, similar to ELE, progression rewards from TDR over-optimistically urge the agent to "keep moving" to advance task progress, which however can negatively impact policy learning. For example, on certain conditions such as Breakout or Freeway, maintaining a stationary position may facilitate catching the ball or avoiding collision more easily, thereby yielding higher returns in the long run. Therefore, we do not include the over-optimistic progression rewards in our design.
(c) Qbert
Figure 9: Atari experiments comparing using discriminative rewards (STG) and using both discriminative rewards and progression rewards (STG*).
In summary, our experimental results provide strong evidence for the ability of STG to learn from visual observation, substantially outperforming baselines in a variety of tasks. The ablation study highlights the importance of the TDR module for temporally aligned representations. However, TDR may not be used to generate progression rewards that drive over-optimistic behaviors.
In this paper, we introduce the State-To-Go (STG) Transformer, offline pretrained to predict latent state transitions in an adversarial way, for learning from visual observation to boost downstream reinforcement learning tasks. Our STG, tested across diverse Atari and Minecraft tasks, demonstrates superior robustness, sample efficiency, and performance compared to baseline approaches. We are optimistic that STG offers an effective solution in situations with plentiful video demonstrations, limited environment interactions, and where labeling action is expensive or infeasible.
In future work, STG is likely to benefit from more powerful large-scale vision foundation models to facilitate generalization across a broader range of related tasks. Besides, it can extend to a hierarchical framework where one-step predicted rewards can be employed for training low-level policies and multi-step rewards for the high-level policy, which is expected to improve performance and solve long-horizon tasks.
This work was supported by NSF China under grant 62250068. The authors would like to thank the anonymous reviewers for their valuable comments.
[37] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning (ICML), 2017.
[38] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104, 2023.
[39] Samuel Cohen, Brandon Amos, Marc Peter Deisenroth, Mikael Henaff, Eugene Vinitsky, and Denis Yarats. Imitation learning from pixel observations for continuous control. In Neural Information Processing Systems (NeurIPS) Workshop on Deep Reinforcement Learning, 2021.
[40] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
[41] Xue Bin Peng, Ze Ma, Pieter Abbeel, Sergey Levine, and Angjoo Kanazawa. Amp: Adversarial motion priors for stylized physics-based character control. ACM Transactions on Graphics (ToG), 40(4):1-20, 2021.
[42] Chao Yang, Xiaojian Ma, Wenbing Huang, Fuchun Sun, Huaping Liu, Junzhou Huang, and Chuang Gan. Imitation learning from observations by minimizing inverse dynamics disagreement. In Neural Information Processing Systems (NeurIPS), 2019.
[43] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning (ICML), 2020.
[44] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning (ICML), 2018.
[45] Linxi Fan, Guanzhi Wang, Yunfan Jiang, Ajay Mandlekar, Yuncong Yang, Haoyi Zhu, Andrew Tang, De-An Huang, Yuke Zhu, and Anima Anandkumar. Minedojo: Building open-ended embodied agents with internet-scale knowledge. In Neural Information Processing Systems (NeurIPS) Datasets and Benchmarks Track, 2022.
[46] Ziluo Ding, Hao Luo, Ke Li, Junpeng Yue, Tiejun Huang, and Zongqing Lu. Clip4mc: An rl-friendly vision-language model for apache. arXiv preprint arXiv:2303.10571, 2023.
[47] Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localization. In International Conference on Computer Vision (ICCV), 2017.
[48] Junhyuk Oh, Yijie Guo, Satinder Singh, and Honglak Lee. Self-imitation learning. In International Conference on Machine Learning (ICML), 2018.
We present our algorithm sketches for STG Transformer offline pretraining and online reinforcement learning with intrinsic rewards respectively.
Algorithm 1 STG Transformer Offline Pretraining
Input: STG Transformer  $T_{\sigma}$  , feature encoder  $E_{\xi}$  discriminator  $D_{\omega}$  , expert dataset  $D^{e} =$ $\{\tau^1,\tau^2,\dots ,\tau^m\} ,\tau^i = \{s_1^i,s_2^i,\dots \}$  , buffer  $\mathcal{B}$  , loss weights  $\alpha ,\beta ,\kappa$  . 1: Initialize parametric network  $E_{\xi},T_{\sigma},D_{\omega}$  randomly.   
2: for  $e\gets 0,1,2\dots$  do  $\triangleright$  epoch   
3: Empty buffer  $\mathcal{B}$    
4: for  $b\gets 0,1,2\dots |\mathcal{B}|$  do  $\triangleright$  batchsize   
5: Stochastically sample state sequence  $\tau^i$  from  $D^{e}$    
6: Stochastically sample timestep  $t$  and  $n$  adjacent states  $\{s_t^i,\ldots ,s_{t + n - 1}^i\}$  from  $\tau^i$    
7: Store  $\{s_t^i,\dots ,s_{t + n - 1}^i\}$  in  $\mathcal{B}$    
8: end for   
9: Update  $D_{\omega}$  ..  $\omega \gets$  clip(  $\omega -\epsilon \nabla_{\omega}\mathcal{L}_{dis}, - 0.01,0.01)$    
10: Update  $E_{\xi}$  and  $T_{\sigma}$  concurrently by minimizing total loss  $\alpha \mathcal{L}_{mse} + \beta \mathcal{L}_{adv} + \kappa \mathcal{L}_{tdr}$    
11: end for
Algorithm 2 Online Reinforcement Learning with Intrinsic Rewards
Input: pretrained  $E_{\xi},T_{\sigma},D_{\omega}$  policy  $\pi_{\theta}$  , MDP  $\mathcal{M}$  , intrinsic coefficient  $\eta$  1: Initialize parametric policy  $\pi_{\theta}$  with random  $\theta$  randomly and reset  $\mathcal{M}$  2: while updating  $\pi_{\theta}$  do  $\triangleright$  policy improvement 3: Execute  $\pi_{\theta}$  and store the resulting  $n$  state transitions  $\{(s,s^{\prime})\}_{t}^{t + n}$  4: Use  $E_{\xi}$  to obtain  $n$  real latent transitions  $\{(e,e^{\prime})\}_{t}^{t + n}$  5: Use  $T_{\sigma}$  to obtain  $n$  predicted latent transitions  $\{(e,\hat{e}^{\prime})\}_{t}^{t + n}$  6: Use  $D_{\omega}$  to calculate intrinsic rewards:  $\Delta_t^{t + n} = \{D_\omega (e,\hat{e} ')\}_{t}^{t + n} - \{D_\omega (e,e')\}_{t}^{t + n}$  7: Perform PPO update to improve  $\pi_{\theta}$  with respect to  $r^i = -\eta \Delta$  8: end while
B.1 Atari
We directly adopt the official default setting for Atari games. Please refer to https://www.gymlibrary.dev/environments/atari for more details.
B.2 Minecraft
Table 1 outlines how we set up and initialize the environment for each harvest task.
Table 1: Environment Setup for Harvest Tasks
Harvest Item	Initialized Tool	Biome
milk	empty bucket	plains
wool	shears	plains
tallgrass	shears	plains
sunflower	diamond shovel	sunflower plains
(a) Plains
Figure 1: Biomes in Minecraft
(b) Sunflower Plains
Biomes. Our method is tested in two different biomes: plains and sunflower plains. Both the plains and sunflower plains offer a wider field of view. However, resources and targets are situated further away from the agent, which presents unique challenges. Figure 1a and 1b show the biomes of plains and sunflower plains respectively.
Observation Space. Despite MineDojo offering an extensive observation space, encompassing RGB frames, equipment, inventory, life statistics, damage sources, and more, we exclusively rely on the RGB information as our observation input.
Action Space. In Minecraft, the action space is an 8-dimensional multi-discrete space. Table 2 lists the descriptions of action space in the MineDojo simulation platform. At each step, the agent chooses one movement action (index 0 to 4) and one optional functional action (index 5) with corresponding parameters (index 6 and index 7).
Table 2: Action Space of MineDojo Environment
Index	Descriptions	Num of Actions
0	Forward and backward	3
1	Move left and right	3
2	Jump, sneak, and sprint	4
3	Camera delta pitch	25
4	Camera delta yaw	25
5	Functional actions	8
6	Argument for “craft”	244
7	Argument for “equip”, “place”, and “destroy”	36
C Offline Pretraining Details
Hyperparameters. Table 3 outlines the hyperparameters for offline pretraining in the first stage.
Network Structure. Different architectures for feature encoding are designed for different environments. In Atari, we stack four gray-scale images of shape (84,84) to form a 4-channel state and use the feature encoder architecture as shown in Figure 2a. In Minecraft, a 3-channel image of shape (160,256,3) is directly regarded as a single state, which is processed by a feature encoder with more convolutional layers and residual blocks to capture more complex features in the ever-changing Minecraft world. The detailed structure of the feature encoder for Minecraft is illustrated in Figure 2b. All discriminators, taking in two 512-dimension embeddings from the feature encoder, follow the MLP structure of FC(1024,512)  $\rightarrow$  FC(512,256)  $\rightarrow$  FC(256,128)  $\rightarrow$  FC(128,64)  $\rightarrow$  FC(64,32)  $\rightarrow$  FC(32,1) with spectral normalization.
Representation Visualization. We draw inspiration from Grad-CAM [47] to visualize the saliency map of offline-pretrained feature encoder to assess the effectiveness and advantages of the representation of STG. Specifically, we compare the visualization results of STG and ELE in the Atari environment as illustrated in Figure 3. Each figure presents three rows corresponding to the features captured by the three layers of the convolutional layers, respectively. The saliency maps demonstrate that STG exhibits a particular focus more on local entities and dynamic scenarios and effectively ignores extraneous distractions. As a result, compared with ELE, STG shows greater proficiency in
Table 3: Hyperparameters for Offline Pretraining
Hyperparameter	Value
STG optimizer	AdamW
Discriminator optimizer	RMSprop
LR	1e-4
GPT block size	128
CSA layer	3
CSA head	4
Embedding dimension	512
Batch size	16
MSE coefficient	0.5
Adversarial coefficient	0.3
TDR coefficient	0.1
WGAN clip range	[-0.01,0.01]
Type of GPUs	A100, or Nvidia RTX 4090 Ti
(b) Minecraft
Figure 2: Feature encoder structure for Atari and Minecraft
identifying information strongly correlated with state transitions, thereby generating higher-quality rewards for downstream reinforcement learning tasks.
(a) Saliency maps of different CNN layers in Breakout
(b) Saliency maps of different CNN layers in Freeway
(c) Saliency maps of different CNN layers in Qbert
(d) Saliency maps of different CNN layers in Space Invaders
Figure 3: Saliency maps (SM) of different CNN layers in Atari tasks. The first two columns display the normalized saliency maps and corresponding observations of STG and the last two columns represent SM and corresponding observations of ELE. Through comparison, STG is better at capturing fine-grained features which are strongly correlated with transitions.
D RL Training Details
The general training hyperparameters of PPO for downstream RL tasks in the second stage are listed in Table 4.
Table 4: General Hyperparameters for PPO
Hyperparameter	Value
Optimizer	Adam
Learning rate	2.5e-4
RL discount factor	0.99
Number of workers (CPU)	1
Parallel GPUs	1
Type of GPUs	A100, or Nvidia RTX 4090 Ti
Minecraft image shape	(160,256,3)
Atari stacked state shape	(84,84,4)
Clip ratio	0.1
PPO update frequency	0.1
Entropy coefficient	0.01