Optimize Rollout Sampling

Optimize Rollout Sampling

Treat rollout work as distribution design, selection, and estimation—not as a single temperature knob.

Start with the right unit

State the sampling unit before proposing an algorithm:

• Token: one vocabulary item inside a model call.
• Reasoning/action step: a thought chunk, tool call, or environment action.
• Trajectory: the complete sequence from prompt to terminal outcome.
• Population: a set of trajectories used for voting, ranking, resampling, or a policy update.

Do not call step-local temperature scaling “trajectory-level power sampling.” Read references/taxonomy.md whenever power sampling, importance sampling, SMC, or off-policy correction is involved.

Run the design workflow

1. Formalize the target

Write down:

1. Base/behavior distribution p(τ|x) or μ(τ|x).
2. Desired target distribution or decision rule.
3. Available signal: none, answer agreement, outcome verifier, process score, environment reward, or human preference.
4. Access: text-only API, token log-probabilities, full logits, model weights, environment snapshots, or a world model.
5. Budget: model calls, generated tokens, wall time, verifier calls, and environment interactions.
6. Goal: maximize single returned answer, coverage/pass@k, expected reward, training signal quality, diversity, or unbiased evaluation.

Keep these goals separate. A sampler that improves oracle coverage can worsen pass@1; a selector that improves measured reward can overfit the verifier; a biased low-variance estimator can be useful for optimization but invalid for evaluation.

2. Select the smallest adequate method

Prefer a baseline ladder: greedy/pass@1 → IID repeated sampling/pass@k → self-consistency or BoN → stepwise search → MCMC/SMC or adaptive tree search. Add a more complex method only when the previous rung exposes a specific limitation.

3. Build a matched-budget experiment

Always compare methods at matched cost. Record at least:

• prompt/task ID, seed, policy/model version, sampler configuration;
• full trajectory or a stable hash, terminal status, reward/correctness;
• generated tokens, wall time, model calls, verifier calls, tool/environment calls;
• behavior log-probability and target log-probability when using IS;
• selected candidate and the rule that selected it.

Use references/protocols.md for experiment grids, agent-state branching, and reporting. Use scripts/analyze_rollouts.py on JSONL logs. Run scripts/power_vs_temperature.py when explaining or checking sequence power.

4. Diagnose before concluding

Report all of the following when available:

• pass@1 and the unbiased pass@k estimator;
• oracle@N separately from selected-answer accuracy;
• majority/self-consistency and Best-of-N accuracy;
• unique-answer and unique-trajectory rates;
• reward-score calibration and reward hacking indicators;
• mean/median tokens, model calls, verifier calls, wall time, and environment calls;
• IS effective sample size (ESS), maximum weight, clipping rate, and estimator type;
• mean across task IDs with task-level bootstrap intervals, not a confidence interval over correlated rollouts.

5. Produce an actionable answer

Return these sections unless the user asks for another format:

1. Formal objective and access assumptions
2. Recommended sampler and why
3. Minimal baseline and ablation grid
4. Implementation sketch tied to an actual framework
5. Metrics, failure criteria, and expected cost
6. What the result would and would not prove

For literature reviews, read references/papers.md. For implementation work, also read references/code-notes.md.

Enforce the non-negotiable checks

• Distinguish q_seq(τ) ∝ p(τ)^α from token-local ancestral sampling at temperature 1/α; prefix normalizers generally make them different.
• Distinguish sampling from reweighting. Importance sampling estimates an expectation under a target distribution; it does not itself produce independent target samples.
• Require target support to be contained in proposal support before using IS. Treat missing behavior log-probabilities as a blocker for exact IS.
• Compute products and weights in log space. Report ESS and weight concentration; never hide them behind a clipped estimate.
• Do not reuse a verifier for search and call its score an unbiased final evaluation. Evaluate with rules, tests, humans, or a separately trained held-out judge.
• Do not average pass@k over tasks with fewer than k samples without explicitly reporting the reduced task set.
• Keep trajectory length normalization explicit. Sum log-probabilities corresponds to sequence probability; mean log-probability changes the target and introduces a different length preference.
• In agent trees, restore both textual context and environment state at every branch. Replaying only the transcript is insufficient for a stochastic or stateful environment.
• Treat correlated branches, shared prefixes, resampled particles, and repeated revisions as non-IID. Do not use the IID pass@k interpretation for them.
• Separate the rollout policy used to collect training data from the inference policy. Record policy lag and the granularity of any correction.

Use the bundled tools

Analyze rollout logs:

python scripts/analyze_rollouts.py rollouts.jsonl --k 1,2,4,8 --power-alpha 2

Expected core fields are prompt_id plus one or more of correct, reward, answer, score, base_logp, behavior_logp, and target_logp. See references/protocols.md for the schema.

Demonstrate the power/temperature distinction:

python scripts/power_vs_temperature.py --alpha 2

Keep the review current

Treat references/papers.md as a curated map through 2026-08-25, not a claim of literal exhaustiveness. For a new review, search forward from the canonical works, prefer primary papers and official repositories, record the search cutoff, and label unreviewed or unofficial code clearly.