Build a Budget-Aware Discovery Agent

Build a Bayesian agent that selects experiments under a fixed budget.

Introduction

30 Second Summary

Research rarely gives you unlimited trials. With only a small budget, each choice determines how much you can discover.

In this project, you will build a budget-aware autonomous discovery agent in Python. Your agent will use Bayesian optimization to choose promising experiments under noise and a fixed budget.

What You'll Build

You will run one benchmark to watch three competing strategies reveal their strengths across a fixed experiment budget.

By the end of this project, you'll have:

  • Run a finite-budget campaign that reveals one noisy outcome at a time while preventing repeat selections.
  • Inspect a Gaussian-process surrogate to compare posterior means with uncertainty across every untested candidate.
  • Benchmark random, greedy, and uncertainty-aware selection under identical budgets across repeated seeds and noise levels. Produce outputs/regret_by_budget.png, outputs/calibration.png, and outputs/findings.md as evidence for your conclusions.
  • Secret Mission: Add Thompson-style posterior sampling. Compare its high-noise performance against the uncertainty-aware policy.

Are there any prerequisites?

You should be comfortable with graduate-level probability or statistics. You also need experience with Python functions, loops, and basic train/test evaluation.

Before We Start

A discovery campaign becomes useful when you treat each result as evidence about where to experiment next. This checkpoint locks in what your agent will optimize under a fixed experiment budget.

Set Up the Reproducible Environment

A benchmark is only useful when another researcher can recreate the environment that produced it. Unknown package versions can turn a strategy comparison into evidence about different software stacks.

In this step, you'll confirm the Python runtime. You'll also build an isolated workspace with pinned dependencies in VS Code.

In this step, get ready to:
  • Confirm that Python 3.11 or later is available.
  • Create the budgeted-discovery-agent workspace in VS Code.
  • Install the pinned dependencies inside an isolated virtual environment.
Confirm your Python version and create the workspace

The pinned version of BoTorch requires Python 3.11 or later. Checking the runtime first prevents dependency failures later.

You'll run the check in PowerShell. This keeps the environment setup aligned with your Windows machine.

  • Press the Windows key to open Windows search.
  • Type PowerShell and press Enter.
  • Check your installed Python version by running:
python --version

What does this command check?

The command prints the Python version connected to the python command. That runtime creates your virtual environment in the next substep.

✔️ I see version 3.11 or higher

Your Python runtime meets the project's requirement. You can continue with the workspace setup.

ⓧ I see an older version

The installed runtime is below the project's minimum version. Install a supported release before creating the virtual environment.

  • Open the official Python releases for Windows page.
  • Choose a Python 3.11 or later release.
  • Select Download Windows installer (64-bit).
  • Run the downloaded installer.
  • Select Install.
  • Close the current PowerShell window after installation finishes.
  • Launch PowerShell again using Windows search.
  • Confirm the updated runtime by running:
python --version

What should the new check prove?

The output should now report Python 3.11 or later. This proves that the python command points to a compatible runtime.

ⓧ Command not found

Windows cannot currently find a Python runtime through the python command. Install a supported release from the official source.

  • Open the official Python releases for Windows page.
  • Choose a Python 3.11 or later release.
  • Select Download Windows installer (64-bit).
  • Run the downloaded installer.
  • Select Install.
  • Close the current PowerShell window after installation finishes.
  • Launch PowerShell again using Windows search.
  • Confirm that Windows can now find Python by running:
python --version

What should the new check prove?

The output should report Python 3.11 or later. This confirms that the runtime is available from PowerShell.

A VS Code workspace gives your files a single project boundary. You'll create that boundary on your Desktop so the folder is easy to find.

  • Press the Windows key to open Windows search.
  • Type VS Code and press Enter.

You'll see the VS Code window. The next actions create your project folder directly from its folder picker.

What if Workspace Trust appears?

VS Code may display a Workspace Trust dialog after you open the folder. You created this empty workspace yourself, so you can trust its contents.

  • Select File from the top menu.
  • Select Open Folder....
  • Select your Desktop in the folder picker.
  • Select New Folder.
  • Enter budgeted-discovery-agent as the folder name.
  • Select Select Folder.
  • Select Yes, I trust the authors if Workspace Trust appears.

You'll see budgeted-discovery-agent at the top of the Explorer view. Your research workspace now has a clear home.

Create the isolated dependency environment

A virtual environment keeps this project's packages separate from other Python projects. The requirements.txt file records the exact dependency versions that belong inside it.

The environment contains PyTorch, BoTorch, GPyTorch, and Matplotlib. Each library has a fixed role in the later benchmark.

  • Select New File in the Explorer view.
  • Enter requirements.txt as the file name.
  • Press Enter.
  • Add the pinned dependencies by pasting:
torch==2.15.0
gpytorch==1.15.2
botorch==0.18.1
matplotlib==3.11.2

What do these pins control?

  • The torch==2.15.0 pin fixes the tensor runtime used by every campaign.
  • The gpytorch==1.15.2 pin fixes the Gaussian-process modeling layer.
  • The botorch==0.18.1 pin fixes the Bayesian optimization interface.
  • The matplotlib==3.11.2 pin fixes the plotting library used for benchmark figures.
  • Press Ctrl+S to save requirements.txt.

You'll see requirements.txt in the Explorer view without an unsaved-change indicator. The dependency record is now stored on disk.

Missing the requirements file?

Check that requirements.txt appears directly inside budgeted-discovery-agent. A nested copy would place the dependency record outside the folder used by PowerShell.

Help me find a misplaced requirements file in VS Code.

✔️ Awesome, I've got everything!

Great. Confirm that requirements.txt is saved directly inside budgeted-discovery-agent.

ⓧ I'd like to double check the full code

Compare your complete requirements.txt file with this reference:

torch==2.15.0
gpytorch==1.15.2
botorch==0.18.1
matplotlib==3.11.2

The file should contain these four lines in this order. It should not contain any extra dependency entries.

The dependency file describes the environment. PowerShell now creates that environment inside the same project folder.

  • Switch back to the PowerShell window from earlier.
  • Move into the Desktop project folder by running:
Set-Location -Path "$env:USERPROFILE\Desktop\budgeted-discovery-agent" -PassThru

What does this command do?

The command changes PowerShell's current location to budgeted-discovery-agent. The -PassThru flag prints the resulting path so you can verify the move.

You'll see a path ending in Desktop\budgeted-discovery-agent. This confirms that the next commands affect the correct folder.

  • Create the project-local virtual environment by running:
python -m venv .venv

What does this command create?

Python creates a .venv directory inside the project folder. Packages installed after activation stay within this environment.

  • Confirm that the environment directory exists by running:
Get-ChildItem -Force

What should the listing contain?

The listing should include .venv and requirements.txt. Together they show that the isolated environment sits beside its dependency record.

You'll see both entries in the project folder listing. The environment exists before any third-party package is installed.

  • Activate the virtual environment by running:
.\.venv\Scripts\Activate.ps1

What does activation change?

Activation points the current PowerShell session at the Python executable inside .venv. Package installation now targets this project instead of the system runtime.

You'll see the environment name at the beginning of the PowerShell prompt. That prompt marker confirms that .venv is active.

Activation blocked?

PowerShell may block the activation script because of the machine's script-execution policy. Check the policy configured by your administrator before changing any security setting.

Help me activate this virtual environment safely.

The installation can take several minutes while Python downloads the pinned packages. A quiet pause during this command does not mean the process has stalled.

  • Install every pinned dependency from requirements.txt by running:
python -m pip install -r requirements.txt

What does this installation use?

The pip installer reads requirements.txt. It installs each recorded version into the active .venv environment.

You'll return to the active PowerShell prompt after the installation completes without a dependency failure. That's the longest setup task finished: the benchmark libraries now live in one isolated environment.

Dependency installation failed?

Confirm that the PowerShell prompt still shows the active environment. Also confirm that requirements.txt contains the four exact pins from the reference tab.

Help me troubleshoot the pinned dependency installation.

Verify the pinned imports

Successful installation output shows that package files were written. An import check proves that the active Python runtime can load all four libraries together.

Version pinning fixes the dependency definitions used by the benchmark. Isolation prevents packages from unrelated projects from silently changing those definitions.

Before you run the final check, do you expect all four imports to complete in the active environment?

  • Verify the complete environment by running:
python -c "import torch, botorch, gpytorch, matplotlib; print('Environment ready')"

What does the final check prove?

Python imports PyTorch, BoTorch, GPyTorch, and Matplotlib in one process. The print statement runs only after every import succeeds.

You'll see Environment ready without an import error. Excellent work. Your reproducible environment is ready to support a controlled benchmark.

Your workspace now has a compatible runtime with isolated dependencies. Next, you'll build the noisy discovery campaign that gives the agent a finite experiment budget.

Build the Noisy Discovery Campaign

Your reproducible environment is ready. Now the discovery agent needs candidates it can test under a fixed experiment budget.

Each observed response includes noise. A limited budget makes every selection consequential.

You'll build a synthetic candidate library with two peaks. Then you'll let a random policy reveal noisy outcomes without testing one candidate twice.

In this step, get ready to:
  • Define a two-parameter candidate grid with a hidden response surface.
  • Reveal noisy outcomes without selecting duplicate candidates.
  • Track simple regret across a random campaign.
Build the candidate library and hidden response

A candidate library turns a continuous design space into a finite set of possible experiments. The hidden response gives each candidate a noise-free score that acts as ground truth.

The campaign uses PyTorch tensors for its candidate coordinates. Double-precision values keep the numerical calculations consistent throughout the benchmark.

  • Create discovery_agent.py inside budgeted-discovery-agent using the new-file icon in the VS Code file sidebar.
  • Add the shared imports plus campaign configuration by pasting this code into discovery_agent.py:
from pathlib import Path

import matplotlib.pyplot as plt
import torch
from botorch.fit import fit_gpytorch_mll
from botorch.models import SingleTaskGP
from gpytorch.mlls import ExactMarginalLogLikelihood


torch.set_default_dtype(torch.float64)

GRID_SIZE = 15
ADDITIONAL_BUDGET = 8
SEEDS = (3, 11, 29)
NOISE_LEVELS = (0.05, 0.15)
POLICIES = ("random", "greedy", "ucb")
UCB_BETA = 2.0
OUTPUT_DIR = Path("outputs")

INITIAL_DESIGNS = torch.tensor(
    [
        [0.10, 0.10],
        [0.15, 0.30],
        [0.30, 0.15],
        [0.35, 0.35],
        [0.25, 0.25],
    ]
)

What does this setup do?

  • The imports provide the numerical tools plus the modeling tools used by the campaign file.
  • The campaign constants define the grid size plus the additional experiment budget.
  • The benchmark constants establish the seeds plus the noise conditions used across policies.
  • The INITIAL_DESIGNS tensor places five starting experiments around the lower-left region.
  • Save discovery_agent.py.
  • Verify that the imports plus constants load by running this command in the activated PowerShell terminal from earlier:
python discovery_agent.py

PowerShell returns to the prompt without a traceback. This confirms that the campaign file can load its dependencies plus configuration.

Seeing an import or syntax error?

  • Confirm that your PowerShell prompt still begins with (.venv).
  • Compare the import names plus brackets with the code block above.
  • Save discovery_agent.py before running it again.

Help me debug the initial discovery_agent.py setup.

The first function creates every pair of parameter values on a square grid. With GRID_SIZE set to 15, the campaign receives 225 possible experiments.

  • Add build_candidate_library below INITIAL_DESIGNS by pasting this code:
def build_candidate_library(grid_size=GRID_SIZE):
    axis = torch.linspace(0.0, 1.0, steps=grid_size)
    first, second = torch.meshgrid(axis, axis, indexing="ij")
    return torch.stack([first.reshape(-1), second.reshape(-1)], dim=-1)

What does this function do?

  • The axis divides each parameter range into evenly spaced values from 0.0 through 1.0.
  • The mesh grid creates every combination of the two parameter values.
  • The stacked result stores one candidate per row with two coordinates per candidate.
  • Save discovery_agent.py.
  • Check that the candidate-library function loads by running:
python discovery_agent.py

PowerShell returns to the prompt without a traceback. The candidate-grid function is now valid Python.

Seeing an error in the grid function?

  • Check that build_candidate_library starts at the left edge of the file.
  • Confirm that each line inside the function uses four spaces of indentation.
  • Check that indexing="ij" includes matching quotation marks.

Help me fix my candidate-library function.

The response surface gives the campaign a lower local peak plus a stronger distant peak. A diagonal ridge adds another structure that the selection policies must learn from limited observations.

  • Add hidden_response below build_candidate_library by pasting this code:
def hidden_response(candidates):
    first = candidates[:, 0]
    second = candidates[:, 1]
    local_peak = 0.85 * torch.exp(
        -((first - 0.25) ** 2 + (second - 0.25) ** 2) / 0.018
    )
    global_peak = 1.20 * torch.exp(
        -((first - 0.78) ** 2 + (second - 0.72) ** 2) / 0.012
    )
    ridge = 0.15 * torch.exp(-((first - second) ** 2) / 0.08)
    return local_peak + global_peak + ridge

How is the response surface shaped?

  • The local peak is centered near (0.25, 0.25) with an amplitude of 0.85.
  • The global peak is centered near (0.78, 0.72) with an amplitude of 1.20.
  • The ridge adds smaller responses near candidates whose two coordinates are similar.
  • Save discovery_agent.py.
  • Check that the hidden-response function loads by running:
python discovery_agent.py

PowerShell returns to the prompt without a traceback. The synthetic response surface is ready for candidate evaluation.

Seeing a bracket or indentation error?

  • Compare the closing parentheses after each peak expression with the code block above.
  • Confirm that all three response components remain inside hidden_response.
  • Check that the final return line has one indentation level.

Help me debug my hidden response surface.

The starting design coordinates must map to actual rows in the finite candidate library. The nearest-grid calculation performs that mapping before any new experiment is selected.

  • Add get_initial_indices below hidden_response by pasting this code:
def get_initial_indices(candidates):
    distances = torch.cdist(INITIAL_DESIGNS, candidates)
    return distances.argmin(dim=1).tolist()

How are the starting experiments matched?

  • The distance matrix compares each requested starting coordinate with every candidate.
  • The minimum-distance index identifies the closest available grid point for each initial design.
  • The returned list becomes the campaign's first set of observed candidate indices.
  • Save discovery_agent.py.
  • Check that the initial-design mapping loads by running:
python discovery_agent.py

PowerShell returns to the prompt without a traceback. The campaign can now map its lower-left starting designs onto the candidate grid.

  • Inspect the center coordinates inside local_peak.
  • Inspect the center coordinates inside global_peak.
  • Compare the amplitudes assigned to the two peaks.
Reveal noisy outcomes without duplicates

A campaign must distinguish tested candidates from remaining candidates. That separation prevents the random policy from spending its budget on duplicate experiments.

  • Add available_indices below get_initial_indices by pasting this code:
def available_indices(candidate_count, observed_indices):
    observed = set(observed_indices)
    return [index for index in range(candidate_count) if index not in observed]

How does the candidate mask work?

  • The set provides fast membership checks for previously observed indices.
  • The returned list contains only indices that have not been observed.
  • Every selection policy can use this list to avoid duplicates.
  • Save discovery_agent.py.
  • Check that the candidate-mask function loads by running:
python discovery_agent.py

PowerShell returns to the prompt without a traceback. The campaign can now separate observed candidates from available candidates.

Seeing a list-comprehension error?

  • Check that observed is created before the return statement.
  • Confirm that the condition reads if index not in observed.

Help me fix my available-candidate mask.

A virtual measurement combines the candidate's hidden truth with random noise. A seeded generator makes the noisy sequence repeatable for a given campaign seed.

  • Add observe below available_indices by pasting this code:
def observe(truth, index, noise_std, generator):
    noise = noise_std * torch.randn((), generator=generator)
    return truth[index] + noise

What does an observation contain?

  • The hidden truth provides the noise-free response for the selected candidate.
  • The noise standard deviation controls how strongly the measurement can vary.
  • The generator makes the random draw reproducible when the same seed is reused.
  • Save discovery_agent.py.
  • Check that the observation function loads by running:
python discovery_agent.py

PowerShell returns to the prompt without a traceback. The campaign can now reveal a noisy measurement for any selected candidate.

Seeing an error in the noise draw?

  • Confirm that torch.randn receives an empty shape written as ().
  • Check that the supplied generator uses the generator keyword.

Help me debug my noisy observation function.

The random policy chooses a position within the remaining-candidate list. It then maps that position back to the original candidate-library index.

  • Add choose_random below observe by pasting this code:
def choose_random(candidate_count, observed_indices, generator):
    remaining = available_indices(candidate_count, observed_indices)
    position = torch.randperm(len(remaining), generator=generator)[0].item()
    return remaining[position]

How does random selection avoid duplicates?

  • The policy starts from the filtered list produced by available_indices.
  • The random permutation selects one position inside that filtered list.
  • The final lookup returns the corresponding index from the original candidate library.
  • Save discovery_agent.py.
  • Check that the random-selection function loads by running:
python discovery_agent.py

PowerShell returns to the prompt without a traceback. Random selection now draws only from candidates that remain unobserved.

Could a duplicate candidate still be selected?

  • Confirm that choose_random passes observed_indices into available_indices.
  • Check that the function returns remaining[position].

Help me trace duplicate random selections.

Run the random campaign and measure regret

Simple regret measures the gap between the best candidate in the full library and the best candidate tested so far. A value closer to zero means the campaign has discovered a stronger candidate.

  • Add simple_regret below choose_random by pasting this code:
def simple_regret(truth, observed_indices):
    observed_tensor = torch.tensor(observed_indices)
    best_possible = truth.max()
    best_found = truth[observed_tensor].max()
    return (best_possible - best_found).item()

What does simple regret reveal?

  • The global maximum represents the strongest candidate available anywhere in the library.
  • The observed maximum represents the strongest candidate found within the current budget.
  • Their difference records how much opportunity remains undiscovered.
  • Save discovery_agent.py.
  • Check that the regret function loads by running:
python discovery_agent.py

PowerShell returns to the prompt without a traceback. The campaign can now score discovery progress after every experiment.

Seeing an indexing error in regret?

  • Confirm that observed_indices is converted into a tensor before indexing truth.
  • Check that both maximum operations include parentheses.

Help me debug my simple-regret calculation.

The campaign loop begins with five noisy initial observations. It then spends eight additional experiments through the random policy while recording regret after every selection.

  • Add run_campaign below simple_regret by pasting this code:
def run_campaign(policy, seed, noise_std):
    torch.manual_seed(seed)
    noise_generator = torch.Generator().manual_seed(seed)
    choice_generator = torch.Generator().manual_seed(seed + 1000)

    candidates = build_candidate_library()
    truth = hidden_response(candidates)
    observed_indices = get_initial_indices(candidates)
    observed_values = [
        observe(truth, index, noise_std, noise_generator)
        for index in observed_indices
    ]
    regrets = [simple_regret(truth, observed_indices)]

    for _ in range(ADDITIONAL_BUDGET):
        next_index = choose_random(
            len(candidates), observed_indices, choice_generator
        )
        observed_indices.append(next_index)
        observed_values.append(
            observe(truth, next_index, noise_std, noise_generator)
        )
        regrets.append(simple_regret(truth, observed_indices))

    return torch.tensor(regrets)

How does the campaign progress?

  • Separate generators control measurement noise plus random candidate selection.
  • The initial design produces the campaign's first five noisy observations.
  • Each loop iteration selects one remaining candidate before storing its noisy outcome.
  • The returned tensor contains the initial regret plus one value for every additional experiment.
  • Save discovery_agent.py.

Before you run the campaign, ask yourself whether eight random experiments are enough to reach the distant upper-right peak.

  • Add a temporary print call below run_campaign that displays the value returned by run_campaign("random", 3, 0.15).
  • Save discovery_agent.py.
  • Run the random campaign in the activated PowerShell terminal with this command:
python discovery_agent.py

Your first discovery campaign is working. You'll see a tensor containing nine regret values.

The first value describes the initial design. Each remaining value records the result after one of the eight additional experiments.

Missing the nine-value tensor?

  • Confirm that the temporary print call sits below the complete run_campaign function.
  • Check that ADDITIONAL_BUDGET remains set to 8.
  • Confirm that regrets.append remains inside the campaign loop.

Help me debug the random campaign output.

  • Delete the temporary print call from the bottom of discovery_agent.py.
  • Save discovery_agent.py.

✔️ Awesome, I've got everything!

Great work. Your saved discovery_agent.py now contains the candidate library plus the complete random campaign without the temporary print call.

ⓧ I'd like to double check the full code

Compare your saved discovery_agent.py with this complete version:

from pathlib import Path

import matplotlib.pyplot as plt
import torch
from botorch.fit import fit_gpytorch_mll
from botorch.models import SingleTaskGP
from gpytorch.mlls import ExactMarginalLogLikelihood


torch.set_default_dtype(torch.float64)

GRID_SIZE = 15
ADDITIONAL_BUDGET = 8
SEEDS = (3, 11, 29)
NOISE_LEVELS = (0.05, 0.15)
POLICIES = ("random", "greedy", "ucb")
UCB_BETA = 2.0
OUTPUT_DIR = Path("outputs")

INITIAL_DESIGNS = torch.tensor(
    [
        [0.10, 0.10],
        [0.15, 0.30],
        [0.30, 0.15],
        [0.35, 0.35],
        [0.25, 0.25],
    ]
)


def build_candidate_library(grid_size=GRID_SIZE):
    axis = torch.linspace(0.0, 1.0, steps=grid_size)
    first, second = torch.meshgrid(axis, axis, indexing="ij")
    return torch.stack([first.reshape(-1), second.reshape(-1)], dim=-1)


def hidden_response(candidates):
    first = candidates[:, 0]
    second = candidates[:, 1]
    local_peak = 0.85 * torch.exp(
        -((first - 0.25) ** 2 + (second - 0.25) ** 2) / 0.018
    )
    global_peak = 1.20 * torch.exp(
        -((first - 0.78) ** 2 + (second - 0.72) ** 2) / 0.012
    )
    ridge = 0.15 * torch.exp(-((first - second) ** 2) / 0.08)
    return local_peak + global_peak + ridge


def get_initial_indices(candidates):
    distances = torch.cdist(INITIAL_DESIGNS, candidates)
    return distances.argmin(dim=1).tolist()


def available_indices(candidate_count, observed_indices):
    observed = set(observed_indices)
    return [index for index in range(candidate_count) if index not in observed]


def observe(truth, index, noise_std, generator):
    noise = noise_std * torch.randn((), generator=generator)
    return truth[index] + noise


def choose_random(candidate_count, observed_indices, generator):
    remaining = available_indices(candidate_count, observed_indices)
    position = torch.randperm(len(remaining), generator=generator)[0].item()
    return remaining[position]


def simple_regret(truth, observed_indices):
    observed_tensor = torch.tensor(observed_indices)
    best_possible = truth.max()
    best_found = truth[observed_tensor].max()
    return (best_possible - best_found).item()


def run_campaign(policy, seed, noise_std):
    torch.manual_seed(seed)
    noise_generator = torch.Generator().manual_seed(seed)
    choice_generator = torch.Generator().manual_seed(seed + 1000)

    candidates = build_candidate_library()
    truth = hidden_response(candidates)
    observed_indices = get_initial_indices(candidates)
    observed_values = [
        observe(truth, index, noise_std, noise_generator)
        for index in observed_indices
    ]
    regrets = [simple_regret(truth, observed_indices)]

    for _ in range(ADDITIONAL_BUDGET):
        next_index = choose_random(
            len(candidates), observed_indices, choice_generator
        )
        observed_indices.append(next_index)
        observed_values.append(
            observe(truth, next_index, noise_std, noise_generator)
        )
        regrets.append(simple_regret(truth, observed_indices))

    return torch.tensor(regrets)

Your controlled discovery campaign can now reveal noisy evidence within a fixed budget. Next, you'll fit a surrogate model and watch pure exploitation focus too heavily on the local peak.

Fit the Surrogate and Test Greedy Selection

Your random campaign can reveal noisy outcomes without repeating an experiment. Now the campaign needs a model that turns those observations into predictions across the candidate library.

You will fit a Gaussian-process surrogate with BoTorch. A greedy policy will trust the largest posterior mean. The key question is whether that confidence leads the campaign to the best region.

In this step, get ready to:
  • Fit a surrogate to the observed coordinates and noisy responses.
  • Inspect posterior means and standard deviations for untested candidates.
  • Evaluate greedy selection with simple regret and posterior coverage.
Fit the Gaussian-process surrogate

A surrogate estimates the hidden response from the experiments observed so far. Its training coordinates use the shape [number_of_observations, 2]. Its observed responses use the shape [number_of_observations, 1].

  • In discovery_agent.py, find the simple_regret function.
  • Add the surrogate-fitting function directly below it by pasting this code:
def fit_surrogate(candidates, observed_indices, observed_values):
    index_tensor = torch.tensor(observed_indices)
    train_x = candidates[index_tensor]
    train_y = torch.stack(observed_values).reshape(-1, 1)
    model = SingleTaskGP(train_X=train_x, train_Y=train_y)
    mll = ExactMarginalLogLikelihood(model.likelihood, model)
    fit_gpytorch_mll(mll)
    return model

What does this code do?

  • The index tensor selects the measured coordinates from the full candidate library.
  • The reshaped response tensor gives each experiment one model output.
  • The SingleTaskGP model learns a posterior over the hidden response.
  • The ExactMarginalLogLikelihood objective lets GPyTorch tune the model parameters from the observed data.
  • Save discovery_agent.py.
  • Fit the surrogate against the five initial designs by running this command in the active PowerShell terminal:
python -c "from discovery_agent import build_candidate_library, fit_surrogate, get_initial_indices, hidden_response
candidates = build_candidate_library()
truth = hidden_response(candidates)
observed_indices = get_initial_indices(candidates)
observed_values = [truth[index] for index in observed_indices]
model = fit_surrogate(candidates, observed_indices, observed_values)
print(model.__class__.__name__)"

What does this check prove?

You should see SingleTaskGP in the terminal. That output confirms the five observed coordinates reached the model fitting code successfully.

Surrogate fit failing?

Confirm that the PowerShell prompt still begins with (.venv). Check that train_y ends with reshape(-1, 1).

Ask for help with the model input shapes.

The fitted model can now describe every untested candidate through a posterior mean and standard deviation. The mean represents predicted response. The standard deviation represents remaining uncertainty.

  • Add the posterior-summary function directly below fit_surrogate by pasting this code:
def posterior_summary(model, candidates, candidate_indices):
    index_tensor = torch.tensor(candidate_indices)
    with torch.no_grad():
        posterior = model.posterior(candidates[index_tensor])
        mean = posterior.mean.squeeze(-1)
        standard_deviation = torch.sqrt(
            posterior.variance.squeeze(-1).clamp_min(1e-12)
        )
    return mean, standard_deviation

How is the posterior summarized?

  • The candidate indices restrict prediction to the requested coordinates.
  • The torch.no_grad() context keeps this prediction check focused on inference.
  • The square root converts posterior variance into standard deviation.
  • The lower clamp keeps the square-root calculation numerically stable.
  • Save discovery_agent.py.
  • Inspect the posterior tensor shapes for every untested candidate by running this command:
python -c "from discovery_agent import available_indices, build_candidate_library, fit_surrogate, get_initial_indices, hidden_response, posterior_summary
candidates = build_candidate_library()
truth = hidden_response(candidates)
observed_indices = get_initial_indices(candidates)
observed_values = [truth[index] for index in observed_indices]
model = fit_surrogate(candidates, observed_indices, observed_values)
remaining = available_indices(len(candidates), observed_indices)
mean, standard_deviation = posterior_summary(model, candidates, remaining)
print(mean.shape)
print(standard_deviation.shape)"

What should you see?

You should see torch.Size([220]) twice. The five initial designs leave 220 untested points in the 225-candidate grid.

Posterior shapes look different?

Check that posterior.mean and posterior.variance both use squeeze(-1). Confirm that remaining comes from available_indices.

Help me diagnose my posterior tensor shapes.

Use posterior means for greedy selection

A greedy policy assigns every remaining candidate its posterior mean as a score. The candidate with the largest score becomes the next experiment.

  • Add the shared proposal function directly below posterior_summary by pasting this code:
def propose_next(
    policy,
    candidates,
    observed_indices,
    observed_values,
    choice_generator,
):
    if policy == "random":
        return choose_random(len(candidates), observed_indices, choice_generator)

    remaining = available_indices(len(candidates), observed_indices)
    model = fit_surrogate(candidates, observed_indices, observed_values)
    mean, standard_deviation = posterior_summary(model, candidates, remaining)

    if policy == "greedy":
        scores = mean
    elif policy == "ucb":
        scores = mean + UCB_BETA * standard_deviation
    else:
        raise ValueError(f"Unknown policy: {policy}")

    selected_position = torch.argmax(scores).item()
    return remaining[selected_position]

How does the policy choose?

  • The random branch preserves the baseline from the previous step.
  • The model-based branches fit the same surrogate from the same observed history.
  • The greedy branch uses only the posterior mean as its score.
  • The final position maps back to the original candidate-library index.
  • In run_campaign, find the current random-selection block:
        next_index = choose_random(
            len(candidates), observed_indices, choice_generator
        )

What is the current limitation?

This block sends every campaign through choose_random. The policy argument currently has no influence over experiment selection.

  • Replace that block with the shared policy call shown here:
        next_index = propose_next(
            policy,
            candidates,
            observed_indices,
            observed_values,
            choice_generator,
        )

What changes in the campaign?

Each campaign now passes the same candidate library and observed history into propose_next. Only the selection rule changes between random and greedy runs.

  • Save discovery_agent.py.

Before you run the greedy diagnostic, where do you think its early coordinates will fall after seeing only the lower-left initial designs?

  • Print all eight greedy proposals by running this diagnostic command:
python -c "import torch
from discovery_agent import ADDITIONAL_BUDGET, build_candidate_library, get_initial_indices, hidden_response, observe, propose_next
torch.manual_seed(3)
noise_generator = torch.Generator().manual_seed(3)
choice_generator = torch.Generator().manual_seed(1003)
candidates = build_candidate_library()
truth = hidden_response(candidates)
observed_indices = get_initial_indices(candidates)
observed_values = [observe(truth, index, 0.15, noise_generator) for index in observed_indices]
for _ in range(ADDITIONAL_BUDGET):
    next_index = propose_next('greedy', candidates, observed_indices, observed_values, choice_generator)
    print(candidates[next_index])
    observed_indices.append(next_index)
    observed_values.append(observe(truth, next_index, 0.15, noise_generator))"

What does the diagnostic reveal?

You should see the early proposed coordinates cluster around the lower-left portion of the grid. The observed local peak gives that neighborhood strong posterior means.

A high posterior mean expresses the model's current belief. Evidence for the global optimum requires measurements from competing regions.

Greedy diagnostic failing?

Confirm that propose_next appears above run_campaign. Check that the greedy branch assigns scores = mean.

Help me trace my greedy-selection diagnostic.

Measure final posterior coverage

Simple regret checks the quality of the best candidate discovered so far. Posterior coverage checks how often the hidden response falls inside the model's two-standard-deviation interval across the whole library.

  • Add the coverage function directly below propose_next by pasting this code:
def two_sigma_coverage(model, candidates, truth):
    all_indices = list(range(len(candidates)))
    mean, standard_deviation = posterior_summary(model, candidates, all_indices)
    lower = mean - 2.0 * standard_deviation
    upper = mean + 2.0 * standard_deviation
    covered = (truth >= lower) & (truth <= upper)
    return covered.to(torch.float64).mean().item()

What does coverage measure?

  • The lower and upper tensors create one uncertainty interval per candidate.
  • The Boolean coverage tensor records whether each hidden response lies inside its interval.
  • The final mean converts those checks into a proportion between 0.0 and 1.0.
  • At the bottom of run_campaign, find the current return line:
    return torch.tensor(regrets)

What does the campaign return now?

The current return value contains only the regret history. It has no final model diagnostic for the candidate library.

  • Replace that return line with the final surrogate fit and coverage calculation shown here:
    final_model = fit_surrogate(candidates, observed_indices, observed_values)
    coverage = two_sigma_coverage(final_model, candidates, truth)
    return torch.tensor(regrets), coverage

Why fit one final surrogate?

The final fit includes every experiment selected during the campaign. The function returns the nine-point regret curve beside the coverage proportion from that final model.

  • Save discovery_agent.py.

Before the final check, do you expect the campaign result to contain one value or two distinct metrics?

  • Run one complete greedy campaign by using this command:
python -c "from discovery_agent import run_campaign
print(run_campaign('greedy', 3, 0.15))"

What should you see?

You should see a pair containing a nine-value PyTorch tensor and one coverage value. The tensor tracks simple regret before selection plus eight additional experiments.

That is the model-based campaign working end to end. Greedy selection now exposes how confident local choices can leave a distant region untested.

Campaign return looks incomplete?

Confirm that run_campaign calls propose_next inside its loop. Check that its final line returns both torch.tensor(regrets) and coverage.

Help me debug my greedy campaign output.

✔️ Awesome, I've got everything!

Your surrogate fits the observed campaign data. Greedy selection now proposes untested candidates from posterior means. Save discovery_agent.py before continuing.

ⓧ I'd like to double check the full code

Compare your complete discovery_agent.py file with this cumulative version:

from pathlib import Path

import matplotlib.pyplot as plt
import torch
from botorch.fit import fit_gpytorch_mll
from botorch.models import SingleTaskGP
from gpytorch.mlls import ExactMarginalLogLikelihood


torch.set_default_dtype(torch.float64)

GRID_SIZE = 15
ADDITIONAL_BUDGET = 8
SEEDS = (3, 11, 29)
NOISE_LEVELS = (0.05, 0.15)
POLICIES = ("random", "greedy", "ucb")
UCB_BETA = 2.0
OUTPUT_DIR = Path("outputs")

INITIAL_DESIGNS = torch.tensor(
    [
        [0.10, 0.10],
        [0.15, 0.30],
        [0.30, 0.15],
        [0.35, 0.35],
        [0.25, 0.25],
    ]
)


def build_candidate_library(grid_size=GRID_SIZE):
    axis = torch.linspace(0.0, 1.0, steps=grid_size)
    first, second = torch.meshgrid(axis, axis, indexing="ij")
    return torch.stack([first.reshape(-1), second.reshape(-1)], dim=-1)


def hidden_response(candidates):
    first = candidates[:, 0]
    second = candidates[:, 1]
    local_peak = 0.85 * torch.exp(
        -((first - 0.25) ** 2 + (second - 0.25) ** 2) / 0.018
    )
    global_peak = 1.20 * torch.exp(
        -((first - 0.78) ** 2 + (second - 0.72) ** 2) / 0.012
    )
    ridge = 0.15 * torch.exp(-((first - second) ** 2) / 0.08)
    return local_peak + global_peak + ridge


def get_initial_indices(candidates):
    distances = torch.cdist(INITIAL_DESIGNS, candidates)
    return distances.argmin(dim=1).tolist()


def available_indices(candidate_count, observed_indices):
    observed = set(observed_indices)
    return [index for index in range(candidate_count) if index not in observed]


def observe(truth, index, noise_std, generator):
    noise = noise_std * torch.randn((), generator=generator)
    return truth[index] + noise


def choose_random(candidate_count, observed_indices, generator):
    remaining = available_indices(candidate_count, observed_indices)
    position = torch.randperm(len(remaining), generator=generator)[0].item()
    return remaining[position]


def simple_regret(truth, observed_indices):
    observed_tensor = torch.tensor(observed_indices)
    best_possible = truth.max()
    best_found = truth[observed_tensor].max()
    return (best_possible - best_found).item()


def fit_surrogate(candidates, observed_indices, observed_values):
    index_tensor = torch.tensor(observed_indices)
    train_x = candidates[index_tensor]
    train_y = torch.stack(observed_values).reshape(-1, 1)
    model = SingleTaskGP(train_X=train_x, train_Y=train_y)
    mll = ExactMarginalLogLikelihood(model.likelihood, model)
    fit_gpytorch_mll(mll)
    return model


def posterior_summary(model, candidates, candidate_indices):
    index_tensor = torch.tensor(candidate_indices)
    with torch.no_grad():
        posterior = model.posterior(candidates[index_tensor])
        mean = posterior.mean.squeeze(-1)
        standard_deviation = torch.sqrt(
            posterior.variance.squeeze(-1).clamp_min(1e-12)
        )
    return mean, standard_deviation


def propose_next(
    policy,
    candidates,
    observed_indices,
    observed_values,
    choice_generator,
):
    if policy == "random":
        return choose_random(len(candidates), observed_indices, choice_generator)

    remaining = available_indices(len(candidates), observed_indices)
    model = fit_surrogate(candidates, observed_indices, observed_values)
    mean, standard_deviation = posterior_summary(model, candidates, remaining)

    if policy == "greedy":
        scores = mean
    elif policy == "ucb":
        scores = mean + UCB_BETA * standard_deviation
    else:
        raise ValueError(f"Unknown policy: {policy}")

    selected_position = torch.argmax(scores).item()
    return remaining[selected_position]


def two_sigma_coverage(model, candidates, truth):
    all_indices = list(range(len(candidates)))
    mean, standard_deviation = posterior_summary(model, candidates, all_indices)
    lower = mean - 2.0 * standard_deviation
    upper = mean + 2.0 * standard_deviation
    covered = (truth >= lower) & (truth <= upper)
    return covered.to(torch.float64).mean().item()


def run_campaign(policy, seed, noise_std):
    torch.manual_seed(seed)
    noise_generator = torch.Generator().manual_seed(seed)
    choice_generator = torch.Generator().manual_seed(seed + 1000)

    candidates = build_candidate_library()
    truth = hidden_response(candidates)
    observed_indices = get_initial_indices(candidates)
    observed_values = [
        observe(truth, index, noise_std, noise_generator)
        for index in observed_indices
    ]
    regrets = [simple_regret(truth, observed_indices)]

    for _ in range(ADDITIONAL_BUDGET):
        next_index = propose_next(
            policy,
            candidates,
            observed_indices,
            observed_values,
            choice_generator,
        )
        observed_indices.append(next_index)
        observed_values.append(
            observe(truth, next_index, noise_std, noise_generator)
        )
        regrets.append(simple_regret(truth, observed_indices))

    final_model = fit_surrogate(candidates, observed_indices, observed_values)
    coverage = two_sigma_coverage(final_model, candidates, truth)
    return torch.tensor(regrets), coverage

How does the complete file fit together?

The first functions build the synthetic campaign and reveal noisy observations. The surrogate functions convert the observed history into posterior predictions.

The policy function selects the next candidate. The campaign loop records regret before returning final posterior coverage.

Your agent can now fit a surrogate and act on its highest predicted response. Next, you will give it a way to value uncertain regions before the local peak consumes its experiment budget.

Add Uncertainty-Aware Exploration

Your greedy agent can now rank candidates with the posterior mean. Its early proposals clustered around the lower-left peak while the uncertain upper-right region remained neglected.

An uncertainty-aware policy rewards candidates that look promising under the Gaussian-process posterior. It also rewards candidates whose outcomes remain uncertain.

In this step, get ready to:
  • Add a posterior-standard-deviation bonus to the candidate score.
  • Route the random, greedy, and UCB policies through the same selection function.
  • Compare matched greedy and UCB campaigns using identical seeds and noise levels.
Add the uncertainty bonus

An upper-confidence score combines two signals. The posterior mean drives exploitation while the posterior standard deviation drives exploration.

The project uses UCB_BETA to control the strength of the uncertainty bonus. A value of 2.0 lets uncertain candidates compete with candidates that already have a high predicted response.

  • In discovery_agent.py, find the greedy scoring branch inside propose_next().
  • Use these lines as your editing anchor:
if policy == "greedy":
        scores = mean

What is the current limit?

The greedy branch ranks every remaining candidate using only mean. A candidate receives no credit for uncertainty about its hidden response.

  • Replace the greedy-only branch with this policy routing:
if policy == "greedy":
        scores = mean
    elif policy == "ucb":
        scores = mean + UCB_BETA * standard_deviation
    else:
        raise ValueError(f"Unknown policy: {policy}")

How Does This Score Work?

  • The mean term favors candidates with strong predicted responses.
  • The standard_deviation term favors candidates whose responses remain uncertain.
  • The UCB_BETA constant controls how strongly uncertainty affects the score.
  • The final else branch stops unknown policy names from producing a misleading selection.
  • Save discovery_agent.py.
Prepare a matched debugger check

A fair comparison holds the initial designs, random seed, noise level, and experiment budget constant. The policy becomes the only intentional difference.

Conditional breakpoints are a little fiddly because the benchmark calls run_campaign() many times. You will use them to pause only for the two matched campaigns you need.

  • In discovery_agent.py, find observed_indices.append(next_index) inside run_campaign().
  • Right-click the editor margin beside that line.
  • Choose the conditional breakpoint option.
  • Enter policy in ("greedy", "ucb") and seed == 3 and noise_std == 0.15 and len(observed_indices) == 5 as the condition.
  • Find return torch.tensor(regrets), coverage near the end of run_campaign().
  • Set a conditional breakpoint on that line.
  • Enter policy in ("greedy", "ucb") and seed == 3 and noise_std == 0.15 as the condition.

You should see two breakpoint indicators in the editor margin. One captures the first proposed candidate while the other captures the completed regret curve.

The debugger in VS Code can evaluate expressions whenever execution pauses. Three watch expressions expose the policy, selected coordinate, and regret history.

  • Press Ctrl+Shift+D to open the Run and Debug view.
  • Add policy to the WATCH section.
  • Add candidates[next_index] to the WATCH section.
  • Add torch.tensor(regrets) to the WATCH section.

You should now see three watch rows. Their values populate when the debugger reaches the matching campaigns.

Compare the matched campaigns

The matched run uses seed 3 with a noise standard deviation of 0.15. Each policy receives the same starting evidence before proposing eight new experiments.

Before you start debugging, which policy do you expect to leave the lower-left neighborhood first?

  • Press F5 to start debugging the active Python file.
  • Choose the Python debugger for the active file if VS Code asks how to start the session.

At the first pause, policy should show greedy. The candidates[next_index] watch shows its first proposed coordinate.

  • Record the first greedy coordinate as your greedy first coordinate.
  • Press F5 to continue to the completed greedy campaign.

At the next pause, torch.tensor(regrets) should contain nine values. The first value comes from the initial designs while each later value follows one additional experiment.

  • Record the greedy values as your greedy regret curve.
  • Press F5 to continue to the first UCB proposal.

The policy watch should now show ucb. Its first coordinate should reach beyond the greedy policy's lower-left neighborhood.

  • Record the first UCB coordinate as your UCB first coordinate.
  • Press F5 to continue to the completed UCB campaign.

The final pause exposes the matched UCB regret curve. Its nine values show how the best discovered response changes as the uncertainty-aware policy spends its budget.

  • Record the UCB values as your UCB regret curve.
  • Press Shift+F5 to stop the debugging session.

You now have direct evidence that the uncertainty bonus changes experiment selection. Compare the two curves experiment by experiment to see when that wider search improves simple regret.

Debugger Not Pausing?

Check that both breakpoint conditions use seed 3 with noise level 0.15. Confirm that the first breakpoint sits on observed_indices.append(next_index).

If a watch expression is unavailable, wait for the breakpoint inside run_campaign(). The required local variables only exist while that function is active.

Help me diagnose my matched greedy and UCB debugger check.

✔️ Awesome, I've got everything!

Your ucb branch now combines posterior mean with posterior standard deviation. You have also compared it with greedy selection under a matched seed, noise level, initialization, and budget.

ⓧ I'd like to double check the full code

Compare your saved discovery_agent.py with this cumulative reference. Pay particular attention to the three policy branches inside propose_next().

from pathlib import Path

import matplotlib.pyplot as plt
import torch
from botorch.fit import fit_gpytorch_mll
from botorch.models import SingleTaskGP
from gpytorch.mlls import ExactMarginalLogLikelihood


torch.set_default_dtype(torch.float64)

GRID_SIZE = 15
ADDITIONAL_BUDGET = 8
SEEDS = (3, 11, 29)
NOISE_LEVELS = (0.05, 0.15)
POLICIES = ("random", "greedy", "ucb")
UCB_BETA = 2.0
OUTPUT_DIR = Path("outputs")

INITIAL_DESIGNS = torch.tensor(
    [
        [0.10, 0.10],
        [0.15, 0.30],
        [0.30, 0.15],
        [0.35, 0.35],
        [0.25, 0.25],
    ]
)


def build_candidate_library(grid_size=GRID_SIZE):
    axis = torch.linspace(0.0, 1.0, steps=grid_size)
    first, second = torch.meshgrid(axis, axis, indexing="ij")
    return torch.stack([first.reshape(-1), second.reshape(-1)], dim=-1)


def hidden_response(candidates):
    first = candidates[:, 0]
    second = candidates[:, 1]
    local_peak = 0.85 * torch.exp(
        -((first - 0.25) ** 2 + (second - 0.25) ** 2) / 0.018
    )
    global_peak = 1.20 * torch.exp(
        -((first - 0.78) ** 2 + (second - 0.72) ** 2) / 0.012
    )
    ridge = 0.15 * torch.exp(-((first - second) ** 2) / 0.08)
    return local_peak + global_peak + ridge


def get_initial_indices(candidates):
    distances = torch.cdist(INITIAL_DESIGNS, candidates)
    return distances.argmin(dim=1).tolist()


def available_indices(candidate_count, observed_indices):
    observed = set(observed_indices)
    return [index for index in range(candidate_count) if index not in observed]


def observe(truth, index, noise_std, generator):
    noise = noise_std * torch.randn((), generator=generator)
    return truth[index] + noise


def choose_random(candidate_count, observed_indices, generator):
    remaining = available_indices(candidate_count, observed_indices)
    position = torch.randperm(len(remaining), generator=generator)[0].item()
    return remaining[position]


def simple_regret(truth, observed_indices):
    observed_tensor = torch.tensor(observed_indices)
    best_possible = truth.max()
    best_found = truth[observed_tensor].max()
    return (best_possible - best_found).item()


def fit_surrogate(candidates, observed_indices, observed_values):
    index_tensor = torch.tensor(observed_indices)
    train_x = candidates[index_tensor]
    train_y = torch.stack(observed_values).reshape(-1, 1)
    model = SingleTaskGP(train_X=train_x, train_Y=train_y)
    mll = ExactMarginalLogLikelihood(model.likelihood, model)
    fit_gpytorch_mll(mll)
    return model


def posterior_summary(model, candidates, candidate_indices):
    index_tensor = torch.tensor(candidate_indices)
    with torch.no_grad():
        posterior = model.posterior(candidates[index_tensor])
        mean = posterior.mean.squeeze(-1)
        standard_deviation = torch.sqrt(
            posterior.variance.squeeze(-1).clamp_min(1e-12)
        )
    return mean, standard_deviation


def propose_next(
    policy,
    candidates,
    observed_indices,
    observed_values,
    choice_generator,
):
    if policy == "random":
        return choose_random(len(candidates), observed_indices, choice_generator)

    remaining = available_indices(len(candidates), observed_indices)
    model = fit_surrogate(candidates, observed_indices, observed_values)
    mean, standard_deviation = posterior_summary(model, candidates, remaining)

    if policy == "greedy":
        scores = mean
    elif policy == "ucb":
        scores = mean + UCB_BETA * standard_deviation
    else:
        raise ValueError(f"Unknown policy: {policy}")

    selected_position = torch.argmax(scores).item()
    return remaining[selected_position]


def two_sigma_coverage(model, candidates, truth):
    all_indices = list(range(len(candidates)))
    mean, standard_deviation = posterior_summary(model, candidates, all_indices)
    lower = mean - 2.0 * standard_deviation
    upper = mean + 2.0 * standard_deviation
    covered = (truth >= lower) & (truth <= upper)
    return covered.to(torch.float64).mean().item()


def run_campaign(policy, seed, noise_std):
    torch.manual_seed(seed)
    noise_generator = torch.Generator().manual_seed(seed)
    choice_generator = torch.Generator().manual_seed(seed + 1000)

    candidates = build_candidate_library()
    truth = hidden_response(candidates)
    observed_indices = get_initial_indices(candidates)
    observed_values = [
        observe(truth, index, noise_std, noise_generator)
        for index in observed_indices
    ]
    regrets = [simple_regret(truth, observed_indices)]

    for _ in range(ADDITIONAL_BUDGET):
        next_index = propose_next(
            policy,
            candidates,
            observed_indices,
            observed_values,
            choice_generator,
        )
        observed_indices.append(next_index)
        observed_values.append(
            observe(truth, next_index, noise_std, noise_generator)
        )
        regrets.append(simple_regret(truth, observed_indices))

    final_model = fit_surrogate(candidates, observed_indices, observed_values)
    coverage = two_sigma_coverage(final_model, candidates, truth)
    return torch.tensor(regrets), coverage


def benchmark():
    results = {
        noise: {policy: [] for policy in POLICIES} for noise in NOISE_LEVELS
    }
    coverages = {
        noise: {policy: [] for policy in POLICIES} for noise in NOISE_LEVELS
    }

    for noise in NOISE_LEVELS:
        for policy in POLICIES:
            for seed in SEEDS:
                regret_curve, coverage = run_campaign(policy, seed, noise)
                results[noise][policy].append(regret_curve)
                coverages[noise][policy].append(coverage)
    return results, coverages


def mean_and_spread(curves):
    stacked = torch.stack(curves)
    mean_curve = stacked.mean(dim=0)
    spread = torch.sqrt(((stacked - mean_curve) ** 2).mean(dim=0))
    return mean_curve, spread


def plot_results(results, coverages):
    OUTPUT_DIR.mkdir(exist_ok=True)
    figure, axes = plt.subplots(
        1, len(NOISE_LEVELS), figsize=(10, 4), layout="constrained"
    )
    experiment_counts = list(range(ADDITIONAL_BUDGET + 1))

    for axis, noise in zip(axes, NOISE_LEVELS):
        for policy in POLICIES:
            mean_curve, spread = mean_and_spread(results[noise][policy])
            axis.plot(
                experiment_counts,
                mean_curve.tolist(),
                marker="o",
                label=policy,
            )
            axis.fill_between(
                experiment_counts,
                (mean_curve - spread).tolist(),
                (mean_curve + spread).tolist(),
                alpha=0.15,
            )
        axis.set(
            xlabel="Additional experiments",
            ylabel="Simple regret",
            title=f"Noise standard deviation = {noise}",
        )
        axis.legend()

    figure.savefig(
        OUTPUT_DIR / "regret_by_budget.png", dpi=150, bbox_inches="tight"
    )

    calibration_figure, calibration_axis = plt.subplots(
        figsize=(7, 4), layout="constrained"
    )
    for noise in NOISE_LEVELS:
        values = [
            sum(coverages[noise][policy]) / len(coverages[noise][policy])
            for policy in POLICIES
        ]
        calibration_axis.plot(
            list(POLICIES), values, marker="o", label=f"noise={noise}"
        )
    calibration_axis.set(
        xlabel="Selection policy",
        ylabel="Two-standard-deviation coverage",
        title="Posterior coverage across the candidate library",
    )
    calibration_axis.legend()
    calibration_figure.savefig(
        OUTPUT_DIR / "calibration.png", dpi=150, bbox_inches="tight"
    )


def write_findings(results, coverages):
    high_noise = NOISE_LEVELS[-1]
    final_regret = {
        policy: mean_and_spread(results[high_noise][policy])[0][-1].item()
        for policy in POLICIES
    }
    best_policy = min(final_regret, key=final_regret.get)
    mean_coverage = {
        policy: sum(coverages[high_noise][policy])
        / len(coverages[high_noise][policy])
        for policy in POLICIES
    }

    content = (
        "# Benchmark findings\n\n"
        f"- Additional experiment budget: {ADDITIONAL_BUDGET}\n"
        f"- Seeds: {list(SEEDS)}\n"
        f"- Noise levels: {list(NOISE_LEVELS)}\n"
        f"- Lowest mean final regret at noise {high_noise}: {best_policy}\n"
        f"- Final regret by policy: {final_regret}\n"
        f"- Two-standard-deviation coverage at noise {high_noise}: "
        f"{mean_coverage}\n\n"
        "Interpretation: compare the regret and coverage together. A policy can "
        "find strong candidates while its uncertainty remains poorly calibrated.\n\n"
        "Limitation: this controlled synthetic campaign supports method comparison, "
        "not a claim of validated performance on laboratory data.\n"
    )
    (OUTPUT_DIR / "findings.md").write_text(content, encoding="utf-8")


def main():
    results, coverages = benchmark()
    plot_results(results, coverages)
    write_findings(results, coverages)
    print("Saved outputs/regret_by_budget.png")
    print("Saved outputs/calibration.png")
    print("Saved outputs/findings.md")


if __name__ == "__main__":
    main()

What Should You Compare?

Confirm that POLICIES contains random, greedy, and ucb. Inside propose_next(), confirm that UCB adds UCB_BETA * standard_deviation to the posterior mean.

Confirm that selected_position comes from the largest score. The return statement must map that position back through remaining so the agent selects an original candidate index.

Your agent now balances promising predictions with unresolved uncertainty. Next, you will test whether that advantage survives repeated seeds, different noise levels, and calibration checks.

Run the Budget and Calibration Benchmark

Your UCB policy has already shown how uncertainty can pull a campaign beyond the greedy policy's local neighborhood. That matched run gave you one useful trajectory.

A single trajectory can still be luck. This step tests every policy across repeated seeds and harsher noise before judging discovery performance with simple regret and posterior calibration.

In this step, get ready to:
  • Run repeated campaigns under identical experiment budgets.
  • Compare discovery performance with uncertainty calibration.
  • Review the generated plots and research findings brief.
Benchmark every policy fairly

The benchmark uses the same INITIAL_DESIGNS for every policy. It also holds ADDITIONAL_BUDGET at 8 experiments.

Each policy runs with SEEDS set to (3, 11, 29). The two NOISE_LEVELS are 0.05 and 0.15.

Why repeat the same campaign?

Random noise can make one policy look stronger in a single campaign. Repeated seeds expose whether that result persists across different observations.

Shared initial designs keep the comparison controlled. Each policy starts with the same evidence before choosing its additional experiments.

Before you run the benchmark, which policy do you expect to finish with the lowest regret under the higher noise level?

  • Launch all repeated campaigns from the activated PowerShell terminal by running:
python discovery_agent.py

What does this command do?

The script runs benchmark() across every policy, seed, and noise level. Each campaign receives the same additional experiment budget.

The script then calls plot_results() and write_findings(). These functions turn the repeated campaigns into two plots and one Markdown brief.

Repeated Gaussian-process fitting keeps the terminal busy between messages. Leave the process running until it prints all three save confirmations.

You'll see Saved outputs/regret_by_budget.png.

You'll also see Saved outputs/calibration.png.

The final confirmation is Saved outputs/findings.md. You now have a reproducible benchmark report built from every campaign.

Benchmark stopped before saving?

  • Confirm that the PowerShell prompt still begins with (.venv).
  • Confirm that the terminal path ends with budgeted-discovery-agent.
  • Compare the final function definitions in discovery_agent.py with the reference tab below.

Help me diagnose why discovery_agent.py stopped before saving its benchmark outputs.

Compare regret with calibration

A strong discovery policy should find high-value candidates within its budget. Its surrogate should also express uncertainty that agrees with the hidden response across the candidate library.

  • Expand the outputs folder in the VS Code Explorer sidebar.

You'll see regret_by_budget.png, calibration.png, and findings.md inside the folder.

Before you inspect the plots, where do you expect repeated-seed variation to appear most clearly?

  • Select outputs/regret_by_budget.png in the Explorer sidebar to compare policy performance across the experiment budget.
  • Select outputs/calibration.png in the Explorer sidebar to compare posterior coverage across policies.

The regret figure has separate panels for noise standard deviations 0.05 and 0.15. Each panel includes random, greedy, and UCB curves with shaded across-seed spread.

The calibration figure shows all three policies at both noise levels. Its vertical axis reports two-standard-deviation posterior coverage.

What do the metrics reveal?

Simple regret compares the best candidate found so far with the best candidate in the full synthetic library. A lower value means the campaign has discovered a stronger candidate.

Two-standard-deviation coverage measures how often the hidden response falls inside the surrogate's uncertainty interval. It catches overconfident or poorly calibrated predictions that regret alone can miss.

Review the generated findings

The plots hold the detailed evidence. The findings brief turns that evidence into a recommendation while preserving the limits of a synthetic experiment.

Before you read the recommendation, which policy do you expect the brief to name at the higher noise level?

  • Select outputs/findings.md in the Explorer sidebar to identify the policy with the lowest mean final high-noise regret.

The brief records the experiment budget, seeds, noise levels, final regret values, and coverage values. It also names the high-noise policy with the lowest mean final regret.

  • Compare the named policy with the final points in the high-noise panel of outputs/regret_by_budget.png.

The recommended policy should match the curve with the lowest mean final regret. This agreement connects the written claim directly to the plotted evidence.

  • Read the final limitation in outputs/findings.md to confirm that the result is framed as a controlled method comparison.

The limitation explains that a synthetic campaign does not validate performance on laboratory data. That distinction keeps the recommendation scientifically defensible.

Why read regret and coverage together?

A policy can find a strong candidate while its surrogate remains poorly calibrated. Low regret alone cannot establish trustworthy uncertainty.

Coverage adds evidence about the surrogate's confidence. Together, the metrics separate successful discovery from reliable uncertainty estimation.

✔️ Awesome, I've got everything!

Your benchmark has generated both plots and the findings brief. Confirm that all three files remain inside the outputs folder.

ⓧ I'd like to double check the full code

Compare your complete discovery_agent.py file with this reference:

from pathlib import Path

import matplotlib.pyplot as plt
import torch
from botorch.fit import fit_gpytorch_mll
from botorch.models import SingleTaskGP
from gpytorch.mlls import ExactMarginalLogLikelihood


torch.set_default_dtype(torch.float64)

GRID_SIZE = 15
ADDITIONAL_BUDGET = 8
SEEDS = (3, 11, 29)
NOISE_LEVELS = (0.05, 0.15)
POLICIES = ("random", "greedy", "ucb")
UCB_BETA = 2.0
OUTPUT_DIR = Path("outputs")

INITIAL_DESIGNS = torch.tensor(
    [
        [0.10, 0.10],
        [0.15, 0.30],
        [0.30, 0.15],
        [0.35, 0.35],
        [0.25, 0.25],
    ]
)


def build_candidate_library(grid_size=GRID_SIZE):
    axis = torch.linspace(0.0, 1.0, steps=grid_size)
    first, second = torch.meshgrid(axis, axis, indexing="ij")
    return torch.stack([first.reshape(-1), second.reshape(-1)], dim=-1)


def hidden_response(candidates):
    first = candidates[:, 0]
    second = candidates[:, 1]
    local_peak = 0.85 * torch.exp(
        -((first - 0.25) ** 2 + (second - 0.25) ** 2) / 0.018
    )
    global_peak = 1.20 * torch.exp(
        -((first - 0.78) ** 2 + (second - 0.72) ** 2) / 0.012
    )
    ridge = 0.15 * torch.exp(-((first - second) ** 2) / 0.08)
    return local_peak + global_peak + ridge


def get_initial_indices(candidates):
    distances = torch.cdist(INITIAL_DESIGNS, candidates)
    return distances.argmin(dim=1).tolist()


def available_indices(candidate_count, observed_indices):
    observed = set(observed_indices)
    return [index for index in range(candidate_count) if index not in observed]


def observe(truth, index, noise_std, generator):
    noise = noise_std * torch.randn((), generator=generator)
    return truth[index] + noise


def choose_random(candidate_count, observed_indices, generator):
    remaining = available_indices(candidate_count, observed_indices)
    position = torch.randperm(len(remaining), generator=generator)[0].item()
    return remaining[position]


def simple_regret(truth, observed_indices):
    observed_tensor = torch.tensor(observed_indices)
    best_possible = truth.max()
    best_found = truth[observed_tensor].max()
    return (best_possible - best_found).item()


def fit_surrogate(candidates, observed_indices, observed_values):
    index_tensor = torch.tensor(observed_indices)
    train_x = candidates[index_tensor]
    train_y = torch.stack(observed_values).reshape(-1, 1)
    model = SingleTaskGP(train_X=train_x, train_Y=train_y)
    mll = ExactMarginalLogLikelihood(model.likelihood, model)
    fit_gpytorch_mll(mll)
    return model


def posterior_summary(model, candidates, candidate_indices):
    index_tensor = torch.tensor(candidate_indices)
    with torch.no_grad():
        posterior = model.posterior(candidates[index_tensor])
        mean = posterior.mean.squeeze(-1)
        standard_deviation = torch.sqrt(
            posterior.variance.squeeze(-1).clamp_min(1e-12)
        )
    return mean, standard_deviation


def propose_next(
    policy,
    candidates,
    observed_indices,
    observed_values,
    choice_generator,
):
    if policy == "random":
        return choose_random(len(candidates), observed_indices, choice_generator)

    remaining = available_indices(len(candidates), observed_indices)
    model = fit_surrogate(candidates, observed_indices, observed_values)
    mean, standard_deviation = posterior_summary(model, candidates, remaining)

    if policy == "greedy":
        scores = mean
    elif policy == "ucb":
        scores = mean + UCB_BETA * standard_deviation
    else:
        raise ValueError(f"Unknown policy: {policy}")

    selected_position = torch.argmax(scores).item()
    return remaining[selected_position]


def two_sigma_coverage(model, candidates, truth):
    all_indices = list(range(len(candidates)))
    mean, standard_deviation = posterior_summary(model, candidates, all_indices)
    lower = mean - 2.0 * standard_deviation
    upper = mean + 2.0 * standard_deviation
    covered = (truth >= lower) & (truth <= upper)
    return covered.to(torch.float64).mean().item()


def run_campaign(policy, seed, noise_std):
    torch.manual_seed(seed)
    noise_generator = torch.Generator().manual_seed(seed)
    choice_generator = torch.Generator().manual_seed(seed + 1000)

    candidates = build_candidate_library()
    truth = hidden_response(candidates)
    observed_indices = get_initial_indices(candidates)
    observed_values = [
        observe(truth, index, noise_std, noise_generator)
        for index in observed_indices
    ]
    regrets = [simple_regret(truth, observed_indices)]

    for _ in range(ADDITIONAL_BUDGET):
        next_index = propose_next(
            policy,
            candidates,
            observed_indices,
            observed_values,
            choice_generator,
        )
        observed_indices.append(next_index)
        observed_values.append(
            observe(truth, next_index, noise_std, noise_generator)
        )
        regrets.append(simple_regret(truth, observed_indices))

    final_model = fit_surrogate(candidates, observed_indices, observed_values)
    coverage = two_sigma_coverage(final_model, candidates, truth)
    return torch.tensor(regrets), coverage


def benchmark():
    results = {
        noise: {policy: [] for policy in POLICIES} for noise in NOISE_LEVELS
    }
    coverages = {
        noise: {policy: [] for policy in POLICIES} for noise in NOISE_LEVELS
    }

    for noise in NOISE_LEVELS:
        for policy in POLICIES:
            for seed in SEEDS:
                regret_curve, coverage = run_campaign(policy, seed, noise)
                results[noise][policy].append(regret_curve)
                coverages[noise][policy].append(coverage)
    return results, coverages


def mean_and_spread(curves):
    stacked = torch.stack(curves)
    mean_curve = stacked.mean(dim=0)
    spread = torch.sqrt(((stacked - mean_curve) ** 2).mean(dim=0))
    return mean_curve, spread


def plot_results(results, coverages):
    OUTPUT_DIR.mkdir(exist_ok=True)
    figure, axes = plt.subplots(
        1, len(NOISE_LEVELS), figsize=(10, 4), layout="constrained"
    )
    experiment_counts = list(range(ADDITIONAL_BUDGET + 1))

    for axis, noise in zip(axes, NOISE_LEVELS):
        for policy in POLICIES:
            mean_curve, spread = mean_and_spread(results[noise][policy])
            axis.plot(
                experiment_counts,
                mean_curve.tolist(),
                marker="o",
                label=policy,
            )
            axis.fill_between(
                experiment_counts,
                (mean_curve - spread).tolist(),
                (mean_curve + spread).tolist(),
                alpha=0.15,
            )
        axis.set(
            xlabel="Additional experiments",
            ylabel="Simple regret",
            title=f"Noise standard deviation = {noise}",
        )
        axis.legend()

    figure.savefig(
        OUTPUT_DIR / "regret_by_budget.png", dpi=150, bbox_inches="tight"
    )

    calibration_figure, calibration_axis = plt.subplots(
        figsize=(7, 4), layout="constrained"
    )
    for noise in NOISE_LEVELS:
        values = [
            sum(coverages[noise][policy]) / len(coverages[noise][policy])
            for policy in POLICIES
        ]
        calibration_axis.plot(
            list(POLICIES), values, marker="o", label=f"noise={noise}"
        )
    calibration_axis.set(
        xlabel="Selection policy",
        ylabel="Two-standard-deviation coverage",
        title="Posterior coverage across the candidate library",
    )
    calibration_axis.legend()
    calibration_figure.savefig(
        OUTPUT_DIR / "calibration.png", dpi=150, bbox_inches="tight"
    )


def write_findings(results, coverages):
    high_noise = NOISE_LEVELS[-1]
    final_regret = {
        policy: mean_and_spread(results[high_noise][policy])[0][-1].item()
        for policy in POLICIES
    }
    best_policy = min(final_regret, key=final_regret.get)
    mean_coverage = {
        policy: sum(coverages[high_noise][policy])
        / len(coverages[high_noise][policy])
        for policy in POLICIES
    }

    content = (
        "# Benchmark findings\n\n"
        f"- Additional experiment budget: {ADDITIONAL_BUDGET}\n"
        f"- Seeds: {list(SEEDS)}\n"
        f"- Noise levels: {list(NOISE_LEVELS)}\n"
        f"- Lowest mean final regret at noise {high_noise}: {best_policy}\n"
        f"- Final regret by policy: {final_regret}\n"
        f"- Two-standard-deviation coverage at noise {high_noise}: "
        f"{mean_coverage}\n\n"
        "Interpretation: compare the regret and coverage together. A policy can "
        "find strong candidates while its uncertainty remains poorly calibrated.\n\n"
        "Limitation: this controlled synthetic campaign supports method comparison, "
        "not a claim of validated performance on laboratory data.\n"
    )
    (OUTPUT_DIR / "findings.md").write_text(content, encoding="utf-8")


def main():
    results, coverages = benchmark()
    plot_results(results, coverages)
    write_findings(results, coverages)
    print("Saved outputs/regret_by_budget.png")
    print("Saved outputs/calibration.png")
    print("Saved outputs/findings.md")


if __name__ == "__main__":
    main()

How is the complete file organized?

The campaign helpers build observations and propose candidates. The benchmark helpers aggregate repeated regret curves and coverage values.

The final functions save the plots and findings brief. The main() function runs that complete workflow.

Secret mission

Add Posterior-Sampling Experiment Selection

Your UCB agent uses the same uncertainty bonus for every decision. This challenge turns the fitted posterior into a randomized experiment-selection rule. You will compare its high-noise behavior with UCB under the same seeds and experiment budget.

Clean Up Your Resources

Clean Up Your Resources

This benchmark runs entirely on your Windows computer with no ongoing cost. Choose whether to keep the local files, pause your session, or delete the project.

Resources you used:

  • The budgeted-discovery-agent folder contains discovery_agent.py, thompson_extension.py, requirements.txt, .venv, and outputs.

Keep everything running

No action is needed. Choose this if you want to rerun the benchmark or continue testing experiment-selection policies.

  • Keep the budgeted-discovery-agent folder in its current location.
  • Preserve outputs/regret_by_budget.png, outputs/calibration.png, outputs/thompson_vs_ucb.png, and outputs/findings.md for future comparisons.
  • Continue using the open VS Code workspace when you want to test another policy.

Pause - I'll come back to this later

End the current local session to free its memory. Your scripts, environment, plots, and findings remain available for later.

  • End the active PowerShell session by clicking the trash can icon in the Terminal panel.
  • Close VS Code.
  • Reopen the existing budgeted-discovery-agent folder in VS Code when you return.
  • Use the activation instruction from Step 1 to restore .venv before rerunning the benchmark.

Delete - I don't want to use this again

This option permanently removes the entire local project. The generated results can be copied to another location first.

  • Copy the outputs folder outside budgeted-discovery-agent if you want to preserve your figures and findings.

What This Removes

The deletion removes discovery_agent.py, thompson_extension.py, requirements.txt, .venv, and outputs. It does not uninstall Python, Git, or VS Code.

  • Right-click budgeted-discovery-agent in the VS Code Explorer sidebar.
  • Choose the menu option that displays this folder in File Explorer.
  • Close VS Code.
  • Use the File Explorer address bar to move up to the folder containing budgeted-discovery-agent.
  • Select the budgeted-discovery-agent folder.
  • Press Shift+Delete to remove the selected folder permanently.
  • Approve the Windows deletion confirmation.
  • Confirm that budgeted-discovery-agent is no longer listed in File Explorer.

That completes the local cleanup. The scripts, local environment, and generated results no longer occupy space on your computer.

Folder Refusing to Delete?

  • Close any File Explorer preview that is displaying a project file.
  • Use Windows Task Manager to end any Python process that is still reading the project files.
  • Ask for help diagnosing the locked project folder.

Nice Work!

Nice Work!

You did it! Your budget-aware autonomous discovery agent now proposes experiments from a Gaussian-process posterior within a fixed budget.

You've learned how to:

  • Build a noisy finite-budget campaign that reveals one outcome at a time. Track simple regret after every selected experiment.
  • Fit a Gaussian-process surrogate with BoTorch plus GPyTorch. Use posterior means and uncertainty to diagnose why greedy selection stays near a local peak.
  • Benchmark random, greedy, plus UCB policies under the same budget. Test robustness across repeated seeds at two noise levels. Save the evidence as regret and calibration plots. Distill the results into outputs/findings.md.
  • Complete the Secret Mission by adding Thompson-style posterior sampling. Compare its high-noise regret against UCB with the same seeds plus budget. Defend the observed behavior without claiming guaranteed superiority.

Ready to quiz yourself?