23 · Controller Abstractions

Part IV — State Uncertainty

Executed in: Julia · Book: ADM Ch. 23 · Lecture: to be added

Big picture

Where this chapter fits

Chapters 20–22 planned over beliefs: to act, the agent kept a probability distribution over states and updated it after every observation. This chapter represents a policy differently, as a finite state controller: a small graph whose nodes each prescribe an action, and whose edges say which node to move to after each observation. The node is the agent’s entire memory. There’s no belief to update, and a controller with a handful of nodes can run forever.

The question: how much memory does a good policy need, and how do we find the best policy with a fixed amount of it?

The short answer: evaluate a controller by solving linear equations, then improve it by adding nodes or by climbing a gradient.

  • Evaluation is exact and cheap: a controller and a POMDP together form a Markov chain on (node, state) pairs.
  • Policy iteration adds nodes that improve on the current ones, and prunes nodes that are no better than others.
  • Nonlinear programming and gradient ascent fix the number of nodes and optimize the action and transition probabilities directly.
ch23 ctrl Controller nodes, actions ψ, successors η eval Evaluation linear equations on (node, state) ctrl->eval with the POMDP alpha Alpha vectors one per node eval->alpha pi Policy iteration add and prune nodes alpha->pi nlp Nonlinear programming fixed size alpha->nlp grad Gradient ascent fixed size alpha->grad better Better controller evaluate again, repeat pi->better grow nlp->better solve grad->better climb
Figure 1: Map of this chapter. A controller is evaluated exactly on the chain of (node, state) pairs; its nodes’ values are alpha vectors. Policy iteration grows the controller; nonlinear programming and gradient ascent optimize one of fixed size.

The key ideas at a glance

  1. A controller is a policy with a finite memory. Each node \(x\) has an action distribution \(\psi(a \mid x)\) and, for each action and observation, a successor distribution \(\eta(x' \mid x, a, o)\). Conditional plans (Chapter 20) are controllers shaped like trees; a controller can loop back, so it can act forever with a few nodes.

  2. Evaluation is a linear system. Pair the controller’s node with the hidden state. The pair \((x, s)\) moves as a Markov chain, so the value \(U(x, s)\) solves one set of linear equations. Each node’s values over states form an alpha vector, so a controller’s value over beliefs is piecewise linear and convex, just as in Chapter 20.

  3. Policy iteration grows the controller. Try every deterministic node “take \(a\), then go to these nodes”: keep the ones that are better somewhere, and let a new node replace any old node it beats in every state. Replacing a node reroutes every link into it, which is how loops form.

  4. Or fix the size and optimize. With \(n\) nodes, the best controller solves a nonlinear program over \(\psi\) and \(\eta\). Gradient ascent climbs the same objective, projecting back onto probability distributions after each step. Both can stop at local optima.

  5. Memory is the limiting resource. How many nodes a good policy needs depends on the problem. For the crying baby, 2 nodes are optimal. For the fork, anything under 4 nodes can do no better than waiting.

Vocabulary

Table 1: Terms used on this page.
Term Plain meaning
Finite state controller A policy given by a graph of nodes, with actions at nodes and observation-labeled edges
Node \(x\) The controller’s memory state (not the environment’s state)
Action distribution \(\psi(a \mid x)\) How the controller chooses its action at node \(x\)
Successor distribution \(\eta(x' \mid x, a, o)\) Which node comes next, after action \(a\) and observation \(o\)
Deterministic controller One action per node and one successor per observation
Dominated node A node whose value is no better than another node’s in every state
Projection onto the simplex The closest probability distribution to a vector

After this chapter you can

  • Read a controller as a policy, run it, and evaluate it exactly by solving linear equations.
  • Explain how each node’s value is an alpha vector, and choose the best starting node for a belief.
  • Run policy iteration for controllers, and explain why pruning with pointwise dominance can leave many nodes.
  • Set up the fixed-size problem as a nonlinear program, and optimize it by projected gradient ascent.
  • Recognize when a controller needs more memory, and when an optimizer is stuck at a local optimum.

Intuition

TipAnalogy

A night nurse doesn’t compute the probability that the baby is hungry. They follow a rule on a card: “If the baby cries, feed it. After a feed, ignore the next cry, it’s probably just fussing.” The card has two states of mind, watching and just fed, and a rule for moving between them. That card is a finite state controller. It’s easy to follow, needs no arithmetic, and, for this problem, it is exactly optimal.

We use the two problems of Chapter 21. The crying baby has an exact solution to compare against. The fork (problem gallery) asks a hiker to pay for directions: its optimal value from 50–50 is 27.115.

A policy that remembers just enough

The two-node rule above, “feed on crying, but not right after a feed”, is worth -24.675 from the 50–50 belief, starting with a feed: the exact optimum. The simpler rule “feed whenever the baby cries” is worth -25.423, 0.75 less. It wastes feeds on the false cries that follow a feed, since a sated baby still cries 10% of the time. See Figure 3.

Values are alpha vectors again

Each node’s values in each state form an alpha vector. A controller started at the best node for a belief earns the upper surface of these lines, as in Chapter 20. For the two-node rule, that surface is the optimal value, to within \(10^{-6}\) at every belief: the exact solution of Chapter 20 has two alpha vectors, and they are this rule’s two nodes. See Equation 1.

Count what you’ve heard

At the fork, the optimal policy asks until the answers lead by two, then goes. A controller does that by counting: 3 asking nodes (even, one ahead “right”, one ahead “left”) and 2 going nodes. With fewer than 4 nodes, no deterministic controller does better than waiting forever, because it can’t remember enough answers to act on them safely. The best 4-node controller, found by trying all 16,777,216 of them, earns 15.33; 5 nodes reach the optimum. See Figure 4.

Growing a controller

Policy iteration starts with one node and keeps adding nodes that improve on the current ones. For the crying baby, it reached the optimal value at 50–50 after 2 iterations. For the fork, after 3. See Equation 3.

Climbing, and getting stuck

Gradient ascent adjusts the probabilities of a fixed-size controller. For the crying baby with 2 nodes, 23 of 30 random starts reached the optimum. At the fork it failed every time: all 30 five-node starts ended by waiting forever. Asking only pays if a later node goes the right way, and a small step can’t change both at once. See Equation 5.

Formalism

Controllers and their value

A controller has nodes \(X\), an action distribution \(\psi(a \mid x)\) and a successor distribution \(\eta(x' \mid x, a, o)\) (ADM §23.1). At node \(x\) it draws \(a \sim \psi(\cdot \mid x)\); the world moves to \(s'\) and emits \(o\); the controller moves to \(x' \sim \eta(\cdot \mid x, a, o)\). The value of being at node \(x\) in state \(s\) is

\[ \begin{aligned} U(x, s) = \sum_a \psi(a \mid x)\Big[ R(s, a) &+ \gamma \sum_{s'} T(s' \mid s, a) \\ &\times \sum_o O(o \mid a, s') \sum_{x'} \eta(x' \mid x, a, o)\, U(x', s') \Big] \end{aligned} \tag{1}\]

In words. A node’s value in a state is the expected reward of its action, plus the discounted value of wherever the world and the controller go next.

Term by term.

  • \(\psi(a \mid x)\) averages over the node’s actions; \(R(s, a)\) is the reward.
  • \(T(s' \mid s, a)\,O(o \mid a, s')\) is the chance of the next state and observation.
  • \(\eta(x' \mid x, a, o)\) is the chance of the next node, which depends on what was observed, never on the hidden state.
  • There are \(|X||\mathcal{S}|\) equations and unknowns: 4 for a 2-node crying baby controller.

Why it holds. The pair (node, state) evolves as a Markov chain: its next value depends only on the current pair. Equation 1 is that chain’s Bellman expectation equation, the same as MDP policy evaluation (Chapter 7) on \(|X||\mathcal{S}|\) states. So it has a unique solution for \(\gamma < 1\), found by one linear solve.

From a belief \(b\), the value of starting at node \(x\), and the best starting node, are (ADM eqs. 23.2–23.3)

\[ U(x, b) = \sum_s b(s)\, U(x, s) = \alpha_x^\top b, \qquad x^* = \arg\max_x \alpha_x^\top b \tag{2}\]

In words. Each node is an alpha vector; start at the one that is best for your belief.

Term by term.

  • \(\alpha_x(s) = U(x, s)\) is node \(x\)’s alpha vector.
  • \(\max_x \alpha_x^\top b\) is piecewise linear and convex in \(b\), like Chapter 20’s value functions.

Why it holds. The controller’s future doesn’t depend on the belief, only on the hidden state, so the value from a belief is the belief-weighted average over states.

Policy iteration

Policy iteration (ADM §23.2; E. A. Hansen, UAI, 1998) alternates evaluation with improvement. Improvement considers every deterministic node “take \(a\), then after \(o\) go to node \(\nu(o)\)”, whose alpha vector is a one-step backup of the current ones:

\[ \alpha_{a,\nu}(s) = R(s, a) + \gamma \sum_{s'} T(s' \mid s, a) \sum_o O(o \mid a, s')\, U(\nu(o), s') \]

There are \(|\mathcal{A}||X|^{|\mathcal{O}|}\) such nodes. Pruning compares nodes state by state:

\[ x \text{ is dominated by } x' \quad\text{if}\quad U(x, s) \le U(x', s) \text{ for every state } s \tag{3}\]

In words. A node that is never better than another, whatever the state, can be removed.

Term by term.

  • A new node dominated by an existing node (or by another new one) is dropped.
  • A new node that dominates an existing node replaces it in place: every link into the old node now reaches the new one. That can create loops.
  • New nodes that are neither dominated nor dominating are added.

Why it holds. Replacing a node by one that is at least as good in every state can’t lower any node’s value, since every path through it does at least as well. The candidates are exactly the alpha vectors of a Bellman backup of the current value function. If every one of them is dominated by an existing node, the backup can’t raise the value at any belief, so the controller’s value already satisfies the Bellman optimality equation: it is optimal. The catch: pointwise dominance is weaker than Chapter 20’s belief-space pruning, so nodes that are never best at any belief can survive. For the crying baby, the controller grew to 58 nodes in 10 iterations, although only 2 are reachable from the best starting node.

Nonlinear programming

Fix the number of nodes and optimize everything at once (ADM §23.3; C. Amato, D. S. Bernstein and S. Zilberstein, Autonomous Agents and Multi-Agent Systems, 2010). Starting at node \(x_1\) from belief \(b\):

\[ \begin{aligned} \max_{U,\,\psi,\,\eta}\quad & \textstyle\sum_s b(s)\, U(x_1, s) \\ \text{subject to}\quad & U \text{ satisfies @eq-eval for every } x, s \\ & \psi(\cdot \mid x) \text{ and } \eta(\cdot \mid x, a, o) \text{ are probability distributions} \end{aligned} \tag{4}\]

In words. Choose the controller’s probabilities, and the values that go with them, to maximize the value at the start.

Term by term.

  • The unknowns are \(U\) (\(|X||\mathcal{S}|\) numbers), \(\psi\) and \(\eta\).
  • The evaluation constraints multiply \(\psi\) and \(\eta\) by \(U\), so they are quadratic: a quadratically constrained linear program.
  • The remaining constraints say each row is nonnegative and sums to 1.

Why it holds. Any feasible point is a real controller and its exact value, so the program’s optimum is the best \(n\)-node controller. But the constraints are not convex, so a local solver can stop at a local optimum. The Python tab solves it with SciPy’s SLSQP after substituting \(U\) from the linear solve.

Gradient ascent

Write the evaluation as a matrix equation on (node, state) pairs (ADM §23.4; N. Meuleau, K.-E. Kim, L. P. Kaelbling and A. R. Cassandra, UAI, 1999):

\[ \begin{aligned} u_\theta &= r_\theta + \gamma T_\theta u_\theta, \qquad r_\theta(x, s) = \sum_a \psi(a \mid x) R(s, a), \\ T_\theta\big((x, s), (x', s')\big) &= \sum_a \psi(a \mid x)\, T(s' \mid s, a) \sum_o O(o \mid a, s')\, \eta(x' \mid x, a, o) \end{aligned} \]

The objective is \(J(\theta) = \beta^\top u_\theta\), where \(\beta\) puts belief \(b\) on node \(x_1\) and zero elsewhere. Its gradient is

\[ \frac{\partial J}{\partial \theta} = v^\top \Big( \frac{\partial r_\theta}{\partial \theta} + \gamma \frac{\partial T_\theta}{\partial \theta}\, u_\theta \Big), \qquad u_\theta = (I - \gamma T_\theta)^{-1} r_\theta, \quad v = (I - \gamma T_\theta)^{-\top} \beta \tag{5}\]

In words. Changing a probability matters in proportion to how often the controller is at that node (the weights \(v\)), times how much the change alters the reward and the value of what comes next.

Term by term.

  • \(u_\theta\) is the value of every (node, state) pair; \(v\) is the expected discounted time spent in each pair, starting from \(\beta\).
  • For \(\psi(a \mid x)\): the derivative of \(r_\theta\) is \(R(s, a)\) in node \(x\)’s rows; the derivative of \(T_\theta\) is the transition that action \(a\) would make from node \(x\).
  • For \(\eta(x' \mid x, a, o)\): only \(T_\theta\) depends on it, through \(\psi(a \mid x)\,T(s' \mid s, a)\,O(o \mid a, s')\).
  • A node that is never visited (\(v = 0\) in its rows) gets no gradient.

Why it holds. Differentiate \((I - \gamma T_\theta)u_\theta = r_\theta\): \((I - \gamma T_\theta)\,\partial u = \partial r + \gamma\,(\partial T)\,u\). Multiply by \(\beta^\top (I - \gamma T_\theta)^{-1} = v^\top\). This is the book’s eq. 23.18 multiplied by \(\beta\), arranged so that one linear solve gives \(v\) for every parameter at once. The code checks it against finite differences.

After each step, every distribution is projected back onto the probability simplex:

\[ \text{proj}(y) = \arg\min_{z \ge 0,\ \mathbf{1}^\top z = 1} \|y - z\|_2 = \max(y - \tau, 0) \tag{6}\]

with the shift \(\tau\) chosen so that the result sums to 1.

In words. Shift every entry down by the same amount, clip at zero, and choose the shift so the result is a distribution.

Term by term.

  • \(\tau\): sort \(y\) in decreasing order, find the largest \(k\) with \(y_{(k)} > (\sum_{i \le k} y_{(i)} - 1)/k\), and set \(\tau = (\sum_{i \le k} y_{(i)} - 1)/k\) (J. Duchi, S. Shalev-Shwartz, Y. Singer and T. Chandra, ICML, 2008).
  • Entries below \(\tau\) become exactly 0: projection can switch actions off completely.

Why it holds. The optimality conditions for the least-squares problem give \(z_i = \max(y_i - \tau, 0)\) with one multiplier \(\tau\) for the sum constraint; the sort finds which entries stay positive.

Worked example

Evaluating the two-node rule

The rule “feed on crying, but not right after a feed” has nodes I (ignore; crying → F, quiet → I) and F (feed; then → I). Feeding always sates the baby, so a feed’s value is its reward plus the value of ignoring a sated baby:

  • \(U(F, \text{sated}) = -5 + 0.9\,U(I, \text{sated})\) and \(U(F, \text{hungry}) = -15 + 0.9\,U(I, \text{sated})\).

Ignoring a sated baby: it stays sated with probability 0.9 (then cries with probability 0.1) or turns hungry (then cries with probability 0.8):

  • \(U(I, \text{sated}) = 0 + 0.9\,\big[0.9\,(0.1\,U(F, \text{sated}) + 0.9\,U(I, \text{sated})) + 0.1\,(0.8\,U(F, \text{hungry}) + 0.2\,U(I, \text{hungry}))\big]\).
  • \(U(I, \text{hungry}) = -10 + 0.9\,\big[0.8\,U(F, \text{hungry}) + 0.2\,U(I, \text{hungry})\big]\), since a hungry baby stays hungry.

Four linear equations, four unknowns. The solution:

Table 2: The two-node rule’s values (alpha vectors).
Node Sated Hungry
I (ignore) -16.31 -38.25
F (feed) -19.67 -29.67

At 50–50, starting at I gives the average of its two values, -27.28, and starting at F gives -24.67. Feed first: that’s the exact optimal value, -24.67.

A projection by hand

A gradient step turns an action distribution into \(y = (0.7, 0.5, -0.1)\), which isn’t a distribution. Sorted, it’s the same order; the running sums are 0.7, 1.2, 1.1.

  • \(k = 1\): \(0.7 > (0.7 - 1)/1 = -0.3\). Yes.
  • \(k = 2\): \(0.5 > (1.2 - 1)/2 = 0.1\). Yes.
  • \(k = 3\): \(-0.1 > (1.1 - 1)/3 = 0.033\)? No.

So \(\tau = 0.1\), and the projection is \((0.6, 0.4, 0)\): the third action is switched off.

The code: project_simplex([0.7, 0.5, -0.1]) = [0.6, 0.4, 0.0].

Interactive

Nodes as belief regions

Choose a controller and step through a simulated run. The chart shows the belief a Chapter 19 filter would hold (black), which the controller never computes. The colored band shows the node the controller is in. Each node ends up standing for a range of beliefs.

Figure 2: One simulated run of the chosen controller: the belief a filter would hold (black line) and the controller’s node at each step (colored band).

Over 20,000 steps, the two-node rule was in node I at beliefs \(P(\text{hungry})\) from 0.0 to 0.032, and in node F from 0.471 to 0.541. At every step, the rule took the action the exact policy takes at the filter’s belief, which is why it is optimal. For the counter, each node always had the same belief: \(P(\text{home right}) =\) 0.5 in the even asking node, 0.85 and 0.15 one answer ahead, and 0.97 and 0.03 when going. The belief depends only on the answers heard since the hike began, and that’s exactly what the nodes count.

Algorithms & code

Controllers: evaluation, policy iteration, brute force, gradient ascent

# Finite state controllers for POMDPs (ADM Ch. 23).
# A POMDP p has p.T[s, a, s′], p.R[s, a], p.O[a, s′, o] and p.γ (as in Ch. 20-22).
# A controller has nodes x = 1..n. At node x it takes action a with probability
# ψ[x, a]; after taking a and observing o it moves to node x′ with probability
# η[x, a, o, x′]. It keeps no belief: the node is its whole memory.
struct Controller
    ψ::Matrix{Float64}            # ψ[x, a]
    η::Array{Float64,4}           # η[x, a, o, x′]
end
nnodes(c::Controller) = size(c.ψ, 1)

# A deterministic controller: action acts[x] at node x, then node next[x][o]
# after observation o (whatever the action, since only one action is ever taken).
function deterministic_controller(p, acts, next)
    n, nA, nO = length(acts), size(p.T, 2), size(p.O, 3)
    ψ, η = zeros(n, nA), zeros(n, nA, nO, n)
    for x in 1:n
        ψ[x, acts[x]] = 1
        for a in 1:nA, o in 1:nO
            η[x, a, o, next[x][o]] = 1
        end
    end
    return Controller(ψ, η)
end

# The controller and the POMDP together form a Markov chain on pairs (x, s),
# stored as index (x - 1) * |S| + s. Returns its transition matrix and reward
# vector (eqs. 23.8-23.9):
#   Tθ[(x,s), (x′,s′)] = Σ_a ψ[x,a] T[s,a,s′] Σ_o O[a,s′,o] η[x,a,o,x′]
#   rθ[(x,s)]          = Σ_a ψ[x,a] R[s,a]
function controller_chain(p, c::Controller)
    nS, nA, nO, n = size(p.T, 1), size(p.T, 2), size(p.O, 3), nnodes(c)
    idx(x, s) = (x - 1) * nS + s
    Tθ, rθ = zeros(n * nS, n * nS), zeros(n * nS)
    for x in 1:n, s in 1:nS, a in 1:nA
        c.ψ[x, a] == 0 && continue
        rθ[idx(x, s)] += c.ψ[x, a] * p.R[s, a]
        for s′ in 1:nS, o in 1:nO, x′ in 1:n
            Tθ[idx(x, s), idx(x′, s′)] += c.ψ[x, a] * p.T[s, a, s′] * p.O[a, s′, o] * c.η[x, a, o, x′]
        end
    end
    return Tθ, rθ
end

# Policy evaluation (eq. 23.1) by one linear solve: u = (I - γ Tθ)⁻¹ rθ.
# Returns U[x, s]; row x is node x's alpha vector (eq. 23.2: U(x, b) = U[x, :] ⋅ b).
function evaluate_controller(p, c::Controller)
    Tθ, rθ = controller_chain(p, c)
    u = (I - p.γ * Tθ) \ rθ
    return permutedims(reshape(u, size(p.T, 1), nnodes(c)))   # reshape fills s fastest
end

# Value at belief b from the best starting node (eq. 23.3), and that node.
function controller_value(p, c::Controller, b; U = evaluate_controller(p, c))
    vals = [U[x, :] ⋅ b for x in 1:nnodes(c)]
    return (u = maximum(vals), x = argmax(vals))
end

# Nodes the controller can ever reach from node x: follow every successor that
# has positive probability after an action that has positive probability.
function reachable_nodes(c::Controller, x)
    seen, stack = Set([x]), [x]
    while !isempty(stack)
        y = pop!(stack)
        for a in axes(c.ψ, 2), o in axes(c.η, 3), y′ in axes(c.η, 4)
            if c.ψ[y, a] > 0 && c.η[y, a, o, y′] > 0 && !(y′ in seen)
                push!(seen, y′); push!(stack, y′)
            end
        end
    end
    return sort(collect(seen))
end

# Run the controller on the true POMDP from node x and state s for `steps`
# steps; returns the discounted return and the (node, action, observation) trace.
function run_controller(p, c::Controller, x, s, steps, rng)
    ret, trace = 0.0, Tuple{Int,Int,Int,Int}[]
    for t in 0:steps-1
        a = rand(rng, Categorical(c.ψ[x, :]))
        s′, r, o = pomdp_step(p, s, a, rng)
        push!(trace, (x, s, a, o))
        ret += p.γ^t * r
        x, s = rand(rng, Categorical(c.η[x, a, o, :])), s′
    end
    return ret, trace
end

# Policy iteration for controllers (§23.2; E. A. Hansen, 1998). Improvement
# adds every deterministic node "take a, then go to next[o]" whose alpha vector
# is not dominated state by state (eq. 23.4). A new node that dominates an old
# node replaces it in place, so every link into the old node now reaches the
# new one: that is how loops form. Stops when no node survives, or at k_max.
# Returns the controller and, per iteration, (nodes, value at b).
function controller_policy_iteration(p, b; k_max = 20, c = nothing, tol = 1e-9, max_nodes = 500)
    nS, nA, nO = size(p.T, 1), size(p.T, 2), size(p.O, 3)
    c === nothing && (c = Controller(fill(1 / nA, 1, nA), ones(1, nA, nO, 1)))   # one uniform node
    history = [(nnodes(c), controller_value(p, c, b).u)]
    for _ in 1:k_max
        U = evaluate_controller(p, c)
        n = nnodes(c)
        # Every candidate: an action and a successor for each observation.
        cands = Tuple{Int,Vector{Int},Vector{Float64}}[]
        for a in 1:nA, next in Iterators.product(fill(1:n, nO)...)
            α = [p.R[s, a] + p.γ * sum(p.T[s, a, s′] * p.O[a, s′, o] * U[next[o], s′]
                                       for s′ in 1:nS, o in 1:nO) for s in 1:nS]
            push!(cands, (a, collect(next), α))
        end
        dominated(α, β) = all(α .<= β .+ tol)                  # α ≤ β in every state
        # Drop candidates no better than an old node, or than an earlier survivor.
        keep = Int[]
        for (i, (_, _, α)) in enumerate(cands)
            any(dominated(α, U[x, :]) for x in 1:n) && continue
            any(dominated(α, cands[j][3]) for j in keep) && continue
            filter!(j -> !dominated(cands[j][3], α), keep)       # α beats earlier survivors
            push!(keep, i)
        end
        isempty(keep) && break                                   # no improvement anywhere
        ψ, η = copy(c.ψ), copy(c.η)
        added = Tuple{Int,Vector{Int}}[]
        for i in keep
            a, next, α = cands[i]
            olds = [x for x in 1:n if dominated(U[x, :], α)]
            if isempty(olds)
                push!(added, (a, next))                          # a genuinely new node
            else
                for x in olds                                    # overwrite dominated nodes
                    ψ[x, :] .= 0; ψ[x, a] = 1
                    η[x, :, :, :] .= 0
                    for a′ in 1:nA, o in 1:nO
                        η[x, a′, o, next[o]] = 1
                    end
                end
            end
        end
        n + length(added) > max_nodes && break
        # Grow the arrays; new nodes point only to old nodes.
        m = n + length(added)
        ψ2, η2 = zeros(m, nA), zeros(m, nA, nO, m)
        ψ2[1:n, :], η2[1:n, :, :, 1:n] = ψ, η
        for (k, (a, next)) in enumerate(added)
            ψ2[n + k, a] = 1
            for a′ in 1:nA, o in 1:nO
                η2[n + k, a′, o, next[o]] = 1
            end
        end
        c = Controller(ψ2, η2)
        push!(history, (nnodes(c), controller_value(p, c, b).u))
    end
    return c, history
end

# Solve M u = r in place by Gaussian elimination with partial pivoting; r
# becomes u. Used by the brute-force search below, which solves millions of
# tiny systems: avoiding allocation makes it about ten times faster.
function solve_inplace!(M, r)
    m = length(r)
    for k in 1:m
        piv = k
        for i in k+1:m
            abs(M[i, k]) > abs(M[piv, k]) && (piv = i)
        end
        if piv != k                                  # swap rows k and piv
            for j in 1:m
                M[k, j], M[piv, j] = M[piv, j], M[k, j]
            end
            r[k], r[piv] = r[piv], r[k]
        end
        for i in k+1:m                               # eliminate below the pivot
            f = M[i, k] / M[k, k]
            f == 0 && continue
            for j in k:m
                M[i, j] -= f * M[k, j]
            end
            r[i] -= f * r[k]
        end
    end
    for k in m:-1:1                                  # back substitution
        t = r[k]
        for j in k+1:m
            t -= M[k, j] * r[j]
        end
        r[k] = t / M[k, k]
    end
    return r
end

# Brute force, for validation: evaluate every deterministic n-node controller
# (|A|^n action choices times n^(n|O|) successor tables) and return the best
# value at belief b (from its best start node), with its actions and successors.
# acts and next are counted like the digits of a mixed-radix number.
function best_deterministic_controller(p, n, b)
    nS, nA, nO = size(p.T, 1), size(p.T, 2), size(p.O, 3)
    m = n * nS
    M, r = zeros(m, m), zeros(m)
    acts, next = ones(Int, n), ones(Int, n * nO)       # next[(x - 1) * nO + o]
    best = (u = -Inf, acts = copy(acts), next = copy(next))
    # Increment a digit vector with base `base`; false once it wraps to all ones.
    function increment!(v, base)
        for i in eachindex(v)
            v[i] < base && (v[i] += 1; return true)
            v[i] = 1
        end
        return false
    end
    while true
        while true
            fill!(M, 0.0)                              # build I - γTθ and rθ
            for x in 1:n, s in 1:nS
                a, i = acts[x], (x - 1) * nS + s
                r[i] = p.R[s, a]
                M[i, i] += 1
                for s′ in 1:nS, o in 1:nO
                    M[i, (next[(x - 1) * nO + o] - 1) * nS + s′] -= p.γ * p.T[s, a, s′] * p.O[a, s′, o]
                end
            end
            solve_inplace!(M, r)                       # r now holds u
            for x in 1:n
                v = 0.0
                for s in 1:nS
                    v += r[(x - 1) * nS + s] * b[s]
                end
                v > best.u + 1e-12 && (best = (u = v, acts = copy(acts), next = copy(next)))
            end
            increment!(next, n) || break
        end
        increment!(acts, nA) || break
    end
    return best
end

# Gradient of J(θ) = β ⋅ u_θ, the value of starting in node 1 with belief b
# (β holds b in node 1's block, zeros elsewhere; eq. 23.11). With Z = I - γTθ,
# u = Z⁻¹rθ and v = Z⁻ᵀβ (how much discounted time is spent in each (x, s)),
#   ∂J/∂θ = v ⋅ (∂rθ/∂θ + γ (∂Tθ/∂θ) u),
# which is eq. 23.18 multiplied by β, arranged to need only two linear solves.
function controller_gradient(p, c::Controller, b)
    nS, nA, nO, n = size(p.T, 1), size(p.T, 2), size(p.O, 3), nnodes(c)
    Tθ, rθ = controller_chain(p, c)
    Z = I - p.γ * Tθ
    β = [x == 1 ? b[s] : 0.0 for s in 1:nS, x in 1:n][:]     # column-major: s fastest
    u, v = Z \ rθ, Z' \ β
    U, V = permutedims(reshape(u, nS, n)), permutedims(reshape(v, nS, n))
    gψ, gη = zeros(n, nA), zeros(n, nA, nO, n)
    for x in 1:n, a in 1:nA, s in 1:nS
        V[x, s] == 0 && continue
        gψ[x, a] += V[x, s] * p.R[s, a]
        for s′ in 1:nS, o in 1:nO
            w = p.γ * V[x, s] * p.T[s, a, s′] * p.O[a, s′, o]
            w == 0 && continue
            for x′ in 1:n
                gψ[x, a] += w * c.η[x, a, o, x′] * U[x′, s′]
                gη[x, a, o, x′] += w * c.ψ[x, a] * U[x′, s′]
            end
        end
    end
    return (J = β ⋅ u, ψ = gψ, η = gη)
end

# Euclidean projection onto the probability simplex (J. Duchi et al., 2008):
# sort, find how many entries stay positive, and shift them all by one amount.
function project_simplex(y)
    μ = sort(y, rev = true)
    cs = cumsum(μ)
    ρ = findlast(j -> μ[j] - (cs[j] - 1) / j > 0, 1:length(μ))
    θ = (cs[ρ] - 1) / ρ
    return max.(y .- θ, 0.0)
end

# A random n-node controller: each distribution drawn uniformly from its simplex
# (normalized exponential draws are uniform on the simplex).
function random_controller(p, n, rng)
    nA, nO = size(p.T, 2), size(p.O, 3)
    unif(k) = (e = -log.(rand(rng, k)); e ./ sum(e))
    ψ = permutedims(reduce(hcat, [unif(nA) for _ in 1:n]))
    η = zeros(n, nA, nO, n)
    for x in 1:n, a in 1:nA, o in 1:nO
        η[x, a, o, :] = unif(n)
    end
    return Controller(ψ, η)
end

# Controller gradient ascent (§23.4): step along the gradient, then project
# each distribution back onto the simplex. Starts from `c` (default: random).
# Returns the controller and J after each step.
function controller_gradient_ascent(p, b, n, rng; α = 0.01, k_max = 2000, c = random_controller(p, n, rng))
    nA, nO, Js = size(p.T, 2), size(p.O, 3), Float64[]
    for _ in 1:k_max
        g = controller_gradient(p, c, b)
        push!(Js, g.J)
        ψ = reduce(vcat, [project_simplex(c.ψ[x, :] .+ α .* g.ψ[x, :])' for x in 1:n])
        η = copy(c.η)
        for x in 1:n, a in 1:nA, o in 1:nO
            η[x, a, o, :] = project_simplex(c.η[x, a, o, :] .+ α .* g.η[x, a, o, :])
        end
        c = Controller(ψ, η)
    end
    push!(Js, controller_gradient(p, c, b).J)
    return c, Js
end
import numpy as np
from scipy.optimize import minimize

from adm.problems.crying_baby import crying_baby

p = crying_baby()                     # 0 = sated / feed / crying
nS, nA, nO = p.T.shape[0], p.T.shape[1], p.O.shape[2]
n = 2                                 # controller nodes
b0 = np.array([0.5, 0.5])


def unpack(theta):
    """Split the parameters into psi[x, a] and eta[x, a, o, x2]."""
    psi = theta[: n * nA].reshape(n, nA)
    eta = theta[n * nA:].reshape(n, nA, nO, n)
    return psi, eta


def evaluate(psi, eta):
    """eq. 23.1 as one linear system on (node, state) pairs; returns U[x, s]."""
    # Tt[x, s, y, t] = sum over a, o of psi[x,a] T[s,a,t] O[a,t,o] eta[x,a,o,y]
    Tt = np.einsum("xa,sat,ato,xaoy->xsyt", psi, p.T, p.O, eta)
    Tt = Tt.reshape(n * nS, n * nS)       # row x*nS + s, column y*nS + t
    r = np.einsum("xa,sa->xs", psi, p.R).reshape(n * nS)
    U = np.linalg.solve(np.eye(n * nS) - p.gamma * Tt, r)
    return U.reshape(n, nS)


def negative_value(theta):
    psi, eta = unpack(theta)
    return -evaluate(psi, eta)[0] @ b0    # start in node 0


# eq. 23.5 with U substituted: every psi and eta row must sum to 1.
rows = [slice(x * nA, (x + 1) * nA) for x in range(n)]
rows += [slice(n * nA + k * n, n * nA + (k + 1) * n) for k in range(n * nA * nO)]
cons = [{"type": "eq", "fun": lambda th, sl=sl: th[sl].sum() - 1} for sl in rows]

best = None
for seed in range(5):                 # a local solver: try several starts
    rng = np.random.default_rng(seed)
    theta0 = np.concatenate([rng.dirichlet(np.ones(nA), n).ravel(),
                             rng.dirichlet(np.ones(n), n * nA * nO).ravel()])
    res = minimize(negative_value, theta0, method="SLSQP",
                   bounds=[(0, 1)] * len(theta0), constraints=cons)
    if best is None or res.fun < best.fun:
        best = res

psi, eta = unpack(best.x)
print(f"best 2-node value at b0: {-best.fun:.3f}")
print("most likely action per node:", [p.actions[a] for a in psi.argmax(axis=1)])
assert abs(-best.fun - (-24.6749)) < 0.01             # the exact optimum
  • Array{Float64,4} is a 4-dimensional array; η[x, a, o, :] is a view of its last axis, like eta[x, a, o, :] in NumPy (but 1-based).
  • permutedims(reshape(u, nS, n)): Julia reshapes column-major, filling the first index fastest, so the vector indexed (x - 1) * nS + s becomes an nS × n matrix whose columns are nodes; permutedims transposes it to U[x, s].
  • Z' \ β solves the transposed system \(Z^\top v = \beta\); \ is np.linalg.solve.
  • Iterators.product(fill(1:n, nO)...) is itertools.product(range(n), repeat=nO): every successor table.
  • increment!(v, base) counts through all digit vectors, like an odometer, without building the whole list.
  • v[i] < base && (v[i] += 1; return true) runs the parenthesized block only if the test is true: a one-line if.

Run it

p = crying_baby()
# Node 1 ignores (crying → node 2, quiet → node 1); node 2 feeds (then → node 1).
c = deterministic_controller(p, [3, 1], [[2, 1], [1, 1]])
U = evaluate_controller(p, c)          # U[x, s]: rows are alpha vectors
controller_value(p, c, [0.5, 0.5])     # best starting node and its value
(u = -24.674934966510335, x = 2)

Results

Controllers for the crying baby

Figure 3: Crying baby. Left: the alpha vectors of two 2-node controllers (dashed: feed whenever crying; solid: never feed twice in a row), against the exact value (black). Right: the value of the controller built by policy iteration, from its best starting node, after 0, 1 and 2 iterations.
  • Two nodes are enough. The rule that never feeds twice in a row is optimal at 50–50 (starting at F), and from its best node at every belief. Brute force over all 2-node and 3-node deterministic controllers agrees: the best value at 50–50 is -24.675 with 2 nodes and -24.675 with 3. The best 1-node controller always feeds: -55.
  • Randomness helps a tiny controller. A 1-node controller that feeds with probability 0.396 and otherwise ignores (found by gradient ascent from all 30 starts) earns -40.23. It can’t remember when it last fed, so it spreads its feeds out at random.
  • Policy iteration finds the optimum, then keeps growing. It reached the optimal value at 50–50 after 2 iterations, with 3 nodes. After 10 iterations it had 58 nodes, of which 2 are reachable from the best starting node. Pointwise pruning can’t remove nodes that are better in some state but never best at any belief.

How much memory does the fork need?

Figure 4: The fork: the best value at 50–50 achievable by a deterministic controller with 1 to 5 nodes (brute force for 1 to 4; the counting controller for 5), against the optimal value (dashed). Left of the dashed line’s reach, the controller can only wait.
  • Below 4 nodes, waiting is best. With 3 nodes, the natural policy “ask once, then go where told” is worth -0.51: an 85%-reliable answer isn’t worth acting on at these stakes, so it loses a little each hike.
  • 4 nodes get part of the way. The best 4-node controller asks, goes left after a single “left”, but asks again after a “right” and goes right only after two in a row. It earns 15.33, from its best starting node.
  • 5 nodes are optimal. The counting controller earns 27.115, which matches Chapter 21’s bounds on the optimum (27.115 from below).
  • Policy iteration finds it too. Starting from one node, it reached the optimum after 3 iterations, with 23 nodes. It stopped after 6 iterations with 242 nodes, of which 5 are reachable from its best starting node.

Gradient ascent: local optima

Figure 5: Controller gradient ascent (step size 0.003, 3000 steps). Left: crying baby, 2 nodes, 30 random starts; the exact optimum is dashed. Right: fork, 5 nodes, 30 random starts (gray) and 10 starts blended 50/50 with the counting controller (blue); the optimum is dashed.
  • The crying baby is easy. With 2 nodes, 23 of 30 starts reached the optimum (within 0.01); with 3 nodes, 16. The others stopped at local optima.
  • The fork has a trap. From random 5-node starts, all 30 of 30 runs ended by waiting forever, with value 0. In the first run’s final controller, 3 nodes are ever visited, and each waits. At those nodes the gradient for waiting is 0, for asking at most -6.7, and for going at most -125.8. Asking alone only costs; it pays only together with a later node that goes the right way.
  • A good start escapes it. Blending the counting controller 50/50 with a random one, 10 of 10 runs climbed to the optimum; at 70/30, 10; at 30/70, 1.

Solved problems

Problem 1: counting controllers. For the fork (\(|\mathcal{A}| = 4\), \(|\mathcal{O}| = 2\)), how many candidate nodes does one policy-iteration step try on a 5-node controller? How many deterministic 4-node controllers are there?

Each candidate is an action and a successor for each observation: \(|\mathcal{A}||X|^{|\mathcal{O}|} = 4 \times 5^2 = 100\). A deterministic \(n\)-node controller chooses an action per node (\(4^n\)) and a successor per node and observation (\(n^{2n}\)): for \(n = 4\), \(256 \times 65{,}536 = 16{,}777{,}216\). That is why brute force stops at 4 nodes here; for \(n = 5\) it would be about \(10^{10}\).

Problem 2: a 1-node controller by hand. A crying baby controller has one node, which always feeds. Find its values and its value at 50–50.

Feeding always leads to a sated baby, and the controller stays in its one node: \(U(\text{sated}) = -5 + 0.9\,U(\text{sated})\), so \(U(\text{sated}) = -50\). \(U(\text{hungry}) = -15 + 0.9\,U(\text{sated}) = -60\). At 50–50: \(-55\), the brute-force result for one node above. Always ignoring is far worse, since the baby stays hungry forever: \(U(\text{hungry}) = -10/(1 - 0.9) = -100\).

Problem 3: why randomize? Why can a 1-node controller that feeds at random beat both “always feed” and “always ignore”? Would randomness help a controller with enough nodes?

With one node, the controller can’t tell a quiet baby from a crying one, or remember its last feed; a deterministic node must do the same thing forever. Random feeding with a moderate probability feeds roughly as often as the baby gets hungry, without paying for a feed every step. It’s a substitute for memory. With enough nodes, a deterministic controller can match the optimal belief-based policy, which is itself deterministic; here 2 nodes did. Randomness helps when memory is scarce.

Problem 4: project a distribution. Project \(y = (0.9, 0.4, -0.2)\) onto the simplex.

Running sums: 0.9, 1.3, 1.1. \(k = 1\): \(0.9 > -0.1\); \(k = 2\): \(0.4 > 0.15\); \(k = 3\): \(-0.2 > 0.033\)? No. So \(\tau = (1.3 - 1)/2 = 0.15\) and the projection is \((0.75, 0.25, 0)\).

Problem 5: why does “ask once, then go” lose? Compute the value of the 3-node fork controller that asks once and goes where told, and explain the result.

Let \(V\) be the value at the asking node. Going after one answer is right with probability 0.85: expected reward \(0.85 \times 10 + 0.15 \times (-50) = 1\), then a new hike starts at the asking node. So \(V = -1 + 0.95\,(1 + 0.95\,V)\), which gives \(V(1 - 0.9025) = -0.05\) and \(V \approx -0.513\). Each hike costs 1 to ask and earns only 1 on average, a step later, so the controller loses slowly and waiting (0) is better. A second answer that agrees raises the chance of being right to about 0.97, which is worth the extra cost: that’s why the optimal controller counts to two.

Recap

  1. A controller is a policy with a finite memory (Equation 1). Why it matters: it runs without belief updates, which is useful when the model is too big to filter, or when the policy must run on simple hardware.
  2. Evaluation is exact: one linear solve on (node, state) pairs, giving an alpha vector per node (Equation 2). Why it matters: controllers slot into the value-function picture of Chapters 20–21; their value over beliefs is a lower bound on the optimal one.
  3. Policy iteration adds and replaces nodes, and reaches the optimum on small problems (Equation 3). Why it matters: replacing dominated nodes creates the loops that make controllers compact, but pointwise pruning can leave many useless nodes.
  4. Fixed-size controllers can be optimized as a nonlinear program or by gradient ascent (Equation 4, Equation 5, Equation 6). Why it matters: the objective is not concave, and at the fork every random start fell into the same trap. Information-gathering behavior needs coordinated changes that local search can’t make one step at a time.
  5. The number of nodes is the real design decision. Why it matters: the crying baby needs 2, the fork needs 5 (with fewer than 4 it can only wait). Too few nodes caps the value; too many makes optimization harder.

Cheat sheet & connections

Table 3: Key equations from this page.
Name Formula Ref
Controller evaluation \(U(x, s) = \sum_a \psi(a \mid x)[R(s, a) + \gamma\sum_{s'} T \sum_o O \sum_{x'} \eta\, U(x', s')]\) Equation 1
Node values over beliefs \(U(x, b) = \alpha_x^\top b\); start at \(\arg\max_x \alpha_x^\top b\) Equation 2
Pointwise dominance \(U(x, s) \le U(x', s)\) for all \(s\) Equation 3
Fixed-size program \(\max \sum_s b(s) U(x_1, s)\) s.t. Equation 1, \(\psi\) and \(\eta\) distributions Equation 4
Controller gradient \(v^\top(\partial r_\theta + \gamma\,\partial T_\theta\, u_\theta)\), \(v = (I - \gamma T_\theta)^{-\top}\beta\) Equation 5
Simplex projection \(\max(y - \tau, 0)\), \(\tau\) so it sums to 1 Equation 6

Common pitfalls

  • Confusing nodes with states. A node is the controller’s memory; the environment’s state stays hidden. \(\eta\) may depend on the observation, never on the state.
  • Starting at the wrong node. A controller’s value at a belief depends on the starting node; choose it with Equation 2.
  • Trusting a local optimum. Gradient ascent and local nonlinear solvers need several starts, or a start from a known good controller.
  • Too few nodes. If the value plateaus at something simple (like waiting), the controller may lack the memory to do better. Brute force on small sizes, or policy iteration, can tell you.
  • Counting nodes that are never used. Report the nodes reachable from the starting node, not the total.

Connections