From Imitation Learning to Online RL
In the previous lecture we saw how imitation learning can produce capable policies by mimicking expert demonstrations. Given a dataset $\mathcal{D}$ of state-action pairs from a demonstrator, we minimize a behavioral cloning objective:
$$\min_\theta \; -\E_{(\mathbf{s},\mathbf{a})\sim\mathcal{D}}\!\bigl[\log \pi_\theta(\mathbf{a}\mid\mathbf{s})\bigr]$$This is simple and scalable, but it has a fundamental limitation: the learned policy can never exceed the demonstrator's performance. It learns what to do from examples, but has no mechanism for improvement through practice.
To go beyond the demonstrator, we need a different paradigm. Rather than copying behavior, we want to discover good behavior through interaction with the environment—collecting our own experience and using reward signals to improve. This is online reinforcement learning.
- Offline: Learning uses only a fixed, pre-collected dataset. No new data is gathered from the learned policy.
- Online: Learning uses new data actively collected by the learned policy as it interacts with the environment.
The RL Objective
Recall the core quantities of reinforcement learning. A state $\mathbf{s}_t$ describes the world at time $t$. An action $\mathbf{a}_t$ is the decision taken at time $t$. A trajectory $\tau$ is a sequence of states and actions: $\tau = (\mathbf{s}_1, \mathbf{a}_1, \mathbf{s}_2, \mathbf{a}_2, \ldots, \mathbf{s}_T, \mathbf{a}_T)$. The reward function $r(\mathbf{s}, \mathbf{a})$ measures how good a state-action pair is. And the policy $\pi(\mathbf{a}\mid\mathbf{s})$ is the behavior we want to learn.
Our goal is to find policy parameters $\theta$ that maximize the expected sum of rewards:
$$\max_\theta \; \E_{\tau \sim p_\theta(\tau)}\!\left[\sum_{t=1}^{T} r(\mathbf{s}_t, \mathbf{a}_t)\right]$$where the trajectory distribution factors as:
$$p(\mathbf{s}_1, \mathbf{a}_1, \ldots, \mathbf{s}_T, \mathbf{a}_T) = p(\mathbf{s}_1)\prod_{t=1}^{T} \pi_\theta(\mathbf{a}_t \mid \mathbf{s}_t)\, p(\mathbf{s}_{t+1} \mid \mathbf{s}_t, \mathbf{a}_t)$$Notice that both the policy $\pi_\theta$ and the environment dynamics $p(\mathbf{s}_{t+1}\mid\mathbf{s}_t,\mathbf{a}_t)$ contribute to the trajectory distribution. We can control only the policy; the dynamics are given by the environment.
The Online RL Loop
Online RL follows a simple iterative structure. First, we initialize the policy (randomly, from imitation learning, or using heuristics). Then we repeat:
- Collect data: Run the current policy $\pi_\theta$ in the environment to gather a batch of trajectories.
- Improve the policy: Use the collected trajectories to update $\theta$ in a direction that increases expected reward.
This lecture focuses on the most direct approach to step 2: computing the gradient of the RL objective with respect to $\theta$ and taking gradient ascent steps. This family of methods is called policy gradients—our first online RL algorithm.
Evaluating and Differentiating the RL Objective
Define the RL objective as:
$$J(\theta) = \E_{\tau \sim p_\theta(\tau)}\!\left[\sum_{t=1}^{T} r(\mathbf{s}_t, \mathbf{a}_t)\right]$$We can evaluate $J(\theta)$ empirically by sampling $N$ trajectories $\{\tau^i\}_{i=1}^N$ from the current policy $\pi_\theta$ and averaging:
$$J(\theta) \approx \frac{1}{N}\sum_{i=1}^{N}\sum_{t=1}^{T} r(\mathbf{s}_{i,t}, \mathbf{a}_{i,t})$$But to improve the policy, we need the gradient $\nabla_\theta J(\theta)$. The challenge is that the expectation is taken over $p_\theta(\tau)$, which itself depends on $\theta$. We cannot simply differentiate through the sampling process. Instead, we use a beautiful algebraic identity.
The Log-Gradient Trick
The key insight comes from rewriting the gradient of a probability in terms of the gradient of its logarithm. For any distribution $p_\theta(\tau)$:
$$p_\theta(\tau)\,\nabla_\theta \log p_\theta(\tau) = p_\theta(\tau)\,\frac{\nabla_\theta p_\theta(\tau)}{p_\theta(\tau)} = \nabla_\theta p_\theta(\tau)$$This identity is sometimes called the log-derivative trick or REINFORCE trick. It allows us to express the gradient of an expectation under a parameterized distribution as an expectation of the gradient of the log-probability, which we can estimate from samples.
Deriving the Policy Gradient Theorem
We now carry out the full derivation, starting from the RL objective written as an integral over trajectories and arriving at a form that can be estimated purely from samples of the policy.
Proof of the Policy Gradient Theorem
Step 1: Write the objective as an integral. Let $r(\tau) = \sum_{t=1}^T r(\mathbf{s}_t, \mathbf{a}_t)$ denote the total reward of trajectory $\tau$. Then:
$$J(\theta) = \E_{\tau \sim p_\theta(\tau)}[r(\tau)] = \int p_\theta(\tau)\, r(\tau)\, d\tau$$Step 2: Differentiate under the integral.
$$\nabla_\theta J(\theta) = \int \nabla_\theta p_\theta(\tau)\, r(\tau)\, d\tau$$Step 3: Apply the log-derivative trick. Since $\nabla_\theta p_\theta(\tau) = p_\theta(\tau)\,\nabla_\theta \log p_\theta(\tau)$:
$$\nabla_\theta J(\theta) = \int p_\theta(\tau)\,\nabla_\theta \log p_\theta(\tau)\, r(\tau)\, d\tau = \E_{\tau \sim p_\theta(\tau)}\!\bigl[\nabla_\theta \log p_\theta(\tau)\, r(\tau)\bigr]$$Step 4: Expand $\log p_\theta(\tau)$. The trajectory probability factors as:
$$p_\theta(\tau) = p(\mathbf{s}_1)\prod_{t=1}^{T} \pi_\theta(\mathbf{a}_t \mid \mathbf{s}_t)\, p(\mathbf{s}_{t+1} \mid \mathbf{s}_t, \mathbf{a}_t)$$Taking the logarithm:
$$\log p_\theta(\tau) = \log p(\mathbf{s}_1) + \sum_{t=1}^{T}\bigl[\log \pi_\theta(\mathbf{a}_t \mid \mathbf{s}_t) + \log p(\mathbf{s}_{t+1} \mid \mathbf{s}_t, \mathbf{a}_t)\bigr]$$Step 5: Differentiate with respect to $\theta$. The initial state distribution $p(\mathbf{s}_1)$ and the transition dynamics $p(\mathbf{s}_{t+1}\mid\mathbf{s}_t,\mathbf{a}_t)$ do not depend on $\theta$, so their gradients vanish:
$$\nabla_\theta \log p_\theta(\tau) = \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(\mathbf{a}_t \mid \mathbf{s}_t)$$Step 6: Substitute back.
$$\nabla_\theta J(\theta) = \E_{\tau \sim p_\theta(\tau)}\!\left[\left(\sum_{t=1}^{T}\nabla_\theta \log \pi_\theta(\mathbf{a}_t \mid \mathbf{s}_t)\right)\left(\sum_{t=1}^{T} r(\mathbf{s}_t, \mathbf{a}_t)\right)\right] \qquad \blacksquare$$A crucial feature of this result is that the gradient does not require knowledge of the environment dynamics $p(\mathbf{s}_{t+1}\mid\mathbf{s}_t,\mathbf{a}_t)$. It depends only on the policy $\pi_\theta$ (which we control) and the rewards $r$ (which we observe). This makes the method model-free.
The REINFORCE Algorithm
The policy gradient theorem gives us a gradient in expectation. To estimate it in practice, we replace the expectation with a sample average over $N$ trajectories collected from the current policy:
$$\nabla_\theta J(\theta) \approx \frac{1}{N}\sum_{i=1}^{N}\left(\sum_{t=1}^{T}\nabla_\theta \log \pi_\theta(\mathbf{a}_{i,t} \mid \mathbf{s}_{i,t})\right)\left(\sum_{t=1}^{T} r(\mathbf{s}_{i,t}, \mathbf{a}_{i,t})\right)$$Combined with gradient ascent on $\theta$, this gives the REINFORCE algorithm (Williams, 1992), also called the vanilla policy gradient.
- Sample a batch of $N$ trajectories $\{\tau^i\}_{i=1}^N$ by running $\pi_\theta(\mathbf{a}_t \mid \mathbf{s}_t)$ in the environment.
- Estimate the gradient: $$\nabla_\theta J(\theta) \approx \frac{1}{N}\sum_{i=1}^{N}\left(\sum_{t=1}^{T}\nabla_\theta \log \pi_\theta(\mathbf{a}_{i,t} \mid \mathbf{s}_{i,t})\right)\left(\sum_{t=1}^{T} r(\mathbf{s}_{i,t}, \mathbf{a}_{i,t})\right)$$
- Update parameters: $\theta \leftarrow \theta + \alpha\,\nabla_\theta J(\theta)$
- Repeat from step 1.
Intuition: Reward-Weighted Imitation
The policy gradient has a striking structural similarity to the behavioral cloning gradient. Recall that the imitation learning gradient is:
$$\nabla_\theta J_{\text{BC}}(\theta) \approx \frac{1}{N}\sum_{i=1}^{N}\left(\sum_{t=1}^{T}\nabla_\theta \log \pi_\theta(\mathbf{a}_{i,t}\mid\mathbf{s}_{i,t})\right)$$The policy gradient is the same expression, but with each trajectory's log-likelihood gradient weighted by the total reward of that trajectory. The core intuition is simple:
- Increase the likelihood of actions from trajectories with high reward.
- Decrease the likelihood of actions from trajectories with negative reward.
In other words: do more of the good stuff, less of the bad stuff. Policy gradients are a formal mathematical realization of trial-and-error learning.
- $\tau^1$: falls backwards (negative reward)
- $\tau^2$: falls forwards (slightly positive reward)
- $\tau^3$: manages to stand still (zero reward)
- $\tau^4$: one small step forward, then falls backward (mixed reward)
- $\tau^5$: one large step backward, then small step forward (negative overall)
Variance Reduction: Causality and Baselines
The REINFORCE gradient estimator is unbiased but suffers from high variance. With a finite batch of trajectories, the gradient estimates can be extremely noisy, leading to slow and unstable learning. Two key techniques reduce this variance without introducing bias: the causality trick (reward-to-go) and baselines.
The Causality Trick (Reward-to-Go)
In the vanilla policy gradient, the log-probability of each action $\mathbf{a}_t$ is multiplied by the total trajectory reward $\sum_{t'=1}^T r(\mathbf{s}_{t'}, \mathbf{a}_{t'})$. But actions taken at time $t$ cannot influence rewards received at earlier times $t' < t$. Including past rewards in the weight only adds noise.
We can tighten the estimator by replacing the total reward with the reward-to-go—the sum of future rewards from time $t$ onward:
$$\nabla_\theta J(\theta) \approx \frac{1}{N}\sum_{i=1}^{N}\sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(\mathbf{a}_{i,t}\mid\mathbf{s}_{i,t})\left(\sum_{t'=t}^{T} r(\mathbf{s}_{i,t'}, \mathbf{a}_{i,t'})\right)$$This is still an unbiased estimator of the policy gradient (the extra past-reward terms had zero expected contribution), but it has strictly lower variance because each action is now weighted only by the consequences it can actually affect.
Baselines
The second source of variance comes from the absolute scale of the rewards. If all trajectories have positive reward (even the bad ones), the gradient will increase the likelihood of all actions—just more so for the good trajectories. This makes learning slow because the gradient signal does not clearly distinguish above-average from below-average behavior.
The fix is to subtract a baseline $b$ from the reward:
$$\nabla_\theta J(\theta) = \E_{\tau \sim p_\theta(\tau)}\!\left[\nabla_\theta \log p_\theta(\tau)\,\bigl(r(\tau) - b\bigr)\right]$$If $b$ is chosen to be the average reward, then above-average trajectories receive positive weight and below-average trajectories receive negative weight, creating a much clearer learning signal.
But is this still a valid gradient? We need to verify that subtracting a constant does not introduce bias.
Proof of unbiasedness
Using the log-derivative trick in reverse:
$$\E\!\bigl[\nabla_\theta \log p_\theta(\tau)\cdot b\bigr] = \int p_\theta(\tau)\,\nabla_\theta \log p_\theta(\tau)\, b\; d\tau = \int \nabla_\theta p_\theta(\tau)\, b\; d\tau$$ $$= b\,\nabla_\theta \int p_\theta(\tau)\, d\tau = b\,\nabla_\theta\, 1 = 0$$The key step is that $\int p_\theta(\tau)\,d\tau = 1$ for any $\theta$ (the trajectory distribution is always normalized), so the gradient of this constant is zero. $\blacksquare$
The simplest and most common baseline is the average reward across the current batch:
$$b = \frac{1}{N}\sum_{i=1}^{N} r(\tau^i)$$Combining the reward-to-go with a baseline, the improved policy gradient estimator becomes:
$$\nabla_\theta J(\theta) \approx \frac{1}{N}\sum_{i=1}^{N}\sum_{t=1}^{T}\nabla_\theta \log \pi_\theta(\mathbf{a}_{i,t}\mid\mathbf{s}_{i,t})\left(\left(\sum_{t'=t}^{T} r(\mathbf{s}_{i,t'}, \mathbf{a}_{i,t'})\right) - b\right)$$Practical Implementation
Implementing the policy gradient estimator naively—computing $\nabla_\theta \log \pi_\theta(\mathbf{a}_{i,t}\mid\mathbf{s}_{i,t})$ individually for every time step of every trajectory—would require $N \times T$ separate backward passes through the policy network, which is inefficient.
The Surrogate Objective
The standard solution is to construct a surrogate objective $\tilde{J}(\theta)$ whose gradient with respect to $\theta$ equals the policy gradient, and then use automatic differentiation to compute it in a single backward pass:
$$\tilde{J}(\theta) = \frac{1}{N}\sum_{i=1}^{N}\left(\sum_{t=1}^{T} \log \pi_\theta(\mathbf{a}_{i,t}\mid\mathbf{s}_{i,t})\right)\left(\left(\sum_{t'=t}^{T} r(\mathbf{s}_{i,t'}, \mathbf{a}_{i,t'})\right) - b\right)$$This is a weighted maximum likelihood objective: it looks like the behavioral cloning loss, but with each log-probability term weighted by the advantage of that trajectory relative to the baseline. When we call loss.backward() in a deep learning framework, we automatically obtain the policy gradient.
Policy Representations
The $\log \pi_\theta(\mathbf{a}\mid\mathbf{s})$ term takes different forms depending on the action space:
- Discrete actions: The policy network outputs logits over actions, and $\log \pi_\theta$ is the cross-entropy loss between the predicted distribution and the actions actually taken.
- Continuous actions: The policy network outputs the mean (and optionally variance) of a Gaussian distribution, and $\log \pi_\theta$ corresponds to a squared error term (from the Gaussian log-likelihood).
The On-Policy Limitation
The policy gradient as derived requires trajectories sampled from the current policy $\pi_\theta$. After each gradient update $\theta \leftarrow \theta + \alpha\,\nabla_\theta J(\theta)$, the parameters change, and the old trajectories are no longer distributed according to the new policy. This means we must discard all collected data and resample after every gradient step.
- On-policy: The policy update uses only data collected from the current policy. After each update, data must be recollected.
- Off-policy: The policy update can reuse data collected from other (e.g., previous) policies.
Vanilla policy gradient (REINFORCE) is an on-policy algorithm. This is its most significant practical limitation: collecting trajectories is often the most expensive part of the RL pipeline (especially in robotics or complex simulations), and on-policy methods waste this data by using it for only a single gradient step. This motivates the development of off-policy policy gradients.
Off-Policy Policy Gradients via Importance Sampling
Can we update our policy using data collected from a different policy? Importance sampling provides a principled way to do this.
Importance Sampling
Importance sampling is a general technique for estimating expectations under one distribution using samples from another. If we want to compute $\E_{x \sim p}[f(x)]$ but can only sample from a proposal distribution $q$, we write:
$$\E_{x \sim p}[f(x)] = \int p(x)\,f(x)\,dx = \int q(x)\,\frac{p(x)}{q(x)}\,f(x)\,dx = \E_{x \sim q}\!\left[\frac{p(x)}{q(x)}\,f(x)\right]$$The ratio $\frac{p(x)}{q(x)}$ is called the importance weight. This is an unbiased estimator as long as $q(x) > 0$ wherever $p(x) > 0$ (the support condition).
Importance Sampling for Trajectories
To apply importance sampling to the RL objective, suppose we have trajectories from an old policy $\bar{\pi}$ but want to evaluate the objective under our current policy $\pi_\theta$. The importance weight for a trajectory is:
$$\frac{p_\theta(\tau)}{\bar{p}(\tau)} = \frac{p(\mathbf{s}_1)\prod_{t=1}^T \pi_\theta(\mathbf{a}_t\mid\mathbf{s}_t)\,p(\mathbf{s}_{t+1}\mid\mathbf{s}_t,\mathbf{a}_t)}{p(\mathbf{s}_1)\prod_{t=1}^T \bar{\pi}(\mathbf{a}_t\mid\mathbf{s}_t)\,p(\mathbf{s}_{t+1}\mid\mathbf{s}_t,\mathbf{a}_t)} = \prod_{t=1}^{T}\frac{\pi_\theta(\mathbf{a}_t\mid\mathbf{s}_t)}{\bar{\pi}(\mathbf{a}_t\mid\mathbf{s}_t)}$$The initial state distribution and transition dynamics cancel out—we do not need a model of the environment. The importance weight is simply a product of action-probability ratios.
This lets us write the off-policy policy gradient. If we want to update parameters $\theta'$ using data from policy $\pi_\theta$:
$$\nabla_{\theta'} J(\theta') = \E_{\tau \sim p_\theta(\tau)}\!\left[\prod_{t=1}^{T}\frac{\pi_{\theta'}(\mathbf{a}_t\mid\mathbf{s}_t)}{\pi_\theta(\mathbf{a}_t\mid\mathbf{s}_t)}\left(\sum_{t=1}^{T}\nabla_{\theta'}\log \pi_{\theta'}(\mathbf{a}_t\mid\mathbf{s}_t)\right)\left(\left(\sum_{t'=t}^{T} r(\mathbf{s}_{t'}, \mathbf{a}_{t'})\right) - b\right)\right]$$The Variance Problem with Trajectory Importance Weights
There is a serious practical problem: the product $\prod_{t=1}^T \frac{\pi_{\theta'}(\mathbf{a}_t\mid\mathbf{s}_t)}{\pi_\theta(\mathbf{a}_t\mid\mathbf{s}_t)}$ involves $T$ terms multiplied together. For long horizons, this product can become astronomically large or vanishingly small, leading to extreme variance in the gradient estimate.
A common approximation avoids this issue by computing importance weights at the individual time step level rather than the trajectory level:
$$\nabla_{\theta'} J(\theta') \approx \frac{1}{N}\sum_{i=1}^{N}\sum_{t=1}^{T}\frac{\pi_{\theta'}(\mathbf{a}_{i,t}\mid\mathbf{s}_{i,t})}{\pi_\theta(\mathbf{a}_{i,t}\mid\mathbf{s}_{i,t})}\,\nabla_{\theta'}\log \pi_{\theta'}(\mathbf{a}_{i,t}\mid\mathbf{s}_{i,t})\left(\left(\sum_{t'=t}^{T} r(\mathbf{s}_{i,t'}, \mathbf{a}_{i,t'})\right) - b\right)$$This per-step importance weight $\frac{\pi_{\theta'}(\mathbf{a}_{i,t}\mid\mathbf{s}_{i,t})}{\pi_\theta(\mathbf{a}_{i,t}\mid\mathbf{s}_{i,t})}$ is a single ratio rather than a product of $T$ ratios, so it is much less likely to explode or vanish. However, this approximation ignores the distribution mismatch over states (only correcting the action distribution), introducing some bias. In practice, when $\theta'$ is close to $\theta$, this approximation works well.
Constraining the Policy Update
Off-policy policy gradients allow multiple gradient steps per batch of data, but if the policy changes too much between data-collection rounds, the importance weights become unreliable and the gradient estimate degrades. The data no longer reflects the states the updated policy would visit, and the approximation breaks down.
A natural solution is to constrain how far the policy can change during each update. One common and principled choice is to bound the KL divergence between the updated policy and the data-collection policy:
$$\E_{\mathbf{s} \sim \pi_\theta}\!\bigl[\Dkl\!\bigl(\pi_{\theta'}(\cdot\mid\mathbf{s})\,\|\,\pi_\theta(\cdot\mid\mathbf{s})\bigr)\bigr] \leq \delta$$This constraint ensures that $\pi_{\theta'}$ does not deviate too far from $\pi_\theta$ in any single update, keeping the importance weights well-behaved. The threshold $\delta$ controls the trade-off between making large steps (fast progress) and staying close to the data-collection policy (reliable gradients).
This idea—constraining the policy change per update while using importance sampling for off-policy corrections—is the foundation of trust region methods such as TRPO (Trust Region Policy Optimization) and the widely-used PPO (Proximal Policy Optimization) algorithm, which we will study in a later lecture.
Summary and Looking Ahead
This lecture developed the theory and algorithms of policy gradients, our first online RL method. The key ideas form a logical chain:
- The RL objective $J(\theta) = \E_{\tau \sim p_\theta(\tau)}[\sum_t r(\mathbf{s}_t, \mathbf{a}_t)]$ asks us to maximize expected cumulative reward over trajectories induced by the parameterized policy.
- The log-derivative trick converts the gradient of an expectation under a parameterized distribution into an expectation of a product—the log-probability gradient times the reward—that we can estimate from samples.
- The policy gradient theorem eliminates the need for a dynamics model: only the policy contributes to the gradient, since the initial state distribution and transition dynamics do not depend on $\theta$.
- REINFORCE is the resulting algorithm: sample trajectories, compute reward-weighted log-probability gradients, and take a gradient ascent step.
- Variance reduction via the causality trick (reward-to-go) and baselines (e.g., average reward) tightens the gradient estimate without introducing bias.
- Off-policy corrections via importance sampling allow reusing data from previous policies, improving sample efficiency at the cost of potential variance from importance weights.
- KL constraints prevent the policy from changing too much between data-collection rounds, keeping the off-policy approximation valid.