One Lens, Many Worlds : A Capability-Typed Interface for World-Model Interpretability

2026-06-07

AI Summary

Existing world model architectures vary significantly and lack a unified framework for interpretability analysis, necessitating redundant implementations of methods across models. This work proposes WorldModelLens, a general-purpose framework that abstracts core operations—such as encoding, state transition, initial state setup, and sampling—into capability-typed interfaces and supports diverse prediction heads. We introduce universal adapters that, for the first time, integrate both reinforcement learning–based and self-supervised world models into a cohesive interpretability framework. By incorporating hooking and caching mechanisms, our system enables time-indexed activation inspection, imagined trajectory generation, and intervention-based replay, with compatibility across prominent architectures including PlaNet, Dreamer, IRIS, and I-JEPA. This design allows interpretability techniques—such as probing, activation patching, sparse autoencoders, and anomaly analysis—to be implemented once and reused across multiple models, substantially enhancing research efficiency in model interpretability.

Abstract

World models are now built on substantially different computational substrates. Latent recurrent state-space models such as PlaNet and the Dreamer family compress observations into recurrent states; token-based models such as IRIS quantize observations into a learned codebook and predict autoregressively with a transformer; and joint-embedding predictive architectures such as I-JEPA predict in a learned latent space with no pixel decoder. The interpretability methods applied to these models, including probing, activation patching, sparse autoencoders, and surprise analysis, share a common set of primitives, yet they are re-implemented from scratch for each architecture because existing hook-and-cache tooling assumes a transformer language model with no notion of actions, environment steps, or imagined rollouts. We argue that this fragmentation reflects the tooling rather than the models, and that the shared structure of world models is captured by a small typed interface. We present WorldModelLens, an open-source interpretability substrate organized around a capability-typed adapter: every model implements four required methods (encode, transition, initial state, sample) and declares a set of optional heads (decode, reward, continue, actor, critic) through an explicit capability descriptor, so that reinforcement-learning and self-supervised world models are first-class without either imitating the other. A single hook and cache layer exposes time-indexed activations, imagination rollouts, and intervention replay over this interface, allowing each analysis to be written once.

Problem

• world models
• interpretability
• tooling fragmentation
• capability-typed interface
• model-agnostic analysis

Innovation

• world models
• interpretability
• capability-typed interface
• WorldModelLens
• model-agnostic analysis