Ahmed Haroon

Optimizing the Policy Directly

Value-based methods learn $Q(s,a)$ and obtain a policy through an argmax. Policy-gradient methods instead represent the policy itself:

\pi_\theta(a\mid s) =\Pr(A_t=a\mid S_t=s;\theta).

We optimize its parameters to maximize expected return:

\boxed{ \theta^*\in\arg\max_\theta J(\theta) }.

Direct policy optimization is useful when the action space is continuous, smooth changes in action probabilities matter, or the best policy is genuinely stochastic. It also avoids solving

\arg\max_a Q(s,a)

every time the agent acts.

For discrete actions, a neural policy commonly produces logits $z_\theta(s,a)$ and applies softmax:

\pi_\theta(a\mid s) =\frac{\exp(z_\theta(s,a))} {\sum_{a'}\exp(z_\theta(s,a'))}.

For continuous actions, a common policy is

a_t\sim\mathcal N \left(\mu_\theta(s_t),\Sigma_\theta(s_t)\right).

Stochasticity is not merely an exploration trick. It is part of the policy class being optimized.

Suppose the same observation can be produced by two hidden states that require different actions. A deterministic observation policy must choose the same action in both. A stochastic policy can assign probability to both possibilities:

\pi_\theta(a\mid o).

We can also condition on a recurrent summary of the observation history:

h_t=f_\theta(h_{t-1},o_t,a_{t-1}), \qquad \pi_\theta(a_t\mid h_t).

The derivation below does not require the policy input to be a fully observed Markov state. It only requires the sampled action probability to be differentiable with respect to $\theta$ .

The Policy Objective

For a finite-horizon trajectory

\tau=(s_0,a_0,r_0,\ldots,s_{H-1},a_{H-1},r_{H-1},s_H),

define the return

R(\tau)=\sum_{t=0}^{H-1}r_t.

The objective is the expected return under trajectories generated by the current policy:

\boxed{ J(\theta) =\mathbb E_{\tau\sim p_\theta(\tau)}[R(\tau)] }.

The trajectory probability factorizes as

p_\theta(\tau) =\rho_0(s_0) \prod_{t=0}^{H-1} \pi_\theta(a_t\mid s_t) P(s_{t+1}\mid s_t,a_t),

where $\rho_0$ is the initial-state distribution and $P$ contains the environment dynamics.

Taking logs turns the product into a sum:

\begin{aligned} \log p_\theta(\tau) &=\log\rho_0(s_0) +\sum_{t=0}^{H-1}\log\pi_\theta(a_t\mid s_t)\\ &\quad+\sum_{t=0}^{H-1} \log P(s_{t+1}\mid s_t,a_t). \end{aligned}

Only the policy depends on $\theta$ . Therefore,

\boxed{ \nabla_\theta\log p_\theta(\tau) =\sum_{t=0}^{H-1} \nabla_\theta\log\pi_\theta(a_t\mid s_t) }.

This is why policy gradients are model-free. The dynamics determine which trajectories occur, but we never differentiate through them.

The Likelihood-Ratio Identity

Once a trajectory has been sampled, its observed rewards are fixed numbers. The dependence on $\theta$ lies in the probability of sampling that trajectory.

For any differentiable positive density,

\nabla_\theta\log p_\theta(x) =\frac{\nabla_\theta p_\theta(x)}{p_\theta(x)}.

Rearranging gives the likelihood-ratio identity

\boxed{ \nabla_\theta p_\theta(x) =p_\theta(x)\nabla_\theta\log p_\theta(x) }.

It converts a derivative of a sampling distribution into the derivative of a log-probability that we can evaluate on samples.

See derivation

Let

F(\theta)=\mathbb E_{x\sim p_\theta}[f(x)]

and assume $f$ does not explicitly depend on $\theta$ . Then

\begin{aligned} \nabla_\theta F(\theta) &=\nabla_\theta\int p_\theta(x)f(x)\,dx\\ &=\int\nabla_\theta p_\theta(x)f(x)\,dx\\ &=\int p_\theta(x) \frac{\nabla_\theta p_\theta(x)}{p_\theta(x)} f(x)\,dx\\ &=\mathbb E_{x\sim p_\theta} \left[f(x)\nabla_\theta\log p_\theta(x)\right]. \end{aligned}

Hence

\boxed{ \nabla_\theta\mathbb E_{x\sim p_\theta}[f(x)] =\mathbb E_{x\sim p_\theta} \left[f(x)\nabla_\theta\log p_\theta(x)\right] }.

Applying the identity to trajectories gives

\begin{aligned} \nabla_\theta J(\theta) &=\mathbb E_{\tau\sim p_\theta} \left[R(\tau)\nabla_\theta\log p_\theta(\tau)\right]\\ &=\mathbb E_{\tau\sim p_\theta} \left[ R(\tau) \sum_{t=0}^{H-1} \nabla_\theta\log\pi_\theta(a_t\mid s_t) \right]. \end{aligned}

Therefore,

\boxed{ \nabla_\theta J(\theta) =\mathbb E_{\tau\sim p_\theta} \left[ \sum_{t=0}^{H-1} \nabla_\theta\log\pi_\theta(a_t\mid s_t) R(\tau) \right] }.

The score

\nabla_\theta\log\pi_\theta(a_t\mid s_t)

says how the sampled action's probability changes, while the return says whether that sampled behavior was useful.

For every state,

\boxed{ \mathbb E_{a\sim\pi_\theta(\cdot\mid s)} \left[\nabla_\theta\log\pi_\theta(a\mid s)\right] =0 }.

See derivation

For a continuous action,

\begin{aligned} \mathbb E_{a\sim\pi_\theta} [\nabla_\theta\log\pi_\theta(a\mid s)] &=\int\pi_\theta(a\mid s) \nabla_\theta\log\pi_\theta(a\mid s)\,da\\ &=\int\nabla_\theta\pi_\theta(a\mid s)\,da\\ &=\nabla_\theta\int\pi_\theta(a\mid s)\,da\\ &=\nabla_\theta1\\ &=0. \end{aligned}

This identity lets us remove irrelevant rewards and subtract baselines without changing the expected gradient.

REINFORCE

Given $N$ trajectories sampled from the current policy, the Monte Carlo estimator is

\boxed{ \widehat{\nabla_\theta J} =\frac1N\sum_{i=1}^{N} R(\tau^{(i)}) \sum_{t=0}^{H-1} \nabla_\theta\log\pi_\theta \left(a_t^{(i)}\mid s_t^{(i)}\right) }.

We perform gradient ascent:

\theta\leftarrow\theta+ \alpha\widehat{\nabla_\theta J}.

This resembles weighted maximum likelihood. Ordinary maximum likelihood uses

\sum_{i,t}\nabla_\theta \log\pi_\theta(a_t^{(i)}\mid s_t^{(i)}),

while REINFORCE weights each score by its observed return. Positive weights increase the sampled action's probability; negative weights decrease it.

This statement is statistical, not deterministic. A good action can receive a poor return because of a later mistake, and a bad action can receive a good return because of luck. These errors only cancel in expectation.

Reward-to-go

The full return credits an action at time $t$ for rewards that occurred before it. Those past rewards cannot have been caused by $a_t$ , so they only add noise.

Define reward-to-go

G_t=\sum_{t'=t}^{H-1}r_{t'}.

We can replace the episode return by $G_t$ without changing the expected gradient:

\boxed{ \nabla_\theta J =\mathbb E \left[ \sum_{t=0}^{H-1} \nabla_\theta\log\pi_\theta(a_t\mid s_t)G_t \right] }.

See derivation

Expand the full-return estimator:

\nabla_\theta J =\mathbb E \left[ \sum_{t=0}^{H-1} \nabla_\theta\log\pi_\theta(a_t\mid s_t) \sum_{t'=0}^{H-1}r_{t'} \right].

For a past reward $r_{t'}$ with $t'<t$ , condition on the history through $s_t$ :

\begin{aligned} &\mathbb E \left[ r_{t'}\nabla_\theta\log\pi_\theta(a_t\mid s_t) \right]\\ &=\mathbb E \left[ r_{t'} \mathbb E_{a_t\sim\pi_\theta(\cdot\mid s_t)} \left[ \nabla_\theta\log\pi_\theta(a_t\mid s_t) \mid s_t,\text{past} \right] \right]\\ &=\mathbb E[r_{t'}\cdot0]\\ &=0. \end{aligned}

Past rewards are fixed before

Only rewards at or after $t$ remain, producing $G_t$ .

For discounted problems, we use

G_t^\gamma =\sum_{t'=t}^{H-1}\gamma^{t'-t}r_{t'}.

Discounting controls the influence of uncertain distant outcomes and makes infinite-horizon returns finite when $\gamma<1$ . It also changes the objective by preferring earlier rewards.

\textbf{REINFORCE}

\begin{aligned} &\text{Initialize policy parameters }\theta.\\[0.25em] &\text{Repeat:}\\[0.25em] &\qquad \text{collect }N\text{ complete trajectories using }\pi_\theta,\\[0.25em] &\qquad G_{H-1}^{(i)}\leftarrow r_{H-1}^{(i)},\\[0.25em] &\qquad G_t^{(i)}\leftarrow r_t^{(i)}+\gamma G_{t+1}^{(i)} \quad\text{backward through each trajectory},\\[0.25em] &\qquad \widehat g\leftarrow\frac1N\sum_i\sum_t G_t^{(i)}\nabla_\theta\log\pi_\theta (a_t^{(i)}\mid s_t^{(i)}),\\[0.25em] &\qquad \theta\leftarrow\theta+\alpha\widehat g. \end{aligned}

REINFORCE is on-policy: the trajectories must come from the policy whose gradient we estimate unless we add an explicit distribution correction.

Baselines and Advantages

Even reward-to-go varies greatly between trajectories. We can subtract a state-dependent baseline:

\boxed{ \nabla_\theta J =\mathbb E \left[ \sum_t \nabla_\theta\log\pi_\theta(a_t\mid s_t) \left(G_t-b(s_t)\right) \right] }.

The baseline supplies a reference. An action is reinforced when its return is better than expected from that state and discouraged when it is worse.

Why a baseline does not change the gradient

An action-independent baseline contributes zero in expectation.

See derivation

At one time step,

B_t =\mathbb E \left[ \nabla_\theta\log\pi_\theta(a_t\mid s_t)b(s_t) \right].

Using iterated expectation,

\begin{aligned} B_t &=\mathbb E_{s_t} \left[ b(s_t) \mathbb E_{a_t\sim\pi_\theta(\cdot\mid s_t)} [\nabla_\theta\log\pi_\theta(a_t\mid s_t)] \right]\\ &=\mathbb E_{s_t}[b(s_t)\cdot0]\\ &=0. \end{aligned}

A constant baseline such as the batch-average return is valid, but it ignores that typical return varies by state. The canonical choice is a learned value function:

V_\phi(s_t)\approx V^{\pi_\theta}(s_t) =\mathbb E_{\pi_\theta}[G_t\mid s_t].

This produces the Monte Carlo advantage estimate

\boxed{ \widehat A_t^{\mathrm{MC}} =G_t-V_\phi(s_t) }.

The policy-gradient estimator becomes

\boxed{ \widehat{\nabla_\theta J} =\frac1N\sum_{i=1}^{N}\sum_t \nabla_\theta\log\pi_\theta (a_t^{(i)}\mid s_t^{(i)}) \widehat A_t^{(i)} }.

An inaccurate state-only baseline does not bias the actor gradient. It simply removes less variance. We fit it to Monte Carlo returns:

\mathcal L_V(\phi) =\frac1M\sum_{i,t} \left(V_\phi(s_t^{(i)})-G_t^{(i)}\right)^2.

As the policy changes, $V^{\pi_\theta}$ changes too, so the baseline must continually track a moving target.

The variance-minimizing baseline

The value function is a strong baseline, but not always the exact minimum-variance choice. Let

z(a,s)=\nabla_\theta\log\pi_\theta(a\mid s).

At a fixed state, minimizing the estimator's second moment gives

\boxed{ b^*(s) =\frac{ \mathbb E[\|z(a,s)\|^2G\mid s] }{ \mathbb E[\|z(a,s)\|^2\mid s] } }.

See derivation

Because the baseline does not change the mean gradient, minimizing variance is equivalent to minimizing

\mathbb E[\|z\|^2(G-b)^2\mid s].

Differentiate with respect to $b$ :

\frac{d}{db} \mathbb E[\|z\|^2(G-b)^2\mid s] =-2\mathbb E[\|z\|^2(G-b)\mid s].

Setting this to zero gives

\mathbb E[\|z\|^2G\mid s] -b^*(s)\mathbb E[\|z\|^2\mid s] =0,

which yields the stated result.

If the squared score norm is approximately independent of return, then

b^*(s)\approx\mathbb E[G\mid s]=V^\pi(s).

This explains why the value function works well in practice.

Advantage normalization is another common heuristic:

\widetilde A_t =\frac{\widehat A_t-\mu_A}{\sigma_A+\varepsilon}.

It improves numerical conditioning and reduces sensitivity to reward scale, but it is a batch-level optimization heuristic rather than a state-dependent baseline derived from the score identity.

Policy Distributions

The score function determines how a weighted sample changes the policy. Its form depends on the action distribution.

Softmax policies

Let a linear preference be

h_\theta(s,a)=\phi(s,a)^\top\theta.

The softmax policy is

\pi_\theta(a\mid s) =\frac{e^{\phi(s,a)^\top\theta}} {\sum_{a'}e^{\phi(s,a')^\top\theta}}.

Its score is

\boxed{ \nabla_\theta\log\pi_\theta(a\mid s) =\phi(s,a) -\mathbb E_{a'\sim\pi_\theta(\cdot\mid s)}[\phi(s,a')] }.

See derivation

Taking logs,

\log\pi_\theta(a\mid s) =\phi(s,a)^\top\theta -\log\sum_{a'}e^{\phi(s,a')^\top\theta}.

The log-partition derivative is

\begin{aligned} \nabla_\theta \log\sum_{a'}e^{\phi(s,a')^\top\theta} &=\frac{ \sum_{a'}e^{\phi(s,a')^\top\theta}\phi(s,a') }{ \sum_{a'}e^{\phi(s,a')^\top\theta} }\\ &=\sum_{a'}\pi_\theta(a'\mid s)\phi(s,a'). \end{aligned}

Subtracting this from $\phi(s,a)$ gives the score.

For neural-network logits $z_j$ ,

\frac{\partial\log\pi(a\mid s)}{\partial z_j} =\mathbb 1[j=a]-\pi(j\mid s).

A positive advantage pushes the sampled action's logit upward and competing logits downward. A negative advantage reverses the direction.

Gaussian policies

For a scalar action with fixed variance,

\pi_\theta(a\mid s) =\mathcal N(a;\mu_\theta(s),\sigma^2).

The score is

\boxed{ \nabla_\theta\log\pi_\theta(a\mid s) =\frac{a-\mu_\theta(s)}{\sigma^2} \nabla_\theta\mu_\theta(s) }.

See derivation

The log-density is

\log\pi_\theta(a\mid s) =-\frac12\log(2\pi\sigma^2) -\frac{(a-\mu_\theta(s))^2}{2\sigma^2}.

Only the quadratic term depends on $\theta$ , so

\begin{aligned} \nabla_\theta\log\pi_\theta(a\mid s) &=-\frac1{2\sigma^2} 2(a-\mu_\theta(s))(-\nabla_\theta\mu_\theta(s))\\ &=\frac{a-\mu_\theta(s)}{\sigma^2} \nabla_\theta\mu_\theta(s). \end{aligned}

For a multivariate Gaussian with fixed covariance $\Sigma$ ,

\nabla_\theta\log\pi_\theta(a\mid s) =\left(\nabla_\theta\mu_\theta(s)\right)^\top \Sigma^{-1}(a-\mu_\theta(s)).

A diagonal Gaussian often learns both $\mu$ and $\ell=\log\sigma$ . For one dimension,

\frac{\partial\log\pi}{\partial\mu} =\frac{a-\mu}{\sigma^2},

and

\boxed{ \frac{\partial\log\pi}{\partial\ell} =\frac{(a-\mu)^2}{\sigma^2}-1 }.

An unexpectedly distant, positively weighted action can increase the variance; an action close to the mean can decrease it. The policy learns both where to act and how much randomness to retain.

The Surrogate Loss

Automatic-differentiation libraries minimize scalar losses, while the desired update is gradient ascent:

\widehat g =\frac1M\sum_{j=1}^{M} \widehat A_j \nabla_\theta\log\pi_\theta(a_j\mid s_j).

We define the actor loss

\boxed{ \mathcal L_{\mathrm{actor}}(\theta) =-\frac1M\sum_{j=1}^{M} \operatorname{stopgrad}(\widehat A_j) \log\pi_\theta(a_j\mid s_j) }.

Then

-\nabla_\theta\mathcal L_{\mathrm{actor}} =\widehat g.

This is a surrogate loss. Its numerical value is not the negative expected return; its gradient is constructed to equal the sampled policy-gradient estimator.

Returns and advantages must be treated as fixed weights during the actor update:

\nabla_\theta\operatorname{stopgrad}(\widehat A_t)=0.

Otherwise gradients can flow through the critic or return computation in ways not justified by the likelihood-ratio derivation.

For variable-length episodes, flatten the valid time steps or use a mask $m_t^{(i)}$ :

\mathcal L_{\mathrm{actor}} =- \frac{ \sum_{i,t}m_t^{(i)} \log\pi_\theta(a_t^{(i)}\mid s_t^{(i)}) \widehat A_t^{(i)} }{ \sum_{i,t}m_t^{(i)} }.

\textbf{Vanilla Policy-Gradient Update}

\begin{aligned} &\text{Roll out }\pi_\theta\text{ to collect an on-policy batch}.\\[0.25em] &\text{Compute discounted reward-to-go }G_t\text{ backward}.\\[0.25em] &\phi\leftarrow\arg\min_\phi\sum_t \left(V_\phi(s_t)-G_t\right)^2.\\[0.25em] &\widehat A_t\leftarrow \operatorname{stopgrad}\left(G_t-V_\phi(s_t)\right).\\[0.25em] &\text{Optionally normalize }\widehat A_t\text{ within the batch}.\\[0.25em] &\mathcal L_{\mathrm{actor}} \leftarrow-\frac1M\sum_t \log\pi_\theta(a_t\mid s_t)\widehat A_t.\\[0.25em] &\theta\leftarrow\theta-\alpha\nabla_\theta \mathcal L_{\mathrm{actor}}. \end{aligned}

Why Policy Gradients Are Difficult

REINFORCE is mathematically clean, but its estimator is noisy. Variance comes from stochastic actions, stochastic dynamics, long horizons, delayed rewards, and using one sampled future to estimate many possible futures.

On-policy sample cost. The gradient is an expectation under the current policy:

\nabla_\theta J(\theta) =\mathbb E_{\tau\sim p_\theta}[\cdots].

After a substantial update, trajectories from $p_{\theta_{\mathrm{old}}}$ are no longer distributed according to $p_{\theta_{\mathrm{new}}}$ . Reusing them naively biases the estimate.

Importance sampling can correct the action-distribution mismatch with

\rho_t(\theta) =\frac{\pi_\theta(a_t\mid s_t)} {\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)}.

But products or large values of these ratios create severe variance. This motivates algorithms that reuse recent data while explicitly limiting policy change.

Sensitivity to step size. A vanilla update is

\theta_{k+1}=\theta_k+\alpha_k\widehat g_k.

If $\alpha_k$ is too small, learning is slow. If it is too large, the policy can collapse after one batch.

Euclidean parameter distance is not a reliable measure of policy change. A small parameter update may drastically change action probabilities in one part of the state space. Adaptive optimizers and normalized advantages help numerical conditioning, but they do not guarantee a safe change in the represented policy.

The Bridge to Actor-Critic

The Monte Carlo advantage

\widehat A_t^{\mathrm{MC}} =G_t-V_\phi(s_t)

still depends on one complete sampled future. A critic can bootstrap from the next state:

\boxed{ \widehat A_t^{\mathrm{TD}} =r_t+\gamma V_\phi(s_{t+1})-V_\phi(s_t) }.

This follows from

Q^\pi(s_t,a_t) =\mathbb E \left[r_t+\gamma V^\pi(s_{t+1})\mid s_t,a_t\right]

and

A^\pi(s_t,a_t)=Q^\pi(s_t,a_t)-V^\pi(s_t).

The actor update becomes

\boxed{ \widehat{\nabla_\theta J} =\frac1M\sum_t \nabla_\theta\log\pi_\theta(a_t\mid s_t) \left[r_t+\gamma V_\phi(s_{t+1})-V_\phi(s_t)\right] }.

The critic averages over possible futures rather than using one full sampled return, so variance is usually much lower. An inaccurate critic introduces bias.

A one-step critic minimizes

\boxed{ \mathcal L_V^{\mathrm{TD}}(\phi) =\mathbb E \left[ \left( V_\phi(s_t) -\operatorname{stopgrad} \left[r_t+\gamma V_\phi(s_{t+1})\right] \right)^2 \right] }.

The target is detached during critic optimization, just as in value-based reinforcement learning.

Method	Actor weight	Bias	Variance
Full-return REINFORCE	$R(\tau)$	Unbiased	Very high
Reward-to-go REINFORCE	$G_t$	Unbiased	High
REINFORCE with baseline	$G_t-V_\phi(s_t)$	Actor remains unbiased	Lower
One-step actor-critic	$r_t+\gamma V_\phi(s_{t+1})-V_\phi(s_t)$	Biased with an inaccurate critic	Usually much lower

Actor-critic methods deliberately trade some estimator bias for lower variance. Proximal and trust-region methods address the second major weakness by controlling how far the policy moves during each update.

Practical Diagnosis

Symptom	Likely cause	Useful response
Returns vary wildly between batches	High Monte Carlo variance	Use reward-to-go, a value baseline, and larger batches
Learning reinforces obviously bad actions	Noisy long-horizon credit	Improve the baseline or use actor-critic estimates
Policy collapses after one update	Policy step is too large	Lower the learning rate or constrain policy change
Learning stalls with tiny gradients	Reward scale or advantage scale is poor	Normalize advantages and inspect reward scaling
Old batches hurt performance	On-policy distribution mismatch	Discard stale data or use explicit importance corrections
Continuous-action variance explodes	Unstable standard deviation	Parameterize log standard deviation and monitor entropy
Actor changes while critic loss backpropagates	Missing gradient separation	Detach advantages and use separate actor/critic losses

The entire derivation can be remembered as one chain:

\begin{aligned} J(\theta) &=\mathbb E_{\tau\sim p_\theta}[R(\tau)]\\ \nabla_\theta J(\theta) &=\mathbb E \left[R(\tau)\nabla_\theta\log p_\theta(\tau)\right]\\ &=\mathbb E \left[ \sum_t \nabla_\theta\log\pi_\theta(a_t\mid s_t)G_t \right]\\ &=\mathbb E \left[ \sum_t \nabla_\theta\log\pi_\theta(a_t\mid s_t) \left(G_t-b(s_t)\right) \right]. \end{aligned}

The likelihood-ratio identity makes the environment differentiable in expectation, temporal causality removes past rewards, and a state-dependent baseline removes predictable variation. Replacing the remaining Monte Carlo return with a learned critic produces actor-critic methods.