Value gradients and costates

We begin with a simple, interactive 2D environment that helps us visually compare the two methods and, hopefully, develop insight into them. We then cover the derivations of vanilla REINFORCE and the value-gradient method. Before we begin, if you need a refresher on policy gradients, see this excellent blog post by Lilian Weng or Spinning Up’s introduction to policy optimization. For a broader picture, the survey by Mohamed et al. (Mohamed et al., 2020) compares score-function and pathwise estimators (the mechanism behind value gradients) well beyond reinforcement learning.

The environment

Consider the 2D system whose drift is shown by the streamlines below. In this environment, we choose an initial state and let the system evolve for 64 noisy steps following the drift. A trajectory receives a reward of +1 if it ends on the right side, and 0 otherwise. The initial state is our only decision, and the reward arrives only after all 64 steps. The natural question is in which direction we should move the starting point to obtain high-return trajectories (here, the reward and the return are the same).

The reward is delayed and sparse, which gives the policy-gradient estimate a high variance. The initial decision (the location of the initial state) passes through many random transitions before a single delayed reward is observed. We represent the policy by this single decision, and the gradient with respect to this decision becomes the policy gradient.

Try dragging the initial point and pressing Resample a few times. You will see that the thick dashed red curve changes with each resample. This curve shows how the starting point moves under successive gradient-ascent steps following the gradient estimate of REINFORCE, where at each step we resample a new set of trajectories to compute the gradient. The solid blue curve follows the value-gradient estimate from the same starting point in the same way.

At this point, you might have noticed that if every sampled trajectory ends in the red region (i.e., zero reward), the REINFORCE estimate is exactly zero. If every trajectory ends in the blue region, every return is one, and the estimate is a plain average of scores. These scores have mean zero, and their average over a finite batch is pure noise. In both cases, the batch contains no contrast between good and bad runs, and REINFORCE learns from exactly that contrast. The value-gradient estimate is near zero in both cases as well, and rightly so.

The main message is that having access to a model (here, the exact model) provides us with higher-quality gradients and hence better decisions. Value gradients are one way to leverage a model to produce better gradient estimates. This contrasts with, for instance, using the model to simulate trajectories and then obtaining gradient estimates from them, as is more common in the literature. As we shall see later, value gradients can also be applied to approximate models (Heess et al., 2015), which makes them practical.

Preliminaries

To explain the two curves, we leave the environment for a while and turn to the general problem it was built from. Here, we introduce some of the notation and the assumptions we will use throughout.

As usual in reinforcement learning (perhaps a little predictably), let’s consider a Markov decision process (MDP) with states $s_t$, actions $a_t$, rewards $r_t \coloneqq r(s_t,a_t)$, and transitions $p(s_{t+1}\mid s_t,a_t)$. Starting from $s_0\sim p_0$, a policy $\pi$ generates a trajectory $\tau=(s_0,a_0,r_0,s_1,a_1,\dots)$. For a discount factor $\gamma\in(0,1)$, the goal is to find a policy that maximizes the expected cumulative reward.

In this problem, the policy $\pi$ influences the objective through the action distribution $\pi(a_t\mid s_t)$. Intuitively, we look for a policy that puts higher probability on actions that lead to better trajectories. We measure the quality of a trajectory by its return, defined as

The expected return from a given initial state defines the value function $V_\pi(s)=\mathbb E_{\pi,p}[G(\tau)\mid s_0=s]$. Averaging this value over $s_0\sim p_0$ gives the objective function $J(\pi)=\mathbb E_{p_0}[V_\pi(s_0)]$ of the policy-optimization problem. Although we define the objective as the expectation of the value function, we will not need to realize a value function explicitly.

For a parameterized policy $\pi_\theta$, write this objective as $J(\theta)=\mathbb E_{p_0,\pi_\theta,p}[G(\tau)]$. Following a gradient-ascent scheme, we will need to compute the gradient of this objective. We will compare two approaches to this computation. Intuitively, one approach updates $\theta$ to make actions in high-return trajectories more likely. The other method will update the parameters to make actions yield a higher return under the same trajectory. The former is the intuition behind REINFORCE, while the latter is the intuition behind value-gradient methods. This distinction may feel subtle at first, but it will become clearer as we derive the two gradients. In the environment, the initial state plays the role of $\theta$, and no actions occur between the decision and the reward.

REINFORCE

To derive the gradient estimated by REINFORCE, we start by differentiating the objective, which is given by

The parameter appears only in the distribution over trajectories. We therefore cannot move the gradient inside the expectation as it stands, because the return of a fixed trajectory does not depend on $\theta$ and its derivative is zero. The log-derivative trick resolves this by rewriting the derivative of a distribution as an expectation under that distribution.

For REINFORCE, we apply the log-trick with $X=\tau$ and $f(\tau)=G(\tau)$. Assuming the reward function has no direct dependence on $\theta$, the return of a fixed trajectory is independent of $\theta$. Thus, we obtain

It remains to compute the trajectory log likelihood. Assuming that the initial-state distribution and transitions do not depend directly on $\theta$, the log trajectory density factors as

Further assuming that $p$ does not depend on $\theta$, the transitions do not contribute to the derivative. Hence, we obtain the gradient

In the equation above, only the policy factors contribute to the derivative. This is the gradient estimated by REINFORCE (Williams, 1992).

We estimate the expectation by averaging return-weighted scores over sampled trajectories. Recall that in the 2D environment, we sample several trajectories and use them to compute the gradient estimates. Intuitively, the derivative acts on the log likelihood, while the return simply weights the gradients. In plain words, REINFORCE makes the sampled actions more likely in proportion to the return of their trajectory.

This is, however, the most naive form of REINFORCE. In practice, we reduce the variance using control variates, or simply a value function as a baseline. We use the plain version throughout to keep the comparison simple. For more details, Spinning Up’s section on baselines walks through both.

Value gradients

Let’s return to the same objective again and reconsider the assumptions we made in REINFORCE. Previously, we used the log-derivative trick to express the gradient of the objective $J(\theta)$ as an expectation of return-weighted scores. That route never looks inside the dynamics. When a differentiable model of the dynamics is available, we can obtain the same gradient with an alternative derivation. We will then show that the estimator of this derivation is the value-gradient estimator.

Let’s consider again the derivative of the objective, given by

Again, the parameter appears only in the distribution over trajectories. With a model, however, we can ask how the sampled trajectory itself would change if the parameters changed. In fact, as the parameters change, so do the policy and hence the sampled actions and states. However, this dependency passes through the sampling process, which is not differentiable as written, and we thus rely on another trick.

Reparameterization puts the parameter dependence inside the expectation, through the function $f$. The noise distribution $q$ stays fixed as $\theta$ changes. If $f$ and $h$ are differentiable, we can move differentiation inside the expectation, and we obtain

To apply reparameterization to a trajectory, we need representations of both the policy and the transitions. Assume that there exist deterministic and differentiable functions $\mu_\theta$ and $f$ satisfying

where the noise variables $\eta_t$ and $\zeta_t$ are independent across steps and of each other, with fixed distributions independent of the states, actions, and parameters. These representations let us rewrite the objective as

where $\varepsilon$ collects the initial state $s_0\sim p_0$ and all noise variables $\eta_t$ and $\zeta_t$ for $t \geq 0$. For fixed $\varepsilon$, define $\tau_\theta(\varepsilon)=(s_0,a_0,r_0,s_1,a_1,\dots)$ recursively by

We hold the initial state and noise variables in $\varepsilon$ fixed and generate the actions, later states, and rewards using this recursion.

We can now differentiate the rewritten objective. Since the distribution of $\varepsilon$ is independent of $\theta$, we can move the derivative inside the expectation, provided the sampled return is differentiable and suitable integrability conditions hold.

To estimate this gradient, we sample $\varepsilon$ and differentiate the return through the recursion. As we vary $\theta$, we change the actions through $\mu_\theta$ and propagate these changes through $f$ to later states and rewards. We then average the return derivatives to obtain a pathwise gradient estimator.

Estimator

Both REINFORCE and value gradients express $\nabla_\theta J(\theta)$ as an expectation of gradients. To turn either expression into an estimator, we replace the expectation with an average over samples. In REINFORCE, the expectation is over the trajectory distribution, and we can thus average the return-weighted scores of sampled trajectories. In value gradients, however, the expectation is over the noise $\varepsilon$, and the estimator needs samples of the noise, from which we then compute the trajectories and their return derivatives. This is impractical, since in most reinforcement-learning settings we observe only sampled trajectories and not the noise that generated them. We now introduce a way to use sampled trajectories in the value-gradient estimator, by inferring the noise that generated them.

Suppose we have a set of trajectories sampled under the policy $\pi_\theta$ and the transition model $p$. The noise of a step is simplest to infer when the observed action and next state determine it uniquely, and we restrict ourselves to this case. We assume that $\eta\mapsto\mu_\theta(s,\eta)$ is invertible for each state $s$ and denote its inverse by $\mu_{s,\theta}^{-1}$. Similarly, we assume that $\zeta\mapsto f(s,a,\zeta)$ is invertible for each pair $(s,a)$ and denote its inverse by $f_{s,a}^{-1}$. For each observed transition $(s,a,s^+)$, we then infer the noise by

In words, $\eta$ is the noise that makes the policy output the observed action, and $\zeta$ is the noise that makes the model produce the observed next state.

We repeat this calculation at every step and collect the inferred noise, together with the initial state, into $\varepsilon_i$ for trajectory $i$. At the parameters that generated the data, the reparameterized trajectory $\tau_\theta(\varepsilon_i)$ reproduces the observed trajectory, and as $\theta$ changes with $\varepsilon_i$ fixed, it shows how this trajectory would change. The gradient $\nabla_\theta G(\tau_\theta(\varepsilon_i))$ is thus a sample of the pathwise gradient $\mathbb E_\varepsilon[\nabla_\theta G(\tau_\theta(\varepsilon))]$ derived above, evaluated along an observed trajectory. When the noise was recorded during sampling, we use it as $\varepsilon_i$ directly.

Once the trajectories are reparameterized with the inferred noise, we can compute the gradient with respect to the policy parameters $\theta$ by simply running backpropagation. For a visual reference, see the computational graph above and switch on Compute gradient. Averaging across the sampled trajectories, we obtain the value-gradient estimate.

Costates

We can also build a recursion, very similar to the Bellman recursion, but this time over the gradient flow. Recall that $V(s)=\mathbb E_\varepsilon[G(\tau_\theta(\varepsilon))\mid s_0=s]$. Although our objective is to obtain $\nabla_\theta J(\theta)$, let us first introduce $\lambda(s)\coloneqq\nabla_s V(s)$, which we call the costate. With the policy-induced dynamics $f^\pi(s,\varepsilon)=f(s,\mu_\theta(s,\eta),\zeta)$ and reward $r^\pi(s,\varepsilon)=r(s,\mu_\theta(s,\eta))$, where $\varepsilon=(\eta,\zeta)$, we then obtain the recursion

Using the costates, we can then compute the gradient with respect to the policy parameters $\theta$ by

We can develop an interesting insight by interpreting the individual dimensions of the costate. Since costates carry sensitivity information, each feature (or dimension) of the costate can be seen as a marginal feature value. That is, each dimension of the costate tells us how influential the corresponding feature of the state is for the overall value. For example, a large $\lambda(s)_i$ indicates that any decision that increases feature $i$ has a high impact on the value.

Where this leads

In my opinion, value gradients and the costate framework offer an exciting direction in model-based reinforcement learning. They provide an elegant way to leverage a differentiable model for credit assignment, with gradient estimates that do not rely solely on statistical assumptions. If you want to read further, the following list provides some directions and references.

• Learned models. The main attraction of value gradients and costates comes when the model is learned. Stochastic value gradients learn the model from data and use it in the recursions of this post, which gives an elegant model-based reinforcement learning framework with recursions of its own (Heess et al., 2015). Dreamer uses the same mechanism over “dreamed” trajectories in a learned latent space (Hafner et al., 2020). In both, the policy is only as good as the model allows, since a fitted model biases the gradient in a way that more samples cannot fix.
• Language models. Asadulaev et al. use value gradients to explain why critic-free methods such as PPO and GRPO work for language-model post-training (Asadulaev et al., 2026). They show that the backward pass of the actor update propagates costates whose expectation is the value gradient, and that for transformer policies over discrete tokens, backpropagation through attention approximates these costates.
• When pathwise gradients fail. The costate recursion multiplies the Jacobians $\partial_s f^\pi$ along a trajectory, and when these Jacobians expand the state, the product grows exponentially with the horizon. In chaotic dynamics, this makes pathwise gradients explode (Metz et al., 2021; Parmas et al., 2018), and with contact or other discontinuities, they can also be biased (Suh et al., 2022).
• The general picture. Mohamed et al. provides an excellent survey of the relationship between score-function, pathwise, and measure-valued estimators across machine learning (Mohamed et al., 2020).

References