14 · The policy-gradient idea
Instead of learning values and acting greedily, optimise the POLICY directly. Parameterise it with a softmax over preferences and push those preferences in the direction that earns more reward.
Parameterise a policy π_θ(a|s) = softmax(θ) and optimise θ by GRADIENT ASCENT. The policy-gradient theorem says: move θ in the direction ∇θ log π_θ(a|s) scaled by the return — push up the preferences of actions that led to high reward.
Without this:
Without policy gradients you can only act greedily on a learned value table, which breaks down for continuous or huge action spaces and can't represent stochastic optimal policies. Policy gradients optimise behaviour directly and underlie PPO, A2C, and RLHF.
Everything so far learned a value (Q or V) and then acted greedily on it. Policy-gradient methods take a different route: they parameterise the policy itself and optimise it directly by gradient ascent on the expected return. No value table is required to act — the policy is the model.
The standard parameterisation for a small discrete action space is a softmax over preferences. Give each action a real-valued preference θ[a] (one number per action, learnable), and define
π(a) = exp(θ[a]) / Σ_b exp(θ[b])
Higher preference → higher probability, but every action keeps a non-zero chance, so the policy is stochastic and can explore.
How do we improve θ? The policy-gradient theorem gives a beautifully simple rule. To make good actions more likely, nudge θ in the direction of the score function — the gradient of the log-probability of the action you took — scaled by how much reward that action earned:
θ ← θ + α · (reward) · ∇θ log π(a)
For a softmax the score has a clean closed form: ∂/∂θ_b log π(a) = 1[b=a] − π(b). In words: push the chosen action's preference up, and push every action's preference down in proportion to its current probability. Multiply by a positive reward and the chosen action becomes more likely; multiply by a negative reward and it becomes less likely.
Below we take one gradient-ascent step on a 3-action softmax policy and watch the rewarded action's probability rise.
Python (in browser)
One policy-gradient step on a softmax policy: θ ← θ + α·reward·∇log π(a). The rewarded action's probability rises; the others fall.
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
In the softmax policy-gradient update θ ← θ + α · reward · ∇θ log π(a), what is the role of the reward?
- Policy-gradient methods parameterise the policy directly (e.g. softmax over preferences θ) and optimise θ by gradient ascent on expected return — no value table needed to act.
- The update is θ ← θ + α · return · ∇θ log π(a): the score ∇θ log π(a) gives the direction, the return scales it.
- For softmax the score is 1[b=a] − π(b): push the chosen action's preference up, push every action down by its probability.
Softmax policies and the score function are the foundation of REINFORCE, A2C, PPO, and the RL step in RLHF that fine-tunes instruction-following LLMs.
If you remove it: Without direct policy optimisation you are stuck with value-greedy agents that struggle with continuous/huge action spaces and cannot represent genuinely stochastic optimal policies.