Transformers and Neural State-Space Architectures › Causal Sequence Modeling

This section explains how transformers are trained as predictive models of trajectories: the causal mask, the autoregressive factorization of the likelihood, and what a transformer trained this way computes when an agent uses it inside a decision loop.

An action-conditioned predictor must not read its future targets. During training, a trajectory with $T$ observed positions provides $T-1$ adjacent next-observation targets. A causal mask excludes attention from position $i$ to $j>i$, usually by assigning those scores negative infinity. Provided an allowed finite score remains in each row, the masked weights are zero. With aligned targets and no other leakage paths, all positions can be supervised in one layerwise-parallel forward pass. This reuses computation across prediction tasks without removing attention's quadratic pairwise arithmetic.

Formally, the model defines an autoregressive factorization of trajectory likelihood, echoing the probabilistic framing of Part II:

$$p(o_{1:T}, a_{1:T}) = \prod_{t=1}^{T} p(o_t \mid o_{<t}, a_{<t}) \; p(a_t \mid o_{\le t}, a_{<t}),$$

where:

• $o_{1:T}$ and $a_{1:T}$: the full sequences of observations and actions from timestep $1$ to $T$
• $p(o_t \mid o_{<t}, a_{<t})$: the probability mass or density of observation $o_t$ given all earlier observations and actions, the world-model factor
• $p(a_t \mid o_{\le t}, a_{<t})$: the probability mass or density of action $a_t$ given the observations seen up to and including time $t$ and all earlier actions, the policy factor

The chain rule uses the ordering $o_1,a_1,o_2,a_2,\ldots,o_T,a_T$. For $t=1$, the observation factor is the marginal $p(o_1)$. Training both observation and action factors would model environment observations and the data-collecting policy, respectively. However, the paired token $(o_t,a_t)$ used below already contains $a_t$: an action head reading that same token would leak its target despite a causal mask. Joint supervision needs separately ordered observation/action tokens or an action-excluded, appropriately shifted input path. The observation-only objective below does not train action probabilities. A behavior-policy estimate can support off-policy weighting only with reliable propensities and suitable support; learning such a head is not itself a validity guarantee.

A next-observation world model maximizes the expected log-likelihood of observations under the data distribution, estimated as the empirical average over $N$ training trajectories, with teacher forcing:

$$\mathcal{L}(\theta) = \frac{1}{N} \sum_{n=1}^{N} \sum_{t=1}^{T-1} \log p_\theta(o_{t+1}^{(n)} \mid o_{\le t}^{(n)}, a_{\le t}^{(n)})$$

where:

• $\mathcal{L}(\theta)$: the empirical trajectory-average log-likelihood, maximized during training; minimizing $-\mathcal{L}$ is equivalent
• $\theta$: the model parameters
• $N$: the number of trajectories; superscript $(n)$ identifies observations and actions from trajectory $n$
• $T$: the common number of observed positions in this equation, giving $T-1$ adjacent prediction targets
• $p_\theta(o_{t+1}^{(n)}\mid o_{\le t}^{(n)},a_{\le t}^{(n)})$: the next-observation distribution conditioned on the recorded prefix and current action

This is a conditional likelihood, not a variational ELBO. The later simulator uses $T$ for its number of transitions and allocates $T+1$ observation slots, including a zero start marker; its target count therefore differs from this equation's convention.

For continuous observations, $p_\theta$ is a density, not the probability of one exact value. The code below minimizes mean squared error. For a scalar Gaussian observation model with the network output as its mean and a fixed positive variance, negative log-likelihood differs from squared error only by a positive constant factor and an additive constant. With the same targets and weighting, the objectives therefore have the same minimizers; this interpretation does not mean that the code learns or calibrates a predictive variance.

Teacher forcing conditions on recorded observations; a recursive open-loop forecast under fixed candidate actions conditions on its own later predictions. Closed-loop control is different: it uses newly observed environment feedback to choose or revise actions. Identical one-step errors need not produce identical rollout errors. To see why, consider the local error recursion $e_{t+1}\approx J_t e_t+\delta_t$, where $J_t$ propagates an existing error and $\delta_t$ is a new disturbance or model mismatch. In the idealized scalar unit-gain additive case with zero initial error, independent zero-mean disturbances with common finite variance $\sigma^2$ give root-mean-square error $\sigma\sqrt H$; a constant bias $b$ instead gives signed error $Hb$. Contractive propagation can bound either effect; unstable propagation can amplify it. Feeding predictions back can introduce dependent errors, but neither correlation nor monotonic growth is inevitable.

There is a second, quieter gap in the same objective: it is an average over the data-collecting distribution, so it weights trajectories exactly as the logging policy produced them. A planning agent, by contrast, will steer its rollouts toward states it finds interesting or valuable, which may be rare or absent in the training data. Predictive fidelity on the data distribution therefore bounds neither fidelity nor uncertainty quality under the agent's own exploration distribution, a fragility we return to in the limitations section.

A causal transformer produces a learned predictive representation, not automatically a Bayesian belief or a statistically sufficient state. Let $\mathcal H_t=(o_{\le t},a_{<t})$ denote available history and $b_t=p(s_t\mid\mathcal H_t)$ the Bayesian posterior in a specified POMDP. Under its transition and observation assumptions, this belief supports prediction for supplied future actions. A learned feature $r_t=g_\theta(\mathcal H_t,a_t)$ always determines its own prediction head, but statistical sufficiency would additionally require $p(o_{t+1}\mid\mathcal H_t,a_t)=p(o_{t+1}\mid r_t)$ for the true data-generating distribution. Finite context, lossy representations, limited capacity, or unsuccessful optimization can break that equality. Retained token representations provide addressable memory from which the network computes features; they are neither the full raw history nor the posterior by definition.

This view is what makes sequence-modeling world models work at all:

• Trajectory Transformer (Janner et al., 2021) discretizes trajectory components including observations, actions, and rewards and fits an autoregressive model. A separate beam-search procedure uses that model for trajectory optimization; action proposals and reward predictions are not the whole planning algorithm.
• Decision Transformer (Chen et al., 2021) produces actions from desired return-to-go and state/action history. It is a return-conditioned policy, not necessarily an environment-transition predictor or an online search planner.
• VideoGPT uses autoregressive prediction of learned video tokens, while Gato applies sequence modeling to multimodal policy outputs. Shared sequence machinery is useful, but modality-specific representations, data, and objectives remain important.

A planner knows each candidate open-loop action sequence it proposes; a feedback policy instead chooses later actions in response to later observations. An action-conditioned model can compute a forecast for either kind of query, but observational conditioning is not automatically a valid intervention model. Appropriate action support and causal assumptions are needed, for example, actions randomized independently of hidden disturbances, or a history sufficient to remove action-outcome confounding. If a logging policy uses hidden information unavailable to the model, its action may reveal that information rather than cause the observed outcome. Our simulator uses independent random controls and a known transition mechanism. System identification likewise may use deliberately designed excitation inputs, not only inputs selected by someone else.