Learning from Visual Observation via Offline Pretrained State-to-Go Transformer
Authors: Bohan Zhou
Ke Li
Jiechuan Jiang
Zongqing Lu
Venue: NeurIPS 2023 poster (2023)
Keywords: Learning from Observations
Offline Learning from Visual Observations
State-to-Go Transformer
TL;DR: We propose a two-stage framework, from offline pretraining to online reinforcement learning, to effectively and effciently learn from visual observation.

Abstract: 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.

Learning from Visual Observation via Offline Pretrained State-to-Go Transformer
Bohan Zhou' Ke Li
Jiechuan Jiang' Zongqing Lu 1.2f
1 School of Computer Science, Peking University
2 Beijing Academy of Artificial Intelligence
Abstract
Learning from visual observation (LfVO), aiming recovering policies from only visual observation data, is promising yet 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 issucs, we propose two-stage framework for learning from visual observation. In the first stage, introduce and pretrain State- to-Go (STG) Transformer offline to predict and differentiate latent transitions 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 Atari and Minecraft show that our proposed method outperforms baselines and in some tasks even achieves performance comparable 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 https://sites, gowgle.come wwt.e///////eeeeeeeeeeeeeceeeeee wwwm/vvww/sgraaasffoottt
Encoder
STG Transformer
Policy We (a.(s.)
et
ē
St KCR tdr
Input
Temporally-aligner Representation Leurning
Transition Discrimination
Intrinsic
PPO
Transfer
r
8t
e
Observe
Och
BLade
Input
Encoder
Stage 1:Offline Pretraining
Stage 2:Online
Figure A two-stage framework for learning from visual observation. The first stage involves three concurrently pretrained components. feature encoder is trained in self-supervised manner to provide easily predicted and temporally aligned representations for stacked-image states. State-to-Gio (STG) Transformer is trained in an adversarial way to accurately predict transitions in latent space. A discriminator is updated simultancously to distinguish state transitions of prediction from expert demonstrations, which provides high-quality intrinsic rewards for downstream online reinforcement learning in the next stage.
1 Corresponding author <zongqing www.ku......uukkuu.uuu...mmm
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Introduction
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 (II). II agents learn policies by imitating expert demonstrations 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-on -nly data about agent behaviors and state transitions for policy learning. Humans naturally learn from visual observation without requiring explicit action guidance, such 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 policy due to the lack of explicit action information. Thus, further study of learning from visual observation (LIVO) 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 be applied to dinensinnni visual observations. Among them, representation-leamiii methods [9, 10, 11] learn effective visual representations and recover intrinsic reward function from them. Adversarial methods [3, 12, 13] learn 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 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 effectively learn from visual observation, we propose a general two-stage framework that leverages visual observations of expert demonstrations 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-aligee representations are learned. Together, a discriminator is trained to differentiate expert transitions, generating intrinsic rewards 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 both sample efficiency and overall performance across various video games, demonstrating the effectiveness of the proposed framework.
and predictable visual
Our main contributions are as follows:
We propose general two-stage framework, providing novel way enable agents to effectively Iearn 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.

---

2 Related Work
Learning from Observation (LfO) is a more challenging setting than imitation learning (IL), in which an agent learns from a set of demonstrated observations complete 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 latent forward model, but the latent actions can sometimes be ambiguous and may not correspond accurately with real actions. GAIFO [3], inspired by [22], 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], shows limited applicability in environments with high- sigh--imensionll observations. Follow-up methods [12, 13] pay more attention to visual-observation environments, but require vector-state in expert observations to either learn feasible policy or proper visual representations. More importantly, learning a discriminator online less sample-efficient, compared to LfO via offline pretraining. A recent attempt [23] demonstrates some progress in action-free offline pretraining, but return-to- -to-gos are indispensable in addition to observations because of upside down reinforcement lcarning (UDRL) framework [24]. Moreover, it only shows satisfactory results in vector vestrr-observatioon environments like MuloCo. 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 [9, 10, 11] learn temporally continuous representations 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 temporal information often results in over-optimistic estimates of task progress [14], potentially misleading RL training.
Transformer in RL. Transformer [29] is widely acknowledged 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 context- conditioned sequential problem to learn an offline policy directly, following the UDRL framework [24]. DT takes states, actions, and return-to- tetuuntto-goss to-gos inputs and autoregressively predicts actions to learn policy. TT predicts the complete sequence dimension by dimension and uses beam search for planning. MGDT [31] samples from a learned return distribution 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 multi-agent reinforcement learning setting.
kind of powerful structure for
3 Methodology
3.1 Preliminaries
Reinforcement Learning. The RL problem can be formulated as a Markov decision process (MDP)
[34], which can be represented by a tuple M =< S, A, P, R, 7/, Po >. S denotes the state space
and A denotes the action space. P S x A x S > [0, 1) is the state transition function and
R S x A - IR is the reward function. y E [0. is the discount factor and Pa SS - [0. 1] represents the initial state distribution. The objective is to find a policy (aass) S - A. which maximizes the expected discounted return:
( = Ppo.etrm(--aa (-1e))
Σ r'r (ot, at)
(1)
3

---

Transformer. Stacked self- attentien layers with residual connections in Transformer is instrumental
in processing long-range dependencies, each of which embeds tokens {2,) ] and outputs n embeddings (<x)== 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 causal self-attention mask to facilitate autoregressive generation. Specifically, each input token z; is mapped to a key k; a query qi, and a value v, through linear transformations, where is obtained by computing the weighted sum of history values Una with attention weights determined by the normalized dot product between the query qi and history keys kis:
2, = coftmax((1] ky })-=) (0 Uy.
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 policy from an expert state sequence dataset S s) s) E S. Denote the transition distribution as pass s'). The objective of LfO can be formulated as distribution matching problem, finding a policy that minimizes the f-divergence between 1e (s, s') induced by the agent and ("(s, s') induced by the
D = (T T² ..... }. =
expert [7]:
JL.O (Tr)
(","""""",,"'"" Dy [nr (s, I| (s,s")]
almost impossible to learn a policy directly from the state-only dataset D However, our delicately designed framework (see Figure 1) effectively captures transition features in expert demon- strations to provide informative guidance for RI. agents, which will be expounded in the following.
is
3.2 Offline Pretraining Framework
STG Transformer is built upon GPT [35] similar to [17), but with a smaller scale and more struc- tural 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 sequence of states.
T.
Causal Self-Attention Module
emb, + pos.
As depicted in Figure 2, first we concatenate a few consecutive image frames in the expert dataset to ap- proximate single state 8t. Then, a sequence of states ( 8t)+n-1) } are encoded into sequence of n token embeddings {et,.... - 1) } by the fea- ture encoder Eg composed of several CNN layers and a single-layer MLP, where e = Eg(st) (st). A group of learnable positional embedding parameters is added to the token embedding sequence to remember tem- poral order. These positional-encodee embeddings are then processed by the causal self-attention mod- ule which excels in incorporating information about the previous state sequence to better capture temporal dependencies, followed by layer normalization. The linear decoder outputs the final latent prediction sequence {é.+1,..., ).+n). Denote the positional encoding, transition predicting, and linear decod- ing model together T,. 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 residual way, which is commonly applied in transition prediction [4] and trajectory forecasting [36] to improve prediction quality.
TDR Predictor P.
Feature Encoder Eg
.8
8
SC
Figure 2: State-to-Go Transformer
For brevity, in further discussion we will refer to T, (ex) directly the predicted ē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 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 few updates. As 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, measure the difference between the distributions of expert and online transitions. Compared to the sigmoid probability limited in [0, 1J. 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 lensen-Shannon divergence or Bailbback-eebble 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 8t, S1+1 are sampled from the expert dataset, then we have e = Egg (st) c++1 Fgg ($t+1), and ē T, (Eg (st)). The WGAN discriminator Da aims to maximize the Wasserstein distance between the distribution of expert transition (Ee, Gt+1) and the distribution of predicted transition (ex, 1.11), while the generator tries to minimize The objective can be formulated as:
min max (e. DDw (Ee (8t), Ee ($t+1)) Dw (Ee (8t) To (Ee (s))))] Eo WEW
(4)
(Dwwwew represents a parameterized family of functions that are 1-Lipschitz, limiting the variation of the gradient. We clamp the weights fixed box (W [-0.01, 0.01]') after each gradient update to have parameters w lie in a compact space. Besides, to suppress the potential pattern collapse, an additional L2 norm penalizes errors in the predicted transitions, constraining all e, and é, in consistent representation space. Thus, the loss functions can be rewritten as follows.
For discriminator:
min Ldie = (s))))))) (Dw (Eg (st), To (Ee (st))) - Dw (Ee (st), Eg (s(+1)))- (5)
wEW
For STG Transformer (generator):
min Lade + Lisee ErI o Dw (Eg (s:), To (Eg (st)))
[To (Eg (s:)) - Ee (81+1) 1|2-
(6)
By such 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.
Cemporally-Aligned Representation Learning. Having a high-quality representation is crucial for latent transition prediction. To ensure the embedding is temporally aligned, we devise self- supervised auxiliary module, named temporal distance regressor (TDR). Since the time span between any two states s, and sy in a state sequence may vary significantly, inspired by [38], we define symlog temporal distance between two embeddings e; = Eg (s;) and ej = Eg (s)): - i|)-
tij sign(j - i) In(1
This bi-symmetric logarithmic distance helps scale the value and accurately capture the fine- tite-grained temporal variation. The TDR module Ps consists of MLPs with ID self-attention for symlog prediction. The objective of TDR is to simply minimize the MSE loss:
min fade = BrinDDs.e..))> (are))rr E,o
PP (Eg (st), Fg (85)) - 14yll2-
Offline Pretraining. In our offline pretraining, the transition predictor T, and transition discriminator D. share the same feature encoder Ee similar to online methods [39], which allows them to both operate in pasily-predictable and temporally-continuon representation space.
At each training step, 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. Lmse. and Lodv concurrently update Eg and T, to provide high-quality visual embeddings as well as accurate embedding prediction. Lidr is responsible for updating the Eg and Po as an auxiliary component, and Latt updates Da,. Algorithm Appendix A details the offline pretraining of the STG Transformer.
5

---

3.3 Online Reinforcement Learning
Intrinsic Reward. For downstream RL tasks, our idea to guide the agent follow the pretrained STG Transformer to match the expert state transition distribution. Unlike [14], our experimental results show that the pertrained 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 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 consequently defined as follows:
rf - D Do (s,) To (Eg (s))) Du (Eg (s:) Eg (9t+1))
(9)
A larger rf means a smaller gap between the current transition and the expert transition.
Online Learning Procedure. Given image observation sequence collected by 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 rt, is used to calculate generalized advantage, based which the agent policy Ty is updated using PPO [40]. It is worth noting that the agent learns the policy merely from intrinsic rewards and environmental rewards are used.
4 Experiments
In this section, we conduct a comprehensive evaluation of our proposed STG on diverse tasks from two environments: classical Atari environment and 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 ad adversarial method GAIFO [3] and its follow-ups AMP [41], IDDM [42], and representation-leariig method ELE [9].
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 align with STG. Through extensive experiments, we answer the following questions:
Is proposed framework effective and efficient in visual environments?
Is our offline pretrained discriminator better than the ones which are trained online?
Does TDR make difference to visual representations? And do we need to add 'progress' rewards, as is done in ELE?
4.1
Atari
Atari Datasets. Atari 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 10s 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 x 10t steps and leverage the trained policy gather approximately 50 trajectories (around 10y 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 baselines in the rest tasks. Table displays the final scores of methods along with expert datasets and PPO learned from environmental rewards. Remarkably, in Breakout

---

STG
ELE
GAIRO
BCO
AMP
IDDM
3X0X
500
2000
100
1900
200
120
1000
0.00
1.70
0000
0.20
Thansiinn
(a) Breakout
(b) Freeway
(c) Qbert
(d) Space Invaders
Figure 3: STG compared with BCO, GAIRO, 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.
820
0.31
8:2
0.30
0.14
1.2
0.30
0.10
0.5
0.10
810
0.1
810
0.30
110
Epecd
Epedi
Spodi
(a) Pick flower
(b) Milk a cow
(c) Harvest tallgrass
(d) Gather wool
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.
and Qbert, the performance of STG surpasses offline datasets. In Breakout, 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 prominent breakthrough 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
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 explore 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, 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 transfer to new states.
4.2 Minecraft
Brief Introduction. Built upon Minecraft, Minedojo [45] provides a simulation platform with thou- sands of diverse open-ended tasks, which propose new challenges for goal completion. Furthermore, the 3D egocentric view and intricate scenes make extremely hard to extract task-relevant visual signals for LfVO. We evaluate STG on four Minecraft tasks including "pick flower". "milk a cow".

---

2012
23
1A
30000
400
200
2000
100
10000
200
BTG
Tronction
x100
Trunsttin
x1
Thunstion
Tronstton
(a) Breakout
Freeway
(c) Qbert
(d) Space Invaders
0.05
0.20
0.20
0.06
0.13
SIG STG.
0.2
STC
0.02
0.00
0.0
0.00
200
750
1200
750
1020
500
2003
750
100
Epoch
Epoch
Epoch
Epoch
(e) Pick a flower
Milk a cow
(g) Harvest tallgrass
(h) Gather wool
Figure 5: The removal of the TDR loss from STG, denoted STG-, induces decline performance
and sample efficiency in Atari and Minecraft.
"harvest tallgrass", and "gather wool". All tasks are sparse-reward, where only 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, utilize the learned policies of these two algorithms to collect 100 trajectories (around x 10* observations).
Performance in Minecraft. The results 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 consistent superiority of STG over ELE. In terms of inverse dynamic module, it only shows 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".
BCO, albeit equipped with an additional online-lcarned
Notably, we observe that in "milk a cow". STG attains 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 tendency for over- optimistic progress estimations and BCO can easily get stuck in local online sub-optimal transitions.
4.3 Ablations
TDR Ablation. Since STG outperforms other baselines in Atari and Minecraft do- main, it raises question regarding the con- tribution of the designed temporal distance regressor. To answer this question, we con- duct an ablation study, denoted as STG-, where we remove the TDR loss Ltdr from STG to figure out how the TDR module contributes to enhancing performance and representation. Consequently, the feature encoder Eg and the STG Transformer T, are pretrained under the guidance lin- ear combination of L.ne and Ludv. 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 substantially superior to STG- in most tasks.
Figure 6: T-SNE visualization of a sampled Qbert tra- jectory 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-Tase 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 multi-task datasets for guiding downstream tasks.
-30
STVV STA AG
3300
500
1000
250
100-
Tranitiion
Thanstion
(a) Breakout
(b) Freeway
(c) Qbert
(d) Space Invaders
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 Liot aLCmse GLadv + KLrrr. We set a = 0 to investigate the contribution of Lmee. denoted as STG(rm MSE), and B = 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(m Adv) in SpaceInvaders are listed in Figure 8a and 8b. The sample efficiency slightly drops without L2 penalty while the final performance declines heavily without WGAN. Thus, we can conclude each item in Ltot makes an indispensable contribution to the final performance.
It worth noting that we arbitrarily choose a 3 0.5 and 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 tune these parameters and leave them consistent across all tasks.
Datasets. Given that the reward guidance thoroughly originates from offline dataset, undoubtly the capacity and the quality of the pretraining dataset have a non-neglig 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 Data) in SpaceInvaders with other settings unchanged. The outcomes of this experiment are presented in Figure Sc. The ablation clearly demonstrates that an increased number demonstrations contributes to enhanced performance. This observation aligns seamlessly with prior investigations. as illustrated in Figure 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 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 down- stream tasks STG, we conduct experiments to figure out whether we can additionally add progres-

---

600
630
(40)
500
500-
140
300-
500
330
500
200
230
200
0.00
0.20
0.500
1.00
0.00
00.25
0.200
0070
1.00
3.00
18.20
00.25
Thorectoo::
Transition
Tranciioo
Trunsition
(a) Ablate removing Limee (b) Ablate removing Lade
(c) Ablate dataset size
(d) Ablate dataset quality
Figure 8: Learning curves of four pre-training ablations: (a) removing Emse in SpaceInvaders; (b) removing Ladv Spacelnvaders; (c) using half dataset to train STG SpaceInvaders; (d) using expert dataset to train STG in Breakout.
sion rewards derived from TDR like ELE further the performance. To this end, train the agent under the guidance both the discriminative and progression rewards from the same pretrained STG Transformer in Atari tasks, denoted as STG*. As illustrated in Figure 9, STG outperforms STG in all tasks. We analyze that, similar to ELE, progression rewards from TDR over-optimisticalyy urge the agent "keep moving" to advance task progress, which however negatively impact policy learning. For example, on certain conditions such as Breakout or Freeway, maintaining a stationary position may facilitate catching the ball avoiding collision returns in the long run. Therefore, we do not include the over-optimistic progression rewards in our design.
easily, thereby yielding higher
300
STC
STC
STG
STC
SIC
Mally WAl
Li
10
Traxiiio:
x10
(a) Breakout
(b) Freeway
(c) Qbert
(d) Space Invaders
Figure Atari experiments comparing using discriminative rewards (STG) and using both discrimi-
native rewards and progression rewards (STG*).
In summary, our experimental results provide strong evidence for the ability of STG learn from visual observation, substantially outperforming basclines in 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.
Conclusion and Future Work
In this paper, introduce the State-To-Goo (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 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 benefit from powerful large-scale vision foundation models to facilitate generalization across a broader range of related tasks. Besides, can extend 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.
Acknowledgments and Disclosure of Funding
This work was supported by NSF China under grant 62250068. The authors would like thank the anonymous reviewers for their valuable comments.
10

---

References
[1] Jack Harmer, Linus Gisslén, Jorge del Val, Henrik Holst, Joakim Bergdahl, Tom Olsson, Kristoffer Sjöö, and Magnus Nordin. Imitation learning with concurrent actions in 3d games. In IEEE Conference on Computational Intelligence and Games (CIG), 2018.
[2] Yuke Zhu, Ziyu Wang, Josh Merel, Andrei Rusu, Tom Erez, Serkan Cabi, Saran Tunyasuvu- nakool, János Kramár, Raia Hadsell, Nando de Freitas, et al. Reinforcement and imitation learning for diverse visuomotor skills. arXiv preprint arXiv: 2018.
[3] Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observa- tion. arXiv preprint arXiv: axXh::180.066158 2018
[4] Ashley Edwards, Himanshu Sahni, Yannick Schroecker, and Charles Isbell. Imitating latent policies from observation. In International Conference on Machine Learning (ICML), 2019.
[5] Nathan Gavenski, Juarez Monteiro, Roger Granada, Felipe Meneguzzi, and Rodrigo C Barros. Imitating unknown policies via exploration. arXiv preprint arXiv:2008.085 05660, 2020.
[6] Rahul Kidambi, Jonathan Chang, and Wen Sun. Mobile: Model-based imitation learning from observation alone. In Neural Information Processing Systems (NeurIPS), 2021.
[7] Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Off-policy imitation learning from observations. In Neural Information Processing Systems (NeurIPS), 2020.
[8] Tanmay Gangwani, Yuan Zhou, and Jian Peng. Imitation learning from observations under transition model disparity. In Neural Information Processing Systems (NeurIPS) Workshop on Deep Reinforcement Learning, 2021.
[9] Jake Bruce, Ankit Anand, Bogdan Mazoure, and Rob Fergus. Learning about progress from experts. In International Conference on Learning Representations (ICLR), 2023.
[10] Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain. Time-contrastive networks: Self-supervised learning from video. In IEEE International Conference on Robotics and Automation (ICRA), 2018.
[11] Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang, and Nando De Freitas. Playing hard exploration games by watching youtube. In Neural Information Processing Systems (NeurIPS), 2018
[12] Faraz Torabi, Garrett Warnell, and Peter Stone. proprioception. In International Joint Conference on Artificial Intelligence (IJCAI), 2019.
Imitation learning from video by leveraging
[13] Haresh Karnan, Garrett Warnell, Faraz Torabi, and Peter Stone. Adversarial imitation learning from video using state observer. In International Conference on Robotics and Automation (ICRA), 2022.
[14] Youngwoon Lee, Andrew Szot, Shao-Hua Sun, and Joseph Lim. Generalizable imitation learning from observation via inferring goal proximity. In Neural Information Processing Systems (NeurIPS), 2021.
[15] Minghuan Liu, Tairan He, Weinan Zhang, Shuicheng Yan, and Zhongwen Xu. Visual imitation learning with patch rewards. arXiv preprint arXiv:2.30..09965 2023.
[16] Haoqi Yuan, Chi Zhang, Hongeheng Wang, Feiyang Xie, Penglin Cai, Hao Dong, and Zongqing Lu. Plan4me: Skill reinforcement learning and planning for open-world minecraft tasks. arXiv preprint arXiv:2303. 16563, 2023.
[17] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Neural Information Processing Systems (NeurIPS), 2021.
[18] Deepak Pathak, Parsa Mahmoudieh, Guanghao Luo, Pulkit Agrawal, Dian Chen, Yide Shentu, Evan Shelhamer, Jitendra Malik, Alexei Efros, and Trevor Darrell. Zero-shot visual imitation. In IEEE Computer Vision and Pattern Recognition (CVPR) Workshops, 2018.

---

[19] Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In International Joint Conference on Artificial Intelligence (IJCAI), 2018.
[20] Bowen Baker, llge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning by watching unlabeled online videos. In Neural Information Processing Systems (NeurlPS), 2022.
[21] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. structured prediction to no-regret online learning. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2011.
A reduction of imitation learning and
[22] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Neural Informa- lion Processing Systems (NeurIPS), 2016.
[23] Deyao Zhu, Yuhui Wang, Jürgen Schmidhuber, and Mohamed Elhoseiny. Guiding online reinforcement learning with action-free offline pretraining, arXiv preprint arXiv:2301.12876. 2023.
[24] Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards just map them to actions. In Neural Information Processing Systems (NeurIPS). 2019.
[25] Alejandro Escontrela, Ademi Adeniji, Wilson Yan, Ajay Jain, Xue Bin Peng, Ken Goldberg, Youngwoon Lee, Danijar Hafner, and Pieter Abbeel. Video prediction models rewards for reinforcement learning. arXiv preprint arXiv:2305. 2023.
[26] Bradly C Stadie, Pieter Abbeel, and Ilya Sutskever. Third- person imitation learning. arXiv preprint arXiv:
2017.
[27] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised represen- tations for reinforcement learning. In International Conference on Machine Learning (ICML), 2020.
[28] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning (ICML), 2021.
[29] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Neural Information Processing Systems (NeurIPS), 2017.
[30] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Neural Information Processing Systems (NeurIPS), 2021.
[31] Kuang-Huei Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadarrama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. In Neural Information Processing Systems (NeurIPS), 2022.
[32] Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In International Conference on Machine Learning (ICML), 2022.
[33] Muning Wen, Jakub Kuba, Runji Lin, Weinan Zhang, Ying Wen, Jun Wang, and Yaodong Yang. Multi-agent reinforcement learning is sequence modeling problem. In Neural Information Processing Systems (NeurIPS), 2022.
[34] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press. 2018.
[35] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training, 2018.
[36] Tim Salzmann, Boris Ivanovic, Punarjay Chakravarty, and Marco Pavone. Trajectron++: Dynamically-feasible trajectory forecasting with heterogeneous data, In European Conference on Computer Vision (ECCV), 2020.
12

---

[37] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. 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.00 04444, 2023.
[39] Samuel Cohen, Brandon Amos, Mare Peter Deisenroth, Mikael Henaff, Eugene Vinitsky, and Denis Yarats. Imitation learning from pixel observations for continuous control. In Neural Information Processing Systems (NeurIPS) Workshop Deep Reinforcement Learning, 2021.
[40] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford. and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1rXiv:1707.0344, 2017.
[41] Xue Bin Peng, Ze Ma, Pieter Abbeel, Sergey Levine, and Angjoo Kanazawa. Amp: Adversarial motion priors for stylized physics-hased 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 disagree- ment. 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. Clip4me: An rl-friendly vision-language model for minecraft. 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 Interna- tional Conference on Machine Learning (ICMI), 2018
13

---

A
Algorithms
We present our algorithm sketches for STG Transformer offline pretraining and online reinforcement leamning with intrinsic rewards respectively.
Algorithm STG Transformer Offline Pretraining
Input: STG Transformer To, feature encoder Eec, discriminator D.. expert dataset D {si sh.... buffer B, loss weights at 3, K. 1: Initialize parametric network Eg, To, D. randomly.
17,,72 ,,.,,,,,, ""...,,,"""""""
2: for 0, 1, 2.. do
D epoch
3: Empty buffer B.
4: for b 0, do
batchsize
5: Stochastically sample state sequence 7' from D Stochastically sample timestep and adjacent states { 8ttn-1) from r B.
6:
7:
Store .
8:
end for
9: Update Du: + clip(w V 0.01, 0.01).
10: Update Egg and Th concurrently by minimizing total loss omms + BLado + KiLedr-
11:
end for
Algorithm 2 Online Reinforcement Learning with Intrinsic Rewards
Input: pretrained Eg, Ta, Da., policy me, MDP M, intrinsic coefficient 7).
1: Initialize parametric policy Ta with random randomly and reset M.
2: while updating my do
policy improvement
3: Execute Ty and store the resulting n state transitions ((8, (()tt
4: Use Eg to obtain n real latent transitions {(c, c')))
5: Use Tn to obtain n predicted latent transitions {(e, 8')]ttn
6: Use D. to calculate intrinsic rewards: A A (Dw e') (Du(e,e))tt.
7: Perform PPO update to improve To with respect rt -nA.
8: end while
B Environment Details
B.1 Atari
We directly adopt the official default setting for Atari games. Please refer to attps://ww: gymlibrary. dev/ wwveeennnnnttaaaeeeeeee atari for more details.
B.2
Minecraft
Environment Settings
Table 1 outlines how we 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
M
tallgrass
shears
plains
sunflower
diamond shovel
sunflower plains
14

---

(a) Plains
(b) Sunflower Plains
Figure Biomes in Minecraft
Biomes. Our method is tested in two different biomes: plains and sunflower plains. Both the plains and sunflower plains offer wider field of view. However, resources and targets situated further away from the agent, which presents unique challenges. Figure la and 1 show the biomes of plains and sunflower plains respectively.
Observation Space. Despite MineDojo offering an extensive observation space, encompassing RGB frames, equipment, inventory, life statisties, 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- 8-dimensional multi-discrete space. Table lists the descriptions action space in the MineDojo simulation platform. each step, the agent chooses one movement action (index 0 to 4) and one optional functional action (index 5) with corresponding parameters (index and index 7).
Table 2: Action Space of MineDojo Environment
Index
Descriptions
Num of Actions
0
Forward and backward
Move left and right
2
Jump, sneak, and sprint
4
3
Camera delta pitch
25
Camera delta yaw Functional actions
25
5
Argument for "craft" Argument for "equip", "place", and "destroy"
244
7
36
C Offline Pretraining Details
Hyperparameters. Table outlines the hyperparameters for offline pretraining in the first stage.
Network Structure. Different architectures for feature encoding are designed for different environ- ments. 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 single state, which is processed by a feature encoder with more convolutional layers and residual blocks capture more complex features 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 with spectral normalization.
Representation Visualization. We draw inspiration from Grad-CAM [47] to visualize the saliency map diftine-prerraiee feature encoder assess the effectiveness and advantages of the repre- sentation of STG. Specifically, we compare the visualization results of STG and ELE in the Atari environment as illustrated Figure 3. Each figure presents three rows corresponding to features captured by the three layers of the convolutional layers, respectively. The saliency maps demonstrate that STG exhibits particular focus more on local entities and dynamic scenarios and effectively ignores extraneous distractions. As result, compared with ELE, STG shows greater proficiency in
15

---

Table 3: Hyperparameters for Offline Pretraining
Hyperparameter
Value
STG optimizer
Adam W
Discriminator optimizer
RMSprop le-4
LR
GPT block size
128
CSA layer
CSA head
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
A 1000, or Nvidia RTX 4090 Ti
(1,64,32,20)
(1,512)
LeakyRelu
Conv:3 Stride:2
ResBlock
BatchNorm
ResBlock
Flatten & Lincur
BatchNorm
Stride:2 Conv:3x3
LeakyRelu
(1,64,7,7)
(1,64, 64, 40)
(1,64,16 16, 601
Relu
Stride:1
LeakyRelu
Conv:3 Stide:2
Conv:3x3
ResBlock
BatchNorm
ResBlock
1 (1,64.9.9)
BatchNorm
Stride:2 Conv:3x3
LeakyRelu
Relu
(1,32,128,80)
(1,64,8,5)
Stride:2
Conv(mouut3.1
Conv:4x4
BurchNonn
1 (1,32,20,20)
LeakyRelu
Conv:3x3 Stride:2
Convtin,utt
BotekNocu
ResBlock
BatchNorm
Conviout.out3.1) ButchNoom
BatchNorm
Relu
Stride:2
Conv:3x3
LeakyRelu
Stride:4
Conv:8*8
Flatten & Linear
1
(1,3,256,160)
(1,128,4,3)
(1,512)
(1,4,4.84)
(a) Atari
(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.
16

---

20
90
91
1914
181
180
191
150
123
100
135
1x -3
55
55
31
34
00
131
170
133
150
02
42
210
6
333
0.0
210
210
46
135
0.0
210
-014
-31
00-
00
01
0.
131
190
150
190
130
Saliency maps of different CNN layers in Breakout
21
18-
131
120
131
120
100
200
65
2100
200
40
125
2011
-20
01
01
150
100
150
190
1.0
0-
00
00
181
190
190
190
200
0
80
333
190
0.0
210
1.00
350
210
40
135
700
210
3.5
500
(b) Saliency maps of different CNN layers in Freeway
17

---

30
96
46
0.4
190
130
090
130
100
1.0
2.01
00
143
131
131
191
190
091
190
200
123
100
0.0
200
200
125
0.0
200
00-
00
00
0.
120
131
0.4
120
131
100
101
100
210
210
(c) Saliency maps of different CNN layers in Qbert
BaGn
00
120
130
190
150
190
190
200
133
190
210
1.00
190
200
40
60
130
100
210
130
191
0
22
2320
0000
10
606
18
121
120
131
0.4
120
190
150
190
150
02
201
2001
0 3320 2
0320
0500
20
0.5
0.
01
91
30
150
150
181
100
180
190
6
123
(90
46
30
125
.00
(d) Saliency maps different CNN layers 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.
18

---

D RL Training Details
The general training hyperparameters of PPO for downstream RL tasks in the second stage are listed
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)
I
Parallel GPUs
Type of GPUs
A100, or Nvidia RTX 4090 Ti (160,256,3)
Minecraft image shape
Atari stacked state shape
(84,84,4)
Clip ratio
0.1
PPO update frequency Entropy coefficient
0.1
0.01
Neither the discriminative reward from STG nor the progression reward from ELE bounded. Therefore, it is reasonable adjust certain hyperparameters to bring the best performance of each algorithm in each task. In Table 5, the coefficient of intrinsic reward r7(7) >> 0) for different baselines tuned balance the value scale and GAE A(O < À < 1) is tuned adjust the impact of intrinsic rewards in different tasks.
Table 5: Specific Hyperparameters for Different Tasks
Task
1SS
HELE
GGAIFO
AGAE
Breakout
0.6
1.0
2.0
0.1
Freeway
2.0
0.1
1.0
0.15
Qbert
5.0
0.05
2.0
0.95
Space Invaders
6.0
0.1
2.0
0.95
Milk a Cow
1.0
0.5
0.8
Gather Wool
10.0
0.1
0.8
Harvest Tallgrass
1.0
0.1
0.95
Pick a Flower
1.0
0.1
0.95
The coefficients of STG* (noted as yp vr*) in four Atari tasks are reported in Table 6.
Table 6: Coefficients for STG* in Atari Tasks
Task
Breakout
0.6
0.01
Freeway
2.0
0.1
Qbert
5.0
0.03
Space Invaders
6.0
0.01
Training Details. For Minecraft tasks, we adopt hybrid approach utilizing both PPO [40] and Self- Imitation Learning (SIL) [48]. Specifically, we store trajectories with high intrinsic rewards in buffer and alternate between PPO and SIL gradient steps during the training process. This approach allows us to leverage the unique strengths of both methods and achieve superior performance compared to utilizing either method alone [46].
19

---

E Further Experiments
Intrinsic Reward Design. In Equation (9), we define our intrinsic reward r¹ the difference between y.uide and pbase:
(Eg (s1+1)) D. (Eg (st (Eg (s:))) rugide
base
(10)
On the one hand, pretrained D. clearly provides informative judgment y-uide of transition quality during online interaction. On the other hand, the baseline reward ylase solely relying on the current state 8t, serves as a baseline to normalize r' to relatively lower level. In this section, we aim to investigate the necessity of incorporating base.
To assess the significance of base, we conducted experiments on the four Atari tasks utilizing only prude the intrinsic reward, which is similar to previous adversarial methods like GAIfO [3]. In order to bring the scale of guuide in line with r', we employ running normalization and bound the values within the range of [-1,]] mitigate the negative influence of outliers. All other settings remain unchanged. We denote this ablation baseline as STG'.
As illustrated in Figure 4, ysuide yields comparable final performance in Breakout and Space Invaders while failing to achieve satisfactory performance in Freeway and Qbert. contrast, by leveraging y base, which provides unique reference from expert datasets for each individual 8g, observe reduced variance and improved numerical stability compared to the running normalization trick that calculates batch means for normalization.
300
A
30000
200
20000
yany
100
10000
200
My
0.50 Transtion
0.25
01.50 Transition
Transition
le7
Trunsition
lo6
le7
107
(a) Breakout
(b) Freeway
(c) Qbert
(d) Space Invaders
Figure 4: Atari experiments comparing using pride (STG') and rd (STG) as intrinsic reward.
Additional Ablations
Given the four times larger datasets, we enlarge the size of the STG Transformer by increasing the number of heads (24) and layers (16) within the multi-head causal self-attention modules, augmenting the model capacity by about four times. Here further assess the efficacy of multi-task adaptation of the STG Transformer Minecraft. We tune the intrinsic coefficient to be 5 for Breakout, Qbert, Spacelnvaders, "milk cow", "harvest tallgras". "pick flower", and 10 for Freeway and "gather wool". STG-Multi results in Minecraft are illustrated in Figure 5. As all four tasks share similar biome in MineDojo, STG-Multi may provide less clear guidance for the agent than STG in downstream tasks, which may cause their discrepancy in performance.
Multi-Task STG.
0170
1.50
6.0
0.12
0.20
0.10
Malis
Epodi
Exoch
Ecoch
E2ooch
(a) Pick flower
(b) Milk a cow
(c) Harvest tallgrass
(d) Gather wool
Figure 5: Multi-task STG (STG-Multi) is pretrained on the whole Minecraft datasets to guide RL
training in downstream tasks..
20

---

G Additional Visualizations
Intrinsic Return. As no environmental rewards participate in updating policy, the ultimate objective of LfVO is to maximize the expectation of cumulative dense intrinsic rewards, namely intrinsic return. Figure shows the learning curves of smoothed intrinsic return in Atari and Minecraft. The rising trend proves that online collected observation distribution is getting closer to expert observation distribution during training, indicating effectiveness of the offiine-peeraiieed Modules.
20
1100
500
1225
1030
41.5
1001
0035
80
236
$2.0
1030
120
(a) Breakout
(b) Freeway
(c) Qbert
(d) Space Invaders
Figure 6: The ultimate objective of LfVO is to maximize the expectation of cumulative dense intrinsic rewards, namely intrinsic return. The rising trend proves that online collected observation distribution is getting closer to expert observation distribution.
Multi- Taati-Traaectories Visualization. As extension to increase diversity, additionally visualize more trajectories SpaceInvaders Figure 7. We randomly sample five trajectories in the expert dataset and use t-SNE visualize the embedding sequences encoded by STG and STG- respectively. Expert trajectories of STG exhibit more continuity in adjacent states compared with STG- This consistent with the visualization results in Figure 6. Furthermore, can be observed that after 10M-step of training, the observation distribution is getting closer to expert, and the patterns of STG closer exert in comparison with STG-. This reflects that WGAN-style training indeed generates meaningful reward signal to imitate the hehavior experts and TDR module accelerate the process.
STG
J
(a) random trajectory
training kk steps
(c) expert datasets
STG-
random trajectory
(b) training 1w steps
(c) expert datasets
Figure 7: T-SNE visualization of embeddings of random trajectory, a rollout trajectory and expert trajectories in SpaceInvaders.
21


Reviews:
[{"review_id": "OIaDJKykig", "reviewer": "Reviewer_9y8A", "summary": "Paper presents a novel method for learning from visual observations. Specifically, a state encoder is first trained in an offline fashion with two tasks: 1) predicting the normalized relative distance between two states; 2) predicting the relative difference between two consecutive states. Then the whole agent is trained using a standard WGAN-based IRL from observation loss. Experiments on Atari and Minecraft shows some promises against baselines and ablations.", "questions": "See [Weaknesses]", "limitations": "I am wondering whether the proposed representations losses have to be tied with the transformer architecture, i.e. will they still work with a RNN-style recurrence?", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "+The paper is overall clear and well-written. The authors did a good job illustrating their main idea with stylish graphics and visualization, which is friendly to non-expert readers.\n\n+The method itself is technically sound with good motivations. Although some design choices are not properly backed by ablations, I think the authors explain the rationales under their losses and model architecture well, mostly relevant to the challenges, ex. distinguishing states.\n\n+On the selected arenas, the method indeed show superior performances against the counterparts and the gap is quite significant, suggesting the effectiveness. The TSNE viz is also quite helpful with better understanding on how the method work.\n\n", "weaknesses": "Having said those above, I do have some major concerns regarding some claims and mostly experiments. I hope the authors can help clarify them in a rebuttal:\n\n-In sec. 3.3, the text reads “That is, the WGAN discriminator can clearly distinguish between the state sequences collected under the learning policy and the expert state sequences, without fine-tuning”. This could be a very strong argument on the proposed representation learning method and the authors seem to derive this from the comparison to GAIfO (w/o the proposed representation learning). But this can be indirect evidence as there might be other factors contributing to the differences, ex. hyper-parameters, additional models, etc. The authors are suggested to include a trajectory visualization that shows if the collected trajectory indeed becomes more distinguishable to expert demonstrations with the proposed representation learning method. \n\n-Baselines; Although I can understand the authors are focusing on IRL-based LfO approach, but indeed I think it is still necessary to include some baselines based off behavior cloning, ex. BCO ([19] in the original paper) and some follow-ups. It would be exciting to see if the proposed representation learning pretraining could help with these BC-based approach as well. Moreover, some additional IRLfO baseline, ex. IDDM ([2] below) and some more recent work are also worth comparing.\n\n-Ablations; The authors only include an ablation that removes the TDR loss from the complete pretraining objective. However, there are many more components and design choices that should be justified by ablations. To name a few:\n- transition discrimination losses, L_{mse} and L_{adv}\n- coefficient of these losses\n\n-Is It possible to use multiple trajectories for the TSNE visualization in Fig. 6? \n\n-Also Fig. 6, I am a bit confused why the state can be clustered with TDR loss. I will expect a task of predicting the distance between the current state and the goal state (ex. as in [1]) to encourage the emergence of clustering based off the distance-to-go. However, the TDR task is about pair-wise distance between two arbitrary states. Can the authors explain why such loss could facilitate the clusters we observed in Fig. 6?\n\n-citations; some relevant papers on learning Minecraft controllers and LfO in general should be cited: [1-2]\n\n[1] https://arxiv.org/abs/2301.10034 \n\n[2] https://arxiv.org/abs/1910.04417  \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for detailed comments and insightful review. We are encouraged that the reviewer finds \"the idea technically sounds well-motivated in distinguishing states\" and \"our rationales are clear\". We address the questions and concerns below.\n\n`Q1.`\"Include some baselines based off behavior cloning and IRLfO baselines.\"\n\n`A1.` We have added baselines including IDM-based method **BCO**, IRLfO method **IDDM** and \"advanced-GAIfO\" **AMP** in Atari. The learning curve of STG outperforms all the five baselines in Breakout, Qbert and Space Invaders. In Freeway, the final performance of STG is comparable to BCO and superior to other baselines.\n\n`Q2.` \"loss design and coefficients should be justified by ablations.\"\n\n`A2.` STG concurrently minimize $\\mathcal{L}\\_{tot}=\\alpha\\mathcal{L}\\_{mse}+\\beta\\mathcal{L}\\_{adv}+\\kappa\\mathcal{L}\\_{tdr}$ during pre-training. \n\nAs shown in **Figure 2(a) and 2(b)** in the attached PDF, lack of $\\mathcal{L}\\_{adv}$ will lead to obvious degradation in performance while lack of $\\mathcal{L}\\_{mse}$ will lead to slight degradation. Along with TDR ablation, we can find that $\\mathcal{L}\\_{mse}$ and $\\mathcal{L}\\_{tdr}$ synergistically contribute to temporally-aligned representation and high-quality sequence prediction, which helps WGAN learn to provide accurate transition discrimination. This is critical for offline learning from visual observations.\n\nAbout loss coefficients, according to the value scales of these losses, we arbitrarily choose $\\alpha=\\beta=0.5$ and $\\kappa=0.1$ . 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.\n\n`Q3.` About visualization: \"Include a trajectory visualization that shows if the collected trajectory indeed becomes more distinguishable to expert demonstrations with the proposed representation learning method\"; \"use multiple trajectories\"; \"why the visualized embeddings can be clustered with TDR loss?\"\n\n`A3.` **Figure 5** in the attached PDF displays more visualizations between STG and STG- without TDR module. To increase diversity, we visualize the trajectories in **SpaceInvaders** rather than Qbert in the paper. We randomly sample **5** trajectories in the expert dataset and use t-SNE to visualize the embedding sequences encoded by STG and STG- respectively. Expert trajectories of STG exhibit more continuity in adjacent states compared with STG-. This is consistent with the visualization results in Figure 6 of the paper. From Figure 5(a) to Figure 5(b) in the attached PDF, after 10M-step-training the observation distribution is getting closer to expert, and the patterns of STG are closer to exert in comparison with STG-. This reflects that our WGAN-style training indeed generates meaningful reward signal to imitate the behavior of experts and TDR module accelerate the process. \n\nAs for the reason for clustering, the phenomena both occur in Qbert and SpaceInvaders. It is probably related to Atari environment, because in a single Atari task, the distinctive observation patterns are actually countable. With effective representation learning techniques, near-expert observations are likely to cluster into some correct \"behavior patterns\" for solving a task. This discovery is consistent with previous work [3]. It also describes the task of reaching a key to a door in Montezuma Revenge can be roughly clustered into four patterns with a time-based representation learning skill easily.\n\n`Q4.` \"citing some relevant papers.\"\n\n`A4.` We thank the reviewer for the valuable suggestions. We will include the references[1,2] in our final version.\n\n`Q5.` \"Do the proposed representations losses still work with a RNN-style recurrence?\"\n\n`A5.` The design of representation learning is orthogonal to our proposed offline pretraining adversarial method.  It will be expected to achieve satisfactory performance no matter applied to RNN-style recurrence or expressive transformer structure.\n\nThanks again for the review! We will implement the feedback in the next version of this paper. Further comments are welcome!\n\n**Reference**\n\n[1] Shaofei Cai, et al. \"Open-World Multi-Task Control Through Goal-Aware Representation Learning and Adaptive Horizon Prediction.\" CVPR 2023.\n\n[2] Yang, Chao, et al. \"Imitation learning from observations by minimizing inverse dynamics disagreement.\" NeurIPS 2019.\n\n[3] Aytar, Yusuf, et al. \"Playing hard exploration games by watching youtube.\" NeurIPS 2018."}, {"review_id": "9Hlct57SEP", "reviewer": "Reviewer_HW2S", "summary": "This paper presents a new method, called state-to-go (STG) transformer, for learning from visual observation (LfVO), which aims to learn strong policies given the access to (i) offline datasets without actions and rewards and (ii) following online interaction. The main idea of STG transformer is to learn representations by predicting the temporal diference and learn adversarial discriminator that discriminates the embeddings predicted from the model and expert states within the dataset. Then the model is trained to achieve the target tasks by learning to optimize the intrinsic rewards given by the discriminator, which is basically trained to follow the expert behaviors within the offline dataset.", "questions": "- Please address the points described in Weaknesses section.\n- Predicted embeddings are used as a 'negative' samples in Equation (5) but it's then used as a 'expert' samples in Equation (9). Is this intended, and if so, how could this be justified?", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Strengths\n- 1. Pre-training the discriminator from offline datasets instead of learning the discriminator using the online samples is a novel and interesting idea, to my best knowledge.\n- 2. Consistent improvement over the baselines on Atari and Minecraft domains.\n- 3. Main description of the method is clear.\n\n\n", "weaknesses": "Weaknesses\n- 1. This paper is missing crucial analysis that investigates the importance of learning the discriminator using the offline samples. Experiments on Figure 5 (STG, and STG- that does not use temporal representation learning) actually shows that removing representation learning makes performance be very similar to baselines that do not depend on pre-training in several tasks. This makes it questionable whether the proposed offline pre-training component of learning the discriminator is indeed important for the performance and make be think that pre-training good representation learning from offline datasets might be the most important factor. Including more investigations, e.g., experiments with pre-trained encoders for all the methods, could might support that the proposed discriminator training scheme and the intrinsic reward from this are indeed important.\n- 2. If the paper wants to claim the newly proposed visual representation learning is one of main contributions, I think proper comparison to prior generic representation learning methods (contrastive learning, masked autoencoding, time contrastive networks, ..) is crucial to claim such a contribution. It makes sense that the proposed method works well but its effectiveness and novelty are not well verified in current status so to be deserved as one of main contributions of the paper, especially considering that using the time index for representation learning is also investigated in prior work ([Cai et al., 2023], [Yun et al., 2022])\n- 3. The paper is misusing the term 'State-to-Go' which is not clear what exactly this means. Decision Transformer uses return-to-go as their inputs because their main idea to learn a return-conditioned policy that enables it to stitch the offline samples. But there is no definition of state-to-go in the paper and it is actually not used as inputs or conditioning the policy in any part of the paper. This makes it very difficult to understand the main message from the paper.\n- 4. The paper often positions itself as a paper that addresses the difficulty of learning from visual observations on video game domains, while prior work only considers the robotics domains (e.g., line 97). But what exactly are the main factors that make video game domains is more challenging and which experimental evidence supports this? Isn't this claim actually narrowing down the scope of the paper? Also, this should be clearly stated in the abstract or the title that the paper mainly wants to address the main challenges for training offline agents in video game domains.\n- 5. This is a minor point, but the paper assumes some background knowledge of prior methods by passing too fastly over the baselines. For instance, the authors might not be familar with ELE and the exact meaning of progression rewards, but the paper assumes such knowledge and does not give a formal definition of progression reward, thus making it difficult to understand the main message from Figure 7 and corresponding subsection.\n\n[Cai et al., 2023] \"Open-world multi-task control through goal-aware representation learning and adaptive horizon prediction.\" CVPR 2023\n\n[Yun et al., 2022] Yun, Sukmin, et al. \"Time is matter: Temporal self-supervision for video transformers.\" ICML 2022", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for detailed comments and valuable feedback, and for pointing out the novelty of our pre-training scheme. We provide a point-by-point response below.\n\n`Q1.`\"Include more investigations, e.g., experiments with pre-trained encoders for all the methods.\" and \"Investigate the importance of learning the discriminator using the offline samples.\"\n\n`A1.` We would like to clarify that in fact we have taken this point into careful consideration. GAIfO online updates discriminator without pre-training. ELE just pre-trains a temporally-aligned encoder with an auxiliary temporal distance regression task. STG, however, distinguishes itself from GAIfO by incorporating representation learning tricks into offline pre-training paradigm and distinguishes itself from ELE with WGAN-style adversarial learning framework. These prominent differences make STG achieve outstanding performance in offline learning from visual observations, and the performance gap between STG and ELE highlights the importance of learning a WGAN discriminator. Besides, in **Figure 2(b)** in the attached PDF, the signal of transition discrimination from WGAN framework is much stronger than L2 prediction error, which fully proves the significance of adversarial learning.\n\n`Q2.`\"comparison to prior generic representation learning methods.\" \n\n`A2.` The main contribution of the paper is to offer insight to offline learning from visual observations instead of online learning a discriminator like most previous algorithms. To make it work better, we find temporal representation learning is indispensable for the discriminator to provide instructive rewards. Thus, we focus on comparing with intrinsic-reward based methods, and generic representation learning methods do not fall into the group. In **Figure 1** in the attached PDF, we additionally compared STG with more related online adversarial methods **AMP, IDDM** that provide intrinsic rewards and an IDM-based method **BCO**. The high scores and sample efficiency show that STG can efficiently provide effective intrinsic rewards for downstream RL tasks.\n\n`Q3.`The meaning of \"State-to-Go\". \n\n`A3.` \"State-to-Go\" (STG) is named as an analogy to \"Return-to-Go\" in Decision Transformer (DT), but distinct from direct network inputs or policy conditions in DT. We use the name for the following reasons:\n\n(1) \"STG\" directly displays the **model property**. We want to stress that STG-Transformer features on modeling pure sequences consisting of continuously adjacent states, rather than modeling (s,a,R) sequence in DT or others in its follow-ups.\n\n(2) \"STG\" emphasizes the **problem setting**. \"State-to-Go\" suggests that we focus on learning from state offline datasets without actions or rewards, which is relatively difficult but practical.\n\n(3) \"STG\" indicates the **algorithm feature**. No matter in offline-pretraining or online learning stage, the key points lie in distinguishing expert or non-expert states. The name STG helps summarize the idea of our algorithm and helps readers understand and memorize our methods. \n\nUpon careful consideration of the reviewer's comments, to address this concern and improve clarity, we will revise the name of our algorithm. \n\n`Q4.`\"what are the main factors that make video game domains is more challenging (compared with robotic control tasks)? \" \n\n`A4.` Video games are challenging in the following aspects: \n\n1. Complex Dynamics \n2. Randomness \n3. Egocentric Visual Observation \n4. Inaccessible Rewards & Actions \n5. Benchmarks \n\nUnlike robotic control tasks with deterministic dynamics and accessible instant feedback, video games like Minecraft typically involve intricate and stochastic dynamics that are hard to model. The reward signals in video games are usually sparse and delayed, and the actions are hard to record in most cases. Besides, observations of robotic control usually come from third-person cameras or even proprioceptive states while video games input egocentric images, which makes LfVO a challenging but practical setting in video game domains. Consequently, video games like Atari and Minecraft have become popular benchmarks to evaluate learning from observations or even multi-modal prior knowledge [1,2,3]. \n\n`Q5.`\"the paper assumes some background knowledge of prior methods by passing too fastly over the baselines and progress reward is not clear.\" \n\n`A5.` We fastly pass over the baselines for sort out the main line: RL -> IL -> LfO -> LfVO -> offline LfVO. To our knowledge, in ELE, progression is monotonous from beginning to goal, so the progress reward reflects the temporal distance between any two key states. We employ the auxiliary task to learn temporall-aligned representation for STG. \n\n`Q6.`\"why predicted embeddings are used as a 'negative' samples in Equation (5) but are used as a 'expert' samples in Equation (9)?\"\n\n`A6.` During pre-training phase, the discriminator should learn to distinguish true expert transitions from fake transitions. Therefore, the expert observations serve as positive samples and generated samples serve as negative ones. During online reinforcement learning, especially in the early stage, online collected transitions $(s_t,s_{t+1})\\sim \\pi_\\theta$ are usually far from expert while $(s_t,\\hat{s}\\_{t+1})$ generated by offline pretrained STG Transformer can work as expert guidance. Under this circumstance, the intrinsic reward $r_i$ measures the gap between the current non-expert transition to an expected high-quality transition. A smaller gap means larger $r_i$.\n\nThanks again for the review! We will implement the feedback in the next version of this paper. Further comments are welcome!\n\n**Reference**\n\n[1] Hafner, Danijar, et al. \"Mastering diverse domains through world models.\" arxiv 2023.\n\n[2] Ding, Ziluo, et al. \"CLIP4MC: An RL-Friendly Vision-Language Model for Minecraft.\" arxiv 2023.\n\n[3] Yuan, Haoqi, et al. \"Plan4mc: Skill reinforcement learning and planning for open-world minecraft tasks.\" arxiv 2023."}, {"review_id": "6jyKQpk9It", "reviewer": "Reviewer_B1XQ", "summary": "This paper studies the problem of learning from visual observation. While most adversarial imitation learning methods learn intrinsic reward online, in order to improve the sample-efficiency of online learning, this work proposes to offline pre-train a state-to-go prediction transformer and a discriminator, both conditioned on current states and operating on next states. Representation learning is also enhanced with a shared encoder and an auxiliary temporal distance regression task. Empirical evaluations are conducted in both Atari and Minecraft domains.", "questions": "1. Some highly relevant baselines should be compared:\n   - BC from observation (BCO) with an IDM: Although the authors claim this method \"often suffers from compounding error\". A recent work VPT [1] has successfully adopted behavior cloning in the challenging Minecraft domain. Since this work focuses on online imitation learning, BCO, which learns IDM from online interactions, should be compared.\n   - Enhanced GAIfO: This paper introduces an offline pretraining framework and State-to-go transformer into adversarial imitation learning. Additional design choices from this work include WGAN and representation learning, which can also equip naive GAIfO. Note that advanced techniques to stabilize adversarial training have been studied by literature. For example, AMP [2] uses a least-squares GAN to mitigate gradient vanishing, similar to this work. Comparison with enhanced GAIfO baselines can further highlight the innovation of the proposed method.\n2. It seems the offline adversarial framework in this paper only handles conditional distribution $p(s_{t+1}|s_t)$ rather than marginal distribution $p(s_t, s_{t+1})$ like GAIL, since STG transformer generates samples conditioned on previous states and discriminator contrasts $(s_t, \\hat s_{t+1})$ and $(s_t, s_{t+1})$. This may introduce a distribution shift problem. Since at the beginning of online RL, states encountered by the agents can easily deviate from expert distribution, both generator and discriminator have to deal with out-of-distribution states $s_t$. Can the authors provide any intuition or empirical evidence in terms of this problem?\n\nIf the authors solve my questions properly, especially for empirical evaluation, I will be happy to increase my rating.\n\n[1] Baker, B., Akkaya, I., Zhokov, P., Huizinga, J., Tang, J., Ecoffet, A., ... & Clune, J. Video pretraining (vpt): Learning to act by watching unlabeled online videos.\n\n[2] Peng, X. B., Ma, Z., Abbeel, P., Levine, S., & Kanazawa, A. Amp: Adversarial motion priors for stylized physics-based character control.\n\n------\n\nUpdate: The rebuttal from the authors, including additional experimental results and a detailed explanation of the reward design, has addressed most of my issues. Thus I raised my score from 3 to 5.", "limitations": "This work has discussed its limitations and future work in the conclusion. There does not seem to be any negative social impact of this paper that should be discussed. ", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. To my knowledge, the proposed offline pretraining framework for adversarial imitation learning is novel and intuitive.\n2. Empirical experiments have been conducted in the challenging Minecraft domain.\n3. This paper is well-written and easy to follow.", "weaknesses": "In my opinion, the most significant weakness is that empirical evaluations are insufficient to support the superiority of the proposed method since some highly relevant but strong baselines are not compared. (see Question 1 below).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for careful review and valuable feedback, and for the positive comments that “offline-pretraining framework for adversarial imitation learning is novel and intuitive”. We address the questions and concerns below.\n\n`Q1.` \"More empirical evaluations to support the superiority of the proposed method.\"\n\n`A1.` We would like to explain the reasons for choosing baselines. The main idea of the paper is to **offline pretrain a State-To-Go model extracting visual transition knowledge to provide intrinsic rewards for downstream online RL tasks**. It is the first attempt to offline learn a reward function for online adversarial imitation learning to our knowledge. Thus, in the paper we focus on comparing with methods whose main job is to recover a reward function from visual observations for model-free RL algorithms. Naturally, the IDM-based method BCO [3] **does not fall into the group**, so we choose GAIfO and ELE. Moreover, both GAIfO and ELE **exhibit better performance compared with BCO in their original papers**.\n\nWe thank the reviewer for the reasonable suggestions. As shown in the attached PDF **Figure 1**, we compare STG with additional 3 baselines including IDM-based method **BCO** [1], \"advanced-GAIfO\" **AMP** [2], and **IDDM** [3] in Atari. Just as the reviewer mentioned, BCO is a strong online IDM-based imitation learning baseline that achieves competitive performance in Breakout, Freeway and SpaceInvaders, and \"advanced-GAIfO\" methods exhibit some advantages compared with GAIfO to some extent. In terms of STG, the final performance is the most competitive among all baselines, and the sample efficiency prominently outperforms most baselines during training. It is worth noting that the final scores of STG in Breakout and Qbert even **exceed expert performance as listed in Table in the paper** while online imitation learning method BCO fails. All these evidences demonstrate STG's remarkable potential of offline learning from visual observations. \n\n\n\n`Q2.` \"The offline adversarial framework in this paper only handles conditional distribution rather than marginal distribution?\"\n\n`A2.` GAIfO [4] optimizes target $\\min_{\\pi \\in \\Pi} \\max\\_{D \\in(0,1)^{{\\mathcal{S}} \\times \\mathcal{S}}} \\mathbb{E}\\_\\pi\\left[\\log \\left(D\\left(s, s^{\\prime}\\right)\\right)\\right]+ \\mathbb{E}\\_{\\pi_E}\\left[\\log \\left(1-D\\left(s, s^{\\prime}\\right)\\right)\\right]$ by online learning a transition discriminator to make $p_\\pi(s_t,s_{t+1})$ closer to $p_{\\pi_E}(s_t,s_{t+1})$. The optimization exactly works on marginal distribution without distribution shift but it suffers from low sample efficiency. To alleviate this problem, we offline pretrain STG-Transformer to predict $\\hat{s}_{t+1}$ given $s_t$. \n\nConsidering policy conducted by a learning agent may differ from expert policy, during online learning there would exist some \"novel\" states out of observation dataset distribution. Nevertheless, our **reward design**, denoted as $r_{t}^{i}=D_{\\omega}\\big( E_{\\xi}\\left( s_t \\right) ,E_{\\xi}\\left( s_{t+1} \\right) \\big)-D_{\\omega}\\big( E_{\\xi}\\left( s_t \\right) ,T_{\\sigma}\\left( E_{\\xi}\\left( s_t \\right) \\right) \\big)$, contributes to mitigating this problem. During pre-training, the discriminator learns to effectively distinguish \"non-expert\" transitions as well as \"fake\" transitions. At the beginning of online RL, when facing an OOD $s_t$ at timestep $t$, the discriminator will assign relatively low scores to both low-quality online transition $(s_t,s_{t+1})$ and less-accurately predicted transition $(s_t,\\hat{s}\\_{t+1})$. Consequently, the intrinsic reward $r_{t}^{i}$ almost gets close to zero, which indicates that the OOD $s_t$ will be attached less importance to during RL optimization. \n\nEmpirical evaluations also support our intuitions. According to **Figure 4 in the attached PDF**, we visualize the sum of intrinsic rewards, i.e. intrinsic return, of four Atari games respectively with 4 different random seeds. The increasing trend of intrinsic return demonstrates the effectiveness of alleviating distribution shift and guiding RL training. In Freeway, the rising trend is less significant because the episodic return of STG efficiently reaches a plateau in the early stages of training. \n\nTo sum up, no matter from intuition or empirical evidence, STG is capable of solving LfVO problems excellently.\n\nThanks again for the review! We will implement the feedback in the next version of this paper. Further comments are welcome!\n\n**Reference**\n\n[1] Torabi Faraz, et al. \"Behavioral Cloning from Observation\". arxiv 2018.\n\n[2] Peng, Xue Bin, et al. \"Amp: Adversarial motion priors for stylized physics-based character control.\" ACM 2021.\n\n[3] Yang, Chao, et al. \"Imitation learning from observations by minimizing inverse dynamics disagreement.\" NeurIPS 2019.\n\n[4] Torabi Faraz, et al. \"Generative adversarial imitation from observation.\" CoRR 2018.\n"}, {"review_id": "tbKuLlw4W0", "reviewer": "Reviewer_JcJo", "summary": "This paper proposes a novel approach to deal with the \"learning from expert visual observation\" setting, lacking action labels. The proposed method first pretrains State-to-Go Transformer (STG), which takes the sequence of encoded visual observations as inputs and predicts next-step embeddings under causal masking. It is trained with adversarial training inspired by WGAN, and with temporally-aligned regressor objectives. The agents are learned, similar to IRL, to follow the action-unlabelled demonstrations. The experiments show the better performance of STG approach on 4 Atari tasks and 4 Minecraft tasks.", "questions": "- Why did the authors use PPO for the online RL phase, instead of other algorithms, such as SAC as the authors used to generate the demonstrations in Atari experiments?\n- Are STG transformer, discriminator, and encoder frozen during online RL, or co-optimized with online interactions?", "limitations": "I think this paper appropriately mentioned the limitations and broader impacts.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "### Originality\n- Incorporating sequence modeling with causal transformer and temporally-aligned regressor into learning from expert visual observation imitation learning seems novel and interesting.\n\n### Quality\n- Overall, the paper is well-written and easy to follow.\n\n### Clarity\n- Figure 1 helps us understand the brief workflow of the proposed method.\n\n### Significance\n- The performance of STG in Atari and Minecraft is better compared to the baselines (ELE, GAIfO, PPO, Expert).\n- It would be an insightful observation to the community that combining temporally-aligned regressor into sequence modeling works really well (Figure 5).\n", "weaknesses": "### Originality\n- Seeing Figure 3,4,5, the results of STG without temporally-aligned regressor is almost the same as ELE. The most important contribution may come from the off-the-shelf module. Maybe ELE + temporally-aligned regressor would be an interesting ablation if applicable.\n\n### Clarity\n- \"State-to-Go\" might be misleading naming, if it is named as an analogy of Return-to-Go in Decision Transformer. The role of State-to-Go Transformer is quite different from Return-to-Go in Decision Transformer. The agent do not leverage any multi-step-extended values or state representations.\n- Minecraft experiments use 50000 demonstrations, while ones on Atari only use 50 demonstrations. There is no description on such a significantly different data size. Also, including the ablation with dataset size is interesting.\n- The authors said, \"Dopamine is not the expert demonstration\" in Section 4.1, but I guess the generated dataset is also not the expert one, seeing Table 4. It is good to include the ablation on dataset quality.\n\n\n### Significance\n- In Minecraft, MineDojo or CLIP4MC builds open-source large-scale pre-trained models that can predict the surrogate reward to measure the similarity between rollouts and textual descriptions. Those are also useful in LfVO settings considered in this paper. In this current situation, the significance of building a small-scale per-task module for surrogate reward is unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for careful review and constructive suggestions, and for praise our originality \"incorporating sequence modeling and temporally-aligned regressor into LfVO\" \"novel and insightful\". We address the questions and concerns below.\n\n`Q1.`\"Maybe ELE + temporally-aligned regressor would be an interesting ablation if applicable.\" \n\n`A1.` In fact, ELE[1] offline learns a temporally-aligned regressor to provide progression rewards for online RL training. STG employs a temporally-aligned regressor to learn high-quality temporally-aligned state representation for transition-prediction module. It offline leans a WGAN discriminator in an adversarial manner to provide intrinsic rewards for downstream RL tasks.\n\n`Q2.`\"The most important contribution may come from the off-the-shelf module?\"\n\n`A2.` We acknowledge that the TDR module indeed plays a vital role in learning temporally-aligned representation in our work. Without TDR module, the STG-Transformer may not be able to effectively capture latent transition patterns, potentially leading to low-quality discrimination rewards (as discussed in **Section 4.3**). We sincerely hope this insight will prove to be valuable and impactful to the community.\n\nBut the most important contribution we want to emphasize is that **unlike most previous adversarial imitation learning methods learning intrinsic rewards online, we offline pretrain STG to effectively and efficiently provide intrinsic rewards guidance for solving difficult visual RL tasks**. We believe this successful attempt highlights the promising future of making an agent \"learn like a human\". After watching tutorial videos, it can explore with experiences and get down to complete unfamiliar tasks.\n\n`Q3.`The name of \"State-to-Go\". \n\n`A3.`\"State-to-Go\" (STG) is named as an analogy to \"Return-to-Go\" in Decision Transformer (DT), but is not intended for values or representations of states. We use the name \"State-to-Go\" for the following reasons:\n\n(1) \"STG\" directly displays the **model property**. We want to stress that STG-Transformer features on modeling pure sequences consisting of continuously adjacent states, rather than modeling (s,a,R) sequence in DT or others in its follow-ups.\n\n(2) \"STG\" emphasizes the **problem setting**. \"State-to-Go\" suggests that we focus on learning from state offline datasets without actions or rewards, which is difficult but practical.\n\n(3) \"STG\" indicates the **algorithm feature**. No matter in offline-pretraining or online learning stage, the key points lie in distinguishing expert or non-expert states. The name STG helps summarize the idea of our algorithm and helps readers understand and memorize our methods. \n\nTo address this concern and improve clarity, we will revise the name in our final version as the reviewer suggests. \n\n`Q4.`\"Minecraft experiments use 50000 demonstrations, while ones on Atari only use 50 demonstrations\". \n\n`A4.`About dataset size, we collect around 50 trajectories, specifically, around 100000 transitions (cf. **Line 254**) in Atari and around 50000 (cf. **Line 299**) transitions in Minecraft. Two datasets are similar in the number of trajectories. **Figure 2(c)** of the attached PDF shows that dataset size indeed makes a difference to the performance.\n\n`Q5.`\"Both Dopamine and generated datasets are not expert. It is good to include the ablation on dataset quality.\"\n\n`A5.`We compare the learning curve between the algorithm [2] for generating Dopamine datasets and our self-implemented SAC. We find that the final performance of [2] drops in the later stages of training and is less competitive than the final performance of SAC. So for Breakout and Freeway we choose to train a SAC agent as expert to provide higher-quality observation datasets. We thank the reviewer for raising an interesting point to ablating dataset quality. As mentioned in Author Rebuttal, **Figure 2(d)** of the attached PDF shows the results ablating dataset quality.\n\n`Q6.`\"The significance of building a small-scale per-task module for surrogate reward.\"\n\n`A6.`In the paper, the reason for pre-training STG in task-specific dataset to solve the corresponding task is to evaluate its capability of offline extracting instructive visual information solely from small observation datasets. To demonstrate that STG can transfer to multi-task setting, we conduct additional experiments by pretraining STG Transformer on Atari datasets encompassing four tasks. The model pre-trained on these multi-task datasets (STG-Multi) is then used to guide RL training for each specific task. To accommodate four-time larger training dataset, we augment the model capacity by increasing the number of heads (24) and layers (16) in self-attention modules. As depicted in **Figure 3** of the attached PDF, the comparable performance reveals the potential of pretraining STG on large-scale multi-task datasets for guiding downstream tasks. With model size scaling up, we firmly believe that STG will become a potent technique to solve larger-scale LfVO tasks, showcasing its versatility and utility in more complex scenarios.\n\n`Q7.`\"Choice of RL algorithms.\"\n\n`A7.`We would like to clarify that any reinforcement method whose target is to maximize the expectation of cumulative discounted rewards can be utilized to generate demonstrations or to optimize the intrinsic return.\n\n`Q8.`\"Are STG transformer, discriminator, and encoder frozen during online RL, or co-optimized with online interactions?\"\n\n`A8.` After pre-training, modules including the discriminator, encoder, and STG transformer are all frozen and concurrently work as a \"reward function\" to provide intrinsic rewards for an online RL agent without fine-tuning.\n\nThanks again for the review! We will implement the feedback in the next version of this paper.\n\n**Reference**\n\n[1] Bruce Jake, et al. \"Learning about progress from experts.\" ICLR 2022.\n\n[2] Rishabh Agarwal, et al. \"An optimistic perspective on offline reinforcement learning\". ICML 2020"}]

Author rebuttal:
We sincerely thank the four reviewers for their thoughtful comments! We have completed some supplements according to reviewers' suggestions, and we summarize the major changes as follows:

1. We have added three baselines including IDM-based method **BCO**, "advanced-GAIfO" **AMP** and IRLfO method **IDDM** in Atari domain.
2. We ablate the loss items and offline dataset to investigate the factors for good performance in Atari domain. 
3. We have conducted multi-task experiments in Atari domain, proving the efficacy of multi-task adaptation of STG.

We provide more details as follows:

1. **About Baselines**

   The main idea we want to convey is to **offline pretrain a State-To-Go model extracting visual transition knowledge to provide intrinsic rewards for downstream online RL tasks**. To our knowledge, it is the first attempt to offline learn a reward function for online adversarial imitation learning. Therefore, in the paper, we focus on comparing methods whose main job is to recover a reward function from visual observations for model-free RL algorithms. Naturally, IDM-based method BCO [3] **does not fall into the group**, so we select two representatives as baselines: 

   (1) LfO baseline GAIfO [1] which recovers rewards from online transitions 

   (2) recent work ELE [2] which derives rewards from offline representation learning.

   Moreover, both GAIfO and ELE **exhibit better performance compared with BCO in their original papers**. Of course, we believe more baselines help verify the effectiveness of STG. So three baselines including IDM-based method **BCO** [3], "advanced-GAIfO" **AMP** [4], and IRLfO method **IDDM** [5] are added. The learning curves of four Atari tasks shown in **Figure 1** in the attached PDF demonstrate the competitive performance of STG.

2. **About Ablations**

   Incorporating TDR with WGAN is one of the most important designs within our paper. We have conducted an ablation study on the key TDR module to prove its importance. **The influence of other factors appears comparatively less substantial in comparison to TDR or has been investigated in previous literature** (like ablating datasets in [1] and [3]). Upon reasonable suggestions of the reviewers, we have expanded our ablation exploration to encompass loss design and datasets. We find modifications to pre-training loss or datasets lead to analogous variations across different environments. Therefore, we choose one Atari task for each ablation. 

   (1) loss design

   During pre-training, STG concurrently minimize $L_{tot}=\alpha L_{mse}+\beta L_{adv}+\kappa L_{tdr}$. We set $\alpha=0$ to investigate the contribution of $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 [6]. The final performance of STG(rm MSE) and STG(rm Adv) in SpaceInvaders are listed in **Figure 2(a) and 2(b)** in the attached PDF. The sample efficiency slightly drops without L2 penalty while the final performance declines heavily without WGAN. Thus, we can conclude each item in $L_{tot}$ makes an indispensable contribution to the final performance. 

   (2) dataset size and quality 

   We conduct an ablation study by reducing the size of offline dataset in SpaceInvaders. We randomly sample 50,000 transitions, i.e. half of the original dataset to train **STG(Half-Data)** to solve with other settings unchanged. The outcomes of this experiment are presented in **Figure 2(c)** of the attached PDF. 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 [1] and Figure 3 of [3]**.

   Furthermore, we ablate dataset quality 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 Breakout task. **Figure 2(d)** in the attached PDF 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. 

3. **About model and dataset scale-up for multi-task"**

   In the paper, we focus on evaluating STG's capability of offline extracting instructive visual information solely from small observation datasets. To demonstrate that STG can transfer to the multi-task setting, we conduct additional experiments by pretraining STG Transformer on Atari datasets encompassing four tasks. The model pre-trained on these multi-task datasets (STG-Multi) is then used to guide RL training for each specific task. To accommodate a four-time larger training dataset, we augment the model capacity by increasing the number of heads (24) and layers (16) in self-attention modules. As depicted in **Figure 3** in the attached PDF, the comparable performance reveals the potential of pretraining STG on large-scale multi-task datasets for guiding downstream tasks. With model size scaling up, we firmly believe that STG will become a potent technique to solve larger-scale LfVO tasks, showcasing its versatility and utility in more complex scenarios.

All these experiments will be incorporated into the next vision of this paper. 

**Reference**

[1] Torabi Faraz, et al. "Generative adversarial imitation from observation." CoRR 2018.

[2] Bruce Jake, et al. "Learning about progress from experts." ICLR 2022.

[3] Torabi Faraz, et al. "Behavioral Cloning from Observation". arxiv 2018.

[4] Peng Xue Bin, et al. "Amp: Adversarial motion priors for stylized physics-based character control." ACM 2021.

[5] Yang Chao, et al. "Imitation learning from observations by minimizing inverse dynamics disagreement." NeurIPS 2019.

[6] Zhu Deyao, et al. "Guiding online reinforcement learning with action-free offline pretraining." arXiv 2023.

Meta-review: The paper presents a new method, called state-to-go (STG) transformer, for learning from visual observations. The paper in its current form is borderline; based on the reviewers' comments and the authors' rebuttal, I believe the paper has addressed several concerns but still has some aspects that could be further refined. 

Reviewer 9y8A expressed concerns about the technical contribution and the limited scope of the proposed method. The authors have appropriately addressed the concerns regarding baselines and ablations, which has led to a positive adjustment in the reviewer's score. However, the reviewer still questions the novelty of the proposed loss function in the broader context of trajectory representation learning. To address this concern, I recommend that the authors emphasize the unique aspects of their loss function within the context of offline pre-training for visual reinforcement learning tasks. While the focus is primarily on Minecraft, discussing potential extensions or applications to other domains, even if speculative, could broaden the paper's appeal.

Reviewer JcJo raised a concern about the terminology used in the paper, particularly "State-to-Go (STG)," and suggested an alternative term such as "OLfVO" (Offline Learning from Visual Observations). The authors should consider this suggestion, especially if it improves clarity and aligns better with the algorithm's nature. Furthermore, the reviewer questioned the significance of a task-specific model in the presence of large-scale pre-trained models in the Minecraft domain. To address this concern, the authors could provide a more explicit comparison with these existing models, highlighting the advantages of their approach in terms of efficiency, performance, or any other relevant metrics.

Reviewer HW2S expressed lingering doubts about the claimed contribution of the proposed representation learning method. While the authors have provided some response to this concern, the reviewer remains uncertain due to the limited number of baselines and tasks evaluated. To address this, I recommend that the authors explore additional baseline methods and tasks to provide a more comprehensive evaluation of their proposed method's performance. This will help strengthen the evidence for the claimed contribution and establish the method's robustness.

In conclusion, the paper shows promise in addressing the challenges of learning from visual observation for open-ended tasks. The authors have effectively responded to several concerns raised by the reviewers, and for this reason my recommendation is for accept. However, there is room for improvement, and the authors should address the remaining issues in time for a camera ready.
Decision: Accept (poster)