Actions, Inverse Dynamics, and Latent Actions › Action-Conditioned Transitions

Michael BrenndoerferJuly 2, 20269 min read
World Models HandbookAll chapters
Part 6 of 12Learning World Models at Scale
Chapter 32Actions, Inverse Dynamics, and Latent Actions › Action-Conditioned Transitions
Reading section 2 of 10 · Action-Conditioned Transitions · 9 minSection 2 of 10 · 9 min

Book contents

Section 2 of 10 · 9 min read

Actions, Inverse Dynamics, and Latent Actions › Action-Conditioned TransitionsOne section · Sections in this chapter2 / 10 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The core object of this chapter is the action-conditioned transition model. It is the function that says, given a state and an action, how the world evolves.

Action-Conditioned Transition Model

Given a state st∈Ss_t \in \mathcal{S} and an action at∈Aa_t \in \mathcal{A}, the transition model is a conditional distribution T(st+1∣st,at)T(s_{t+1} \mid s_t, a_t) over the next state. In a latent world model, the same structure applies with a latent state hth_t and decoder or encoder in between: T(ht+1∣ht,at)T(h_{t+1} \mid h_t, a_t).

The role of the action is easy to get wrong. It looks like just another input. Stack it next to the state, feed the concatenation into a network, and out comes a prediction. In an ordinary sensor-fusion sense, that is what it is. But the action is conceptually distinct from the state in a way that matters for what you can do with the model afterward.

Consider two features in the state: the color of a traffic light and the steering angle of your car. Both are numbers. But only one of them responds to your decisions. The traffic light is observed. The steering angle is chosen. A world model that treats these as interchangeable covariates can predict the future, but it cannot plan, because it has no way to ask "what happens if I steer left instead of right?" The action input is the entry point for intervention, not just for correlation.

This resembles the distinction between conditioning and intervention in causal inference. Watching a traffic light turn red is not the same as making it turn red. An action-conditioned model exposes the interface needed to compare hypothetical action sequences, but conditioning alone does not make those comparisons causal. They are valid intervention forecasts only to the extent that the action variable has the intended semantics and the training data and model account for relevant confounding, hidden state, and distribution shift. Models that treat actions as decorative inputs fail even the weaker predictive test.

What Makes an Action an Action

An action has a few properties that distinguish it from a state feature.

  • It indexes a decision. A policy or controller chooses it, although a learner may later observe it in a log. Passive video usually does not expose that decision variable.
  • It has an addressable interface. Many engineered systems use a lower-dimensional action than state, such as joint commands or game controls, but structured actions can also be large. What matters for planning is that the interface can be enumerated, sampled, or optimized.
  • It is applied at a specific time. Actions carry a time index and have a rate. The rate may not match the observation rate, and it may not even be constant across time.
  • It is assigned to a transition. In an idealized Markov model, ata_t is the control applied during the transition from sts_t to st+1s_{t+1}. Real actuators can have delay, smoothing, and persistent effects; the state or action history must represent those effects if the model is to remain Markovian.

Some systems blur these lines. The joint torque commanded by a controller is an action; the resulting joint velocity is a state. But velocity is also what the low-level controller reacts to. Where the boundary lies depends on which loop you are modeling. In a hierarchical world model, high-level actions might be subgoals (target positions), and low-level actions might be motor commands. The same physical robot can therefore have several valid action interfaces, and which one you pick determines what kind of planning and learning you can do above it.

This layering mirrors how many real controllers are organized. A high-level planner chooses a waypoint, a mid-level controller produces a reference trajectory, and a low-level servo loop produces torque. Each level sees the level below it as "the action." The same world model could be trained at any of these levels, and the latent action space it discovers will inherit the semantics of whichever level the data was logged at. A model trained on high-rate torque commands has a different action interface from one trained on lower-rate waypoint commands. Its dimension may be larger or smaller depending on the controller, but the semantics and temporal resolution are not interchangeable.

Action Spaces and Representations

The set of possible actions, denoted A\mathcal{A}, comes in several flavors, and each one has consequences for how you build the transition model.

  • Discrete: at∈{1,…,K}a_t \in \{1, \dots, K\}, where KK is the number of discrete actions. Typical for Atari and Minecraft. Represented as a one-hot vector or an index. A categorical policy can treat action selection as classification, while the transition model conditions on the selected category.
  • Continuous: at∈Rma_t \in \mathbb{R}^m, where mm is the action dimension. Typical for robotic torque or velocity control. Represented directly as real numbers. A differentiable model can provide action gradients where it is locally smooth, but continuous spaces remain difficult to cover uniformly.
  • Mixed: some dimensions discrete, some continuous. Common in real interfaces like keyboard-plus-mouse, where the keyboard is discrete and mouse motion is continuous.
  • Structured: sequences, images, or graphs. Auto-regressive interfaces and API calls live here. Here the "action" is a program or a message, and the transition model has to interpret it.
  • Action chunks: sometimes the natural unit is a short sequence at:t+Ha_{t:t+H} rather than a single step, where HH is the chunk horizon. Action Chunking with Transformers predicts such sequences. Chunking reduces the effective decision rate and can represent temporally coordinated behavior, although the best chunk length depends on the data and controller.

The representation matters because scaling and geometry affect optimization. Continuous actions are often rescaled to a fixed range. Discrete actions are one-hot or embedded. Continuous action sequences are sometimes quantized into discrete tokens so the same transformer backbone used for language or images can process them, an idea we will revisit in Part IX: Foundation and World-Action Models.

Practical Structure of Action-Conditioned Models

In practice, most latent world models condition on actions in one of a few standard ways.

  • Concatenation: the action vector is concatenated with the latent state and passed through the transition network. The network then has to learn how to use the action, which is easier when actions explain a large share of the target variation.
  • FiLM / conditioning: the action modulates intermediate activations through learned scale and shift parameters. This feature-wise affine transformation can gate, amplify, suppress, or shift intermediate features.
  • Cross-attention: the action is treated as a token and the transition attends to it. This shows up in transformer world models. Attention lets the model decide, per state, how much to listen to the action, which is helpful when the action's relevance varies with context.
  • Action as an event: instead of conditioning, the action drives a step function on the state, useful for switch-like effects such as discrete modes.

The choice rarely changes what the model can represent, but it changes what it tends to learn, and it changes training stability. Conditioning too weakly, for example by treating the action as one among many inputs, often leads to a model that ignores actions entirely because ignoring them is a good local optimum when most of the variance in the next state comes from the state itself. This local optimum is easy to fall into, because the residual action-driven variance is often a small fraction of the total next-state variance.

This failure mode is worth watching for. A quick test is to shuffle actions across time steps and see whether predictions degrade. If they do not, the model has no measurable predictive dependence on the recorded actions in that evaluation distribution. A second test is to feed alternative actions at the same state and check whether the predicted futures diverge. Divergence shows sensitivity, not causal correctness; agreement with held-out interventions is the stronger test. A third test is to compute the gradient of the predicted next state with respect to the action; near-zero gradients diagnose local insensitivity in differentiable models.

The following diagnostic turns the shuffle test into a visual comparison. It simulates a small linear system, computes next-state prediction error under three conditions, and shows that breaking the action-transition correspondence removes most of the action's predictive benefit.

In[3]:
Code
import numpy as np

rng_demo = np.random.default_rng(123)
dt_demo = 0.1
A_demo = np.array([[1.0, dt_demo], [0.0, 0.93]])
B_demo = np.array([[0.0], [dt_demo]])

n_traj_demo, T_demo = 200, 50
states_demo = np.zeros((n_traj_demo, T_demo + 1, 2))
actions_demo = rng_demo.normal(0.0, 1.0, size=(n_traj_demo, T_demo, 1))
states_demo[:, 0] = rng_demo.normal(0.0, 0.3, size=(n_traj_demo, 2))
process_noise_demo = 0.05
for t in range(T_demo):
    states_demo[:, t + 1] = (
        states_demo[:, t] @ A_demo.T + actions_demo[:, t] @ B_demo.T
    )
    states_demo[:, t + 1] += rng_demo.normal(
        0.0, process_noise_demo, size=(n_traj_demo, 2)
    )

s_t_demo = states_demo[:, :-1].reshape(-1, 2)
s_next_demo = states_demo[:, 1:].reshape(-1, 2)
a_flat = actions_demo.reshape(-1, 1)
a_shuffled = rng_demo.permutation(a_flat)

mse_state_only = float(np.mean((s_t_demo @ A_demo.T - s_next_demo) ** 2))
mse_true_action = float(
    np.mean((s_t_demo @ A_demo.T + a_flat @ B_demo.T - s_next_demo) ** 2)
)
mse_shuffled_action = float(
    np.mean((s_t_demo @ A_demo.T + a_shuffled @ B_demo.T - s_next_demo) ** 2)
)
Out[4]:
Visualization
Bar chart of prediction error with state-only, true, and shuffled actions.
Next-state prediction error under three input conditions. The state-only baseline ignores actions, the true-action model uses the recorded controls, and the shuffled-action model breaks the correspondence between actions and transitions.