Emergence of Abstract State Representations in Embodied Sequence Modeling (2023)

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.

## introduction sequence modeling, the task of predicting the next symbol based on the given context (e.g., words, pixels), is a simple yet versatile objective. leveraging this sequence modeling objective and the transformer @xcite , language models have achieved impressive capabilities on various tasks, such as code writing @xcite , mathematical problem solving @xcite , and question answering @xcite . 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 a sequence modeling problem @xcite . although these models show impressive 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 encode the likelihood of an action based on the given contexts. if the latter is true, we would need to re-examine the underlying reasons for the strong empirical success of the sequence models. however, 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 representations emerge in embodied sequence modeling. we use the babyai @xcite environment to design a "blindfolded" navigation task, where only the initial layout of the environment, a 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 representations and then picking next moves based on these representations. although the model has no prior knowledge about the intermediate layouts of the environment (i.e., non-initial layouts), our probing results demonstrate that a trained model learns internal representations of the intermediate layouts. in this navigation task, there are three components -layouts of the environment, language instructions, and action sequences -each of which can impact the emergence of abstract state representations in embodied sequence modeling. we further raise the question: what role do language instructions play in the emergence of such representations? we ablate with regards to the presence or absence of language instructions during model training, and train probes to recover the intermediate environmental layouts. we observe that language instructions indeed facilitate the emergence of abstract environmental representations. to sum up, we present three contributions: (1) we demonstrate that abstract state representations do emerge via embodied sequence modeling in the babyai navigation task; (2) we study the impacts of language instructions in the emergence of such representations, and find that language instructions can strengthen such emergence in embodied sequence modeling; (3) we show that a model trained with only the initial environment layout learns internal representations of the environment and can perform competitively against a model trained with all the available information, which supports the promise of sequence modeling objectives for decision making problems. ## conclusion in this paper, we provide a series of analyses on whether and how state representations emerge in models trained with the embodied sequence model-ing objective in a language-conditioned navigation task. we find that even when intermediate states (i.e., non-initial/last states) are missing during models' pretraining, the models can still learn to infer the intermediate states, implying that these models automatically learn internal representations of intermediate states. based on this observation, we further investigate the impact of each input component during pretraining on the emergence of internal 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 information. last, we show the importance of emergent state representations with the finding that the model trained without intermediate states performs competitively against the model trained with complete state information. we hope our work has provided useful insights of the capabilities of sequence modeling for decision making in complex environments and values for future designs of decision making via sequence modeling.