Writing · · 10 min read
The Latent Action
How a world model can learn what a surgeon did without action labels: tokenize the frames, infer the transition, and test whether the action is real.
#world-models#surgical-ai#video#generative-models
A world model needs three things: a state, an action, and a rule for what happens next.
Surgical video gives the first and the third. It gives frames, and it gives the next frames. It does not give the action. The surgeon's hand moves, the tissue deforms, the tool disappears behind an organ, and nowhere in the file is there a label that says what happened.
The latent action is a way of inventing that label from the video itself.
The simplest model: supervised action labels
The naive model assumes action labels exist. A surgeon annotates each clip frame by frame, gesture by gesture: grasp, pull, cut, cauterise. A supervised model maps frames to labels. A dynamics model maps states and labels to the next state.
The simplest program written against that model looks like this:
# Labels are the bottleneck.
# Each clip needs a surgeon to annotate frame by frame.
def supervised_step(frame, label):
next_frame, _ = dynamics(frame, label)
return next_frameConsider a laparoscopic clip. At time , a tool is touching a piece of tissue. At time , the tissue has moved. What was the action?
It could be a grasp, a pull, a cut, a cauterisation, a camera movement, or a combination. The pixels show the consequence. They do not show the intention. Supervised annotation is expensive, slow, and inconsistent across annotators.
The unsupervised alternative is to learn a small code that explains the transition. The code is not a human label. It is a compressed difference between two observations. The only test of the code is whether it is useful for predicting the next frame.
A toy where the answer is known
Before surgical video, a world small enough to check by hand. A frame is a 1-D strip of 16 pixels with one bright pixel. A hidden action moves it left, keeps it still, or moves it right. The video is saved; the actions are thrown away.
import numpy as np
W = 16 # frame width in pixels
MOVES = np.array([-1, 0, 1]) # hidden actions: left, stay, right
def render(p):
x = np.zeros(W); x[p] = 1.0
return x
def make_video(T, seed=0):
rng = np.random.default_rng(seed)
p, frames, labels = W // 2, [], []
for _ in range(T):
a = rng.integers(3)
frames.append(render(p)); labels.append(a)
p = int(np.clip(p + MOVES[a], 0, W - 1))
frames.append(render(p))
return np.array(frames), np.array(labels) # labels are hidden from the learnerThe question is whether the three actions can be recovered from frame pairs alone. The simplest learner is hard EM. Give each of codes its own linear dynamics map. Assign every frame pair to the code whose map predicts the next frame best. Refit each map on the pairs it owns. Repeat.
def fit_latent_actions(frames, K, iters=20, seed=0):
rng = np.random.default_rng(seed)
X = np.hstack([frames[:-1], np.ones((len(frames) - 1, 1))]) # x_t plus bias
Y = frames[1:] # x_{t+1}
z = rng.integers(K, size=len(X)) # random codes
for _ in range(iters):
maps = []
for k in range(K):
idx = z == k
if idx.sum() == 0: # dead code: re-seed it
idx = rng.random(len(X)) < 0.05
M, *_ = np.linalg.lstsq(X[idx], Y[idx], rcond=None)
maps.append(M)
err = np.stack([((X @ M - Y) ** 2).sum(1) for M in maps], 1)
z = err.argmin(1) # E-step: best code per pair
return z, err.min(1).mean()
def purity(z, labels, n_labels):
# Fraction of pairs whose code's majority label equals their own label.
return sum(np.bincount(labels[z == k], minlength=n_labels).max()
for k in np.unique(z)) / len(z)Sweeping on 5,000 transitions, best of five seeds, gives:
| codes | prediction error | action purity | next-position purity |
|---|---|---|---|
| 1 | 0.632 | 0.335 | 0.094 |
| 2 | 0.275 | 0.665 | 0.144 |
| 3 | 0.000 | 0.953 | 0.150 |
| 4 | 0.000 | 0.953 | 0.166 |
| 8 | 0.000 | 0.568 | 0.322 |
| 16 | 0.000 | 0.621 | 0.476 |
Three regimes are visible. Below three codes the model cannot explain the video. At three or four codes the prediction error hits zero and the codes line up with the hidden actions; the missing 5% are pairs at the walls, where "left" and "stay" produce identical frames and no learner could tell them apart. At eight or sixteen codes the prediction error is still zero, but the codes have drifted toward describing where the pixel ends up rather than how it moved. Prediction error alone cannot tell the good regime from the bad one.
That table is the whole post in miniature. The rest scales it up to pixels and surgeons.
Deriving the structure: states as tokens, actions as bottlenecks
Two structural decisions follow. First, raw pixels are the wrong substrate for a dynamics model. Pixels are large, noisy, and entangled with appearance. A model that predicts pixels directly must relearn appearance at every step.
Second, an unconstrained code for the transition will memorize the next frame instead of describing the action. A code large enough to hold directly frees the dynamics model from using at all. The action becomes a copy operation.
The fixes come in sequence: discretize the state so a transformer can model it, then bottleneck the action so it can only describe what changed.
Step 1: tokenize the frames with a VQ-VAE
The first component is a VQ-VAE. An encoder maps a frame to a latent grid; each vector in the grid is replaced by the nearest entry in a learned codebook; a decoder reconstructs the frame from the quantized grid.
The output is a grid of discrete tokens. That is useful for two reasons. It compresses the video into a small vocabulary that a transformer can model. It separates the visual content from the temporal model: the dynamics model operates on tokens, not raw pixels.
The reconstruction loss trains the encoder and decoder. The codebook loss and the commitment loss keep the encoder from ignoring the codebook. A typical objective is:
where is the encoder output, is the selected codebook entry, and is the stop-gradient operator. The first term is reconstruction. The second moves the codebook toward the encoder. The third moves the encoder toward the codebook.
def vq_step(x, encoder, decoder, codebook, beta=0.25):
z_e = encoder(x) # continuous grid
e, indices = codebook.nearest(z_e) # quantized grid
x_hat = decoder(e)
recon = ((x - x_hat) ** 2).mean()
cb_loss = ((z_e.detach() - e) ** 2).mean()
commit_loss = beta * ((z_e - e.detach()) ** 2).mean()
return recon + cb_loss + commit_loss, indicesBottleneck after Step 1: tokens describe appearance. Nothing yet describes a transition. The temporal model still lacks an action.
Step 1b: watch the codebook, not just the loss
A VQ-VAE can reach a low reconstruction loss while using a handful of its codes. The rest are dead: never selected, never updated. Two numbers expose it — utilization and perplexity of the code histogram.
import numpy as np
def codebook_health(indices, K):
# indices: integer array of selected codes over a batch of frames
counts = np.bincount(indices.ravel(), minlength=K)
p = counts / counts.sum()
nz = p[p > 0]
perplexity = float(np.exp(-(nz * np.log(nz)).sum()))
return {"utilization": float((counts > 0).mean()), "perplexity": perplexity}
def restart_dead_codes(codebook, z_e_flat, counts, rng):
# Replace unused entries with random encoder outputs from the batch.
dead = np.flatnonzero(counts == 0)
codebook[dead] = z_e_flat[rng.integers(len(z_e_flat), size=len(dead))]
return codebookA perplexity close to means codes are used evenly. A perplexity of 30 on a 1,024-entry codebook means the vocabulary is mostly decoration. The same check applies later to the action codebook, where collapse to one code means "the action does nothing".
Step 2: infer the action from a frame pair
The second component is the latent action model. It takes two consecutive frames and produces a discrete action code:
The dynamics model then predicts the next frame from the current frame and the action:
The two models are trained together. The action model is not given the true action; it is given the pair of frames. The only pressure on is that it must help the dynamics model predict .
def joint_step(x_t, x_next, action_model, dynamics):
z = action_model(x_t, x_next) # discrete action code
x_hat = dynamics(x_t, z) # prediction
loss = ((x_hat - x_next) ** 2).mean() # prediction pressure only
return loss, zThis is an information bottleneck. If is too small, it cannot carry enough information about the transition. If it is too large, it can encode the next frame directly, and the dynamics model learns to ignore the current frame. The useful regime is in between: a compact code that captures what changed, not everything that is present.
The Genie framework introduced this idea for games: a latent action model infers a discrete action from a pair of frames, and a dynamics model learns to generate the next frame conditioned on that action. It kept the action vocabulary deliberately tiny — eight codes — so that each code had to mean something a player could choose. The same architecture can be applied to surgical video, where the "action" is a gesture rather than a button press.
Step 2b: size the bottleneck in bits
Bottleneck: "small enough" needs a number. The toy shows the failure at , which is exactly the number of possible next frames.
Fix: compare the information the code can carry with the information in the next frame. A code from a vocabulary of carries at most bits per transition.
| setting | next-frame information | action code capacity |
|---|---|---|
| toy strip, | bits | 1.6 bits |
| toy strip, | 4 bits | 4 bits — enough to copy |
| 16×16 tokens, 1,024-entry codebook | up to 2,560 bits | — |
| same, 8 action codes | — | 3 bits |
import math
def copy_risk(frame_tokens, codebook_size, action_codes):
# Ratio of action capacity to the worst-case information in a frame.
frame_bits = frame_tokens * math.log2(codebook_size)
action_bits = math.log2(action_codes)
return action_bits / frame_bits
print(copy_risk(256, 1024, 8)) # ~0.001: the code must summarise, not copyThe rule is not "as small as possible". Too few codes and distinct gestures share a code, as at in the toy. The search is over a narrow band, and the evaluation below is what chooses within it.
What the code can learn
The latent action space is not guaranteed to align with human semantics. It can learn:
- Tool motion: a translation, a rotation, a grasp.
- Camera motion: a pan, a zoom, a change in viewpoint.
- Tissue deformation: a pull, a release, a recoil.
- Lighting and smoke: a change in illumination, a plume from cautery.
- Compression artefacts: a change in the encoder's reconstruction.
Only some of these are surgical actions. The model does not know the difference. It knows which differences are useful for predicting the next frame.
That is the central risk. A latent action that predicts the next frame perfectly may be a camera-motion detector. A latent action that clusters by gesture may be a semantic representation. The two are not the same.
Step 3: test whether the action is real
Bottleneck: video quality alone cannot distinguish a semantic action from a camera-motion detector. Evaluation must go beyond reconstruction.
Fix: apply four tests, each targeting a different failure mode.
Intervention. Hold the frame fixed, change the action code, and generate. If the code is meaningful, different codes should produce different, plausible motions. If the code is ignored, the outputs will be identical.
def intervention_gap(frame, codes, dynamics):
outputs = [dynamics(frame, z) for z in codes]
# Pairwise distance between outputs conditioned on different codes.
return max(((outputs[i] - outputs[j]) ** 2).mean()
for i in range(len(outputs))
for j in range(i + 1, len(outputs)))Consistency. Apply the same action code to similar frames. The generated motion should be similar. Apply different codes to the same frame. The outputs should differ in a structured way.
Probing. Collect a small labelled set of surgical gestures. Train a linear classifier or a nearest-neighbour probe on the latent codes. If the codes carry gesture information, the probe should do better than chance. That result does not prove causality, but it is evidence.
import numpy as np
def nearest_centroid_probe(codes_train, y_train, codes_test, y_test):
# codes: (n, d) continuous action embeddings before quantization
classes = np.unique(y_train)
centroids = np.stack([codes_train[y_train == c].mean(0) for c in classes])
d = ((codes_test[:, None, :] - centroids[None]) ** 2).sum(-1)
acc = (classes[d.argmin(1)] == y_test).mean()
chance = np.bincount(y_test).max() / len(y_test)
return acc, chance # report both; accuracy alone hides class imbalanceRun the same probe against a nuisance label — camera motion, smoke present, frame brightness. A code that predicts camera motion better than gesture is the toy's failure in surgical form.
Physical plausibility. Ask a surgeon to review generated clips. Does the tissue move like tissue? Does the tool interact with the anatomy correctly? Does the model invent bleeding, perforation, or motion that violates the scene?
Downstream transfer. Train a policy inside the world model and evaluate it on a real or simulated task. If the world model is useful, the policy should learn faster or better than a policy trained without it. That is the hardest test and the most convincing one.
The paper and the claim
The Surgical Vision World Model adapted this recipe to surgical video. The model learns discrete visual tokens, infers latent actions from frame pairs, and generates action-controllable video. It was presented at the MICCAI 2025 workshop on Data Engineering in Medical Imaging, and the code is public.
The claim is not that the model understands surgery. The claim is that useful action representations can be learned from unlabeled video, and that the learned actions can be tested. The first part is an engineering result. The second part is the research programme.
What to measure
- Reconstruction loss for the VQ-VAE. High loss means tokens describe appearance poorly.
- Codebook utilization. The fraction of codebook entries ever selected. Collapse means the vocabulary is wasted.
- Intervention gap. Pairwise distance between outputs from different action codes on the same frame.
- Probe accuracy. Gesture classification accuracy from latent codes against a small labelled set, next to chance.
- Nuisance probe accuracy. The same probe against camera motion or smoke. It should be lower than the gesture probe.
- Action-code perplexity. Collapse to one code means the action does nothing; perplexity near the vocabulary size with low purity means it does too much.
- Downstream gain. Policy performance after training inside the world model versus a baseline.
- Plausibility rate. The fraction of generated clips a surgeon judges anatomically plausible.
A checklist
- Tokenize before modeling dynamics. Tokens separate appearance from time.
- Bottleneck the action. A code large enough to hold the next frame is a copy, not an action. Count the bits.
- Start on a toy with known actions. If the pipeline cannot recover left, stay, and right, it will not recover a grasp.
- Train the action model and the dynamics model together. The only supervision is prediction.
- Intervene before claiming semantics. Fixed frames with varied codes expose ignored codes.
- Probe with a small labelled set. Cluster structure is evidence, not proof.
- Test downstream transfer last. A world model that does not help a policy is decoration.
A latent action is a bet that the difference between two frames is the right abstraction for an action. The bet may be wrong. The difference may include the camera, the light, the smoke, and the patient's breathing. But the bet is testable, and that is what makes it science rather than decoration.
The scalpel does not say what it is doing. The model has to infer the sentence from the movement. The latent action is the first draft of that sentence.