Emergence of Abstract State Representations in Embodied Sequence Modeling
Authors: Tian Yun
Zilai Zeng
Kunal Handa
Ashish V Thapliyal
Bo Pang
Ellie Pavlick
Chen Sun
Venue: EMNLP 2023 Main (2023)
Keywords: Interpretability and Analysis; Decision Making via Sequence Modeling; Language Grounding to Vision and Beyond
TL;DR: We investigate the impacts of states and natural language instructions on what knowledge Transformer-based sequence models can learn.

Abstract: Decision making via sequence modeling aims to mimic the success of language models, where actions taken by an embodied agent are modeled as tokens to predict. Despite their promising performance, it remains unclear if embodied sequence modeling leads to the emergence of internal representations that represent the environmental state information. A model that lacks abstract state representations would be liable to make decisions based on surface statistics which fail to generalize. We take the BabyAI environment, a grid world in which language-conditioned navigation tasks are performed, and build a sequence modeling Transformer, which takes a language instruction, a sequence of actions, and environmental observations as its inputs. In order to investigate the emergence of abstract state representations, we design a "blindfolded" navigation task, where only the initial environmental layout, the language instruction, and the action sequence to complete the task are available for training. Our probing results show that intermediate environmental layouts can be reasonably reconstructed from the internal activations of a trained model, and that language instructions play a role in the reconstruction accuracy. Our results suggest that many key features of state representations can emerge via embodied sequence modeling, supporting an optimistic outlook for applications of sequence modeling objectives to more complex embodied decision-making domains.

Emergence of Abstract State Representations in Embodied Sequence Modeling
Tian Yun
Zilai Zeng"
Kunal Handa"
Ashish V. Thapliyal*
Bo Pang"
Ellie Pavlick
Chen Sun
Brown University,
Google Research
(tian_yun, zilai_zeng, chensun)ebrown.ede edu
Abstract
have achieved impressive capabilities on various tasks, such as code writing (Chen al., 2021b), mathematical problem solving (Lewkowycz et al., 2022), and question answering (Sanh et al., 2022).
Decision making via sequence modeling aims mimic the success of language models, where actions taken by an embodied agent are modeled as tokens to predict. Despite their promising performance, it remains unclear if embodied sequence modeling leads to the emer- gence of internal representations that represent the environmental state information. A model that lacks abstract state representations would be liable make decisions based surface statistics which fail to generalize. We take the BabyAl environment, a grid world which language-con eonniigicood navigation tasks per- formed, and build sequence modeling Trans- former, which takes language instruction, sequence of actions, and environmental obser- vations as its inputs. In order to investigate the emergence of abstract state representations, design "blindfolded" navigation task, where only initial environmental layout, the lan- guage instruction, and the action sequence to complete the task are available for training. Our probing results show that intermediate environ- mental layouts can be reasonably reconstructed from the internal activations a trained model, and that language instructions play a role the reconstruction accuracy. Our results suggest that many key features of state representations emerge via embodied sequence modeling, supporting an optimistic outlook for applica- tions of sequence modeling objectives to more complex embodied Secksinnmmkking domains."
Motivated by the success of sequence modeling in language models, decision making via sequence modeling treats each action of an embodied agent as a token, and formulates problems traditionally solved via reinforcement learning (RL) as quence modeling problem (Chen al., 202 la; Jan- ner et al., 2021). Although these models show im- pressive performance, it remains unclear whether such embodied sequence modeling leads to the emergence of internal representations that encode the structure of the external environment, or rather that the models simply learn to predict next actions based on "surface statistics" correlations that en- code the likelihood of action based the given contexts. If the latter true, we would need re-examine the underlying reasons for the strong empirical success of the sequence models. How- ever, if the sequence models do encode an implicit world model, it can help us better understand their effectiveness for decision making tasks.
In this work, we study whether abstract state rep- resentations emerge in embodied sequence model- ing. We use the BabyAl -Coivalier-Boisvertt et al.. 2019) environment to design "blindfolded" nav- igation task, where only the initial layout of the environment, language instruction, and an action sequence which completes the task are available during model pretraining. This agent can follow the instructions by memorizing the action sequences it has seen during pretraining, or by constructing and updating abstract environmental representa- tions and then picking next moves based these representations. Although the model has no prior knowledge about the intermediate layouts of the environment (i.e., non-initial layouts), our prob- ing results demonstrate that trained model learns internal representations of the intermediate layouts.
1 Introduction
Sequence modeling, the task of predicting the next symbol based on the given context (e.g., words, pix- els), is simple yet versatile objective. Leveraging this sequence modeling objective and the Trans- former (Vaswani al., 2017), language models
Project webpage and other resources are available at: https: settps:/aabie /tt////bbaaattttaaaaaaas senect-stttt-seeeeee stattracttsttte--smmdellle
*: First two authors contribute equally. Part of the work was performed while T'Y was a student rescarcher at Google.

---

Language
Instruction: "go to the grey key"
(A)
(B)
(C)
9
0
"Purn Lefe"
0
Move Forward"
0
0
0
8
Figure A demonstration of a language-condittoned navigation task in BabyAI GoToLocal level (Chevalier- Boisvert et al., 2019) where agent (i.c., red triangle) needs navigate the target objects specified in the given instructions. This is an example trajectory where the agent achieves the goal of "go to the grey key": (A) The board is randomly initialized with the agent and objects. (B) The agent navigates in the environment to try to approach the target object. (C) The agent keeps navigating until it reaches the target object.
In this navigation task, there are three compo- nents - layouts of the environment, language in- structions, and action sequences each of which can impact the emergence of abstract state repre- sentations in embodied sequence modeling. We further raise the question: what role do language instructions play in the emergence of such repre- sentations? We ablate with regards to the presence or absence of language instructions during model training, and train probes to recover the interme- diate environmental layouts. We observe that lan- guage instructions indeed facilitate the emergence of abstract environmental representations.
2021; Chen et al., 2021a; Li et al., 2022; Zeng et al.,
2023; Wu et al., 2023) formulates decision mak- ing problems using sequence modeling by treating actions tokens. Although decision making via sequence modeling has shown impressive perfor- mance downstream tasks, sequence modeling objective does explicitly drive the models learn abstract and updatable representations of en- vironments, which have been shown to essential for solving decision making problems (Yang and Nachum, 2021; Hafner et al., 2020). Thus, in or- der to further understand why decision making via sequence modeling is successful despite minimal inductive biases, in this work examine whether the sequence modeling objective leads to the emer- gence of such representations of environments.
To sum up, present three contributions: (1) we demonstrate that abstract state representations emerge via embodied sequence modeling in BabyAl navigation task; (2) we study the impacts of language instructions in the emergence of such representations, and find that language instructions strengthen such emergence in embodied se- quence modeling; (3) show that model trained with only the initial environment layout learns in- ternal representations of the environment and can perform competitively against model trained with all the available information, which supports the promise of sequence modeling objectives for deci- sion making problems.
Probing Transformer-based architectures. Though sequence modeling has been shown to be successful in language modeling (Raffel et al., 2020; Sanh et al., 2022), image modeling (Doso- vitskiy et al., 2021), 1), and decision making (Janner et al., 2021: Li et al., 2022: Pashevich et al., 2021), remains unknown whether such training results in models that encode rich abstractions or rather simply model surface statistics of their domains. Probing is one standard tool (Tenney et al., 2019; Li et al., 2023; Alain and Bengio, 2016) used to understand the internal representations of neural networks. A probe to predict some salient features from the internal activations of model. If a trained probe can accurately predict a feature of interest, it suggests
2 Related Work
classifier that is trained
Decision Making via Sequence Modeling. In- spired by the success of sequence modeling in lan- guage models, some recent work (Janner et al.,

---

above objective over dataset D = {(T k N which holds set of (trajectory, goal) pairs. Each trajectory 7K includes a sequence of states {s: and actions {a,)==1 that records how the agent suc- cessfully reaches the corresponding goal g where T is the length of a trajectory.
that a representation of this feature is encoded in the internal activations of this model.
.9" k-11
Such work has suggested that sequence model- ing of language tokens leads to representations of abstract concepts such as syntax (Hewitt and Liang, 2019; Tenney et al., 2019) and even non-linguistic notions such as of colors and directions (Patel and Pavlick, 2021; Abdou et al., 2021). Such findings suggest that sequence modeling over embodied do- mains (like we study here) could similarly result in abstractions over states and actions. We utilize probes to recover the intermediate layouts of the environment to understand whether decision mak- ing via sequence modeling can automatically learn grounded representations of the environments.
3.2 BabyAl
We seek to understand whether embodied sequence modeling leads to the emergence of abstract inter- nal representations of the environment. We use the Baby AI (Chevalier-Boisvee et al., 2019) 1) environ- ment as a testbed to investigate this. The goal in this environment is to control an agent to navigate and interact in an N X N grid to follow the given language instructions (e.g., "go to the yellow key"). The state of Baby is a symbolic observation of the grid, where each grid cell represented by two attributes: the type and color of the object. At each timestep, the agent can take an action (c.g., "turn left", "move forward") based on the current state of the grid (Figure 1). Since this environment is synthetic, we can generate arbitrary amount of data and ensure that there is no overlap. This en- vironment also has a fixed mapping between the intermediate layout and pair of initial layout and action sequence, which facilitates our following use of probes to reconstruct the intermediate states.
Othello-GPT (Li et al., 2023) shows that GPT (Radford al., 2019) learned model Oth- ello game moves learns internal representations of board states. That work assumes fixed initial layout, which is different from most real-world tasks. In our work, we use the BabyAl (Chevalier- Boisvert et al., 2019) environment, which considers random initial layouts, to investigate whether se- quence modeling leads to emergence of grounded representations of environments.
3 Embodied Sequence Modeling
3.1 Problem Setup
In this work, consider the following scenario: given a task specified by natural language instruc- tions, an agent needs to predict next actions based on past experiences (e.g. previous states, actions taken) to complete the task. This process can be formulated as embodied sequence modeling, where each input (e.g. states, actions, language instruc- tions) is represented in the form of tokens. The sequence first defines task via a language instruc- tion and the initial environmental layout. During training, it then demonstrates how the task can be completed via a sequence of actions (and optionally the corresponding intermediate states). Formally, we define embodied sequence modeling as next (action) token prediction task:
3.3 Sequence Modeling Setups in BabyAI
In this section, we introduce two sequence mod- eling setups BabyAl, namely the Regular setup and the Blindfolded setup.
The objective is to predict next actions based on all the information at the current and the previous steps provided explicitly including intermediate states 8; which are computed using the Baby environment: P P(orlg, St, ht-1), where 1411 (st, a.)/=l. Note that s; is computed exter- nal to the sequence model by interacting with the BabyAl Al environment given h4-1-
Regular setup:
"Blindfolded" setup:
The objective is to pre- diet next actions based on the language instruc- tion, prior actions, and only the initial state: P(arlg, 81, (a,)i=)) In this case, the current state is not fed to the model explicitly. By comparing task completion performance of this setup against the regular setup, we hope to gain insights into whether these sequence models are able to inter- nally reconstruct the intermediate states implicitly in order to determine the right next actions to take.
max log Po(aelg, St, h+-1) 0
which aims maximize the likelihood of the next action token a, given the state token at current timestep 8t, language instruction tokens g and past experiences h1-1 (s, a.) We train our se- quence modeling framework by optimizing the

---

a
2
linear decoder
Trans former
emb.
+
pos.
enc
g
g
S
BOG
S
S
T-1
go to the red ball
Language Instruetion
state
Action
Figure 2: overview model architecture. Our Transformer- sequence model trains on a sequence of states, actions and language instruction tokens, with the objective of predicting next action tokens given past experiences and the language instruction.
4 Methods
embedding vector. To encode state, we apply two-layer convolutional neural network with same padding, flatten the output feature map across all grid cells, and project it d-dimensional em- bedding with 3-layer MLP. We use sinusoidal embedding to encode positional and temporal in- formation for the language, action, and state tokens. Our sequence model, a Transformer with causal at- tention masks, is trained end-to sod-to-end/n by minimizing the cross-entropy loss between the predicted and ground-truth actions.
4.1
Architecture and Model training
We build our embodied sequence modeling frame- work with Transformer-basee architecture, in which we can process heterogeneous data (e.g. nat- ural language instructions, states and actions) as sequences and autoregressively model the state- action trajectories with causal attention masks.
Figure 2 illustrates the regular setup where the model is given the current state and the full his- tory state-action pairs, along with a natural language instruction, to predict the next action. We denote the input I to our model as tuple (g1sm: $1:7, $1:7, (1:-1) of the goal, state and action sequence, where the goal g is composed of nat- prepend
Prior works (Chen et al., 2021a; Janner et al., 2021) formulate sequential decision making as a sequence modeling problem. They train a policy with GPT architecture (Radford et al., 2019), which shares the same objective (i.c. next action token prediction) ours. Thus, both regular and "blind- folded" sequence modeling (see Section 3.3) can serve as policies conditioned on natural language instructions, by setting include all the observed states and actions up to the current time step We can then evaluate the sequence models using the standard Baby AI decision making benchmarks.
ural language tokens. Additionally, a [BOS] token the action sequence facilitate autoregressive action prediction:
I (9t.... 81, [BOS], 82, a1 .... Sy, aT-1)
To employ a neural network for sequence model- ing, we need to first embed the language, state, and action inputs into the model's latent space of dimen- sion d. We learn two embedding lookup tables to project the language tokens and action inputs into fixed-length vectors of dimension d, respectively. Each state observation is represented by N N grid cells. Each cell is described by the discrete object type and color attributes, which are embed- ded separately and concatenated into a fixed-length
4.2 Probes for Internal Representations of Environmental Layouts
Once we have models trained to perform language-conditionee navigation task, train sim- ple classifiers as probes (Tenney al., 2019; Alain and Bengio, 2016) to explore whether the internal activations of trained model contain representa-

---

Boisvert et al., 2019) to generate trajectories. We select the trajectories that successfully solve the final tasks, and discard the rest.
P-ansfo-ner layer 1el
E11
Complete-State Model simulates the regular setup of sequence modeling in BabyAl (Section 3.3). Thus, this model predicts next actions based on all the information at current and previous steps:
tcanstormez Layer
Prediated State
Bur St
P(ay (a gg. St, h4-1), where 14-1 = (3,, as)i=l
Figure 3: overview of our probing pipeline. We train 2-layer MLP probe to reconstruct intermediate states from the internal activations in our model.
Missing-State Model simulates the "blindfolded" setup. Therefore, this model predicts next actions based on the language instruction, prior actions, and only the initial state: P(algg $1, (a))==1). In practice, we observe this vanilla setup is too chal- lenging as the model lacks sufficient information to understand the language instructions. We therefore choose to additionally provide the last state right before the agent solves the task for 50% of the time during training. During evaluation, we always only pass in initial states.
tions of abstract envirommental states (c.g., current board layouts). If the trained probes can accurately reconstruct the current state when is not explic- itly given to the model (e.g., in the "blindfolded" setup). it suggests that abstract state representations emerge during model's pretraining. As illus- denote the internal activations
trated in Figure of the t-th action token after the I-th Transformer layer as zz. During training, the internal activa- tions are passed into a 2-layer MLP to reconstruct the state S1+1, and the objective is to minimize the cross-entropy loss between the reconstructed and ground-truth states.
5.2 Probing Setup
Data for probing. To train the prohes, we re- use the 1M trajectories used for pretraining the sequence models. We then collect 100K trajecto- ries each for the validation and test to evaluate the trained probes. Following, we pass the trajecto- ries of each split, paired with corresponding goals, to the pretrained and extract internal activations for action tokens. For each input, we randomly sample timestep t and retrieve t-th action token's 's internal activations 2 in each Transformer layer l. To focus whether Missing-State model learns internal representations of intermediate states, we only sample from / otherwise the Missing-Stee model can directly access ground-truth state information through 81.
5 Experimental Setup
5.1 Pretraining Data & Model Variants
Siuplete-/Missingg -Sate models,
Data generation. We select GoToLocal (Figure 1) and MiniBossLevel (Figure Al in Appendix) al., 2019).
levels in Al (Chevalier-Boisver The GoToLocal environment has a layout of 8 grids (including the walls) single room. Each task involves a single goal, namely navigating to the target object. To validate the generality of our observations, we also use the MiniBossLevel en- vironment (see Figure Al in the appendix), which has layout of 9 x 9 grids separated into four rooms. The rooms connected by lockable doors. The tasks are compositional, which may involve solving multiple sub-goals sequentially (e.g., tog- gle the door then pick up an object). In order to solve the tasks in both environments, the sequence model needs to ground the object descriptions in the language instructions into object attributes and their locations in the "grid-world" environment.
Probe metrics. We design three types of metrics to cover different aspects of states: (1) Agent-Loc: Agent locations are represented as discrete (x, y) coordinates on the grid. Location prediction can thus be framed as classification (e.g. x 6 classes for GoToLocal since we ignore the walls), from which we report the classification accuracy. (2) Board- Typed-yype/Cllo:: We measure the probing classi- fiers' performance predicting the object type and color attributes for all cells on the grid. An empty cell has the object type background. GoTooocal has types and colors, while MiniBossI.er has 7 types (door is the extra object) and 6 colors. We report the mean average precision (mAP) over all
For both environments, we create 1M training trajectories respectively, each of which consists of a language instruction, a state sequence, and an action sequence. We use the BOT agent (Chevalier-

---

Random Init.
Agent-Loc
Board-Type
Board-Color
Neighhor-Type
Neighbor- Color
Initial State
n/a
7.4
82.2
80.1
39.8
23.2
10.2
27.4
21.2
47.7
31.7
Complete-Stute Model
97.8
56.6
48.5
84.4
83.4
10.2
19.7
47.8
31.8
Missing-State Model
99.5
64.4
67.4
75.2
75.1
Table I: Probing results comparing Missing- and Complete-State models GoToLocal. We report accuracy for Agent-Loc, and mean Average Precision for Board-Type/-Colee and Weighbor-Tspe/-Cole Type/-Color. For each metric, we show the best performance over Transformer layers. We observe that compared to Complete-Star model which achieves high mAP reconstructing local state information (i.e. Begghbor-Type--Colory Missing-Star State model has a better understanding of global state information (i.e. Agent-Loc and Board- SaandTyype/-Cooor, "Random Init." refers to a randomly initialized sequence model without training.
object types and colors respectively, account for the class imbalance (e.g.. most cells belong to the background). (3)
Results
We first investigate how different input modalities impact the emergence of abstract state represen- tations. We then evaluate model performance on language-condition
We report the same mAP but only consider the 3 x neigh- borhood of the agent. This is to test whether the model can track the change in the neighborhood of agent, which are presumably more important for decision making than other cells the board.
navigation tasks and demon- strate that Missing-Sues -Saate model performs compet- itively when compared to Complete-Stee model which is trained with all available information.
Randomly initialized baseline is probe trained sequence model with random weights. If the trained probes can outperform this bascline, im- plies the effectiveness of embodied sequenee mod- eling on the emergence of state representations.
6.1
Role of Intermediate State Observations in the Emergence of State Representations
We first perform prohing on the internal activations of pretrained Complete-State model and pre- trained Missing-State model. If probe can rately recover certain state properties, we consider the information as being encoded in the internal activations. Note that we expect the internal acti- vations of a pretrained Complete-Stur Saplttte-Sttteeaee model to encode state representations, since it has access to the current states as part of its input.
Initial-state baseline to always predict the ini- tial state, without considering the actions taken by the agent. Intuitively, this baseline would perform strongly for Board-Type and Yoard-Color since only the agent and the objects it has interacted with change from the initial states. However, we expect it to perform poorly on inferring agent location and its neighboring objects.
Results on GoToLocal. As shown in Table 1, after pretraining, both Complete- State model and Missing-State model can significantly outperform randomly initialized baselines probing state properties. For Complete-state model, while the current state is explicitly fed the model as part We train the Transformer model with Adam op- timizer. We choose the set of hyperparameters based on validation performances for the corre- sponding tasks. In the probing experiment, we use 2-layer MLP probe and optimize it with Adam with learning rate of le-4. All experiments, in- cluding Complete- and Missing-State model pre- training and probe training, are performed on single NVIDIA GeForee RTX 3090. More details hyperparameters are elaborated in Appendix A.
5.3 Model Hyperparameters
We instantiate our embodied sequence model with a causal Transformer. It has layers, 8 attention heads, and uses a 768-dimensional hidden size.
of its input, the poor performance of the randomly initialized weights confirms that this information not trivially preserved in its internal activations the model needs to be properly pretrained for that to happen. For Missing-State model, even though the current states are not explicitly given, they can be reconstructed reasonably well from the model's internal representations.
In fact, probes of Missing-Ster State model perform better on reconstructing global state information

---

Agent Location
Board Color
Neighbor
0.54
0.62
0.30
0.56
6.94
0.76
0.56
CampunssSeee Mookl
Mining-Ste Stae Modk!
Figure 4: Probing results of the impact of intermediate observations over Transformer Layers GoTolo
Random Init.
Agent-Loc
Board-Type
Board-Color Neighhor-Typp
Neighhor-
Initial State
n/a
19.1
78.5
85.1
34,9
24.2
V
19.5
26.6
32.5
49.7
35.1
Complete-State Model
99.7
66.9
44.4
98.5
71.9
15.8
26.3
32.5
49.9
35.2
Missing-Statc Model
99.0
61.6
41.1
74.7
47.1
Table 2:
Probing results MiniBossLevvel. We report accuracy for Agent-Loc, and mean Average Precision For each metric, we take the best performance over different
for Bound-Type//Colar and Neighhor-Type//Cone Transformer layers. We observe that Missing-State model performs worse than Complete-State model, but much better than the random baseline. We attribute the performance drop from Complete- State to Missing- State the complexity of MiniBossI environment, which makes the emergence of state representations more difficult.
(i.e.. Lgent-Loc and Board-Type/-Color) than those of Complete-State model. We speculate one pos- sible reason is that, in Complete-Stane model, the prediction of next action could "access" the cur- rent state which is encoded as part of the input sequence, whereas for Missing-state model, the in- ternal activation of the last action has to encode this information. Perhaps related to this hypothe- sis, the internal activations (of the last action) in Complete-State model can focus more on the infor- mation (c.g., the local neighborhood of the agent) critical for the path planning that complements to the global layout information. This may explain why observe probes of Complete-State model achieving better performance on local state infor- mation reconstruction (i.e. Neighhor- Brighhor-TTpe/-Caree Type/-Color).
inferring the intermediate states.
Results on MiniBossLevel. Based on Table 2, we observe that the probes of both pretrained Complete- and Missing-State models can recon- struct the intermediate states reasonably well, in that both outperform the randomly initialized base- lines significantly on the probing performance. We find that Complete-Stee model can capture both cal and global state information better than Missing- State model, which we conjecture might be due to the fact that Complete-Stee the current state at each timestep, and the environ- ment layout is more complex than GoToLocal In both GoCCl and MiniBoss the probes of Missing-Stee State models can effectively recover the termediate states signicantly better than the random baseline, which demonstrates the emergence of ab- stract intermediate states Missing-Statt models.
model has access to
Figure 4 illustrates probing accuracy over dif- ferent Transformer layers. For both models, we observe obvious trends on all metrics. In general, from bottom to top Transformer layers, Complete- State model encode less global state information, and Missing-Stade State model encode less local state in- formation but more global state information. One hypothesis for such trends is that the top Trans- former layers in Complete-State model mainly work as policy for decision making, and those layers in Missing-Sae State model mainly works on
6.2 Role of Language Instructions in the Emergence of State Representations
To study how language instructions impact the emergence of internal representations of states, we pretrain a Complete-State model and Missing- State model without passing language instruc- tions. Since models trained this way cannot per- form iinguage-conditione navigation task, we en-

---

Random Init.
Instruction
Agent-Loc
Board Type
Board-Color
Neighbor Neighbor- Color
State
n/a
n/a
82.2
80.1
23.2
/
10.2
27.4
21.2
47.7
31.7
Complete-State Model
72.2
39.2
20.3
63.5
42.9
V
97.8
56.6
48.5
84.4
83.4
/
24.4
19.7
47.8
Missing-State Model
88.1
44.3
21.4
62.0
40.9
99.5
64.4
67.4
75.2
75.1
Table 3: Probing results showing the impact of language instructions on the emergence of abstract state representa- tions GoCal. We report accuracy for AgentLoc,, and mean Average Precision for Boant-Type/-Coot and Neighbor-TNe Tepe--ppp//ooon For each metric, we take the best performance over different Transformer layers. We observe that training without language instructions leads to performance drop on state information recovery. "Instruction" column specifies whether the probed sequence model trained with language instructions.
sure that both models converge on the training tra- jectories based training loss curves. Next, use the same probing pipeline as Scction 6.1 to explore internal representations of states learned by models trained without language instructions.
Evaluation Benchmark. We create 108 evalu- ation tasks for GoToLocal, and 512 evaluation tasks for MiniBossLevel. each task corresponds to a distinct random seed, which defines the en- vironmental layout and the instruction. The num- ber of training trajectories and the number of eval- uation tasks are comparable to what were origi- nally used by BT To ensure that our evalu- ation tasks test trained model's generalization behavior,
Results. As shown in Table 3, both Complete- State model and Missing-State model trained with- out language instruction can recover state infor- mation significantly better than models without pretraining. This another piece of evidence that sequence modeling is capable of learning internal representations of states. However, when the lan- guage instruction is absent during pretraining, the models struggle recover state information curately the models trained with language in- structions, which reflects that language instructions play an important role in the emergence of internal state representations. We also notice that language instructions are essential for learning object types and colors of the cells. Since this information only occurs in language instructions, the models trained without instructions cannot reconstruct object types colors of the grid cells as accurately as the mod- els trained with instructions. In all, we observe that language instructions are essential for the emer- gence of abstract state representations.
hold out subset of objects and grid sub-region for GoToLocal when creating the training set, and only use them in the evaluation set (see Appendix B.1 for more details). Since Mini- BossLevel's s task space is much larger, we directly perform random sampling when creating the eval- uation set and make sure no duplicate tasks and object layouts appear in the training dataset.
Evaluation Metries. We consider: (1) Success rate: the proportion of tasks that the agent able accomplish within 64 steps. (2) Presstoo-uccess the average number of steps the agent takes when is able to follow a given instruction successfully.
Results on GoToLocal. We train Complete- Missing- State models with 3 random initializa- tions and report the average performance. We also report the average performance of ran- dom policy over seeds. Figure 5a, we ob- serve that Missing-State model performs competi- tively against Complete-State model, even though Missing-Starn State model only has access to the non- initial/last states during its training. This aligns with our expectation that model which learns internal representations of states can utilize such representations to predict next actions.
6.3 Blindfolded Navigation Performance
Experimental results in Sections 6.1 and 6.2 demon- strate that embodied sequence modeling does lead emergence of internal state representations. We further explore the practical implications of the emergence of internal state representations by eval- uating pretrained Complete- and Missing-Ster State models on both GoTal and MiniBossLevvl.
Results on MiniBossLevel. In Figure 5b, we observe that the success rate of Missing-Star

---

35
25
23
Number otspp
20
85-
15
10
5-
Randert
Stmplde-Statz
Missing-State
Randum
Complcte-Stute State
Misscing State
100
I
90
60
40
20
20
Randoes
CompleteSSaatz
Missing-Sute
Complete
State
(a) GoToLocal
(b) MiniBossLLvvel
Figure 5: Evaluation Performance GoToloca (a) and MiniBosslLevel (b). Higher to success suggest better performance. In both levels, Complete-State and Missing-State models can significantly outperform random policy. We also observe that, without intermediate states, Missing-Star State models can still achieve competitive performance against Complete-Stan models on both metries.
and fewer steps
model is lower than that of Complete-Stee -tatte model, while both are significantly higher than the ran- dom policy. This is in line with our probing re- sults on GoTooLocal where Missing-State model is able to reconstruct intermediate states better than the randomly initialized baseline, but worse than Complete-S model. We also notice that the ex- perimental results reflect MiniBossLevel a more challenging environment, as the success rate of Complete-Stare -Saate model MiniBossL is signifi- cantly lower than that of Complete-Stare model in GoToLocal. In all, in both GoToLocal and Mini- BossLeevel, Missing-State models perform competi- tively against Complete-Strre Statte models, which serves further evidence that emergence of internal rep- resentations of states are important for next action prediction emboded sequence modeling.
ing objective a language- task. We find that even when intermediate states (i.e., non-initial/last states) are missing during mod- els' pretraining. the models can still learn to infer the intermediate states, implying that these models automatically learn internal representations of in- termediate states. Based on this observation, we further investigate the impact of each input com- ponent during pretraining on the emergence ternal state representations for sequence modeling. We observe that language instructions are essential for the models to understand basic affordances of objects and for recovering complete state informa- tion. Last, show the importance of emergent state representations with the finding that the model trained without intermediate states performs com- petitively against the model trained with complete state information. We hope our work has provided useful insights of the capabilities of sequence mod- cling for decision making in complex environments and values for future designs of decision making via sequence modeling.
language-conditioned conditioned navigation
7 Conclusion
this paper, we provide a series of analyses on whether and how state representations emerge in models trained with the embodied sequence model-

---

Limitations
Brockman, et al. 2021b. Evaluating large lan- guage models trained on code. arXiv preprint arXiv:2107.088 03374.
We perform sequence modeling pretraining and probing experiments on GoTooLoca and Mini- BossLevel levels. Future work includes evaluating our sequence models in more challenging BabyAT levels or even other environments that simulate real-world settings. Thus, the models need to un- derstand more diverse text-object correlations and infer more involved intermediate states by learning complex internal world model. While we believe methodology can be applied to more challeng- ing environments, our architecture assumes lan- guage instructions as inputs, which may limit the possible choices of environments.
Maxime Chevalier Chevalier-Boisvert, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. 2019. Babyai: A plat- form to study the sample efficiency grounded lan- guage learning. In Proceedings of The 7th Interna- lional Conference on Learning Representations.
Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghami, Matthias Minderer. Georg Heigold. Sylvain Gelly, et al. 2021. An image is worth 16x16 words: Transformers for image recognition at scale. In Proceedings of the 9th International Conference on Learning Representations.
A Acknowledgements
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mo- hammad Norouzi. 2020. Dream to control: Learning behaviors by latent imagination. Proceedings of the 8th International Conference on Learning Repre- sentations.
We would like thank the anonymous reviewers for their detailed and constructive comments. We are very grateful to Xi Chen, Scbastian Goodman, and Radu Soricut for useful discussions. We appre- ciate Kenneth Li for helpful feedback on probing experiments and further conversation about poten- tials of the project. We thank Haotian Fu, Apoorv Khandelwal, Michael Lepori, Calvin Luo, and Jack Merullo for the help the project. Part of the work was done while Tian Yun was a student re- searcher at Google Research. This project in part supported by Samsung Advanced Institute of Technology, and a Richard B. Salomon Faculty Research Award for C.S.
John Hewitt and Percy Liang. 2019. Designing and in- terpreting probes with control tasks. In Proceedings of the 2019 Empirical Methods in Natural Language Processing.
Michael Janner, Qiyang Li, and Sergey Levine. 2021. Offline reinforcement learning as one big sequence modeling problem. In Proceedings of the 35th Ad- vances in Neural Information Processing Systems.
Aitor Lewkowyez, Anders Andreassen, David Dohan, Ethan Dyer, Hemryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al. 2022. Solving quantitative rea- soning problems with language models. Proceed- ings of the 36th Advances in Neural Information Pro- cessing Systems.
References
Kenneth Li, Aspen K Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. 2023. Emergent world representations: Exploring sequence model trained on a synthetic task. Proceedings of the 11th International Conference Learning Representations.
Mostafa Abdou, Artur Kulmizev, Daniel Hershcovich, Stella Frank, Ellie Pavlick, and Anders Søgaard. 2021. Can language models encode perceptual struc- ture without grounding? case study in color. Pro- ceedings of the 25th Conference on Computational Natural Language Learning.
Shuang Li. Xavier Puig, Chris Paxton, Yilun Du, Clin- ton Wang, Linxi Fan, Tao Chen, De-An Huang, Ekin Akyürek, Anima Anandkumar, et al. 2022. Pre- trained language models for interactive decision- making. In Proceedings of the 36th Advances in Neural Information Processing Systems.
Guillaume Alain and Yoshua Bengio. 2016. Under- standing intermediate layers using linear classifier probes. ArXiv, abs/1610.01644.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas. and Igor Mordatch. 202 Decision trans- former: Reinforcement learning via sequence model- ing. In Proceedings of the 35th Advances in Neura! Information Processing Systems.
Alexander Pashevich, Cordelia Schmid, and Chen Sun. 2021. Episodic transformer for vision-and-language navigation. In Proceedings of the IEEE/CVF Inter- national Conference on Computer Vision.
Roma Patel and Ellie Pavlick. 2021. Mapping language models to grounded conceptual spaces. In Proceed- ings of the 9th International Conference on Learning Representations.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg

---

Alec Radford, Jeffrey Wu, Rewon Child, David Luan. Dario Amodei, lyya Sutskever, et al. 2019. Language models are unsupervised multitask learners. OpenAl blog, 188):
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020. Exploring the limits of transfer learning with unified text-to-text trans- former. The Journal of Machine Learning Research, 21(():5485-5551
Victor Sanh, Albert Webson, Colin Raffel, Stephen H Bach, Lintang Sutawika, Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Teven Ie Scao, Arun Raja, al. 2022. Multitask prompted training enables zero- shot task generalization. Proceedings of the 10th International Conference on Learning Represenia- tions.
Ian Tenney, Patrick Xia, Berlin Chen, Alex Wang, Adam Poliak, Thomas McCoy, Najoung Kim, Benjamin Van Durme, Samuel R Bowman, Dipanjan Das, et al. 2019. What do you learn from context? probing sentence structure in contextualized word repre- sentations. In Proceedings of the 7th International Conference Learning Representations.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Proceedings of the 31th Advances in Neural Information Processing Systems.
Philipp Wu, Arjun Majumdar, Kevin Stone, Yixin Lin, Igor Mordatch, Pieter Abbeel, and Aravind Ra- jeswaran. 2023. Masked trajectory models for pre- diction, representation, and control. Proceedings of the 40th International Conference on Machine Learning.
Mengjiao Yang and Ofir Nachum. 2021. Representation matters: Offline pretraining for sequential decision making, In Proceedings of the 38th International Conference Machine Learning.
Zilai Zeng. Ce Zhang, Shijie Wang, and Chen Sun. 2023. condiiiond predictive coding for offline rein- forcement learning. In Proceedings of the 37th Ad- Neural Information Processing Systems.

---

A Hyperparameters
Hyperparameters
Values
# MLP Layers
2
Hyperparameters
Values
Embedding Dim
768
# Transformer Layers
Learning Rate
le-4
Embedding Dim
768
Optimizer
Adam
Symbolic Embedding Dim
32
Epochs
200
# Attention Heads
Batch Size
128
Nonlinearity
GELU
Optimizer
Adam
Table A3: Hyperparameters for Prohing Experiments
3e-5 Complete-Statee Model
Learning Rate
le-5 Missing-Sui -aasin Model 0.1
Dropout
target objects and their locations. Out of 18 color- object combinations, we withhold 3 color-object combinations in training trajectories, but use them in trajectories. Further, 3 color-object combi- nations can be placed anywhere in the grid. The remaining color-object combinationswill not be placed predefined quadrant in training trajec- tories, but is allowed to be placed in this quadrant in test trajectories. Based on the held-out rules, we randomly produce a test set of 108 trajectories, where there are 6 trajectories for each of the 18 color-object combinations (i.e., there are 6 colors, and 3 object types for the target object candidates). This set of evaluation trajectories is thus designed to check for compositional generalizability over object types and object colors.
50 Complete-State Model
Epuchs
100 Missing-Stale State Model
Batch Size
128
Context Size
64
Max Step Threshold
64
Max Episode Length
64
Max Goal Length
Table A1: Hyperparameters for pretraining on GoToLo-
cal
Hyperparameters
Values
# Transformer Layers Embedding Dim
768
Symbolic Embedding Dim
# Attention Heads Nonlinearity
GELU Adam
Optimizer
B.2 Roles of Actions in the Emergence of Abstract State Representations
3e-5 Complete-Statt Model
Learning Rate
le-4 Missing-State Model 0.1
In Sections 6.1 and 6.2, we have looked into the roles of intermediate states and language instruc- tions in emergent internal representations of states for action sequence modeling, What if we take one more step further by removing everything except actions? Under this scenario, learning infer inter- mediate states is quite challenging, since the initial layouts will vary in our language igation task. To deal with this issuc, we fix the initial layout and investigate how much state infor- mation can be recovered from internal activations of models trained with only actions. We hypothe- size language instructions would introduce signals about complete information (i.e., colors and object types) of target objects. To test this, we further explore the scenario when both language instruc- tions and actions are available during the models' pretraining in this fixed layout setting.
Dropout
50 Complete-Start Model
Epochs
150 Missing-State Model
Batch Size
128
Context Size
Max Step Threshold
64
Max Episode Length
64
Max Goal Length
50
nav-
Table A2: Hyperparameters for pretraining on Mini-
BossLevel
B Additional Details and Experiments: GoToLocal
B.1 Evaluation Set Generation
Since GoToLocal enviromment is relatively simple, 1M training trajectories may cover significant fraction of different possible configurations, thus leading to leakage from the training set to the test set and hindering the measurement of generaliza- tion. Thus, we design a held-out set based on the
Setup. We train two models: Action-Oni Oollon-Only model, which observes only actions, and Instruction- Action model, which observes both language in-

---

Language Instruction: "open the red door behind you and put the key next to the red ball"
(A)
(B)
(C)
"Cucn LAC
"Rove FoiNard"
"Pickup Key"
"BUt Ray
OF
Figure A1: A demonstration of al., 2019) where agent (i.e., red triangle) needs complete multiple tasks sequentially. This is example trajectory where the agent achieves the goal of "open the red door behind and put the key next to the red ball": (A) The board with four rooms randomly initialized with agent and several objects. (B) The agent navigates on the board approach and pick up the key. (C) The agent keeps navigating until puts the key next the red ball.
navigation task in BabyAI MiniBoss level (Chevalier- Boisvert
Agent-Loc Loc
Neighbor Neighbor-Occupancy
Srighbor-bbjer Neighbor Color
Initial State
12.1
50.0
32.8
16.7
Action-Only Model
44.0
55,8
42.4
Instruction-Action Model
46.2
88.1
61.3
45.1
Table A4: Probing results showing the impact of both actions and language instructions the emergence of abstract state representations in the fixed layout setting of GoTal. We report accuracy for Agent-Loc,, and mean Average Precision for SeaehhoooTyype-Coooe --Cooo//-Occppancy For each metric, we take the best performance over different Transformer layers. When state information is provided, language instructions can still help models learn internal representations states. Especially, we find the performance gaps are smaller on Board/Neighbor- Socebb-r--ccucccy and grow larger when the information recover is related to object types and colors, which only included in language instructions.
structions and actions. For model pretraining, we generate 10K trajectories on a fixed initial layout in our language-conditionee navigation task, where each trajectory represents path from agent's spawn location a target object. For probing experiments, we generate IK trajectories as test trajectories. We then randomly sample 10 steps from each training trajectory and 10 steps from each test trajectory as the training and test data We focus on recovering the cells in agent's 3 x in Section 5.2) since this is a fixed initial layout setting and recovering the whole grid would be much easier.
information for the cells, but also the occupancy of the cells (i.e., whether a cell is occupied by an object, agent, wall). We denote this metric Btighbor-Occupancy
Results. Based the experimental results shown in Table A4, we observe that when state information is completely missing, model trained with language instructions learns internal represen- tations states better than a model trained without instructions when provided with all types of state information. However, the gaps are small when the state information is related to the basic affordances of objects (i.e., occupancy of cells), but the gaps significantly large when the state information is related to colors and object types of the cells. This aligns with our expectations, since there is no signal containing the color or object type when the model is trained without language instructions. In all, we find that, the fixed layout
3 neighborhood (Neighbor-TType/Coll
Additional probe metrics. To understand the role of language, we not only probe for color/object
"To mitigate the concern of data leakage in training and test sets, we fixed the first 6 actions during the generation of trajectories in test set. With this approach, theoretically, the overlap ratio between training set and test set zr
x 0.14%.

---

Agent-Lor
Bourd-Tye
Bound-Color
Neighbor Type Neighbur-Cood Color
Complete-Siee Model Random Init.
10.2
27.4
21.2
47.7
31.7
Complete- State Model State Encoder
100.0
73.0
78.7
85.5
84,7
Complete-State Model
97.8
56.6
48.5
84.4
83.4
Missing-State Model Random Init.
10.2
24.4
19.7
47.8
31,8
Missing-State Model State Encoder
26.6
41.7
46.6
46.6
33.2
Missing-Sten Statin Model
99.5
64.4
67.4
75.2
75.1
Table A5: Probing results showing the role of Transformer layers in learning abstract state representations on Mouul-Type/-Coblee and Neighbor-
GoTol. report accuracy for Agent-Loc, and mean Average Precision different Transformer layers. We observe that
Type/-Color. For each metric, take the best performance Transformer layers of Complete-State model function as policy, while those of Missing-Stud model function to infer intermediate states." "Random Init." stands for randomly initialized baselines. "State Encoder" stands for state encoder baselines.
scenario, language instructions are beneficial for the emergence of understanding basic affordances of objects, and are particularly essential for the emergence learning grounded information, such colors and object types.
-Color). This suggests that the Transformer block retains local information that is critical for predic- tions, and learns to discard irrelevant information from the encoded states. Note that the state encoder baseline of Complete-Sten -tate model serves as an "up- per bound" since has access to the per-step state information as its input.
B.3 Role of Intermediate State Observations over all Transformer Layers
For Missing-Statt model, the state encoder base- line performs poorly since the state encoder does not have access to the action sequence but only has the initial board information, which is not suffi- cient infer the intermediate states. The gains over the state encoder baseline are consistent for global and local states, indicating that (1) the Transformer block encodes intermediate states; (2) simply car- rying over the initial state to future steps is not sufficient, even for the global board information.
Figure A3. illustrates the probing accuracy different transformer layers on all metrics. We ob- serve that probes trained on Missing-State model can achieve higher accuracy on global state recov- ery (i.e. Agent-Loc and Board- Board-Type/- Type/-Colon) while Complete-State model pays more attention to local state information (i.e. Neighbor-Typer Type/-Coll Color).
B.4 Role of the Transformer Block in Embodied Sequence Models
Finally, we validate if the abstract state representa- tions emerge from the Transformer block, or simply from the encoded representation of the observed states. We establish a state encoder baseline, which is probe trained on the outputs directly obtained from the state encoder (i.e., before passing into Transformer block). Note that the state encoder baseline of Complete-State model takes the cor- responding state for each time step as the input; while that Missing-Statt model always takes in the initial state as the input for any timesteps.
Table AS reports results for state encoder base- lines on pretrained Complete- and Missing-Sue -tate models. We start with observations on Complete- State model: compared to the state encoder base- line, the reconstruction accuracy in Transformer block significantly lower for global board infor- mation -Board-Type -Color), but remains simi- lar for local board information (Neighbor-Tyy Type/ /

---

Agent Location
Agent I Location
1.00
Eammatine wil Inatnction
0.66
0.955
0.99
0.62
0.925
0.900
0.68
0.815
0.850
058
indraction
Color
Board
0.45
0.5
30.0
1135
0.4
0.30
0.3
0.30
Board Type
Board Type
0.1
0.55
0.60
mAP
0.55
0.45
0.90
0.40
Nelghbor Color
Neighbor Color
3.75
0.5
0.7
0.65
0.66
0.55
0.5
0.45
0.6
Neighbor Type
Neighhor Type
1185
0.50
0.72
314
0.75
0.6
0.64
0.65
2
Traraormeraase
Trandomna Laya
Figure A2: Full Probing Results of the Impact of Instructions over Transformer Layers on GoTolocal, Left: Complete State Model. Right: Missing State Model.

---

Agent Location
0.98
Acouasy
0.96
0,91
CompleteSture Model
Miccing- State Medal
5
Board Color
0.65
0.60
0.55
6.90
0.45
0.40
Board Type
0.61
0.62
0.60
0.58
0.56
0.52
Neighbor Color
0.61
6.82
0.80
WAP
0.78
6.74
6.72
Neighbor Type
0.64
6.60
0.78
0.76
6.71
Transfrema Layer
Figure A3: GoToLocal: Probing results of the impact of intermediate states over Transformer Layers on all metrics.


Reviews:
[{"review_id": "vHnLaSYlhC", "reviewer": "Reviewer_89YP", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to 89YP\n\nThank you for your constructive comments and additional references. Here, we summarize your major concerns/questions and will incorporate your comments into our final manuscript.\n\n---\n>**Q1. The underlying motivation can be strengthened in the subsection “A debate about sequence modeling” in Section 2. In particular, Bender and Koller (2020) argue in the context of language modeling, which is different from embodied sequence modeling. I would argue that pretraining RL agents using sequence modeling on the same environment does allow them to capture meaning. I think that by saying, e.g., “we take Bender and Koller’s view for language/pixel modeling, but test this view empirically for embodied sequence modeling”, the motivation can be cleaned up.**\n\nA1. Thank you! We appreciate you pointing this out, and agree that the current wording is not a fair characterization of that debate. We will reword the section titled “a debate about sequence modeling” to read as follows:\n\n*Though sequence modeling has been shown to be successful in language modeling [1, 2], image modeling [3], and decision making [4, 5, 6], it remains unknown whether such training results in models that encode rich abstractions or rather simply model surface statistics of their domains. Recent work has suggested that sequence modeling of language tokens leads to representations of abstract concepts such as syntax [7, 8] and even non-linguistic notions such as of colors and directions [9]. Such findings suggest that sequence modeling over embodied domains (like we study here) could similarly result in abstractions over states and actions.*\n\n*A related debate concerns whether sequence modeling over text is sufficient to encode meaning [10]; however that work is concerned with the dependence on text, not with sequence modeling itself. Thus, our work is related but makes a complementary point, namely, that sequence modeling over non-text data can lead to good abstractions over the non-linguistic world.*\n\n[1] Raffel, Colin, et al. \"Exploring the limits of transfer learning with a unified text-to-text transformer.\" JMLR 2020.\n\n[2] Sanh, Victor, et al. \"Multitask prompted training enables zero-shot task generalization.\" arXiv preprint arXiv:2110.08207\n\n[3] Dosovitskiy, Alexey, et al. \"An image is worth 16x16 words: Transformers for image recognition at scale.\" arXiv preprint arXiv:2010.11929\n\n[4] Janner, Michael, Qiyang Li, and Sergey Levine. \"Offline reinforcement learning as one big sequence modeling problem.\" NeurIPS 2021.\n\n[5] Chen, Lili, et al. \"Decision transformer: Reinforcement learning via sequence modeling.\" NeurIPS 2021.\n\n[6] Li, Shuang, et al. \"Pre-trained language models for interactive decision-making.\" NeurIPS 2022.\n\n[7] Hewitt, John, and Percy Liang. \"Designing and interpreting probes with control tasks.\" arXiv preprint arXiv:1909.03368\n\n[8] Tenney, Ian, et al. \"What do you learn from context? probing for sentence structure in contextualized word representations.\" arXiv preprint arXiv:1905.06316\n\n[9] Patel, Roma, and Ellie Pavlick. \"Mapping language models to grounded conceptual spaces.\" ICLR 2021.\n\n[10] Bender, Emily M., and Alexander Koller. \"Climbing towards NLU: On meaning, form, and understanding in the age of data.\" ACL 2020.\n\n\n---\n>**Q2. Missing some key references.**\n\nA2. We appreciate these additional references and will incorporate them into our camera-ready.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper is timely and sheds light on how a relatively new, and potentially widely used, new RL paradigm works.\n2. The contribution is easy to understand, the paper very readable, and experiments thorough (e.g., the ablation on language instructions), and the results convincing.\n\nOverall, I find that the contribution of this paper is clear, straightforward, and technically sound.", "reasons_to_reject": "Overall, I found this to be a strong paper with a straightforward contribution. However, it is missing some references to past work and its underlying motivation can be made more convincing.\n\n1. The underlying motivation can be clarified. \nThe section \"A debate about sequence modeling\" can be strengthened. In particular, it's not obvious that the view in Bender and Koller (2020) applies to embodied sequence modeling for decision making. Bender and Koller argue that, in the context of __language modeling__, form prediction without external grounding does not lead to meaning. In the language modeling case, however, there is a crucial separation between what the models are trained on (text corpus; form prediction) and how language is __used__ by humans (in the real world; grounded meaning). In the case of embodied sequence modeling, there is __no clear separation__ between the (pre-)training data and the use; I've understood that the sequence modeling pre-training and actual deployment of the model occur in the same environment. Therefore, I would argue that, intuitively, pre-training RL agents using sequence modeling _on the same environment_ does allow them to capture meaning. In addition, the Bender and Koller (2020) argument has since been refuted; meaning doesn't need grounded referents (Piantadosi & Hill, 2023), and language model representations are able to predict external, visual representations (Ilharco et al., 2021). I think that by saying, e.g., \"we take Bender and Koller's view for language/pixel modeling, but test this view empirically for embodied sequence modelling\", the motivation can be cleaned up.\n\n2. Missing some key references \n\n- [Systematic Generalization on gSCAN: What is Nearly Solved and What is Next?](https://aclanthology.org/2021.emnlp-main.166) (Qiu et al., EMNLP 2021)\n\n- [Probing Contextual Language Models for Common Ground with Visual Representations](https://aclanthology.org/2021.naacl-main.422) (Ilharco et al., NAACL 2021)\n\n- [A Benchmark for Systematic Generalization in Grounded Language Understanding](https://arxiv.org/abs/2003.05161) (Ruis et al., NeurIPS 2020)\n\n- [Representation Learning for Grounded Spatial Reasoning](https://aclanthology.org/Q18-1004) (Janner et al., TACL 2018)", "questions_to_authors": "", "paper_topic_and_main_contributions": "While sequence modeling for offline RL is highly performant (Janner et al. 2021a; Chen et al., 2021a), it is yet unclear whether deep neural agents trained in this way form grounded representations of their environment, given they are trained to predict \"surface statistics\" in the sequence modeling objective. The authors investigate whether grounded representations emerge in the sequence modeling task. Here, \"emergence of grounded representations\" is measured by the ability of an MLP probe to predict the current world state from agents' internal activations. They find that (1) embodied sequence modeling indeed gives rise to grounded representations; (2) inclusion of language instructions in agent input strengthens this emergence; (3) agents given only the initial state and its own action history can represent intermediate world states. That is, sequence modeling objectives indeed give rise to grounded representations.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "\n"}, "extra_text": {}}, {"review_id": "uuAshh4tsn", "reviewer": "Reviewer_zk2R", "summary": "", "questions": "What is the dynamics of the emergence of the grounded representation? Do they develop gradually throughout the training? What are the possible benefits of having the grounded representation for this task and could those be quantified (that could be an additional metric that could complement the probing approach that the authors used)? For example, could grounding help with generalization or transfer learning? ", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to zk2R\n\nThank you for your constructive comments. Here, we summarize your major concerns/questions and will incorporate your comments into our final manuscript.\n\n\n--- \n**Q1. The evaluations are relatively limited and restricted to a single environment.**\n\nA1. We agree that evaluating more environments is helpful for supporting our conclusions. In response to your suggestion, we provide additional experiments on a more complex level – MiniBossLevel in BabyAI. MiniBossLevel involves compositional tasks, more language varieties, more objects (e.g, doors), and more actions (e.g., pick up, toggle). We conduct the same experiments (Table 1 and Figure 5) as in our submission and report the results in the tables below.\n\n|                      | Random init. | Agent-Loc | Board-Type | Board-Color | Neighbor-Type | Neighbor-Color |\n|----------------------|:------------:|:---------:|:----------:|:-----------:|:-------------:|:--------------:|\n| Complete-State Model | $\\checkmark$ |    19.5   |    26.6    |     32.5    |      49.7     |      35.1      |\n| Complete-State Model |              |    99.7   |    66.9    |     44.4    |      98.5     |      71.9      |\n| Missing-State Model | $\\checkmark$ |    15.8   |    26.3    |     32.5    |      49.9     |      35.2      |\n| Missing-State Model |              |    99.0   |    61.6    |     41.1    |      74.7     |      47.1      |\n\n\nWe observe that both pretrained Complete- and Missing-State models can reconstruct the intermediate states reasonably well, in that both outperform the _randomly initialized_ baselines significantly on the probing performance. We observe that Complete-State model can capture both local and global state information better than Missing-State model, which we conjecture might be due to that Complete-State model has access to the current state at each time step, and the environment layout is more complex than GoToLocal. Last, we observe that the patterns over Transformer layers also show consistent trends (i.e., increase/decrease over layers) as in GoToLocal (the environment reported in the submission).\n\nWe observe that compared to the probing results shown in GoToLocal, the Missing-State model performs worse than the Complete-State model on MiniBossLevel environment on all the probe metrics. One possible cause of this is the much increased complexity of the environment. To further explore the practical implications of the emergence of internal state representations and to better understand the complexity of the environment, we report the agent success rates of the Complete-State and Missing-State models in the table below:\n\n|  | Number of steps to success | Success Rate (%) |\n|----------------------|:--------------:|:------------:|\n| Random policy | 27.2 (± 6.7) | 0.9 (± 1.2) |\n| Complete-State Model | 18.9 (± 0.9) | 79.6 (± 3.8) |\n| Missing-State Model | 18.1 (± 1.7) | 50.0 (± 2.3) |\n\nWe observe that the success rate of the Missing-State model is indeed worse than the Complete-State model, while both are significantly higher than the random policy. This is in line with our probing results that the Missing-State model is able to reconstruct intermediate states better than the random baseline, but worse than the Complete-State model. Overall, we observe that the MiniBossLevel is a much more challenging environment, as the success rate of the complete-state model is significantly lower than that of the original environment.\n\nWe will update the MiniBossLevel results and visualize the trends over Transformer layers in the appendix of the final version.\n\n\n---\n>**Q2. While the authors justified the probing technique used for the evaluation, it would be interesting to have additional ways to explore the grounding (perhaps some evaluation of transfer learning or visualization of the attention layer).**\n\nA2. Thank you for suggesting complementary evaluation proxies! We agree that both suggested approaches may be alternative approaches to explore grounding-ness, but both approaches require non-trivial designs. For transfer learning in decision making domain, it is non-trivial to find two tasks which can highlight generalizability of the models. For visualization of the attention weights, previous work (Kobayashi et al. 2020) has shown that the attention weights themselves do not necessarily reflect the importance of tokens.\n\nCompared to the above two proxies, probing is a more direct way to study our research question, as we hope to evaluate whether grounded representations emerged from embodied sequence modeling in a quantitative way. In all, we welcome any of your suggestions on design choices for the evaluation of transfer learning and any further discussions.\n\n[1] Kobayashi et al., 2020. Attention is Not Only a Weight: Analyzing Transformers with Vector Norms. EMNLP 2020.\n\n  \n---\n>**Q3. What is the dynamics of the emergence of the grounded representation? Do they develop gradually throughout the training?**\n\nA3. We provide additional experiments to investigate the dynamics of the emergence of the grounded representation. Below tables show the probing results on intermediate checkpoints throughout the pretraining (i.e., 20% means pretraining with 20% of the total training steps). The bold columns are the best performant checkpoints reported in the submission. We observe that for both models the internal state representations emerge rapidly in the early stage of pre-training, and the probing results saturate as early as 20% of the total steps, subject to some minor fluctuations. Because of the difficulty of “blindfolded” setup, the success rate of Missing-State models converges slower than that of Complete-State models. We also notice the interesting trend that for Missing-State models, the emergence of internal state representations happens at 20%, while the success rate keeps increasing. In general, we observe that the emergence of internal grounded representations is positively correlated with successful task completion.\n\n| Complete-State Model | 0% | 20% | **40%** | 60% | 80% | Final ckpt |\n|----------------------|:----:|:----:|:---------:|:----:|:----:|:----------:|\n| Agent Location | 10.2 | 99.1 | **97.6** | 96.3 | 95.1 | 94.5 |\n| Neighbor Type | 47.7 | 84.0 | **84.4** | 84.6 | 84.4 | 84.3 |\n| Success Rate (valid) | 28.7 | 91.7 | **100.0** | 97.2 | 97.2 | 97.2 |\n\n---\n\n| Missing-State Model | 0% | 20% | 40% | 60% | 80% | **Final ckpt** |\n|----------------------|:----:|:----:|:----:|:----:|:----:|:--------------:|\n| Agent Location | 10.2 | 99.8 | 99.7 | 99.6 | 99.6 | **99.5** |\n| Neighbor Type | 47.8 | 74.1 | 74.6 | 74.9 | 75.3 | **75.2** |\n| Success Rate (valid) | 28.7 | 72.2 | 72.2 | 86.1 | 83.3 | **91.7** |\n---\n\n\n>**Q4. What are the possible benefits of having the grounded representation for this task and could those be quantified (that could be an additional metric that could complement the probing approach that the authors used)? For example, could grounding help with generalization or transfer learning?**\n\nA4. Having grounded representations enables the agent to grasp a better understanding of environment layouts (e.g. the obstacles and target objects) given the language instruction, which allows it to plan an optimal path that completes the goals. We agree that “quantifying how much grounding-ness is needed for a task” can help us better understand the importance of grounding. However, to formally study this problem, we would need to define the quantitative metrics to measure the grounding-ness of states and carefully design experiments for evaluations, both of which are non-trivial.\n\nWe also believe that the investigation about the relationship between grounding and generalization or transfer learning is a very interesting direction for our future work, which can further demonstrate the significance of grounded representations for generalizability of embodied sequence modeling.\n\nWe welcome any suggestions on related works and experimental designs for measuring the grounding-ness needed for a task, or for studying the relationship between grounded representations and generalizability of embodied sequence modeling.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The paper is well-written and straightforward. The authors found meaningful differences in terms of grounded representations depending on whether they used instructions or missing/complete state model. ", "reasons_to_reject": "The evaluations are relatively limited and restricted to a single environment. While the authors justified the probing technique used for the evaluation, it would be interesting to have additional ways to explore the grounding (perhaps some evaluation of transfer learning or visualization of the attention layer).  ", "questions_to_authors": "What is the dynamics of the emergence of the grounded representation? Do they develop gradually throughout the training? What are the possible benefits of having the grounded representation for this task and could those be quantified (that could be an additional metric that could complement the probing approach that the authors used)? For example, could grounding help with generalization or transfer learning? ", "paper_topic_and_main_contributions": "This work investigates whether sequence models form representations that are grounded in the environment. Experiments are done on BabyAI environment that combines text instructions and artificial agents that are navigating a gridworld. The authors use a probing technique based on previous work to evaluate the grounding and compare different sequence models (e.g. with and without language).  They found that having instructions was particularly useful for forming grounded representations. Also, they found that the \"missing state model\" (model that predicts the next actions based on the language instruction, prior actions, and only the initial state) learned representations that were more grounded than representations learned with \"complete state model\".\n", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "LdFpfK6uRS", "reviewer": "Reviewer_GDPh", "summary": "", "questions": "- Question A: Does the probe prediction included the empty cell? Thus might the probe always predict empty cell (the majority class) and still achieve a high mAP (what would be the value for this behavior)? Having the exact formula used for the mAP would be helpful (or a reference to the utility packaged used).\n- Question B: Do you have a single probe for each cell on the grid? (Figure 3 might be misleading, a single probe computing the whole layout?)\n- Question C: For the Neighbor prediction, is the agent in the middle of the 3x3 grid?\n- Question D: Do you use partial observation or is the agent given the whole grid at each time-step?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to GDPh\n\nThank you for your constructive review and also the detailed editorial suggestions. We will incorporate your editorial suggestions in the final version. We respond to your major questions and concerns below.\n\n\n--- \n>**Q1. The contribution of this work would be stronger when the evaluation had been done for more than a single environment in the minigrid environment.**\n\nA1. We agree that evaluating more environments is helpful for supporting our conclusions. In response to your suggestion, we provide additional experiments on a more complex level – MiniBossLevel in BabyAI. MiniBossLevel involves compositional tasks, more language varieties, more objects (e.g, doors), and more actions (e.g., pick up, toggle). We conduct the same experiments (Table 1 and Figure 5) as in our submission and report the results in the tables below.\n\n|                      | Random init. | Agent-Loc | Board-Type | Board-Color | Neighbor-Type | Neighbor-Color |\n|----------------------|:------------:|:---------:|:----------:|:-----------:|:-------------:|:--------------:|\n| Complete-State Model | $\\checkmark$ |    19.5   |    26.6    |     32.5    |      49.7     |      35.1      |\n| Complete-State Model |              |    99.7   |    66.9    |     44.4    |      98.5     |      71.9      |\n| Missing-State Model | $\\checkmark$ |    15.8   |    26.3    |     32.5    |      49.9     |      35.2      |\n| Missing-State Model |              |    99.0   |    61.6    |     41.1    |      74.7     |      47.1      |\n\n\nWe observe that both pretrained Complete- and Missing-State models can reconstruct the intermediate states reasonably well, in that both outperform the _randomly initialized_ baselines significantly on the probing performance. We observe that Complete-State model can capture both local and global state information better than Missing-State model, which we conjecture might be due to that Complete-State model has access to the current state at each time step, and the environment layout is more complex than GoToLocal. Last, we observe that the patterns over Transformer layers also show consistent trends (i.e., increase/decrease over layers) as in GoToLocal (the environment reported in the submission).\n\nWe observe that compared to the probing results shown in GoToLocal, the Missing-State model performs worse than the Complete-State model on MiniBossLevel environment on all the probe metrics. One possible cause of this is the much increased complexity of the environment. To further explore the practical implications of the emergence of internal state representations and to better understand the complexity of the environment, we report the agent success rates of the Complete-State and Missing-State models in the table below:\n\n|  | Number of steps to success | Success Rate (%) |\n|----------------------|:--------------:|:------------:|\n| Random policy | 27.2 (± 6.7) | 0.9 (± 1.2) |\n| Complete-State Model | 18.9 (± 0.9) | 79.6 (± 3.8) |\n| Missing-State Model | 18.1 (± 1.7) | 50.0 (± 2.3) |\n\nWe observe that the success rate of the Missing-State model is indeed worse than the Complete-State model, while both are significantly higher than the random policy. This is in line with our probing results that the Missing-State model is able to reconstruct intermediate states better than the random baseline, but worse than the Complete-State model. Overall, we observe that the MiniBossLevel is a much more challenging environment, as the success rate of the complete-state model is significantly lower than that of the original environment.\n\nWe will update the MiniBossLevel results and visualize the trends over Transformer layers in the appendix of the final version.\n\n\n---\n>**Q2. That sequence models are able to learn strong world representations has been already convincingly shown in other work, most famously in the World Models paper of Ha and Schmidhuber (2018) for the Doom and CarRacing environments.**\n\n\nA2. The blindfolded setup proposed by our paper studies the **emergence** of world representations via embodied sequence modeling, given only access to the initial state and the action sequences. In contrast, previous work focuses on _explicitly learning_ world representations with a sequence model, assuming full access to the entire states and action sequences of an agent. For example, World Models (Ha and Schmidhuber, 2018) leverages a variational autoencoder to learn a latent state representation via reconstructing each observed state in the training sequences; it then leverages a recurrent neural network to predict the next state representation given the history information.\n\nIn short, prior work **explicitly learns** state representations and their dynamics via state reconstruction, whereas we study the **emergence** of state representations via next action prediction (for both regular and blindfolded setups), without using any state reconstruction objectives.\n\n  \n---\n>**Q3. I find the results (in Table 1) contradict the conclusion that \"the models can still learn to infer the intermediate states\" (L554). Actually, the results show that the model seems rather sensitive to the state information that is directly passed as an input or not. The Missing-State model achieves higher (non-optimal) scores on \"Board\" information than on the \"Neighbor\" information because this information is given initially and carried over to further timesteps. Similarly, the Complete-State model achieves lower scores at global information than on \"Neighbor\" information because it is highly influenced by the given state at a certain timestep. The more interesting baseline would have been the CNN representations.**\n\n\nA3. We would like to clarify that the conclusion that “the models can still learn to infer the intermediate states” (L554) means that Missing-State models can reconstruct the intermediate states significantly better than the randomly initialized baselines. Plus, we treat Complete-State models’ performance as potential upper bounds, since the models can directly access the ground truth intermediate states.\n\nWe agree that probing the representations output by the CNN state encoder is another interesting baseline. We established the state encoder baselines for both Missing- and Complete-State models, in which the probes are trained on the outputs from the CNN state encoders (i.e., before passing into Transformer blocks). Note that the state encoder baseline of Complete-State model takes as input the corresponding state for each time step; while that of Missing-State model always takes in the initial state input for any time step.\n\nBelow are the tables for state encoder baselines on pretrained Complete- and Missing-State models. Compared with Table 1, for Missing-State model, it performs poorly since its state encoder has no access to the action sequence and solely initial board information is not sufficient to infer the intermediate states. The gains over the state encoder baseline are consistent for global and neighbor states, indicating that simply carrying over the initial state to future steps is not sufficient, even for the global board information.\n\n| | Random Init. | State Encoder Baseline | Agent-Loc | Board-Type | Board-Color | Neighbor-Type | Neighbor-Color |\n|----------------------|:------------:|:----------------------:|:---------:|:----------:|:-----------:|:-------------:|:--------------:|\n| Missing-State Model | $\\checkmark$ | | 10.2 | 24.4 | 19.7 | 47.8 | 31.8 |\n| Missing-State Model | | $\\checkmark$ | **26.6** | **41.7** | **46.6** | **46.6** | **33.2** |\n| Missing-State Model | | | 99.5 | 64.4 | 67.4 | 75.2 | 75.1 |\n\nFor Complete-State model, we observe that the reconstruction accuracy goes down from state encoder baseline to Transformer layers, which reflects that the Transformer layers function as a policy, and they learn to discard irrelevant information from the encoded states. For example, we can see that the global board information (Board-Type / -Color) drops a lot from the state encoder baseline, while the neighbor information remains similar, which implies that Complete-State model learns to localize. Please note that the state encoder baseline serves as an \"upper bound\" since it has access to the per-step state information as its input. We will include and discuss this set of experimental results in the appendix.\n\n| | Random Init. | State Encoder Baseline | Agent-Loc | Board-Type | Board-Color | Neighbor-Type | Neighbor-Color |\n|----------------------|:------------:|:----------------------:|:---------:|:----------:|:-----------:|:-------------:|:--------------:|\n| Complete-State Model | $\\checkmark$ | | 10.2 | 27.4 | 21.2 | 47.7 | 31.7 |\n| Complete-State Model | | $\\checkmark$ | **100.0** | **73.0** | **78.7** | **85.5** | **84.7** |\n| Complete-State Model | | | 97.8 | 56.6 | 48.5 | 84.4 | 83.4 |\n\n\n---\n>**Q4. The paper could be more precise in what it actually targets to measure. Of course the model learns \"some\" useful internal state representation. Are these internal representations the ground truth states? The \"internal representation of a state\" is only measured via the proxy of probes. Are these internal representations properly grounded? The reported results are not specifically related to the given instruction (Can the agent identify the targeted object?).**\n\n\nA4. You raise a very interesting point, to which we are sympathetic and which we hope to explore in follow up work. As we understand it, by “properly grounded”, you are asking whether the signal we are detecting via the probes is causally tied to the input and causally implicated in the model’s behavior, as opposed to being “epiphenomenal” so to speak?\n\nAssuming we understand correctly, we agree that probes are a proxy measure. However, the results of probing experiments are still highly informative as 1) they represent a first step which opens the door for future work (i.e., if the probing results are not positive, there is little reason to pursue more advanced methods) and 2) probing classifiers are currently the only standard, widely used method for inspecting internal states. Newer methods are currently being proposed which support causal intervention which we hope to explore/extend in the future. These methods include, for example, INLP [1] and its successors, gradient-based methods [2], and subnetwork/circuit discovery [3,4]. Each of these has potential but is still under active development. Using such methods in our setting warrants a paper in itself, to address all the concerns about these newer methods. Thus, exploring these ideas is done best as a study that follows first showing that the probing experiments yield positive findings (i.e., the current paper).\n\n[1] Ravfogel, Shauli, et al. \"Null it out: Guarding protected attributes by iterative nullspace projection.\" arXiv preprint arXiv:2004.07667 (2020)\n\n[2] Tucker, Mycal, Peng Qian, and Roger Levy. \"What if this modified that? syntactic interventions via counterfactual embeddings.\" arXiv preprint arXiv:2105.14002 (2021).\n\n[3] Conmy, Arthur, et al. \"Towards automated circuit discovery for mechanistic interpretability.\" arXiv preprint arXiv:2304.14997 (2023)\n\n[4] Lepori, Michael A., Thomas Serre, and Ellie Pavlick. \"Break it down: evidence for structural compositionality in neural networks.\" arXiv preprint arXiv:2301.10884 (2023)\n\n---\n>**Q5. Does the probe prediction include the empty cell? Thus might the probe always predict empty cell (the majority class) and still achieve a high mAP (what would be the value for this behavior)? Having the exact formula used for the mAP would be helpful (or a reference to the utility packaged used).**\n\nA5. Yes, the probe predictions include the empty cells and the empty cells are likely to occupy the majority of the board, which in part makes the output classes imbalanced. Since accuracy would be a misleading metric because of such an imbalance, we use mAP to deal with the skewed distributions of output classes. When the probe always predicts empty cell, it may get high average precision on “Empty Cell” class but much lower scores on the rest of the classes, which eventually leads to a low mAP. Specifically, we use *sklearn.metrics.average_precision_score()* to compute this (i.e., this can be found in the *vlm_rule_learning/model/probe_model.py* of the submitted code as well). We will add this detail to the appendix.\n\n\n---\n>**Q6. Do you have a single probe for each cell on the grid? (Figure 3 might be misleading, a single probe computing the whole layout?)**\n\nA6. No, we use a single probe (2-layer MLP) for the whole board (or 3x3 neighbor grid). More implementation details can be found in *vlm_rule_learning/model/probe_utils.py* of the submitted code.\n\n\n---\n>**Q7. For the Neighbor prediction, is the agent in the middle of the 3x3 grid?**\n\nA7. Yes, for the neighbor prediction, the agent is in the middle of the 3x3 grid.\n\n---\n>**Q8. Do you use partial observation or is the agent given the whole grid at each time-step?**\n\nA8. We use the observation of the whole grid at each time-step.\n\n\n---\n>**Q9. L410: \"use 2 learning rates\" -- How's that?**\n\nA9. In the experiments, we observed that the optimal hyperparameters for Complete- and Missing-State models could be different, mainly in terms of learning rate. Thus, we swept different sets of hyperparameters in the pre-training and used the best model checkpoints for probing experiments. We will add the hyperparameters for each model in the appendix and will release the data for pretraining and probing for reproducibility purpose.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- A novel \"blindfolded\" navigation task that directly aims to test for an agents ability to learn useful intermediate representations of the world (in the minigrid domain).\n- A throughout introspection of probing performances in different layers of the Transformer (Figure 4).\n- A nice ablation that shows that language instructions impact to probe performances (Table 2).", "reasons_to_reject": "- The contribution of this work would be stronger when the evaluation had been done for more than a single environment in the minigrid environment. That sequence models are able to learn strong world representations has been already convincingly shown in other work, most famously in the World Models paper of Hu and Schmidhuber (2018) for the Doom and CarRacing environments. \n- I find the results (in Table 1) contradict the conclusion that \"the models can still learn to infer the intermediate states\" (L554). Actually, the results show that the model seems rather sensitive to the state information that is directly passed as an input or not. The Missing-State model achieves higher (non-optimal) scores on \"Board\" information than on the \"Neighbor\" information because this information is given initially and carried over to further timesteps. Similarly, the Complete-State model achieves lower scores at global information than on \"Neighbor\" information because it is highly influenced by the given state at a certain timestep. The more interesting baseline would have been the CNN representations.\n- The paper could be more precise in what it actually targets to measure. Of course the model learns \"some\" useful internal state representation. Are these internal representations the ground truth states? The \"internal representation of a state\" is only measured via the proxy of probes. Are these internal representations properly grounded? The reported results are not specifically related to the given instruction (Can the agent identify the targeted object?).", "questions_to_authors": "- Question A: Does the probe prediction included the empty cell? Thus might the probe always predict empty cell (the majority class) and still achieve a high mAP (what would be the value for this behavior)? Having the exact formula used for the mAP would be helpful (or a reference to the utility packaged used).\n- Question B: Do you have a single probe for each cell on the grid? (Figure 3 might be misleading, a single probe computing the whole layout?)\n- Question C: For the Neighbor prediction, is the agent in the middle of the 3x3 grid?\n- Question D: Do you use partial observation or is the agent given the whole grid at each time-step?", "paper_topic_and_main_contributions": "The authors hypothesise that instruction-guided sequence modelling for RL leads to world-grounded internal representations within the agents. To test this they propose a \"blindfolded\" navigation task in which the agent is given the state of the environment only initially. The authors assume that the model would have to approximate intermediate states -- that would be usually given during the sequence learning itself -- to be successful in this task.\n\nThe authors compare Transformer-based sequence learners (i) that have seen intermediate states during training with those (ii) that have not seen intermediate states. They evaluate the agent performances with respect to the ability of probes to reconstruct the environment's layout (colours and objects at each cell). They find that (a) models without the intermediate states succeed about the same amount of times as the others, (b) that giving language instructions is improving the metric (mAP) and (c) that probes on the \"intermediate\"-conditioned model are better in predicting agent neighbours and probes on the \"only initially\"-conditioned models are better in predicting global cells.", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "- L010: \"surface statistics fail to generalize\" - I think what is meant is rather the other way around: an agent with a good world model would be able to generalize better. Otherwise I do not fully understand this claim as this is what we usually (successfully) do in deep learning: modelling \"surface\" statistics (the data). But yes, biases in the data are problematic as these might be exploited by the learning algorithm to achieve a low loss.\n- L045: \"an embodied\" agent -- actually none of the cited or related work is about embodied agents (robots in real or virtual environments; the red triangle is maybe the weakest possible proxy for embodiment) and I would rather connect this work with the term \"Offline Reinforcement Learning\". But okey, the paper wants to propose this as a new term in L181.\n- L072: \"pretraining\" -- For me it feels that pre-training and training is used not appropriately. I would would rather frame it as \"training\" the agent and then \"evaluate\" with probing (instead of \"training the probes\" L89). Otherwise L177ff \"During training\" is confusing; same with L277 \"trained end-to-end\". \n- L201: For better understanding, the important term \"layout\" could be connected with the grid formalism. The term \"layout\" is actually never properly defined.\n- L226: \"is computed external\" -- is given by the environment\n- Figure 2: The display of action space is misleading. Actually there exists only turn left, turn right, move forward in minigrid. But the picture shows \"down\".\n- L264ff: The text could be a bit clearer about \"the nature\" of the symbols. I guess here it is color and object type. Ok, this is done later in L380.\n- L305: \"2-layer MLP to reconstruct the current state\" -- This could be more elaborate.\n- L319ff: \"there are 1M trajectories\" -- Where are these coming from? \n- Chapter 5 reads harder than previous chapters. There are multiple syntax errors, missing \"the\".\n- L410: \"use 2 learning rates\" -- How's that?\n- Table 1+2 could be merged (only having Table 2)"}}]

Author rebuttal:

Meta-review: The paper is written well and clear to read. The introduced problem and method are novel, and the experiments are thorough. The related work could be updated with some references suggested by the reviewers. The paper would also be stronger if it considered more than one environment for evaluation.
Decision: Accept-Main