Surya Dantuluri blog2

Policy Gradients

Explanation and derivation of policy gradients, from the objective to a sample estimate.

· 3 min read · Draft

Material adapted from Sergey Levine’s CS294-112 Fall 2018 and Josh Achiam’s Spinning Up With RL

Correction — October 2026: This 2019 draft is preserved here with corrected notation and a completed derivation below. The draft label remains; no original implementation was recovered.

What is a Policy Gradient and how it is useful?

The goal in reinforcement learning is to find a policy with high expected return. Policy-gradient methods adjust a policy’s action probabilities using rewards from sampled trajectories. Following the gradient can improve the policy, but noisy estimates and local optima mean that reaching an optimal policy is not guaranteed.

Derivation

Consider a stochastic policy πθ\pi_{\theta} and a trajectory τ=(s0,a0,…,sT−1,aT−1,sT)\tau = (s_0, a_0, \ldots, s_{T-1}, a_{T-1}, s_T) with TT actions. We use the finite-horizon, undiscounted return:

R(τ)=∑t=0T−1r(st,at)R(\tau) = \sum_{t=0}^{T-1} r(s_t,a_t)

The goal is to choose the policy parameters that maximize expected return:

θ⋆∈argmax⁡θ J(θ)\theta^\star \in \underset{\theta}{\operatorname{argmax}}\, \color{red}J(\theta)

Here J(θ)\color{red}J(\theta) is the expectation of the total reward. Writing pθ(τ)p_{\theta}(\tau) for the trajectory distribution gives:

J(θ)=Eτ∼pθ[R(τ)]=∫pθ(τ)R(τ) dτ\color{red}J(\theta) \color{black}= \color{orange}\mathbb{E}_{\tau\sim p_{\theta}}[R(\tau)] \color{black}= \int \color{blue}p_{\theta}(\tau)\color{green}R(\tau)\color{black}\,d\tau

Assume the initial-state distribution, environment transitions, and reward function do not depend on θ\theta. Also assume the policy is differentiable and we can move the derivative inside the integral. Then:

∇θJ(θ)=∇θ∫pθ(τ)R(τ) dτ=∫∇θpθ(τ)R(τ) dτ\begin{aligned} \nabla_{\theta}\color{red}J(\theta) \color{black}&= \nabla_{\theta}\int \color{blue}p_{\theta}(\tau)\color{green}R(\tau)\color{black}\,d\tau \\ &= \int \color{black}\nabla_{\theta}\color{blue}p_{\theta}(\tau)\color{green}R(\tau)\color{black}\,d\tau \end{aligned}

For trajectories with positive probability, the log-derivative identity is:

∇θpθ(τ)=pθ(τ)∇θlog⁡pθ(τ)\nabla_{\theta}\color{blue}p_{\theta}(\tau) \color{black}= \color{blue}p_{\theta}(\tau)\color{black}\nabla_{\theta}\color{yellow}\log p_{\theta}(\tau)

Substituting it into the integral:

∇θJ(θ)=∫pθ(τ)∇θlog⁡pθ(τ)R(τ) dτ\nabla_{\theta}J(\theta) = \int \color{blue}p_{\theta}(\tau)\color{black}\nabla_{\theta}\color{yellow}\log p_{\theta}(\tau)\color{green}R(\tau)\color{black}\,d\tau

Going back to expectation form:

∇θJ(θ)=Eτ∼pθ ⁣[∇θlog⁡pθ(τ)R(τ)]\nabla_{\theta}J(\theta) = \mathbb{E}_{\tau\sim p_{\theta}}\!\left[\nabla_{\theta}\color{yellow}\log p_{\theta}(\tau)\color{green}R(\tau)\color{black}\right]

The trajectory probability factors into the initial-state distribution ρ0\rho_0, the environment transitions PP, and the policy:

pθ(τ)=ρ0(s0)∏t=0T−1P(st+1∣st,at)πθ(at∣st)p_{\theta}(\tau) = \rho_0(s_0)\prod_{t=0}^{T-1} P(s_{t+1}\mid s_t,a_t)\pi_{\theta}(a_t\mid s_t)

Taking the log turns the product into a sum:

log⁡pθ(τ)=log⁡ρ0(s0)+∑t=0T−1 ⁣[log⁡P(st+1∣st,at)+log⁡πθ(at∣st)]\begin{aligned} \log p_{\theta}(\tau) &= \log \rho_0(s_0) \\ &\quad+ \sum_{t=0}^{T-1}\!\left[\log P(s_{t+1}\mid s_t,a_t) + \log \pi_{\theta}(a_t\mid s_t)\right] \end{aligned}

For a fixed trajectory, the initial-state and transition terms are independent of θ\theta. Their gradients vanish. The policy term depends on θ\theta and remains:

∇θlog⁡pθ(τ)=∑t=0T−1∇θlog⁡πθ(at∣st)\nabla_{\theta}\log p_{\theta}(\tau) = \sum_{t=0}^{T-1}\nabla_{\theta}\log \pi_{\theta}(a_t\mid s_t)

Putting this into the expectation completes the basic policy-gradient derivation:

∇θJ(θ)=Eτ∼pθ ⁣[∑t=0T−1∇θlog⁡πθ(at∣st)R(τ)]\nabla_{\theta}J(\theta) = \mathbb{E}_{\tau\sim p_{\theta}}\!\left[\color{yellow}\sum_{t=0}^{T-1}\nabla_{\theta}\log \pi_{\theta}(a_t\mid s_t)\color{green}R(\tau)\color{black}\right]

Given NN independent trajectories sampled from the current policy, a sample estimate is:

g^=1N∑i=1N[∑t=0T−1∇θlog⁡πθ(at(i)∣st(i))]R(τ(i))\widehat{g} = \frac{1}{N}\sum_{i=1}^{N}\left[\sum_{t=0}^{T-1}\nabla_{\theta}\log \pi_{\theta}(a_t^{(i)}\mid s_t^{(i)})\right]R(\tau^{(i)})

We can use this estimate for a gradient-ascent step, θ←θ+αg^\theta \leftarrow \theta + \alpha\widehat{g}, with step size α>0\alpha > 0. A finite sample is noisy; an individual update is not guaranteed to improve return.

The corrected derivation follows Spinning Up’s policy-gradient chapter and Sergey Levine’s Fall 2018 lecture 5. The original draft’s promised implementation was not included in the recovered article.