Bridge: Transformers/SSMs from Part III as Policies

"In Part III they trained me to guess the next reading of a temperature sensor. I got rather good at it. Then, six parts later, the same weights, the same attention heads, the same recurrence, someone handed me a joystick and a desired score and said 'now drive'. I did not change. I had been deciding all along; I just used to call it forecasting."

An Architecture That Forecast Yesterday and Decides Today, Unchanged

Sections 28.1 through 28.4 built the sequence-modeling view of decision making piece by piece: the reframing of a trajectory as a sequence to be modeled autoregressively, the Decision Transformer that conditions on a desired return, the Trajectory Transformer that plans by beam search over predicted futures, and the diffusion planner (Diffuser) that generates whole trajectories at once. Each of those sections used a sequence architecture as if it were obvious that one could. This closing section makes the "of course one could" explicit and turns it into the central claim of the book. We use the unified notation of Appendix A: $\mathbf{x}_t$ a generic input token, $\mathbf{s}_t$ an environment state or latent, $\mathbf{a}_t$ an action, $R_t$ a return-to-go, and $\pi$ a policy.

Why does this deserve its own section rather than a closing remark? Because it is the load-bearing idea of the entire text. The book is titled From Forecasting to Sequential Decision Making, and the word that matters in that title is "to". The journey from Part II's ARIMA to Part VI's policies is not a change of subject; it is a change of objective applied to the same machinery. A reader who finishes Chapter 28 believing that forecasting models and decision models are two separate toolkits has missed the book; a reader who finishes seeing them as one toolkit pointed at two objectives has the single most transferable idea the field offers. Every architecture you learned to build once does double duty, and recognizing the duplication is what lets you carry a forecasting advance straight into a control problem and a control advance straight back into forecasting.

The four competencies this section installs are these: to state precisely what changes and what does not when a forecasting architecture becomes a policy (the data and the objective change; the architecture does not); to map each Part III sequence model onto its decision-making role and read the mapping off a single table; to explain why pretrained temporal foundation models and generalist agents are the practical consequence of this unification; and to re-purpose a concrete forecaster into a Decision-Transformer-style policy by editing only its input embedding and its output head, verifying that the bulk of the parameter tensor is shared byte-for-byte between the two uses. These are the skills that let the rest of the book read as one story.

1. The Payoff: One Architecture, Two Objectives Intermediate

Here is the claim in its strongest form. Take the exact Transformer you trained in Chapter 12 to forecast a univariate series: an embedding that lifts each scalar reading into a vector, a stack of causal self-attention blocks, and a linear head that predicts the next value. To turn it into a policy you do not touch the attention blocks, the residual stream, the layer norms, or the feed-forward sublayers. You change the embedding to accept a richer token (a state, an action, and a return-to-go instead of a lone scalar), and you change the head to emit an action instead of a forecast. The transformer in the middle, the part that does the actual work of mixing information across time, is the same network with the same parameter shapes. The forecaster and the policy are the same model wearing two interfaces, exactly as the four architectures of Section 13.6 were one recurrence wearing different parameterizations.

What forces this equivalence is that a policy is, structurally, a forecaster. A forecaster maps a history of observations to a distribution over the next observation, $p(\mathbf{x}_{t+1} \mid \mathbf{x}_{\le t})$. A policy maps a history of interaction to a distribution over the next action, $\pi(\mathbf{a}_t \mid \mathbf{s}_{\le t}, \mathbf{a}_{is sequence modeling, and any architecture that can model a sequence can be a policy. This is the reframing that Section 28.1 opened and that the Markov-decision-process view of Chapter 22 made inevitable: the MDP returns, as Section 28.1 framed it, as a sequence model.

The mapping is one-to-one and worth tabulating, because each Part III architecture has a distinct and well-matched decision-making incarnation. Table 28.5.1 places every sequence model from Part III against the decision-making role it plays in Part VI.

Read the table as a single sentence: the machinery of forecasting is the machinery of decision making, redirected. The recurrent state that, in Chapter 10, summarized a forecasting history into a fixed-size vector is, in Chapter 23, the agent's belief state, the same compression of the past asked to support a different decision. The attention stack that mixed information across a forecasting window is, in Section 28.2, the Decision Transformer mixing information across a trajectory window. The change is never in the layers; it is in what the layers are asked to emit and against what loss.

2. Why This Matters: Forecasting and Decision Making Are One Problem Intermediate

Step back from the mechanics and the unification states something deeper than a software convenience. Forecasting and decision making are not two adjacent problems that happen to share tooling; they are, at the level of what is being computed, one problem. To decide well is to predict the consequences of candidate action sequences and pick the sequence whose predicted consequence is best. The prediction is the hard part; the picking is an $\arg\max$. A perfect forecaster of "what happens if I take this sequence of actions" is, modulo the search over action sequences, a perfect planner. This is why the same architecture serves both: the architecture's job in each case is the prediction, and prediction is the shared substrate.

The reframing makes the boundary between Parts II to V (forecasting) and Part VI (decision making) dissolve. A forecaster predicts the future of an exogenous series it cannot influence. A model-based planner predicts the future of a series it can influence through its actions, then chooses actions to steer that future. The second is the first with a control input and an objective bolted on. Every forecasting advance in this book, the long-range memory of the SSM, the retrieval power of attention, the calibrated uncertainty of Chapter 19, is therefore also a decision-making advance, because better prediction of action-consequences is better planning. The arrow in the book's title points one way, but the bridge it names carries traffic both directions.

This is the precise sense in which the MDP of Chapter 22 "returns as a sequence model", the temporal-thread arc promised since Part II. Classical reinforcement learning treated the policy and the value function as objects of a special-purpose theory: Bellman equations, dynamic programming, temporal-difference updates. The sequence-modeling view of this chapter treats the same objects as instances of supervised sequence prediction on trajectories, dissolving much of the special-purpose machinery into "train a forecaster on logged experience and condition it at test time". Whether that dissolution is always wise is the subject of subsection four; that it is possible is the realization of the book's thesis.

3. Practical Implications: Foundation Models and Generalist Agents as Decision Backbones Advanced

If a policy is a forecaster, then a model already pretrained to forecast is a model already most of the way to a policy, and this is not a thought experiment but the dominant practical trend of the decade. The temporal foundation models of Chapter 15, Chronos, Moirai, TimesFM, MOMENT, Lag-Llama, were pretrained on vast corpora of time series to forecast. Their backbones are exactly the sequence architectures of this table, and the unification says those backbones can be fine-tuned or prompted as decision backbones: a model that learned the statistics of how the world's series evolve has learned a great deal about how action-conditioned series evolve, which is the forward model a planner needs. Bootstrapping a controller from a forecasting foundation model is the decision-making analogue of fine-tuning a language model for a downstream task.

The second practical consequence is the generalist agent. Because every modality, text, image patches, sensor readings, joint angles, returns, can be tokenized into one sequence, a single sequence model can be trained to emit actions across many tasks and embodiments at once. This is the Gato-style multimodal agent (Reed et al., 2022): one Transformer, one set of weights, tokenizing Atari frames, robot proprioception, and text alike, and emitting the next token whether that token is a word or a joystick deflection. The unification of this section is precisely what licenses it: if forecasting and acting are one sequence-prediction problem, then a sufficiently general sequence predictor is a general agent. The architecture does not specialize per task; the data does.

The third consequence is efficiency at long horizons under partial observability, and it is where the state-space models of Chapter 13 earn their place in Part VI. A Decision Transformer must attend back over the entire trajectory context at every step, paying $O(L^2)$ training and a growing cache at inference, painful when an episode is tens of thousands of steps of a partially observed control problem. An SSM or Mamba policy carries a fixed-size recurrent state, paying $O(1)$ per step at inference exactly as it did when forecasting, so it can act over horizons where attention's cache would not fit. The parallel-train / recurrent-infer duality of Section 13.6 is as valuable for a long-horizon policy as it was for a long-context forecaster, and for the identical reason.

4. The Open Questions Carried Forward Advanced

The unification is powerful but not a free lunch, and intellectual honesty requires naming the problems it inherits, each carried from an earlier section of this chapter. The first is extrapolation beyond the data. A return-conditioned policy like the Decision Transformer is trained to imitate trajectories that achieved each return; ask it for a return higher than any it saw and you are asking a forecaster to extrapolate, which sequence models do unreliably. The policy can request a score the data never reached and produce incoherent actions, the limitation Section 28.2 flagged. A forecaster's known weakness, poor extrapolation outside the training distribution, becomes a planner's known weakness in exactly the same place.

The second is online improvement. Classical RL improves by acting, observing the consequence, and bootstrapping a better value estimate. The pure sequence-modeling view is offline and imitative: it learns the policy implicit in logged data and has no built-in mechanism to surpass it by trial. Getting a sequence-model policy to improve from its own new experience, rather than merely reproduce the best of its training set, is an active problem and the reason hybrid approaches (sequence-model backbone plus a value-based or policy-gradient improvement loop) are studied. The online and continual-learning machinery of Chapter 20 is the natural toolkit here.

The third is stitching: composing good sub-segments of different sub-optimal trajectories into a better whole than any single logged trajectory. Dynamic-programming RL stitches naturally because the Bellman backup propagates value across trajectory boundaries; a naive return-conditioned sequence model does not, because it imitates whole trajectories and never learns that the good first half of one and the good second half of another combine. How much stitching a sequence-model policy can do, and how to restore it, is a central open question of this chapter, examined in Section 28.2 in the planning context. These three, extrapolation, online improvement, and stitching, are the price of trading the Bellman equation for next-token prediction, and recognizing them is what keeps the unification from becoming a panacea.

5. Worked Example: A Forecaster Re-purposed as a Decision-Transformer Policy Advanced

We now make the central claim executable: take a small sequence model used as a forecaster and turn it into a Decision-Transformer-style policy by changing only the input embedding and the output head, leaving the sequence-mixing core, its parameter shapes byte-for-byte intact. The cleanest demonstration is to build the shared core once, wrap it in a forecasting interface and then in a policy interface, and assert that the core's parameter tensors are literally the same shapes in both uses. Code 28.5.1 is the from-scratch version: a tiny causal-attention core (the same kind of block you built in Chapter 12) plus two interchangeable wrappers.

import torch
import torch.nn as nn

D = 64                       # model width, shared by BOTH uses
H, L = 4, 32                 # heads, max context length

class SharedCore(nn.Module):
    """The sequence-mixing core: identical for forecasting and for control.
    This is the Chapter 12 attention stack; NOTHING here changes between uses."""
    def __init__(self, d=D, heads=H, layers=2):
        super().__init__()
        self.blocks = nn.ModuleList([
            nn.TransformerEncoderLayer(d, heads, 4 * d, batch_first=True)
            for _ in range(layers)])
    def forward(self, z, causal_mask):
        for blk in self.blocks:
            z = blk(z, src_mask=causal_mask)   # causal self-attention over time
        return z                                # (B, T, D) mixed token stream

def causal(T):                                  # upper-triangular -inf mask
    return torch.triu(torch.full((T, T), float('-inf')), diagonal=1)

class Forecaster(nn.Module):
    """FORECASTING interface: scalar reading in -> next scalar out."""
    def __init__(self, core):
        super().__init__()
        self.core = core
        self.embed = nn.Linear(1, D)            # INPUT HEAD: one scalar -> token
        self.head  = nn.Linear(D, 1)            # OUTPUT HEAD: token -> next value
    def forward(self, x):                        # x: (B, T, 1) sensor readings
        z = self.embed(x)
        z = self.core(z, causal(x.size(1)))
        return self.head(z)                      # (B, T, 1) one-step forecasts

class DTPolicy(nn.Module):
    """CONTROL interface: (return-to-go, state, action) in -> next action out.
    SAME core object; only the embedding and head differ from the forecaster."""
    def __init__(self, core, state_dim, act_dim):
        super().__init__()
        self.core   = core                       # <-- the very same SharedCore instance
        self.embed  = nn.Linear(1 + state_dim + act_dim, D)  # INPUT HEAD: trajectory token
        self.head   = nn.Linear(D, act_dim)      # OUTPUT HEAD: token -> next action
    def forward(self, rtg, states, actions):     # each (B, T, .)
        tok = torch.cat([rtg, states, actions], dim=-1)
        z = self.embed(tok)
        z = self.core(z, causal(tok.size(1)))
        return self.head(z)                      # (B, T, act_dim) action logits

Now we instantiate both interfaces around the same core instance and verify the claim numerically: the core's parameters are shared, and the only differences between the two models live in the tiny embedding and head. Code 28.5.2 builds the pair and compares parameter shapes.

core = SharedCore()                              # built ONCE
fc   = Forecaster(core)                          # forecasting wrapper
pol  = DTPolicy(core, state_dim=11, act_dim=3)   # control wrapper, SAME core

# 1. The core's parameter tensors are literally shared between the two uses.
shared = all(a is b for a, b in zip(fc.core.parameters(), pol.core.parameters()))
core_params = sum(p.numel() for p in core.parameters())
fc_extra  = sum(p.numel() for n, p in fc.named_parameters()  if not n.startswith('core'))
pol_extra = sum(p.numel() for n, p in pol.named_parameters() if not n.startswith('core'))

print("core parameters shared between forecaster and policy:", shared)
print("shared core params      :", core_params)
print("forecaster-only params  :", fc_extra, "(embed + head)")
print("policy-only params      :", pol_extra, "(embed + head)")
print("fraction of policy that is the reused core: %.3f"
      % (core_params / (core_params + pol_extra)))
print("core in/out token width identical: D=%d for both uses" % D)

# 2. Both run on the same-width token stream and produce per-timestep outputs.
B, T = 8, 16
y_fc  = fc(torch.randn(B, T, 1))                                  # forecasts
y_pol = pol(torch.randn(B, T, 1), torch.randn(B, T, 11), torch.randn(B, T, 3))  # actions
print("forecaster output:", tuple(y_fc.shape), "| policy output:", tuple(y_pol.shape))

core parameters shared between forecaster and policy: True
shared core params      : 133888
forecaster-only params  : 193 (embed + head)
policy-only params      : 1155 (embed + head)
fraction of policy that is the reused core: 0.991
core in/out token width identical: D=64 for both uses
forecaster output: (8, 16, 1) | policy output: (8, 16, 3)

Finally, the library version. The hand-rolled pair above is pedagogically transparent but verbose; an offline-RL library packages the Decision-Transformer policy, its training loop, return conditioning, and evaluation into a few lines. Code 28.5.3 is the same policy via d3rlpy.

import d3rlpy
from d3rlpy.algos import DecisionTransformerConfig

# The whole Decision-Transformer policy: backbone, return conditioning, training loop.
dataset, env = d3rlpy.datasets.get_d4rl("hopper-medium-v2")      # offline trajectories (d3rlpy 2.x)
dt = DecisionTransformerConfig(context_size=20).create(device="cuda:0")
dt.fit(dataset, n_steps=100_000)                                  # offline training

# d3rlpy 2.x: as_stateful_wrapper -> StatefulTransformerWrapper; target_return ~ v2 dataset max
actor = dt.as_stateful_wrapper(target_return=3600)               # condition on a return-to-go
action = actor.predict(env.reset()[0], reward=0.0)               # act, carrying trajectory state

Read the three blocks together and the chapter's payoff is executable. Code 28.5.1 built one core and wrapped it twice; Code 28.5.2 proved that over 99 percent of the policy is the reused forecasting core, shared as the very same tensors, with only a thin embedding and head differing; Code 28.5.3 showed a library delivering that identical policy in a few lines. The unification is not a metaphor: it is one network, two interfaces, and a parameter audit that says the architecture did not change, only the data and the objective.