Learning Lifted World Models with Transformers

Recent work has studied whether transformers trained on sequential data learn reusable world models. In this paper, we address this question for lifted models, i.e., abstract world models that generalize across different numbers of objects. To this end, we leverage STRIPS, a logic-based formalism for specifying lifted, interpretable world models. We propose two transformer architectures for learning STRIPS models: the Lifted STRIPS (LS) Transformer, with built-in symbolic structure and parameters that map directly to STRIPS; and the Lifted Stick-Breaking (LB) Transformer, an architecture closer to a standard Transformer that replaces softmax with stick-breaking attention and uses randomized object embeddings. The transformers are trained in a supervised manner on action traces under two supervision settings. Under applicability supervision, each action in a sequence is labeled as applicable or inapplicable given the initial state and preceding actions, and no subsequent states are observed. Under final-state supervision, each trace contains an initial state, an applicable action sequence, and the resulting final state, while intermediate states remain unobserved. Experiments on five domains show that both transformers achieve high accuracy and generalize to longer traces and larger numbers of objects. Additionally, the extracted STRIPS models are fully interpretable and can be used for planning with off-the-shelf symbolic planners.