Sample-or-Observe Sites

Some programs describe a process that also runs for real: an agent calling tools, a controller acting on a plant, a policy in an environment. Such a program has sites whose values the model draws when it simulates, and the world supplies when the program is executed. There, the handler takes the world's value and scores it like an observation.

This page shows the pattern, where each score belongs, and what to watch for. It is a legitimate choice for a handler to make, not a misuse of sample: a sample site says this value is random under the model, and the handler decides where the value comes from.

One program, three interpreters

The program is the same in all three. Only the handler changes.

The program

A tool-use loop. At each step a decision site picks a tool from the agent's habit. Unless the tool ends the episode, an outcome site says whether the call succeeded:

/// The tools the agent can call. Calling `RESPOND` ends the episode.
const TOOLS: [&str; 3] = ["lookup", "update", "respond"];
const RESPOND: usize = 2;

/// The agent's habit at a step: its probability of calling each tool.
fn habit(step: usize) -> Vec<f64> {
    if step == 0 {
        vec![0.8, 0.1, 0.1]
    } else {
        vec![0.3, 0.3, 0.4]
    }
}

/// How often each tool succeeds, as far as the model knows.
fn reliability(tool: usize) -> f64 {
    [0.95, 0.7, 1.0][tool]
}

/// One episode, from `step` on: a decision site picks a tool, and unless it
/// ends the episode, an outcome site says whether the call succeeded. The
/// outcome's distribution carries its tool as metadata, so a handler that
/// executes the call knows which call to make.
fn episode(step: usize, max_steps: usize) -> Model<Vec<(usize, bool)>> {
    if step == max_steps {
        return pure(Vec::new());
    }
    let decide = Categorical::new(habit(step)).unwrap();
    sample(addr!("decide", step), decide).bind(move |tool| {
        if tool == RESPOND {
            return pure(Vec::new());
        }
        let outcome = WithMeta::new(Bernoulli::new(reliability(tool)).unwrap(), tool);
        sample(addr!("outcome", step), outcome).bind(move |ok| {
            episode(step + 1, max_steps).map(move |rest| {
                let mut steps = vec![(tool, ok)];
                steps.extend(rest);
                steps
            })
        })
    })
}

The outcome's distribution is wrapped in WithMeta, which carries the tool as metadata. A handler that executes the call reads it at the site, so it needs no state of its own saying which decision the outcome belongs to.

Simulate

PriorHandler draws every site, decisions and outcomes alike. This is the model "dreaming" an episode:

    // Simulate: the model draws every site, decisions and outcomes alike.
    let mut rng = StdRng::seed_from_u64(7);
    let (dreamt, simulated) = run(
        PriorHandler {
            rng: &mut rng,
            trace: Trace::default(),
        },
        episode(0, max_steps),
    );
    println!("simulated: {}", show(&dreamt, max_steps));
    println!(
        "  log prior {:.3}, log likelihood {:.3}",
        simulated.log_prior, simulated.log_likelihood
    );

Execute

To execute, a handler makes the call and takes the world's answer. With Delegate, it overrides only the outcome sites and leaves the decisions to PriorHandler:

/// A stand-in for the real world: each call succeeds or fails as scripted.
struct World {
    script: Vec<bool>,
    calls: Vec<usize>,
}

impl World {
    fn call(&mut self, tool: usize) -> bool {
        let ok = self.script[self.calls.len() % self.script.len()];
        self.calls.push(tool);
        ok
    }
}

/// Executes a program's outcome sites against the world. Decision sites are
/// left to the inner `PriorHandler`, which draws them from the habit.
struct Execute<'a> {
    world: &'a mut World,
    /// Each executed step's surprise, `-ln p(outcome)` under the model.
    surprise: &'a mut Vec<f64>,
}

impl<'r, R: RngCore> Overrides<PriorHandler<'r, R>> for Execute<'_> {
    fn on_sample_bool(
        &mut self,
        inner: &mut PriorHandler<'r, R>,
        addr: &Address,
        dist: &dyn Distribution<bool>,
    ) -> bool {
        let tool = *dist
            .downcast_ref::<WithMeta<Bernoulli, usize>>()
            .expect("outcome sites carry their tool")
            .meta();
        // The value comes from outside, so it is scored like an observation:
        // into the likelihood, and recorded so that the run can be replayed
        // and audited.
        let ok = self.world.call(tool);
        let logp = dist.log_prob(&ok);
        inner.trace.log_likelihood += logp;
        inner
            .trace
            .insert_choice(addr.clone(), ChoiceValue::Bool(ok), logp);
        self.surprise.push(-logp);
        ok
    }
}

    // Execute: decisions still come from the habit, but each outcome is the
    // world's answer to a real call.
    let mut world = World {
        script: vec![true, false, true],
        calls: Vec::new(),
    };
    let mut surprise = Vec::new();
    let mut rng = StdRng::seed_from_u64(7);
    let handler = Delegate::with(
        PriorHandler {
            rng: &mut rng,
            trace: Trace::default(),
        },
        Execute {
            world: &mut world,
            surprise: &mut surprise,
        },
    );
    let (done, executed) = run(handler, episode(0, max_steps));
    println!("executed:  {}", show(&done, max_steps));
    println!(
        "  log prior {:.3} (the decisions), log likelihood {:.3} (the world's answers)",
        executed.log_prior, executed.log_likelihood
    );
    let per_step: Vec<String> = surprise.iter().map(|s| format!("{s:.2}")).collect();
    println!("  surprise per call, in nats: {}", per_step.join(", "));

The world's answer is recorded in the trace with its log-probability, like any choice, so the run can be replayed and audited. Its score goes to log_likelihood, because the model did not draw it: it is evidence about the model.

Surprise. -ln p(outcome) is the step's surprise, in nats: how unexpected the world's answer was under the model. A call the model thought would succeed 95% of the time costs 0.05 nats when it does, and 3.0 when it fails. Summed over an episode, surprise measures how well the model predicts the world; a run of high-surprise steps says the model no longer fits.

Audit

ScoreGivenTrace re-scores a recorded trace under the same program: here the run just executed, but it could be a log from production:

    // Audit: score the recorded run under the same program.
    let (_, audited) = run(
        ScoreGivenTrace {
            base: executed.clone(),
            trace: Trace::default(),
        },
        episode(0, max_steps),
    );
    for (addr, choice) in &executed.choices {
        assert!((audited.choices[addr].logp - choice.logp).abs() < 1e-12);
    }
    assert!((audited.total_log_weight() - executed.total_log_weight()).abs() < 1e-9);
    println!(
        "audited:   same {} sites, total log weight {:.3}; ScoreGivenTrace puts all of it in the prior ({:.3})",
        audited.choices.len(),
        audited.total_log_weight(),
        audited.log_prior
    );

Each site's Choice::logp and the total log weight match the executed run. The split does not: ScoreGivenTrace treats every sample site alike and puts all of it in log_prior. Compare per-site scores or totals across interpreters, not the prior and likelihood separately.

The accounting: prior or likelihood?

The rule: a value the handler draws from the site's distribution is scored in log_prior; a value the world supplies is scored in log_likelihood.

It matters wherever the two are used apart:

• Importance weights. With the prior as the proposal, a run's weight is its likelihood (plus factors): the prior terms cancel against the proposal. If the world's answers were put in log_prior, every executed run would get the same weight whatever the world said.
• Model evidence. Averaging the likelihood over runs drawn from the prior estimates the evidence, p(data). The world's answers are the data.
• Decisions. A decision the handler drew from the habit is part of the proposal, and its log-probability is the propensity of the action taken. Logged, it is what off-policy evaluation reweights by. A decision made outside the model, such as an LLM's choice, is a value from the world like an outcome. Score it as likelihood to measure how well the habit predicts the agent.

When to use observe instead

Use observe when the value is known before the run: data that the program is conditioned on, fixed when the model is built.

observe(addr!("y", i), Normal::new(mu, sigma).unwrap(), data[i])

A sample-or-observe site is for values that exist only during the run, because they depend on what was decided earlier. The model cannot build an observe for an outcome of a call it has not made yet.

Pitfalls

• Values outside the support. The world can return a value the model thought impossible. A tool whose reliability is 1.0 that fails scores -inf, and the run's weight with it. Keep probabilities off 0 and 1 (clamp to [1e-6, 1 - 1e-6]) where the world can surprise the model, or treat -inf as a signal to stop trusting the model.
• Addresses must not depend on who supplied a value. A site's address must be the same whether it is simulated, executed or audited, or ScoreGivenTrace will not find the recorded value. Build addresses from the program's own state (the step index), never from where the value came from.
• State between sites belongs in the program. If an outcome site needs to know which decision it belongs to, put that in the site: in its address or its distribution's metadata (WithMeta), as above. A variable kept in the handler (last_tool) breaks as soon as a site is skipped or replayed.
• Record what you score. A handler that supplies a value must insert the choice into its trace with its log-probability. Otherwise replay samples a fresh value at that site, and an audit panics on the missing value.

See also

• Custom Handlers: writing handlers, and Delegate for overriding a few sites.
• RFC-001's spike, whose harness handler executes outcome sites this way, and stretto's audit, which scores an agent's recorded sessions under a flow with ScoreGivenTrace.