Prelude · Reinforcement learning

Objectives

Policy-gradient methods over features — a softmax policy, REINFORCE with a baseline, entropy, and the clipped PPO gradient.

open Prelude.Objectivesbrings in Prelude.Approximation, Prelude.Policies, Prelude.Kernels, Prelude.LinearAlgebraObjectives.hb →

Prelude.Objectives improves a policy directly instead of learning values and reading a policy off them. The policy is a softmax over linear preferences: for features x (n×p, one row per sampled state) and weights w (p×A), \(\pi(a \mid s_i) \propto \exp(\phi_i^\top w_{\cdot a})\). A batch also carries ha, the one-hot matrix of actions taken (n×A), and adv, an advantage or return for each sample (n) — for example from gae or episodeReturns in Returns.

Every gradient here is the score of Policies, scoreOf \(= \text{one-hot} - \pi\), multiplied by something. REINFORCE multiplies it by the advantage. PPO multiplies it by the advantage, the probability ratio, and a 0/1 mask: the clip does not scale a sample down, it switches it off when its ratio has left the trust region in the direction that would help. The gradients are for ascent; apply them with ascended.

A one-state bandit’s preference update and a deep actor’s gradient are the same expression over different features, so a bandit is simply x equal to a column of ones. The sample model Actor.hb uses this module.

Notation. \(X \in \mathbb{R}^{n \times p}\) (the argument x) has rows \(\phi_i^{\top}\); \(W \in \mathbb{R}^{p \times A}\) is w; \(H \in \{0, 1\}^{n \times A}\) (the argument ha) has \(H_{ia} = [\,a = a_i\,]\) for the action \(a_i\) taken in sample \(i\); and \(\hat A \in \mathbb{R}^n\) is adv. \([\,\cdot\,]\) is the Iverson bracket, and \(\operatorname{safe}(d) = d + [\,d = 0\,]\) is the guarded divisor from Linear algebra.

open Prelude.Objectives

A softmax policy over features

policyOf

policyOf(x, w) · x features matrix n×p · w weights matrix p×A → policy rows matrix n×A

Let \(X \in \mathbb{R}^{n \times p}\) (the argument x) be the features and \(W \in \mathbb{R}^{p \times A}\) (the argument w) the weights, and let \(Z = X W\). Then

\[\operatorname{policyOf}(X, W)_{ia} = \pi_W(a \mid s_i) = \frac{e^{Z_{ia}}}{\sum_{b=1}^{A} e^{Z_{ib}}} ,\]

computed stably by subtracting each row’s maximum before exponentiating. The softmax policy for each sampled state.

Intuition Each action gets a linear preference score from the state’s features, and the softmax turns scores into probabilities. Differences in preference become odds ratios, so the policy is always stochastic and changes smoothly with the weights — which is what makes it differentiable.

Use it when you learn a policy directly over discrete actions: REINFORCE, actor-critic or PPO actors, or a gradient bandit (features a single column of ones). Sample actions from its rows with caller-supplied randomness, and record the chosen action as a one-hot row of ha. For a policy read off learned values instead, use the policies of Prelude.Policies.

policy-gradientdecision-makingexploration

def policyOf(x, w) ≔ softmaxRows(qLinear(x, w))

REINFORCE with a baseline

\[\nabla_w J \approx \frac{1}{n} \sum_i \phi_i \big(\mathbf{1}_{a_i} - \pi(\cdot \mid s_i)\big)^\top \hat A_i\]

Subtracting a baseline from the returns lowers the variance of this estimate without changing its mean, because the score has expectation zero under the policy.

baselined

baselined(ret) · ret returns vector n → centred returns vector n

Let \(G \in \mathbb{R}^n\) (the argument ret) be the returns. Then, for \(i = 1, \dots, n\),

\[\operatorname{baselined}(G)_i = G_i - \frac{1}{n} \sum_{j=1}^{n} G_j .\]

The returns minus their batch mean, the simplest baseline.

Intuition Raw returns are often all positive, so every sampled action gets pushed up and the signal is buried in noise. Subtracting the batch average asks instead “was this sample better or worse than usual?” — the expected gradient is unchanged, but its variance drops.

Use it when you run REINFORCE without a critic and want a cheap variance reduction. It needs a batch of more than one sample, since a single return centres to zero. Dividing by the batch standard deviation as well is common practice, though this function does not. When you have a critic, a state-dependent baseline such as gae from Returns is better.

policy-gradientbias-variance-tradeoffcredit-assignment

def baselined(ret)           ≔ ret − einsum("i->", ret) / rows(ret)

reinforce

reinforce(x, w, ha, adv) · x features matrix n×p · w weights matrix p×A · ha one-hot actions matrix n×A · adv advantages vector n → gradient matrix p×A

Let \(X \in \mathbb{R}^{n \times p}\), \(W \in \mathbb{R}^{p \times A}\), \(H \in \{0, 1\}^{n \times A}\) (the argument ha) and \(\hat A \in \mathbb{R}^n\) (the argument adv), and let \(\pi = \operatorname{policyOf}(X, W)\). Then

\[\operatorname{reinforce}(X, W, H, \hat A)_{ja} = \frac{1}{n} \sum_{i=1}^{n} X_{ij}\, \big(H_{ia} - \pi_{ia}\big)\, \hat A_i .\]

The REINFORCE policy-gradient estimate, averaged over the batch.

Intuition Make actions that turned out well more likely and those that turned out badly less likely. For a softmax policy, “more likely” is the direction one-hot minus \(\pi\): raise the chosen action, lower the rest in proportion to their probability. The advantage scales and signs the push, and the features say which weights carry it.

Use it when you have complete episodes (or advantage estimates) collected by the current policy and want the simplest unbiased policy gradient. Apply it with ascended. The estimate is noisy, so centre returns (baselined) or use gae advantages. It must be recomputed after every update, since the data are only valid for the policy that collected them; to reuse a batch for several steps, use ppo.

policy-gradientcredit-assignment

def reinforce(x, w, ha, adv) ≔ einsum("ip,ia,i->pa", x, scoreOf(policyOf(x, w), ha), adv) / rows(x)

ascended

ascended(w, grad, rate) · w weights matrix p×A · grad gradient matrix p×A · rate step size scalar → weights matrix p×A

Let \(W \in \mathbb{R}^{p \times A}\) (the argument w), \(\nabla \in \mathbb{R}^{p \times A}\) (the argument grad) and \(\alpha > 0\) (the argument rate). Then

\[\operatorname{ascended}(W, \nabla, \alpha) = W + \alpha\, \nabla .\]

A gradient-ascent step.

Intuition Walk uphill: the gradients in this module point toward more expected return, so they are added, not subtracted. Keeping the step as its own definition makes the sign convention hard to get wrong.

Use it when you apply reinforce or ppo to the weights. The step size matters more than in supervised learning because the data change with the policy: small values (around 0.01 or below, relative to feature scale) are safer, and too large a step can collapse the policy onto one action. For loss gradients meant for descent, subtract instead.

policy-gradientoptimisation

def ascended(w, grad, rate)  ≔ w + rate · grad

entropyPush

entropyPush(x, w) · x features matrix n×p · w weights matrix p×A → mean entropy scalar

Let \(X \in \mathbb{R}^{n \times p}\) and \(W \in \mathbb{R}^{p \times A}\), and let \(\pi = \operatorname{policyOf}(X, W)\). Then

\[\operatorname{entropyPush}(X, W) = -\frac{1}{n} \sum_{i=1}^{n} \sum_{a=1}^{A} \pi_{ia} \ln\big(\pi_{ia} + [\,\pi_{ia} = 0\,]\big) .\]

The mean entropy of the policy over the batch, in nats. It is the value of the entropy bonus used to discourage premature certainty, not its gradient.

Intuition How undecided the policy still is, averaged over states: \(\ln A\) nats when every action is equally likely, 0 when it always picks the same one. Policy gradients tend to become certain too early and stop exploring, and this number shows that happening.

Use it when you monitor training — a fast drop toward zero warns of premature convergence — or report the entropy term of an objective such as PPO’s. It returns the value only, not a gradient, so adding it to reinforce or ppo does not act as an entropy bonus; if entropy collapses, take smaller steps or fewer updates per batch.

explorationpolicy-gradientmodel-evaluation

def entropyPush(x, w)        ≔ einsum("i->", hedging(policyOf(x, w))) / rows(x)

PPO

PPO compares the current weights w with the weights old that collected the batch through the ratio \(\rho_i = \pi_w(a_i \mid s_i) / \pi_{\text{old}}(a_i \mid s_i)\), and maximises

\[\frac{1}{n}\sum_i \min\big(\rho_i \hat A_i,\ \operatorname{clip}(\rho_i, 1-\varepsilon, 1+\varepsilon)\,\hat A_i\big).\]

A sample contributes gradient only while the unclipped term is the smaller one.

ratioOf

ratioOf(x, w, old, ha) · x features matrix n×p · w current weights matrix p×A · old behaviour weights matrix p×A · ha one-hot actions matrix n×A → ratios vector n

Let \(X \in \mathbb{R}^{n \times p}\), \(W, W^{\text{old}} \in \mathbb{R}^{p \times A}\) (the arguments w and old) and \(H \in \{0, 1\}^{n \times A}\) (the argument ha). Let \(c_i(V) = \operatorname{chanceOn}(X, V, H)_i = \pi_V(a_i \mid s_i)\). Then

\[\operatorname{ratioOf}(X, W, W^{\text{old}}, H)_i = \rho_i = \frac{c_i(W)}{c_i(W^{\text{old}}) + [\,c_i(W^{\text{old}}) = 0\,]} .\]

The probability ratio of each taken action under the current and the old policy.

Intuition How much more (or less) likely the current policy is to take the action the old policy actually took. It is the importance weight that lets data collected by the old policy stand in for data from the new one: 1 means nothing has changed for that sample.

Use it when you reuse a batch for several updates (PPO, other off-policy corrections), or monitor how far the policy has drifted — ratios far from 1 mean the batch is going stale. Keep old fixed while you iterate over the batch. An old probability of exactly zero is guarded to avoid division by zero, but such a sample could not have been collected anyway.

policy-gradientoff-policy-learning

def ratioOf(x, w, old, ha) ≔ chanceOn(x, w, ha) / safe(chanceOn(x, old, ha))

clipped

clipped(rt, eps) · rt ratios vector n · eps clip range ε scalar → clipped ratios vector n

Let \(\rho \in \mathbb{R}^n\) (the argument rt) be the ratios and \(0 < \varepsilon < 1\) (the argument eps). Then, entrywise,

\[\operatorname{clipped}(\rho, \varepsilon)_i = \operatorname{clip}(\rho_i, 1 - \varepsilon, 1 + \varepsilon) = \min\big(\max(\rho_i, 1 - \varepsilon),\ 1 + \varepsilon\big) .\]

The ratio clipped to \([1-\varepsilon, 1+\varepsilon]\).

Intuition A leash on the ratio: the policy may change a sample’s probability by at most a factor \(1 \pm \varepsilon\) before the objective stops rewarding further change. It is the “trust region” of PPO, expressed as a clamp rather than a constraint.

Use it when you evaluate the clipped PPO objective, or build kept. \(\varepsilon = 0.2\) is the usual choice; 0.1 is more conservative. The clamp works for any vector, so it also serves as a general entrywise clamp to \([1 - \varepsilon, 1 + \varepsilon]\).

policy-gradienttraining-stability

def clipped(rt, eps)       ≔ min(max(rt, 1 − eps), 1 + eps)

kept

kept(rt, adv, eps) · rt ratios vector n · adv advantages vector n · eps clip range ε scalar → mask vector n

Let \(\rho \in \mathbb{R}^n\) (the argument rt), \(\hat A \in \mathbb{R}^n\) (the argument adv) and \(0 < \varepsilon < 1\) (the argument eps). Then

\[\operatorname{kept}(\rho, \hat A, \varepsilon)_i = \big[\, \rho_i\, \hat A_i \le \operatorname{clip}(\rho_i, 1 - \varepsilon, 1 + \varepsilon)\, \hat A_i \,\big] .\]

It is 1 where the sample is inside the trust region, or outside it in the direction that hurts, and 0 where the clip has switched it off.

Intuition The pessimistic minimum in PPO is the same as a switch per sample. If the policy has already moved past the clip in the direction the advantage wants, the sample is switched off — no more credit for going further. If it has moved the wrong way, the sample stays on so the mistake can be corrected.

Use it when you compute the PPO gradient (as ppo does) or diagnose training: the share of zeros is the clip fraction, and a large one (say above a third) suggests steps that are too big or too many epochs over one batch.

policy-gradienttraining-stabilitymasking

def kept(rt, adv, eps)     ≔ rt · adv ≤ clipped(rt, eps) · adv

ppo

ppo(x, w, old, ha, adv, eps) · x features matrix n×p · w current weights matrix p×A · old behaviour weights matrix p×A · ha one-hot actions matrix n×A · adv advantages vector n · eps clip range ε scalar → gradient matrix p×A

Let \(X\), \(W\), \(W^{\text{old}}\), \(H\), \(\hat A\) and \(\varepsilon\) be as above, let \(\pi = \operatorname{policyOf}(X, W)\), \(\rho = \operatorname{ratioOf}(X, W, W^{\text{old}}, H)\) and \(\kappa = \operatorname{kept}(\rho, \hat A, \varepsilon)\). Then

\[\operatorname{ppo}(X, W, W^{\text{old}}, H, \hat A, \varepsilon)_{ja} = \frac{1}{n} \sum_{i=1}^{n} X_{ij}\, \big(H_{ia} - \pi_{ia}\big)\, \hat A_i\, \rho_i\, \kappa_i .\]

The gradient of the clipped PPO objective: the REINFORCE expression with each sample weighted by its ratio and the kept mask. When w equals old every ratio is 1 and it equals reinforce.

Intuition REINFORCE, reweighted by the ratio so an old batch can drive several updates, with the clip mask stopping any sample from pulling the policy too far from the one that collected it. The result is large, stable improvement steps without an explicit trust-region solve.

Use it when collecting experience is expensive and you want to squeeze several gradient steps out of each batch: freeze old, take a few passes of ascended with this gradient, then collect again. Typical settings are \(\varepsilon = 0.2\) with gae advantages normalised per batch. It equals reinforce on the first step, when w is old. There is no value loss or entropy gradient here; train the critic separately.

policy-gradienttraining-stabilityoff-policy-learning

def ppo(x, w, old, ha, adv, eps) ≔ einsum("ip,ia,i->pa", x, scoreOf(policyOf(x, w), ha), adv · ratioOf(x, w, old, ha) · kept(ratioOf(x, w, old, ha), adv, eps)) / rows(x)

Building blocks

  • chanceOn(x, w, ha) — Let \(X \in \mathbb{R}^{n \times p}\), \(W \in \mathbb{R}^{p \times A}\) and \(H \in \{0, 1\}^{n \times A}\), with \(\pi = \operatorname{policyOf}(X, W)\). Then \(\operatorname{chanceOn}(X, W, H)_i = \sum_{a=1}^{A} H_{ia}\, \pi_{ia} = \pi_W(a_i \mid s_i)\). The probability each sample’s taken action has under the weights w.

Examples

The examples below live in one file, Gradients.hb, so the class is Gradients. Vectors and matrices are passed as Tensor.

A REINFORCE step on a two-armed bandit

A bandit has one feature equal to 1. With zero weights the policy is 50/50; arm 0 was pulled with advantage 2, so the gradient is \(2 \cdot ([1, 0] - [0.5, 0.5]) = [1, -1]\): raise arm 0, lower arm 1. At the start of training PPO gives the same gradient.

open Prelude.Objectives

fn softmaxPolicy(x, w) ↦ policyOf(x, w)
fn reinforceGradient(x, w, ha, adv) ↦ reinforce(x, w, ha, adv)
fn ppoGradient(x, w, old, ha, adv, eps) ↦ ppo(x, w, old, ha, adv, eps)
using MLambda.Hilbert.Runtime;
using Examples;

var x     = (Tensor)new double[,] { { 1 } };
var w     = (Tensor)new double[,] { { 0, 0 } };
var taken = (Tensor)new double[,] { { 1, 0 } };

var policy = Gradients.SoftmaxPolicy(x, w);
var grad   = Gradients.ReinforceGradient(x, w, taken, (Tensor)new[] { 2d });
var ppo    = Gradients.PpoGradient(x, w, w, taken, (Tensor)new[] { 2d }, 0.2);

Output (verified):

policyOf = shape [1, 2] values [0.5, 0.5]
reinforce = shape [1, 2] values [1, -1]
ppo = shape [1, 2] values [1, -1]

When the PPO clip switches a sample off

A ratio of 1.5 with \(\varepsilon = 0.2\) is outside the trust region. With a positive advantage, pushing further would help, so the sample is switched off (\(3 > 2.4\)). With a negative advantage the unclipped term is already the smaller one (\(-3 \le -2.4\)), so it is kept.

open Prelude.Objectives

fn trustMask(rt, adv, eps) ↦ kept(rt, adv, eps)
using Examples;

var helpful = Gradients.TrustMask(1.5,  2.0, 0.2);
var harmful = Gradients.TrustMask(1.5, -2.0, 0.2);

Output (verified):

kept+ = 0
kept- = 1

Centring returns with a baseline

open Prelude.Objectives

fn withBaseline(ret) ↦ baselined(ret)
using MLambda.Hilbert.Runtime;
using Examples;

var centred = Gradients.WithBaseline((Tensor)new[] { 1d, 2d, 3d });

Output (verified):

baselined = [-1, 0, 1]