Control as Inference and Active Inference › From Cost Minimization to Beliefs About Optimality

Recall the picture from Classical Planning and Optimal Control: a Markov decision process with dynamics $p(s_{t+1} \mid s_t, a_t)$ and a cost or reward signal, solved by dynamic programming to produce a policy $\pi(a \mid s)$ that maximizes expected discounted return. The solution is prescriptive. It tells you what to do. The Bellman equation is a command: at each state, take the best action, where "best" is defined with all the force of an argmax operator. There is no room in that formalism for hesitation between two nearly equivalent routes, no place to record that three different behaviors were all pretty good, and no mechanism by which acting stochastically could ever be rational if acting deterministically is available.

Control as inference is descriptive first and prescriptive second. We start with bare dynamics and no special status for actions: a trajectory $\tau = (s_1, a_1, s_2, a_2, \ldots, s_T)$ is simply drawn from a generative model in which states evolve and actions occur according to some vague prior. Then we bolt on one extra piece of evidence: a binary variable $\mathcal{O}_t$ for each timestep that says whether time $t$ was "optimal". Conditioning on the observation that every timestep was optimal converts the generative model into a distribution over near-optimal behavior. Solving the control problem now looks like a familiar statistical task: compute (or approximate) a posterior. Notice what has happened to the direction of reasoning: instead of pushing preferences forward through time to decide what to do, we assume good behavior exists and reason backward to discover what it looks like.

Why go through this indirection? It would be fair to object that we are making life complicated on purpose: we had a working optimization framework, and we wrapped it in probabilistic clothing. The probabilistic formulation adds capabilities the bare optimization framework structurally lacked. Three reasons recur throughout the literature.

• Planning becomes inference tooling. Any approximate posterior machinery, message passing, variational inference, Monte Carlo sampling, becomes a candidate planner. This is a substantial practical dividend. The community has spent decades building fast, scalable, well-understood inference algorithms for graphical models; the moment planning is phrased as posterior computation, that entire library of machinery applies directly to decision-making without modification.
• Exploration becomes automatic. The posterior over optimal trajectories is multimodal: it spreads mass over all behaviors that plausibly earned high return, rather than collapsing onto a single deterministic winner. We will see that this yields maximum-entropy reinforcement learning. Nobody writes "add noise for exploration" into the algorithm; the hedging is intrinsic to conditioning on soft evidence, and its magnitude is set by a single interpretable dial.
• Agents that reason about their own models get a unified currency. As we saw in Bayesian Filtering and Belief States, perception is already posterior inference over latent state. Control as inference makes action the same kind of computation over the same kind of object. Perception and action stop being two modules glued together and become two uses of one calculus, which is precisely the unification that makes world-model-based agents conceptually coherent.

Let us make the construction precise. Fix a finite horizon $T$. The generative model over trajectories has the usual chain structure, with one addition per timestep: an optimality node $\mathcal{O}_t$ that depends on $(s_t, a_t)$ and is treated as observed. This dependence structure matters: $\mathcal{O}_t$ hangs off the state-action pair like a sensor reading, since optimality is judged by what happened at that step alone, mediated through nothing else in the graph.

$$p(\tau, \mathcal{O}_{1:T}) = p(s_1) \prod_{t=1}^{T} p(a_t)\, p(s_{t+1} \mid s_t, a_t)\, p(\mathcal{O}_t \mid s_t, a_t).$$

Read the graph left to right: latent states chain forward, actions are drawn from a (typically uniform) prior, and each pair of state and action emits evidence about its own optimality. Nothing in this graph knows about rewards yet. It describes a random system bumbling through states taking arbitrary actions, exactly the setup from which HMM smoothing begins. The reward enters through one modeling choice, a single sentence linking evidence to reward:

$$p(\mathcal{O}_t = 1 \mid s_t, a_t) \;\propto\; \exp\!\left(\frac{r(s_t, a_t)}{\beta}\right),$$

Here $\propto$ denotes an unnormalized factor: the omitted multiplier is independent of $(s_t,a_t)$ at each timestep. That convention is what lets the product of omitted multipliers be absorbed into one trajectory-level normalizer. The function $r$ is the familiar reward and $\beta > 0$ is a temperature: small $\beta$ means optimality evidence is sharp and demanding, while large $\beta$ tolerates sloppy, high-entropy behavior. Equivalently, $1/\beta$ is the inverse temperature. The terminology comes from statistical physics, where temperature governs how strongly a Boltzmann distribution concentrates on low-energy configurations; here it plays the same role with reward standing in for negative energy. The message-passing section below shows how the exponential form permits backward recursions and how a common maximum-entropy relaxation restores the Bellman machinery from the previous chapter.

Conditioning on all optimality nodes gives the core object of the whole framework:

$$p(\tau \mid \mathcal{O}_{1:T}) = \frac{p(\tau)\, p(\mathcal{O}_{1:T} \mid \tau)}{p(\mathcal{O}_{1:T})} \;\propto\; p(\tau) \exp\!\left(\frac{1}{\beta} \sum_{t=1}^{T} r(s_t, a_t)\right).$$

This is worth pausing on, because the entire chapter lives inside this equation. Under the unnormalized-factor convention above, all state-action-independent multipliers are absorbed into $p(\mathcal{O}_{1:T})$, leaving the dynamics prior tilted by the exponential of trajectory return. High-return trajectories occupy most of the posterior mass; mediocre ones retain some; terrible ones are exponentially suppressed but never eliminated. Because exponential weighting never assigns exactly zero probability to anything, even catastrophically bad trajectories survive with vanishing-but-positive weight, which will matter when we discuss exploration and robustness later. The sum in the exponent is the plain additive return from Part II Chapter 1, which is precisely why reward and optimality evidence can be identified without changing the underlying decision problem. Had we chosen any other encoding, the quantity being traded off across timesteps would not be return, and we would have silently invented a different objective.

One honesty checkpoint before proceeding. Treating an event like "this trajectory had return $R$" as having probability $\exp(R/\beta)$ stretches the word probability. As written, the exponential is an unnormalized weight: it is not probability mass, and calling it a likelihood would require specifying a base measure and normalization. (A continuous density may exceed one, so magnitude alone is not the objection.) The clean return tilt also relies on the omitted factor being independent of state and action; normalizing the binary conditional separately at every $(s,a)$ would generally introduce state-action-dependent constants and change the posterior. The partition function $p(\mathcal{O}_{1:T})$ is mathematically well defined, but it need not be the physical chance of witnessing a repeatable optimality event. The construction survives because we use relative weights within the posterior, and those ratios remain well defined regardless of whether the absolute scale has a physical interpretation. The practical boundary is therefore simple: treat the exponential as an unnormalized tilt, not as a calibrated physical probability, and test the induced ranking under the stochastic outcomes that matter.

Computing the full trajectory posterior by enumerating trajectories is hopeless, since the number of length-$T$ trajectories grows exponentially in both horizon and branching factor; even our humble $5 \times 5$ grid admits $4^{140} \approx 1.9 \times 10^{84}$ action sequences within the rollout horizons we will sample shortly. But we do not need to enumerate anything. A backward message produces the causally conditioned action distribution $p(a_t \mid s_t, \mathcal{O}_{t:T})$ used for control. Full smoothing marginals under $\mathcal{O}_{1:T}$ combine that backward message with a forward message, exactly as HMM smoothing combines both directions without enumerating state sequences. The remarkable fact is what the backward messages turn out to be.

Define dimensionless backward log-messages $\bar Q_t(s,a)=\log p(\mathcal{O}_{t:T}\mid s_t{=}s,a_t{=}a)$ and $\bar V_t(s)=\log\sum_a p(a)\exp(\bar Q_t(s,a))$. Log-space converts products of evidence into sums. For stochastic transitions, exact marginalization gives

$$\bar Q_t(s,a) = \frac{r(s,a)}{\beta} + \log \mathbb{E}_{s'\sim p(\cdot\mid s,a)} \left[\exp\!\big(\bar V_{t+1}(s')\big)\right].$$

The logarithm sits outside the next-state expectation. Replacing this risk-sensitive state aggregation with a linear expectation yields the familiar entropy-regularized Bellman operator used by maximum-entropy reinforcement learning,

$$\bar Q(s,a) = \frac{\mathbb{E}[r(s,a,s')]}{\beta} + \gamma\,\mathbb{E}_{s'}[\bar V(s')], \qquad \bar V(s)=\log\sum_{a'}\exp\big(\bar Q(s,a')\big).$$

When $\gamma=1$, replacing the log-of-expectation state update by a linear expectation is exact for deterministic transitions. The stationary operator displayed here uses $\gamma<1$ later in the grid world, so even with deterministic dynamics it should be read as the fixed point of a discounted maximum-entropy control objective, not as the exact backward-message recursion of the finite optimality graph. Its action aggregation is an actual log-sum-exp, so as $\beta\to0^+$ the induced policy concentrates on hard-max actions; as $\beta$ grows, inferior actions receive more mass. Encoding a discounted return directly in the optimality graph would instead use time-dependent evidence factors such as $\exp(\gamma^{t-1}r_t/\beta)$; the resulting exact recursion carries that time dependence and is not the stationary $\gamma$-on-value operator above.

We can watch this interpolation happen directly. Take four illustrative action values with the spread produced by our grid world, compute the implied Boltzmann weights across a range of temperatures, and plot each weight against its value: as $\beta$ shrinks, mass collapses onto the best action; as it grows, all actions converge toward equal weight.