---
title: "Policy Gradient"
description: "Watch returns change action probabilities with REINFORCE"
source: https://vizlearn.app/ai/reinforcement/policy-gradient
access: premium (interactive player requires Lifetime; this doc is free)
---

# Policy Gradient

Policy gradient methods learn which actions to take from rewards, adjusting action probabilities to maximize expected return. REINFORCE uses the return from a sampled episode to reinforce or discourage its actions. A correct update can still make success less likely: in the first episode here, the chance of reaching the goal falls from 25% to 24.6%.

Imagine learning a route from a single verdict at the destination. You took a right, then a left, and missed the goal. That verdict does not tell you which turn caused the failure. REINFORCE assigns the return to both decisions. The useful first turn can get penalized along with the second one.

This is the problem policy gradients address: you can try actions and observe rewards, but have no labeled correct action for each state. The policy supplies differentiable action probabilities even when the environment's choices and outcomes are discrete.

## How Policy Gradient Works

A **policy** maps a state to a distribution over actions. At state `s`, write the probability of choosing right as `p = π(R | s)`. Left has probability `1 − p`. Sampling from this distribution allows either action to occur; taking the largest probability every time would be a different decision rule.

Here the environment has two decisions. From `s0`, left leads to `sL` and right leads to `sR`. Each of those states has another left/right choice, which ends the episode. Only the path **R → R** reaches the goal and earns +1. All other paths earn the configured failure reward, initially −1. The first transition always pays 0.

Each state owns one parameter, θ. Its right-action probability is the sigmoid of that parameter:

`p = 1 / (1 + exp(−θ))`

At θ = 0, the actions are equally likely. Increasing θ makes right more likely; decreasing it favors left. This is also a two-action softmax with logits `[0, θ]`. Unlike a neural policy, these three state parameters do not share weights.

### A return weights a direction

The **return** `G_t` is the sum of rewards from decision `t` onward. With no discounting and only a terminal reward, both decisions in this example receive the same return. In a longer task with intermediate rewards, their returns can differ.

For a binary action, encode right as `a = 1` and left as `a = 0`. The derivative of the sampled action's log probability is:

`∂ log π(a | s) / ∂θ = a − p`

That derivative points toward making the sampled action more likely. For L at `p = 0.5`, it is `0 − 0.5 = −0.5`: reducing θ favors L. For R, it is `1 − 0.5 = +0.5`: increasing θ favors R. No probability changes yet; this first step only finds a direction.

REINFORCE then weights that direction by the return. Positive return keeps it, negative return reverses it, and zero return removes the contribution. The parameter moves in a gradient **ascent** step:

`g = G_t × (a − p)`

`θ ← θ + ηg`

For a positive learning rate and this independent-state parameterization, positive return increases the sampled action's probability, negative return decreases it, and zero return leaves it unchanged. Both action probabilities still sum to 1. Increasing right's probability necessarily decreases left's.

The gradient is evaluated using the probability that generated the action. The policy stays fixed during an episode, and both contributions are computed from those saved probabilities. The display updates the last state's coordinate and then the root's coordinate; because the two states have independent parameters, these writes give the same result as applying the episode's gradient vector at once.

### Why a discrete action can still supply a gradient

You do not need to differentiate a sampled left turn or the environment's reward function. The policy assigns a differentiable probability to each possible action. The log-derivative identity rewrites the derivative of expected return as an expectation involving those log probabilities. An episode supplies a sample of that expectation. This is the basis of the [policy-gradient derivation in Spinning Up](https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html).

The environment map is available to you as an explanation. The updater receives the sampled actions, their probabilities, and the observed return. It never receives a target label saying "choose R" or the exact gradient of the chance of reaching the goal.

## Policy Gradient Step by Step

The default run uses learning rate **0.5**, failure reward **−1**, initial parameters `[0, 0, 0]`, and four episodes. A fixed random seed makes the draws reproducible. Initially, the chance of reaching the goal is `0.5 × 0.5 = 25%`.

### Episode 1: a failed route penalizes a useful action

The first draw is approximately `u = 0.254283`. It lies below the root's right-action probability of 0.5, so the policy chooses R. At `sR`, the next draw is `0.541644`, above 0.5, so the policy chooses L. The path RL fails and returns −1.

At the second decision, L means `a = 0`. Its log-probability gradient is −0.5, which points toward choosing L more often. The negative return reverses that direction:

`g_sR = −1 × (0 − 0.5) = +0.5`

`θ_sR = 0 + 0.5 × 0.5 = 0.25`

Right's probability at `sR` rises to 56.2%; the sampled left action falls to 43.8%. The negative return reduced the probability of the action that was taken.

Now apply the same return to the first decision. R means `a = 1`:

`g_s0 = −1 × (1 − 0.5) = −0.5`

`θ_s0 = 0 + 0.5 × (−0.5) = −0.25`

Right's probability at the root falls to 43.8%. Yet right is required to reach the goal. The new chance of success is `0.4378235 × 0.5621765 ≈ 24.6%`, slightly below the original 25%. The failed path became less likely, but so did the only successful path.

### Episode 2: reinforcing success

The next draws, approximately `0.000823` and `0.407636`, fall within the right-action intervals at both states. This episode takes RR and earns +1.

The two gradients are now positive: about `0.5621765` at `s0` and `0.4378235` at `sR`. After updating, the root's right-action probability is about 50.8%, and `sR`'s is about 61.5%. Their product gives a goal probability of **31.2%**.

The last two episodes both sample LL and receive −1. Their updates move probability away from L at the visited states. The complete ledger is:

| Episode | Path | Return | Goal probability after updating |
| --- | --- | ---: | ---: |
| Initial | — | — | 25.0% |
| 1 | RL | −1 | 24.6% |
| 2 | RR | +1 | 31.2% |
| 3 | LL | −1 | 35.1% |
| 4 | LL | −1 | 39.3% |

These percentages are exact policy evaluations rounded for display, not the fraction of episodes that succeeded. One of the four sampled episodes succeeded, while the final policy would succeed with probability about 39.3% on a new episode. Four episodes do not establish convergence.

### What to check in a larger task

In this tiny environment, you can compute the final policy's success probability by multiplying two known probabilities. In a larger environment, freeze a policy and evaluate it over multiple fresh episodes. Track performance across those evaluations: the return from one training episode is too little evidence to judge the policy, and older training episodes may have used different policies.

If the policy does not move, distinguish a missing learning signal from a disabled update. The **No reward for a miss** preset shows `G = 0` removing the gradient. **Freeze the policy** keeps the gradient but sets `η = 0`. These failures look alike in the probability bars; the arithmetic tells you where learning stopped.

## Time and Space Complexity

For a tabular binary-action policy with `S` states and an episode of `T` decisions, sampling the episode and accumulating its policy-gradient contributions have **time complexity O(T)**. Computing reward-to-go with a backward accumulation is also O(T). Policy parameters and the saved episode require O(S + T) space.

A neural policy adds the cost of forward and backward passes. The animation also stores copies of its three-state policy for seeking; those display snapshots are separate from REINFORCE's storage requirements.

## Edge Cases

- **Zero return.** Set the failure reward to 0. The first RL episode still fails, but now both gradients are zero. At the default η = 0.5, the second episode samples RR and learns from +1; the remaining two failures again leave the policy unchanged. A lack of feedback does not mean success is impossible to discover later.
- **Zero learning rate.** Set η to 0. Returns and gradients are still computed, but every parameter update is zero. The policy stays at 50/50.
- **Unvisited states.** An episode provides no contribution for a state it never visited. The first two default episodes leave `sL` unchanged. Neural policies with shared parameters need not preserve the outputs at unvisited states this way.
- **Larger updates.** η scales the logit change for a fixed sampled episode; a larger η moves it further only where the gradient is nonzero. Later samples may then change because the policy changed. No setting in this control range is a universal stability boundary for policy gradients.
- **Sparse feedback.** With zero reward on a miss and no baseline, unsuccessful episodes provide no gradient. Longer tasks can require many samples before finding a rewarding trajectory. The two-decision environment keeps that issue small enough to inspect.

## Common Mistakes

- **Treating the return as a label for each action.** A failed route does not prove every decision on it was bad. The root-R update in episode 1 is a concrete counterexample.
- **Expecting every update to improve performance.** The sampled gradient is noisy. The default run starts with a valid update that lowers the chance of success.
- **Confusing a gradient's sign with the sampled action's probability change.** A positive θ-gradient favors R. If the sampled action was L and the return was negative, that same positive gradient reduces the sampled action's probability.
- **Using descent with the wrong sign.** Maximizing return uses ascent. An optimizer that minimizes a loss can instead minimize negative return-weighted log probability. Omitting that minus sign reverses the intended update.
- **Reusing an episode indefinitely.** REINFORCE samples under the current policy. Repeatedly optimizing the same old episode is not the same estimator evaluated on fresh on-policy data.
- **Comparing raw returns across different reward settings.** A failure reward of −1 and a failure reward of 0 define different numerical objectives. Here both favor the same RR route, but their expected rewards have different scales and offsets. Compare goal probability when checking the runs.

## From REINFORCE to a Baseline

An action-independent baseline replaces the return weight with `G_t − b(s_t)`. This subtracts a reference value without changing the expected policy gradient, and a suitable baseline can reduce its variance. The sign then means better or worse than the baseline; a positive raw return alone no longer determines the update direction.

A learned state-value estimate is a common baseline. Actor-critic methods add a value estimator and often bootstrap their learning signal. That extension is described in the [VPG reference](https://spinningup.openai.com/en/latest/algorithms/vpg.html); the animation here uses zero-baseline REINFORCE. For the neighboring pieces, [Gradient Descent](https://vizlearn.app/ai/neural-networks/gradient-descent) explains step size, [Softmax & Cross-Entropy](https://vizlearn.app/ai/neural-networks/softmax-cross-entropy) connects logits to categorical probabilities, and [Backpropagation](https://vizlearn.app/ai/neural-networks/backpropagation) follows gradients through shared neural weights.

## A Note on Simplification

This is a two-step tabular explainer, not a substitute for formal reinforcement-learning study. It uses terminal rewards, no discounting, one episode per update, and independent state logits. It omits learned critics, shared neural parameters, batching, entropy regularization, and long horizons. The exact goal-probability calculation is a diagnostic shown to you; it is never used to train the policy.