Title: Pre-training Contextualized World Models with In-the-wild Videos for Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Model-based reinforcement learning

world model

pre-training
Abstract: Unsupervised pre-training methods utilizing large and diverse datasets have achieved tremendous success across a range of domains. Recent work has investigated such unsupervised pre-training methods for model-based reinforcement learning (MBRL) but is limited to domain-specific or simulated data. In this paper, we study the problem of pre-training world models with abundant in-the-wild videos for efficient learning of downstream visual control tasks. However, in-the-wild videos are complicated with various contextual factors, such as intricate backgrounds and textured appearance, which precludes a world model from extracting shared world knowledge to generalize better. To tackle this issue, we introduce Contextualized World Models (ContextWM) that explicitly separate context and dynamics modeling to overcome the complexity and diversity of in-the-wild videos and facilitate knowledge transfer between distinct scenes. Specifically, a contextualized extension of the latent dynamics model is elaborately realized by incorporating a context encoder to retain contextual information and empower the image decoder, which encourages the latent dynamics model to concentrate on essential temporal variations. Our experiments show that in-the-wild video pre-training equipped with ContextWM can significantly improve the sample efficiency of MBRL in various domains, including robotic manipulation, locomotion, and autonomous driving. Code is available at this repository: https://github.com/thuml/ContextWM.

Pre-training Contextualized World Models with In-the-wild Videos for Reinforcement Learning

Abstract

Unsupervised pre-training methods utilizing large and diverse datasets have achieved tremendous success across a range of domains. Recent work has investigated such unsupervised pre-training methods for model-based reinforcement learning (MBRL) but is limited to domain-specific or simulated data. In this paper, we study the problem of pre-training world models with abundant in-the-wild videos for efficient learning of downstream visual control tasks. However, inthe-wild videos are complicated with various contextual factors, such as intricate backgrounds and textured appearance, which precludes a world model from extracting shared world knowledge to generalize better. To tackle this issue, we introduce Contextualized World Models (ContextWM) that explicitly separate context and dynamics modeling to overcome the complexity and diversity of in-the-wild videos and facilitate knowledge transfer between distinct scenes. Specifically, a contextualized extension of the latent dynamics model is elaborately realized by incorporating a context encoder to retain contextual information and empower the image decoder, which encourages the latent dynamics model to concentrate on essential temporal variations. Our experiments show that in-the-wild video pre-training equipped with ContextWM can significantly improve the sample efficiency of MBRL in various domains, including robotic manipulation, locomotion, and autonomous driving. Code is available at this repository: https://github.com/thuml/ContextWM .

1 Introduction

Model-based reinforcement learning (MBRL) holds the promise of sample-efficient learning for visual control. Typically, a world model [20, 40] is learned to approximate state transitions and control signals of the environment to generate imaginary trajectories for planning [22] or behavior learning [71]. In the wake of revolutionary advances in deep learning, world models have been realized as action-conditional video prediction models [34] or latent dynamics models [22, 21]. However, given the expressiveness and complexity of deep neural networks, the sample efficiency of MBRL can still be limited by the failure to learn an accurate and generalizable world model efficiently.

Pre-training and fine-tuning paradigm has been highly successful in computer vision [82, 26] and natural language processing [57, 13] to fast adapt pre-trained representations for downstream tasks, while learning tabula rasa is still dominant in MBRL. Recent work has taken the first step towards pre-training a world model, named Action-free Pre-training from Videos (APV) [68]. However, it has been conducted by pre-training on domain-specific and carefully simulated video datasets, rather than abundant in-the-wild video data [17, 18, 30, 83]. Prior attempts of leveraging real-world video data result in either underfitting for video pre-training or negligible benefits for downstream visual control tasks [68]. Against this backdrop, we naturally ask the following question:

Can world models pre-trained on diverse in-the-wild videos benefit sample-efficient learning of downstream visual control tasks?

World

Model

Human

Object

Interaction

Videos

Real-world

Driving

Videos

Human

Motion

Videos

In-the-wild Video Data

Pre-training

Visual Control Tasks

Fine-tuning

Robotic Manipulation

Robotic Locomotion

Autonomous Driving

C

C

C

It opens up a way to utilize the vast amount of videos available on the Internet and thereby fully release the potential of this methodology by applying it at scale. Large and diverse data spanning various scenes and tasks can provide world knowledge that is widely generalizable and applicable to a variety of downstream tasks. For instance, as depicted in Figure 1, world models for robotic manipulation tasks can probably benefit not only from videos of humans interacting with target objects in diverse ways but also from motion videos that embody rigorous physical rules. Analogous to APV [68], we brand our paradigm as In-the-wild Pre-training from Videos (IPV).

However, in-the-wild videos are of inherent high complexity with various factors, such as intricate backgrounds, appearance, and shapes, as well as complicated dynamics. Coarse and entangled modeling of context and dynamics can waste a significant amount of model capacity on modeling low-level visual details of what is there and prevent world models from capturing essential shared knowledge of what is happening . Biological studies reveal that in natural visual systems, ∼ 80% retinal ganglion cells are P-type operating on spatial detail and color, while ∼ 20% are M-type operating on temporal changes [28, 44]. Partially inspired by this, we suggest that leveraging IPV for world models calls for not only appropriate data and model scaling [35] but more essentially, dedicated design for context and dynamics modeling.

In this paper, we present the Contextualized World Model (ContextWM), a latent dynamics model with visual observations that explicitly separates context and dynamics modeling to facilitate knowledge transfer of semantically similar dynamics between visually diverse scenes. Concretely, a contextualized latent dynamics model is derived to learn with variational lower bound [33, 38] and elaborately realized by incorporating a context encoder to retain contextual information and a parsimonious latent dynamics model to concentrate on essential temporal variations. Moreover, dual reward predictors are introduced to enable simultaneously enhancing task-relevant representation learning as well as learning an exploratory behavior. The main contributions of this work are three-fold:

From a data-centric view, we systematically study the paradigm of pre-training world models with in-the-wild videos for sample-efficient learning of downstream tasks, which is much more accessible and thus makes it possible to build a general-purpose world model.

We propose Contextualized World Models (ContextWM), which explicitly model both the context and dynamics to handle complicated in-the-wild videos during pre-training and also encourage task-relevant representation learning when fine-tuning with MBRL.

Our experiments show that equipped with our ContextWM, in-the-wild video pre-training can significantly improve the sample efficiency of MBRL on various domains.

2 Related Work

Pre-training in RL. Three categories of pre-training exist in RL: unsupervised online pre-training, offline pre-training, and visual representation pre-training. The first two are both domain-specific, which learn initial behaviors, primitive skills, or domain-specific representations by either online environment interaction [55, 39] or offline demonstrations [41, 3]. Recent work has explored general-

purpose visual pre-training using in-the-wild datasets [69, 36, 77, 53, 86, 49, 45] to accelerate policy learning in model-free manners. Plan2Explore [66], as a model-based method, pre-trains its world model and exploratory policy simultaneously by online interaction. Ye et al. [81] and Xu et al. [78] build upon EfficientZero [80] and pre-train their world models using offline experience datasets across Atari games. To the best of our knowledge, APV [68] is the first to pre-train world models with datasets from different domains, though the datasets are manually simulated demonstrations by scripted policies across tasks from RLBench [32], which still lacks diversity and scale.

Visual control with in-the-wild videos. While several existing works leverage demonstration videos for learning visual control [43, 84, 62], rare efforts have been made to utilize off-the-shelf video datasets from the Internet. Shao et al. [70] and Chen et al. [9] learn reward functions using videos from the Something-Something dataset [17]. R3M [49] leverages diverse Ego4D dataset [18] to learn reusable visual representations by time contrastive learning. Other applications with such videos include video games [7, 6], visual navigation [8], and autonomous driving [88]. We instead pre-train world models by video prediction on in-the-wild video datasets.

World models for visual RL. World models that explicitly model the state transitions and reward signals are widely utilized to boost sample efficiency in visual RL. A straightforward method is to learn an action-conditioned video prediction model [50, 34] to generate imaginary trajectories. Ha and Schmidhuber [20] propose to first use a variational autoencoder to compress visual observation into latent vectors and then use a recurrent neural network to predict transition on compressed representations. Dreamer [22, 21, 23], a series of such latent dynamics models learning via image reconstruction, have demonstrated their effectiveness for both video games and visual robot control. There also have been other methods that learn latent representations by contrastive learning [51, 11] or value prediction [63, 25]. Recent works also use Transformers [75] as visual encoders [67] or dynamics models [10, 47, 58] that make world models much more scalable and can be complementary with our effort to pre-train world models with a vast amount of diverse video data.

Concurrent to our work, SWIM [46] also pre-trains a world model from large-scale human videos and then transfers it to robotic manipulation tasks in the real world. However, it designs a structured action space via visual affordances to connect human videos with robot tasks, incorporating a strong inductive bias and thus remarkably limiting the range of pre-training videos and downstream tasks.

3 Background

Problem formulation. We formulate a visual control task as a partially observable Markov decision process (POMDP), defined as a tuple ( O , A , p, r, γ ) . O is the observation space, A is the action space, p ( o t | o <t , a <t ) is the transition dynamics, r ( o ≤ t , a <t ) is the reward function, and γ is the discount factor. The goal of MBRL is to learn an agent π that maximizes the expected cumulative rewards E p,π [ ∑ T t =1 γ t -1 r t ] , with a learned world model (ˆ p, ˆ r ) approximating the unknown environment. In our pre-training and fine-tuning paradigm, we also have access to an in-the-wild video dataset D = { ( o t ) T t =1 } without actions and rewards to pre-train the world model.

Dreamer. Dreamer [21, 23, 24] is a visual model-based RL method where the world model is formulated as a latent dynamics model [16, 22] with the following four components:

The representation model, also known as posterior of z t , approximates latent state z t from previous state z t -1 , previous action a t -1 and current observation o t , while the transition model, also known as prior of z t , predicts it directly from z t -1 and a t -1 . The overall models are jointly learned by minimizing the negative variational lower bound (ELBO) [33, 38]:

Behavior learning of Dreamer can be conducted by actor-critic learning purely on imaginary latent trajectories, for which we refer readers to Hafner et al. [21] for details.

Action-free Pre-training from Videos (APV). To leverage action-free video data, APV [68] first pre-trains an action-free variant of the latent dynamics model (with representation q θ ( z t | z t -1 , o t ) , transition p θ (ˆ z t | z t -1 ) and image decoder p θ (ˆ o t | z t ) ) as a video prediction model, which drops action conditions and the reward predictor from Eq. (1) and Eq. (2). When fine-tuned for MBRL, APV stacks an action-conditional model (with representation q ϕ ( s t | s t -1 , a t -1 , z t ) and transition p ϕ (ˆ s t | s t -1 , a t -1 ) ) on top of the action-free one. Action-free dynamics models and image decoders are initialized with pre-trained weights θ from video pre-training. Furthermore, to utilize pre-trained representations for better exploration, APV introduces a video-based intrinsic bonus r int t measuring the diversity of visited trajectories by random projection and nearest neighbors. The overall models during fine-tuning are optimized by minimizing the following objective:

4 Contextualized World Models

In this section, we present Contextualized World Models (ContextWM), a framework for both actionfree video prediction and visual model-based RL with explicit modeling of context and dynamics from observations. Our method introduces: (i) a contextualized extension of the latent dynamics model (see Section 4.1), (ii) a concrete implementation of contextualized latent dynamics models tailored for visual control (see Section 4.2). We provide the overview and pseudo-code of our ContextWM in Figure 3 and Appendix A, respectively.

4.1 Contextualized Latent Dynamics Models

While latent dynamics models have been successfully applied for simple and synthetic scenes, in-the-wild images and videos are naturalistic and complex. Our intuition is that there are two groups of information in the sequence of observations, namely time-invariant context and timedependent dynamics [76, 12, 42, 15]. The context encodes static information about objects and concepts in the scene, such as texture, shapes, and colors, while the dynamics encode the temporal transitions of the concepts, such as positions, layouts, and motions. To overcome the complexity of in-the-wild videos, it is necessary to explicitly represent the complicated contextual information, thus shared knowledge with respect to dynamics can be learned and generalized across distinct scenes.

We thus propose a contextualized extension of latent dynamics models [16, 22], extending Eq. (1). As shown in

⋯

⋯

a

1

a

2

a

3

a

T

r

1

r

2

r

T

r

3

z

1

z

2

z

3

z

T

o

1

o

2

o

3

o

T

c

⋯

⋯

Figure 2, the probabilistic model generates the observation o t conditioned not only on the current latent dynamics z t but also on a context variable c that can include rich information (solid lines in Figure 2), while the variational inference model approximates the posterior of the latent dynamics z t conditioned on observations o t and previous states z t -1 , without referring to the context (dotted lines in Figure 2). The proposed design allows contextualized image decoders to reconstruct diverse and complex observations using rich contextual information, expanding their ability beyond the expressiveness of latent dynamics variables z t . On the other hand, through the variational information bottleneck [5, 1], latent dynamics variables z t are encouraged to selectively exclude contextual information and only capture essential temporal variations that are not included in the context c . Furthermore, the design of context-unaware latent dynamics inference encourages representation and transition models to learn this temporal information within a high-level semantic space, in contrast to merely differentiating with the context frame in the low-level visual space, facilitating the acquisition of representations more transferable and robust to intricate contexts.

<latexi sh

1_b

64="pPcruK3yZWI

Ew0

/

2B

8A

>

V

NS

U

q

9L

n

R

D

k

7

G

g

vH

T

5

+

Q

M

f

m

d

o

C

OJ

z

X

Y

F

j

<latexi sh

1_b

64="

o3/

LA

u

I

GmCJ

D

nk

T

>

B

c

V

SgN

EOy

r

fUY9

P

K

2

Q

8

5

H

v

d

Z

w

F

p

R

0

W

j

+

q

M

X

<latexi sh

1_b

64="

u

oY

jm

Z

DpSPfB9

V8

>A

X

c

L

gN

EOy

r

U

K

2

Q

I5

HJ

n

d

/7C

w

Fv

3

k

G

q

T

W

0

R

+z

M

<latexi sh

1_b

64="mDoM

8dw ku

z

LJ/

G2+g fQ

>A

B

3

c

V

S

N

EOy

r

UY9

P

K

I

Zp

W

5

F

q

0

X

H

C

v

R

T

n7

j

<latexi sh

1_b

64="vj

0F

2D

V/Op

CUW

k

LI

>A

B

n

c

TgJ

E

z

+

9

H

R

ryP

M8

Q

m5

Z

w

d

X

N

o

q

7

K

3

S

u

G

f

Y

<latexi sh

1_b

64="3P

L

wS

IDf

R/d

o

9CZz

>A

B8X

c

V

gN

EO

G

Mr

H

Y

j0

W

y

nJ

k

T

K

q

u

7

2

F

U

v

p

m

5

+

Q

<latexi sh

1_b

64="mpKkB/U5

nDu

Y

g

f

OH

F

>A

c

V

LS

M

3

W

r

JN

y

z

RPq

d

GZ

E

Q

C

2

9

v

w

o

X

7

I

0

T

j

8

<latexi sh

1_b

64="V

UuN

9Jy

o

Y

p0

M

>A

B

n

c

DLTg

EO

F

+I

jH

R

Z

PGK

w

2d

3

58

7

rv

W

X

Q

/

S

m

k

C

f

<latexi sh

1_b

64="ZP

S+rgVF

NE

ABw

W

dL

/cY

>

8X

D

Oy fU

9

o

K

z

I

Hp

m

n

5C

q3

j

ku

R

v

Q

7

T

G

M

J

0

2

<latexi sh

1_b

64="

9K

7gYSZTp

RqP/B

2

X

y

>

3

c

VDL

N

F

r

fU

du ov

+

0

8w

E

G

M

I

j

H

5

O

W

z

Q

k

C

n

J

<latexi sh

1_b

64="pPcruK3yZWI

Ew0

/

2B

8A

>

V

NS

U

q

9L

n

R

D

k

7

G

g

vH

T

5

+

Q

M

f

m

d

o

C

OJ

z

X

Y

F

j

<latexi sh

1_b

64="3+gC72

uzY

E

k

JA

oR9

>

B8X

c

VDLS

N

Oy

r

fU

P

K

Q

I5

H

m

n

d

p

/

w

Fv

G

jq

T

W

Z

0

M

<latexi sh

1_b

64="

o3/

LA

u

I

GmCJ

D

nk

T

>

B

c

V

SgN

EOy

r

fUY9

P

K

2

Q

8

5

H

v

d

Z

w

F

p

R

0

W

j

+

q

M

X

<latexi sh

1_b

64="IrjPuS7

J3zQ

k

pg

Z+

/

>A

B

n

c

VDL

N

EOy fUY9

o

K

2

8

5

H

m

d

w

Fv

G

q

T

W

0

R

C

X

M

<latexi sh

1_b

64="

u

oY

jm

Z

DpSPfB9

V8

>A

X

c

L

gN

EOy

r

U

K

2

Q

I5

HJ

n

d

/7C

w

Fv

3

k

G

q

T

W

0

R

+z

M

<latexi sh

1_b

64="mDoM

8dw ku

z

LJ/

G2+g fQ

>A

B

3

c

V

S

N

EOy

r

UY9

P

K

I

Zp

W

5

F

q

0

X

H

C

v

R

T

n7

j

<latexi sh

1_b

64="

K/7pg

QC

o

XPFYk

r

>A

B8

c

DLS

N

EOz

G

M

H

d3

j0

W

nJ

Z

R

w

I

T

q

u

U

v

9

m

+

f

5

<latexi sh

1_b

64="wVNPT

dnOF

0

H

j

39

cp zBY

>A

DLSg

E

y

r

fU

o

K

2

Q

8

I5

J

m

7

v

k

G

q

W

Z

u

R

C

+

X

M

/

<latexi sh

1_b

64="vj

0F

2D

V/Op

CUW

k

LI

>A

B

n

c

TgJ

E

z

+

9

H

R

ryP

M8

Q

m5

Z

w

d

X

N

o

q

7

K

3

S

u

G

f

Y

<latexi sh

1_b

64="L3+A

SV

voFdD

E

QY

9HU

0

>

B

n

c

gN

Oz

G

Mr

8

j

Cy

J

k

ZmRX

/w

I

T

K

q

2

7

W

p

f

P

u

5

<latexi sh

1_b

64="3P

L

wS

IDf

R/d

o

9CZz

>A

B8X

c

V

gN

EO

G

Mr

H

Y

j0

W

y

nJ

k

T

K

q

u

7

2

F

U

v

p

m

5

+

Q

<latexi sh

1_b

64="mpKkB/U5

nDu

Y

g

f

OH

F

>A

c

V

LS

M

3

W

r

JN

y

z

RPq

d

GZ

E

Q

C

2

9

v

w

o

X

7

I

0

T

j

8

<latexi sh

1_b

64="p/q+X

U

R

wYT9

Sk2NFdJocg

>A

B8

VDL

EOy

r

f

P

K

z

I

H

m

Zn

W

5C

3

j

u

v

Q

7

G

M

0

<latexi sh

1_b

64="q9AIS

U

urW

MT

X

P+

c7z

>

B

n

VDL

gN

EO

G

H

Y

8

d3

j0

Cy

J

k

/w oK

v

2

p

Q

f

F

5

<latexi sh

1_b

64="V

UuN

9Jy

o

Y

p0

M

>A

B

n

c

DLTg

EO

F

+I

jH

R

Z

PGK

w

2d

3

58

7

rv

W

X

Q

/

S

m

k

C

f

<latexi sh

1_b

64="KWYXZ

AFrc32P

o

B

S/fQ

>

n

VDL

gN

EOy

U

9

I

5

8

q

k

u

R

C

v

d

7

T

G

M

+0

J

<latexi sh

1_b

64="ZP

S+rgVF

NE

ABw

W

dL

/cY

>

8X

D

Oy fU

9

o

K

z

I

Hp

m

n

5C

q3

j

ku

R

v

Q

7

T

G

M

J

0

2

<latexi sh

1_b

64="

9K

7gYSZTp

RqP/B

2

X

y

>

3

c

VDL

N

F

r

fU

du ov

+

0

8w

E

G

M

I

j

H

5

O

W

z

Q

k

C

n

J

<latexi sh

1_b

64="po

HR

/fOI

9XcgN

Cm0F

TDQ

>A

B8

V

LS

E

y

r

UY

P

K

2

5

J

n

d

7

w

v

3

k

G

jq

W

Z

u

M

+

z

<latexi sh

1_b

64="Mu

+YOA02U

vC

7y

Xjn3oF

fQ

>

B

c

VDLSgN

E

r

9

P

K

z

I

Z

W

5

8

q

/k

R

d

T

G

J

<latexi sh

1_b

64="kE7fK

0

Jn

A

Tgu

D

H

+

>

B8X

c

V

LS

N

Oy

r

UY9

oP

2

Q

I5

m

d

p

/

C

w

Fv

3

G

jq

W

Z

R

M

z

<latexi sh

1_b

64="YSIW7

D

Ug

A

+yQ8jm

k

z

M

>

B

X

c

V

L

N

EO

r

f

9

oP

K

2

5

HJ

n

d

p

/

C

w

Fv

3

G

q

T

Z

0

u

R

Multi-scale

Context

Cross Attention

o

1

o

2

o

3

z

1

z

2

z

3

ˆ

o

1

ˆ

o

2

ˆ

o

3

o

˜

t

(a) Action-free video pre-training

Action-

conditioned

Action-free

Multi-scale

Context

z

1

z

2

z

3

o

1

o

2

o

3

o

˜

t

ˆ

o

1

ˆ

o

2

ˆ

o

3

ˆ

o

4

ˆ

r

1

ˆ

r

2

ˆ

r

3

s

1

s

2

s

3

a

1

a

2

a

3

ˆ

s

4

ˆ

r

4

(b) Action-conditioned fine-tuning with MBRL

As derived in Appendix B, the overall model can be learned with ELBO of conditional log probability ln p θ ( o 1: T , r 1: T | a 1: T , c ) , without the need to model the complex distribution of contexts p ( c ) :

4.2 Contextualized World Model Architectures

We then introduce a concrete implementation of contextualized latent dynamics models, specifically tailored for visual control. Our approach builds upon Dreamer [21], which utilizes a Recurrent State Space Model (RSSM) [22]. To enable action-free pre-training from videos and action-conditioned fine-tuning for MBRL, we also incorporate the stacked latent model from APV [68] into our design. Notably, our work distinguishes itself from previous research through the elaborate design of contextualized components and a dual reward predictor structure that enhances task-relevant representation learning. An overview of the overall architecture is illustrated in Figure 3.

Context formulation. There exist various choices of contextual information to be conditioned on, including text descriptions, pre-trained representations, semantic maps, or more sophisticated structured data [59]. As our focus is on end-to-end visual control, we choose the simplest one solely based on visual observations: we consider the context c as a single frame of observation randomly sampled from the trajectory segment o 1: T : c . = o ˜ t , ˜ t ∼ Uniform { 1 , 2 , . . . , T } . It is assumed that contextual information lies equally in each frame, and by randomly selecting a context frame, our context encoder should learn to be robust to temporal variations.

Multi-scale cross-attention conditioning. Given the context in the form of a frame of the same size as the reconstruction ˆ o t , a U-Net [60] architecture would be one of the most appropriate choices to empower a contextualized image decoder p θ (ˆ o t | z t , c ) , as its multi-scale shortcuts help to propagate the information directly from context to reconstruction. Moreover, in our case, the image decoder learns by also conditioning on latent dynamics z t , as shown in Figure 3. While a conventional U-Net architecture incorporates shortcut features directly by concatenation or summation and thus forces a spatial alignment between them, temporal variations such as motions or deformations cannot be neglected in our case. Inspired by recent advances in generative models [59, 4], we augment the decoder feature X ∈ R c × h × w with the context feature Z ∈ R c × h × w by a cross-attention mechanism

[75], shown as follows (we refer to Appendix C.3 for details):

Dual reward predictors. For simplicity of behavior learning, we only feed the latent variable s t into the actor π ( a t | s t ) and critic v ( s t ) , thus introducing no extra computational costs when determining actions. However, this raises the question of whether the latent variable s t contains sufficient taskrelevant information for visual control and avoids taking shortcuts through the U-Net skip connections. For example, in robotics manipulation tasks, the positions of static target objects may not be captured by the time-dependent s t , but by the time-invariant context c . We remark that reward predictors p ϕ ( r t | s t ) play a crucial role in compelling the latent variable s t to encode task-relevant information, as the reward itself defines the task to be accomplished. Nevertheless, for hard-exploration tasks such as Meta-world [85], the video-based intrinsic bonus [68] may distort the exact reward regressed for behavioral learning and constantly drift during the training process [61], making it difficult for the latent dynamics model to capture task-relevant signals. Therefore, we propose a dual reward predictor structure comprising a behavioral reward predictor that regresses the exploratory reward r t + λr int t for behavior learning, and additionally, a representative reward predictor that regresses the pure reward r t to enhance task-relevant representation learning [31].

Overall objective. The model parameters of ContextWM can be jointly optimized as follows:

When pre-trained from in-the-wild videos, ContextWM can be optimized by minimizing an objective that drops action-conditioned dynamics and reward predictors from Eq. (6), similar to APV [68]. For behavior learning, we adopt the actor-critic learning scheme of DreamerV2 [23]. For a complete description of pre-training and fine-tuning of ContextWM for MBRL, we refer to Appendix A.

5 Experiments

We conduct our experiments on various domains to evaluate In-the-wild Pre-training from Videos (IPV) with Contextualized World Models (ContextWM), in contrast to plain world models (WM) used by DreamerV2 and APV 2 . Our experiments investigate the following questions:

Can IPV improve the sample efficiency of MBRL?

How does ContextWM compare to a plain WM quantitatively and qualitatively?

What is the contribution of each of the proposed techniques in ContextWM?

How do videos from different domains or of different amounts affect IPV with ContextWM?

5.1 Experimental Setup

Visual control tasks. As shown in Figure 4, we formulate our experiments on various visual control domains. Meta-world [85] is a benchmark of 50 distinct robotic manipulation tasks and we use the same six tasks as APV [68]. DMC Remastered [19] is a challenging extension of the widely used robotic locomotion benchmark, DeepMind Control Suite [73], by expanding a complicated graphical variety. We also conduct an autonomous driving task using the CARLA simulator [14], where an agent needs to drive as far as possible along Town04's highway without collision in 1000 timesteps. See Appendix C.2 for details on visual control tasks.

1

Pre-training datasets. We utilize multiple in-the-wild video datasets which can potentially benefit visual control. Something-Something-v2 (SSv2) [17] dataset contains 193K videos of people interacting with objects. Human3.6M [30] dataset contains videos of human poses over 3M frames under 4 different viewpoints. YouTube Driving [88] dataset collects 134 real-world driving videos, over 120 hours, with various weather conditions and regions. We also merge the three datasets to construct an assembled dataset for the purpose of pre-training a general world model. See Figure 1 for example video frames and Appendix C.1 for details on data preprocessing.

Implementation details. For the newly introduced hyperparameter, we use β r = 1 . 0 in Eq. (6) for all tasks. To ensure a sufficient capacity for both plain WM and ContextWM pre-trained on diverse in-the-wild videos, we implement visual encoders and decoders as 13-layer ResNets [27]. Unless otherwise specified, we use the same hyperparameters with APV. See Appendix C.3 for more details.

Evaluation protocols. Following Agarwal et al. [2] and APV, we conduct 8 individual runs for each task and report interquartile mean with bootstrap confidence interval (CI) for individual tasks and with stratified bootstrap CI for aggregate results.

5.2 Meta-world Experiments

SSv2 pre-trained results. Figure 5a shows the learning curves on six robotic manipulation tasks from Meta-world. We observe that in-the-wild pre-training on videos from SSv2 dataset consistently improves the sample efficiency and final performance upon the DreamerV2 baseline. Moreover, our proposed ContextWM surpasses its plain counterpart in terms of sample efficiency on five of the six tasks. Notably, the dial-turn task proves challenging, as neither method is able to solve it. These results demonstrate that our method of separating context and dynamics modeling during pre-training and fine-tuning can facilitate knowledge transfer from in-the-wild videos to downstream tasks.

Ablation study. We first investigate the contribution of pre-training in our framework and report the aggregate performance with and without pre-training at the top of Figure 5b. Our results show that a plain WM seldom benefits from pre-training, indicating that the performance gain of IPV with plain WM is primarily due to the intrinsic exploration bonus. This supports our motivation that the complex contexts of in-the-wild videos can hinder knowledge transfer. In contrast, ContextWM significantly improves its performance with the aid of video pre-training. Additionally, we evaluate the contribution of the proposed techniques in ContextWM, as shown at the bottom of Figure 5b. We experiment with replacing the cross-attention mechanism (Eq. (5)) with simple concatenation or removing the dual reward predictor structure. Our results demonstrate that all these techniques contribute to the performance of ContextWM, as both variants outperform the plain WM.

Effects of dataset size. To investigate the effects of pre-training dataset size, we subsample 1.5k and 15k videos from SSv2 dataset to pre-train ContextWM. Figure 7a illustrates the performance of ContextWM with varying pre-training dataset sizes. We find that pre-training with only a small subset of in-the-wild videos can almost match the performance of pre-training with the full data. This is probably because, despite the diversity of contexts, SSv2 dataset still lacks diversity in dynamics patterns as there are only 174 classes of human-object interaction scenarios, which can be learned by proper models with only a small amount of data. It would be interesting for future work to explore whether there is a favorable scaling property w.r.t. pre-training dataset size when pre-training with more sophisticated video datasets with our in-the-wild video pre-training paradigm.

Effects of dataset domain. In Figure 7b, we assess pre-training on various video datasets, including RLBench [32] videos curated by APV [68] and our assembled data of three in-the-wild datasets. We observe that while pre-training always helps, pre-training from a more similar domain, RLBench, outperforms that from SSv2. Nevertheless, simulated videos lack diversity and scale and make it difficult to learn world models that are broadly applicable. Fortunately, in-the-wild video pre-training

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Lever Pull

0

5

10

15

20

25

Environment Steps (

×

)

Drawer Open

0

5

10

15

20

25

Environment Steps (

×

)

Door Lock

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Button Press Topdown Wall

0

5

10

15

20

25

Environment Steps (

×

)

Reach

DreamerV2

IPV w/ Plain WM

IPV w/ ContextWM

0

5

10

15

20

25

Environment Steps (

×

)

(a) Learning curves

(b) Ablation study

can continuously improve its performance with a more diverse assembled dataset, suggesting a promising scalable alternative for domain-specific pre-training.

5.3 DMControl Remastered Experiments

In the top row of Figure 6a, we present the learning curves of IPV from SSv2 dataset and DreamerV2 on the DMC Remastered locomotion tasks. Remarkably, we observe that pre-training with the SSv2 dataset can significantly enhance the performance of ContextWM, even with a large domain gap between pre-training and fine-tuning. This finding suggests that ContextWM effectively transfers shared knowledge of separately modeling context and dynamics. In contrast, a plain WM also benefits from pre-training, but it still struggles to solve certain tasks, such as hopper stand. These results also suggest that our ContextWM could be a valuable practice for situations where visual generalization is critical, as ContextWM trained from scratch also presents a competitive performance in Figure 6b.

Effects of dataset domain. Figure 7b demonstrates the performance of pre-training on human motion videos from the Human3.6M dataset. However, we observe a negative transfer, as pre-training from Human3.6M leads to inferior performance compared to training from scratch. We argue that the Human3.6M dataset is collected in the laboratory environment, rather than truly in-the-wild . These results support our motivation that pre-training data lacking in diversity can hardly help learn world models that are generally beneficial. Additionally, we experiment with pre-training on the assembled data of three video datasets but find no significant improvement over SSv2 pre-training.

5.4 CARLA Experiments

The bottom row of Figure 6a displays the learning curves of IPV from the SSv2 dataset and DreamerV2 on the CARLA driving task under different weather and sunlight conditions. Similar to the DMC Remastered tasks, we find that MBRL benefits from pre-training on the SSv2 dataset, despite a significant domain gap. We observe that in almost all weather conditions, IPV with ContextWM learns faster than plain WM at the early stages of training and can also outperform plain WM in terms of the final performance. Minor superiority of ContextWM over plain WM also suggests that in autonomous driving scenarios, a single frame may not contain sufficient contextual information, and more sophisticated formulations of contextual information may further enhance performance.

Effects of dataset domain. We then assess the performance of ContextWM on CARLA by pretraining it from alternative domains. As shown in Figure 7b, we observe that pre-training from

(a) Learning curves

(b) Effects of pre-training

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Meta-world

w/ full SSv2 videos

w/ 15k SSv2 videos

w/ 1.5k SSv2 videos

w/o Pre-train

(a) Effects of dataset size

(b) Effects of dataset domain

the YouTube Driving dataset or a combination of three in-the-wild datasets can both improve upon learning from scratch. However, neither can significantly surpass the performance achieved by pre-training from SSv2, despite a narrower domain gap between YouTube Driving and CARLA. We conjecture that this could be attributed to the higher complexity of the YouTube Driving dataset in comparison to SSv2. It would be promising to explore the potential of scaling ContextWM further to capture valuable knowledge from more complex videos.

5.5 Qualitative Analysis

Video prediction. We visually investigate the future frames predicted by a plain WM and ContextWM on SSv2 in Figure 8a. We find that our model effectively captures the shape and motion of the object, while the plain one fails. Moreover, we also observe that cross-attentions in our model (Eq. (5)) successfully attend to varying spatial positions of the context to facilitate the reconstruction of visual details. This shows our model works with better modeling of context and dynamics.

Video representations. Wesample video clips of length 25 with two distinct labels ( push something from right to left and from left to right ) and visualize the averaged model states of the sampled videos using t-SNE [74] in Figure 8b. Note that we do not utilize any labels of the videos in pre-training. Video representations of plain WM may entangle with extra contextual information and thus are not sufficiently discriminative to object motions. However, ContextWM which separately models context and dynamics can provide representations well distributed according to different directions of motion.

truth

Predicted

(ContextWM)

Context

Attention

Predicted

(Plain WM)

(a) Video prediction

(b) Video representations

truth

IPV w/

ContextWM

Context

Shift

IPV w/

Plain WM

Reconstruction

Context

Open-loop Prediction

(c) Compositional decoding

Compisitional decoding. We conduct a compositional decoding analysis by sampling a random frame from another trajectory to replace the original context and leaving dynamics the same. As shown in Figure 8c, ContextWM correctly combines the new context with the original dynamics information. These results show that our model, fine-tuned on the DMCR domain, has successfully learned disentangled representations of contexts and dynamics. In contrast, the plain WM suffers from learning entangled representations and thus makes poor predictions about the transitions.

6 Discussion

This paper presents Contextualized World Models (ContextWM), a framework for both action-free video prediction and visual model-based RL. We apply ContextWM to the paradigm of In-the-wild Pre-training from Videos (IPV), followed by fine-tuning on downstream tasks to boost learning efficiency. Experiments demonstrate the effectiveness of our method in solving a variety of visual control tasks from Meta-world, DMC Remastered, and CARLA. Our work highlights not only the benefits of leveraging abundant in-the-wild videos but also the importance of innovative world model design that facilitates knowledge transfer and scalable learning.

Limitations and future work. One limitation of our current method is that a randomly selected single frame may not sufficiently capture complete contextual information of scenes in real-world applications, such as autonomous driving. Consequently, selecting and incorporating multiple context frames as well as multimodal information [59] for better context modeling need further investigation. Our work is also limited by medium-scale sizes in terms of both world models and pre-training data, which hinders learning more broadly applicable knowledge. An important direction is to systematically examine the scalability of our method by leveraging scalable architectures like Transformers [47, 67] and massive-scale video datasets [18, 48]. Lastly, our work focuses on pretraining world models via generative objectives, which utilize the model capacity inefficiently on image reconstruction to overcome intricate contexts. Exploring alternative pre-training objectives, such as contrastive learning [51, 11] or self-prediction [65, 80], could further release the potential of IPV by eliminating heavy components on context modeling and focusing on dynamics modeling.

Acknowledgments

We would like to thank many colleagues, in particular Haixu Wu, Baixu Chen, and Jincheng Zhong, for their valuable discussion. This work was supported by the National Key Research and Development Plan (2020AAA0109201), the National Natural Science Foundation of China (62022050 and 62021002), and the Beijing Nova Program (Z201100006820041).

References

Alessandro Achille and Stefano Soatto. Emergence of invariance and disentanglement in deep representations. The Journal of Machine Learning Research , 19(1):1947-1980, 2018.

Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. In NeurIPS , 2021.

Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. OPAL: Offline primitive discovery for accelerating offline reinforcement learning. In ICLR , 2021.

Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. In NeurIPS , 2022.

Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410 , 2016.

Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang, and Nando De Freitas. Playing hard exploration games by watching youtube. In NeurIPS , 2018.

Bowen Baker, Ilge Akkaya, Peter Zhokhov, 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 NeurIPS , 2022.

Matthew Chang, Arjun Gupta, and Saurabh Gupta. Semantic visual navigation by watching youtube videos. In NeurIPS , 2020.

Annie S Chen, Suraj Nair, and Chelsea Finn. Learning generalizable robotic reward functions from "in-the-wild" human videos. In RSS , 2021.

Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. TransDreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.09481 , 2022.

Fei Deng, Ingook Jang, and Sungjin Ahn. DreamerPro: Reconstruction-free model-based reinforcement learning with prototypical representations. In ICML , 2022.

Emily Denton and Vighnesh Birodkar. Unsupervised learning of disentangled representations from video. In NeurIPS , 2017.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In NAACL , 2019.

Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. CARLA: An open urban driving simulator. In CoRL , 2017.

Jean-Yves Franceschi, Edouard Delasalles, Mickaël Chen, Sylvain Lamprier, and Patrick Gallinari. Stochastic latent residual video prediction. In ICML , 2020.

Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. DeepMDP: Learning continuous latent space models for representation learning. In ICML , 2019.

Raghav Goyal, Samira Ebrahimi Kahou, Vincent Michalski, Joanna Materzynska, Susanne Westphal, Heuna Kim, Valentin Haenel, Ingo Fruend, Peter Yianilos, Moritz Mueller-Freitag, et al. The "something something" video database for learning and evaluating visual common sense. In ICCV , 2017.

Kristen Grauman, Andrew Westbury, Eugene Byrne, Zachary Chavis, Antonino Furnari, Rohit Girdhar, Jackson Hamburger, Hao Jiang, Miao Liu, Xingyu Liu, et al. Ego4D: Around the world in 3,000 hours of egocentric video. In CVPR , 2022.

Jake Grigsby and Yanjun Qi. Measuring visual generalization in continuous control from pixels. arXiv preprint arXiv:2010.06740 , 2020.

David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In NeurIPS , 2018.

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In ICLR , 2020.

Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In ICML , 2019.

Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In ICLR , 2021.

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

Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In ICML , 2022.

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In CVPR , 2020.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR , 2016.

David H Hubel and Torsten N Wiesel. Receptive fields and functional architecture in two nonstriate visual areas (18 and 19) of the cat. Journal of neurophysiology , 28(2):229-289, 1965.

Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML , 2015.

Catalin Ionescu, Dragos Papava, Vlad Olaru, and Cristian Sminchisescu. Human3.6M: Large scale datasets and predictive methods for 3d human sensing in natural environments. TPMAI , 2013.

Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In ICLR , 2017.

Stephen James, Zicong Ma, David Rovick Arrojo, and Andrew J Davison. RLBench: The robot learning benchmark & learning environment. RA-L , 2020.

Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning , 37:183-233, 1999.

Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Modelbased reinforcement learning for atari. In ICLR , 2020.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020.

Apoorv Khandelwal, Luca Weihs, Roozbeh Mottaghi, and Aniruddha Kembhavi. Simple but effective: Clip embeddings for embodied ai. In CVPR , 2022.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In ICLR , 2014.

Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Pieter Abbeel. URLB: Unsupervised reinforcement learning benchmark. In NeurIPS , 2021.

Yann LeCun. A path towards autonomous machine intelligence. preprint posted on openreview , 2022.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Yingzhen Li and Stephan Mandt. Disentangled sequential autoencoder. In ICML , 2018.

YuXuan Liu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine. Imitation from observation: Learning to imitate behaviors from raw video via context translation. In ICRA , 2018.

Margaret Livingstone and David Hubel. Segregation of form, color, movement, and depth: anatomy, physiology, and perception. Science , 240(4853):740-749, 1988.

Arjun Majumdar, Karmesh Yadav, Sergio Arnaud, Yecheng Jason Ma, Claire Chen, Sneha Silwal, Aryan Jain, Vincent-Pierre Berges, Pieter Abbeel, Jitendra Malik, et al. Where are we in the search for an artificial visual cortex for embodied intelligence? arXiv preprint arXiv:2303.18240 , 2023.

Russell Mendonca, Shikhar Bahl, and Deepak Pathak. Structured world models from human videos. In RSS , 2023.

Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. In ICLR , 2023.

Antoine Miech, Dimitri Zhukov, Jean-Baptiste Alayrac, Makarand Tapaswi, Ivan Laptev, and Josef Sivic. HowTo100M: Learning a text-video embedding by watching hundred million narrated video clips. In ICCV , 2019.

Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3M: A universal visual representation for robot manipulation. In CoRL , 2022.

Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard L Lewis, and Satinder Singh. Actionconditional video prediction using deep networks in atari games. In NeurIPS , 2015.

Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by latent imagination without reconstruction. In ICRA , 2021.

Minting Pan, Xiangming Zhu, Yunbo Wang, and Xiaokang Yang. Iso-Dream: Isolating and leveraging noncontrollable visual dynamics in world models. In NeurIPS , 2022.

Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsurprising effectiveness of pre-trained vision models for control. In ICML , 2022.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-performance deep learning library. In NeurIPS , 2019.

Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In ICML , 2017.

Georgios Pavlakos, Xiaowei Zhou, Konstantinos G Derpanis, and Kostas Daniilidis. Coarse-tofine volumetric prediction for single-image 3d human pose. In CVPR , 2017.

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

Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. In ICLR , 2023.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In CVPR , 2022.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional networks for biomedical image segmentation. In MICCAI , 2015.

Lukas Schäfer, Filippos Christianos, Josiah P Hanna, and Stefano V Albrecht. Decoupled reinforcement learning to stabilise intrinsically-motivated exploration. In AAMAS , 2022.

Karl Schmeckpeper, Annie Xie, Oleh Rybkin, Stephen Tian, Kostas Daniilidis, Sergey Levine, and Chelsea Finn. Learning predictive models from observation and interaction. In ECCV , 2020.

Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature , 588(7839):604-609, 2020.

John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. Highdimensional continuous control using generalized advantage estimation. In ICLR , 2016.

Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In ICLR , 2021.

Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In ICML , 2020.

Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control. In CoRL , 2022.

Younggyo Seo, Kimin Lee, Stephen L James, and Pieter Abbeel. Reinforcement learning with action-free pre-training from videos. In ICML , 2022.

Rutav Shah and Vikash Kumar. RRL: Resnet as representation for reinforcement learning. In ICML , 2021.

Lin Shao, Toki Migimatsu, Qiang Zhang, Karen Yang, and Jeannette Bohg. Concept2Robot: Learning manipulation concepts from instructions and human demonstrations. IJRR , 2021.

Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In Machine learning proceedings 1990 , pages 216-224. Elsevier, 1990.

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

Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.00690 , 2018.

Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of machine learning research , 9(11), 2008.

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

Ruben Villegas, Jimei Yang, Seunghoon Hong, Xunyu Lin, and Honglak Lee. Decomposing motion and content for natural video sequence prediction. In ICLR , 2017.

Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06173 , 2022.

Yifan Xu, Nicklas Hansen, Zirui Wang, Yung-Chieh Chan, Hao Su, and Zhuowen Tu. On the feasibility of cross-task transfer with model-based reinforcement learning. In ICLR , 2023.

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. In ICLR , 2022.

Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. In NeurIPS , 2021.

Weirui Ye, Yunsheng Zhang, Pieter Abbeel, and Yang Gao. Become a proficient player with limited data through watching pure videos. In ICLR , 2023.

Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson. How transferable are features in deep neural networks? In NeurIPS , 2014.

Fisher Yu, Haofeng Chen, Xin Wang, Wenqi Xian, Yingying Chen, Fangchen Liu, Vashisht Madhavan, and Trevor Darrell. BDD100K: A diverse driving dataset for heterogeneous multitask learning. In CVPR , 2020.

Tianhe Yu, Chelsea Finn, Annie Xie, Sudeep Dasari, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot imitation from observing humans via domain-adaptive meta-learning. In RSS , 2018.

Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In CoRL , 2020.

Zhecheng Yuan, Zhengrong Xue, Bo Yuan, Xueqian Wang, Yi Wu, Yang Gao, and Huazhe Xu. Pre-trained image encoder for generalizable visual reinforcement learning. In NeurIPS , 2022.

Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In ICLR , 2021.

Qihang Zhang, Zhenghao Peng, and Bolei Zhou. Learning to drive by watching youtube videos: Action-conditioned contrastive policy pretraining. In ECCV , 2022.

Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data augmentation. In AAAI , 2020.

A Pseudocode

For clarity, we first elaborate on the optimization objective for contextualized world models during action-free pre-training, which drops action-conditioned dynamics and reward predictors from the fine-tuning objective (Eq. (6)):

For behavior learning, we adopt the same actor-critic learning scheme based purely on imaginary latent trajectories ˆ s τ , ˆ a τ , ˆ r τ with horizon H , as DreamerV2 [21, 23]. The critic v ξ ( s ) is learned by regressing the λ -target [72, 64]:

where sg is a stop gradient function. And the actor π ψ ( a | s ) is learned by maximizing the imagined return by back-propagating the gradients through the learned world model, with an entropy regularizer:

Overall, a complete description of pre-training and fine-tuning of ContextWM for MBRL is presented in Algorithm 1.

B Derivations

The variational bound for contextualized latent dynamics models p ( o 1: T , r 1: T | a 1: T , c ) and a variational posterior q ( z 1: T | a 1: T , o 1: T ) follows from importance weighting and Jensen's Inequality as shown 3 ,

where

Algorithm 1 Contextualized World Models with In-the-wild Pre-training from Videos

1: // Action-free pre-training with in-the-wild videos

2: Initialize parameters θ of action-free dynamics model, image encoder, and decoder randomly

3: Load in-the-wild video dataset D

4: for each training step do

5: // WORLD MODEL PRE-TRAINING

6: Sample random minibatch { o 1: T } ∼ D

8: Update world model by minimizing L CWM-pt ( θ ) in Eq. (7)

7: Get context c ← o ˜ t , ˜ t ∼ Uniform { 1 , T }

9: end for

10: // Action-conditioned fine-tuning with MBRL

11: Load pre-trained parameters θ of action-free dynamics model, image encoder, and decoder

12: Initialize parameters ϕ, φ of action-conditioned dynamics model and reward predictors randomly

13: Initialize parameters ψ, ξ of actor π ψ ( a | s ) and critic v ξ ( s )

14: Initialize replay buffer B with random seed episodes

15: for each timestep t do

16: // COLLECT TRANSITIONS

17: Get representation z t ∼ q θ ( z t | z t -1 , a t -1 , o t ) , s t ∼ q ϕ ( s t | s t -1 , a t -1 , z t )

18: Get action a t ∼ π ψ ( a t | s t )

19: Add transition { ( o t , a t , r t ) } to replay buffer B

20: // WORLD MODEL FINE-TUNING

21: Sample random minibatch { ( o 1: T , a 1: T , r 1: T ) } ∼ B

23: Update world model by minimizing L CWM ( ϕ, φ, θ ) in Eq. (6)

22: Get context c ← o ˜ t , ˜ t ∼ Uniform { 1 , T }

24: // BEHAVIOR LEARNING

25: Imagine future rollouts { ˆ s τ , ˆ a τ , ˆ r τ + λ ˆ r int τ } using world model and actor

26: Update actor and critic by minimizing objectives in Eq. (8) and Eq. (10)

27: end for

Overall, we have a lower bound on the log-likelihood of the data:

We can also calculate that the tightness of the bound is KL [ q ( z | a, o ) ∥ p ( z | a, o, r, c )] .

C Experimental Details

C.1 Data Preprocessing

We use three common video datasets for in-the-wild video pre-training. All inputs are resized to 64 × 64 pixels.

Something-Something-v2. The Something-Something-v2 dataset 4 [17] shows footage of humans interacting with common everyday objects. We use all videos in the training set but filter out videos less than 25 frames long, resulting in a total of 162K videos for pre-training.

Human3.6M. The Human3.6M dataset 5 [30] contains videos of various activities performed by 11 human subjects. We use the particular dataset processed by Pavlakos et al. [56], which includes a total of 840 videos with 210 scenarios and 4 different viewpoints for each scenario. Each video has around 500 frames. To make sure the images used for pre-training better fit our downstream tasks, we place the subject at the center of each frame by cutting a bounding box calculated from the 3D joint positions provided by the dataset. In addition, a padding of 200 pixels is performed around

the bounding box so that the subject is properly positioned. The image clipped using the padded bounding box is then resized to 64 × 64 pixels.

YouTube Driving. The YouTube Driving dataset 6 [88] consists of real-world driving videos with different road and weather conditions. As videos in this dataset usually have a short opening, the first 100 frames of each video are skipped. The original video is downsampled 5 times. We employ all 134 videos in our pre-training process, with each video containing roughly 10000 to 20000 frames.

Assembled Dataset. We compose the three datasets mentioned above into an assembled dataset. On each sample, we uniformly choose a dataset from the three choices and sample videos from that dataset accordingly.

C.2 Benchmark Environments

Meta-world. Meta-world [85] is a benchmark of 50 distinct robotic manipulation tasks. Following APV [68], we use a subset of six tasks, namely Lever Pull, Drawer Open, Door Lock, Button Press Topdown Wall, Reach, and Dial Turn. In all tasks, the episode length is 500 steps without any action repeat. The action dimension is 4, and the reward ranges from 0 to 10. All methods are trained over 250K environment steps, which is consistent with the setting of APV [68].

DMC Remastered. The DMC Remastered (DMCR) Suite [19] is a variant of the DeepMind Control Suite [73] with randomly generated graphics emphasizing visual diversity. On initialization of each episode for both training and evaluation, the environment makes a random sample for 7 factors affecting visual conditions, including floor texture, background, robot body color, target color, reflectance, camera position, and lighting. Our agents are trained and evaluated on three tasks: Cheetah Run, Hopper Stand, and Walker Run. All variation factors are used except camera position in the Hopper Stand task, where we find it too difficult for the agent to learn when the camera is randomly positioned and rotated. Following the common setup of DeepMind Control Suite [21, 79], we set the episode length to 1000 steps with an action repeat of 2, and the reward ranges from 0 to 1. All methods are trained over 1M environment steps.

CARLA. CARLA [14] is an autonomous driving simulator. We conduct a task using a similar setting as Zhang et al. [87], where the agent's goal is to progress along a highway in 1000 time-steps and avoid collision with 20 other vehicles moving along. The maximum episode length is set to 1000 without action repeat. Possible reasons to end an episode early include deviation from the planned route, driving off the road, and being stuck in traffic due to collisions. Each action in the CARLA environment is made up of a steering dimension and an accelerating dimension (or braking for negative acceleration). In addition to the reward function in Zhang et al. , we add a lane-keeping reward to prevent our agent from driving on the shoulder. Our reward function is as follows:

The first term represents the effective distance traveled on the highway, where v ego is the velocity vector projected onto the highway's unit vector ˆ u h , multiplied by a discretized time-step ∆ t = 0 . 05 and a lane-keeping term | center | to penalize driving on the shoulder. The impulse term represents collisions during driving, while the steer ∈ [ -1 , 1] term prevents excessive steering. The hyperparameters λ 1 and λ 2 are set to 10 -3 and 1 , respectively. The central camera of the agent, observed as 64 × 64 pixels, is used as our observation. We use three different weather and sunlight conditions offered by CARLA for training and evaluation: ClearNoon, WetSunset, and dynamic weather. In the dynamic weather setting, a random weather and sunlight condition is sampled at the start of each episode and changes realistically during the episode. To aggregate per-weather scores, we normalize raw episode returns of each weather to comparable ranges. Namely, we linearly rescale episode returns from [ -400 , 1200] under ClearNoon, [ -400 , 800] under WetSunset, and [ -400 , 600] under dynamic weather to a unified range of [0 , 1] . All methods are trained over 150K environment steps.

C.3 Model Details

Overall model. We adopt the stack latent model introduced by APV [68] and extend it into a Contextualized World Model (ContextWM) with the following components:

Specifically, the action-free representation model utilizes a ResNet-style image encoder to process observation input o t . Moreover, the contextualized image decoder also incorporates a similar context encoder to extract contextual information from the context frame c and augment the decoder features with a cross-attention mechanism. We elaborate on the details below.

Image encoders and decoder. To ensure a sufficient capacity for both plain WM and ContextWM pre-trained on diverse in-the-wild videos, we implement image encoders and decoders as 13-layer ResNets [27]. For the sake of simplicity, we only show the architecture of the image encoder in Table 1, as the architecture of the context encoder is exactly the same, and the architecture of the image decoder is symmetric to the encoders. In the decoder, we use nearest-neighbor upsampling for unpooling layers. The outputs of the last residual block of two stages in the context encoder (stage2 and stage3) before average pooling (thus in the shape of 16 × 16 and 8 × 8 respectively) are passed to the corresponding residual block of the image decoder and used to augment the incoming decoder features with cross-attention.

Context cross-attention. The input features X ∈ R c × h × w of residual blocks in the image decoder are augmented with shortcut features Z ∈ R c × h × w of the same shape from the context encoder, by a cross-attention mechanism as Eq. (5). Here we elaborate on the details. After Z is flattened into a sequence of tokens ∈ R hw × c , to mitigate the quadratic complexity of cross-attention, we random

mask 75% of the tokens, resulting tokens of the length ⌊ 1 4 hw ⌋ . Moreover, we use learned position embedding ∈ R hw × c initialized with zeros for both Q and K,V , respectively.

Context augmentation. We randomly sample a frame of observation from the trajectory segment as the input to the context encoder. To prevent the image decoder from taking a shortcut by trivially copying the context frame and thus impeding the training process, we apply random Cutout [89] as image augmentation on the context frame for all Meta-world tasks (except drawer open), as Meta-World observations have a static background and a fixed camera position. Note that we do not utilize any augmentation for DMCR and CARLA tasks. Our comparison with plain WM is fair since we do not introduce any augmentation into the image encoder of the latent dynamics models, thus the agent determines its actions solely based on the original observation.

Latent dynamics model. The stacked latent model from APV [68] is adopted to enable action-free pre-training and action-conditioned fine-tuning, where both the action-free and action-conditioned models are built upon the discrete latent dynamics model introduced in DreamerV2 [23], where the latent state consists of a deterministic part and a discrete stochastic part. Inputs to the latent dynamics model are representations from the image encoder. Following APV [68], we use 1024 as the hidden size of dense layers and the deterministic model state dimension.

Reward predictors, actor, and critic. Both of the dual reward predictors, as well as actor and critic, are all implemented as 4-layer MLPs with the hidden size 400 and ELU activation function, following DreamerV2 [23] and APV [68].

C.4 Hyperparameters

For the newly introduced hyperparameter, we use β r = 1 . 0 for all tasks. We set β z = 1 . 0 for pre-training, and β z = 0 . 0 , β s = 1 . 0 for fine-tuning, following APV [68]. Despite we adopt a deeper architecture for both image encoders and decoders, we find that hyperparameters are relatively robust to architectures, and the new deeper architecture under the same hyperparameters can match or even slightly improve the performance originally reported by APV. Thus, we use the same hyperparameters with APV. For completeness, we report important hyperparameters in Table 2.

C.5 Qualitative Analysis Details

We explain the visualization scheme for Figure 8. For video prediction in Figure 8a, we sample a video clip of length 25 and let the model observe the first 15 frames and then predict the future 10 frames open-loop by the latent dynamics model. The context frame is selected as the last of the observed 15 frames. To visualize context attention, we compute the cross-attention weight across the 16 × 16 feature maps without random masking and plot a heatmap of the weights averaged over attention heads and over attention targets of a small region (for example, 3 × 3 ) on the decoder feature map. For video representations in Figure 8b, we sample video clips of length 25 and compute the averaged model states avg( z 1: T ) of dimension 2048 for each video clip, which are then visualized by t-SNE [74].

C.6 Computational Resources

We implement all the methods based on PyTorch [54] and train with automatic mixed precision. In terms of parameter counts, ContextWM consists of 26M and 47M parameters for video pre-training and MBRL, respectively, while a plain WM consists of 24M and 44M parameters. In terms of training time, it takes ∼ 37 hours for pre-training of ContextWM over 600K iterations, and finetuning of ContextWM with MBRL requires ∼ 24 hours for each run of Meta-world experiments over 250K environment steps, ∼ 23 hours for each run of DMC Remastered experiments over 1M environment steps and ∼ 16 hours for each run of CARLA experiments over 150K environment steps, respectively. Although ContextWM introduces new components of the model, we find it does not significantly impact the training time, since the context encoder only needs to operate on one frame of observations, and the cross-attention between features also accounts for a relatively small amount of computation. In terms of memory usage, Meta-world experiments require ∼ 18GB GPU memory, and DMC Remastered and CARLA experiments require ∼ 7GB GPU memory, thus the experiments can be done using typical 24GB and 12GB GPUs, respectively.

D Additional Experimental Results

D.1 Comparison with Additional Baselines

We compare our approach with a state-of-the-art model-free RL method DrQ-v2 [79] and a modelbased RL method Iso-Dream [52]. We adapt their official implementations 78 to all the tasks in our experiments and present the results in Figure 9. As shown, our method consistently outperforms stateof-the-art baselines, which demonstrates the merit of pre-training from in-the-wild videos. Model-free methods such as DrQ-v2 have been shown to perform worse than model-based methods [68, 79], and our results are consistent with this. We also remark that our method differs from Iso-Dream [52] since we separate modeling of context and dynamics at the semantic level while Iso-Dream isolates controllable and noncontrollable parts at the pixel level. Thus, contextual information such as body color has to be modeled by the dynamics branch in Iso-dream, which probably results in worse decoupling of context and dynamics as well as inferior performance for challenging tasks. Moreover, compared to Iso-Dream, our model makes a simpler adjustment to Dreamer architecture.

(a) Aggregate results

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Lever Pull

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Drawer Open

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Door Lock

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Button Press Topdown Wall

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Reach

DrQ-v2

Iso-Dream

IPV w/ ContextWM

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

(b) Meta-world

0.0

0.2

0.4

0.6

0.8

1.0

Environment Steps (

×

)

0

100

200

300

400

500

600

700

DMCR Cheetah Run

0.0

0.2

0.4

0.6

0.8

1.0

Environment Steps (

×

)

0

100

200

300

400

500

600

700

0.0

0.2

0.4

0.6

0.8

1.0

Environment Steps (

×

)

0

100

200

300

400

500

600

DMCR Walker Run

DrQ-v2

Iso-Dream

IPV w/ ContextWM

(c) DMC Remastered

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Environment Steps (

×

)

400

200

0

200

400

600

800

1000

1200

1400

CARLA Clear Noon

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Environment Steps (

×

)

400

200

0

200

400

600

800

1000

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Environment Steps (

×

)

400

200

0

200

400

600

CARLA Dynamic Weather

DrQ-v2

Iso-Dream

IPV w/ ContextWM

(d) CARLA

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Lever Pull

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Drawer Open

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Door Lock

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Button Press Topdown Wall

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

Reach

IPV

IPV (Enc / Dec Only)

0

5

10

15

20

25

Environment Steps (

×

)

0

20

40

60

80

100

Success Rate (%)

D.2 Effects of Dynamics Knowledge

To investigate the importance of dynamics knowledge captured by pre-trained ContextWM, we report the performance of IPV with ContextWM when only the pre-trained parameters of the visual encoders and decoder are transferred. Figure 10 indicates that only transferring visual encoders and decoder performs worse on four of six tasks. These results align with the observation of APV [68] and highlight that pre-trained world models are capable to capture essential dynamics knowledge beyond learning visual representations solely [49, 53], further boosting visual control learning.

E Additional Discussions

Particularly strong performance gains on DMCR tasks. Our method obtains considerable performance gains on DMCR tasks. The main reason is that DMCR is a purposefully designed benchmark, which measures visual generalization and requires the agent to extract task-relevant information as well as ignore visual distractors [19]. As demonstrated in Figure 8c, our ContextWM has the advantage of separately modeling contexts (task-irrelevant in DMCR) and dynamics (taskrelevant in DMCR), which avoids wasting the capacity of dynamics models in modeling low-level visual details. Furthermore, pre-training with in-the-wild videos enables our models to eliminate diverse distractors and capture shared motions, which is essential for visual generalization in RL. In contrast, a plain WM needs to model complicated contexts and dynamics in an entangled manner, which adds difficulty to dynamics learning and behavior learning on these features. We also emphasize that, motivated by separating contexts and enhancing temporal dynamics modeling, our proposed IPV w/ ContextWM is a general-purpose framework and, as shown in our experiments, can obtain adequate performance gains on various benchmarks that have more complicated entangling of contexts and dynamics beyond DMCR.

Discussion with Transformer-based world models. While a cross-attention mechanism is leveraged for context information conditioning, we still use an RNN-based latent dynamics model from Dreamer [21] to model history observations. More powerful sequential backbones such as Transformers are orthogonal to our primary technical contribution of explicit context modeling and conditioning. Nevertheless, exploring Transformer-based world models [10, 47, 58] is valuable, not only because of its favorable scalability but also of its flexibility to incorporate pre-training on action-free videos by simply changing conditioning variables and to condition the contextualized decoder in a more native way. Overall, the combination of our pre-training framework with a Transformer architecture holds great potential, and we will delve deeper into this aspect in future work.

F Additional Visualizations

Video prediction. We present additional showcases of video prediction by plain WM and ContextWM, respectively, on three video datasets, in Figure 11. On the challenging Something-Something-v2 and YouTube Driving datasets, our model provides better prediction quality in comparison to a plain WM and well captures dynamics information (e.g., how objects are moving) with concentrated cross-attention. However, on the Human3.6M dataset, although our model better predicts human motions, we find its context cross-attention is relatively divergent. This result supports our conjecture that our model pre-trained on the Human3.6M dataset lacking in diversity may suffer from overfitting, which results in insufficient modeling of context and dynamics as well as inferior performance when fine-tuned with MBRL.

truth

Predicted

(ContextWM)

Context

Attention

Predicted

(Plain WM)

(a) Something-Something-V2

truth

(ContextWM)

Attention

(Plain WM)

truth

Predicted

(ContextWM)

Context

Attention

Predicted

(Plain WM)

Predicted

Context

Predicted

Video representation. We present in Figure 12 additional t-SNE visualization similar to Figure 8. Our ContextWM learns well-distributed representations that clearly separate motion directions while plain WM fails.

Compositional decoding. Figure 13 includes additional compositional decoding analysis on the DMCR Walker Run task.

truth

IPV w/

ContextWM

Context

Shift

IPV w/

Plain WM

Reconstruction

Context

Open-loop Prediction

AA

G Broader Impact

Advancements in world models and training them with in-the-wild videos can raise concerns about the generation of deepfakes, realistic synthetic videos that mimic real events or individuals. Additionally, advancements in MBRL can also increase automation in domains like robotic manipulation and autonomous driving, bringing efficiency and safety benefits but also leading to job displacement and socioeconomic consequences. However, it is important to note that as the field is still in its early stages of development and our model is only a research prototype, the aforementioned negative social impacts are not expected to manifest in the short term.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a method for learning world models from in-the-wild videos. By utilizing a context encoder to capture contextual information, the proposed method explicitly models both the context and dynamics to facilitate knowledge transfer across scenes. Experiments are performed on various simulation benchmarks such as Meta-world, DMC Remastered, DeepMind Control Suite, and CARLA.\n\nSTRENGTHS:\nThe motivation for learning better world models by disentangling context information and dynamics is clear and seems reasonable to me; \n\nThe proposed method for learning context information is simple and straightforward; \n\nThe paper is well-organized and easily read.\n\nWEAKNESSES:\nThe proposed framework is based on the assumption that the context information lies equally in each frame, however, it is very likely that some moving objects might be occluded at some time. There are no specific designs for handling these situations; \n\nThe experiments mainly validate the sample efficiency of the RL training process, there is no sufficient ablation study on the learned context information; \n\nCompared to the vanilla WM baseline, the performance of RL training on the Meta-world benchmark seems not very impressive; \n\nThe predicted video frames seem not very promising; \n\nThere is no discussion of limitations and failure cases.\n\nQUESTIONS:\nThe authors propose to randomly select a frame for predicting the context feature, have the authors tried to use a fixed frame (e.g., the first frame or the last frame of the video clip)? Would the context feature extracted from the different frames be consistent with each other? And how would choosing different frames affect the final performance of the RL training? \n\nIt is assumed that the learned context feature captures the static properties of objects (e.g., colors, shapes), would modifying the context feature allow us to generate diverse video frames (e.g., changing the color of the object)? It would be interesting to show more temporal consistent video frames by editing the context feature.\n\nLIMITATIONS:\nThere is no discussion of limitations and failure cases.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nLearning a world-model that can generalize to different domains and tasks is difficult. The authors enhanced an existing framework for pre-training world models using in-the-wild videos, which can be fine-tuned on downstream tasks. In particular, the authors introduce a contextual encoder which helps in disentangling temporal dynamics from static contexts. Additionally, they include a cross-attention mechanism and dual reward predictors to improve the learning of task-relevant representations.\n\nSTRENGTHS:\nThe manuscript is very well written and structured. \n\n●\tThe idea of using a context to encode static information is novel, well-motivated and it shows good results in the DMC remastered task\n\n●\tLearning world models on in the wild videos is hard (and so far does not help downstream tasks, as shown by Seo et al. [49]) - The extension of the Action-Free Pre-training from Videos approach to in-the-wild videos is convincing and the context modulation via a U-Net for reconstruction is innovative. \n\n●\tThe authors perform several relevant ablations to illustrate the role of the different proposed components.\n\nWEAKNESSES:\nOverall, the context is only used together with the latent dynamic to decode the image. Is there a clear reason why the context is not taken into account for predicting the reward (in Fig. 3b the arrow of the context that goes to the reward seems to be misleading)? As the authors stated, the context might implicitly encode some important information about the task, e.g., the static position of an object, which can be helpful during the fine-tuning phase.\n\nQUESTIONS:\n●\tWhat does the context encode? The motivation and the qualitative analysis are convincing, but it would also be useful to test this further. This could be done via a decoding analysis. In particular, it would be very interesting to see the difference between the context and the dynamic in the case of the DMC remastered task, where the contextWM provides a significant advantage. \n\n●\tWhy do you think is the gap in performance to prior methods particularly strong for DMCR?\n\n●\tLooking at figure 6b (bottom), the performance of the pre-trained ContextWM on the SSv2 dataset for the CARLA driving task seems to outperform the one that is not-pretrained. However, the effects of the dataset domain of figure 7b (right) shows that there is almost no difference between the one pre-trained on SSv2 and the one without pre-training. What am I missing?\n\n●\tDoes the contextWM help to achieve better generalization? Could one, e.g., purposefully change the color, size or shape of the object for the meta-world and achieve better performance than the WM?\n\nLIMITATIONS:\nFine", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies whether large-scale in-the-wild datasets can be used to pre-train world models for efficient downstream reinforcement learning. Specifically, they introduce Contextualized World Models (ContextWM), an architecture specifically designed to learn to separate context and dynamics modeling. Their experimentation reveals that the proposed methodology outperforms DreamerV2 on a variety of downstream tasks.\n\nSTRENGTHS:\n- Well-written. The paper is well-written, and the figures aid in the understanding of the methodology.\n- Variety of experiments. The authors conduct a variety of experiments, comparing not just to DreamerV2, but also analyzing the effects of pre-training, architecture choices, dataset domain, etc. Furthermore, both quantitative and qualitative comparisons are included.\n- Strong experimental. Relative to DreamerV2, the proposed methodology achieves a strong performance -- the gap seems to be particularly large for the DMC Remastered tasks.\n\nWEAKNESSES:\n- Missing comparison to prior work. APV [49] is the most similar prior work which this work builds off of (and seems to be the SOTA in this space), and yet the proposed method is not compared to APV. Many of the tasks used, hyper-parameters, and evaluation protocols adopted are from APV, allowing for a comparison, yet somehow, this comparison is omitted. Looking at the results plots in the APV paper, visually, the performance of ContextWM seems similar to that of APV (and in some cases clearly worse, e.g. in the dial turn task). While one may claim that a comparison to APV may be unfair because both the data and the model would be different, this would still reveal whether the ability to use a larger amount of in-the-wild data as well as the changes to the model architecture are actually beneficial. Furthermore, training APV on this large-scale data and seeing whether ContextWM outperforms it would be an experiment which would reveal whether the architectural changes proposed are significant. Finally, APV primarily compares to DreamerV2 as this was the SOTA when APV was published -- now, DreamerV3 seems to be the SOTA, so a comparison to DreamerV3 should be conducted.\n- Incomprehensive ablations. It is not clear which one environment the ablation study is conducted on. It would be a lot more convincing if the ablation study was done across tasks, and if the trends held true across tasks as this would alleviate concerns of cherry-picking a task.\n\nQUESTIONS:\n- Is it possible to qualitatively demonstrate that ContextWM is able to separate context and dynamics modeling?\n- For some tasks, pre-training and choice of data to pre-train with makes a huge difference, e.g. DMC Re-mastered. For others, not so much, e.g. CARLA. Is there a sense of why this is the case?\n\nLIMITATIONS:\nThere is no discussion on the limitations of the work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a Contexturelized World Model with In-the-wild Video Pretraining, which extends recently proposed action-free pre-training from videos (APV) to the case of contextualized video-prediction models. Specifically, they propose to sample a randomly chosen frame and use it as \"contextualized information\" for the prediction model. The \"contextualized\" information is combined into the prediction model via multi-scale cross-attention mechanisms. Besides, during the model-based RL phase, they propose to predict both pure task reward and the sum of the task reward plus a weighted intrinsic reward. Empirical validations on several well-known benchmarks (Meta-world, DMC, and CARLA) show the superior performance of the proposed method.\n\nSTRENGTHS:\n1. The paper is generally well-written and easy to follow. \n2. While the proposed method is simple, experimental results show a stable performance gain by the proposed methods. \n3. Besides the main contribution, the paper also provides an interesting analysis of the choice of pre-training datasets.\n\nWEAKNESSES:\n1. Technical novelty of the proposed method is not high. The main proposal is to add contextualized information to facilitate a better prediction. Besides, the paper does not discuss how the proposed approach (random sample selection) incorporates contextualized information well. \n\nEspecially, I'm not sure we can call the random variable $c$ a contextualized vector since it is computed via only single frames rather than the full context of the trajectories. In this sense, I'm wondering what happens if we predict/reconstruct observations using the same multi-scale cross-attention architecture but with the image from the same time step for each frame. In the setup, the model does not use contextualized information but has similar architecture to the proposed method, and thus more appropriate for the baseline. \n\n2. Lack of discussion with transformer-based world models. As shortly discussed, several studies incorporate transformer architecture in the prediction model. Since the transformer architecture predicts the future via autoregressive fashion, I think it is more natural to handle the contextualized information. However, the current manuscripts lack discussion on this point and comparison with transformer-based world models\n\nQUESTIONS:\n1. Did you try methods other than random sample selection to compute contextualized information? \n\n2. Did you try a comparison with transformer-based architecture? \n\n3. The logic behind several sentences is hard to capture.  Could you please explain more about these sentences?\n\n- In line 182, \"Nevertheless, ..... Therefore, we propose a dual reward predictor\" => Why does the dual predictor resolve the issue? Why not just balance lambda depending on the task? \n- In line 234, \"indicating that the performance gain of IPV with vanilla WM is primarily due to the intrinsic exploration bonus. \" => The logic is unclear. \n\n[Minor]\n- In eq. 6, How do you compute KL when t=1? \n- In Fig. 5-b, Pre: O should be Pre: ✓ or with Pre or something like that. \n- If I correctly understood, Fig. 3 is a bit misleading as a context variable $c$ is directly input into each frame, rather then recurrent prediction.\n\nLIMITATIONS:\nA limitation section should be added.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis submission presents the contextualized world model (ContextWM), a framework for leveraging in-the-wild videos for pre-training of a world model to be used in model-based reinforcement learning. Following the work from Seo et al. (2022), the authors pre-train an action-free version of the recurrent state-space model (RSSM) with two important modifications: 1) a context encoder processes a randomly sampled frame of the input video to provide context features, which the decoder can directly access to better reconstruct static visual details, enabling the dynamics model, on which the decoder is also conditioned, to focus on temporally varying information. 2) the authors also opt for using a dual-reward predictor during fine-tuning, which predicts the pure task reward in addition to the combined task and video-based intrinsic novelty reward proposed by Seo et al. (2022). This facilitates task-relevant representation learning.\nThe ContextWM is evaluated on Meta-world, the remastered DeepMind Control Suite, and a task with varying weather conditions in the CARLA driving simulator. In most benchmarks, ContextWM shows significant improvements in terms of sample efficiency or final performance. An ablation study and other analytical experiments further show the effectiveness of ContextWM and its design decisions.\n\nSTRENGTHS:\nThe proposed contextualized world model is novel and the design decisions are well motivated.\nFor the most part the description of the method and experimental setup is very clear. The performance improvements are significant and in some cases very impressive. The qualitative analysis provides some interesting insights, for instance the clear separation of video representations for two videos with contrastive labels in Figure 8b.\n\nWEAKNESSES:\n1. As the main contribution the cross-attention to the context should be explained in a bit more detail. Since the authors mention U-Nets, I wonder whether the decoder attends to context features at each of the corresponding three stages shown in Table 1 in the supplementary material or only at the decoder input level?\n\nQUESTIONS:\n2. Why did you choose BatchNorm instead of LayerNorm? Did you experiment with both?\n1. maybe use a check mark in the legend of Figure 6b instead of the circle to indicate \"with pretraining\".\n1. It'd be very interesting to see more examples with other contrastive labels for the Video representations experiment in Section 5.5.\n\n## Acknowledgement of rebuttal\nI have read the rebuttal, other reviews. My relatively minor concerns have been addressed, given that the authors have provided some requested clarifications, a discussion of limitations, and additional insightful experiments. I strongly believe this paper should be accepted.\n\nLIMITATIONS:\nThe discussion of limitations is very *limited*. The discussion section mostly mentions future work on scaling and exploring other pre-training objectives.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}