Title: Learning from Visual Observation via Offline Pretrained State-to-Go Transformer
Venue: NeurIPS 2023
Accepted: yes
Keywords: Learning from Observations

Offline Learning from Visual Observations

State-to-Go Transformer
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

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 Stateto-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.

STG Transformer

Encoder

Temporally-aligned

Policy

Encoder

Intrinsic

Observe

Transfer

PPO

Stage 1:Offline Pretraining

Stage 2:Online RL

Input

Input

Transition

Discrimination

1 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 (IL). IL agents learn policies by imitating expert demonstrations in a data-driven manner rather than through trial-and-error processes. It has been proven effective in various domains, including games [1] and robotics [2].

However, acquiring demonstrated actions can be expensive or impractical, e.g. , from videos that are largely available though, leading to the development of learning from observation (LfO) [3, 4, 5, 6, 7, 8, 9]. This line of research utilizes observation-only data about agent behaviors and state transitions for policy learning. Humans naturally learn from visual observation without requiring explicit action guidance, such as beginners in video games improving their skills by watching skilled players' recordings. However, LfO agents face challenges in extracting useful features from raw visual observations and using them to train a policy due to the lack of explicit action information. Thus, further study of learning from visual observation (LfVO) has the potential to grant agents human-like learning capabilities.

In this paper, we investigate a reinforcement learning setting in which agents learn from visual observation to play challenging video games, such as Atari and Minecraft. Existing LfO approaches like [4, 5, 6, 7, 8] primarily apply to vector-observation environments such as MuJoCo. Others like [3, 9, 10, 11, 12, 13, 14] explicitly consider or can be applied to high-dimensional visual observations. Among them, representation-learning methods [9, 10, 11] learn effective visual representations and recover an intrinsic reward function from them. Adversarial methods [3, 12, 13] learn an expertagent observation discriminator online to directly indicate visual differences. However, noises or local changes in visual observations may easily cause misclassification [15]. In [12, 13], additional proprioceptive features ( e.g. , joint angles) are used to train a discriminator, which are unavailable in environments that only provide visual observations. Moreover, as these methods require online training, sample efficiency is much lower compared to offline learning. Goal-oriented methods, like [14], evaluate the proximity of each visual observation to expert demonstrations or predefined goals. However, defining explicit goals is often impractical in open-ended tasks [16]. Furthermore, the continuity of observation sequences in video games cannot be guaranteed due to respawn settings or unanticipated events.

To address these limitations and hence enable RL agents to effectively learn from visual observation, we propose a general two-stage framework that leverages visual observations of expert demonstrations to guide online RL. In the first stage, unlike existing online adversarial methods, we introduce and pretrain State-to-Go (STG) Transformer , a variant of Decision Transformer (DT) [17], for offline predicting transitions in latent space. In the meanwhile, temporally-aligned and predictable visual representations are learned. Together, a discriminator is trained to differentiate expert transitions, generating intrinsic rewards to guide downstream online RL training in the second stage. That said, in the second stage, agents learn merely from generated intrinsic rewards without environmental reward signals. Our empirical evaluation reveals significant improvements in both sample efficiency and overall performance across various video games, demonstrating the effectiveness of the proposed framework.

Our main contributions are as follows:

We propose a general two-stage framework, providing a novel way to enable agents to effectively learn from visual observation. We introduce State-to-Go Transformer, which is pretrained offline merely on visual observations and then employed to guide online reinforcement learning without environmental rewards.

We simultaneously learn a discriminator and a temporal distance regressor for temporallyaligned 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 to complete a task without the assistance of action or reward guidance. Many existing works [5, 18, 19, 20] attempt to train an inverse dynamic model (IDM) to label observation-only demonstrations with expert actions, enabling behavior cloning. However, these methods often suffer from compounding error [21]. On the other hand, [4] learns a latent policy and a latent forward model, but the latent actions can sometimes be ambiguous and may not correspond accurately with real actions. GAIfO [3], inspired by [22], is an online adversarial framework that couples the process of learning from expert observations with RL training. GAIfO learns a discriminator to evaluate the similarity between online-collected observations and expert demonstrations. Although helpful in mitigating compounding error [3], it shows limited applicability in environments with high-dimensional observations. Follow-up methods [12, 13] pay more attention to visual-observation environments, but require vector-state in expert observations to either learn a feasible policy or proper visual representations. More importantly, learning a discriminator online is less sample-efficient, compared to LfO via offline pretraining. A recent attempt [23] demonstrates some progress in action-free offline pretraining, but return-to-gos are indispensable in addition to observations because of upside down reinforcement learning (UDRL) framework [24]. Moreover, it only shows satisfactory results in vector-observation environments like MuJoCo. In this work, we focus on learning from offline visual observations (LfVO), a more general and practical setting, which is also considered by the concurrent work VIPER [25]. STG and VIPER both offline pretrain a reward function and then use it for online reinforcement learning. STG learns discriminative intrinsic rewards in WGAN style while VIPER leverages logarithm predictive probability as guidance.

Visual Representation Learning in RL. High-quality visual representations are crucial for LfVO. Many previous works [26, 27, 28, 9, 10, 11] have contributed to this in various ways. For example, [26] employs GANs to learn universal representations from different viewpoints, and [27, 28] learn representations via contrastive learning to associate pairs of observations separated by a short time difference. In terms of LfVO, a wide range of methods such as [9, 10, 11] learn temporally continuous representations in a self-supervised manner and utilize temporal distance to assess the progress of demonstrations. Nevertheless, in games like Atari, adjacent image observations may exhibit abrupt or subtle changes due to respawn settings or unanticipated events, not following a gradual change along the timeline. Moreover, over-reliance on temporal information often results in over-optimistic estimates of task progress [14], potentially misleading RL training.

Transformer in RL. Transformer [29] is widely acknowledged as a kind of powerful structure for sequence modeling, which has led to domination in a variety of offline RL tasks. Decision Transformer (DT) [17] and Trajectory Transformer (TT) [30] redefine the offline RL problem as a contextconditioned sequential problem to learn an offline policy directly, following the UDRL framework [24]. DT takes states, actions, and return-to-gos as inputs and autoregressively predicts actions to learn a policy. TT predicts the complete sequence dimension by dimension and uses beam search for planning. MGDT [31] samples from a learned return distribution to avoid manually selecting expert-level returns as DT. ODT [32] extends DT to bridge the gap between offline pretraining and online fine-tuning. MADT [33] extends DT to a multi-agent reinforcement learning setting.

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 , γ, ρ 0 > . S denotes the state space and A denotes the action space. P : S × A × S → [0 , 1) is the state transition function and R : S × A → R is the reward function. γ ∈ [0 , 1] is the discount factor and ρ 0 : S → [0 , 1] represents the initial state distribution. The objective is to find a policy π ( a | s ) : S → A , which maximizes the expected discounted return:

Transformer. Stacked self-attention layers with residual connections in Transformer is instrumental in processing long-range dependencies, each of which embeds n tokens { x i } n i =1 and outputs n embeddings { z i } n i =1 of the same dimensions considering the information of the whole sequence. In this study, we utilize the GPT [35] architecture, an extension of the Transformer model, that incorporates a causal self-attention mask to facilitate autoregressive generation. Specifically, each input token x i is mapped to a key k i , a query q i , and a value v i through linear transformations, where z i is obtained by computing the weighted sum of history values v 1: i , with attention weights determined by the normalized dot product between the query q i and history keys k 1: i :

The GPT model only attends to the previous tokens in the sequence during training and inference, thereby avoiding the leakage of future information, which is appropriate in state prediction.

Learning from Observation. The goal is to learn a policy from an expert state sequence dataset D e = { τ 1 , τ 2 , . . . , τ m } , τ i = { s i 1 , s i 2 , . . . , s i n } , s i j ∈ S . Denote the transition distribution as µ ( s, s ′ ) . The objective of LfO can be formulated as a distribution matching problem, finding a policy that minimizes the f -divergence between µ π ( s, s ′ ) induced by the agent and µ e ( s, s ′ ) induced by the expert [7]:

It is almost impossible to learn a policy directly from the state-only dataset D e . However, our delicately designed framework (see Figure 1) effectively captures transition features in expert demonstrations to provide informative guidance for RL agents, which will be expounded in the following.

3.2 Offline Pretraining Framework

STG Transformer is built upon GPT [35] similar to DT [17], but with a smaller scale and more structural modifications to better handle state sequence prediction tasks. Unlike DT, in our setting, neither the action nor the reward can be accessible, so the STG Transformer primarily focuses on predicting the next state embedding given a sequence of states.

As depicted in Figure 2, first we concatenate a few consecutive image frames in the expert dataset to approximate a single state s t . Then, a sequence of n states { s t , . . . , s t + n -1 } are encoded into a sequence of n token embeddings { e t , . . . , e t + n -1 } by the feature encoder E ξ composed of several CNN layers and a single-layer MLP, where e t = E ξ ( s t ) . A group of learnable positional embedding parameters is added to the token embedding sequence to remember temporal order. These positional-encoded embeddings are then processed by the causal self-attention module which excels in incorporating information about the previous state sequence to better capture temporal

emb. + pos. enc.

decoder

84

84

4

Feature Encoder

TDR Predictor

……

dependencies, followed by layer normalization. The linear decoder outputs the final latent prediction sequence { ˆ e t +1 , . . . , ˆ e t + n } . Denote the positional encoding, transition predicting, and linear decoding model together as T σ . It is worth noting that instead of predicting the embeddings of the next state sequence directly, we predict the embedding change and combine it with token embeddings in a residual way, which is commonly applied in transition prediction [4] and trajectory forecasting [36] to improve prediction quality.

For brevity, in further discussion we will refer to T σ ( e t ) directly as the predicted ˆ e t +1 .

Expert Transition Discrimination. Distinguishing expert transiting patterns is the key to leveraging the power of offline expert datasets to improve sample efficiency in downstream RL tasks. Traditional online adversarial methods [3, 12, 13] employ a discriminator to maximize the logarithm probability of transitions sampled from expert datasets while minimizing that from transitions collected online,

which is often sample-inefficient in practice. Moreover, in the case of visual observation, the traditional discriminator may rapidly and strictly differentiate expert transitions from those collected online within a few updates. As a result, the collected observations will be assigned with substantially low scores, making it arduous for policy improvement.

To overcome these limitations, we draw inspiration from WGAN [37] and adopt a more generalized distance metric, known as the Wasserstein distance, to measure the difference between the distributions of expert and online transitions. Compared to the sigmoid probability limited in [0 , 1] , the Wasserstein distance provides a wider range and more meaningful measure of the difference between two transition distributions, as it captures the underlying structure rather than simply computing the probability. More importantly, unlike traditional online adversarial methods like GAIfO [3] that use the Jensen-Shannon divergence or Kullback-Leibler divergence, the Wasserstein distance is more robust to the issues of vanishing gradients and mode collapse, making offline pretraining possible. Specifically, two temporally adjacent states s t , s t +1 are sampled from the expert dataset, then we have e t = E ξ ( s t ) , e t +1 = E ξ ( s t +1 ) , and ˆ e t +1 = T σ ( E ξ ( s t )) . The WGAN discriminator D ω aims to maximize the Wasserstein distance between the distribution of expert transition ( e t , e t +1 ) and the distribution of predicted transition ( e t , ˆ e t +1 ) , while the generator tries to minimize it. The objective can be formulated as:

{ D ω } ω ∈W represents a parameterized family of functions that are 1-Lipschitz, limiting the variation of the gradient. We clamp the weights to a fixed box ( W = [ -0 . 01 , 0 . 01] l ) after each gradient update to have parameters w lie in a compact space. Besides, to suppress the potential pattern collapse, an additional L 2 norm penalizes errors in the predicted transitions, constraining all e t and ˆ e t in a consistent representation space. Thus, the loss functions can be rewritten as follows.

For discriminator:

For STG Transformer (generator):

By such an approach, the discriminator can distinguish between expert and non-expert transitions without collecting online negative samples, providing an offline way to generate intrinsic rewards for downstream reinforcement learning tasks.

Temporally-Aligned Representation Learning. Having a high-quality representation is crucial for latent transition prediction. To ensure the embedding is temporally aligned, we devise a selfsupervised auxiliary module, named temporal distance regressor (TDR). Since the time span between any two states s i and s j in a state sequence may vary significantly, inspired by [38], we define symlog temporal distance between two embeddings e i = E ξ ( s i ) and e j = E ξ ( s j ) :

This bi-symmetric logarithmic distance helps scale the value and accurately capture the fine-grained temporal variation. The TDR module P ϕ consists of MLPs with 1D self-attention for symlog prediction. The objective of TDR is to simply minimize the MSE loss:

Offline Pretraining. In our offline pretraining, the transition predictor T σ and transition discriminator D ω share the same feature encoder E ξ similar to online methods [39], which allows them to both operate in an easily-predictable and temporally-continuous representation space.

At each training step, a batch of transitions is randomly sampled from the expert dataset. The model is trained autoregressively to predict the next state embedding without accessing any future information. When backpropagating, L mse and L adv concurrently update E ξ and T σ to provide high-quality visual embeddings as well as accurate embedding prediction. L tdr is responsible for updating the E ξ and P ϕ as an auxiliary component, and L dis updates D ω . Algorithm 1 in Appendix A details the offline pretraining of the STG Transformer.

3.3 Online Reinforcement Learning

Intrinsic Reward. For downstream RL tasks, our idea is to guide the agent to follow the pretrained STG Transformer to match the expert state transition distribution. Unlike [14], our experimental results show that the 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 as the intrinsic reward for online RL. Moreover, we do not use 'progress' like what is done in [9]. This is because, in games with multiple restarts, progress signals can easily be inaccurate and hence mislead policy improvement, while the WGAN discriminator mastering the principle of transitions can often make the correct judgment. The intrinsic reward at timestep t is consequently defined as follows:

A larger r i t means a smaller gap between the current transition and the expert transition.

Online Learning Procedure. Given an image observation sequence collected by an agent, the feature encoder first generates corresponding visual representations, followed by the STG Transformer predicting the latent one-step expert transition. Then the discriminator compares the difference between real transitions and predicted transitions. Their Wasserstein distances, as intrinsic rewards r i , is used to calculate generalized advantage, based on which the agent policy π θ is updated using PPO [40]. It is worth noting that the agent learns the policy merely from intrinsic rewards and environmental rewards are not used.

4 Experiments

In this section, we conduct a comprehensive evaluation of our proposed STG on diverse tasks from two environments: a classical Atari environment and an open-ended Minecraft environment. Among the three mainstream methods mentioned in Section 1, goal-oriented methods are not appropriate for comparison because there is no pre-defined target state. Therefore, we choose IDM-based method BCO [19], online adversarial method GAIfO [3] and its follow-ups AMP [41], IDDM [42], and representation-learning method ELE [9].

For each task, we conduct 4 runs with different random seeds and report the mean and standard deviation. The same network architecture is applied for each algorithm to maintain consistency. Similar to [39], the discriminator and policy network share the same visual encoder for all online adversarial methods, and we choose one-step prediction for ELE to align with STG. Through extensive experiments, we answer the following questions:

Is our proposed framework effective and efficient in visual environments?

Is our offline pretrained discriminator better than the ones which are trained online?

Does TDR make a difference to visual representations? And do we need to add 'progress' rewards, as is done in ELE?

4.1 Atari

Atari Datasets. Atari is a well-established benchmark for visual control tasks and also a popular testbed for evaluating the performance of various LfVO algorithms. We conduct experiments on four Atari games: Breakout, Freeway, Qbert, and Space Invaders. To ensure the quality of expert datasets, two approaches are utilized to collect expert observations. For Qbert and SpaceInvaders, we collect the last 10 5 transitions (around 50 trajectories) from Google Dopamine [43] DQN replay experiences. For Breakout and Freeway, we find the quality of transitions from Dopamine is far from expert. Therefore, we alternatively train a SAC agent [44] from scratch for 5 × 10 6 steps and leverage the trained policy to gather approximately 50 trajectories (around 10 5 transitions) in each game to construct the expert dataset.

Performance in Atari Games. As illustrated in Figure 3, STG is on par with BCO in Freeway but outperforms all baselines in the rest tasks. Table 1 displays the final scores of all methods along with expert datasets and PPO learned from environmental rewards. Remarkably, in Breakout

STG

ELE

GAIfO

AMP

IDDM

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Transition

×

10 7

0

100

200

300

(a)

Breakout

0

1

2

3

4

5

Transition

×

10 6

0

5

10

15

20

(b)

Freeway

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

0

10000

20000

30000

(c)

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

100

200

300

400

500

600

(d)

Space Invaders

STG

ELE

GAIfO

AMP

IDDM

STG

ELE

0

200

400

600

800

1000

Epoch

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

(a)

Pick a flower

0

200

400

600

800

1000

Epoch

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

(b)

Milk a cow

0

200

400

600

800

1000

Epoch

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

(c)

Harvest tallgrass

0

200

400

600

800

1000

Epoch

0

.

00

0

.

02

0

.

04

0

.

06

(d)

Gather wool

and Qbert, the performance of STG surpasses offline datasets. In Breakout, we observe that STG learns to consecutively bounce the ball up into the top gaps to hit the upper bricks to obtain high intrinsic rewards, while online adversarial methods and ELE fail. In Qbert, STG achieves a prominent breakthrough in the later stages, substantially outperforming other baselines.

Analysis. During the training process, we observe that online adversarial methods tend to easily get stuck in a suboptimal policy and struggle to explore a better policy. This is because the discriminator can easily distinguish between the visual behavior of the expert and the imitator based on relatively insignificant factors within just a few online interactions. In contrast, STG learns better temporallyaligned representations in an offline manner, enabling the discriminator to detect more substantial differences. Besides, instead of relying on probability, STG employs the Wasserstein distance metric to provide more nuanced and extensive reward signals. Consequently, even without fine-tuning during the online RL process, STG can offer valuable guidance to the RL agent. Additionally, in Breakout and Freeway we observe that ELE drops in final performance primarily due to the over-optimistic progress, which will be further investigated in Section 4.3. In comparison, STG ensures precise expert transition prediction and discriminative transition judgment, avoiding over-optimistically driving the agent to transfer to new states.

4.2 Minecraft

Brief Introduction. Built upon Minecraft, Minedojo [45] provides a simulation platform with thousands of diverse open-ended tasks, which propose new challenges for goal completion. Furthermore, the 3D egocentric view and intricate scenes make it extremely hard to extract task-relevant visual signals for LfVO. We evaluate STG on four Minecraft tasks including 'pick a flower', 'milk a cow',

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Transition

×

10 7

0

100

200

300

STG

STG-

(a)

Breakout

0

2

4

Transition

×

10 6

0

5

10

15

20

STG

STG-

(b)

Freeway

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

0

10000

20000

30000

STG

STG-

(c)

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

200

400

STG

STG-

(d)

Space Invaders

0

250

500

750

1000

Epoch

0.00

0.05

0.10

0.15

0.20

STG

STG-

(e)

Pick a flower

0

250

500

750

1000

Epoch

0.00

0.05

0.10

0.15

0.20

0.25

0.30

STG

STG-

(f)

Milk a cow

0

250

500

750

1000

Epoch

0.0

0.1

0.2

0.3

STG

STG-

(g)

Harvest tallgrass

0

250

500

750

1000

Epoch

0.00

0.02

0.04

0.06

0.08

STG

STG-

(h)

Gather wool

'harvest tallgrass', and 'gather wool'. All tasks are sparse-reward, where only a binary reward is emitted at the end of the episode, thus the performance is measured by success rates.

Minecraft Datasets. Recently, various algorithms, e.g. , Plan4MC [16] and CLIP4MC [46] have been proposed for Minecraft tasks. To create expert datasets, for each task, we utilize the learned policies of these two algorithms to collect 100 trajectories (around 5 × 10 4 observations).

Performance in Minecraft. The results on Atari show that online adversarial methods can hardly tackle LfVO tasks well. Therefore, in Minecraft, we focus on comparing BCO, ELE, and STG. As depicted in Figure 4, the success rates across four Minecraft tasks reveal a consistent superiority of STG over ELE. In terms of IDM-based BCO, albeit equipped with an additional online-learned inverse dynamic module, it only shows a negligible advantage over STG in 'gather wool' and fails to achieve satisfactory performance in the rest three tasks, and even suffers from a severe performance drop in 'harvest tallgras'.

Notably, we observe that in 'milk a cow', STG attains a success rate approaching 25%, significantly eclipsing the 5% success rate of ELE or BCO. The reasons for this stark contrast are not yet entirely elucidated. However, a plausible conjecture could be attributed to the task's primary objective, i.e. locating the cow. STG, which is proficient in memorizing high-quality transitions, can effectively accomplish this subgoal, while ELE may lose the intended viewpoint due to its tendency for overoptimistic 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 domain, it raises a question regarding the contribution of the designed temporal distance regressor. To answer this question, we conduct an ablation study, denoted as STG-, where we remove the TDR loss L tdr from STG to figure out how the TDR module contributes to enhancing performance and representation. Consequently, the feature encoder E ξ and the STG Transformer T σ are pretrained under the guidance of a lin-

STG

STG-

ear combination of L mse and L adv . We keep the other settings fixed and leverage the pretrained model to tackle the respective downstream tasks. The results are shown in Figure 5, where STG is substantially superior to STG- in most tasks.

In order to figure out the underlying reasons for their discrepancy in performance, we compare the visualization of embeddings encoded by STG and STG-. We randomly select an expert trajectory from Qbert and utilize t-SNE projection to visualize their embedding sequences. As illustrated in Figure 6, the embeddings learned by STG exhibit remarkable continuity in adjacent states, significantly different from the scattered and disjoint embeddings produced by STG-. The superior temporal alignment of the STG representation plays a critical role in capturing latent transition patterns, thereby providing instructive information for downstream RL tasks. For more multi-trajectory visualizations, see Appendix G.

Multi-Task STG. In the previous experiments, the demonstrations used for pretraining come from agents solving the same tasks, it raises a question on how well do STG generalize to a variety of downstream tasks. To this end, we pretrain a new instance of the STG Transformer named STG-Multi , with the same network architecture, on the whole Atari datasets encompassing all four downstream tasks to further assess the efficacy of multi-task adaptation. Considering the four times increase in datasets, we enlarge the model capacity by about four times via stacking more multi-head causal self-attention modules. More details are available in Appendix E. As shown in Figure 7, STG-Multi shows comparable performance across four Atari games. These results strongly reveal the potential of pretraining STG on multi-task datasets for guiding downstream tasks.

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Transition

×

10 7

0

100

200

300

STG

STG-Multi

(a)

Breakout

0

1

2

3

4

5

Transition

×

10 6

0

5

10

15

20

STG

STG-Multi

(b)

Freeway

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

0

10000

20000

30000

STG

STG-Multi

(c)

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

100

200

300

400

500

600

STG

STG-Multi

(d)

Space Invaders

Pretraining Loss Design. During pre-training, STG concurrently minimizes L tot = α L mse + β L adv + κ L tdr . We set α = 0 to investigate the contribution of L mse , denoted as STG(rm MSE) , and β = 0 to investigate the contribution of WGAN, denoted as STG(rm Adv) . For STG(rm Adv), we can only use prediction error as intrinsic rewards like what has been done in [23]. The final performance of STG(rm MSE) and STG(rm Adv) in SpaceInvaders are listed in Figure 8a and 8b. The sample efficiency slightly drops without L 2 penalty while the final performance declines heavily without WGAN. Thus, we can conclude each item in L tot makes an indispensable contribution to the final performance.

It is worth noting that we arbitrarily choose α = β = 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 to 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-negligible impact on performance of LfVO tasks. We first conduct extra ablation studies on reduction of dataset size. 50,000 transitions, i.e. half of the original dataset are randomly sampled to train STG(Half-Data) in SpaceInvaders with other settings unchanged. The outcomes of this experiment are presented in Figure 8c. The ablation clearly demonstrates that an increased number of demonstrations contributes to enhanced performance. This observation aligns seamlessly with prior investigations, as illustrated in Figure 4 of [3] and Figure 3 of [19].

Furthermore, we ablate the quality of dataset in Breakout. We manually select 50 trajectories with an episodic return greater than 250 in Breakout to construct expert datasets to train STG(Expert) to solve the Breakout game. Figure 8d shows that higher-quality offline datasets contribute to better sample efficiency. Meanwhile, it also reflects that STG can achieve excellent performance learning from sub-optimal observations.

Progression Reward. Considering ELE follows the same scenario of pretraining to enhance downstream tasks as STG, we conduct experiments to figure out whether we can additionally add progres-

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

100

200

300

400

500

600

STG(rm MSE)

STG

(a)

Ablate removing

L

mse

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

100

200

300

400

500

600

STG

(b)

Ablate removing

L

adv

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

100

200

300

400

500

600

STG(Half-Data)

STG

(c) Ablate dataset size

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Transition

×

10 7

0

100

200

300

STG(Expert)

STG

(d)

Ablate dataset quality

sion rewards derived from TDR like ELE to further the performance. To this end, we train the agent under the guidance of both the discriminative and progression rewards from the same pretrained STG Transformer in Atari tasks, denoted as STG* . As illustrated in Figure 9, STG outperforms STG* in all tasks. We analyze that, similar to ELE, progression rewards from TDR over-optimistically urge the agent to "keep moving" to advance task progress, which however can negatively impact policy learning. For example, on certain conditions such as Breakout or Freeway, maintaining a stationary position may facilitate catching the ball or avoiding collision more easily, thereby yielding higher returns in the long run. Therefore, we do not include the over-optimistic progression rewards in our design.

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Transition

×

10 7

0

100

200

300

STG

STG*

(a)

Breakout

0

2

4

Transition

×

10 6

0

5

10

15

20

STG

STG*

(b)

Freeway

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

0

10000

20000

30000

STG

STG*

(c)

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

200

400

STG

STG*

(d)

Space Invaders

In summary, our experimental results provide strong evidence for the ability of STG to learn from visual observation, substantially outperforming baselines in a variety of tasks. The ablation study highlights the importance of the TDR module for temporally aligned representations. However, TDR may not be used to generate progression rewards that drive over-optimistic behaviors.

5 Conclusion and Future Work

In this paper, we introduce the State-To-Go (STG) Transformer, offline pretrained to predict latent state transitions in an adversarial way, for learning from visual observation to boost downstream reinforcement learning tasks. Our STG, tested across diverse Atari and Minecraft tasks, demonstrates superior robustness, sample efficiency, and performance compared to baseline approaches. We are optimistic that STG offers an effective solution in situations with plentiful video demonstrations, limited environment interactions, and where labeling action is expensive or infeasible.

In future work, STG is likely to benefit from more powerful large-scale vision foundation models to facilitate generalization across a broader range of related tasks. Besides, it can extend to a hierarchical framework where one-step predicted rewards can be employed for training low-level policies and multi-step rewards for the high-level policy, which is expected to improve performance and solve long-horizon tasks.

Acknowledgments and Disclosure of Funding

This work was supported by NSF China under grant 62250068. The authors would like to thank the anonymous reviewers for their valuable comments.

References

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.

Yuke Zhu, Ziyu Wang, Josh Merel, Andrei Rusu, Tom Erez, Serkan Cabi, Saran Tunyasuvunakool, János Kramár, Raia Hadsell, Nando de Freitas, et al. Reinforcement and imitation learning for diverse visuomotor skills. arXiv preprint arXiv:1802.09564 , 2018.

Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observation. arXiv preprint arXiv:1807.06158 , 2018.

Ashley Edwards, Himanshu Sahni, Yannick Schroecker, and Charles Isbell. Imitating latent policies from observation. In International Conference on Machine Learning (ICML) , 2019.

Nathan Gavenski, Juarez Monteiro, Roger Granada, Felipe Meneguzzi, and Rodrigo C Barros. Imitating unknown policies via exploration. arXiv preprint arXiv:2008.05660 , 2020.

Rahul Kidambi, Jonathan Chang, and Wen Sun. Mobile: Model-based imitation learning from observation alone. In Neural Information Processing Systems (NeurIPS) , 2021.

Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Off-policy imitation learning from observations. In Neural Information Processing Systems (NeurIPS) , 2020.

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.

Jake Bruce, Ankit Anand, Bogdan Mazoure, and Rob Fergus. Learning about progress from experts. In International Conference on Learning Representations (ICLR) , 2023.

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.

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.

Faraz Torabi, Garrett Warnell, and Peter Stone. Imitation learning from video by leveraging proprioception. In International Joint Conference on Artificial Intelligence (IJCAI) , 2019.

Haresh Karnan, Garrett Warnell, Faraz Torabi, and Peter Stone. Adversarial imitation learning from video using a state observer. In International Conference on Robotics and Automation (ICRA) , 2022.

Youngwoon Lee, Andrew Szot, Shao-Hua Sun, and Joseph J Lim. Generalizable imitation learning from observation via inferring goal proximity. In Neural Information Processing Systems (NeurIPS) , 2021.

Minghuan Liu, Tairan He, Weinan Zhang, Shuicheng Yan, and Zhongwen Xu. Visual imitation learning with patch rewards. arXiv preprint arXiv:2302.00965 , 2023.

Haoqi Yuan, Chi Zhang, Hongcheng Wang, Feiyang Xie, Penglin Cai, Hao Dong, and Zongqing Lu. Plan4mc: Skill reinforcement learning and planning for open-world minecraft tasks. arXiv preprint arXiv:2303.16563 , 2023.

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.

Deepak Pathak, Parsa Mahmoudieh, Guanghao Luo, Pulkit Agrawal, Dian Chen, Yide Shentu, Evan Shelhamer, Jitendra Malik, Alexei A Efros, and Trevor Darrell. Zero-shot visual imitation. In IEEE Computer Vision and Pattern Recognition (CVPR) Workshops , 2018.

Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In International Joint Conference on Artificial Intelligence (IJCAI) , 2018.

Bowen Baker, Ilge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. In Neural Information Processing Systems (NeurIPS) , 2022.

Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2011.

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Neural Information Processing Systems (NeurIPS) , 2016.

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.

Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards-just map them to actions. In Neural Information Processing Systems (NeurIPS) , 2019.

Alejandro Escontrela, Ademi Adeniji, Wilson Yan, Ajay Jain, Xue Bin Peng, Ken Goldberg, Youngwoon Lee, Danijar Hafner, and Pieter Abbeel. Video prediction models as rewards for reinforcement learning. arXiv preprint arXiv:2305.14343 , 2023.

Bradly C Stadie, Pieter Abbeel, and Ilya Sutskever. Third-person imitation learning. arXiv preprint arXiv:1703.01703 , 2017.

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning (ICML) , 2020.

Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning (ICML) , 2021.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Neural Information Processing Systems (NeurIPS) , 2017.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Neural Information Processing Systems (NeurIPS) , 2021.

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.

Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In International Conference on Machine Learning (ICML) , 2022.

Muning Wen, Jakub Kuba, Runji Lin, Weinan Zhang, Ying Wen, Jun Wang, and Yaodong Yang. Multi-agent reinforcement learning is a sequence modeling problem. In Neural Information Processing Systems (NeurIPS) , 2022.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training, 2018.

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.

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning (ICML) , 2017.

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 , 2023.

Samuel Cohen, Brandon Amos, Marc Peter Deisenroth, Mikael Henaff, Eugene Vinitsky, and Denis Yarats. Imitation learning from pixel observations for continuous control. In Neural Information Processing Systems (NeurIPS) Workshop on Deep Reinforcement Learning , 2021.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Xue Bin Peng, Ze Ma, Pieter Abbeel, Sergey Levine, and Angjoo Kanazawa. Amp: Adversarial motion priors for stylized physics-based character control. ACM Transactions on Graphics (ToG) , 40(4):1-20, 2021.

Chao Yang, Xiaojian Ma, Wenbing Huang, Fuchun Sun, Huaping Liu, Junzhou Huang, and Chuang Gan. Imitation learning from observations by minimizing inverse dynamics disagreement. In Neural Information Processing Systems (NeurIPS) , 2019.

Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning (ICML) , 2020.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning (ICML) , 2018.

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.

Ziluo Ding, Hao Luo, Ke Li, Junpeng Yue, Tiejun Huang, and Zongqing Lu. Clip4mc: An rl-friendly vision-language model for minecraft. arXiv preprint arXiv:2303.10571 , 2023.

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.

Junhyuk Oh, Yijie Guo, Satinder Singh, and Honglak Lee. Self-imitation learning. In International Conference on Machine Learning (ICML) , 2018.

A Algorithms

We present our algorithm sketches for STG Transformer offline pretraining and online reinforcement learning with intrinsic rewards respectively.

Algorithm 1 STG Transformer Offline Pretraining

Algorithm 2 Online Reinforcement Learning with Intrinsic Rewards

Input: pretrained E ξ , T σ , D ω , policy π θ , MDP M , intrinsic coefficient η .

2: while updating π

1: Initialize parametric policy π θ with random θ randomly and reset M .

3: Execute π θ and store the resulting n state transitions { ( s, s ′ ) } t + n t .

θ do ▷ policy improvement

4: Use E ξ to obtain n real latent transitions { ( e, e ′ ) } t + n t .

6: Use D ω to calculate intrinsic rewards: ∆ t + n t = { D ω ( e, ˆ e ′ ) } t + n t -{ D ω ( e, e ′ ) } t + n t .

5: Use T σ to obtain n predicted latent transitions { ( e, ˆ e ′ ) } t + n t .

7: Perform PPO update to improve π θ with respect to r i = -η ∆ .

8: end while

B Environment Details

B.1 Atari

We directly adopt the official default setting for Atari games. Please refer to https://www. gymlibrary.dev/environments/atari for more details.

B.2 Minecraft

Environment Settings

Table 1 outlines how we set up and initialize the environment for each harvest task.

Harvest Item

Initialized Tool

Biome

milk

empty bucket

plains

wool

shears

plains

tallgrass

shears

plains

sunflower

diamond shovel

sunflower plains

(a)

Plains

(b)

Sunflower Plains

Biomes. Our method is tested in two different biomes: plains and sunflower plains. Both the plains and sunflower plains offer a wider field of view. However, resources and targets are situated further away from the agent, which presents unique challenges. Figure 1a and 1b show the biomes of plains and sunflower plains respectively.

Observation Space. Despite MineDojo offering an extensive observation space, encompassing RGB frames, equipment, inventory, life statistics, damage sources, and more, we exclusively rely on the RGB information as our observation input.

Action Space. In Minecraft, the action space is an 8-dimensional multi-discrete space. Table 2 lists the descriptions of action space in the MineDojo simulation platform. At each step, the agent chooses one movement action (index 0 to 4) and one optional functional action (index 5) with corresponding parameters (index 6 and index 7).

C Offline Pretraining Details

Hyperparameters. Table 3 outlines the hyperparameters for offline pretraining in the first stage.

Network Structure. Different architectures for feature encoding are designed for different environments. In Atari, we stack four gray-scale images of shape (84,84) to form a 4-channel state and use the feature encoder architecture as shown in Figure 2a. In Minecraft, a 3-channel image of shape (160,256,3) is directly regarded as a single state, which is processed by a feature encoder with more convolutional layers and residual blocks to capture more complex features in the ever-changing Minecraft world. The detailed structure of the feature encoder for Minecraft is illustrated in Figure 2b. All discriminators, taking in two 512-dimension embeddings from the feature encoder, follow the MLP structure of FC(1024,512) → FC(512,256) → FC(256,128) → FC(128,64) → FC(64,32) → FC(32,1) with spectral normalization.

Representation Visualization. We draw inspiration from Grad-CAM [47] to visualize the saliency map of offline-pretrained feature encoder to assess the effectiveness and advantages of the representation of STG. Specifically, we compare the visualization results of STG and ELE in the Atari environment as illustrated in Figure 3. Each figure presents three rows corresponding to the features captured by the three layers of the convolutional layers, respectively. The saliency maps demonstrate that STG exhibits a particular focus more on local entities and dynamic scenarios and effectively ignores extraneous distractions. As a result, compared with ELE, STG shows greater proficiency in

Stride:4

Conv:8

×

8

Relu

(1,4,84,84)

Stride:2

Conv:4

×

4

Relu

(1,32,20,20)

Stride:1

Conv:3

×

3

Relu

(1,64,9,9)

(1,64,7,7)

(1,512)

&

Linear

(a)

Atari

(1,3,256,160)

Stride:4

Conv:8

×

8

LeakyRelu

BatchNorm

ResBlock

(1,32,128,80)

Stride:2

Conv:3

×

3

LeakyRelu

BatchNorm

ResBlock

(1,64, 64, 40)

Stride:2

Conv:3

×

3

LeakyRelu

BatchNorm

ResBlock

(1,64,32,20)

Conv:3

×

3

Stride:2

LeakyRelu

BatchNorm

(1,64,16,10)

Conv:3

×

3

Stride:2

LeakyRelu

BatchNorm

ResBlock

ResBlock

Conv:3

×

3

Stride:2

LeakyRelu

BatchNorm

(1,64,8,5)

(1,128,4,3)

Stride:2

Conv:3

×

3

LeakyRelu

BatchNorm

ResBlock

(1,512)

Conv(in,out,3,1)

Conv(out,out,3,1)

BatchNorm

Relu

BatchNorm

Conv(in,out,3,1)

BatchNorm

&

Linear

(b)

Minecraft

identifying information strongly correlated with state transitions, thereby generating higher-quality rewards for downstream reinforcement learning tasks.

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

(a)

Saliency maps of different CNN layers in Breakout

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

E91

5

三

5

E91

5

163

5

ArTnVSION

ATIVSION

ACTIMSION

ATiVSION

ACTIMSION

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

(c)

Saliency maps of different CNN layers in Qbert

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

0

40

80

120

160

0

30

60

90

120

150

180

210

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

40

80

120

160

0

30

60

90

120

150

180

210

ARRA

沈R民片

D3D

三

DEED

ARRR

三

三

XXK

AAAR

三

代虎片

(d)

Saliency maps of different CNN layers in Space Invaders

D RL Training Details

The general training hyperparameters of PPO for downstream RL tasks in the second stage are listed in Table 4.

Neither the discriminative reward from STG nor the progression reward from ELE is bounded. Therefore, it is reasonable to adjust certain hyperparameters to bring out the best performance of each algorithm in each task. In Table 5, the coefficient of intrinsic reward η ( η > 0) for different baselines is tuned to balance the value scale and GAE λ (0 < λ < 1) is tuned to adjust the impact of intrinsic rewards in different tasks.

The coefficients of STG* (noted as ηr i + νr ∗ ) in four Atari tasks are reported in Table 6.

Training Details. For Minecraft tasks, we adopt a hybrid approach utilizing both PPO [40] and SelfImitation Learning (SIL) [48]. Specifically, we store trajectories with high intrinsic rewards in a 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].

E Further Experiments

Intrinsic Reward Design. In Equation (9), we define our intrinsic reward r i as the difference between r guide and r base :

On the one hand, pretrained D ω clearly provides informative judgment r guide of transition quality during online interaction. On the other hand, the baseline reward r base , solely relying on the current state s t , serves as a baseline to normalize r i to a relatively lower level. In this section, we aim to investigate the necessity of incorporating r base .

To assess the significance of r base , we conducted experiments on the four Atari tasks utilizing only r guide as the intrinsic reward, which is similar to previous adversarial methods like GAIfO [3]. In order to bring the scale of r guide in line with r i , we employ running normalization and bound the values within the range of [ -1 , 1] to mitigate the negative influence of outliers. All other settings remain unchanged. We denote this ablation baseline as STG'.

As illustrated in Figure 4, r guide yields comparable final performance in Breakout and Space Invaders while failing to achieve satisfactory performance in Freeway and Qbert. In contrast, by leveraging r base , which provides a unique reference from expert datasets for each individual s t , we observe reduced variance and improved numerical stability compared to the running normalization trick that calculates batch means for normalization.

(a)

Breakout

(b)

Freeway

(c)

(d)

Space Invaders

300

STG'

STG

20

30000

STG'

STG

600

STG'

STG

200

15

STG'

20000

400

10

STG

100

5

10000

200

0

0

0

0.0

0.5

1.0

1.5

2.0

0

2

0.00

0.25

0.50

0.75

1.00

0.00

0.25

0.50

0.75

1.00

Transition

1e7

Transition

1e6

Transition

1e7

Transition

1e7

F Additional Ablations

Multi-Task STG. 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 we further assess the efficacy of multi-task adaptation of the STG Transformer in Minecraft. We tune the intrinsic coefficient η to be 5 for Breakout, Qbert, SpaceInvaders, 'milk a cow', 'harvest tallgras', 'pick a flower', and 10 for Freeway and 'gather wool'. STG-Multi results in Minecraft are illustrated in Figure 5. As all four tasks share a 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.

0

200

400

600

800

1000

Epoch

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

STG

STG-Multi

(a)

Pick a flower

0

200

400

600

800

1000

Epoch

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

STG

STG-Multi

(b)

Milk a cow

0

200

400

600

800

1000

Epoch

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

STG

STG-Multi

(c)

Harvest tallgrass

0

200

400

600

800

1000

Epoch

0

.

00

0

.

02

0

.

04

0

.

06

STG

STG-Multi

(d)

Gather wool

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 6 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 the effectiveness of the offline-pretrained Modules.

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

50

100

150

(a)

Breakout

0

2

4

Transition

×

10 6

220

240

260

(b)

Freeway

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

0

5000

10000

15000

(c)

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Transition

×

10 7

50

100

150

(d)

Space Invaders

Multi-Trajectories Visualization. As an extension to increase diversity, we additionally visualize more trajectories in SpaceInvaders in Figure 7. We randomly sample five 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. Furthermore, it can be observed that after 10M-step of 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.

STG

(a) random trajectory

(b) training 1kw steps

(c) expert datasets

STG-

(a) random trajectory

(b) training 1kw steps

(c) expert datasets

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n### 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\nWEAKNESSES:\n### 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nI think this paper appropriately mentioned the limitations and broader impacts.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nIn 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).\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThis 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nStrengths\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\nWEAKNESSES:\nWeaknesses\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\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nPaper 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.\n\nSTRENGTHS:\n+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\nWEAKNESSES:\nHaving 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\nQUESTIONS:\nSee [Weaknesses]\n\nLIMITATIONS:\nI 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?", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-11"}