Visual walkthrough — Trust Region Policy Optimization (TRPO)
We will use only a 2-dimensional parameter space the whole time: . Two numbers. That is enough to see everything; the real thing just has millions of dimensions with the same geometry.
Step 1 — The map: what "parameter space" even means
WHAT. Our neural-network policy is controlled by numbers . Pick two of them, and , and treat them as a point on a flat map. Every point on this map = one policy = one way of acting.
WHY. We are about to talk about "taking a step", "a direction", "how far". Those words only make sense once you have a place to stand and a plane to walk on. So first we draw the ground.
PICTURE. The blue dot is our current policy (the one we already ran and collected data from). The whole plane is every other policy we could jump to.
Step 2 — The reward we want to climb (the surrogate hill)
WHAT. For each point on the map we can (approximately) ask: "how much better is this policy than ?" That number is the surrogate objective from the parent note. Draw it as height above the map — a hill.
- — the advantage (see Advantage Estimation (GAE)): positive = good action, negative = bad action.
- — the importance ratio (see Importance Sampling): reuses old samples for a new policy by re-weighting them.
- — just "average". Nothing scary: add up over your data, divide by count.
WHY draw it as a hill? Because "make the policy better" now becomes the single visual instruction "walk uphill". Everything after this is how far uphill we dare to walk.
PICTURE. Contour lines = lines of equal height, like a topographic map. Closer lines = steeper slope. Our dot sits on the slope; up-and-to-the-right is higher.
Step 3 — The gradient: which way is uphill right here
WHAT. Stand on the dot. The direction of steepest uphill is a little arrow called the gradient, written .
- (read "nabla" / "gradient of") — the tool that answers "which way is up?". Why this tool? Because we want to increase and literally packages the uphill direction and steepness into one arrow.
- — a partial derivative: "if I nudge only a hair, how fast does the height change?" The other coordinate held still.
WHY. This is exactly plain Policy Gradient Methods: it tells us the best local direction. If we followed only we'd be doing vanilla policy gradient — and hit the collapse problem the parent warned about.
PICTURE. The coral arrow points across the contour lines (steepest way to higher ground). It is the naive answer we are about to correct.
Step 4 — The catch: not all steps are equal (distribution vs. parameters)
WHAT. Moving a fixed distance in does not change the policy by a fixed amount. Some directions barely change how the policy behaves; others change it wildly for the same footstep.
WHY this matters. Our real safety rule is "don't change the policy's behaviour too much" — because a wildly changed policy collects wildly different data and can collapse. So a plain circle "" around our dot is the wrong notion of "close".
PICTURE. Two footsteps of equal length in : the mint one moves along a "cheap" direction (policy barely changes), the coral one along an "expensive" direction (policy changes a lot). Same length, very different danger.
Step 5 — The right ruler: KL divergence and its ellipse
WHAT. KL Divergence measures how different two probability distributions are — here, the old policy vs. the candidate new policy at a state. TRPO's trust region says:
- — average (over states) policy difference, in nats.
- — the size of our leash, e.g. . Why a constraint, not a penalty? The parent's exact penalty coefficient is so huge it forces microscopic steps; a hard cap we can tune is practical.
Now the magic. For small steps, KL is well-approximated by a quadratic bowl:
- — the step vector (where we go, relative to the dot).
- — the Fisher Information Matrix: the curvature of the KL bowl at . Big curvature in a direction = KL grows fast there = that direction is "expensive".
Why quadratic and not linear? Because at the policy equals itself, so KL is and its slope is (you can't be "a little different from yourself"). The first non-zero term is the quadratic one — that is precisely why , the second-order term, is the star.
PICTURE. The constraint is an ellipse (a squashed circle) around the dot. Cheap directions → the ellipse stretches far. Expensive directions → the ellipse is short. This ellipse is the trust region.
Step 6 — The problem, drawn: uphill arrow meets the ellipse
WHAT. Put Steps 3 and 5 together. We want to climb (go along ) but stay inside the KL ellipse. Formally:
- — a dot product: how much of our step points along the uphill arrow . Bigger = more height gained. Why a dot product? It's exactly "projection of the step onto the uphill direction" — the linear stand-in for the hill near the dot.
WHY. We reach the highest point we can touch without leaving the ellipse. The answer will sit on the ellipse's boundary (using less of the leash would waste climbing room).
PICTURE. The coral , the ellipse, and the winning point where the highest contour of just kisses the ellipse.
Step 7 — Solving it: the natural gradient falls out
WHAT. At the touching point, the uphill arrow must point straight out of the ellipse (perpendicular to its boundary) — otherwise you could slide along the boundary and climb more. Writing "gradient of objective = × gradient of constraint" (the Lagrange condition):
- — the natural gradient: the raw uphill arrow bent by the ellipse's shape. It shortens expensive directions and lengthens cheap ones. This is the corrected direction we hunted for in Step 4.
- — a scalar we still must pin down.
Fix by demanding we sit exactly on the boundary, . Substituting and solving:
Term by term:
- — direction (natural gradient).
- — the "length² in the ellipse's ruler" of that direction.
- — the exact rescale so the step lands on the KL boundary, using every drop of the leash and not one nat more.
PICTURE. The straightened, correctly-scaled arrow reaching the ellipse edge — the same arrow the parent note handed you, now earned.
Step 8 — Degenerate & edge cases (the map at its limits)
WHAT. What if inputs go extreme? Each has a picture-answer.
- (flat hill). The uphill arrow vanishes → . We stand still. Correct: if nothing improves locally, don't move.
- isotropic, (ellipse is a circle). Then : the natural gradient is just the ordinary gradient, rescaled. TRPO gracefully reduces to plain gradient ascent when all directions cost the same.
- A very cheap direction ( tiny along one axis). The ellipse stretches enormously there; blows that component up — TRPO happily takes a big step where the policy barely reacts. This is the feature, not a bug.
- Approximation lies (true KL > ). The quadratic bowl was only local; far out it can under-estimate. So TRPO backtracks: shrink until the measured KL and the surrogate truly rose.
PICTURE. Three mini-panels: circle-ellipse (case 2), stretched ellipse (case 3), and the backtracking dots walking back inside the true region (case 4).
The one-picture summary
Everything at once: the hill (surrogate ), the naive uphill arrow , the KL ellipse (trust region shaped by ), the bent-and-scaled natural-gradient step landing on the boundary, and a backtracked point pulled inside the true region.
Recall Feynman: tell it back in plain words
We drew every policy as a dot on a map, and "how good it is" as a hill we want to climb. The steepest-up arrow is the ordinary gradient — but footsteps aren't fair: some directions change the policy's behaviour a lot, some barely at all. So we drew a special leash — an ellipse — that measures distance the honest way (how different the two policies act), using the Fisher matrix as the ruler. We then asked: climb as high as possible without leaving the ellipse. The answer bends into (the natural gradient) and stretches it to just touch the ellipse edge — that stretch factor is the square-root term. If the hill is flat we don't move; if all directions cost the same the ellipse is a circle and we're back to plain gradient ascent; and because our ellipse is only a local guess, we walk the step back until the real leash is satisfied. That's TRPO.
Recall
What geometric object is the TRPO trust region on the parameter map? ::: An ellipse , shaped by the Fisher matrix . Why is the KL constraint quadratic (uses ) instead of linear? ::: At the KL is and its slope is ; the first non-zero term is the second-order (curvature) one, given by . What does multiplying by accomplish geometrically? ::: It bends the raw uphill arrow to respect the ellipse — shrinking expensive directions, stretching cheap ones — giving the natural gradient. What does the scalar do? ::: Rescales the step so it lands exactly on the KL boundary, using all of the leash . When does TRPO collapse to ordinary gradient ascent? ::: When (all directions equally costly), so .