Lecture 19 · Part 5

Control as Inference

Levine reframes the control problem as probabilistic inference in a graphical model with binary optimality variables, where the trajectory posterior weights paths by exponentiated reward. Exact inference gives a soft value iteration — softmax in place of max — but with an optimism bias under stochastic dynamics; variational inference fixes it and lands exactly on maximum-entropy RL, whose practical faces are soft Q-learning and soft actor-critic.

▶ Lecture videos (5 parts) Slides (current revision, PDF) ↗

Why hard optimality can’t explain real behavior

The lecture opens with a question about modeling: does optimal control provide a reasonable model of human and animal behavior? A monkey moving a cursor to a target usually reaches it — but by a different route each time, sometimes lazily indirect. The mistakes aren’t uniform noise: with practice they “tend to occur in ways that matter less for the success of the task.” The monkey knows the manner of reaching the goal barely affects the reward, so it lets that part be sloppy.

The RL framework we’ve built cannot express this. For any objective that is linear in the state–action marginal, every fully observed MDP admits a deterministic optimal policy — so randomness can never be rational under hard optimality. Yet natural agents faced with the same situation twice don’t do the same thing, and “good behavior is still the most likely” without being the only behavior. We need a probabilistic model of near-optimal behavior, and Levine flags three reasons to want one: it explains sub-optimal demonstrations (crucial for the inverse RL lecture that follows), it lets us attack control with the whole toolbox of inference algorithms, and it explains why stochastic behavior helps exploration and transfer.

A graphical model in which inference is control

Take the usual MDP chain of states and actions with transitions $p(s_{t+1} \mid s_t, a_t)$ — that alone carries “no assumption of optimal behavior.” To encode why an agent picks one action over another, add one node per time step:

The choice looks arbitrary — Levine admits as much twice — but it earns its keep immediately. Bayes’ rule on the trajectory posterior gives

$p(\tau \mid \mathcal{O}_{1:T}) \propto p(\tau) \prod_{t} \exp(r(s_t, a_t)) = p(\tau)\, \exp\!\left(\sum_t r(s_t, a_t)\right)$

where $p(\tau)$ carries only the dynamics and initial state. With deterministic dynamics, $p(\tau)$ is an indicator of physical consistency, and the picture is exactly the sloppy monkey: the highest-reward trajectory is the most likely, equal- reward routes are chosen at random, and probability decays exponentially as reward drops — so the agent will wander on the way to the goal but will not fail to reach it.

Exact inference: backward messages and soft value functions

Three inference problems matter in this chain-structured model, and they’re solved the way you’d solve an HMM — by message passing. The backward message $\beta_t(s_t, a_t) = p(\mathcal{O}_{t:T} \mid s_t, a_t)$ is the probability of being optimal from now to the end. Insert $s_{t+1}$ and integrate it out; conditional independence of future optimality variables given $s_{t+1}$ factorizes the result into the transition, the current $\exp(r)$ term, and the next message, giving the recursion (run backward from $\beta_T = p(\mathcal{O}_T \mid s_T, a_T)$ ):

$\beta_t(s_t, a_t) = p(\mathcal{O}_t \mid s_t, a_t)\, \mathbb{E}_{s_{t+1} \sim p(s_{t+1} \mid s_t, a_t)}\!\left[\beta_{t+1}(s_{t+1})\right], \qquad \beta_t(s_t) = \mathbb{E}_{a_t \sim p(a_t \mid s_t)}\!\left[\beta_t(s_t, a_t)\right]$

The action prior $p(a_t \mid s_t)$ — how likely actions are before you know whether the agent is optimal, emphatically not the policy — is assumed uniform, and this loses nothing: any non-uniform prior can be folded into the reward as $\tilde{r} = r + \log p(a_t \mid s_t)$ and the machinery rerun with a uniform prior.

Now define $V_t(s_t) = \log \beta_t(s_t)$ and $Q_t(s_t, a_t) = \log \beta_t(s_t, a_t)$ and the recursion becomes eerily familiar:

$V_t(s_t) = \log \int \exp(Q_t(s_t, a_t))\, da_t, \qquad Q_t(s_t, a_t) = r(s_t, a_t) + \log \mathbb{E}\!\left[\exp(V_{t+1}(s_{t+1}))\right]$

The first equation is a log-sum-exp — a soft max: as the $Q$ values grow, it approaches $\max_a Q_t$ . Value iteration’s hard max over actions has been softened, “which kind of makes sense — we want to soften our notion of optimality so that slightly sub-optimal things are still possible.” The second equation is a Bellman backup when the dynamics are deterministic (one term in the expectation, log and exp cancel), but under stochastic dynamics the log-expected-exp is dominated by the luckiest next state — an optimism problem we’ll return to.

The policy is the posterior $p(a_t \mid s_t, \mathcal{O}_{1:T})$ . Past optimality variables are d-separated by $s_t$ , and applying Bayes’ rule to numerator and denominator simultaneously collapses everything into a ratio of messages:

$\pi(a_t \mid s_t) = \frac{\beta_t(s_t, a_t)}{\beta_t(s_t)} = \exp(Q_t(s_t, a_t) - V_t(s_t)) = \exp(A_t(s_t, a_t))$

Forward messages $\alpha_t(s_t) = p(s_t \mid \mathcal{O}_{1:t-1})$ come from the mirror-image recursion, and their product with backward messages gives state marginals: $p(s_t \mid \mathcal{O}_{1:T}) \propto \beta_t(s_t)\, \alpha_t(s_t)$ . Picture two cones — one radiating backward from the goal (states from which you can still reach it), one forward from the start (states reachable with high reward) — with near-optimal behavior living in their intersection. Human reaching experiments actually show this cigar-shaped state distribution: tight at start and goal, widest in the middle. Forward messages aren’t needed for the policy, but they are the workhorse of inverse RL in the next lecture.

The optimism problem, and the variational fix

Worked example — The lottery ticket

Suppose an action buys a lottery ticket: probability $0.001$ of a huge payoff, $0.999$ of nothing. The expected value is small — but $\log \mathbb{E}[\exp(V_{t+1})]$ is dominated by the jackpot state, so the backup scores the ticket highly. The inference question “given that you obtained high reward, what was your action probability?” is subtly wrong: knowing you got a million dollars makes it more likely you played the lottery, but that doesn’t make playing the lottery a good idea. Formally, conditioning on optimality changes the posterior dynamics — $p(s_{t+1} \mid s_t, a_t, \mathcal{O}_{1:T}) \ne p(s_{t+1} \mid s_t, a_t)$ — and the model bends the world to agree with the evidence.

What we actually want is “given that you obtained high reward, what was your action probability, given that your transition probabilities did not change?” That is a variational inference problem: find $q(s_{1:T}, a_{1:T})$ close to the posterior $p(s_{1:T}, a_{1:T} \mid \mathcal{O}_{1:T})$ but constrained to the true dynamics. The trick is a peculiar variational family where most of $q$ is fixed to match $p$ :

$q(s_{1:T}, a_{1:T}) = p(s_1) \prod_t p(s_{t+1} \mid s_t, a_t)\, q(a_t \mid s_t)$

Only the action conditional is learned. With $x = \mathcal{O}_{1:T}$ and $z = (s_{1:T}, a_{1:T})$ , this is exactly the $q(z) \approx p(z \mid x)$ setup from the variational inference lecture, so write the evidence lower bound:

The variational lower bound is maximum-entropy RL

Substituting $q$ into $\log p(x) \ge \mathbb{E}_{z \sim q}[\log p(x, z) - \log q(z)]$ , the initial-state and transition log-probabilities appear in both $\log p$ and $\log q$ and cancel — that is why the family was chosen — leaving

$\log p(\mathcal{O}_{1:T}) \ge \sum_t \mathbb{E}_{(s_t, a_t) \sim q}\!\left[r(s_t, a_t) + \mathcal{H}(q(a_t \mid s_t))\right]$

Because $q$ shares the true dynamics, the first term is precisely the RL objective; the second adds per-step action entropy. Maximizing the bound means: maximize reward and stay as random as you can while doing it — the sub-optimal monkey, derived rather than assumed.

Optimizing the bound by dynamic programming from the last step backward uses one recurring fact: an objective of the form $\mathbb{E}_q[\text{something}] - \mathbb{E}_q[\log q]$ is maximized by $q \propto \exp(\text{something})$ . Each step therefore gives $q(a_t \mid s_t) = \exp(Q_t(s_t, a_t) - V_t(s_t))$ with $V$ the log-integral-exp normalizer — and, crucially, the backup for $Q$ is now the regular Bellman backup $Q_t = r + \mathbb{E}[V_{t+1}]$ , “not optimistic anymore.” The result is soft value iteration: identical to value iteration except $V(s) \leftarrow \operatorname{soft\,max}_a Q(s, a)$ replaces the hard max. Variants follow mechanically: a discount $\gamma$ (equivalently a $1-\gamma$ chance of death), an explicit temperature, and an infinite-horizon version that just runs the backup to convergence.

Algorithms: soft Q-learning, entropy-regularized policy gradients

The variational picture instantiates directly as deep RL algorithms.

Code walkthrough — Soft Q-learning

Take standard Q-learning and change one line — the target’s hard max becomes a soft max:

# standard target:  y = r(s, a) + gamma * max_a' Q_target(s', a')
# soft target:
y = r(s, a) + gamma * logsumexp(Q_target(s_next, :))   # soft max over a'

# everything else is DQN:
# 1. act, store (s, a, s', r) in replay buffer R
# 2. sample a mini-batch from R
# 3. compute targets y with the *target* network (soft max)
# 4. regress Q_theta toward y
# 5. lazily copy theta into the target network

The recovered policy is $\pi(a \mid s) = \exp(Q_\phi(s, a) - V(s)) = \exp(A(s, a))$ — the Boltzmann-over-advantage policy, which is the solution to the corresponding variational inference problem.

Alternatively, attack the bound directly with policy gradients: $J(\theta) = \sum_t \mathbb{E}_{\pi_\theta}[r(s_t, a_t)] + \mathcal{H}(\pi_\theta(a_t \mid s_t))$ is the standard objective plus an entropy bonus — the widely used entropy-regularized policy gradient, which combats premature collapse of policy entropy (on-policy exploration lives entirely in that stochasticity). Its gradient has a pleasant surprise: differentiating through the entropy contributes a constant $-1$ inside the parentheses, which the baseline property lets you drop — so the entire modification is subtracting $\log \pi$ from the reward, nothing more. Substituting $\log \pi = Q - V$ and crunching further, the entropy-regularized policy gradient comes out looking almost exactly like the soft Q-learning update — differing only in an extra $-\nabla V$ term on one side and a soft max on the other — a striking structural connection between the two families.

The practical culmination is soft actor-critic: an off-policy actor-critic where the critic update is the ordinary policy Q-function backup with a $-\log \pi$ entropy term added to the target, and the actor update maximizes $\mathbb{E}[\hat{Q}^\pi_\phi(s, a)] + \mathcal{H}(\pi_\theta)$ . In variational language, the critic step is message passing in the variational family and the actor step fits $q$ to the posterior; in practice it’s a simple, very effective algorithm — “one of the most widely used off-policy continuous control algorithms today.”

What soft optimality buys you

Levine closes with evidence that this is more than a modeling nicety. An ant robot facing two passages to a goal: a conventional algorithm commits early to whichever passage its Q-function happens to rank higher, and can commit wrong; a policy proportional to $\exp(Q)$ keeps probability mass on both peaks until one is clearly better (Haarnoja et al., soft Q-learning with deep energy-based policies). Pre-trained with the underspecified reward “run fast in any direction,” a soft-Q ant runs in every direction — entropy demands it — and then fine-tunes into a specific corridor far faster than DDPG, which must first unlearn its one committed direction. Soft actor-critic trained a real robot arm to stack Lego blocks robustly under physical perturbation, and taught the Minitaur quadruped a real-world walking gait that generalized to slopes and stairs it never trained on.

The suggested readings trace the lineage: Todorov’s linearly solvable MDPs and human motor control work, Kappen’s “optimal control as graphical model inference,” Ziebart’s maximum-entropy inverse RL, the soft Q-learning and soft actor-critic papers, and Levine’s 2018 tutorial “Reinforcement Learning and Control as Probabilistic Inference.”

Back of the chalk tray — what to remember

Hard optimality admits deterministic optimal policies, so it cannot rationalize the stochastic, sloppy-where-it-doesn’t-matter behavior of real agents; a probabilistic notion of near-optimality can.
Adding binary optimality variables with $p(\mathcal{O}_t \mid s_t, a_t) = \exp(r(s_t, a_t))$ makes the trajectory posterior $p(\tau \mid \mathcal{O}_{1:T}) \propto p(\tau) \exp(\sum_t r)$ — the best trajectory is most likely, worse ones exponentially less so.
Exact inference is message passing: log backward messages are soft $Q$ and $V$ functions, value iteration reappears with a softmax in place of the max, and the policy is $\exp(A_t(s_t, a_t))$ .
Exact inference is optimistic under stochastic dynamics — it lets evidence of success bend the transition probabilities, like inferring you should play the lottery because you know you won.
Restricting a variational distribution to the true dynamics and optimizing the ELBO fixes the optimism and yields exactly maximum-entropy RL: expected reward plus action entropy, solved by soft value iteration with a non-optimistic Bellman backup.
Practical instantiations — soft Q-learning, entropy-regularized policy gradients, and soft actor-critic — buy better exploration, principled tie-breaking, robustness, and faster fine-tuning, and underlie one of the most used continuous-control algorithms.

Check yourself

Derivation-style exercises in the lecture’s spirit — work them on paper.

Problem 19.1 trajectory posterior

Starting from the graphical model with $p(\mathcal{O}_t \mid s_t, a_t) = \exp(r(s_t, a_t))$ , derive $p(\tau \mid \mathcal{O}_{1:T}) \propto p(\tau) \exp(\sum_t r(s_t, a_t))$ . Then specialize to deterministic dynamics: for two feasible trajectories $\tau_1, \tau_2$ with total rewards $R_1 > R_2$ , what is the ratio of their posterior probabilities, and what does this predict about the monkey choosing among equal-reward routes?

Show solution

By the definition of conditional probability, $p(\tau \mid \mathcal{O}_{1:T}) = p(\tau, \mathcal{O}_{1:T}) / p(\mathcal{O}_{1:T})$ . The numerator factorizes over the model’s CPDs: $p(\tau) \prod_t p(\mathcal{O}_t \mid s_t, a_t) = p(\tau) \prod_t \exp(r(s_t, a_t))$ , and a product of exponentials is the exponential of the sum, giving $p(\tau)\exp(\sum_t r)$ ; the denominator is a constant in $\tau$ , hence the proportionality. With deterministic dynamics $p(\tau)$ is an indicator of physical consistency, so for two feasible trajectories the ratio is $p(\tau_1 \mid \cdot)/p(\tau_2 \mid \cdot) = \exp(R_1 - R_2)$ — probability decays exponentially in the reward gap. If $R_1 = R_2$ the ratio is 1: equal-reward routes are chosen uniformly at random, which is exactly the observed variety in how the monkey reaches the target while still (almost) never missing it, since missing costs a large reward gap.

Problem 19.2 softmax backup

Define $V_t(s_t) = \log \beta_t(s_t)$ and $Q_t(s_t, a_t) = \log \beta_t(s_t, a_t)$ . (a) Derive $V_t(s_t) = \log \int \exp(Q_t(s_t, a_t))\, da_t$ from the message recursion with a uniform action prior. (b) For discrete actions with values $Q = (10, 9, 2)$ , compute the soft max and compare it to the hard max; repeat for $10 \times Q$ . What does this say about the effect of reward magnitude?

Show solution

(a) The state message is $\beta_t(s_t) = \mathbb{E}_{a_t \sim p(a_t \mid s_t)}[\beta_t(s_t, a_t)]$ ; with a uniform prior this is a constant times $\int \beta_t(s_t, a_t)\, da_t$ . Taking logs and dropping the additive constant (it shifts all values equally and cancels in the policy $\exp(Q - V)$ ), $V_t = \log \int \exp(Q_t)\, da_t$ . (b) $\log(e^{10} + e^9 + e^2) = 10 + \log(1 + e^{-1} + e^{-8}) \approx 10.31$ , versus a hard max of 10 — the soft max exceeds the max slightly, dominated by the largest term. For $(100, 90, 20)$ : $100 + \log(1 + e^{-10} + e^{-80}) \approx 100.00005$ , essentially the hard max. Scaling rewards up (equivalently, lowering the temperature) makes the soft max converge to the hard max, which is why soft optimality “reduces to classical hard optimality as the reward magnitude increases.”

Problem 19.3 optimism bias

An action leads to a jackpot state with $V_{t+1} = 10$ with probability $0.001$ and to a state with $V_{t+1} = 0$ otherwise. Compare the exact-inference backup $\log \mathbb{E}[\exp(V_{t+1})]$ with the standard backup $\mathbb{E}[V_{t+1}]$ . Which inference question does the first answer, and how does the variational formulation change the backup?

Show solution

Standard: $\mathbb{E}[V_{t+1}] = 0.001 \times 10 = 0.01$ — nearly worthless. Exact inference: $\log(0.001\, e^{10} + 0.999\, e^0) = \log(22.03 + 0.999) \approx \log(23.03) \approx 3.14$ — hundreds of times larger, dominated entirely by the rare lucky outcome. The log-exp backup answers “given that you obtained high reward, what happened?” — and conditioned on success, the posterior transition probability of the jackpot state is inflated: the model changes the dynamics to agree with the evidence, which real dynamics are not allowed to do. The variational formulation fixes the dynamics inside $q$ to the true $p(s_{t+1} \mid s_t, a_t)$ , and the resulting dynamic program replaces the backup with the ordinary $Q_t = r + \mathbb{E}[V_{t+1}]$ , scoring the action at its honest $0.01$ while keeping the soft max over actions (which is the part we actually wanted).

Problem 19.4 ELBO cancellation

Write out $\log p(\mathcal{O}_{1:T}) \ge \mathbb{E}_{z \sim q}[\log p(x, z) - \log q(z)]$ for $x = \mathcal{O}_{1:T}$ , $z = (s_{1:T}, a_{1:T})$ , and $q = p(s_1) \prod_t p(s_{t+1} \mid s_t, a_t)\, q(a_t \mid s_t)$ . Show exactly which terms cancel and why the specific choice of variational family is what makes the bound collapse to reward plus entropy.

Show solution

$\log p(x, z) = \log p(s_1) + \sum_t \log p(s_{t+1} \mid s_t, a_t) + \sum_t \log p(\mathcal{O}_t \mid s_t, a_t)$ and $\log q(z) = \log p(s_1) + \sum_t \log p(s_{t+1} \mid s_t, a_t) + \sum_t \log q(a_t \mid s_t)$ . Because $q$ was defined with the same initial-state distribution and the same transitions as $p$ , the $\log p(s_1)$ terms and every transition term appear identically in both and cancel in the difference. What remains is $\mathbb{E}_q[\sum_t \log p(\mathcal{O}_t \mid s_t, a_t) - \log q(a_t \mid s_t)]$ , and substituting $\log p(\mathcal{O}_t \mid s_t, a_t) = r(s_t, a_t)$ and using linearity gives $\sum_t \mathbb{E}_{(s_t, a_t) \sim q}[r(s_t, a_t) + \mathcal{H}(q(a_t \mid s_t))]$ . Had we let the variational family learn its own dynamics, the transition terms would not cancel, the bound would contain dynamics ratios, and — worse — the optimizer could exploit them, reintroducing exactly the optimism the constraint was built to remove. Fixing the dynamics is simultaneously what simplifies the algebra and what makes the answer the true expected reward under real transitions.

Problem 19.5 action prior folding

Levine claims a non-uniform action prior $p(a_t \mid s_t)$ can be absorbed into the reward, so assuming a uniform prior loses no generality. Prove it: show that running the backward pass with prior $p(a_t \mid s_t)$ and reward $r$ gives the same value function as running it with a uniform prior and reward $\tilde{r}(s_t, a_t) = r(s_t, a_t) + \log p(a_t \mid s_t)$ .

Show solution

With a non-uniform prior, the state message is $\beta_t(s_t) = \int p(a_t \mid s_t)\, \beta_t(s_t, a_t)\, da_t$ , so in log space $V_t(s_t) = \log \int \exp\!\big(Q_t(s_t, a_t) + \log p(a_t \mid s_t)\big) da_t$ where $Q_t = r + \log \mathbb{E}[\exp(V_{t+1})]$ . Define $\tilde{Q}_t(s_t, a_t) = r(s_t, a_t) + \log p(a_t \mid s_t) + \log \mathbb{E}[\exp(V_{t+1})]$ — i.e., the $Q$ you would compute from reward $\tilde{r} = r + \log p(a_t \mid s_t)$ . Then $Q_t + \log p(a_t \mid s_t) = \tilde{Q}_t$ exactly, so $V_t(s_t) = \log \int \exp(\tilde{Q}_t)\, da_t$ : the same log-sum-exp a uniform-prior backward pass would produce from $\tilde{r}$ (up to the constant a uniform prior contributes, which cancels in the policy). Since the $V$ ‘s agree at every step by backward induction (the base case $T$ is the same identity), the two runs produce identical value functions and identical policies $\pi = \exp(\tilde{Q} - V)$ . Hence any prior preference — e.g., a lazy monkey favoring small actions — is just a reward shaping term, and the uniform-prior assumption is fully general.