Skip to content

← JAX Stochasticity step 17 of 25

Hard Primitives

REINFORCE with Baseline

Jump to the task

Why this matters

A baseline changes the weight of each sampled score-function gradient from reward to reward - baseline. With a fixed baseline independent of the sampled action, the expected gradient stays the same: the expected score vector is zero. A well-chosen baseline can reduce variance, while a poorly chosen one can increase it.

This exercise uses a supplied scalar baseline. In a policy conditioned on a state, a baseline may depend on that state while remaining independent of the sampled action. The baseline is treated as a fixed weight when computing the policy gradient.

Worked mini-example

With logits [0, 0], each of two actions has probability 0.5. Their score vectors are [0.5, -0.5] and [-0.5, 0.5]. Let rewards be [10, 0] and the baseline be 5:

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

Every draw gives the same gradient vector in this example, so the variance is zero. Without a baseline, a draw gives either [5, -5] or [0, 0]. Both estimators have exact expectation [2.5, -2.5], but a finite batch without a baseline need not equal that expectation.

Problem

Implement reinforce_with_baseline(seed, logits, reward_table, baseline, n_samples):

  1. Create jax.random.PRNGKey(int(seed)) and sample actions with jax.random.categorical(key, logits, shape=(int(n_samples),)).
  2. For each action a, compute the gradient of jax.nn.log_softmax(logits)[a] with respect to logits.
  3. Multiply that vector by reward_table[a] - baseline.
  4. Average the weighted vectors over samples.

logits and reward_table are finite float32 vectors of equal nonzero length K. Rewards are fixed with respect to logits. baseline is a finite scalar independent of the sampled actions. seed is a nonnegative integer-valued number; 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,).

The fixtures compare this seeded finite-batch estimate. Setting baseline to zero reproduces the ordinary estimator for the same draw. Other baseline values preserve its expectation, not necessarily its value on one batch.

Common pitfalls

  • An action-dependent baseline generally changes the expected gradient unless the estimator includes a suitable correction.
  • A baseline computed from the same sampled rewards is not automatically independent of each action. Subtracting the same batch's mean reward gives a biased ordinary estimate; this exercise supplies the baseline.
  • The mean reward is not generally the best constant baseline. To minimize the sum of component variances, with score vector s, the optimum is E[reward * ||s||^2] / E[||s||^2] when the denominator is nonzero. It equals the mean reward when the score norm is constant.
  • Weight an (n_samples, K) score array with rewards[:, None], then average along the sample axis. Do not subtract baseline from the logits.
  • Mark seed and n_samples static if wrapping this implementation in jax.jit, because Python int(...) needs concrete values.