Fall 2026
  • Discord
  • Gradescope
  • Syllabus
  • Spring 2026

On this page

  • Differentiating an Expected Return
  • Trajectory Probabilities
  • REINFORCE
  • Reward to Go

Policy Gradients

Suppose we can run an agent’s policy and observe its rewards, but cannot differentiate through the environment. We can still differentiate the probabilities the policy assigns to its actions. The log-derivative identity uses these probabilities to estimate a gradient from sampled runs, giving us the REINFORCE algorithm.

How can sampled rewards tell us which way to update a policy?


Differentiating an Expected Return

Let \(\mathcal A\) be a fixed finite set of available actions. A policy network with parameters \(\boldsymbol\theta\in\mathbb R^p\) assigns a logit \(z_a(s)\) to each action \(a\in\mathcal A\) at state \(s\), then maps these logits to probabilities \[ \pi_{\boldsymbol\theta}(a\mid s) =\frac{e^{z_a(s)}}{\sum_{j\in\mathcal A}e^{z_j(s)}}. \] The agent samples \(a_t\sim\pi_{\boldsymbol\theta}(\cdot\mid s_t)\). Because this action is discrete, there is no ordinary derivative of the sampled action with respect to \(\boldsymbol\theta\).

For this derivation, fix an integer horizon \(T\geq1\) and consider trajectories \(\tau=(s_0,a_0,r_0,\ldots,s_{T-1},a_{T-1},r_{T-1},s_T)\). If an episode terminates sooner, extend it to \(T\) decisions with an absorbing terminal state and zero rewards after termination. For \(0\leq\gamma\leq1\), define the trajectory return and objective by \[ G(\tau)=\sum_{t=0}^{T-1}\gamma^t r_t, \qquad J(\boldsymbol\theta) =\mathbb E_{\tau\sim p_{\boldsymbol\theta}}[G(\tau)]. \] The parameters affect \(J\) by changing the distribution \(p_{\boldsymbol\theta}\) over trajectories. We therefore differentiate trajectory probabilities rather than environment transitions or measured rewards.

Claim: Let \(p_{\boldsymbol\theta}\) be a differentiable probability mass function or density whose support does not depend on \(\boldsymbol\theta\), and let \(X\) have this distribution. If \(f(x)\) has no additional dependence on \(\boldsymbol\theta\) and differentiation may pass through the sum or integral, then \[ \nabla_{\boldsymbol\theta}\mathbb E[f(X)] =\mathbb E\!\left[ f(X)\nabla_{\boldsymbol\theta}\log p_{\boldsymbol\theta}(X) \right]. \]

Proof

For a discrete random variable, let \(\mathcal X=\{x:p_{\boldsymbol\theta}(x)>0\}\) denote its fixed support. Then \[ \begin{aligned} \nabla_{\boldsymbol\theta}\mathbb E[f(X)] &=\nabla_{\boldsymbol\theta}\sum_{x\in\mathcal X} p_{\boldsymbol\theta}(x)f(x)\\ &=\sum_{x\in\mathcal X} f(x)\nabla_{\boldsymbol\theta}p_{\boldsymbol\theta}(x)\\ &=\sum_{x\in\mathcal X} p_{\boldsymbol\theta}(x)f(x) \frac{\nabla_{\boldsymbol\theta}p_{\boldsymbol\theta}(x)} {p_{\boldsymbol\theta}(x)}\\ &=\mathbb E\!\left[ f(X)\nabla_{\boldsymbol\theta}\log p_{\boldsymbol\theta}(X) \right]. \end{aligned} \] The density case replaces the sum by an integral.

The value \(f(x)\) can come from a black box. Only the log-probability of the sampled outcome must be differentiable.

Trajectory Probabilities

Write \(P(s_0)\) for the initial-state law and \(P(r_t,s_{t+1}\mid s_t,a_t)\) for the environment’s joint reward and next-state law, using densities when the variables are continuous. A trajectory’s probability or density factors as \[ p_{\boldsymbol\theta}(\tau) =P(s_0)\prod_{t=0}^{T-1} \pi_{\boldsymbol\theta}(a_t\mid s_t) P(r_t,s_{t+1}\mid s_t,a_t). \]

Initial State
fixed distribution
×
Policy Actions
depend on θ
×
Environment Transitions
fixed conditional laws

Claim: If the initial-state and environment laws do not depend on \(\boldsymbol\theta\), and the log-derivative identity’s assumptions hold for the trajectory distribution, then \[ \nabla_{\boldsymbol\theta}J =\mathbb E\!\left[ G(\tau)\sum_{t=0}^{T-1}\nabla_{\boldsymbol\theta} \log\pi_{\boldsymbol\theta}(a_t\mid s_t) \right]. \]

Proof

Apply the log-derivative identity to the trajectory: \[ \nabla_{\boldsymbol\theta}J =\mathbb E\!\left[ G(\tau)\nabla_{\boldsymbol\theta}\log p_{\boldsymbol\theta}(\tau) \right]. \] Taking the logarithm of the factorization gives \[ \begin{aligned} \log p_{\boldsymbol\theta}(\tau) &=\log P(s_0) +\sum_{t=0}^{T-1}\log\pi_{\boldsymbol\theta}(a_t\mid s_t)\\ &\quad+\sum_{t=0}^{T-1}\log P(r_t,s_{t+1}\mid s_t,a_t). \end{aligned} \] The first and third terms contain no policy parameters, so \[ \nabla_{\boldsymbol\theta}\log p_{\boldsymbol\theta}(\tau) =\sum_{t=0}^{T-1}\nabla_{\boldsymbol\theta} \log\pi_{\boldsymbol\theta}(a_t\mid s_t). \] Substitution proves the claim.

The environment supplies sampled states and rewards. The policy supplies the derivatives used in the gradient estimate.

REINFORCE

For \(N\geq1\) independent trajectories \(\tau_1,\ldots,\tau_N\) sampled from the current policy, write \(s_{i,t}\) and \(a_{i,t}\) for the state and action at time \(t\in\{0,\ldots,T-1\}\) in trajectory \(i\in\{1,\ldots,N\}\). The policy-gradient identity gives the estimator \[ \widehat{\nabla J} =\frac1N\sum_{i=1}^N G(\tau_i) \sum_{t=0}^{T-1}\nabla_{\boldsymbol\theta} \log\pi_{\boldsymbol\theta}(a_{i,t}\mid s_{i,t}). \] Linearity of expectation shows that \[ \mathbb E[\widehat{\nabla J}] =\nabla_{\boldsymbol\theta}J. \] Gradient ascent then updates \[ \boldsymbol\theta \leftarrow \boldsymbol\theta+\alpha\widehat{\nabla J}. \] This sample, evaluate, and update procedure is REINFORCE.

For a softmax policy, the probability vector \(\mathbf p_t\) and one-hot sampled action \(\mathbf y_t\) each have \(|\mathcal A|\) coordinates, one per action, matching the logit vector \(\mathbf z\). Then \[ \nabla_{\mathbf z}\log\pi_{\boldsymbol\theta}(a_t\mid s_t) =\mathbf y_t-\mathbf p_t. \] This is the negative of the cross-entropy gradient from logistic regression. For a direct step in these logits, multiplying by a positive return increases the sampled action’s probability. A negative return decreases it. A policy network passes this score gradient through its shared weights, so contributions from different decisions can compete. The sign is the opposite of the loss gradient because we maximize expected return.

The explorer takes one step directly in a fixed vector of action logits.

A softmax action score is one-hot action minus probability. A positive return increases the selected action’s probability in a direct logit ascent step. A negative return reverses that step.

Return scales the correction to the sampled action. No derivative of the environment is used.

Reward to Go

The total return weights each action by rewards that occurred both before and after it. Earlier rewards cannot have been caused by the current action and can be removed without changing the expectation.

Claim: For the fixed horizon \(T\), define the return from time \(t\in\{0,\ldots,T\}\) by \[ G_t=\sum_{k=0}^{T-t-1}\gamma^k r_{t+k}. \] An empty sum is zero, so \(G_T=0\). For \(J=\mathbb E[G_0]\), the gradient is \[ \nabla_{\boldsymbol\theta}J =\mathbb E\!\left[ \sum_{t=0}^{T-1}\gamma^tG_t \nabla_{\boldsymbol\theta}\log \pi_{\boldsymbol\theta}(a_t\mid s_t) \right]. \]

Proof

For \(t=0,\ldots,T-1\), write \[ \mathbf u_t =\nabla_{\boldsymbol\theta}\log \pi_{\boldsymbol\theta}(a_t\mid s_t). \] The history \(\mathcal H_t=(s_0,a_0,r_0,\ldots,s_{t-1},a_{t-1},r_{t-1},s_t)\) contains everything observed before choosing \(a_t\), with \(\mathcal H_0=(s_0)\). Conditional on this history, \[ \begin{aligned} \mathbb E[\mathbf u_t\mid\mathcal H_t] &=\sum_{a\in\mathcal A}\pi_{\boldsymbol\theta}(a\mid s_t) \nabla_{\boldsymbol\theta}\log \pi_{\boldsymbol\theta}(a\mid s_t)\\ &=\sum_{a\in\mathcal A}\nabla_{\boldsymbol\theta} \pi_{\boldsymbol\theta}(a\mid s_t)\\ &=\nabla_{\boldsymbol\theta}1 =\mathbf0. \end{aligned} \] The earlier discounted reward \[ B_t=\sum_{k=0}^{t-1}\gamma^k r_k \] is fixed given \(\mathcal H_t\) and equals zero at \(t=0\) by the empty-sum convention, so \[ \mathbb E[B_t\mathbf u_t] =\mathbb E\!\left[ B_t\mathbb E[\mathbf u_t\mid\mathcal H_t] \right] =\mathbf0. \] Because \(G_0=B_t+\gamma^tG_t\), \[ \mathbb E[G_0\mathbf u_t] =\mathbb E[\gamma^tG_t\mathbf u_t]. \] Replacing each term in the trajectory policy gradient proves the result.

Variance.

REINFORCE is unbiased, but trajectories vary because of sampled actions, environment transitions, and rewards. The estimator multiplies a random return by a random score vector, so its variance can be large.

Adding a constant to every return leaves the exact gradient unchanged because the expected policy score is zero. It can nevertheless increase the variance of sampled gradients. Multiplying rewards by a positive constant preserves the best policy but scales both the gradient and its noise.

The next reading uses baselines to remove an action-independent return level while preserving the policy gradient.