This is a personal paper summary of the paper "Learning to Model the World with Language", https://arxiv.org/abs/2308.01399 Moreover, this is the combined version: last week's slides + detailed methods. (https://www.slideshare.net/Seungjoon1/230906-paper-summary-learning-to-world-model-with-language-publicpdf) So, unity and coherence can be small. Some of the contents may be incorrect. Please send me an email if you want to contact me: sjlee1218@postech.ac.kr (for correction or addition of materials, ideas to develop this paper, or others).

1. Machine
Learning
LABoratory
Seungjoon Lee. 2023-09-15. sjlee1218@postech.ac.kr
Learning to Model the World
with Language
Paper Summary
1

2. Contents
• Introduction
• Methods
• Experiments
• ‘Dynamics’ in DynaLang
2

3. Caution!!!
• This is the material I summarized a paper at my personal research meeting.
• Some of the contents may be incorrect!
• Some contributions, experiments are excluded intentionally, because they
are not directly related to my research interest.
• Methods are simpli
fi
ed for easy explanation.
• Please send me an email if you want to contact me: sjlee1218@postech.ac.kr
(for correction or addition of materials, ideas to develop this paper, or others).
3

4. Intro
4

5. Situations
• Most language-conditioned RL methods only use language as instructions
(eg. “Pick the blue box”)
• However, language does not always match the optimal action.
• Therefore, mapping language only to actions is a weak learning signal.
5
“Put the bowls away”

6. Complication
• On the other hand, human can predict the future using language.
• Human can predict environment dynamics (eg. “wrenches tightens nuts.”)
• Human can predict the future observations (eg. “the paper is outside.”)
6

7. Questions & Hypothesis
• Question:
• If we let reinforcement learning predict the future using language, will its
performance improve?
• Hypothesis:
• Predicting the future representation provides a rich learning signal for
agents of how language relates to the world.
7

8. Contributions
• What’s done:
• DynaLang makes a language-conditioned world model, which can be
trained by self-supervised manner.
• So what?
• The self-supervised world model enables training in sparse-reward envs,
and text-only pretraining without actions or task rewards.
• This shows learning language dynamics helps to make useful
representation for RL.
8

9. Why is This New?
• DynaLang makes a language-conditioned world model which learns the
dynamics of language and image.
• Previous works use language to make language-conditioned policies or to
make additional rewards.
• The world model can be trained in self-supervised manner.
• This enables to make useful feature representation in sparse-reward envs, and
text-only pretraining without actions or task rewards.
9

10. Methods
10

11. Problem Setting
• Observation: , where is an image, is a language token.
• An agent chooses action , then environment returns:
• reward ,
• a
fl
ag whether the episode continues ,
• and next observation .
•
The agent’s goal is to maximize
ot = (xt, lt) xt lt
at
rt+1
ct+1
ot+1
E
[
T
∑
t=1
γt−1
rt
]
11

12. Method Outline
• DynaLang components
• World model: encodes current image obs and language into representation.
• RL agent: using encoded representation, acts to maximize the sum of
discounted reward.
12

13. Method - World Model
Outline
• World model components:
• Encoder - Decoder: learns to represent the current state.
• Sequence model: learns to predict the future state representation.
13

14. Method - World Model
Base model (previous work)
• DynaLang = Dreamer V3 + language.
• Dreamer V3 learns to compute compact representations of current state, and
learns how these concepts change by actions.
14
Architecture of Dreamer V3

15. Method - World Model
Incorporation of language
• DynaLang incorporates language into the encoder-decoder of Dremer V3.
• By this, DynaLang gets representations unifying visual observations and
languages.
15

16. Method - World Model
Prediction of the future
• DynaLang predicts future representation using the sequence model, like
Dreamer V3.
• Future representation prediction lets the agent extract the information from
language, relating to the dynamics of multiple modalities.
16

17. Method - World Model
Model Losses
• World model loss: , where
• Image loss
• Language loss or
• Reward loss
• Continue loss
• Regularizer , where sg is stop-gradient
• Future prediction loss
Lx + Ll + Lr + Lc + Lreg + Lpred
Lx = || ̂
xt − x||2
2
Ll = categorical_cross_entropy( ̂
lt, lt) Ll = || ̂
lt − lt ||2
Lr = ( ̂
rt − rt)2
Lc = binary_cross_entropy( ̂
ct, ct)
Lreg = βreg max(1,KL[zt |sg( ̂
zt)])
Lpred = βpred max(1,KL[sg(zt), ̂
zt])
17

18. Language into the World Model
• Questions to address:
• How are languages tokenized and fed into world model?
• What is the language embedding the world model use? Embedding from
pretrained language model (LM)? Or embedding from scratch?
• Answer:
• DynaLang uses existing tokenizer, and feeds pretrained embedding or one-
hot encoded embedding into the world model.
18

19. Language to World Model
Pretrained LM
• DynaLang uses existing tokenizer and pretrained LM encoder[T5]. (In except the
HomeGrid env)
19
Sentence
T5
Tokenizer
Tokens
T5
Encoder embedding
Rn
DynaLang
Language
Encoder
(MLP)
embedding
Rk
Fixed Learnable
[T5]: https://arxiv.org/abs/1910.10683

20. Language from World Model
Pretrained LM
20
• Dynalang makes embedding from decoder close to embedding from the
pretrained LM encoder.
• Loss = ||lDynaLang − lpretrained ||2
world model
embedding
z DynaLang
Language
Decoder
(MLP)
embedding
from
decoder
Rn
embedding
from
LM encoder
Rn
DynaLang
Language
Encoder
(MLP)
embedding
from
encoder
Rk
World
Model
Encoder

21. Language to World Model
One-hot encoder
• On the other hand, DynaLang also can use one-hot encoder with T5
tokenizer. (In HomeGrid env)
21
Sentence
T5
Tokenizer
Tokens
One-hot
Encoder
DynaLang
Language
Encoder
(MLP)
embedding
Rk
Fixed Learnable
0
0
1
0
.
.
.

22. Method - RL Agent
Outline
• The used RL agent is a simple actor critic agent.
• Actor:
• Critic:
• Note that the RL agent is not conditioned on language directly.
π(at |zt, ht)
V(ht, zt)
22

23. Method - RL Agent
Environment interaction
• The RL agent interacts with environment using the encoded representation
and history .
zt
ht
23

24. Method - RL Agent
Training
• Let , the estimated discounted sum of
future rewards.
• Critic loss:
• Actor loss: , maximizing the return estimate
• The agent is trained only using imagined rollout generated by the world model.
• The agent is trained by the action of the agent and the predicted states, rewards.
Rt = rt + γct ((1 − λ)V (zt+1, ht+1) + λRt+1)
Lϕ = (Vϕ(zt, ht) − Rt)
2
Lθ = − (Rt − V(zt, ht)) log πθ(at |ht, zt)
24

25. Experiments
25

26. Diverse Types of Language
Questions
• Questions to address:
• Can DynaLang use diverse types of language along with instruction?
• If can, does it improve task performance?
26

27. Diverse Types of Language
Setup
• Env: HomeGrid
• multitask grid world where agents receive task
instruction in language but also language hints.
• Agents gets a reward of 1 when a task is completed,
and then a new task is sampled.
• Therefore, agents must complete as many tasks
as possible before the episode terminates in 100
steps.
•
27
HomeGrid env. Agents receive 3 typess of hints.

28. Diverse Types of Language
Results
• Baselines: model-free o
ff
-policy algorithms, IMPALA, R2D2.
• Simply image embeddings, language embeddings are conditioned to policy.
• DynaLang solves more tasks with hints, but simple language-conditioned RL
get worse with hints.
28
HomeGrid training performance after 50M steps (2 seeds)

29. World Model with Sparse/No Rewards
• DynaLang learns to extract features in self-supervised manner.
• By encoder-decoder structure, and future predictive objective.
• Because of this learning method, DynaLang can make useful embedding even
in environments with sparse reward and no rewards.
• Existing language-conditioned policy methods cannot make useful
embedding without rewards. Because their language encoder is trained by
reward signal.
29

30. World Model with Sparse Rewards
Setup
• Env: Messenger
• grid world where agents should deliver a message
while avoiding enemies using text manuals.
• Agents must understand manuals and relate them to
the environment to achieve high score.
30
Messenger env. Agent get text manuals.

31. World Model with Sparse Rewards
Results
• EMMA is added to be compared:
• Language + gridworld speci
fi
c method, model-free language-conditioned policy.
• Only DynaLang can learn from S3, the most di
ffi
cult setting.
• Adding future prediction helps the training more than only action generation.
• However, the authors do not include ablation studies which exclude the future
prediction loss from their architecture.
31
Messenger training performance (2 seeds). S1 is most easy, S3 is most hard.

32. World Model with Sparse Rewards
Results
• DynaLang learns sparse-reward Messenger S3(hard), outperforming EMMA.
• EMMA is a model-free special architecture designed for Messenger env.
• Messenger S3 is a di
ffi
cult game, because it have many entities, and entities have
same appearance and di
ff
erent roles and movements.
32

33. World Model with No Rewards
Text only pretaining
• Self-supervised manner allows text-only o
ffl
ine pertaining.
• By zeroing out the other irrelevant loss, and ignoring actions.
• Existing model free language-conditioned methods cannot be pretrained with
action-free and reward-free data.
33

34. World Model with No Rewards
Text only pretaining
• Text-only pretraining of the world model shows improvements of training
performance in Messenger S2 env.
• Learned language dynamics helps to make useful representation for RL.
34
T5 tokenizer + T5 pretrained LM encoder
T5 tokenizer + one-hot encoder (no pretraining)
T5 tokenizer + one-hot encoder + pretraining with Messenger manuals
T5 tokenizer + one-hot encoder + pretraining with domain general TinyStories dataset
(short stories generated by GPT-3.5 and GPT-4)

35. ‘Dynamics’ in DynaLang
35

36. What is the ‘Dynamics’ DynaLang Learns?
• World model dynamics = language dynamics + visual game dynamics.
• DynaLang learns the dynamics of language, relating it to the dynamics of
visual game.
• Evidence:
• DynaLang can generate texts.
• DynaLang can do embodied question answering.
36

37. Language Dynamics
Evidence 1 - text generation
• DynaLang is trained to predict next language token of TinyStories dataset.
• Below is the example of 10-token generations conditioned on prompts.
• The tokens are predicted by DynaLang and decoded by T5 Tokenizer.
37

38. Language Dynamics
Evidence 2 - embodied question answering (EQA)
• The authors introduces new benchmark, LangRoom.
• An agent gets a question for the color of an object.
• The agent should move to the correct object and say the correct color.
• The agent should understand the object name, relating it to the visual observation.
• Action space: movement and 15 color tokens.
• However, there is already an EQA benchmark[EQA]…
38
[EQA]: https://arxiv.org/abs/1711.11543