Better World Models Can Lead to Better Post-Training Performance

TL;DR

The paper addresses how explicit world-model objectives shape internal representations and downstream post-training performance in Transformer-based planning tasks. Using a controlled 2x2x2 Rubik's Cube setup, it compares standard next-token training with explicit state-prediction pretraining and joint objectives, and then applies GRPO for post-training. Probing and causal-intervention analyses reveal that explicit world-model training yields more decodable and steerable latent cube states, and that higher world-model quality amplifies gains from GRPO, especially on harder states. The findings suggest that sharpening latent state representations can meaningfully boost planning capabilities in sequence-modeling tasks and inform training pipelines for structured reasoning tasks.

Abstract

In this work we study how explicit world-modeling objectives affect the internal representations and downstream capability of Transformers across different training stages. We use a controlled 2x2x2 Rubik's Cube and ask: (1) how does explicitly pretraining a world model affect the model's latent representations, and (2) how does world-model quality affect the model's performance after reinforcement learning post-training? We compare standard next-token prediction to two explicit world-modeling strategies -- (i) state-prediction pretraining and (ii) a joint state-prediction + next-token objective -- and assess task performance after Group Relative Policy Optimization (GRPO) is applied as post-training. We evaluate the representation quality with linear probes and causal interventions. We find that explicit world-modeling yields more linearly decodable and causally steerable state representations. More importantly, we find that improved state representations lead to higher gains for GRPO, especially on harder cube states. Our results indicate that sharpening state representations can improve the effectiveness of post-training for sequence-planning tasks.

Figures (4)

• Figure 1: Distribution of cube complexities (i.e., number of moves away from being solved).
• Figure 2: Probing Accuracy. Explicitly training a world model (i.e., Pre-Train or Joint Train) leads to improved cube state representations.
• Figure 3: Intervention success rates.Left: A successful intervention means that the model predicts a good move for the alternative target cube state $\mathcal{T}$ instead of the original cube state $\mathcal{S}$. Right: Distributional steering measured as the change in total probability mass on the set of target-good moves (post-intervention minus pre-intervention). Higher intervention success rates suggest that the model relies more on the latent cube state representations. We see that models that were explicitly trained on world modeling have higher intervention success rates.
• Figure 4: Task accuracy after GRPO. Applying GRPO instead of each training strategy (FT, Pretrain + FT, Joint Train) leads to improved results, especially for more complex cube states that require longer rollouts (solid lines vs. dashed lines). Models that were explicitly trained on world modeling also see higher gains from GRPO (orange curves vs. green curves vs. blue curves).