What it actually does
Play the episode to the end, add up what you actually received, and use that as the value of every state-action pair you passed through. Average over enough episodes and the averages converge on the truth, because they are averages of exactly the thing being estimated. Two consequences follow and they are both unusual. The estimate is unbiased and it never assumes the state is Markov — it does not predict the future from the current state, it measures the future that occurred, so a state representation that hides something still produces correct values for the process as experienced. Nothing else in this category can say that. The costs are equally plain. You learn nothing until the episode terminates, which rules out continuing tasks entirely, and the return of a long episode is the sum of many random quantities, so the variance is large and shrinks only as the square root of the episode count. The subject that actually dominates any implementation, though, is neither: it is that a greedy policy visits a vanishing fraction of the state-action pairs, so most of the table is never estimated at all and the improvement step chooses between numbers that were never measured.
- Every episode terminates, since the return is undefined otherwise and the method has no substitute for it
- Every state-action pair is visited infinitely often, which exploration must actually deliver rather than be assumed to
- The environment is stationary across the episodes being averaged, or the mean describes a mixture of environments
- For off-policy estimates, the behaviour policy has positive probability on everything the target policy would do — positivity, which is an identification condition rather than a data-volume concern
- Notably NOT assumed: that the state is Markov. The realized return does not depend on the state summarizing the history, which is this method's distinctive advantage
Fixed-point condition
Read the two halves separately, because they have different characters and the entry is usually taught as though they had one. The evaluation half is not a fixed-point equation at all: Q is an ordinary expectation and the estimator is a sample mean of realized returns, which makes it unbiased and makes its error fall as one over the square root of the number of visits. Nothing in it refers to Q on the right-hand side, which is exactly what "no bootstrapping" means and exactly why the Markov property is not required — the return that was received does not care whether the state summarized the history. The control half is a fixed-point iteration in the usual sense: evaluate, act greedily, repeat, with the same policy improvement theorem doing the work as in dynamic programming. Putting them together produces the method's defining tension. Greedy improvement makes the policy deterministic, a deterministic policy visits a shrinking set of state-action pairs, and pairs that are never visited have no sample mean to compare — so the improvement step ranks numbers that were never measured. Every variant here is a different answer to that: exploring starts, which assumes a privilege real environments do not grant; epsilon-soft policies, which converge to the best epsilon-soft policy rather than to the optimal one; and off-policy importance sampling, which learns about the greedy policy from a policy that explores, at the cost of weights whose variance compounds along the trajectory.
- the realized return — the actual discounted sum of rewards from t to the end of the episode, not an estimate of it
- an expectation estimated by a plain sample mean, with no bootstrapping and therefore no bias
- the terminal step; the method is defined only for episodes that reach one
- the greedy improvement, which is where the exploration problem enters — it can only rank pairs that were visited
How it is minimized
Generalized policy iteration where evaluation is a sample mean of realized returns — incremental, episode by episode, with no model and no bootstrapping
- visit count, which makes the update an exact running mean rather than an exponential one
- the step size; decaying as one over n is what makes this a sample average and what satisfies Robbins-Monro
- the error against the realized return — no bootstrapped term appears anywhere in it
- the common substitute for 1/N, which tracks a changing policy at the cost of never converging
The incremental form is algebraically identical to storing every return and averaging, and it is worth writing that way because the choice of step size is then visibly the choice of what is being estimated. One over N is an exact running mean: it weights every episode equally and converges to the true expectation, which is what the theory assumes. A fixed alpha is an exponential moving average, weighted toward recent episodes, which does not converge to anything — and that is frequently the correct choice, because the policy is changing and the returns from twenty policies ago are estimates of a different quantity. Stating which of the two is in use decides whether the run has a convergence guarantee or a tracking behaviour, and implementations often do not say. Three further choices sit alongside it. First-visit against every-visit: first-visit averages one return per episode per pair and is unbiased; every-visit uses all of them, is biased at finite samples because the returns within an episode are correlated, and is consistent and usually lower variance — the bias is real and rarely large enough to matter. On-policy against off-policy: an epsilon-soft policy keeps exploring and converges to the best epsilon-soft policy, which is genuinely not the optimal policy, and the gap is proportional to epsilon. Off-policy learning corrects for the difference with importance weights, and the choice between ordinary and weighted importance sampling is the clearest bias-variance trade in this reference — ordinary is unbiased with unbounded variance, weighted is biased and consistent with dramatically lower variance, and weighted is what anyone actually uses.
epsilon
Exploration rate of the behaviour policy. On-policy, it also bounds how far from optimal the converged policy is, so it must be annealed rather than tuned (1.0 annealed toward 0.01)
step size (1/N or fixed alpha)
One over N is an exact sample mean and converges; a fixed alpha tracks a moving target and does not. Which one is in use decides whether guarantees apply (1/N, or 0.01 to 0.1 fixed)
first-visit or every-visit
First-visit is unbiased, every-visit is biased at finite samples, consistent, and usually lower variance. Rarely the thing that decides a result (either)
discount (gamma)
Sets the horizon, and here it also caps variance: returns from far in the future are down-weighted, so a smaller gamma is a variance reduction as well as a shorter horizon (0.9 to 1.0 for episodic tasks)
importance-sampling form
Ordinary is unbiased with unbounded variance; weighted is biased, consistent and usable. The choice is the whole off-policy design (weighted, in practice)
episodes per improvement
How much evaluation before acting greedily again. One episode is the usual choice and it works, which is the practical form of generalized policy iteration (1 to 100)
On-policy Monte Carlo control converges to the optimal policy under GLIE — greedy in the limit with infinite exploration — which means every pair is visited infinitely often and the policy becomes greedy asymptotically, usually by annealing epsilon as one over the episode count. Hold epsilon fixed instead, as almost every implementation does, and the method converges to the optimal EPSILON-SOFT policy, which is a different and worse policy; the gap is small but it is real and it is systematically unreported. The failures that actually stop a run are elsewhere. Exploration starvation is the first: a greedy policy visits a vanishing fraction of the pairs, and an unvisited pair has no estimate at all, so the improvement step is ranking numbers that were never measured — the symptom is a policy that locks in early and never revisits a decision. Variance is the second, and it grows with episode length because the return is a sum over the whole remaining trajectory, so long-horizon problems need many more episodes than their state count suggests. Off-policy makes this sharper: importance weights are products along the trajectory, so their variance grows geometrically, and ordinary importance sampling has infinite variance in ordinary circumstances rather than pathological ones. And positivity binds absolutely — a pair the behaviour policy never takes has no importance-weighted estimate at any sample size, which is not slow convergence but an unidentified quantity.
Per episode: O(T) to compute the returns by one backward pass, and O(T) updates. Memory is O(|S|·|A|) for the table plus visit counts, with no need to store trajectories once returns are accumulated incrementally. The binding cost is episodes rather than arithmetic — the standard error falls as one over the square root of visits, so an extra digit of precision costs a hundred times the episodes, and the constant scales with episode length through the variance of the return.
How, where, and why it is used
Time-Series Forecasting
Not applicableThere is no action whose consequences the return could measure, so there is nothing to average over and nothing to improve; a forecast does not alter the trajectory it is predicting.
Anomaly Detection
Not applicableDetection scores an observation without acting, so no episode, no return and no policy exist; a high-variance value estimate reports an under-sampled state-action pair, never an unusual observation.
Optimization
Primary fitOptimize a policy by sampling complete episodes and averaging realized returns, with no model of the dynamics and no bootstrapped estimate anywhere in the update. The improvement step is ordinary generalized policy iteration; what is unusual is that the evaluation half is an unbiased sample mean.
- Simulation optimization, where episodes are cheap to generate and the dynamics are implicit in the simulator rather than written down
- The bias-variance trade in its clearest form: unbiased high-variance Monte Carlo against biased low-variance bootstrapping
- Ordinary against weighted importance sampling, which is the same trade again with the roles of the two estimators swapped
- Exploration as a first-class design problem, visible here because nothing else is in the way of it
Worth studying carefully despite rarely being the method of choice, because it isolates three things that every later algorithm entangles. It shows what bootstrapping actually buys and costs — comparing this against temporal-difference learning on the same problem is the cleanest demonstration of bias against variance available, and the fact that Monte Carlo needs no Markov assumption while temporal-difference learning does is the part usually left out. It shows that importance sampling has the same structure and the opposite sign: the unbiased estimator is the unusable one, and weighted importance sampling is preferred precisely because a little bias buys a great deal of variance. And it makes exploration impossible to ignore, since a greedy policy simply never visits most pairs and there is no function approximator to paper over the gap. The practical verdict is that it is rarely the fastest method and is often the most trustworthy one when a state representation is suspect.
- Confirm every episode terminates; the method is undefined for continuing tasks and an infinite episode is a hang rather than a slow run
- State which step size is in use: one over N converges to the expectation, a fixed alpha tracks a moving policy and does not converge
- Anneal epsilon rather than fixing it, or accept convergence to the best epsilon-soft policy instead of the optimal one
- Prefer weighted importance sampling off-policy, and report the effective sample size of the weights alongside any estimate
Realized return under the learned policy against a baseline, with standard errors rather than point estimates — the variance is the defining property of the method and reporting a mean without it hides the thing that matters. Track state-action visitation coverage alongside it, since a policy that looks good and visited a tenth of the table is reporting the coverage rather than the policy.
- A fixed epsilon left in place, which converges to a policy that is knowably suboptimal by a margin proportional to it
- Unvisited pairs ranked by the improvement step, which compares numbers that were never measured
- Ordinary importance sampling off-policy, whose variance is unbounded in ordinary rather than pathological cases
- Reporting a mean return without its standard error, when the variance is the property that distinguishes this method
Causal Inference & Uplift
ViableOff-policy Monte Carlo evaluation and inverse propensity weighting are the same estimator. A logged trajectory is reweighted by the ratio of target-policy to behaviour-policy probabilities at every step, and the weighted average of realized returns estimates the value of a policy that was never run.
- Off-policy policy evaluation from logged decisions, which is sequential IPTW under a different name
- Weighted against ordinary importance sampling as the bias-variance choice the causal literature calls stabilized weights
- Positivity as a hard identification condition rather than a data-sufficiency concern
- Dynamic treatment regimes, where the sequential structure and the estimator are both exactly this
Included because the correspondence is exact rather than suggestive, and because importing the causal literature's care about it is the fastest available upgrade to an off-policy evaluation. Every property transfers. Ordinary importance sampling is the Horvitz-Thompson estimator: unbiased, and with variance that is unbounded whenever a behaviour probability can be small. Weighted importance sampling is the Hajek estimator: biased, consistent, and the one everyone actually uses, for the same reason stabilized weights are standard in the causal literature. Positivity has identical force — a state-action pair the logging policy never took is unidentified at any sample size, and no estimator recovers it. And the horizon makes it worse than the single-timepoint case: weights are products along the trajectory, so their variance compounds geometrically and an estimate over more than a handful of steps is usually dominated by two or three trajectories. The honest limit is that reinforcement learning assumes no unobserved confounding, which is the problem causal inference exists to address, so this is a shared estimator rather than a shared solution.
- Log the behaviour probability with every decision, or the weights cannot be formed and nothing downstream is identified
- Check positivity per state-action pair and report violations as unidentified rather than as small numbers
- Report effective sample size for the weights; it collapses fast with horizon and is the honest measure of how much data an estimate used
- Prefer weighted importance sampling, and consider doubly robust variants beyond a few steps
Interval estimates rather than point estimates, since the variance is usually the finding. Where any randomization exists in the logging policy, exploit it — deliberate exploration in the logs is worth more than any estimator, and a small amount of it makes the difference between an identified quantity and a guess.
- Unlogged propensities, which make every later off-policy question unanswerable
- Ordinary importance sampling over a long horizon, where a handful of trajectories dominate the estimate
- Positivity violations reported as small weights rather than as unidentified quantities
- Assuming away confounding because the reinforcement-learning framing does not mention it
Control & Operations
AdaptedApplicable wherever a simulator produces complete episodes: run the simulator, average realized returns, improve the policy. The method needs no model in its own right because the simulator is the model, which is a distinction worth keeping straight.
- Simulation-based policy search where a high-fidelity simulator exists and its dynamics are not available in closed form
- Episodic operations problems — a shift, a production run, a delivery route — that terminate naturally
- Settings where the state representation is known to be incomplete and an unbiased estimate is preferred to a faster biased one
- Evaluating a candidate policy from historical episodes before it is allowed to act
A real but secondary fit, and the boundary is specific. Where the dynamics are documented, dynamic programming solves the problem exactly and no sampling is warranted — that check comes first. Where they are not but a simulator exists, this is a reasonable and unusually trustworthy choice, because the estimate is unbiased and does not assume the operational state is Markov, which operational states frequently are not: a queue length omits the maintenance schedule, a stock level omits the supplier's backlog. Temporal-difference methods are biased under exactly that aliasing while Monte Carlo is not. The cost is episodes, and it is steep: variance grows with episode length, so a long production horizon needs far more simulator runs than the state count suggests, and a slow simulator makes this the wrong tool regardless of its statistical merits.
- Check for a known kernel first; where the dynamics are documented this is the wrong method
- Confirm episodes terminate, and cap them explicitly if the simulator can produce a run that does not
- Prefer this over bootstrapping where the state is known to alias, since the bias there is real and this method does not have it
- Budget episodes from the return variance rather than from the state count, because the horizon drives it
Simulated return with confidence intervals, then a sim-to-real gap estimate before anything acts. Compare against the incumbent policy on the same episodes rather than on fresh ones, which removes a great deal of the variance that would otherwise obscure the comparison.
- Sampling a problem whose dynamics were already available in closed form
- Non-terminating episodes, which are a hang rather than a slow run
- Budgeting episodes from state count, when the return variance and the horizon are what set the requirement
- A simulator whose gap to reality exceeds the gap between the candidate policies
Recommendation & Ranking
AdaptedA session is an episode and the return is session-level engagement, so realized returns are directly available from logs. The natural use is off-policy evaluation: reweight logged sessions by the ratio of the candidate policy to the policy that served them, and estimate what the candidate would have achieved.
- Off-policy evaluation of a candidate ranker from logged sessions before any experiment
- Session-level objectives that a next-click model cannot express at all
- Comparing candidate policies on the same logged sessions, which removes most of the variance from the comparison
- Exploration budgeting, where a small amount of deliberate randomization makes later evaluation identified
The honest fit is as an evaluation method rather than a learning one, and that is where it earns its place. Learning a ranker this way is impractical: the action space is the catalogue, and Monte Carlo needs every state-action pair visited many times, which is hopeless for anything but a tiny candidate set. Evaluating one is entirely practical and is standard industry work, because sessions terminate, returns are observable, and logged propensities — where anyone thought to record them — make weighted importance sampling available. The failure modes are the domain's usual ones sharpened by the horizon. Weights are products over a session, so a long session contributes an estimate dominated by a few trajectories, and the effective sample size is frequently a small fraction of the logged volume. Unlogged propensities make the whole exercise impossible after the fact, which is the most common and most expensive mistake in this area, since it cannot be repaired retrospectively.
- Log the serving policy's probability with every impression; it cannot be recovered later and everything depends on it
- Truncate or cap weights explicitly and report the cap, since an uncapped product over a long session is unusable
- Evaluate candidates on the same logged sessions, which cancels much of the variance in a comparison
- Reserve a small randomized traffic slice; deliberate exploration is what makes later evaluation identified rather than optimistic
Weighted importance-sampled return with an effective-sample-size figure attached, then an online test. The offline number is a filter that rejects bad candidates, not evidence that a good one will win — and reporting it without the effective sample size overstates it considerably.
- Unlogged propensities, which cannot be reconstructed and make the estimate unidentified
- Uncapped importance weights over a long session, where a handful of sessions dominate
- Treating an offline estimate as a decision when its interval spans the incumbent
- Attempting to learn rather than evaluate, where the catalogue-sized action space makes coverage hopeless
What it costs to run
Cheap per episode — one backward pass over the trajectory and one update per visited pair — and expensive in episodes, which is the only cost that matters. The standard error falls as one over the square root of visits, so an extra digit of precision costs a hundredfold more episodes, and the constant grows with episode length because the return is a sum over the whole remaining trajectory. A simulator is effectively mandatory. Illustrative, not a measured benchmark.
A table lookup and an argmax over the action row: microseconds, and completely inspectable. Every action has a number and, unusually, that number has a standard error attached, since it is a sample mean with a visit count — which makes "how confident is this recommendation" a question with an actual answer here.
Continuous while a simulator is available. Against a real system, the visit counts are the thing to watch rather than the calendar: a changed environment invalidates the averages, and because the estimates are means over the whole history they are slow to forget — which is the case for a fixed step size rather than one over N, and it should be a deliberate decision rather than a default.
- Visit counts per state-action pair, which is the direct measure of how much of the table is actually estimated rather than guessed
- Return variance per state, since a pair with a high-variance return needs far more episodes than its neighbours and the mean alone hides that
- Effective sample size of importance weights wherever off-policy estimates are used; it collapses with horizon and is the honest denominator
- Realized return against the table's predicted value, which detects a moved environment directly
- Fraction of episodes that reach termination, because a capped episode contributes a truncated return that biases everything it touches
- A fixed epsilon converges to the optimal epsilon-soft policy, not the optimal policy. The gap is proportional to epsilon, real, and almost never reported
- One over N and a fixed alpha are different estimators: the first converges to the expectation, the second tracks and never converges. Implementations frequently do not say which is in use
- An unvisited state-action pair has no estimate at all, so the improvement step ranks numbers that were never measured — initialization silently becomes policy
- Episodes must terminate. A capped episode contributes a truncated return, which biases the mean toward whatever the cap implies, and the bias is invisible in the value table
- Ordinary importance sampling has unbounded variance in ordinary circumstances. Weighted importance sampling is biased and is the one to use
- Behaviour probabilities must be logged at decision time; they cannot be reconstructed afterwards, and without them no off-policy estimate is identified
- Standard errors are available here almost for free, since every entry is a sample mean with a count. Reporting the mean alone throws away the method's one statistical advantage
Pros and cons, in context
The value estimate is unbiased, because it averages realized returns rather than bootstrapped ones
The error is pure variance and falls as one over the square root of visits, so more episodes always help and there is no systematic error waiting underneath. Nothing that bootstraps can say this
It does not require the state to be Markov
The return that occurred does not care whether the state summarized the history, so a representation that aliases still yields correct values for the process as experienced. This is the method's strongest and least advertised property, and it is exactly where temporal-difference learning is biased
No model of the environment is needed, and no bootstrapped estimate either
The whole method is a sample mean with a policy attached, which makes it easy to reason about and nearly impossible to get subtly wrong in the way bootstrapping methods are
Every entry carries a visit count, so confidence is available for free
A sample mean with a count has a standard error, so "how well do we know this action" is answerable per state. Most methods here offer nothing comparable, and it is a natural exploration signal
Individual state values can be estimated independently of the rest
No backup chain means the value of one state does not depend on estimating its neighbours, so a subset of interest can be evaluated without solving the whole problem — genuinely useful when only a few states matter
Nothing is learned until the episode ends, so continuing tasks are excluded entirely
Not a slow update but no update: the return does not exist until termination. Capping a long episode substitutes a truncated return, which introduces exactly the bias the method was chosen to avoid
Variance grows with episode length, and it is the binding cost
The return is a sum over the whole remaining trajectory, so a long-horizon problem needs far more episodes than its state count suggests. This is the practical reason temporal-difference methods displaced it
A greedy policy visits a vanishing fraction of the state-action pairs
Unvisited pairs have no estimate, so the improvement step compares numbers that were never measured and initialization silently becomes policy. Every variant of the method is a different answer to this
With a fixed epsilon it converges to the optimal epsilon-soft policy, not the optimal one
A knowable, unreported suboptimality proportional to epsilon. Annealing fixes it in theory and is skipped in practice more often than not
Off-policy importance weights are products along the trajectory, so their variance compounds geometrically
Ordinary importance sampling has unbounded variance in ordinary cases, and even weighted importance sampling sees its effective sample size collapse over a long horizon
Sample efficiency is poor compared with bootstrapping methods on the same problem
One over the square root of visits, with a large constant. The trade is real — it buys unbiasedness and Markov independence — but on a well-specified Markov problem temporal-difference learning simply wins
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""Monte Carlo control, transcribed from the objective.
No library. Everything is a sample mean:
Q(s,a) = E[ G_t | S_t = s, A_t = a ], G_t = sum_k gamma^k R_{t+k+1}
Four things to read for.
1. No bootstrapping appears anywhere. compute_returns() sums rewards
that actually occurred, and nothing on the right-hand side of an
update refers to Q. That is what makes the estimate unbiased, and it
is also why the Markov property is never needed -- see
aliasing_demo() below, which is the clearest statement of this
method's real advantage.
2. The exploration problem is the subject, not a footnote. A greedy
policy visits a vanishing fraction of the pairs, so most of the table
is never estimated and the improvement step ranks numbers that were
never measured. exploring_starts_control() and
epsilon_soft_control() are the two classical answers, and the second
converges to the optimal EPSILON-SOFT policy rather than the optimal
one -- a real, knowable, usually unreported gap.
3. ordinary_importance_sampling() and weighted_importance_sampling()
are the bias-variance trade with the usual roles reversed: the
unbiased estimator is the unusable one.
4. Every entry carries a visit count, so a standard error is available
for free. Almost no other method here offers that.
"""
import math
import random
Transition = tuple[int, int, float] # (state, action, reward)
Episode = list[Transition]
def compute_returns(episode: Episode, gamma: float) -> list[float]:
"""G_t for every step, by one backward pass.
Written backwards because it is the only way that is linear: G_t =
R_{t+1} + gamma * G_{t+1}, so walking from the end reuses the work.
Computing each G_t forwards would be quadratic in the episode
length, which is a real cost on the long episodes this method
already struggles with.
"""
returns = [0.0] * len(episode)
running = 0.0
for index in reversed(range(len(episode))):
running = episode[index][2] + gamma * running
returns[index] = running
return returns
class ActionValues:
"""Q and its visit counts, as a running mean.
The count is not bookkeeping. Q(s,a) += (G - Q)/N is algebraically
identical to averaging every stored return, so the estimator IS a
sample mean -- which is what makes it unbiased and what gives every
entry a standard error.
Substituting a fixed step size turns this into an exponential moving
average, which tracks a changing policy and converges to nothing.
Both are defensible; which one is in use decides whether the
convergence theory applies, and implementations routinely fail to
say.
"""
def __init__(self, n_states: int, n_actions: int, initial: float = 0.0) -> None:
self.q = [[initial] * n_actions for _ in range(n_states)]
self.counts = [[0] * n_actions for _ in range(n_states)]
# Sum of squared returns, for the standard error. Costs one
# multiply per update and is almost never kept.
self.sum_squares = [[0.0] * n_actions for _ in range(n_states)]
def update(self, state: int, action: int, target: float) -> None:
self.counts[state][action] += 1
count = self.counts[state][action]
error = target - self.q[state][action]
self.q[state][action] += error / count
self.sum_squares[state][action] += target * target
def standard_error(self, state: int, action: int) -> float:
"""How well this entry is actually known.
A sample mean with a count has a standard error, so "how
confident is this recommendation" has an answer here that most
methods in this category cannot provide. It is also a natural
exploration signal: visit the pairs whose estimates are worst.
"""
count = self.counts[state][action]
if count < 2:
return float("inf")
mean = self.q[state][action]
variance = max(self.sum_squares[state][action] / count - mean * mean, 0.0)
return math.sqrt(variance / count)
def greedy_action(self, state: int) -> int:
row = self.q[state]
best = max(row)
return min(index for index, value in enumerate(row) if value == best)
def coverage(self) -> float:
"""Fraction of the table that has any estimate at all.
The number that decides whether a learned policy means
anything. A policy derived from a table that is a tenth visited
is reporting its initialization for the other nine tenths.
"""
visited = sum(1 for row in self.counts for count in row if count > 0)
total = len(self.counts) * len(self.counts[0])
return visited / total
def first_visit_updates(episode: Episode, returns: list[float]) -> list[tuple[int, int, float]]:
"""One return per (state, action) pair per episode.
Unbiased, because the returns being averaged are independent across
episodes. Every-visit uses all of them, which makes returns within
an episode correlated and the estimator biased at finite samples --
consistent, usually lower variance, and the bias is rarely the thing
that decides a result.
"""
seen: set[tuple[int, int]] = set()
updates = []
for index, (state, action, _) in enumerate(episode):
key = (state, action)
if key in seen:
continue
seen.add(key)
updates.append((state, action, returns[index]))
return updates
def every_visit_updates(episode: Episode, returns: list[float]) -> list[tuple[int, int, float]]:
return [(state, action, returns[index]) for index, (state, action, _) in enumerate(episode)]
def epsilon_greedy_action(values: ActionValues, state: int, epsilon: float, rng: random.Random) -> int:
if rng.random() < epsilon:
return rng.randrange(len(values.q[state]))
return values.greedy_action(state)
def run_episode(
step,
start_state: int,
choose_action,
max_steps: int = 10_000,
) -> Episode:
"""Generate one episode.
The step cap is not a safeguard to be quietly hit. A capped episode
contributes a TRUNCATED return, which biases every mean it touches
toward whatever the cap implies -- and that bias is invisible in the
value table. If this cap fires regularly, the task is not episodic
and this method does not apply to it.
"""
episode: Episode = []
state = start_state
for _ in range(max_steps):
action = choose_action(state)
next_state, reward, done = step(state, action)
episode.append((state, action, reward))
if done:
return episode
state = next_state
raise RuntimeError(
"episode hit the step cap; a truncated return biases every mean it enters, "
"so this is a modelling problem rather than a limit to raise"
)
def epsilon_soft_control(
step, start_states: list[int], n_states: int, n_actions: int,
gamma: float, episodes: int, epsilon: float, seed: int = 0,
):
"""On-policy control with a persistently exploring policy.
The honest caveat lives in the return value. With epsilon FIXED this
converges to the optimal epsilon-soft policy, which is not the
optimal policy: the agent keeps taking random actions forever, and
the values it converges to reflect that. The gap is proportional to
epsilon, it is real, and it is almost never reported.
Annealing epsilon toward zero as episodes accumulate -- the GLIE
condition -- recovers the optimal policy. It is one line and it is
skipped more often than not.
"""
rng = random.Random(seed)
values = ActionValues(n_states, n_actions)
for _ in range(episodes):
start = rng.choice(start_states)
episode = run_episode(
step, start, lambda state: epsilon_greedy_action(values, state, epsilon, rng)
)
returns = compute_returns(episode, gamma)
for state, action, target in first_visit_updates(episode, returns):
values.update(state, action, target)
return {
"values": values,
"policy": [values.greedy_action(state) for state in range(n_states)],
"coverage": values.coverage(),
"note": "fixed epsilon converges to the optimal epsilon-soft policy, not the optimal one",
}
def exploring_starts_control(
step, n_states: int, n_actions: int, gamma: float, episodes: int, seed: int = 0
):
"""Control with every (state, action) pair as a possible start.
This solves the exploration problem by assumption: if every pair can
begin an episode, every pair is visited and the improvement step
always has something to compare. It also assumes a privilege real
environments almost never grant -- the ability to reset into an
arbitrary state with an arbitrary first action.
Kept because it is the cleanest way to see that the exploration
problem is separable from the learning one. Remove it and everything
below becomes about coverage rather than about estimation.
"""
rng = random.Random(seed)
values = ActionValues(n_states, n_actions)
for _ in range(episodes):
start_state = rng.randrange(n_states)
first_action = rng.randrange(n_actions)
chosen = {"first": True}
def choose(state: int) -> int:
if chosen["first"]:
chosen["first"] = False
return first_action
return values.greedy_action(state)
episode = run_episode(step, start_state, choose)
returns = compute_returns(episode, gamma)
for state, action, target in first_visit_updates(episode, returns):
values.update(state, action, target)
return {
"values": values,
"policy": [values.greedy_action(state) for state in range(n_states)],
"coverage": values.coverage(),
}
def ordinary_importance_sampling(
episodes: list[Episode], behaviour_probs: list[list[float]],
target_probs: list[list[float]], gamma: float,
) -> dict[str, float]:
"""The unbiased off-policy estimator, and the unusable one.
V = (1/n) * sum_i rho_i * G_i, rho_i = prod_t pi(a|s)/b(a|s)
Unbiased, and with variance that is unbounded rather than merely
large: the weight is a PRODUCT along the trajectory, so one step
where the behaviour policy was unlikely to do what the target policy
wants inflates a whole episode's contribution. Over a long horizon
the estimate is routinely dominated by two or three trajectories.
This is the Horvitz-Thompson estimator. The causal-inference
literature met the same problem first and reached the same
conclusion.
"""
weighted = []
for episode, behaviour, target in zip(episodes, behaviour_probs, target_probs):
ratio = 1.0
for step_index in range(len(episode)):
if behaviour[step_index] <= 0.0:
# Positivity violation: this trajectory is not
# identified under the target policy, at any sample
# size. Not a small number -- an absent one.
ratio = float("nan")
break
ratio *= target[step_index] / behaviour[step_index]
if ratio != ratio: # NaN check without importing math.isnan
continue
weighted.append(ratio * compute_returns(episode, gamma)[0])
if not weighted:
return {"estimate": float("nan"), "effective_sample_size": 0.0}
mean = sum(weighted) / len(weighted)
return {"estimate": mean, "used_episodes": float(len(weighted))}
def weighted_importance_sampling(
episodes: list[Episode], behaviour_probs: list[list[float]],
target_probs: list[list[float]], gamma: float,
) -> dict[str, float]:
"""The biased, consistent, usable one.
V = sum_i rho_i G_i / sum_i rho_i
Normalizing by the weights rather than by the count bounds the
estimate inside the range of observed returns, which removes the
unbounded variance at the cost of a bias that vanishes as the sample
grows. This is the Hajek estimator, and it is what anyone actually
uses -- the same conclusion the causal literature reaches about
stabilized weights.
The effective sample size is reported because it is the honest
denominator: a thousand logged episodes whose weights concentrate on
three of them is an estimate from three episodes.
"""
ratios = []
returns = []
for episode, behaviour, target in zip(episodes, behaviour_probs, target_probs):
ratio = 1.0
identified = True
for step_index in range(len(episode)):
if behaviour[step_index] <= 0.0:
identified = False
break
ratio *= target[step_index] / behaviour[step_index]
if not identified:
continue
ratios.append(ratio)
returns.append(compute_returns(episode, gamma)[0])
total = sum(ratios)
if total <= 0.0:
return {"estimate": float("nan"), "effective_sample_size": 0.0}
estimate = sum(r * g for r, g in zip(ratios, returns)) / total
# Kish's effective sample size: (sum w)^2 / sum w^2.
effective = (total * total) / sum(r * r for r in ratios)
return {
"estimate": estimate,
"effective_sample_size": effective,
"logged_episodes": float(len(ratios)),
}
def aliasing_demo(gamma: float = 0.9, episodes: int = 20_000, seed: int = 0):
"""Why this method survives a state that is not Markov.
Two underlying states are shown to the agent as one. Monte Carlo
averages the returns that ACTUALLY OCCURRED from that observation,
so it converges to the correct expected return for the process as
experienced -- which is the right answer to the question the agent
can actually ask.
A bootstrapping method updates toward its own estimate at the next
observation, and that estimate belongs to the aliased state rather
than to the real one, so it converges to something that is not the
value of anything. This is the single strongest argument for Monte
Carlo and it is usually omitted.
"""
rng = random.Random(seed)
# Observation 0 is really two states: one that pays 1, one that
# pays -1, chosen with equal probability. The true expected return
# from the observation is 0.
def step(state: int, action: int):
if state == 0:
hidden = rng.choice([1, -1])
return 1, float(hidden), True
return 1, 0.0, True
values = ActionValues(2, 1)
for _ in range(episodes):
episode = run_episode(step, 0, lambda state: 0)
returns = compute_returns(episode, gamma)
for state, action, target in first_visit_updates(episode, returns):
values.update(state, action, target)
return {
"estimate": values.q[0][0],
"standard_error": values.standard_error(0, 0),
"true_expected_return": 0.0,
"note": "unbiased despite the observation aliasing two different states",
}
Per episode: one O(T) backward pass for the returns and O(T) updates, with no trajectory storage once returns are accumulated incrementally. Illustrative, not a measured benchmark: the cost that decides everything is episodes rather than arithmetic — the standard error falls as one over the square root of visits, so an extra digit costs a hundredfold more episodes, and the constant scales with episode length through the variance of the return.
How this language actually wants it written.
"""Monte Carlo control with the estimator choices made explicit.
Four things the literal version decided silently:
* the step size, where 1/N is a sample mean that converges and a
fixed alpha is a moving average that does not;
* the exploration schedule, where a fixed epsilon converges to the
optimal epsilon-soft policy rather than the optimal one;
* what a truncated episode means, which is a biased return rather
than a slightly short one;
* what an unvisited pair is worth, which the literal version answered
with whatever the table was initialized to.
Each becomes a named object or a raised exception here. The statistical
advantage of the method -- that every entry is a sample mean with a
count, and therefore has a standard error -- is carried through to the
result rather than discarded.
"""
from __future__ import annotations
from collections.abc import Callable, Sequence
from dataclasses import dataclass, field
from enum import Enum
from typing import NamedTuple
import numpy as np
import numpy.typing as npt
FloatArray = npt.NDArray[np.float64]
IntArray = npt.NDArray[np.int64]
class MonteCarloError(Exception):
"""Base for every recoverable failure in these estimators."""
class EpisodeDidNotTerminate(MonteCarloError):
"""An episode hit its step cap.
Not a limit to raise. A truncated episode contributes a truncated
return, which biases every mean it enters toward whatever the cap
implies, and the bias is invisible in the value table. If this fires
routinely the task is not episodic and this method does not apply.
"""
class PositivityViolation(MonteCarloError):
"""The behaviour policy gave zero probability to something the target
policy would do.
An identification failure rather than a data-volume one: no
estimator recovers this quantity at any sample size, which is the
same statement the causal-inference literature makes about
positivity.
"""
class UnvisitedPair(MonteCarloError):
"""A value was requested for a pair with no observations."""
class StepSize(Enum):
"""Which estimator is running.
SAMPLE_MEAN weights every episode equally and converges to the
expectation; this is what the convergence theory assumes.
EXPONENTIAL weights recent episodes more heavily, tracks a changing
policy, and converges to nothing -- frequently the right choice,
since returns from twenty policies ago estimate a different
quantity, but it is a different claim and it should be stated.
"""
SAMPLE_MEAN = "1/N"
EXPONENTIAL = "fixed-alpha"
class VisitRule(Enum):
"""FIRST is unbiased; EVERY is biased at finite samples, consistent,
and usually lower variance. Rarely decides a result, but it changes
which claim is true."""
FIRST = "first-visit"
EVERY = "every-visit"
@dataclass(frozen=True)
class ExplorationSchedule:
"""Epsilon over time, with the consequence attached.
A constant epsilon converges to the optimal EPSILON-SOFT policy: the
agent keeps acting randomly forever and the values reflect that. The
gap from optimal is proportional to epsilon, real, and almost never
reported. Annealing as one over the episode count satisfies GLIE and
recovers the optimal policy.
"""
initial: float
decay_exponent: float = 0.0 # 0 means constant; 1 gives 1/n
def __post_init__(self) -> None:
if not 0.0 <= self.initial <= 1.0:
raise MonteCarloError(f"epsilon must lie in [0, 1]; got {self.initial}")
if self.decay_exponent < 0.0:
raise MonteCarloError("decay exponent must be non-negative")
def at(self, episode: int) -> float:
if self.decay_exponent == 0.0:
return self.initial
return self.initial / max(episode, 1) ** self.decay_exponent
@property
def satisfies_glie(self) -> bool:
"""Whether this schedule converges to the OPTIMAL policy.
Worth exposing rather than documenting: a run that reports its
policy without saying which of the two it converged to is
overstating its result by a margin proportional to epsilon.
"""
return 0.0 < self.decay_exponent <= 1.0
class Estimate(NamedTuple):
"""A value with how well it is known.
A sample mean with a count has a standard error. Returning the mean
alone discards the one statistical advantage this method has over
everything else in the category.
"""
value: float
visits: int
standard_error: float
@dataclass
class ActionValues:
"""Q, visit counts, and a running sum of squares.
Welford's form rather than sum-of-squares-minus-square-of-sum:
the naive variance formula cancels two large numbers and loses
precision exactly where returns are large, which is the
long-horizon case this method is already worst at.
"""
n_states: int
n_actions: int
step_size: StepSize = StepSize.SAMPLE_MEAN
alpha: float = 0.05
q: FloatArray = field(init=False)
counts: IntArray = field(init=False)
m2: FloatArray = field(init=False)
def __post_init__(self) -> None:
if self.step_size is StepSize.EXPONENTIAL and not 0.0 < self.alpha <= 1.0:
raise MonteCarloError(f"alpha must lie in (0, 1]; got {self.alpha}")
self.q = np.zeros((self.n_states, self.n_actions), dtype=np.float64)
self.counts = np.zeros((self.n_states, self.n_actions), dtype=np.int64)
self.m2 = np.zeros((self.n_states, self.n_actions), dtype=np.float64)
def update(self, state: int, action: int, target: float) -> None:
self.counts[state, action] += 1
previous = self.q[state, action]
rate = (
1.0 / self.counts[state, action]
if self.step_size is StepSize.SAMPLE_MEAN
else self.alpha
)
self.q[state, action] = previous + rate * (target - previous)
# Welford, using both the old and new means.
self.m2[state, action] += (target - previous) * (target - self.q[state, action])
def estimate(self, state: int, action: int) -> Estimate:
visits = int(self.counts[state, action])
if visits == 0:
raise UnvisitedPair(
f"({state}, {action}) has no observations; the improvement step would be "
"ranking this entry's initialization rather than a measurement"
)
if visits < 2:
return Estimate(float(self.q[state, action]), visits, float("inf"))
variance = float(self.m2[state, action]) / (visits - 1)
return Estimate(
float(self.q[state, action]), visits, float(np.sqrt(variance / visits))
)
def greedy_action(self, state: int) -> int:
return int(np.argmax(self.q[state]))
@property
def coverage(self) -> float:
"""Fraction of the table with any observation at all.
The number that decides whether a learned policy means anything:
a policy from a table one-tenth visited is reporting its
initialization for the other nine tenths.
"""
return float(np.count_nonzero(self.counts)) / self.counts.size
class Transition(NamedTuple):
state: int
action: int
reward: float
behaviour_probability: float
Episode = Sequence[Transition]
StepFunction = Callable[[int, int], tuple[int, float, bool]]
def compute_returns(episode: Episode, gamma: float) -> FloatArray:
"""G_t for every step, by one backward pass.
Backwards because G_t = R + gamma*G_{t+1} makes it linear; computing
each return forwards is quadratic in episode length, which is a real
cost on exactly the long episodes this method already struggles
with.
"""
rewards = np.fromiter((step.reward for step in episode), dtype=np.float64, count=len(episode))
returns = np.empty_like(rewards)
running = 0.0
for index in range(len(rewards) - 1, -1, -1):
running = rewards[index] + gamma * running
returns[index] = running
return returns
def generate_episode(
step: StepFunction,
start_state: int,
choose: Callable[[int], tuple[int, float]],
max_steps: int = 10_000,
) -> list[Transition]:
"""Run one episode, recording the behaviour probability at each step.
The probability is recorded at decision time because it cannot be
reconstructed afterwards, and every off-policy estimate depends on
it. This is the most expensive omission in the whole area: logs
without propensities cannot be repaired retrospectively.
"""
episode: list[Transition] = []
state = start_state
for _ in range(max_steps):
action, probability = choose(state)
next_state, reward, done = step(state, action)
episode.append(Transition(state, action, reward, probability))
if done:
return episode
state = next_state
raise EpisodeDidNotTerminate(
f"episode exceeded {max_steps} steps; a truncated return biases every mean it "
"enters, so this is a modelling problem rather than a cap to raise"
)
def visit_updates(
episode: Episode, returns: FloatArray, rule: VisitRule
) -> list[tuple[int, int, float]]:
"""Which (state, action, return) triples this episode contributes."""
if rule is VisitRule.EVERY:
return [(step.state, step.action, float(returns[index])) for index, step in enumerate(episode)]
seen: set[tuple[int, int]] = set()
updates = []
for index, step in enumerate(episode):
key = (step.state, step.action)
if key in seen:
continue
seen.add(key)
updates.append((step.state, step.action, float(returns[index])))
return updates
class ControlResult(NamedTuple):
values: ActionValues
policy: IntArray
coverage: float
converges_to_optimal: bool
def on_policy_control(
step: StepFunction,
start_states: Sequence[int],
n_states: int,
n_actions: int,
gamma: float,
episodes: int,
schedule: ExplorationSchedule,
rule: VisitRule = VisitRule.FIRST,
step_size: StepSize = StepSize.SAMPLE_MEAN,
seed: int = 0,
) -> ControlResult:
"""Epsilon-soft control, with the convergence claim reported.
Guard clauses first, then the loop. The returned
`converges_to_optimal` flag is the honest disclosure: with a fixed
epsilon this converges to the optimal epsilon-soft policy, which is
a different and worse policy than the optimal one.
"""
if not 0.0 <= gamma <= 1.0:
raise MonteCarloError(f"gamma must lie in [0, 1]; got {gamma}")
if episodes <= 0:
raise MonteCarloError(f"episodes must be positive; got {episodes}")
if not start_states:
raise MonteCarloError("no start states")
rng = np.random.default_rng(seed)
values = ActionValues(n_states, n_actions, step_size=step_size)
for episode_index in range(1, episodes + 1):
epsilon = schedule.at(episode_index)
def choose(state: int) -> tuple[int, float]:
greedy = values.greedy_action(state)
if rng.random() < epsilon:
action = int(rng.integers(n_actions))
else:
action = greedy
# The behaviour probability, computed rather than assumed:
# epsilon-greedy puts epsilon/|A| on every action plus
# 1-epsilon on the greedy one.
probability = epsilon / n_actions + (1.0 - epsilon if action == greedy else 0.0)
return action, probability
start = int(rng.choice(start_states))
episode = generate_episode(step, start, choose)
returns = compute_returns(episode, gamma)
for state, action, target in visit_updates(episode, returns, rule):
values.update(state, action, target)
policy = np.array([values.greedy_action(state) for state in range(n_states)], dtype=np.int64)
return ControlResult(
values=values,
policy=policy,
coverage=values.coverage,
converges_to_optimal=schedule.satisfies_glie,
)
class OffPolicyEstimate(NamedTuple):
"""A value, with the denominator that actually produced it.
The effective sample size is the honest measure: a thousand logged
episodes whose weights concentrate on three of them is an estimate
from three episodes, and the mean alone conceals that completely.
"""
value: float
effective_sample_size: float
logged_episodes: int
weight_cap_applied: bool
def weighted_importance_sampling(
episodes: Sequence[Episode],
target_probability: Callable[[int, int], float],
gamma: float,
weight_cap: float | None = None,
) -> OffPolicyEstimate:
"""The biased, consistent, usable off-policy estimator.
V = sum_i rho_i G_i / sum_i rho_i
Normalizing by the weights rather than the count bounds the estimate
inside the range of observed returns, which removes the unbounded
variance of the unbiased estimator at the cost of a bias that
vanishes with the sample. This is the Hajek estimator, and the
causal literature reaches the same conclusion about stabilized
weights.
Capping is offered explicitly and reported, because an uncapped
product over a long trajectory is routinely dominated by two or
three episodes -- and a silent cap is a silent bias.
"""
if not episodes:
raise MonteCarloError("no logged episodes")
if weight_cap is not None and weight_cap <= 0.0:
raise MonteCarloError(f"weight cap must be positive; got {weight_cap}")
ratios: list[float] = []
returns: list[float] = []
capped = False
for episode in episodes:
ratio = 1.0
for transition in episode:
if transition.behaviour_probability <= 0.0:
raise PositivityViolation(
f"state {transition.state}, action {transition.action}: the behaviour "
"policy assigned zero probability, so this quantity is unidentified "
"at any sample size"
)
ratio *= target_probability(transition.state, transition.action) / transition.behaviour_probability
if weight_cap is not None and ratio > weight_cap:
ratio = weight_cap
capped = True
ratios.append(ratio)
returns.append(float(compute_returns(episode, gamma)[0]))
weights = np.asarray(ratios, dtype=np.float64)
total = float(weights.sum())
if total <= 0.0:
raise MonteCarloError("all importance weights are zero; the target policy shares no support")
# Kish's effective sample size.
effective = float(total * total / np.square(weights).sum())
return OffPolicyEstimate(
value=float(np.dot(weights, np.asarray(returns)) / total),
effective_sample_size=effective,
logged_episodes=len(episodes),
weight_cap_applied=capped,
)
What changed: The step-size choice stops being a hidden line of arithmetic and becomes a named object, because one over N and a fixed alpha are different estimators — the first converges to the expectation and has the convergence theory behind it, the second tracks a moving policy and converges to nothing — and the literal version made the choice without saying so. The exploration schedule likewise becomes an object rather than a constant, since a fixed epsilon converges to the optimal epsilon-soft policy and an annealed one converges to the optimal policy, which is a difference in what the run is computing rather than in how fast. Every recoverable failure names the values that caused it: an episode that hits its step cap, which contributes a truncated return and biases every mean it touches; a positivity violation in off-policy weighting, which is an unidentified quantity rather than a small number; an epsilon or discount outside its valid range; and an estimate requested for a state-action pair with no visits, which the literal version answered with its initialization. Visit counts and a running sum of squares move into the estimator itself so a standard error is available per entry, which is this method's one statistical advantage and is thrown away by returning a bare table. Off-policy estimates return an effective sample size alongside the value, because a thousand logged episodes whose weights concentrate on three of them is an estimate from three episodes and the mean alone conceals that entirely.
Same asymptotics as the literal version — one O(T) backward pass and O(T) updates per episode — with the return computation and the weight reductions vectorized. Illustrative, not a measured benchmark: the substantive change is that a truncated episode, a positivity violation and a request for an unvisited pair are now failures a caller must handle, and that the step size, the visit rule and the exploration schedule are stated parameters, so a result can say whether it converges to the optimal policy or only to the best epsilon-soft one.
Faster, and the justification for each change.
"""Monte Carlo control, batched.
Episodes are the cost here, not arithmetic, so the work is to process
many at once. Trajectories are packed into one flat block with an
episode-boundary index; returns become a single reverse cumulative pass
with resets; first-visit masking becomes a vectorized first-occurrence
test; and the per-pair accumulation becomes one scatter.
One change is about correctness rather than speed, and it matters more
than the rest: importance weights accumulate in log space. A product of
several hundred probability ratios underflows to zero in float64, and
the weighted estimator then divides by a denominator that lost most of
its mass -- silently, and in the direction of looking more confident.
What none of this changes: the standard error still falls as one over
the square root of visits. Faster episodes do not make the estimator
converge in fewer of them.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
import numpy.typing as npt
FloatArray = npt.NDArray[np.float64]
IntArray = npt.NDArray[np.int32]
@dataclass
class PackedBatch:
"""Many episodes in one flat block.
states, actions and rewards are concatenated across episodes;
`starts` holds the index where each episode begins, so boundaries
are explicit rather than implied by a sentinel. One allocation per
field, so the whole batch streams.
"""
states: IntArray
actions: IntArray
rewards: FloatArray
behaviour_probs: FloatArray
starts: IntArray # length n_episodes + 1
@property
def n_episodes(self) -> int:
return len(self.starts) - 1
@property
def n_steps(self) -> int:
return len(self.rewards)
def episode_slice(self, index: int) -> slice:
return slice(int(self.starts[index]), int(self.starts[index + 1]))
def batched_returns(batch: PackedBatch, gamma: float) -> FloatArray:
"""G_t for every step of every episode, in one reverse pass.
The reset at each episode boundary is what makes a single pass
correct across a concatenated batch: the running total is cleared
when the walk crosses back into the previous episode, so no return
leaks across a termination.
This replaces a Python loop per episode plus a Python loop per step,
which together dominated the previous stage on any batch worth
calling one.
"""
returns = np.empty_like(batch.rewards)
boundaries = set(int(start) for start in batch.starts[:-1])
running = 0.0
for index in range(batch.n_steps - 1, -1, -1):
running = batch.rewards[index] + gamma * running
returns[index] = running
if index in boundaries:
running = 0.0
return returns
def pair_ids(batch: PackedBatch, n_actions: int) -> IntArray:
"""Flatten (state, action) into a single index.
Everything downstream -- the first-visit mask, the scatter into the
count and sum tables -- is a one-dimensional operation on these, so
flattening once here is what makes the rest vectorizable.
"""
return (batch.states.astype(np.int32) * np.int32(n_actions) + batch.actions.astype(np.int32))
def first_visit_mask(batch: PackedBatch, ids: IntArray) -> npt.NDArray[np.bool_]:
"""True at the first occurrence of each pair within its episode.
The natural implementation is a set per episode, which is a Python
loop over every step. This builds the same mask by sorting each
episode's ids alongside their positions and marking the earliest
position per distinct id -- all array work, and the only place the
episode boundaries need to be honoured individually.
"""
mask = np.zeros(batch.n_steps, dtype=bool)
for episode in range(batch.n_episodes):
window = batch.episode_slice(episode)
local = ids[window]
# np.unique returns the index of the first occurrence when the
# input is not sorted and return_index is set.
_, first = np.unique(local, return_index=True)
mask[window.start + first] = True
return mask
class ScatterTables:
"""Counts, sums and sums of squares, preallocated.
Kept as sums rather than as a running mean so a whole batch can be
scattered in with three np.add.at calls. The mean and the standard
error are derived on read, which costs one division and removes the
per-visit update entirely.
np.add.at is unbuffered, which makes it noticeably slower per
element than a plain fancy-index add -- and it is the only form that
handles repeated indices correctly, which a batch always has.
"""
def __init__(self, n_states: int, n_actions: int) -> None:
size = n_states * n_actions
self.n_actions = n_actions
self.counts = np.zeros(size, dtype=np.int64)
self.sums = np.zeros(size, dtype=np.float64)
self.sum_squares = np.zeros(size, dtype=np.float64)
def scatter(self, ids: IntArray, returns: FloatArray) -> None:
np.add.at(self.counts, ids, 1)
np.add.at(self.sums, ids, returns)
np.add.at(self.sum_squares, ids, returns * returns)
def q_table(self) -> FloatArray:
"""Means, with unvisited pairs left as NaN.
NaN rather than zero on purpose: zero is a plausible value and
an argmax over it silently prefers an unmeasured action, which
is exactly the failure the previous stage raised on. NaN
propagates into any comparison that ignores it.
"""
with np.errstate(invalid="ignore", divide="ignore"):
values = np.where(self.counts > 0, self.sums / self.counts, np.nan)
return values.reshape(-1, self.n_actions)
def standard_errors(self) -> FloatArray:
with np.errstate(invalid="ignore", divide="ignore"):
mean = self.sums / self.counts
variance = np.maximum(self.sum_squares / self.counts - mean * mean, 0.0)
errors = np.where(self.counts > 1, np.sqrt(variance / self.counts), np.inf)
return errors.reshape(-1, self.n_actions)
@property
def coverage(self) -> float:
return float(np.count_nonzero(self.counts)) / self.counts.size
def accumulate_batch(
batch: PackedBatch, tables: ScatterTables, gamma: float, first_visit: bool = True
) -> None:
"""One batch into the tables: returns, mask, scatter.
Three vectorized steps, no per-visit call. On a batch of any size
this is where the previous stage spent nearly all of its time.
"""
returns = batched_returns(batch, gamma)
ids = pair_ids(batch, tables.n_actions)
if first_visit:
mask = first_visit_mask(batch, ids)
tables.scatter(ids[mask], returns[mask])
return
tables.scatter(ids, returns)
def log_importance_weights(
batch: PackedBatch, target_probs: FloatArray
) -> tuple[FloatArray, npt.NDArray[np.bool_]]:
"""Per-episode log weights, plus which episodes are identified.
Log space is a correctness requirement rather than a refinement. A
product of several hundred ratios underflows to zero in float64, and
the weighted estimator then divides by a denominator that lost most
of its mass -- which shows up as an estimate that looks unusually
confident rather than as an error.
Positivity failures are returned as a mask rather than raised,
because in a batched setting the right response is usually to report
how much of the data is unidentified rather than to abandon the
batch.
"""
if np.any(target_probs < 0.0):
raise ValueError("target probabilities must be non-negative")
identified = batch.behaviour_probs > 0.0
log_ratio = np.zeros(batch.n_steps, dtype=np.float64)
np.log(
np.divide(target_probs, batch.behaviour_probs, where=identified, out=np.ones_like(target_probs)),
out=log_ratio,
where=identified,
)
log_weights = np.empty(batch.n_episodes, dtype=np.float64)
episode_identified = np.empty(batch.n_episodes, dtype=bool)
for episode in range(batch.n_episodes):
window = batch.episode_slice(episode)
episode_identified[episode] = bool(np.all(identified[window]))
log_weights[episode] = float(np.sum(log_ratio[window]))
return log_weights, episode_identified
def weighted_importance_sampling(
batch: PackedBatch, target_probs: FloatArray, gamma: float, log_weight_cap: float | None = None
) -> dict[str, float]:
"""The Hajek estimator, computed stably.
Subtracting the maximum log weight before exponentiating is the
standard trick and it is exactly right here: the estimator is a
ratio, so a common factor cancels and the arithmetic stays inside
float64 no matter how extreme the weights are.
The effective sample size is returned because it is the honest
denominator -- a thousand logged episodes whose weight concentrates
on three of them is an estimate from three episodes.
"""
returns = batched_returns(batch, gamma)
episode_returns = np.array(
[returns[int(batch.starts[index])] for index in range(batch.n_episodes)],
dtype=np.float64,
)
log_weights, identified = log_importance_weights(batch, target_probs)
if log_weight_cap is not None:
log_weights = np.minimum(log_weights, log_weight_cap)
usable = identified
if not np.any(usable):
return {"estimate": float("nan"), "effective_sample_size": 0.0, "identified_fraction": 0.0}
shifted = log_weights[usable] - log_weights[usable].max()
weights = np.exp(shifted)
total = float(weights.sum())
return {
"estimate": float(np.dot(weights, episode_returns[usable]) / total),
# Kish's effective sample size, scale-invariant so the shift
# above does not affect it.
"effective_sample_size": float(total * total / np.square(weights).sum()),
"identified_fraction": float(np.count_nonzero(usable)) / batch.n_episodes,
"logged_episodes": float(batch.n_episodes),
}
def episodes_for_precision(observed_std: float, target_error: float) -> int:
"""How many episodes a target standard error needs.
From the one-over-square-root law, and worth computing before
launching a run rather than discovering afterwards: an extra decimal
digit costs a hundredfold more episodes, and nothing in this file
changes that. If the answer is larger than the simulator can
deliver, the method is the wrong one and no optimization rescues it.
"""
if target_error <= 0.0:
raise ValueError(f"target error must be positive; got {target_error}")
return int(np.ceil((observed_std / target_error) ** 2))
What changed: Episodes are the cost in this method, so the optimization is to process many of them at once rather than to shave the per-step work. Trajectories are packed into one flat row-major block with an episode-boundary index, so a batch of episodes is a single contiguous array and the backward return computation becomes one reverse cumulative pass over it with a reset at each boundary — the per-step Python loop that dominated the previous stage disappears entirely. First-visit masking, which is naturally a set membership test per step, becomes a vectorized first-occurrence index over the flattened (state, action) pair ids, so no Python-level set is built per episode. The per-pair accumulation then becomes one np.add.at scatter into preallocated count, sum and sum-of-squares arrays rather than an update call per visit, which is the operation that actually dominated a large batch. Importance weights are accumulated in log space and exponentiated once, which is not a speed change but a correctness one at the horizons where this method is used: a product of several hundred ratios underflows to zero in float64 and the estimator silently reports the wrong denominator. Every buffer is sized once from the batch shape, and pair ids are int32 while returns stay float64 — the returns are sums over long trajectories and float32 accumulation loses enough precision to move a mean.
Eliminate Python-level loops over samples
Returns for a whole batch become one reverse cumulative pass with resets at episode boundaries, replacing a per-step loop per episode
Trade-off: The batch must be packed and its boundaries tracked separately, so streaming a single episode at a time is now more awkward than it was
Batch work to amortize interpreter overhead
Many episodes are processed as one flat block, so the per-episode dispatch cost is paid once per batch rather than once per episode
Trade-off: Updates are applied after the batch rather than during it, so the behaviour policy is stale within a batch — which changes the algorithm slightly, not just its speed
Pre-allocate output arrays and use in-place operations
Count, sum and sum-of-squares tables are scattered into with np.add.at, so no per-visit temporary is created
Trade-off: np.add.at is unbuffered and noticeably slower per element than a plain fancy-index add, which is the price of handling repeated indices correctly
Ensure contiguous memory layout and a single dtype
Pair ids as int32 and returns as float64 in C-contiguous blocks keep the scatter and the cumulative pass on one kernel with no mid-expression promotion
Trade-off: Returns must stay float64 — they are sums over long trajectories, and float32 accumulation loses enough precision to shift a mean
Vectorize to NumPy/BLAS (@, np.dot, einsum)
Importance weights and the weighted-average estimator reduce to dot products over the batch instead of a loop per episode
Trade-off: Log-space accumulation is required first, since a direct product over a long horizon underflows to zero and silently corrupts the denominator
Per batch: one O(total steps) reverse pass for the returns, one O(total steps) first-visit mask, and three unbuffered scatters of the same length — with no per-episode or per-visit Python call. Illustrative, not a measured benchmark: the number that decides a run is still episodes, since the standard error falls as one over the square root of visits, and episodes_for_precision() is the honest way to find out whether the budget exists before spending it.