Skip to content

← JAX Stochasticity step 15 of 25

Medium Primitives

REINFORCE Gradient Estimator

Jump to the task

Why this matters

Suppose a policy chooses an action from softmax(logits) and receives a fixed reward from a table. We want the gradient of its expected reward with respect to the logits. REINFORCE estimates that gradient by sampling actions and multiplying each reward by the gradient of its log probability.

The reward need not be differentiable with respect to the action. The derivative follows the action's probability, not the discrete selection. Here the reward table is independent of logits; a reward that depends directly on the parameters would require an additional derivative term.

Worked mini-example

For probabilities p, the score vector for action a is one_hot(a, K) - p. With logits [0, 0], rewards [10, 0], and equal action probabilities:

  • Action 0 contributes 10 * [0.5, -0.5] = [5, -5].
  • Action 1 contributes 0 * [-0.5, 0.5] = [0, 0].

Averaging over the distribution gives [2.5, -2.5] exactly. A sampled batch gives that value only if half its actions are action 0. For example, 54 action-0 draws out of 100 give [2.7, -2.7].

import jax
import jax.numpy as jnp

logits = jnp.array([0.0, 0.0])
action = 0
score = jax.grad(lambda l: jax.nn.log_softmax(l)[action])(logits)
contribution = 10.0 * score  # [5, -5]

Problem

Implement reinforce_grad(seed, logits, reward_table, n_samples):

  1. Create jax.random.PRNGKey(int(seed)).
  2. Draw jax.random.categorical(key, logits, shape=(int(n_samples),)).
  3. Compute each sampled action's score vector with respect to logits.
  4. Multiply each vector by that action's reward and average over samples.

logits and reward_table are finite float32 vectors with the same nonzero length K. Rewards are fixed with respect to logits. seed is a nonnegative integer-valued number, and n_samples is a positive integer-valued number; both may arrive as floats. Assume intermediate values remain finite. Return a float32 array of shape (K,) containing the sampled estimate.

This is a gradient for increasing expected reward. A minimization optimizer would use its negative. The fixtures use the specified seeded draw, so do not replace it with the exact expectation's gradient.

Common pitfalls

  • Differentiate with respect to the floating-point logits. An integer action may be an index or a separate nondifferentiated argument; its presence does not prevent jax.grad from differentiating logits.
  • Use log probabilities, not raw logits. Their score vectors differ.
  • jax.vmap can compute a score vector for each action. For scores of shape (n_samples, K), multiply by rewards[:, None] before averaging over axis zero.
  • Under exact sampling, this fixed-reward estimator is unbiased, but one batch can be noisy. More samples or a suitable baseline can help; neither makes every finite estimate equal to the exact gradient.
  • Mark seed and n_samples static under jax.jit because the specified Python conversions need concrete values.