Returns
Discounted returns for every step of a trajectory at once, split at episode boundaries, with TD errors and GAE.
Prelude.Returns turns a reward sequence into the quantities reinforcement learning is built on: the discounted return from every step onward, the same return cut at episode boundaries, one-step temporal-difference errors and generalised advantage estimates (GAE). Every definition works on a whole trajectory at once — one triangular contraction computes all \(T\) suffix returns — so no loop over time steps is written.
Conventions: rw is the reward vector T in time order, g is the discount \(\gamma \in [0, 1)\), and done is a vector T of 0/1 flags where done[t] = 1 means step \(t\) was the last step of its episode. A batch that spans several episodes is fine: pass the flags, and returns and advantages will never cross from one episode into the next. Discount 0 is handled correctly (a myopic return is just the reward).
The value vectors v and vnext in the TD definitions are your critic’s estimates of the state at each step and of the state that followed it. The factor \((1 - \text{done})\) that stops bootstrapping from a state the episode never reached is built in.
Notation. Steps are positions counted from zero, \(t, u \in \{0, 1, \dots, T-1\}\), and \([\,\cdot\,]\) is the Iverson bracket (1 if true, 0 otherwise). Powers follow the convention \(\gamma^0 = 1\), also at \(\gamma = 0\).
open Prelude.Returns
Running totals and positions
Small helpers over a trajectory: a running sum and the upper-triangular “from here onward” mask that the discounting uses.
cumulate
cumulate(row) · row values vector T → running totals vector T
Let \(x = (x_0, \dots, x_{T-1}) \in \mathbb{R}^T\) (the argument row). Then, for \(k = 0, \dots, T-1\),
The cumulative sum.
Intuition Entry \(k\) is “everything collected so far”: a triangular mask of ones picks steps \(0..k\) and one contraction adds them, so every running total is computed at once rather than by a loop that carries an accumulator.
Use it when you need a running balance, a count of events up to each step (sum a 0/1 flag vector, as ended does to number episodes), or an empirical CDF from sorted weights. It builds a T×T mask, so memory grows as \(T^2\).
running-totalsevent-countingepisode-segmentation
def cumulate(row) ≔ einsum("t,tk->k", row, triangle(laid(ones(row)), 0))
Building blocks
turns(rw)— Let \(r \in \mathbb{R}^T\). Then \(\operatorname{turns}(r)_t = t\) for \(t = 0, \dots, T-1\). The step indices as a vector.onward(rw)— Let \(r \in \mathbb{R}^T\). Then \(\operatorname{onward}(r)_{tu} = [\,u \ge t\,]\) for \(t, u = 0, \dots, T-1\). A T×T 0/1 matrix marking pairs (t, u) with u ≥ t.
Discounting
The return from step \(t\) is
\[G_t = \sum_{u \ge t} \gamma^{\,u-t}\, r_u ,\]computed for every \(t\) as a matrix of decay factors contracted against the reward vector.
returns
returns(rw, g) · rw rewards vector T · g discount scalar → suffix returns vector T
Let \(r \in \mathbb{R}^T\) (the argument rw) be the rewards and \(\gamma \in [0, 1]\) (the argument g) the discount. Then, for \(t = 0, \dots, T-1\),
The discounted return from every step to the end of the trajectory, ignoring episode boundaries (use episodeReturns when the trajectory contains several episodes).
Intuition A reward now is worth more than the same reward later, and \(\gamma\) is the exchange rate: one step of delay multiplies its value by \(\gamma\). \(G_t\) answers “how good was it to be at step \(t\), counting everything that followed?” It is the Monte Carlo target that a value function learns to predict.
Use it when you have one complete episode and need a learning target per step: REINFORCE weights, Monte Carlo value regression, or scoring a rollout. The trajectory must end where the task ends, since rewards after \(T\) are taken as zero. For a batch that joins several episodes use episodeReturns. For long or unfinished episodes, a bootstrapped target (tdErrors, gae) has far less variance.
credit-assignmentvalue-targetspolicy-gradient
def returns(rw, g) ≔ einsum("tu,u->t", decayFrom(g, rw), rw)
returnOf
returnOf(rw, g) · rw rewards vector T · g discount scalar → return scalar
Let \(r \in \mathbb{R}^T\) (the argument rw) and \(\gamma \in [0, 1]\) (the argument g). Then
The discounted return of the whole sequence from its first step.
Intuition One number that grades a whole trajectory: the rewards weighted by \(1, \gamma, \gamma^2, \dots\) and added. It equals returns(rw, g) at step 0, but costs a single dot product instead of a T×T matrix.
Use it when you are comparing whole rollouts, such as the score of a policy candidate, an episode-level learning curve, or a net present value of a cash-flow stream (set \(\gamma = 1/(1 + \text{rate})\)). When you need a value at every step, use returns.
policy-evaluationrollout-scoringpresent-value
def returnOf(rw, g) ≔ einsum("t,t->", rw, g ^ turns(rw))
horizon
horizon(g) · g discount scalar → effective horizon scalar
Let \(\gamma \in [0, 1)\) (the argument g). Then
The effective horizon: roughly how many steps ahead a discount of \(\gamma\) looks. The divisor is not guarded, so \(\gamma = 1\) gives infinity.
Intuition The weights \(\gamma^t\) add up to \(1/(1-\gamma)\), as if the agent looked exactly that many steps ahead with full weight and then stopped. \(\gamma = 0.9\) “sees” about 10 steps, \(\gamma = 0.99\) about 100.
Use it when you are choosing a discount. Match the horizon to how long an action’s consequences take to show up in reward, and make sure trajectories (or bootstrapping) cover at least that many steps. It also tells you the scale of returns: a reward of 1 per step gives values near horizon(g).
hyperparameter-choicediscount-tuning
def horizon(g) ≔ 1 / (1 − g)
forever
forever(rw, g) · rw a reward per step scalar or tensor · g discount scalar → value same shape as rw
Let \(r\) (the argument rw) be a reward, or a tensor of rewards taken entrywise, and \(\gamma \in [0, 1)\) (the argument g). Then
The value of receiving the same reward on every step forever.
Intuition A geometric series in closed form: a constant stream \(r, r, r, \dots\) is worth \(r \cdot \operatorname{horizon}(\gamma)\). It is the value of a state that loops on itself.
| Use it when you need a reference value: a sanity bound for a critic ($$ | V | \le r_{\max}/(1-\gamma)$$), the value to bootstrap an absorbing state with a constant reward, or optimistic initialisation of a value table. For a perpetuity in finance, it is the price of a constant payment. |
value-boundsabsorbing-statesoptimistic-initialisation
def forever(rw, g) ≔ rw / (1 − g)
Building blocks
lagOf(rw)— Let \(r \in \mathbb{R}^T\). Then \(\operatorname{lagOf}(r)_{tu} = (u - t)\,[\,u \ge t\,]\) for \(t, u = 0, \dots, T-1\). The T×T matrix of lags, zero below the diagonal (masked before the power so a discount of 0 never produces 0 to a negative power).decayFrom(g, rw)— Let \(\gamma \in [0, 1]\) and \(r \in \mathbb{R}^T\), and \(L = \operatorname{lagOf}(r)\). Then \(\operatorname{decayFrom}(\gamma, r)_{tu} = \gamma^{L_{tu}}\,[\,u \ge t\,] = \gamma^{\,u-t}\,[\,u \ge t\,]\). The T×T matrix of decay factors for u ≥ t, zero elsewhere.
Episode boundaries
A return must not cross the end of an episode. The flags in done are turned into an “episode number” per step, and pairs of steps from different episodes are masked out of the decay matrix.
episodeReturns
episodeReturns(rw, g, done) · rw rewards vector T · g discount scalar · done end-of-episode flags vector T → returns vector T
Let \(r \in \mathbb{R}^T\) (the argument rw) be the rewards, \(\gamma \in [0, 1]\) (the argument g) the discount and \(d \in \{0, 1\}^T\) (the argument done) the end-of-episode flags. Let \(\epsilon_t = \sum_{u < t} d_u\) be the episode number of step \(t\). Then, for \(t = 0, \dots, T-1\),
The discounted return from every step to the end of its own episode.
Intuition Replay buffers and vectorised environments glue episodes end to end. A plain suffix sum would let the next episode’s rewards leak backwards into this one, so this function numbers the episodes and masks out every pair of steps that belong to different ones.
Use it when a batch holds several episodes, which is almost always true in practice. Pass the done flags exactly as the environment reported them. With a single episode it gives the same result as returns.
credit-assignmentbatched-trajectoriesepisode-segmentation
def episodeReturns(rw, g, done) ≔ einsum("tu,u->t", decayWithin(g, rw, done), rw)
Building blocks
ended(done)— Let \(d \in \{0, 1\}^T\). Then \(\operatorname{ended}(d)_t = \operatorname{cumulate}(d)_t - d_t = \sum_{u < t} d_u\). For each step, the number of episodes that finished strictly before it.sameRun(done)— Let \(d \in \{0, 1\}^T\) and \(\epsilon = \operatorname{ended}(d)\). Then \(\operatorname{sameRun}(d)_{tu} = [\,\epsilon_t = \epsilon_u\,]\). A T×T 0/1 matrix marking pairs of steps in the same episode.decayWithin(g, rw, done)— Let \(\gamma \in [0, 1]\), \(r \in \mathbb{R}^T\) and \(d \in \{0, 1\}^T\). Then \(\operatorname{decayWithin}(\gamma, r, d) = \operatorname{decayFrom}(\gamma, r) \odot \operatorname{sameRun}(d)\), with entries \(\gamma^{\,u-t}\,[\,u \ge t\,]\,[\,\epsilon_u = \epsilon_t\,]\). The decay matrix restricted to pairs in the same episode.
Temporal-difference errors and advantages
The one-step TD error is
\[\delta_t = r_t + \gamma\,(1 - d_t)\, V(s_{t+1}) - V(s_t),\]and GAE sums those errors decayed by \(\gamma\lambda\) within an episode:
\[\hat A_t = \sum_{u \ge t} (\gamma\lambda)^{u-t}\, \delta_u .\]\(\lambda = 0\) gives the one-step TD error and \(\lambda = 1\) the Monte Carlo advantage.
tdErrors
tdErrors(rw, v, vnext, g, done) · rw rewards vector T · v value of each state vector T · vnext value of each successor vector T · g discount scalar · done end-of-episode flags vector T → TD errors vector T
Let \(r \in \mathbb{R}^T\) (the argument rw) be the rewards, \(v \in \mathbb{R}^T\) the value of each state, \(v' \in \mathbb{R}^T\) (the argument vnext) the value of each successor, \(\gamma \in [0, 1]\) (the argument g) the discount and \(d \in \{0, 1\}^T\) (the argument done) the end-of-episode flags. Then, for \(t = 0, \dots, T-1\),
The one-step TD error for every transition. A terminal transition does not bootstrap.
Intuition The critic’s surprise: what it predicted for this state (\(v\)) against what one step of reality says (\(r + \gamma v'\)). A positive \(\delta\) means the step went better than expected. It is the basic learning signal of TD learning: nudge \(v\) toward the target, and reinforce actions with positive \(\delta\).
Use it when you train a critic (TD(0), actor-critic) or need a low-variance but biased advantage. The error depends on only one reward, so it is far less noisy than a Monte Carlo return, and it is only as good as the critic. For a tunable mix of bias and variance, feed it into gae.
value-learningadvantage-estimationbias-variance-tradeoff
def tdErrors(rw, v, vnext, g, done) ≔ rw + g · (1 − done) · vnext − v
gae
gae(rw, v, vnext, g, lam, done) · rw rewards vector T · v state values vector T · vnext successor values vector T · g discount scalar · lam GAE parameter λ scalar · done end-of-episode flags vector T → advantages vector T
Let \(r, v, v' \in \mathbb{R}^T\), \(\gamma \in [0, 1]\) and \(d \in \{0, 1\}^T\) be as in tdErrors, let \(\lambda \in [0, 1]\) (the argument lam), let \(\delta = \operatorname{tdErrors}(r, v, v', \gamma, d)\), and let \(\epsilon_t = \sum_{u < t} d_u\). Then, for \(t = 0, \dots, T-1\),
Generalised advantage estimates for every step, never crossing an episode boundary.
Intuition An advantage asks “was this action better than usual here?” One TD error answers with little noise but trusts the critic. The full return trusts nothing but is noisy. GAE averages all the n-step answers in between with weights decaying by \(\lambda\), so \(\lambda\) moves you along that bias-variance axis.
Use it when you train policy-gradient methods such as PPO, A2C or TRPO. \(\lambda \approx 0.95\) with \(\gamma \approx 0.99\) is the usual starting point. Lower \(\lambda\) if updates are noisy, and raise it if the critic is poor. Normalise the result (mean 0, std 1) per batch before using it as a weight.
advantage-estimationpolicy-gradientbias-variance-tradeoff
def gae(rw, v, vnext, g, lam, done) ≔ einsum("tu,u->t", decayWithin(g · lam, rw, done), tdErrors(rw, v, vnext, g, done))
lambdaWeight
lambdaWeight(lam, n) · lam λ scalar · n step count scalar or vector → weight same shape as n
Let \(\lambda \in [0, 1)\) (the argument lam) and \(n \in \{1, 2, \dots\}\) (or a vector of step counts, taken entrywise). Then
The weight that the \(\lambda\)-return gives to the \(n\)-step return; over \(n = 1, 2, \dots\) the weights sum to one.
Intuition The \(\lambda\)-return is not one lookahead but a blend of all of them: the 1-step return gets weight \(1-\lambda\) and each longer one a factor \(\lambda\) less. \(1/(1-\lambda)\) is the average lookahead length, just as horizon is for \(\gamma\).
Use it when you build a \(\lambda\)-return or TD(\(\lambda\)) target from explicit n-step returns, when you truncate the blend at \(N\) steps (put the leftover mass \(\lambda^{N-1}\) on the last one), or when you want to show how much a given \(\lambda\) trusts long rollouts.
value-targetsbias-variance-tradeoffeligibility-traces
def lambdaWeight(lam, n) ≔ (1 − lam) · lam^(n − 1)
Examples
The examples below live in one file, Rewards.hb, so the class is Rewards. Vectors are passed as Tensor because these formulas change rank internally.
Discounted returns from every step
Three rewards of 1 at \(\gamma = 0.5\): the return from the start is \(1 + 0.5 + 0.25 = 1.75\).
open Prelude.Returns
fn suffixReturns(rw, g) ↦ returns(rw, g)
fn discountedTotal(rw, g) ↦ returnOf(rw, g)
fn effectiveHorizon(g) ↦ horizon(g)
using MLambda.Hilbert.Runtime;
using Examples;
var rewards = (Tensor)new[] { 1d, 1d, 1d };
var all = Rewards.SuffixReturns(rewards, 0.5);
var total = Rewards.DiscountedTotal(rewards, 0.5);
var h = Rewards.EffectiveHorizon(0.9);
Output (verified):
returns = [1.75, 1.5, 1]
returnOf = 1.75
horizon = 10
Returns in a batch that holds two episodes
Four rewards of 1, where the first episode ends after step 1. Each episode’s first step is worth \(1.5\) and its last step \(1\) — the second episode’s rewards do not leak into the first.
open Prelude.Returns
fn perEpisodeReturns(rw, g, done) ↦ episodeReturns(rw, g, done)
using MLambda.Hilbert.Runtime;
using Examples;
var byEpisode = Rewards.PerEpisodeReturns(
(Tensor)new[] { 1d, 1d, 1d, 1d }, 0.5, (Tensor)new[] { 0d, 1d, 0d, 0d });
Output (verified):
episodeReturns = [1.5, 1, 1.5, 1]
Generalised advantage estimates
Two steps, the second terminal, with critic values 0.5 everywhere except a successor value of 9 after the terminal step (which must be ignored). The TD errors are \(1 + 0.9 \cdot 0.5 - 0.5 = 0.95\) and \(1 - 0.5 = 0.5\); with \(\gamma\lambda = 0.45\) the first advantage is \(0.95 + 0.45 \cdot 0.5 = 1.175\).
open Prelude.Returns
fn advantageEstimate(rw, v, vnext, g, lam, done) ↦ gae(rw, v, vnext, g, lam, done)
using MLambda.Hilbert.Runtime;
using Examples;
var advantages = Rewards.AdvantageEstimate(
(Tensor)new[] { 1d, 1d }, // rewards
(Tensor)new[] { 0.5, 0.5 }, // V(s)
(Tensor)new[] { 0.5, 9d }, // V(s') — the 9 follows a terminal step
0.9, 0.5,
(Tensor)new[] { 0d, 1d }); // done
Output (verified):
gae = [1.175, 0.5]