CausalWAM: Learning the Causes of Actions in World Action Models

CausalWAM. Learning the Causes of Actions in World Action Models. By Juntao Cheng and four others, from the University of Illinois Urbana-Champaign and UC San Diego. Published on the thirtieth of October.

World action models, or WAMs, learn latent variables from observations and use them for two jobs: predicting how the scene evolves, and generating actions. The two jobs ask for different things. Prediction has to cover the whole scene, while only a portion of the scene causes the actions. A WAM fitted to demonstrations is not, in general, guaranteed to generate actions solely from the latent causes of the demonstrated actions, which the authors call the action-relevant latent variables. It could imitate just as well while entangling those variables with nuisance variables, and fitting the demonstrations alone would not tell the two cases apart.

The authors' response is to separate what causes action from what prediction depends on, and their theorem says when that separation can be recovered. Under a sparsity constraint, the latent variables used to generate actions are an invertible function of the action-relevant latent variables alone. On top of that, the true structure is identified up to a permutation of the latent variables: the right variables are recovered, though possibly in a different order.

The abstract does not walk through the proof, so the argument can only be described by its two ends. At one end sits the sparsity constraint. At the other is a guarantee about the action pathway: whatever the model uses to act depends on the action-relevant variables and on nothing else, up to an invertible transformation. The constraint is what rules out the entangled solutions that imitation by itself would happily accept.

The requirements are deliberately few. The result needs neither sufficient temporal dependence nor any assumption on how the unobserved regimes evolve. It also holds when the action-relevant latent variables are statistically dependent on the rest of the latents, so the nuisance variables do not have to be independent of the causes. The assumption that carries the weight is the sparsity constraint itself.

From there the theory becomes a model. CWAM, or Causal WAM, builds the constraint into a pretrained WAM by adding a sparse routing path to its action head, and leaves the world modeling objective unchanged. Experiments on synthetic data and on real-world robots validate the theory and show improved efficiency in closed-loop success. The prediction side still sees the whole scene. The action side is held, by a sparse path, to the part of the scene that causes action.