Are Good Generators Good Decision-Makers?

The Answer: yes — with ADIGen + Retargeting.

Land the ball on the goal lineout-of-distribution gust · goal 3.4

ConditionalADIGenRetargeting (ours)

ImaginedReal

Top row: each generator imagines its chosen push landing the ball on the goal line. Bottom row: in the real simulator, the conditional and ADIGen pushes leave the ball short of the line; only the retargeted push lands on it.

miss 1.50miss 1.38miss 0.07

Imagined: every generator above pictures its chosen push landing on the goal line.
Real: in the simulator, only ours lands it — miss 0.07 vs. 1.50 and 1.38.
Why: generating realistic futures is not the same as choosing good actions, so we retarget the generator toward the decisions it has to support.

Problem

The data: logs show what happened, not what each action would do; actions can be high-dimensional; and policies change across environments.
The objective: a policy optimized against a poorly-extrapolating generator can be (wrongly) overconfident.

Fix

Train a doubly robust, invariant generator (ADIGen), learn a pilot policy from its rollouts, then retarget the generator toward that policy and relearn.

Key results

Video: realistic PhyWorld rollouts, and better decisions from them — on out-of-distribution gusts, regret 0.74 → 0.28 (conditional vs. retargeted).
Cells: retargeting recovers the rescue peak of all 3 active compounds from only 10% of the high-dose wells.
Theory: retargeting removes the worst-case density ratio \(\rho\) from the regret bound, beating the conditional target by a factor \(\sqrt{\rho}/2\).

1The problem: logs are not experiments

Counterfactual generators, from physics video to virtual cells, are increasingly used to choose interventions: for each candidate action, generate what would happen next, and act on the best imagined future — without trying every option for real.

Two motivating cases. Robotics: a generator imagines where an object goes under each push, but the logs contain hard pushes only on heavy objects, so the model thinks hard pushes don't send things far. Medicine: a generator forecasts each patient's response to each dose, but high doses were mostly given to younger patients, so the model recommends them for older patients too. — **Figure 1: Generators are increasingly used in decision-making — and the logs mislead them.** Robotics: hard pushes were logged only on heavy objects, so the model concludes hard pushes don't send things far. Medicine: high doses went mostly to younger patients, so the model recommends them for older patients too.

The failure has two parts, both visible in these cases:

The data. Logs show what happened, not what each action would do: the robot's logs pair hard pushes with heavy objects, the clinic's pair high doses with younger patients, so the model learns the pairing rather than the push or the dose. On top of that, actions can be high-dimensional, and policies change across environments.
The objective. The planner then asks the model about exactly the actions it never saw — hard pushes on light objects, high doses for older patients — and a poorly-extrapolating generator can be (wrongly) overconfident there.

Let's illustrate this with a simpler case (because we like soccer). Suppose you are trying to take a free kick. Here, the wall covers the near post, the keeper guards the far side, and a gust blows across the pitch. The planner cannot try every kick for real, with candidates — over the wall, near post, keeper's side. A generator will imagine the next few seconds of footage, and the kick whose imagined future scores is taken. The plan is only as good as the generator is at the kicks being compared.

Score the free kick: from the current frame, the model generates short rollouts for three candidate kicks — over the wall scores, near post is blocked, keeper's side is saved — and the planner acts on the best imagined future. — **Planning with a generator.** For each possible kick, the model generates what the camera would see next, and the planner acts on the best imagined future — without trying every kick for real.

In the kick logs, both problems are measurable:

Strong kicks were taken in strong gusts, so conditioning on the kick learns the choosing habit, not the kick's effect.
The winning kick in a strong gust — aim into the wind — was barely ever taken (as little as \(\eta = 0.002\) of the logs), and there the generator's predictions are guesses.

Two-by-two grid of logged kicks by gust strength and aim: 990 clips of aiming straight in a weak gust, but only 2 clips of aiming into a strong gust — the kick that actually scores there. — **The kick we need is the one the logs barely show.** Logged kicks by gust and aim: 990 clips aim straight in a weak gust; the winning kick in a strong gust — aim into the wind, let it curl in — was taken twice.

Planning can turn these flaws into bias: the optimizer ranks actions by the generator's imagination, so it may fall to over-optimism in low-support regions, or, actions whose outcomes the generator can only guess at (Figure 5). Global metrics, averaged under the logging distribution, will not uncover this.

2Policy Learning with Retargeted Counterfactual Generation

Two steps: ADIGen fixes the data problem; retargeting fixes the objective problem.

Step 1 — ADIGen: debias the generator

ADIGen (Automatic, Debiased, Invariant Generation) trains the generator on a three-part loss — guess, correct the guess with real data, stay invariant across environments. The weights \(\alpha\) are learned by Riesz regression, so no density ratio is ever inverted, even for high-dimensional joint interventions.

\[ \underbrace{\int \zeta_\psi(\theta;\bX,\ba)\,\nu(d\ba\mid\bX)}_{\text{guess under target law }\nu} + \underbrace{\alpha(\bA,\bX,E)}_{\text{Riesz weights}}\, \underbrace{\big[\ell(\theta;Y,\bA)-\zeta_\psi(\theta;\bX,\bA)\big]}_{\text{correct with real data}} + \lambda\,\underbrace{\mathrm{Pen}_{\mathcal E}(\theta)}_{\text{invariance}} \]

Three panels: kicks planned with a conditional model land wide of the goal; a reweighted model's kicks scatter; ADIGen's kicks land on the goal. — **What debiasing does.** Dots: where the ball really lands when a planner kicks using each trained model. The conditional model's kicks land wide, reweighting alone scatters them, ADIGen's land on goal.

Step 2 — Retargeting: spend the accuracy where the decision is

The generator only needs to be accurate where the policy will act, so retrain it there:

Propose. Imagine rollouts with ADIGen and fit a regularized pilot policy \(\tilde\pi \propto \pi_b\, e^{\hat\mu/\beta}\); the KL tilt \(\beta\) discourages over-extrapolation.
Retarget. Retrain the generator at the mixture \(\pi_\lambda = \lambda\,\tilde\pi + (1-\lambda)\,\pi_b\): optimism there is checked against real data, and the worst case stays controlled.
Refit. Fit the final policy \(\hat\pi\) from the retargeted generator's rollouts.

Three panels: a pilot policy's guess misses in reality where there is no data; retraining at the pilot teaches the generator the kick lands wide; the refit policy aims into the gust and curls in. — **What retargeting does.** Propose: the pilot's guess misses in reality, where there is no data. Retarget: retrained there, the generator learns the kick lands wide. Refit: the final policy aims into the gust and curls in.

3Why retargeting helps

Debiasing works (Thm. 1). ADIGen's error has the doubly robust product form: its two learned ingredients — the reweighting and the outcome model — multiply, so a mistake in one is forgiven if the other is accurate. Nothing is ever divided by a small probability, so training stays stable.
Regret balances two errors (Thm. 2). Regret is bounded by the generator's error inflated by two opposing terms. The transfer factor \(\rho(\hat\pi\,\|\,\nu)\) shrinks: as the training law \(\nu\) moves toward the pilot, the policy acts closer to where the generator trained. The critical radius grows: as \(\nu\) departs the logger, the weights \(\nu/\pi_b\) cut the effective sample to \(n/(1+\chi^2)\).
The mixture balances them (Thms. 3–4). At \(\pi_\lambda = \lambda\tilde\pi + (1-\lambda)\pi_b\) the transfer factor falls to \(1/\lambda\) while the weight cost shrinks to \(\lambda^2\chi^2(\tilde\pi\|\pi_b)\). Balancing the two with \(\lambda^\star\) beats the conditional approach by a factor \(\sqrt{\rho}/2\).

In one line: a conditional generator spends its error budget where the logs are; a retargeted one spends it where the decision will be made.

4Cell Painting (RxRx19a)

Setup. RxRx19a Cell Painting: microscopy images of SARS-CoV-2-infected VERO cells, each well treated with one of 32 compounds at one of 8 doses (256 arms). The decision is which compound and dose to give; its utility is a rescue score read from the images — 0 means the well looks infected, 1 means it looks like healthy mock cells. Rescue only happens at high dose, and that is where we bias the logs: wells at ≥ 1 µM are thinned to a fraction \(\eta\), low doses are left untouched. As \(\eta\) falls, the data says less and less about precisely the doses worth choosing.

Mock-infected VERO cells: dense monolayer. — **Mock**uninfected

Infected VERO cells, no compound: sparse, rounded cells. — **Mock**uninfected

Dose-response curves for GS-441524, Remdesivir and Chloroquine at the overlap chosen with the slider: real wells rise sharply at high dose; at low overlap the conditional and ADIGen policies stay flat while retargeted ADIGen rises toward the real peak. — **Figure 3: Dose–response as the logs thin out (drag the slider).** Black: real held-out wells; shaded: the ≥ 1 µM band the logger thins to a fraction \(\eta\). At η = 0.1 the conditional and ADIGen policies read the response as flat and miss the rescue; retargeted ADIGen still finds the peak.

Four panels against overlap for RxRx19a: regret, generator RMSE, peak recovery and regret reduction from retargeting, with bootstrap bands. — **Figure 4: RxRx19a across overlap η.** The sweep above summarized. Panels: regret; generator RMSE globally (solid) and at \(\pi^\star\) (dotted); peak recovery (predicted / true rescue height); regret reduction from retargeting. Retargeting helps most in the sparse-log regime (\(\eta \le 0.5\)) and costs nothing once the logs cover the high-dose band.

5Video generation (PhyWorld)

Setup. A PhyWorld ball receives a first push and a gust of wind; the decision is a second push, chosen so the ball lands on a goal line. Outcomes are rendered videos: each generator (conditional, ADIGen, retargeted ADIGen) is a flow-matching video DiT in the latent space of a frozen VAE, and the policy rolls out a bank of candidate pushes and picks the one it believes lands closest. The logs carry both biases of Section 1: the first push is correlated with the gust (confounding), and scenes with strong gusts (\(|g| \in [0.50, 0.85]\), beyond the training range \(|g| \le 0.25\)) enter the logs only with probability \(\eta\) (coverage).

5.1 Out-of-distribution decisions

Out of distribution: each scene has a gust stronger than any in training, seen in only 0.2% of logged scenes (\(\eta = 0.002\)). Top row: what each generator imagines for the push it chose. Bottom row: that push in the real simulator. The black line is the goal. Over the OOD scene set, mean regret drops from 0.74 (conditional) to 0.28 (retargeted).

Scene 236: all three generators imagine a hit; in the simulator only the retargeted push reaches the goal line. — **Figure 5: Every generator imagines success. Only the retargeted one achieves it.** Badges show the real landing miss of each method's chosen push. In every scene all three generators imagine a hit, but only the retargeted push lands. In scenes 18 and 230, debiasing alone (ADIGen) gets partway there and retargeting closes the gap. Scenes chosen for illustration.

Scene 28: matching imagined rollouts; in the simulator only the retargeted model lands on the goal. — **Figure 5: Every generator imagines success. Only the retargeted one achieves it.** Badges show the real landing miss of each method's chosen push. In every scene all three generators imagine a hit, but only the retargeted push lands. In scenes 18 and 230, debiasing alone (ADIGen) gets partway there and retargeting closes the gap. Scenes chosen for illustration.

5.2 Results across overlap and dimension

Line plots for PhyWorld: regret, model error, self-evaluation error and regret reduction against overlap and number of balls. — **Figure 6: PhyWorld at a matched training budget.** Left: overlap \(\eta\) (\(m = 1\)). Right: number of balls \(m\) (\(\eta = 0.005\)). At low overlap retargeting has the lowest regret, then ADIGen, then the conditional model. It also has the smallest self-evaluation error: it knows how good its own plan is.

Takeaways

Video rollouts can be generated for decision-making, not just realism: debiased, retargeted generators choose better actions from the same logs.
Treatment effects are recovered from thinned logs: the dose–response peak of all 3 active compounds, from only 10% of high-dose wells.
Retargeting spends the generator's accuracy on the decisions it supports; the price of staying conditional is a factor \(\propto \sqrt{\rho}\), the worst-case density ratio.

BibTeX

@misc{kim2026goodgenerators,
  title         = {Are Good Generators Good Decision-Makers? Policy Learning for
                   General Interventions via Retargeted Counterfactual Generation},
  author        = {Kim, Raphael C. and Zhu, Jingsen and Zabih, Ramin and
                   Santacatterina, Michele},
  year          = {2026},
  eprint        = {2606.07399},
  archivePrefix = {arXiv},
  primaryClass  = {stat.ML},
  doi           = {10.48550/arXiv.2606.07399}
}