In the previous chapter, we explored various approaches to approximate dynamic programming, focusing on ways to handle large state spaces through function approximation. However, these methods still face significant challenges when dealing with large or continuous action spaces. The need to maximize over actions during the Bellman operator evaluation becomes computationally prohibitive as the action space grows.
This chapter explores a natural evolution of these ideas: rather than exhaustively searching over actions, we can parameterize and directly optimize the policy itself. We begin by examining how fitted Q methods, while effective for handling large state spaces, still struggle with action space complexity.
Recall that in fitted Q methods, the main idea is to compute the Bellman operator only at a subset of all states, relying on function approximation to generalize to the remaining states. At each step of the successive approximation loop, we build a dataset of input state-action pairs mapped to their corresponding optimality operator evaluations:
While this strategy allows us to handle very large or even infinite (continuous) state spaces, it still requires maximizing over actions (maxa∈A) during the dataset creation when computing the operator L for each basepoint. This maximization becomes computationally expensive for large action spaces. A natural improvement is to add another level of optimization: for each sample added to our regression dataset, we can employ numerical optimization methods to find actions that maximize the Bellman operator for the given state.
The above pseudocode introduces a generic maximize routine which represents any numerical optimization method that searches for an action maximizing the given function. This approach is versatile and can be adapted to different types of action spaces. For continuous action spaces, we can employ standard nonlinear optimization methods like gradient descent or L-BFGS (e.g., using scipy.optimize.minimize). For large discrete action spaces, we can use integer programming solvers - linear integer programming if the Q-function approximator is linear in actions, or mixed-integer nonlinear programming (MINLP) solvers for nonlinear Q-functions. The choice of solver depends on the structure of our Q-function approximator and the constraints on our action space.
This process is computationally intensive. A natural question is whether we can “amortize” some of this computation by replacing the explicit optimization for each sample with a direct mapping that gives us an approximate maximizer directly.
For Q-functions, recall that the operator is given by:
If q∗ is the optimal state-action value function, then v∗(s)=maxaq∗(s,a), and we can derive the optimal policy directly by computing the decision rule:
Note that π⋆ is implemented by our maximize numerical solver in the procedure above. A practical strategy would be to collect these maximizer values at each step and use them to train a function approximator that directly predicts these solutions. Due to computational constraints, we might want to compute these exact maximizer values only for a subset of states, based on some computational budget, and use the fitted decision rule to generalize to the remaining states. This leads to the following amortized version:
An important observation about this procedure is that the policy d(s;w) is being trained on a dataset Dd containing optimal actions computed with respect to an evolving Q-function. Specifically, at iteration n, we collect pairs (s′,as′∗) where as′∗=argmaxaq(s′,a;θn). However, after updating to θn+1, these actions may no longer be optimal with respect to the new Q-function.
A natural approach to handle this staleness would be to maintain only the most recent optimization data. We could modify our procedure to keep a sliding window of K iterations, where at iteration n, we only use data from iterations max(0, n-K) to n. This would be implemented by augmenting each entry in Dd with a timestamp:
where t indicates the iteration at which the optimal action was computed. When fitting the policy network, we would then only use data points that are at most K iterations old:
This introduces a trade-off between using more data (larger K) versus using more recent, accurate data (smaller K). The choice of K would depend on how quickly the Q-function evolves and the computational budget available for computing exact optimal actions.
Now the main issue with this approach, apart from the intrinsic out-of-distribution drift that we are trying to track, is that it requires “ground truth” - samples of optimal actions computed by the actual solver. This raises a natural question: how few samples do we actually need? Could we even envision eliminating the solver entirely? What seems impossible at first glance turns out to be achievable. The intuition is that as our policy improves at selecting actions, we can bootstrap from these increasingly better choices. As we continuously amortize these improving actions over time, it creates a virtuous cycle of self-improvement towards the optimal policy. But for this bootstrapping process to work, we need careful management. Move too quickly and the process may become unstable. Let’s examine how this balance can be achieved.
In this section, we consider deterministic parametrized policies of the form d(s;w) which directly output an action given a state. This approach differs from stochastic policies that output probability distributions over actions, making it particularly suitable for continuous control problems where the optimal policy is often deterministic. We’ll see how fitted Q-value methods can be naturally extended to simultaneously learn both the Q-function and such a deterministic policy.
When actions are continuous, a∈Rd, extracting a greedy policy from a Q-function becomes computationally expensive. Consider a robot arm control task where the action is a d-dimensional torque vector. To act greedily given Q-function q(s,a;θ), we must solve:
where A⊂Rd is a continuous set (often a box or polytope). This requires running an optimization algorithm at every time step. For neural network Q-functions, this means solving a nonlinear program whose objective involves forward passes through the network.
This computation happens at inference time. After training converges, we deploy the agent and it must select actions in real-time. Running interior-point methods or gradient-based optimizers at every decision creates unacceptable latency, especially in high-frequency control where decisions occur at 100Hz or faster.
The solution is to amortize the optimization cost by learning a separate policy network d(s;w) that directly outputs actions. During training, we optimize w so that d(s;w)≈argmaxaq(s,a;θ) for states we encounter. At deployment, action selection reduces to a single forward pass through the policy network: a=d(s;w). The computational cost of optimization is paid during training (where time is less constrained) rather than at inference.
This introduces a second approximation beyond the Q-function. We now have two function approximators: a criticq(s,a;θ) that estimates values, and an actord(s;w) that selects actions. The critic is trained using Bellman targets as in standard fitted Q-iteration. The actor is trained to maximize the critic:
where the expectation is over states in the dataset or replay buffer. This gradient ascent pushes the actor toward actions that the critic considers valuable. By the chain rule, this equals (∇aq(s,a;θ)∣a=d(s;w))⋅(∇wd(s;w)), which can be efficiently computed via backpropagation through the composition of the two networks.
Neural Fitted Q-iteration for Continuous Actions (NFQCA)¶
To develop this approach, let’s first consider an idealized setting where we have access to q⋆, the optimal Q-function. Then we can state our goal as finding policy parameters w that maximize q⋆ with respect to the actions chosen by our policy across the state space:
However, it’s computationally infeasible to satisfy this condition for every possible state s, especially in large or continuous state spaces. To address this, we assume a distribution of states, denoted μ(s), and take the expectation, leading to the problem:
However in practice, we do not have access to q∗. Instead, we need to approximate q∗ with a Q-function q(s,a;θ), parameterized by θ, which we will learn simultaneously with the policy function d(s;w). Given a samples of initial states drawn from μ, we then maximize this objective via a Monte Carlo surrogate problem:
When using neural networks to parametrize q and d, we obtain the Neural Fitted Q-Iteration with Continuous Actions (NFQCA) algorithm proposed by Hafner & Riedmiller (2011).
In practice, both the fit and minimize operations above are implemented using gradient descent. For the Q-function, the fit operation minimizes the mean squared error between the network’s predictions and the target values:
For the policy update, the minimize operation uses gradient descent on the composition of the “critic” network q and the “actor” network d. This results in the following update rule:
where α is the learning rate. Both operations can be efficiently implemented using modern automatic differentiation libraries and stochastic gradient descent variants like Adam or RMSProp.
Just as DQN adapted Neural Fitted Q-Iteration to the online setting, DDPG Lillicrap et al. (2015) extends NFQCA to learn from data collected online. Like NFQCA, DDPG simultaneously learns a Q-function and a deterministic policy that maximizes it, but differs in how it collects and processes data.
Instead of maintaining a fixed set of basepoints, DDPG uses a replay buffer that continuously stores new transitions as the agent interacts with the environment. Since the policy is deterministic, exploration becomes challenging. DDPG addresses this by adding noise to the policy’s actions during data collection:
where N represents exploration noise drawn from an Ornstein-Uhlenbeck (OU) process. The OU process is particularly well-suited for control tasks as it generates temporally correlated noise, leading to smoother exploration trajectories compared to independent random noise. It is defined by the stochastic differential equation:
where μ is the long-term mean value (typically set to 0), θ determines how strongly the noise is pulled toward this mean, σ scales the random fluctuations, and dWt is a Wiener process (continuous-time random walk). For implementation, we discretize this continuous-time process using the Euler-Maruyama method:
where Δt is the time step and ϵt∼N(0,1) is standard Gaussian noise. Think of this process like a spring mechanism: when the noise value Nt deviates from μ, the term θ(μ−Nt)Δt acts like a spring force, continuously pulling it back. Unlike a spring, however, this return to μ is not oscillatory - it’s more like motion through a viscous fluid, where the force simply decreases as the noise gets closer to μ. The random term σΔtϵt then adds perturbations to this smooth return trajectory. This creates noise that wanders away from μ (enabling exploration) but is always gently pulled back (preventing the actions from wandering too far), with θ controlling the strength of this pulling force.
The policy gradient update follows the same principle as NFQCA:
We then embed this exploration mechanism into the data collection procedure and use the same flattened FQI structure that we adopted in DQN. Similar to DQN, flattening the outer-inner optimization structure leads to the need for target networks - both for the Q-function and the policy.
Twin Delayed Deep Deterministic Policy Gradient (TD3)¶
While DDPG provided a foundation for continuous control with deep RL, it suffers from similar overestimation issues as DQN. TD3 Fujimoto et al. (2018) addresses these challenges through three key modifications: double Q-learning to reduce overestimation bias, delayed policy updates to reduce per-update error, and target policy smoothing to prevent exploitation of Q-function errors.
Similar to Double Q-learning, TD3 decouples selection from evaluation when forming the targets. However, instead of intertwining the two existing online and target networks, TD3 suggests learning two Q-functions simultaneously and uses their minimum when computing target values to help combat the overestimation bias further.
Furthermore, when computing target Q-values, TD3 adds small random noise to the target policy’s actions and clips it to keep the perturbations bounded. This regularization technique essentially implements a form of “policy smoothing” that prevents the policy from exploiting areas where the Q-function may have erroneously high values:
While DDPG used the OU process which generates temporally correlated noise, TD3’s authors found that simple uncorrelated Gaussian noise works just as well for exploration. It is also easier to implement and tune since you only need to set a single parameter (σexplore) for exploration rather than the multiple parameters required by the OU process (θ, μ, σ).
Finally, TD3 updates the policy network (and target networks) less frequently than the Q-networks, typically once every d Q-function updates. This helps reduce the per-update error and gives the Q-functions time to become more accurate before they are used to update the policy.
Adapting the intuition of NFQCA to the smooth Bellman optimality equations leads us to the soft actor-critic algorithm Haarnoja et al. (2018). To understand this connection, let’s first examine how the smooth Bellman equations follow from entropy regularization.
Consider the standard Bellman operator augmented with an entropy term. The smooth Bellman operator Lβ takes the form:
where H(π)=−Ea∼π[logπ(a∣s)] represents the entropy of the policy. To find the solution to the optimization problem embedded in the operator Lβ, we set the functional derivative of the objective with respect to the decision rule to zero:
Fitted Q-Iteration for the Smooth Bellman Equations¶
Following the principles of fitted value iteration, we can approximate approximate the effect of the smooth Bellman operator by computing it exactly at a number of basepoints and generalizing elsewhere using function approximation. Concretely, given a collection of states si and actions ai, we would compute regression target values:
The expectation over next states can be handled through Monte Carlo estimation using samples from the environment: given a transition (si,ai,si′), we can approximate:
However, we still face the challenge of computing the integral over actions. This motivates maintaining separate function approximators for both Q and V, using samples from the current policy to estimate the value function:
By maintaining both approximators, we can estimate targets using sampled actions from our policy. Specifically, if we have a transition (si,ai,si′) and sample ai′∼d(⋅∣si′;ϕ), our target becomes:
This approach exists only due to the dual representation of the smooth Bellman equations as an entropy-regularized problem, which transforms the intractable log-sum-exp into a form we can estimate efficiently through sampling.
The entropy-regularized objective and the smooth Bellman equation are mathematically equivalent. However, both formulations face a practical challenge: they require evaluating an intractable integral due to the Boltzmann distribution. Soft Actor-Critic (SAC) addresses this problem by approximating the optimal policy with a simpler, more tractable Gaussian distribution. Given the optimal soft policy:
This approximation task naturally raises the question of how to measure the “closeness” between the target Boltzmann distribution and a candidate Gaussian approximation. Following common practice in deep learning, we employ the Kullback-Leibler (KL) divergence as our measure of distributional distance. To find the best approximation, we minimize the KL divergence between our policy and the optimal policy, using our current estimate qθ of q∗:
However, an important question remains: how can we solve this optimization problem when it involves the intractable partition function Z(s)? To see this, recall that for two distributions p and q, the KL divergence takes the form DKL(p∥q)=Ex∼p[logp(x)−logq(x)]. Let’s denote the target Boltzmann distribution based on our current Q-estimate as:
One last challenge remains: ϕ appears in the distribution underlying the inner expectation, as well as in the integrand. This setting departs from standard empirical risk minimization (ERM) in supervised learning where the distribution of the data (e.g., cats and dogs in image classification) remains fixed regardless of model parameters. Here, however, the “data” - our sampled actions - depends directly on the parameters ϕ we’re trying to optimize.
This dependence prevents us from simply using sample average estimators and differentiating through them, as we typically do in supervised learning. The challenge of correctly and efficiently estimating such derivatives has been extensively studied in the simulation literature under the umbrella of “derivative estimation.” SAC adopts a particular solution known as the reparameterization trick in deep learning (or the IPA estimator in simulation literature). This approach transforms the problem by pushing ϕ inside the expectation through a change of variables.
To address this, we can express our Gaussian policy through a deterministic function fϕ that transforms noise samples to actions:
Now ϕ appears only in the integrand through the function fϕ, not in the sampling distribution. The objective involves two terms. First, the log-probability of our Gaussian policy has a simple closed form:
Second, ϕ enters through the composition of q⋆ with fϕ: q⋆(s,fϕ(s,ϵ)). The chain rule for this composition would involve derivatives of both functions. While this might be problematic if the Q-factors were to come from outside of our control (ie. not in the computational graph), but since SAC learns it simultaneously with the policy, then we can simply compute all required derivatives through automatic differentiation.
This composition of policy and value functions - where fϕ enters as input to qθ - directly parallels the structure we encountered in deterministic policy methods like NFQCA and DDPG. In those methods, we optimized:
Thus, rather than learning a single action for each state as in DDPG, we learn a function that transforms random noise into actions, explicitly parameterizing a distribution over actions while maintaining the same underlying principle of differentiating through composed policy and value functions.
Derivative Estimation for Stochastic Optimization¶
Consider optimizing an objective that involves an expectation:
While we can compute this exactly for the Gaussian example, this is often impossible for more general problems. We might then be tempted to approximate our objective using samples:
The issue here is that while the first term could be numerically integrated using the Monte Carlo, the second one can’t as it’s not an expectation.
Would there be a way to transform our objective in such a way that the Monte Carlo estimator for the objective could be differentiated directly while ensuring that the resulting derivative is unbiased? We will see that there are two main solutions to that problem: by doing a change of measure, or a change of variables.
The so-called “score function” derivative estimator is obtained for the choice of q(x)=p(x;θ), where the ratio simplifies to 1 and its derivative becomes the score function:
Therefore if we want to sample from some target distribution p(x;θ), we can do so by first sampling from a simple base distribution q(ϵ) (like a standard normal) and then transforming those samples through a carefully chosen function g. If g(⋅,θ) is invertible, the change of variables formula tells us how these distributions relate:
For example, if we want to sample from any multivariate Gaussian distributions with covariance matrix Σ and mean μ, it suffices to be able to sample from a standard normal noise and compute the linear transformation:
When we need samples constrained to an interval [a,b], we can use the truncated normal distribution. To sample from it, we transform uniform noise through the inverse cumulative distribution function (CDF) of the standard normal:
When we need samples in the unit interval [0,1], a natural choice might be the Beta distribution. However, its inverse CDF doesn’t have a closed form. Instead, we can use the Kumaraswamy distribution as a convenient approximation, which allows for a simple reparameterization:
When sampling from a categorical distribution with probabilities {πi}, one approach uses Gumbel(0,1) noise combined with the argmax of log-perturbed probabilities:
This approach, known in machine learning as the Gumbel-Max trick, relies on sampling Gumbel noise from uniform random variables through the transformation gi=−log(−log(ui)) where ui∼Uniform(0,1). To see why this gives us samples from the categorical distribution, consider the probability of selecting category i:
P(argmaxj(logπj+gj)=i)=P(logπi+gi>logπj+gj for all j=i)=P(gi−gj>logπj−logπi for all j=i)
Since the difference of two Gumbel random variables follows a logistic distribution, gi−gj∼Logistic(0,1), and these differences are independent for different j (due to the independence of the original Gumbel variables), we can write:
The last equality requires some additional algebra to show, but follows from the fact that these probabilities must sum to 1 over all i.
While we have shown that the Gumbel-Max trick gives us exact samples from a categorical distribution, the argmax operation isn’t differentiable. For stochastic optimization problems of the form:
we need g to be differentiable with respect to θ. This leads us to consider a continuous relaxation where we replace the hard argmax with a temperature-controlled softmax:
The resulting distribution over the probability simplex is called the Gumbel-Softmax (or Concrete) distribution. The temperature parameter τ controls the discreteness of our samples: smaller values give samples closer to one-hot vectors but with less stable gradients, while larger values give smoother gradients but more diffuse samples.
To get an analytical expression for the derivative, first note that we can factor out θ to obtain J(θ)=θE[x2] where x∼N(θ,1). By definition of the variance, we know that Var(x)=E[x2]−(E[x])2, which we can rearrange to E[x2]=Var(x)+(E[x])2. Since x∼N(θ,1), we have Var(x)=1 and E[x]=θ, therefore E[x2]=1+θ2. This gives us:
For concreteness, we fix θ=1.0 and analyze samples drawn using Monte Carlo estimation with batch size 1000 and 1000 independent trials. Evaluating at θ=1 gives us dθdJ(θ)∣∣θ=1=1+3(1)2=4, which serves as our ground truth against which we compare our estimators:
First, we consider the naive estimator that incorrectly differentiates the Monte Carlo approximation:
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
key = jax.random.PRNGKey(0)
# Define the objective function f(x,θ) = x²θ where x ~ N(θ, 1)
def objective(x, theta):
return x**2 * theta
# Naive Monte Carlo gradient estimation
@jax.jit
def naive_gradient_batch(key, theta):
samples = jax.random.normal(key, (1000,)) + theta
# Use jax.grad on the objective with respect to theta
grad_fn = jax.grad(lambda t: jnp.mean(objective(samples, t)))
return grad_fn(theta)
# Score function estimator (REINFORCE)
@jax.jit
def score_function_batch(key, theta):
samples = jax.random.normal(key, (1000,)) + theta
# f(x,θ) * ∂logp(x|θ)/∂θ + ∂f(x,θ)/∂θ
# score function for N(θ,1) is (x-θ)
score = samples - theta
return jnp.mean(objective(samples, theta) * score + samples**2)
# Reparameterization gradient
@jax.jit
def reparam_gradient_batch(key, theta):
eps = jax.random.normal(key, (1000,))
# Use reparameterization x = θ + ε, ε ~ N(0,1)
grad_fn = jax.grad(lambda t: jnp.mean(objective(t + eps, t)))
return grad_fn(theta)
# Run trials
n_trials = 1000
theta = 1.0
true_grad = 3 + theta**2
keys = jax.random.split(key, n_trials)
naive_estimates = jnp.array([naive_gradient_batch(k, theta) for k in keys])
score_estimates = jnp.array([score_function_batch(k, theta) for k in keys])
reparam_estimates = jnp.array([reparam_gradient_batch(k, theta) for k in keys])
# Create violin plots with individual points
plt.figure(figsize=(12, 6))
data = [naive_estimates, score_estimates, reparam_estimates]
colors = ['#ff9999', '#66b3ff', '#99ff99']
parts = plt.violinplot(data, showextrema=False)
for i, pc in enumerate(parts['bodies']):
pc.set_facecolor(colors[i])
pc.set_alpha(0.7)
# Add box plots
plt.boxplot(data, notch=True, showfliers=False)
# Add true gradient line
plt.axhline(y=true_grad, color='r', linestyle='--', label='True Gradient')
plt.xticks([1, 2, 3], ['Naive', 'Score Function', 'Reparam'])
plt.ylabel('Gradient Estimate')
plt.title(f'Gradient Estimators (θ={theta}, true grad={true_grad:.2f})')
plt.grid(True, alpha=0.3)
plt.legend()
# Print statistics
methods = {
'Naive': naive_estimates,
'Score Function': score_estimates,
'Reparameterization': reparam_estimates
}
for name, estimates in methods.items():
bias = jnp.mean(estimates) - true_grad
variance = jnp.var(estimates)
print(f"\n{name}:")
print(f"Mean: {jnp.mean(estimates):.6f}")
print(f"Bias: {bias:.6f}")
print(f"Variance: {variance:.6f}")
print(f"MSE: {bias**2 + variance:.6f}")
The numerical experiments coroborate our theory. The naive estimator consistently underestimates the true gradient by 2.0, though it maintains a relatively small variance. This systematic bias would make it unsuitable for optimization despite its low variance. The score function estimator corrects this bias but introduces substantial variance. While unbiased, this estimator would require many samples to achieve reliable gradient estimates. Finally, the reparameterization trick achieves a much lower variance while remaining unbiased. While this experiment is for didactic purposes only, it reproduces what is commonly found in practice: that when applicable, the reparameterization estimator tends to perform better than the score function counterpart.
Score Function Gradient Estimation in Reinforcement Learning¶
Let G(τ)≡∑t=0Tr(st,at) be the sum of undiscounted rewards in a trajectory τ. The stochastic optimization problem we face is to maximize:
Therefore, only the policy depends on w. When taking the logarithm of this product, we get a sum where all the w-independent terms vanish. The final estimator samples trajectories under the distribution p(τ;w) and computes:
This is a direct application of the score function estimator. However, we rarely use this form in practice and instead make several improvements to further reduce the variance.
Given the Markov property of the MDP, rewards rk for k<t are conditionally independent of action at given the history ht=(s0,a0,...,st−1,at−1,st). This allows us to only need to consider future rewards when computing policy gradients.
The condition independence assumption means that the term Eτ[∑t=0T∇wlogd(at∣st;w)∑k=0t−1rk] vanishes. To see this, let’s factor the trajectory distribution as:
The benefit of this estimator compared to the naive one is that it generally has less variance. More formally, we can show that this estimator arises from the application of a variance reduction technique known as the Extended Conditional Monte Carlo Method.
A control variate is a zero-mean random variable that we subtract from our estimator to reduce variance. Given an estimator Z and a control variate C with E[C]=0, we can construct a new unbiased estimator:
In the reinforcement learning setting, we usually choose Ct=∇wlogd(at∣st;w) as our control variate at each timestep. For a given state st, our estimator at time t is:
In practice when implementing this estimator, we won’t have access to the true value function. So as we did earlier for NFQCA or SAC, we commonly learn that value function simultaneously with the policy. Do do so, we could either using a fitted value approach, or even more simply just regress from states to sum of rewards to learn what Williams 1992 called a “baseline”:
When implementing this algorithm nowadays, we always using mini-batching to make full use of our GPUs. Therefore, a more representative variant for this algorithm would be:
The non-batched, purely online, GAE(0) algorithm then becomes:
This version was first derived by Richard Sutton in his 1984 PhD thesis. He called it the Adaptive Heuristic Actor-Critic algorithm. As far as I know, it was not derived using the score function method outlined here, but rather through intuitive reasoning (great intuition!).
Sutton 1999 provided an expression for the gradient of the infinite discounted return with respect to the parameters of a parameterized policy. Here is an alternative derivation by considering a bilevel optimization problem:
The Implicit Function Theorem states that if there is a solution to the problem F(v,w)=0, then we can “reparameterize” our problem as F(v(w),w) where v(w) is an implicit function of w. If the Jacobian ∂v∂F is invertible, then:
Here we made it clear in our notation that the derivative must be evaluated at root (v(w),w) of F. For the remaining of this derivation, we will drop this dependence to make notation more compact.
Applying this to our case with F(v,w)=(I−γPd)v−rd:
This is the policy gradient theorem, where xα(s) is the discounted state visitation distribution and the term in parentheses is the state-action value function q(s,a).
Normalized Discounted State Visitation Distribution¶
The discounted state visitation xα(s) is not normalized. Therefore the expression we obtained above is not an expectation. However, we can tranform it into one by normalizing by 1−γ. Note that for any initial distribution α:
Now we have expressed the policy gradient theorem in terms of expectations under the normalized discounted state visitation distribution. But what does sampling from μα means? Recall that xα⊤=α⊤(I−γPd)−1. Using the Neumann series expansion (valid when ∥γPd∥<1, which holds for γ<1 since Pd is a stochastic matrix) we have:
shows that μα is a mixture distribution: with probability 1−γ you draw a state from the initial distribution α (reset), and with probability γ you follow the policy dynamics Pd from the current state (continue). This interpretation directly connects to the geometric process: at each step you either terminate and resample from α (with probability 1−γ) or continue following the policy (with probability γ).
import numpy as np
def sample_from_discounted_visitation(
alpha,
policy,
transition_model,
gamma,
n_samples=1000
):
"""Sample states from the discounted visitation distribution.
Args:
alpha: Initial state distribution (vector of probabilities)
policy: Function (state -> action probabilities)
transition_model: Function (state, action -> next state probabilities)
gamma: Discount factor
n_samples: Number of states to sample
Returns:
Array of sampled states
"""
samples = []
n_states = len(alpha)
# Initialize state from alpha
current_state = np.random.choice(n_states, p=alpha)
for _ in range(n_samples):
samples.append(current_state)
# With probability (1-gamma): reset
if np.random.random() > gamma:
current_state = np.random.choice(n_states, p=alpha)
# With probability gamma: continue
else:
# Sample action from policy
action_probs = policy(current_state)
action = np.random.choice(len(action_probs), p=action_probs)
# Sample next state from transition model
next_state_probs = transition_model(current_state, action)
current_state = np.random.choice(n_states, p=next_state_probs)
return np.array(samples)
# Example usage for a simple 2-state MDP
alpha = np.array([0.7, 0.3]) # Initial distribution
policy = lambda s: np.array([0.8, 0.2]) # Dummy policy
transition_model = lambda s, a: np.array([0.9, 0.1]) # Dummy transitions
gamma = 0.9
samples = sample_from_discounted_visitation(alpha, policy, transition_model, gamma)
# Check empirical distribution
print("Empirical state distribution:")
print(np.bincount(samples) / len(samples))
While the math shows that sampling from the discounted visitation distribution μα would give us unbiased policy gradient estimates, Thomas (2014) demonstrated that this implementation can be detrimental to performance in practice. The issue arises because terminating trajectories early (with probability 1−γ) reduces the effective amount of data we collect from each trajectory. This early termination weakens the learning signal, as many trajectories don’t reach meaningful terminal states or rewards.
Therefore, in practice, we typically sample complete trajectories from the undiscounted process (i.e., run the policy until natural termination or a fixed horizon) while still using γ in the advantage estimation. This approach preserves the full learning signal from each trajectory
and has been empirically shown to lead to better performance.
This is one of several cases in RL where the theoretically optimal procedure differs from the best practical implementation.
In this section, we’ll explore how incorporating a model of the dynamics can help us design better policy gradient estimators. Let’s begin with a pure model-free approach that uses a critic to maximize a deterministic policy:
While w doesn’t appear in the outer expectation over initial states, it affects the inner expectation over next states—a distributional dependency. As a result, the product rule of calculus yields two terms: the first being an expectation, and the second being problematic for sample average estimation. However, we have tools to address this challenge: we can either apply the reparameterization trick to the dynamics or use score function estimators.
For the reparameterization approach, assuming s′=f(s,a,ξ;w) where ξ is the noise variable:
We’ve now developed a hybrid algorithm that combines model-based and model-free approaches while integrating derivative estimators for stochastic dynamics with a deterministic policy parameterization. While this hybrid estimator remains relatively unexplored in practice, it could prove valuable for systems with specific structural properties.
Consider a robotics scenario with hybrid continuous-discrete dynamics: a robotic arm operates in continuous space but interacts with discrete object states. While the arm’s policy remains differentiable (∇wd), the object state transitions follow categorical distributions. In this case, reparameterization becomes impractical, but the score function approach is viable if we can compute ∇wlogp(s′∣s,d(s;w)) from the known transition model. Similar structures arise in manufacturing processes, where machine actions might be continuous and differentiable, but material state transitions often follow discrete steps with known probabilities. Note that both approaches require knowledge of transition probabilities and won’t work with pure black-box simulators or systems where we can only sample transitions without probability estimates.
Another dimension to explore in our algorithm design is the number of steps we wish to unroll our model. This allows us to better understand and control the bias-variance tradeoffs in our methods.
The questions of derivative estimators only arise with stochastic dynamics. For systems with deterministic dynamics and a deterministic policy, the one-step gradient unroll simplifies to:
where s′=f(s,d(s;w)) is the deterministic next state. Notice the resemblance between this expression and that obtained for ∇wJDPMB-R above. They are essentially the same except that in the reparameterized case, the dynamics have made the dependence on the noise variable explicit and the outer expectation has been updated accordingly. This similarity arises because differentiation through reparameterized dynamics models is, in essence, backpropagation: it tracks the propagation of perturbations through a computation graph—which we refer to as a stochastic computation graph in this setting.
Still under this simplified setting with deterministic policies and dynamics, we can extend the expression for the gradient through n-steps of model unroll, leading to:
with base case ∇ws0=0 since the initial state does not depend on the policy parameters.
At both ends of the spectrum, we have that for n=0, we fall back to the pure critic NFQCA-like approach, and for n=∞, we don’t bootstrap at all and are fully model-based without a critic. The pure model-based critic-free approach to optimization is what we may refer to as a backpropagation-based policy optimization (BPO).
But just as backpropagation through RNNs or very deep networks can be challenging due to exploding and vanishing gradients, “vanilla” Backpropagation Policy Optimization (BPO) without a critic can severely suffer from the curse of horizon. This is because it essentially implements single shooting optimization. Using a critic can be an effective remedy to this problem, allowing us to better control the bias-variance tradeoff while preserving gradient information that would be lost with a more drastic truncated backpropagation approach.
The stochastic value gradient framework of Heess (2015) applies the recipe for policy optimization with a model using reparameterized dynamics and action selection via randomized policies. In this setting, the stochastic policy model based reparameterized estimator over n steps is
While we could implement this expression for the gradient ourselves, this approach is much easier, less error-prone, and most likely better optimized for performance when using automatic differentiation. Given a set of rollouts (for which we know the primitive noise variables), then we can compute the monte carlo objective:
where the states are generated recursively using the known noise variables: starting with initial state s0m, each subsequent state is computed as si+1m=f(sim,d(sim,ϵim;w),ξim). Thus, a trajectory is completely determined by just the sequence of noise variables:(ϵ0m,ξ0m,ϵ1m,ξ1m,...,ϵn−1m,ξn−1m,ϵnm) where ϵim are the action noise variables and ξim are the dynamics noise variables.
The choice of unroll steps n gives us precise control over the balance between model-based and critic-based components in our gradient estimator. At one extreme, setting n=0 yields a purely model-free algorithm known as SVG(0) (Heess et al., 2015), which relies entirely on the critic for value estimation:
In the 2015 paper, the authors make a specific choice to reparameterize both the dynamics and action selections using normal distributions. For the policy, they use:
The method we’ve presented so far assumes we have direct access to the noise variables ϵ and ξ used to generate trajectories. This works well in the on-policy setting where we generate our own data. However, in off-policy scenarios where we receive externally generated trajectories, we need to infer these noise variables—a process the authors call noise inference.
For the Gaussian case discussed above, this inference is straightforward. Given an observed scalar action at and the current policy parameters w, we can recover the action noise ϵt through inverse reparameterization:
This simple inversion is possible because the Gaussian reparameterization is an affine transformation, which is invertible as long as σ(st;w) is non-zero for scalar actions and Σ(st,at;ϕ) is full rank for vector-valued states. The same principle extends naturally to vector-valued actions, where σ would be replaced by a full covariance matrix.
More generally, this idea of invertible transformations can be extended far beyond simple Gaussian reparameterization. We could consider a sequence of invertible transformations:
This is the idea behind normalizing flows: a series of invertible transformations that can transform a simple base distribution (like a standard normal) into a complex target distribution while maintaining exact invertibility.
The noise inference in this case would involve applying the inverse transformations:
More expressive policies and dynamics models capable of capturing multimodal distributions
Exact likelihood computation through the change of variables formula (can be useful for computing the log prob terms in entropy regularization for example)
Precise noise inference through the guaranteed invertibility of the flow
As far as I know, this approach has not been explored in the literature.
At first glance, SAC and DPG might appear to be fundamentally different approaches to policy optimization. SAC begins with the principle of entropy maximization and policy distribution matching through KL divergence minimization, while DPG directly optimizes a deterministic policy to maximize expected Q-values. However, we can show that DPG is a special case of SAC when the temperature parameter goes to zero.
Trust region methods in optimization approximate the objective function with a simpler local model within a region where we “trust” this approximation to be good. This brings about the need to define what we mean by a local region, and therefore to pick a geometry which suits our problem.
In standard optimization in the Euclidean space on Rn, at each iteration k, we create a quadratic approximation around the current point xk:
However, when optimizing over probability distributions p(x;θ), the Euclidean geometry becomes unnatural. Instead, the Kullback-Leibler divergence provides a more natural mean of measuring proximity:
Hafner, R., & Riedmiller, M. (2011). Reinforcement learning in feedback control: Challenges and benchmarks from technical process control. Machine Learning, 84(1–2), 137–169. 10.1007/s10994-011-5235-x
Lillicrap, T. P., Hunt, J. J., Pritzel, A., Heess, N., Erez, T., Tassa, Y., Silver, D., & Wierstra, D. (2015). Continuous Control with Deep Reinforcement Learning. arXiv Preprint arXiv:1509.02971.
Fujimoto, S., Hoof, H., & Meger, D. (2018). Addressing Function Approximation Error in Actor-Critic Methods. International Conference on Machine Learning (ICML), 1587–1596.
Haarnoja, T., Zhou, A., Abbeel, P., & Levine, S. (2018). Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. Proceedings of the 35th International Conference on Machine Learning (ICML), 1861–1870.