Trajectories
Rolling out episodes of a tensor MDP under a policy from caller-supplied uniform draws, plus inverse-CDF sampling and replay gathers.
Prelude.Trajectories simulates episodes of a finite MDP described by the tensors of Decisions — transitions p (S×A×S) and rewards r (S×A) — under a policy pol (S×A). Sampling is an inverse CDF (a cumulative sum, a comparison and a count) against uniform draws you supply, so an episode is a deterministic function of its draws: seed your generator in C# and two runs, or two executors, agree entry for entry.
Each turn uses two draws: one chooses the action, one chooses where it lands. For an episode of \(T\) turns pass draws of shape T×2 (column 0 for the action, column 1 for the transition) with values in \([0, 1)\).
The episode is recorded in a matrix tr of shape (T+1)×4 whose columns are state, action, reward, successor. Before the rollout, fill it with zeros except tr[0, 0], the start state. Turn \(t\) writes row \(t\) and copies the successor into row \(t+1\)’s state column; the last row is never turned, so a (T+1)×4 matrix holds exactly \(T\) turns. Every turn is expanded in the generated code, so a doorway that calls rollout produces a large generated C# file that takes a while to compile the first time.
Notation. All positions — states, actions, rows and columns — are counted from zero. \([\,\cdot\,]\) is the Iverson bracket (1 if true, 0 otherwise), \(e_k\) is the basis vector with 1 at position \(k\), and \(E^{(t,c)}\) is the matrix with 1 at row \(t\), column \(c\) and 0 elsewhere. \(P \in [0, 1]^{S \times A \times S}\), \(r \in \mathbb{R}^{S \times A}\) and the row-stochastic \(\pi \in [0, 1]^{S \times A}\) are the arguments p, r and pol; \(D \in [0, 1)^{T \times 2}\) is draws and \(X \in \mathbb{R}^{(T+1) \times 4}\) is tr.
open Prelude.Trajectories
Sampling a row
drawnFrom
drawnFrom(row, u) · row probabilities vector K · u uniform draw scalar → index scalar
Let \(w \in \Delta^{K-1}\) (the argument row) be a probability vector with cumulative sums \(c_k = \sum_{j=0}^{k} w_j\) for \(k = 0, \dots, K-1\), and let \(u \in [0, 1)\) be a uniform draw. Then
Samples an index from a discrete distribution by inverse CDF: the number of cumulative probabilities that are at most \(u\). The result is clamped to \(K - 1\), so a draw of exactly 1 (or a cumulative sum that rounds below 1) never indexes past the end.
Intuition Lay the probabilities end to end along \([0, 1)\) as segments of their own length; a uniform draw lands in exactly one segment, and larger probabilities catch more draws. Counting how many segment ends lie at or below \(u\) gives that segment’s index without any search or branch.
Use it when you sample a categorical outcome — an action from a policy row, a next state from a transition row, a class from predicted probabilities — reproducibly from a seeded uniform. Zero-probability entries are never picked. The row must sum to 1: any missing mass (including an all-zero row) goes to the last index through the clamp. Use pickOf when you need the result one-hot.
random-variate-generationsimulationreproducibility
def drawnFrom(row, u) ≔ min(einsum("k->", cumulate(row) ≤ u), rows(row) − 1)
pickOf
pickOf(row, u) · row probabilities vector K · u uniform draw scalar → one-hot vector K
Let \(w \in \Delta^{K-1}\) (the argument row) and \(u \in [0, 1)\). Then, for \(k = 0, \dots, K-1\),
that is \(\operatorname{pickOf}(w, u) = e_{\operatorname{drawnFrom}(w, u)}\): the sampled index as a one-hot vector.
Intuition The same draw as drawnFrom, delivered as a selector instead of a number: a vector with a single 1 that can be contracted against a table to read the sampled row or cell.
Use it when the sampled choice feeds straight into more tensor arithmetic — reading a reward or next-state row for the sampled action, or building the one-hot action ha for scoreOf or nudged. When you want to record or return the index itself, use drawnFrom.
random-variate-generationone-hot-encodingaction-selection
def pickOf(row, u) ≔ positions(row) = drawnFrom(row, u)
Reading and writing one cell
cellAt
cellAt(tr, t, c) · tr matrix T×C · t row scalar · c column scalar → value scalar
Let \(X \in \mathbb{R}^{T \times C}\) (the argument tr), \(t \in \{0, \dots, T-1\}\) and \(c \in \{0, \dots, C-1\}\). Then
The entry at row t, column c.
Intuition Indexing written as arithmetic: build a one-hot for the row and one for the column and contract the matrix against both. Every entry except the chosen one is multiplied by zero, so the index can itself be a computed value, such as the state recorded in the previous turn.
Use it when you read one field of a recorded trajectory (the state, action, reward or successor of turn \(t\)) or one uniform from the draws matrix inside a Hilbert definition. A row or column number outside the matrix gives 0 silently. To read whole rows at once use pickedRows.
table-indexingsimulation
def cellAt(tr, t, c) ≔ einsum("tc,t,c->", tr, turnIds(tr) = t, colIds(tr) = c)
cellPut
cellPut(tr, t, c, v) · tr matrix T×C · t row scalar · c column scalar · v value scalar → matrix T×C
Let \(X \in \mathbb{R}^{T \times C}\) (the argument tr), \(t \in \{0, \dots, T-1\}\), \(c \in \{0, \dots, C-1\}\) and \(v \in \mathbb{R}\). Then
Adds v to the entry at row t, column c. On a cell that is still zero, which is how rollout uses it, this writes the value.
Intuition A tensor language has no assignment, so a “write” returns a new matrix that differs from the old one by \(v\) in a single cell. The shape never changes, which is what lets a loop carry the whole trajectory while it fills it in.
Use it when you record values into a fixed-size log inside a loop, turn by turn, as rollout does. It adds rather than overwrites: write each cell once into a zero-initialised matrix, or subtract the old value first (cellPut(tr, t, c, v − cellAt(tr, t, c))) to replace it.
simulationtable-indexing
def cellPut(tr, t, c, v) ≔ tr + v · einsum("t,c->tc", turnIds(tr) = t, colIds(tr) = c)
Building blocks
turnIds(tr)— Let \(X \in \mathbb{R}^{T \times C}\). Then \(\operatorname{turnIds}(X)_t = t\) for \(t = 0, \dots, T-1\). The row indices of a matrix.colIds(tr)— Let \(X \in \mathbb{R}^{T \times C}\). Then \(\operatorname{colIds}(X)_c = c\) for \(c = 0, \dots, C-1\). The column indices of a matrix.
One turn of the world
These definitions make up a single turn of the rollout; you normally call rollout rather than any of them. In each item let \(t\) be a turn, \(\sigma = X_{t0}\) the state recorded in row \(t\), \(\alpha = \operatorname{drawnFrom}(\pi_{\sigma \cdot}, D_{t0})\) the sampled action and \(\lambda = \operatorname{drawnFrom}(P_{\sigma \alpha \cdot}, D_{t1})\) the sampled successor, where \(\pi_{\sigma \cdot}\) is row \(\sigma\) of \(\pi\) and \(P_{\sigma \alpha \cdot}\) the next-state distribution.
Building blocks
rowFor(p, hs, ha)— Let \(P \in [0, 1]^{S \times A \times S}\), \(h^{s} \in \{0, 1\}^{S}\) and \(h^{a} \in \{0, 1\}^{A}\) one-hot. Then \(\operatorname{rowFor}(P, h^{s}, h^{a})_{t'} = \sum_{s, a} P_{sat'}\, h^{s}_s\, h^{a}_a\). The next-state distribution for a one-hot state and action.payFor(r, hs, ha)— Let \(r \in \mathbb{R}^{S \times A}\) and \(h^{s}, h^{a}\) one-hot. Then \(\operatorname{payFor}(r, h^{s}, h^{a}) = \sum_{s, a} r_{sa}\, h^{s}_s\, h^{a}_a\). The reward for a one-hot state and action.hereOf(tr, t)— Let \(X \in \mathbb{R}^{(T+1) \times 4}\) and \(t\) a row. Then \(\operatorname{hereOf}(X, t) = X_{t0} = \sigma\). The state recorded in row t.hereHit(r, tr, t)— Let \(r \in \mathbb{R}^{S \times A}\), \(X\) and \(t\) as above. Then \(\operatorname{hereHit}(r, X, t)_s = [\,s = X_{t0}\,]\) for \(s = 0, \dots, S-1\), that is \(e_{\sigma}\). That state as a one-hot vector.actRow(r, pol, tr, t)— With \(r\), \(\pi\), \(X\) and \(t\) as above, \(\operatorname{actRow}(r, \pi, X, t)_a = \sum_{s} \pi_{sa}\, [\,s = \sigma\,] = \pi_{\sigma a}\). The policy’s action probabilities in that state.actHit(r, pol, draws, tr, t)— With \(D\) as above, \(\operatorname{actHit}(r, \pi, D, X, t) = \operatorname{pickOf}(\pi_{\sigma \cdot}, D_{t0}) = e_{\alpha}\). The sampled action as a one-hot vector, using draw column 0.paidAt(r, pol, draws, tr, t)— With \(r\), \(\pi\), \(D\), \(X\), \(t\), \(\sigma\) and \(\alpha\) as above, \(\operatorname{paidAt}(r, \pi, D, X, t) = \operatorname{payFor}(r, e_{\sigma}, e_{\alpha}) = r_{\sigma \alpha}\). The reward of the sampled action.landsAt(p, r, pol, draws, tr, t)— With \(P\), \(r\), \(\pi\), \(D\), \(X\), \(t\), \(\sigma\), \(\alpha\) and \(\lambda\) as above, \(\operatorname{landsAt}(P, r, \pi, D, X, t) = \operatorname{drawnFrom}\big(\operatorname{rowFor}(P, e_{\sigma}, e_{\alpha}), D_{t1}\big) = \lambda\). The sampled successor state, using draw column 1.took(r, pol, draws, tr, t)— With \(r\), \(\pi\), \(D\), \(X\), \(t\) and \(\alpha\) as above, \(\operatorname{took}(r, \pi, D, X, t) = \sum_{a=0}^{A-1} a\, (e_{\alpha})_a = \alpha\). The sampled action as a number.advanced(p, r, pol, draws, tr, t)— With \(P\), \(r\), \(\pi\), \(D\), \(X\) and \(t\) as above, \(\operatorname{advanced}(P, r, \pi, D, X, t) = X + \alpha\, E^{(t,1)} + r_{\sigma \alpha}\, E^{(t,2)} + \lambda\, E^{(t,3)} + \lambda\, E^{(t+1,0)}\), where \(\sigma\), \(\alpha\) and \(\lambda\) are all read from the incoming \(X\). The trajectory with turn t written: action, reward and successor in row t, successor in row t + 1’s state column.
The rollout
rollout
rollout(p, r, pol, draws, tr) · p transitions tensor S×A×S · r rewards matrix S×A · pol policy matrix S×A · draws uniforms matrix T×2 · tr start trajectory matrix (T+1)×4 → episode matrix (T+1)×4
Let \(P \in [0, 1]^{S \times A \times S}\), \(r \in \mathbb{R}^{S \times A}\) and \(\pi \in [0, 1]^{S \times A}\) be the MDP and policy, \(D \in [0, 1)^{T \times 2}\) (the argument draws) uniform draws, and \(X_0 \in \mathbb{R}^{(T+1) \times 4}\) (the argument tr) zero except for the start state in entry \((0, 0)\). Let \(X_{j+1} = \operatorname{advanced}(P, r, \pi, D, X_j, j)\) for \(j = 0, 1, \dots, T-1\). Then
Runs \(T\) turns from the start state and returns the filled trajectory. Rows \(0, \dots, T-1\) hold (state, action, reward, successor); row \(T\) holds only the final state. The episode is not cut at a terminal state: an absorbing terminal simply repeats with zero reward. Feed the reward column to Returns to compute returns.
Intuition The agent and the world taking turns, written down: in each row look up the state, sample an action from the policy, pay its reward, sample where it lands, and carry that state to the next row. Because all randomness comes from the draws you pass, the same draws always give the same episode.
Use it when you simulate a known or learned MDP: Monte Carlo evaluation of a policy (many rollouts with fresh draws, then average returnOf), generating experience for testing a learner, or reproducing a bug exactly from its seed. Draws must lie in \([0, 1)\) and tr needs \(T + 1\) rows. The generated code grows with \(T\), so keep episodes short. Terminal states are not flagged; build done yourself before calling episodeReturns. For the exact value without sampling use evaluated.
simulationmonte-carlo-estimationpolicy-evaluationreproducibility
def rollout(p, r, pol, draws, tr) ≔ rolled(rows(tr) − 1, p, r, pol, draws, tr)
Building blocks
rolled(k, p, r, pol, draws, tr)— Let \(k \in \mathbb{N}\), \(P\), \(r\), \(\pi\), \(D\) as above and \(X \in \mathbb{R}^{R \times 4}\). Then \(\operatorname{rolled}(0, P, r, \pi, D, X) = X\) and \(\operatorname{rolled}(k, P, r, \pi, D, X) = \operatorname{rolled}\big(k - 1, P, r, \pi, D, \operatorname{advanced}(P, r, \pi, D, X, R - 1 - k)\big)\) for \(k \ge 1\). The last k turns of the rollout (turns \(R-1-k, \dots, R-2\)), applied in order; used byrollout.
Replay
pickedRows
pickedRows(x, ids) · x data matrix n×m · ids row numbers vector k → selected rows matrix k×m
Let \(X \in \mathbb{R}^{n \times m}\) (the argument x) be the data and \(\iota \in \{0, \dots, n-1\}^{k}\) (the argument ids) row numbers. Then, for \(l = 1, \dots, k\) and \(j = 1, \dots, m\),
Gathers rows of x by index, as a one-hot contraction — for example to draw a minibatch from a replay buffer with indices sampled in C#.
Intuition A gather as matrix multiplication: one-hot each requested index against the data’s row axis and multiply, so each output row is a copy of the row it names. Indices may repeat, which is exactly sampling with replacement.
Use it when you draw a minibatch from a replay buffer, reorder data by a shuffled permutation, or select the transitions of chosen episodes. Sample the indices in C# (uniformly, or by priority). An index outside \(0..n-1\) yields a row of zeros rather than an error. It builds a k×n one-hot, so memory grows with batch size times buffer size.
experience-replayminibatchingtable-indexing
def pickedRows(x, ids) ≔ einsum("ij,ki->kj", x, marked(positions(einsum("ij->i", ones(x))), ids))
Examples
The examples below live in one file, Episodes.hb, so the class is Episodes. Everything is passed as Tensor.
Roll out an episode
The MDP of the Decisions examples: two states, action 0 stays, action 1 moves, and staying in state 1 pays 1. Under the optimal policy (move from state 0, stay in state 1) and starting in state 0, three turns give: move to 1 with reward 0, then stay twice with reward 1.
open Prelude.Trajectories
fn episodeOf(p, r, pol, draws, tr) ↦ rollout(p, r, pol, draws, tr)
using MLambda.Hilbert.Runtime;
using Examples;
var p = new double[2, 2, 2];
p[0, 0, 0] = 1; p[0, 1, 1] = 1; p[1, 0, 1] = 1; p[1, 1, 0] = 1;
var r = (Tensor)new double[,] { { 0, 0 }, { 1, 0 } };
var policy = (Tensor)new double[,] { { 0, 1 }, { 1, 0 } };
var draws = (Tensor)new double[,] { { 0.5, 0.5 }, { 0.5, 0.5 }, { 0.5, 0.5 } }; // T = 3
var start = (Tensor)new double[4, 4]; // T + 1 rows, start state 0
var episode = Episodes.EpisodeOf((Tensor)p, r, policy, draws, start);
Output (verified; rows are state, action, reward, successor):
rollout = shape [4, 4] values [0, 1, 0, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 0, 0]
Sample from a discrete distribution
With probabilities \([0.2, 0.5, 0.3]\) the cumulative sums are \([0.2, 0.7, 1.0]\); a draw of 0.6 passes one of them, so index 1 is sampled.
open Prelude.Trajectories
fn sampledIndex(row, u) ↦ drawnFrom(row, u)
using MLambda.Hilbert.Runtime;
using Examples;
var index = Episodes.SampledIndex((Tensor)new[] { 0.2, 0.5, 0.3 }, 0.6);
Output (verified):
drawnFrom = 1
Gather a minibatch from a replay buffer
Rows 2 and 0 of a three-row buffer.
open Prelude.Trajectories
fn replayRows(x, ids) ↦ pickedRows(x, ids)
using MLambda.Hilbert.Runtime;
using Examples;
var buffer = (Tensor)new double[,] { { 1, 2 }, { 3, 4 }, { 5, 6 } };
var batch = Episodes.ReplayRows(buffer, (Tensor)new[] { 2d, 0d });
Output (verified):
pickedRows = shape [2, 2] values [5, 6, 1, 2]