Control as Inference and Active Inference › Implementation: Planning as Message Passing
World Models HandbookAll chapters Control as Inference and Active Inference › Implementation: Planning as Message PassingOne section · Sections in this chapter6 / 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
We will build the grid-world generative model, run soft Bellman iteration, sample trajectories from the induced policy, sweep the temperature, and finish with the toy salience computation. Every quantity is computed on the fly, and each figure renders from precomputed values only. The correspondence to keep in mind as the code unfolds: transitions is the generative model's dynamics, soft_value_iteration is the discounted maximum-entropy control surrogate motivated by backward message passing, and the softmax policy approximates the corresponding optimality-conditioned action marginal rather than reproducing the exact finite-horizon posterior.
Building the Grid World Model
First the environment: transition tensor, reward vector, start and goal.
import numpy as np
from scipy.special import logsumexp
N = 5 # grid dimension
NS = N * N # number of states
UP, RIGHT, DOWN, LEFT = 0, 1, 2, 3
MOVES = {UP: (-1, 0), RIGHT: (0, 1), DOWN: (1, 0), LEFT: (0, -1)}
START, GOAL = 20, 4 # indices of start and dock tiles
STEP_COST, GOAL_REWARD, SLIP = -0.2, 10.0, 0.15
GAMMA = 0.96
assert (N, START, GOAL) == (5, 20, 4)
assert np.isclose(SLIP, 0.15)
assert np.isclose((1.0 - SLIP + SLIP / 4) ** 8, 0.8875**8)
assert 70 == 5 * 7 * 2 # C(8, 4) monotone shortest routes
def idx(row, col):
return row * N + col
transitions = np.zeros((NS, 4, NS))
rewards_state = np.full(NS, STEP_COST)
rewards_state[GOAL] = GOAL_REWARD
for row in range(N):
for col in range(N):
s = idx(row, col)
if s == GOAL: # dock is absorbing
transitions[s, :, GOAL] = 1.0
continue
for a, (dr, dc) in MOVES.items():
nr, nc = min(max(row + dr, 0), N - 1), min(max(col + dc, 0), N - 1)
transitions[s, a, idx(nr, nc)] += 1.0 - SLIP
for sr_, sc_ in MOVES.values(): # slips scatter uniformly
wr, wc = (
min(max(row + sr_, 0), N - 1),
min(max(col + sc_, 0), N - 1),
)
transitions[s, a, idx(wr, wc)] += SLIP / len(MOVES)Each row of transitions[s, a] is a probability vector, so one matrix-vector product computes expectations over successors. That single property makes the soft Bellman update a one-liner shortly, and it is the reason dense tensor representations remain useful for small discrete worlds even when fancier structures exist. Note the absorbing construction at the goal: the planner explicitly zeros the goal state's immediate self-loop reward, while the simulator enforces the same terminal convention by breaking as soon as the sampled rollout reaches the dock. Both pieces of terminal bookkeeping are necessary to prevent repeated docking rewards.
Soft Bellman Iteration as a Control Surrogate
Now the planner. Dividing rewards by the temperature creates dimensionless soft values when is measured in reward units per nat. This implements the linear-expectation maximum-entropy Bellman surrogate distinguished from exact stochastic message passing above.
def soft_value_iteration(beta, gamma=GAMMA, max_iter=4000):
"""Fixed point of the soft Bellman operator at temperature beta."""
imm = transitions @ (rewards_state / beta) # expected r/beta per (s, a)
imm[GOAL, :] = 0.0 # dock self-loop adds nothing
Q = imm.copy()
V = logsumexp(Q, axis=1)
for it in range(max_iter):
future = transitions @ V # E[V(s') | s, a]
Q_new = imm + gamma * future
V_new = logsumexp(Q_new, axis=1)
delta = np.max(np.abs(V_new - V))
Q, V = Q_new, V_new
if delta < 1e-7 * (1.0 + np.max(np.abs(V))):
break
return Q, it + 1
def policy_from_Q(Q):
"""Stable softmax over actions: the posterior action marginals."""
z = Q - Q.max(axis=1, keepdims=True)
w = np.exp(z)
return w / w.sum(axis=1, keepdims=True)The planner and simulator use the same reward-on-arrival convention. The matrix product in imm averages rewards_state[s_next] over each successor distribution, and the rollout below adds that same successor-state reward after sampling s_next.
Notice what the routine does not contain: no explicit maximization, no argmax, no sorting. The entire prescriptive content of optimal control has been replaced by a smooth, contraction-friendly fixed point. The contraction argument carries over too: because log-sum-exp is a smooth max, it inherits the max operator's monotonicity properties, so the soft Bellman operator still converges geometrically, just as the hard version did in Chapter 4. Sweeping the finite temperature range below shows a sharp posterior at low and progressively flatter, but not necessarily uniform, action probabilities at high .
betas = np.array([0.05, 0.10, 0.25, 0.50, 1.00, 2.00, 4.00])
results = {}
for b in betas:
Q_b, iters_b = soft_value_iteration(b)
results[b] = {"Q": Q_b, "policy": policy_from_Q(Q_b), "iters": iters_b}
b_mid = 0.25
Q_mid = results[b_mid]["Q"]
pol_mid = results[b_mid]["policy"]
V_field = logsumexp(Q_mid, axis=1).reshape(N, N)
modal_prob = pol_mid.max(axis=1).reshape(N, N)
dock_value_expected = np.log(len(MOVES)) / (1.0 - GAMMA)
assert np.isclose(V_field[GOAL // N, GOAL % N], dock_value_expected, rtol=1e-5)solved 7 soft control problems beta= 0.1: start-state action probs = [0.493 0.493 0.007 0.007] beta= 2.0: start-state action probs = [0.269 0.269 0.231 0.231]
At the printed start-state policy assigns each to up and right, or in total to the two dockward moves. At it assigns , a spread below across all four actions. The assertions bind both descriptions to the executed policy rather than to an informal temperature intuition. Nothing about the environment changed between the two prints; only the strength of our claim that the behavior was optimal did.
The figure below overlays the soft state value field with arrows pointing to each state's highest-probability action, rendered at a mid-range temperature where the posterior is informative but visibly soft.
Three readings confirm the theory. Arrows along the top and right walls are bold: the posterior strongly favors dockward moves without ever applying an explicit maximum. Arrows in the interior are lighter: several directions compete there because slipping sideways costs little distance, so multiple actions carry comparable posterior credibility rather than a single winner with token runner-ups. Values rise through the approach states, but the starred absorbing dock is a local dip: its self-loop earns no further reward while the log-sum-exp action aggregation retains an entropy baseline. With four equivalent self-loop actions, the fixed point is in reward-over-temperature units, independent of under this convention—not zero. Because predecessor actions differ in their probability of reaching the dock, this pinned value can also change their soft action probabilities. The heat map and policy should therefore be interpreted under this explicit terminal convention. Darker arrows carry stronger posterior preference; lighter ones mark states where several actions remain credible.
Sampling From the Trajectory Posterior
Exact ancestral sampling from the full trajectory posterior requires a forward pass over the grid accumulating posterior weight followed by a backward pass drawing states, standard HMM smoothing transplanted to this graph: the forward-backward algorithm computes filtering and smoothing messages just as it did for hidden states, with optimality evidence playing the role that observation emissions played before. A useful, much simpler surrogate is to roll out the Markov policy induced by the soft Bellman fixed point in the true environment. Those trajectories approximate ; they are not exact ancestral samples of it, because the update above uses a discounted linear expectation over stochastic successors rather than exact log-message passing over a fixed horizon. For a five-by-five grid the surrogate is informative, but the shortcut should be used with care on longer horizons or more stochastic dynamics. We roll 400 simulated shifts per temperature and tally which tiles got visited.
def roll_out(policy, n_ep=400, max_steps=140, seed=11):
rng = np.random.default_rng(seed)
returns, visits = [], np.zeros(NS)
for _ in range(n_ep):
s, ret = START, 0.0
for _t in range(max_steps):
visits[s] += 1
a = rng.choice(4, p=policy[s])
s_next = rng.choice(NS, p=transitions[s, a])
ret += rewards_state[s_next]
s = s_next
if s == GOAL:
visits[s] += 1
break
returns.append(ret)
return float(np.mean(returns)), visits / visits.sum()
mean_return_low, occ_low = roll_out(results[0.10]["policy"], seed=23)
mean_return_high, occ_high = roll_out(results[2.00]["policy"], seed=29)
stats_rows = []
for b in betas:
ret_b, _occ_b = roll_out(
results[b]["policy"], seed=int(31 + 7 * len(stats_rows))
)
p_nonterm = results[b]["policy"][[s for s in range(NS) if s != GOAL]]
ent_b = float(
-(p_nonterm * np.log(np.clip(p_nonterm, 1e-12, None)))
.sum(axis=1)
.mean()
)
stats_rows.append({"beta": float(b), "return": ret_b, "entropy": ent_b})
betas_arr = np.array([r["beta"] for r in stats_rows])
rets_arr = np.array([r["return"] for r in stats_rows])
ents_arr = np.array([r["entropy"] for r in stats_rows])
interior_states = [
row * N + col for row in range(1, N - 1) for col in range(1, N - 1)
]
interior_fraction_low = float(occ_low[interior_states].sum())
peak_low = divmod(int(np.argmax(occ_low)), N)
peak_high = divmod(int(np.argmax(occ_high)), N)
assert len(betas_arr) == 7
assert np.all(np.diff(rets_arr) < 0)
assert np.all(np.diff(ents_arr) > 0)low-beta (0.1): mean return 8.37, peak occupancy 0.110 high-beta (2.0): mean return 5.99, peak occupancy 0.136 low-beta interior visit fraction: 0.411 peak cells (row, col): low=(4, 0), high=(4, 0) beta | return | entropy -------------------------- 0.05 | 8.35 | 0.404 0.10 | 8.34 | 0.432 0.25 | 8.25 | 0.587 0.50 | 7.95 | 0.844 1.00 | 7.13 | 1.102 2.00 | 5.85 | 1.268 4.00 | 3.43 | 1.347
The occupancy contrast tells the story at a glance, so we render the two fields side by side.
At the route posterior leans toward the short-wall routes, but it is far from a ribbon: with slips dispersing probability uniformly in every direction, roughly 40% of visits land in the interior and all four sides of the perimeter carry comparable traffic. At the ribbon dissolves into mist; the warmest cell is the start corner, where steps accumulate before diffuse wandering reaches the dock, whose tile has relatively low occupancy because episodes terminate there. Neither picture is "correct". They are answers to different questions about how good a shift needs to be before we call it optimal, and choosing is choosing that question. This is perhaps the most practically important sentence in the worked example: practitioners who treat the temperature as a nuisance hyperparameter miss that it encodes an explicit stance on how much performance must be sacrificed for coverage, a stance better made deliberately than defaulted into.
The Temperature Dial
Having characterized two endpoints, let us measure the full tradeoff curve that the dial governs.
The two panels trace one phenomenon from opposite sides. Return erodes steadily as climbs because diffuse posteriors waste steps; entropy climbs because soft posteriors hedge. Read together, they price the exchange: each nat of behavioral diversity costs a measurable slice of return, and the curve tells you the going rate in this warehouse. There is no universally best dial setting. A deployed fleet wants low temperature, since its job is delivery, not data collection; an agent collecting training data, or facing a world model whose errors it has not yet discovered, wants headroom, and the entropy regularization view makes that headroom a hyperparameter with measurable consequences. Next chapter's model-based RL algorithms inherit exactly this knob, dressed as the entropy coefficient of soft actor-critic, and the curves above are precisely the sort of diagnostic used to tune it.
Information as Fuel: A Toy Expected-Free-Energy Comparison
Finally, the promised glimpse of salience, the one ingredient temperature-based schemes lack. Consider a two-layer mystery. Treat the unknown valve configuration as a two-valued model parameter : under , curtain 0's valve succeeds with probability (and curtain 1's with ); under the roles reverse. The robot assigns equal prior probability to and , so the two curtains are tied initially; Strategy A breaks that tie by choosing curtain 0. One optional inspection gesture reveals the parameter instantly, at the cost of consuming that turn, which would otherwise have been a chance to activate a valve. Successes matter later, and each activated turn's outcome can itself shift beliefs slightly, which is why Strategy A gains a little information incidentally as a byproduct of acting.
We score two candidate strategies with a transparent surrogate inspired by two expected-free-energy components: expected negative log likelihood of preferred success outcomes and expected information gain about the unknown model parameter . The toy assumes the inspection itself is unambiguous and does not claim to reproduce the full risk-plus-ambiguity decomposition above.
PRIOR = 0.5 # belief in H0
SUCC_GOOD, SUCC_BAD = 0.9, 0.1
def kl_posterior_vs_prior(posterior_h0):
"""KL( Bernoulli(posterior_h0) || Bernoulli(PRIOR) )."""
p = np.clip(posterior_h0, 1e-12, 1 - 1e-12)
q = PRIOR
return float(p * np.log(p / q) + (1 - p) * np.log((1 - p) / q))
def success_probability(theta_h0, curtain):
"""Success probability for a known hypothesis and selected curtain."""
matched = (theta_h0 and curtain == 0) or (not theta_h0 and curtain == 1)
return SUCC_GOOD if matched else SUCC_BAD
def update_belief(prior_h0, curtain, success):
"""Bayesian update after one valve outcome."""
like_h0 = success_probability(True, curtain)
like_h1 = success_probability(False, curtain)
if not success:
like_h0, like_h1 = 1 - like_h0, 1 - like_h1
evidence = prior_h0 * like_h0 + (1 - prior_h0) * like_h1
return prior_h0 * like_h0 / evidence
# Strategy A: skip inspection, activate favored curtain, then react to the outcome.
succ_first = PRIOR * SUCC_GOOD + (1 - PRIOR) * SUCC_BAD # 0.5
beliefs = {
True: PRIOR * SUCC_GOOD / succ_first, # see success
False: PRIOR * (1 - SUCC_GOOD) / (1 - succ_first),
} # see failure
succ_second = np.mean(
[
max(
h * SUCC_GOOD + (1 - h) * SUCC_BAD,
h * SUCC_BAD + (1 - h) * SUCC_GOOD,
)
for h in beliefs.values()
]
) # 0.82
surprise_A = -np.log(succ_first) - np.log(succ_second)
# Mutual information between theta and the complete two-outcome history.
# The second curtain is chosen adaptively from the first posterior.
gain_A = 0.0
for theta_h0, prior_theta in [(True, PRIOR), (False, 1 - PRIOR)]:
for first_success in [True, False]:
p_first = success_probability(theta_h0, curtain=0)
p_first = p_first if first_success else 1 - p_first
posterior_1 = update_belief(PRIOR, curtain=0, success=first_success)
curtain_2 = 0 if posterior_1 >= 0.5 else 1
for second_success in [True, False]:
p_second = success_probability(theta_h0, curtain=curtain_2)
p_second = p_second if second_success else 1 - p_second
posterior_2 = update_belief(
posterior_1, curtain=curtain_2, success=second_success
)
gain_A += (
prior_theta
* p_first
* p_second
* kl_posterior_vs_prior(posterior_2)
)
g_A = surprise_A - gain_A
reward_A = succ_first + succ_second
# Strategy B: inspect (perfect reveal), then activate informed curtain once.
gain_B = np.log(2.0) # posterior collapses to certainty
surprise_B = -np.log(SUCC_GOOD) # informed activation still fails 10%
g_B = surprise_B - gain_B
reward_B = SUCC_GOOD
comparison = {
"strategy": [
"activate twice (never inspect)",
"inspect once, then activate",
],
"expected_reward": [reward_A, reward_B],
"surprise": [surprise_A, surprise_B],
"info_gain": [gain_A, gain_B],
"expected_free_energy": [g_A, g_B],
}
report_rows = [
f"{name:<30} reward {rew:.3f} | surprise {sup:.3f} | gain {gai:.3f} | G {efe:+.3f}"
for name, rew, sup, gai, efe in zip(
comparison["strategy"],
comparison["expected_reward"],
comparison["surprise"],
comparison["info_gain"],
comparison["expected_free_energy"],
)
]
chosen = comparison["strategy"][
int(np.argmin(comparison["expected_free_energy"]))
]
richest = comparison["strategy"][int(np.argmax(comparison["expected_reward"]))]activate twice (never inspect) reward 1.320 | surprise 0.892 | gain 0.514 | G +0.377 inspect once, then activate reward 0.900 | surprise 0.105 | gain 0.693 | G -0.588 lowest expected free energy : inspect once, then activate highest expected reward : activate twice (never inspect)
The same discrepancy renders cleanly: Strategy A wins on reward yet loses on expected free energy. Strategy B already has lower expected surprise, and its perfect inspection yields more information than the complete two-outcome history under Strategy A, widening the advantage.
strategy_labels = ["A: activate twice", "B: inspect first"]
surprise_bars = np.array(comparison["surprise"])
gain_bars = np.array(comparison["info_gain"])
efe_bars = np.array(comparison["expected_free_energy"])
x_positions = np.arange(len(strategy_labels))
Strategy A accumulates substantial expected surprise from an uninformed first activation while harvesting information incidentally from both outcomes. Strategy B converts one sacrificial turn into a full bit of certainty. Its lower expected surprise already makes its expected free energy smaller; the taller information-gain bar widens that margin even though the reward ledger prefers A.
Read the discrepancy between the last lines slowly, because it is the entire philosophical difference between the two frameworks compressed into one toy. Strategy A, staying home and firing valves, collects more expected reward ( versus ). Yet expected free energy crowns Strategy B for two reasons: its expected surprise is already lower, largely because Strategy A's first activation is an uninformed 50/50 choice, and perfect inspection contributes a full bit of certainty ( nats), more than the information in Strategy A's complete adaptive outcome history. Strategy A does react to its first outcome and therefore selects the better curtain on the second turn; its cost is the uninformed first attempt, not two consecutive choices of the wrong curtain. Strategy B's informed activation achieves the ceiling allowed by physics, since even perfect knowledge leaves the ten-percent failure rate untouched. In the surrogate above, surprises are treated as additive across the turn structure, a deliberate simplification for the sake of transparency, but the selection outcome is robust to those details.
Does this make active inference right and reward maximization naive? No. It makes them disagree on the exchange rate between knowledge and satisfaction. Reinforcement learning prices information explicitly: an inspection is worth it exactly when later decisions improve enough to repay the lost turn, and with a long horizon, Strategy A's follow-up gains close most of the gap, a calculation explored deeply in Model-Based Reinforcement Learning and in Part VII. Active inference fixes the exchange rate internally through the units of free energy itself, which yields sweeping, seemingly gratuitous curiosity in unfamiliar worlds and persistent sampling even at stable optima. Which behavior you want is a design decision about the deployment environment, not a theorem. A factory floor punishes gratuitous curiosity; a research probe earns its budget with it.