Which gene would you assay next?

The agent starts with no test results at all. At each turn it either buys one more gene expression measurement or commits to a diagnosis. Every measurement costs budget, and a wrong diagnosis costs more, so it has to work out both which gene is worth buying for this patient and when it has seen enough to stop. The network below runs in your browser, so you are watching the trained policy decide, not a recording.

Loading the trained policy…

Keyboard: S step · Space play or pause · R reset · N next patient

Candidate panel

0 of 40 measured

Grey genes have not been ordered. The bar under a measured gene shows how far its expression sits from the cohort mean.

What the agent is weighing

action values

The value the network assigns to each action it may still take. Illegal actions (already measured, or unaffordable) are left out.

Belief

The posterior over tumour type given only what has been measured so far. The marked row is the patient's true label, revealed here for inspection.

Budget

Genes ordered

0

Confidence

-

Predicted return

C51

The agent does not predict a single number for its best action. It predicts a distribution over returns across 51 atoms, and the dashed line is the mean an ordinary Q network would report on its own. A wide spread means the outcome is still genuinely uncertain.

Decision log

Same patient, four strategies

every strategy commits with the same classifier, so only the ordering and stopping differ

Runs all four policies to completion on the patient currently selected.

Strategy	Genes ordered	Cost	Diagnosis	Correct	Confidence	Order

Training

Two agents were trained on the same environment: a Rainbow DQN built from the six components of Hessel et al. and a PPO agent with action masking. Curves are the mean over seeds with a bootstrap band. Evaluation runs the greedy policy over the whole validation split at every checkpoint rather than over sampled episodes, so the curve is not moved around by exploration noise.

Diagnostic accuracy

Held-out validation accuracy through training.

Genes ordered per patient

How many measurements the policy buys before committing.

Rainbow components

each variant removes exactly one component; higher is better

Final return by variant

Interquartile mean over seeds with a 95% stratified bootstrap interval. Bars that overlap the full agent's interval are components this task does not lean on.

Accuracy against measurements

The same variants, showing where the differences come from.

Variant	Return	Accuracy	Genes	Seeds

Analysis

The question is whether an adaptive ordering beats a fixed one at equal cost. The curves below separate the two halves of the problem: how good the acquisition order is, and how good the stopping rule is.

Accuracy against genes ordered

Each policy is forced to keep measuring, with its stopping rule disabled, and the classifier is scored after every acquisition. This isolates ordering quality.

Cost against accuracy

Operating points on the test split. Agent points come from training runs that priced measurements differently; baseline points come from their own stopping rules.

Operating points

held-out test split, identical classifier for every row

Strategy	Accuracy	Balanced acc.	Genes ordered	Cost	Return

Confusion matrix

Rows are the true tumour type, columns the agent's diagnosis.

Which gene gets ordered first

How often each gene is the agent's opening move, and how often it is ordered at all. A flat first-pick distribution would mean the agent adapts its opening to the patient, which it cannot do before seeing anything, so this is concentrated by construction.

How the opening moves branch

every test patient, folded into a tree by what the agent ordered

Band width is the share of patients taking that route. Colour marks how the previous gene read, blue for below the cohort mean and orange for above. A fixed panel would draw one chain across this figure, so the width here shows how much the policy adapts to the patient in front of it.

Calibration at the stopping point

Stated confidence against how often it was right, measured at the moment the agent chose to stop. Points near the diagonal mean the confidence can be taken at face value, which matters more than raw accuracy if anything is built on top of it.

Spread across seeds

The share of training runs scoring above a given return. One curve sitting above another everywhere is a stronger claim than a difference in averages, since it holds at every threshold rather than at a single summary point.

Classifier accuracy against random subsets

The reference curve. This is what a randomly chosen panel of k genes buys you, and it is the line the agent has to beat with a chosen panel of the same size.

Worked examples

Full traces from the test split, one per tumour type.

Patient	True type	Genes ordered	Diagnosis	Sequence

Method

The decision problem

A patient arrives with no test results. Ordering the full gene panel would classify them accurately, but panels cost money and time, and most of the genes carry no information once a few others have been read. The useful question is which gene to order next given what has come back so far, and when the evidence is enough to stop.

That is a partially observed Markov decision process. The hidden state is the patient's full expression profile together with their tumour type; the agent sees only the entries it has paid to reveal.

Reward

A measurement costs its price scaled by cost_scale. A correct diagnosis pays +1 and a wrong one -1. The ratio between those two settings is what decides how test-hungry the optimal policy is, which is why the frontier in the analysis comes from re-training at different prices rather than from thresholding one agent.

Action masking

The legal action set shrinks as the episode runs: a gene cannot be bought twice, and it cannot be bought at all once the remaining budget will not cover it. Masking is applied to the greedy action, to the bootstrap action inside the temporal difference target, and to the policy distribution, so the agent never assigns value to an action it cannot take.

Agent

The main agent is a Rainbow DQN written from scratch. All six components are present and each is a configuration flag, which is what makes the ablation on the training page a genuine one rather than a comparison against differently tuned code.

Double Q learning picks the bootstrap action with the online network and scores it with the target network.
Dueling heads split state value from action advantage.
Prioritised replay samples by temporal difference error through a sum tree, corrected by importance weights.
Multi-step returns accumulate three steps of reward before bootstrapping.
NoisyNet layers replace epsilon-greedy with learned parameter noise.
Distributional C51 predicts the return distribution over 51 atoms instead of its mean.

Network

Commitment is a small fixed choice and is scored from the pooled state by an ordinary head. Measurement is a choice over the panel, so it uses a pointer-style head: the state produces a query vector and the value of ordering gene j is the inner product of that query with gene j's own representation. The head's size does not grow with the panel.

Belief model

A separate classifier is trained on random measurement subsets, so it stays calibrated at every level of missingness. It supplies the posterior shown on the live page and the one the baselines commit with. The agent never receives it as input, so any margin the agent holds over the greedy baseline comes from looking further ahead, not from a better classifier.

Evaluation

Every policy, learned or scripted, is scored by the same driver over the entire test split rather than over sampled episodes. Aggregates across seeds are interquartile means with stratified bootstrap intervals; at a handful of seeds a plain mean is moved too easily by one run that collapsed or one that got lucky.

Gene selection and standardisation are fit on the training split alone. The panel, the belief model and the agent never see validation or test patients during fitting.

Data

On the cost model

Measurement costs are configurable. The default profile gives every gene a cost of one, which makes the budget a plain count of measurements. The alternative profiles are synthetic. They are not taken from any real assay price list, and exist only to make the budget allocation problem non-uniform.

Reproducing

make data
make belief
make experiments
make dashboard