Prelude · Reinforcement learning

Approximation

Value learning when the table will not fit — linear semi-gradient TD, exact LSTD and the DQN target matrix.

open Prelude.Approximationbrings in Prelude.Networks, Prelude.Returns, Prelude.Policies, Prelude.Kernels, Prelude.LinearAlgebraApproximation.hb →

Prelude.Approximation replaces the value table of Temporal with a function of a feature vector. Each state is described by \(p\) features; a batch of \(n\) transitions arrives as x (n×p, the features of each state) and xnext (n×p, the features of each successor), with rewards rw (n), terminal flags done (n) and discount g. A linear state value uses weights w of shape p; a linear action value uses p×A.

Three methods are provided. Semi-gradient TD(0) takes a gradient step on the TD error while treating the bootstrapped target as a constant. LSTD jumps straight to the least-squares TD fixed point in a single linear solve. As in Decisions, the LSTD system \(\sum \phi(\phi - \gamma\phi')^\top\) is not symmetric and solve is conjugate gradient, so it is solved through its normal equations with a tiny ridge of \(10^{-8}\); expect agreement to about five significant digits, not exact equality.

The DQN target needs no gradient code of its own: aimed builds a target matrix equal to the online network’s outputs everywhere except the taken action’s column, so the squared error is zero on every other action and the ordinary two-layer network steps of Prelude.Networks (stepFirst, stepSecond) become the masked DQN step. The target reads a separate set of held (target-network) weights, so nothing is differentiated through it. The sample model Deep.hb is built this way.

Notation. \(X \in \mathbb{R}^{n \times p}\) (the argument x) has rows \(\phi_i^{\top}\), the features of the state of transition \(i\); \(X' \in \mathbb{R}^{n \times p}\) (the argument xnext) has rows \(\phi_i'^{\top}\), the features of its successor. \(r \in \mathbb{R}^n\) is rw, \(d \in \{0, 1\}^n\) is done, and \(\gamma \in [0, 1]\) is g.

open Prelude.Approximation

Linear value and action value

worthOf

worthOf(x, w) · x features matrix n×p · w weights vector p → values vector n

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

\[\operatorname{worthOf}(X, w) = X w, \qquad \operatorname{worthOf}(X, w)_i = \hat V(s_i) = \phi_i^{\top} w = \sum_{j=1}^{p} X_{ij}\, w_j .\]

The linear state value of each row.

Intuition Instead of storing one value per state, describe each state by features and let its value be a weighted sum of them. States that share features share value, so experience in one state generalises to similar states never visited.

Use it when the state space is too large for a table, or continuous: tile codings, radial basis features, polynomial or hand-built features. With one-hot features it reproduces a table exactly. It is the critic that semiStep trains and lstd solves for. For a value per action, use qLinear; for nonlinear values, the network of Prelude.Networks.

value-learningpolicy-evaluationlarge-state-spaces

def worthOf(x, w) ≔ einsum("ip,p->i", x, w)

qLinear

qLinear(x, w) · x features matrix n×p · w weights matrix p×A → action values 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. Then

\[\operatorname{qLinear}(X, W) = X W, \qquad \operatorname{qLinear}(X, W)_{ia} = \hat Q(s_i, a) = \sum_{j=1}^{p} X_{ij}\, W_{ja} .\]

The linear action values.

Intuition One linear value function per action, sharing the same state features: column \(a\) of the weights scores action \(a\). The result is a table of action values, one row per state, computed from features instead of looked up.

Use it when you choose actions from a large or continuous state space with a small discrete action set — linear Q-learning or SARSA, or an action-value critic. Feed the rows to highest or to a policy from Prelude.Policies to act. When states are described well enough by features alone and you only need state values, worthOf is smaller.

value-learningdecision-makinglarge-state-spaces

def qLinear(x, w) ≔ einsum("ip,pa->ia", x, w)

Semi-gradient TD

\[w \leftarrow w + \frac{\alpha}{n} \sum_i \delta_i\,\phi_i, \qquad \delta_i = r_i + \gamma(1-d_i)\,\phi'^\top_i w - \phi_i^\top w\]

semiStep

semiStep(w, x, xnext, rw, g, done, rate) · w weights vector p · x features matrix n×p · xnext successor features matrix n×p · rw rewards vector n · g discount scalar · done terminal flags vector n · rate step size α scalar → weights vector p

Let \(w \in \mathbb{R}^p\) be the weights, \(X, X' \in \mathbb{R}^{n \times p}\) (the arguments x and xnext), \(r \in \mathbb{R}^n\) (the argument rw), \(\gamma \in [0, 1]\) (the argument g), \(d \in \{0, 1\}^n\) (the argument done) and \(\alpha > 0\) (the argument rate). Let \(\delta_i = r_i + \gamma\,(1 - d_i)\, \phi_i'^{\top} w - \phi_i^{\top} w\). Then

\[\operatorname{semiStep}(w, X, X', r, \gamma, d, \alpha) = w + \frac{\alpha}{n}\, X^{\top} \delta = w + \frac{\alpha}{n} \sum_{i=1}^{n} \delta_i\, \phi_i .\]

One batch step of semi-gradient TD(0), averaging over the batch.

Intuition Nudge each feature’s weight in proportion to the TD error of the states where that feature is active. It is “semi” because the target \(r + \gamma \hat V(s')\) also depends on \(w\), but is treated as fixed; following the full gradient would drag the target toward the guess and learns worse.

Use it when transitions arrive continually and you want cheap, incremental learning: \(O(np)\) per step, averaged over the batch. Keep rate small relative to the feature scale, since large features with a large step diverge. With linear features and on-policy data it converges; off-policy data with bootstrapping can diverge. For a fixed batch where you want the answer directly, use lstd.

value-learningpolicy-evaluationonline-learning

def semiStep(w, x, xnext, rw, g, done, rate) ≔ w + rate · einsum("ip,i->p", x, semiErrors(x, xnext, rw, g, done, w)) / rows(x)

Building blocks

  • semiErrors(x, xnext, rw, g, done, w) — Let \(X, X' \in \mathbb{R}^{n \times p}\), \(r \in \mathbb{R}^n\), \(\gamma \in [0, 1]\), \(d \in \{0, 1\}^n\) and \(w \in \mathbb{R}^p\). Then \(\operatorname{semiErrors}(X, X', r, \gamma, d, w)_i = \delta_i = r_i + \gamma\,(1 - d_i)\,(X' w)_i - (X w)_i\). The TD error of each transition under the linear value, via tdErrors.

LSTD

LSTD solves \(A w = b\) with

\[A = \sum_i \phi_i (\phi_i - \gamma \phi'_i)^\top, \qquad b = \sum_i \phi_i r_i .\]

lstd

lstd(x, xnext, rw, g) · x features matrix n×p · xnext successor features matrix n×p · rw rewards vector n · g discount scalar → weights vector p

Let \(X, X' \in \mathbb{R}^{n \times p}\) (the arguments x and xnext), \(r \in \mathbb{R}^n\) (the argument rw) and \(\gamma \in [0, 1]\) (the argument g). Let \(A = X^{\top}(X - \gamma X') \in \mathbb{R}^{p \times p}\) and \(b = X^{\top} r \in \mathbb{R}^p\). Then

\[\operatorname{lstd}(X, X', r, \gamma) = \operatorname{solve}\big(A^{\top} A + 10^{-8}\, I_p,\; A^{\top} b\big) \approx \big(A^{\top} A + 10^{-8} I_p\big)^{-1} A^{\top} b .\]

The least-squares TD fixed point, by one conjugate-gradient solve of the ridge-regularised normal equations; for a well-conditioned non-singular \(A\) this is \(A^{-1} b\) to about five significant digits. lstd takes no terminal flags: for a transition into a terminal state, pass a row of zeros in xnext.

Intuition Semi-gradient TD, run to convergence, settles where the average TD update is zero. That point is the solution of a linear system, so instead of taking thousands of small steps LSTD builds the system from the batch and solves it once — no step size to tune.

Use it when you evaluate a fixed policy from a stored batch and have a modest number of features: building the system is \(O(np^2)\) and the solve \(O(p^3)\), so it suits tens to hundreds of features, not millions. It is very data-efficient. Mark terminal transitions with a zero row in xnext, since there are no done flags. Solving through the normal equations squares the condition number, so poorly scaled or redundant features cost accuracy; for large \(p\) or streaming data use semiStep.

policy-evaluationvalue-learninglinear-systems

def lstd(x, xnext, rw, g) ≔ solve(einsum("uv,uw->vw", lstdLeft(x, xnext, g), lstdLeft(x, xnext, g)) + 1e-8 · identity(lstdLeft(x, xnext, g)), einsum("uv,u->v", lstdLeft(x, xnext, g), lstdRight(x, rw)))

Building blocks

  • lstdLeft(x, xnext, g) — Let \(X, X' \in \mathbb{R}^{n \times p}\) and \(\gamma \in [0, 1]\). Then \(\operatorname{lstdLeft}(X, X', \gamma) = A = X^{\top}(X - \gamma X') = \sum_{i=1}^{n} \phi_i\, (\phi_i - \gamma \phi'_i)^{\top}\). The p×p matrix A.
  • lstdRight(x, rw) — Let \(X \in \mathbb{R}^{n \times p}\) and \(r \in \mathbb{R}^n\). Then \(\operatorname{lstdRight}(X, r) = b = X^{\top} r = \sum_{i=1}^{n} \phi_i\, r_i\). The vector b.

The DQN target matrix

The network is the two-layer ReLU network of Prelude.Networks, outputs(x, w1, w2), whose outputs are the action values n×A. w1, w2 are the online weights and h1, h2 the held target-network weights; ha is the one-hot matrix of actions taken (n×A).

aimed

aimed(x, xnext, ha, rw, g, done, w1, w2, h1, h2) · x features matrix n×p · xnext successor features matrix n×p · ha one-hot actions matrix n×A · rw rewards vector n · g discount scalar · done terminal flags vector n · w1, w2 online weights · h1, h2 target weights → targets matrix n×A

Let \(X, X' \in \mathbb{R}^{n \times p}\) (the arguments x and xnext), \(H \in \{0, 1\}^{n \times A}\) (the argument ha) the one-hot actions taken, \(r\), \(\gamma\) and \(d\) as above, \(W_1, H_1 \in \mathbb{R}^{p \times h}\) and \(W_2, H_2 \in \mathbb{R}^{h \times A}\) (the arguments w1, h1, w2, h2) the online and held weights. Let \(O = \operatorname{relu}(X W_1)\, W_2\) be the online outputs and \(y_i = r_i + \gamma\,(1 - d_i) \max_{a} \big(\operatorname{relu}(X' H_1)\, H_2\big)_{ia}\). Then

\[\operatorname{aimed}(X, X', H, r, \gamma, d, W_1, W_2, H_1, H_2)_{ia} = O_{ia}\,(1 - H_{ia}) + H_{ia}\, y_i .\]

The regression target for one DQN step: the online outputs, with the taken action’s entry replaced by the bootstrap target from the held network. Use it as y in the network’s gradient steps.

Intuition Only the action actually taken has a new target; for every other action the network is told to predict what it already predicts, so its error — and its gradient — there is zero. That turns Q-learning into ordinary regression. The target uses a separate, slowly updated copy of the network so it does not shift with every step.

Use it when you train a deep Q-network: pass the result as y to the steps of Prelude.Networks, and copy the online weights into h1, h2 every few hundred to few thousand steps. ha must be one-hot per row. The max over the held network’s values tends to overestimate; nothing here corrects that. For linear features without a network, qLinear with a TD update is simpler.

value-learningvalue-targetslarge-state-spaces

def aimed(x, xnext, ha, rw, g, done, w1, w2, h1, h2) ≔ outputs(x, w1, w2) · (1 − ha) + ha · lifted(ha, heldBest(xnext, h1, h2, rw, g, done))

Building blocks

  • heldBest(xnext, h1, h2, rw, g, done) — Let \(X' \in \mathbb{R}^{n \times p}\), \(H_1 \in \mathbb{R}^{p \times h}\), \(H_2 \in \mathbb{R}^{h \times A}\), and \(r\), \(\gamma\), \(d\) as above. Then \(\operatorname{heldBest}(X', H_1, H_2, r, \gamma, d)_i = r_i + \gamma\,(1 - d_i) \max_{a} \big(\operatorname{relu}(X' H_1)\, H_2\big)_{ia}\). The DQN bootstrap target for each transition, from the held network.

Examples

The first two examples live in one file, Features.hb, so the class is Features. Everything is passed as Tensor.

LSTD recovers a tabular value

With one-hot features, linear values are a table. Both transitions lead to state 1, and only state 1 pays a reward of 1, so \(w_1 = 1 + 0.9 w_1 = 10\) and \(w_0 = 0.9 w_1 = 9\). The solve agrees to about five digits.

open Prelude.Approximation

fn linearFixedPoint(x, xnext, rw, g) ↦ lstd(x, xnext, rw, g)
fn linearWorth(x, w) ↦ worthOf(x, w)
using MLambda.Hilbert.Runtime;
using Examples;

var x     = (Tensor)new double[,] { { 1, 0 }, { 0, 1 } };
var xnext = (Tensor)new double[,] { { 0, 1 }, { 0, 1 } };
var w     = Features.LinearFixedPoint(x, xnext, (Tensor)new[] { 0d, 1d }, 0.9);
var v     = Features.LinearWorth(x, (Tensor)new[] { 9d, 10d });

Output (verified, rounded to six digits):

lstd = [8.999984, 9.999982]
worthOf = [9, 10]

One semi-gradient TD step

From zero weights the TD errors are \([0, 1]\); with \(\alpha = 0.5\) averaged over two samples, the second weight moves to 0.25.

open Prelude.Approximation

fn tdWeights(w, x, xnext, rw, g, done, rate) ↦ semiStep(w, x, xnext, rw, g, done, rate)
using MLambda.Hilbert.Runtime;
using Examples;

var x     = (Tensor)new double[,] { { 1, 0 }, { 0, 1 } };
var xnext = (Tensor)new double[,] { { 0, 1 }, { 0, 1 } };
var w1    = Features.TdWeights((Tensor)new[] { 0d, 0d }, x, xnext, (Tensor)new[] { 0d, 1d }, 0.9, (Tensor)new[] { 0d, 0d }, 0.5);

Output (verified):

semiStep = [0, 0.25]

Build the DQN regression target

A one-sample batch on the two-layer network of Prelude.Networks with identity first-layer weights and second-layer weights [[1, 2], [3, 4]], the same for the online and held networks. The online outputs at x are \([1, 2]\) and the held outputs at xnext are \([3, 4]\). Action 0 was taken with reward 1 at \(\gamma = 0.5\), so its entry becomes \(1 + 0.5 \cdot 4 = 3\) and action 1 keeps its own output 2; after a terminal transition the entry is just the reward. The first line shows qLinear on one-hot features, which reads the weight rows.

open Prelude.Approximation

fn linearActions(x, w) ↦ qLinear(x, w)
fn dqnTargets(x, xnext, ha, rw, g, done, w1, w2, h1, h2) ↦ aimed(x, xnext, ha, rw, g, done, w1, w2, h1, h2)
using MLambda.Hilbert.Runtime;
using Examples;

var x     = (Tensor)new double[,] { { 1, 0 } };
var xnext = (Tensor)new double[,] { { 0, 1 } };
var taken = (Tensor)new double[,] { { 1, 0 } };
var w1    = (Tensor)new double[,] { { 1, 0 }, { 0, 1 } };
var w2    = (Tensor)new double[,] { { 1, 2 }, { 3, 4 } };

var q        = Targets.LinearActions((Tensor)new double[,] { { 1, 0 }, { 0, 1 } }, w2);
var target   = Targets.DqnTargets(x, xnext, taken, (Tensor)new[] { 1d }, 0.5, (Tensor)new[] { 0d }, w1, w2, w1, w2);
var terminal = Targets.DqnTargets(x, xnext, taken, (Tensor)new[] { 1d }, 0.5, (Tensor)new[] { 1d }, w1, w2, w1, w2);

This example lives in its own file, Targets.hb, so the class is Targets.

Output (verified):

qLinear = shape [2, 2] values [1, 2, 3, 4]
aimed = shape [1, 2] values [3, 2]
aimed terminal = shape [1, 2] values [1, 2]