TLDR: A new research paper explores whether transformer networks can learn propositional STRIPS world models by framing the task as a next token prediction problem. It introduces a specialized “STRIPS Transformer” architecture that can faithfully represent and learn these planning models from positive and negative action traces. The experiments demonstrate high accuracy and the ability to extract interpretable symbolic STRIPS models from the learned transformer parameters, even generalizing to longer, unseen action sequences.
The field of artificial intelligence is constantly pushing boundaries, and a significant area of research revolves around whether large language models (LLMs) and transformer architectures truly understand the world or merely predict the next token in a sequence. A recent paper delves into this fundamental question by exploring how transformers can learn ‘world models’ in a very specific and interpretable setting: propositional STRIPS planning.
Understanding World Models and Next Token Prediction
At its core, true understanding in AI is believed to involve the ability to make meaningful predictions from a model of the world. While LLMs excel at next token prediction, it’s not always clear if this implies a genuine grasp of underlying world dynamics. Previous work, like ‘OthelloGPT’, attempted to test this by training a transformer to predict legal moves in Othello and then examining if the network’s internal representations encoded the game board’s state. The results were intriguing, suggesting some form of state representation was learned.
However, learning world models from next token prediction is challenging because it requires simultaneously learning a state representation and how those states change. While transformers have shown capabilities in this area, their learned representations can sometimes be inconsistent, failing when applied to new tasks.
Introducing STRIPS World Models
This new research, titled From Next Token Prediction to (STRIPS) World Models – Preliminary Results, tackles this problem by focusing on propositional STRIPS world models. STRIPS (STanford Research Institute Problem Solver) is a classic AI planning formalism that defines problems in terms of states (propositions), actions (with preconditions, add effects, and delete effects), and goals. This structured environment provides a concrete setting to test if transformers can learn robust and interpretable world models.
The authors cast the task of learning these STRIPS world models as a supervised next token prediction problem. Here, the ‘tokens’ are actions, and an action is considered valid if the hidden effects of previous actions don’t make its preconditions false. The paper aims to show that a suitable transformer architecture can accurately represent these propositional STRIPS world models and learn them from sequences of valid (positive) and invalid (negative) actions.
The STRIPS Transformer Architecture
A key insight of the paper is the connection between transformers and formal languages. It leverages B-RASP, a boolean version of the RASP language, which provides a logical abstraction of transformers. Crucially, the language of valid action traces in a STRIPS domain is known to be ‘star-free’, a class of languages that B-RASP programs (and thus hard-attention transformers) are expressive enough to recognize.
The researchers developed a specialized ‘STRIPS Transformer’ architecture. This multi-head architecture dedicates one attention head per domain atom (proposition). It uses a differentiable approach, replacing Boolean operations with real-valued tensor operations, making it trainable with gradient descent. The parameters of this transformer directly correspond to the preconditions, add effects, and delete effects of actions in the STRIPS model.
The learning task involves uncovering the propositions and the structure of actions (preconditions and effects) from a dataset of positive and negative action traces. Positive traces consist entirely of applicable actions, while negative traces contain a single inapplicable action at the end. The model is trained to predict whether each action in a sequence is applicable, effectively learning the underlying STRIPS rules.
Experimental Results and Interpretability
The experiments were conducted across five propositional planning domains, including a simple custom domain, grounded Blocksworld versions (with 2 and 3 blocks), and grounded Ferry domains (with 1 and 2 cars). The results were promising: the STRIPS Transformer successfully learned in all domains. As the size of the training dataset increased, the model achieved perfect training accuracy and, significantly, 100% test accuracy, even on traces considerably longer than those seen during training.
This high accuracy indicates that the transformer learned a domain equivalent to the hidden ground-truth model. Furthermore, a crucial aspect of this work is the interpretability of the learned models. The paper shows that the learned real-valued parameters can be binarized (converted to 0s and 1s) to extract an interpretable STRIPS model. If the binarized model correctly classifies all traces, the extracted STRIPS model is equivalent to the true hidden model.
For instance, in the ‘simple’ domain, the recovered STRIPS model matched the ground truth in nine out of ten training runs with 200 samples, and perfectly matched it with 500 samples. This demonstrates that the approach not only learns to predict but also yields a symbolic, human-understandable representation of the world model.
Also Read:
- TransZero: Accelerating AI Planning with Parallel Tree Expansion
- Advancing Robot Generalization Through Preserved Vision-Language Representations
Future Directions
This research provides strong preliminary results for learning propositional STRIPS world models from action traces using transformers. The authors plan to extend this work to learning ‘lifted’ STRIPS models, which involve parameterized actions and objects, a more complex and general form of planning. They also aim to bridge the gap between their transformer-based approach and symbolic algorithms like SIFT, exploring how the robustness of transformers can be combined with the formal guarantees of symbolic methods.
