Study notes: How Does a Policy Improve?

Write T_{\pi'} for the Bellman operator that backs continuation values up through one decision under \pi'. The assumption is T_{\pi'}V^\pi\ge V^\pi. The operator is monotonic, because it averages with nonnegative probabilities and multiplies by \gamma\ge0, so applying it repeatedly gives V^\pi\le T_{\pi'}V^\pi\le T_{\pi'}^2V^\pi\le\cdots. It is also a contraction: each backup shrinks the largest difference between two bounded value functions by at least a factor of \gamma. After n decisions the effect of reverting to the old policy is at most \gamma^n times the largest difference between V^\pi and V^{\pi'}, which vanishes. The chain converges to V^{\pi'}, so V^{\pi'}\ge V^\pi.

The expectation is over states visited by the new policy, and the data usually come from the old one. Replacing d^{\pi'} by d^\pi gives a manageable surrogate. It is an approximation whose quality depends on how much the behavior changes. A trust region helps keep the data distribution relevant to the policy about to be deployed.

• Supply simulated transitions for value learning: Dyna, and SimPLe at the scale of Atari.
• Train an actor through imagined trajectories: Dreamer.
• Evaluate candidate actions during planning: model-predictive control, or the tree search of AlphaZero and MuZero.

Planning makes extra computation available at deployment, and the computation must finish before the next action is due. Dreamer’s actor amortizes instead: a training cost is paid so that repeated decisions become cheaper. Amortization does not guarantee that the cheap action matches what a planner with more time would choose.

The weights are (e^4,e^4,1) normalized, about (0.495,\,0.495,\,0.009). The average command is 0: straight into the obstacle, the worst candidate. A lower temperature moves the remaining weight onto the two tied winners and leaves the average at 0. An expressive distribution over plans can preserve alternatives that a single mean collapses. Selecting, mixing and fitting candidates are different operations even when their scores are identical.

Setting the derivative to zero with a multiplier for normalization gives q^*(a)=p(a)e^{Q(a)/\eta}/Z with Z=\sum_b p(b)e^{Q(b)/\eta}. It is a softmax of \log p(a)+Q(a)/\eta, so a rare action needs enough value to overcome its small prior. As \eta grows the update approaches the reference; as it shrinks, probability concentrates on the best supported action. The optimal value of the objective is \eta\log Z.

The weights p(a)e^{2Q(a)} are about (0.6,\,2.217,\,5.460), so q^*\approx(0.073,\,0.268,\,0.660). The expected value rises from 0.5 to about 1.587. Greedy improvement would reach 2 and discard both alternatives.

Minimizing D_{\mathrm{KL}}(p\|q) is mass-covering: the expectation is under p, so q must have mass wherever p does, and it spreads over both peaks and the valley between. Minimizing D_{\mathrm{KL}}(q\|p) is mode-seeking: the expectation is under q, so q avoids the valley and settles on one peak. If q’s family can represent p exactly, both give q=p.

The temperature is the Lagrange multiplier of the constraint, found by minimizing g(\eta)=\eta\epsilon+\eta\log\sum_a p(a)e^{Q(a)/\eta}. Since g'(\eta)=\epsilon-D_{\mathrm{KL}}(q_\eta\|p), the minimization adjusts the temperature until the update spends exactly the allowed divergence. If the best action already fits within the budget, the constraint is inactive.

It stays at zero: putting mass there makes D_{\mathrm{KL}}(q\|p) infinite. Regularization toward a prior preserves experience and cannot produce alternatives outside the prior’s support.

By Bayes’ rule p_\theta(\tau\mid O=1)=p_\theta(\tau)e^{R(\tau)/\alpha}/\mathbb E_{p_\theta}[e^{R/\alpha}]: the current trajectory distribution reweighted toward higher returns. For any q with the same support, Jensen’s inequality gives \log p_\theta(O=1)\ge\mathbb E_q[\log p(O=1\mid\tau)]-D_{\mathrm{KL}}(q\|p_\theta). Multiplied by \alpha this is \mathbb E_q[R(\tau)]-\alpha D_{\mathrm{KL}}(q\|p_\theta), up to the constant -C: the return-minus-divergence objective. The gap is D_{\mathrm{KL}}(q\,\|\,p_\theta(\tau\mid O=1)), so the bound is tight when q is the posterior.

The E-step fixes \theta and sets q to the posterior over trajectories. The M-step fixes q and adjusts \theta to raise the bound. Its policy-dependent part is \max_\theta\mathbb E_{\tau\sim q}\bigl[\sum_t\log\pi_\theta(a_t\mid s_t)\bigr]: supervised learning on targets constructed to favor good behavior. After the E-step the bound equals the evidence at the old parameters, the M-step raises the bound, and the evidence at the new parameters is at least the new bound. If the E-step is restricted or approximate, the starting bound need not be tight, and raising a loose bound does not establish that the evidence rose. Sampling error and incomplete neural optimization add further qualifications.

First, \alpha\log\mathbb E_{p_\theta}[e^{R/\alpha}] is not the expected return. For large \alpha it is the mean plus \operatorname{Var}(R)/(2\alpha), so unrestricted conditioning can favor a risky lottery whose lucky outcomes receive enormous weight. Second, a trajectory posterior can change the apparent dynamics: among journeys conditioned on success, the gate appears open more often than it opens. The robot controls its actions, not the gate. The repair restricts q(\tau) to \rho_0(s_0)\prod_t q(a_t\mid s_t)P(s_{t+1}\mid s_t,a_t), which can change action probabilities and keeps the initial-state distribution and the transitions fixed. The optimizer of this restricted ELBO generally does not equal the unrestricted posterior.

Flow conservation: \sum_a\nu(s',a)=\sum_{s,a}\nu(s,a)P(s'\mid s,a), so the probability flowing into a state equals the state’s marginal probability. Without it the optimizer could put all its mass at the delivery destination without paying for the journeys that arrive there. Raising one \nu(s,a) contributes -V(s) through the marginal at s and the expected successor multiplier through the incoming flows. With the reward, the coefficient of \nu(s,a) is \delta_V(s,a)=r(s,a)+\sum_{s'}P(s'\mid s,a)V(s')-V(s). For fixed V and \eta this is the tilt with state-action pairs as candidates: \nu^*(s,a)\propto\mu(s,a)e^{\delta_V(s,a)/\eta}. Here V starts as a set of constraint multipliers, prices that enforce consistency between states, and not as an estimate of a fixed policy’s return.

Each sample gets a weight proportional to e^{Q/\eta}: 1 for a wait, 7.39 for a detour, 54.6 for the shortcut. The three groups total (6,\,22.2,\,54.6), which normalizes to (0.073,\,0.268,\,0.660): the tilt q^*. Sampling from \pi_i already represents the reference factor. Multiplying by the density again weights the actions by \pi_i^2e^{Q/\eta}, about (0.229,\,0.423,\,0.348). The reference is counted twice, and the shortcut loses its lead.

One gradient step in policy parameters need not simultaneously discover good actions, preserve sensible exploration and represent the improved behavior. The E-step constructs a better distribution from estimated values, and the M-step solves a supervised fitting problem. The quality of the estimate and the quality of the fit are distinct responsibilities. The fit maximizes weighted likelihood, which minimizes D_{\mathrm{KL}}(q_i\|\pi_\theta): the mass-covering direction, with a further trust region on how far the parameterized policy moves. A successful M-step learns what distinguishes this junction from other junctions. A poor one memorizes a background texture or pushes all similar scenes toward the shortcut.

Subtracting V(s) multiplies every action weight at s by the same factor, and the factor cancels when each state’s distribution is normalized separately. Weighted regression does not normalize state by state: it fits one network across states with the unnormalized weights e^{\hat A(s,a)/\beta}. With raw returns the second junction’s weights are larger by e^{100/\beta}, and that junction dominates the dataset-wide objective. Subtracting each junction’s baseline removes the common offset. The intended conditional choices stay the same; the allocation of the learner’s finite capacity changes. Weight clipping can further break the cancellation.

The means are 2 and 1. Exponentiating them gives e^2\approx7.39 against e\approx2.72, so the certain reward gets about 0.73 of the probability. Weighting realized returns gives e^2 against 0.9+0.1e^{10}\approx2204, so the lottery gets about 0.997. Conditioning on a return of 10 selects only lottery actions, and playing them again leaves the winning probability at 0.1. Conditioning can select favorable environmental outcomes that a policy cannot cause. MPO’s critic-based step exponentiates an estimated expected value; AWR with sampled returns can give unusually lucky outcomes disproportionately large weights.

With a multiplier c for normalization, the derivative is Q(a)+\lambda p(a)/q(a)-c. Setting it to zero gives q^*(a)=\lambda p(a)/(c-Q(a)) with c>\max_aQ(a): a rational expression, not an exponential one. For the robot, q^*=(0.3/2.06994,\;0.15/1.06994,\;0.05/0.06994)\approx(0.1449,\,0.1402,\,0.7149). The tilt gave (0.0725,\,0.2678,\,0.6597). Both favor the shortcut and keep different amounts of the other actions. Equal penalty coefficients do not imply equal divergence budgets, so the comparison is illustrative.

Since 1+N(s,a)=\hat\pi(a\mid s)\,(|\mathcal A|+N), the score is Q(s,a)+\lambda_N\,p(a\mid s)/\hat\pi(a\mid s) with \lambda_N=c_{\mathrm{puct}}\sqrt N/(|\mathcal A|+N). The derivative of the reversed-KL objective is Q(a)+\lambda p(a)/q(a)-c, and c is common to all actions. With q=\hat\pi and \lambda=\lambda_N, PUCT selects the action with the largest derivative, and visiting it increases its share of the counts. Values and counts evolve during search, so this is not an exact solve for one fixed objective. Final training targets can also use different visit-count conventions.

Rearranging the optimum gives r(x,y)=\beta\log\frac{\pi^*(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}+\beta\log Z(x). In the difference r(x,y_+)-r(x,y_-) the term \beta\log Z(x) cancels. Replacing \pi^* with the trainable \pi_\theta gives the margin m_\theta=\beta\Bigl[\log\frac{\pi_\theta(y_+\mid x)}{\pi_{\mathrm{ref}}(y_+\mid x)}-\log\frac{\pi_\theta(y_-\mid x)}{\pi_{\mathrm{ref}}(y_-\mid x)}\Bigr], and the loss is -\log\sigma(m_\theta) averaged over labeled pairs. At \pi_\theta=\pi_{\mathrm{ref}} the margin is zero and the predicted preference probability is 1/2. The KL-regularized optimum supplies the parameterization; the implemented loss fits preferences directly, without a separately trained reward model.

We can ask whether the simulated experience is faithful, whether the values are correct, whether the preference signal is useful, and whether the actor represents the resulting distribution. At the gate:

• The model hallucinates an opening. The critic may correctly evaluate the hallucinated world.
• The model predicts the closure, and the critic assigns a large terminal value to waiting there.
• Both predictions are sound, and the fitting objective puts too little weight on the rare closed-gate examples.
• Perception mistakes which junction the robot occupies, so the actor acts from the wrong latent state.

A successful system depends on all the parts, and an observed failure need not indict all of them.

The weights are (0.6e^{10},\,0.3e^{2},\,0.1e^{4})\approx(13216,\,2.2,\,5.5), so q^*\approx(0.9994,\,0.0002,\,0.0004). By the critic’s scores the expected value rose from 3.5 to about 5.0. The true expected value fell from 0.5 to about 0.001. The optimization functioned as specified. What failed is the premise that its scores measure the desired consequences. A KL constraint limits the size of a step relative to its reference: it does not calibrate the critic, constrain every future state equally, or ensure that the reference is appropriate. Repeated small steps can also add up to a large departure.

Where do the candidates come from? How are they evaluated? How is an improved distribution built from them? How is the improvement retained?

The exponential tilt is an exact finite-action optimizer under stated assumptions. The steps away from it:

• Replacing values with a neural critic, as in MPO’s E-step.
• Sampling a few candidates, as MPO and POCO do at each state.
• Omitting a state-dependent normalizer, which changes a shared-parameter regression objective, as in AWR.
• Reversing a KL, which changes its optimizer, as in the analysis of AlphaZero’s search.
• Replacing likelihood with a generative surrogate, which changes what can be inferred from an EM argument, as in POCO’s flow-matching loss.