Staying on Track: Efficient Trajectory Discovery with Adaptive Batch Sampling

This paper proposes a trajectory-oriented Bayesian optimization method that utilizes Gaussian process surrogates incorporating both input parameters and random seeds, combined with an adaptive Thompson Sampling algorithm, to efficiently identify data-consistent trajectories in stochastic epidemic models while outperforming traditional parameter-only inference approaches.

The Gist

Imagine you are trying to tune a very complex, expensive radio to find a specific song playing in a storm.

The Problem: The Static of Chance
Most traditional methods for tuning this radio (which represents a computer simulation of something like a virus spreading) only listen to the average sound. They say, "On average, the music sounds like this, so let's adjust the knobs to match the average."

But here's the catch: This radio is broken. Every time you press play, the static changes slightly, and the song sounds different, even if you leave the knobs in the exact same position. Sometimes the song is clear; sometimes it's garbled. If you only listen to the average, you might find a setting that sounds "okay" on average, but it might never actually produce something close to the specific clear version of the song you need to hear.

In the world of science, this is called a stochastic model. The "knobs" are the parameters (like how fast a virus spreads), and the "static" is the random chance (who meets whom, who gets sick first).

The Old Way: Guessing the Average
Old methods would try to find the "best average setting." They would run the simulation 100 times with the same settings, average the results, and say, "This is our best guess."

• The Flaw: This is like trying to find a specific person in a crowd by looking at a blurry photo of the whole group. You might know where the group is, but you can't find the specific person you need to talk to.

The New Way: "Staying on Track"
The authors of this paper propose a smarter way called Trajectory-Oriented Discovery. Instead of just looking for the average, they want to find the exact combinations of 'knobs' AND 'random static' that produce results closer to reality.

Think of it like this:

1. The Radio (The Simulation): It's expensive to run (takes a lot of time and money).
2. The Goal: Find specific recordings ('trajectories') that match a real-life event (like a real epidemic curve).
3. The Secret Sauce: They treat the "random static" (the seed number) not as noise to be ignored, but as a second set of knobs to be tuned.

How They Do It: The Adaptive Search
They use a clever robot assistant (an algorithm called Bayesian Optimization) to do the tuning. Here is how the robot works, using a "Smart Map" analogy:

• The Map (The Grid): Imagine a giant map of all possible knob settings. The robot needs to check points on this map to see if they produce a good song.
• The Old Robot (Fixed Grid): A dumb robot would check every square on a grid, like mowing a lawn in straight lines. It wastes time checking empty, grassy fields (bad settings) and might miss the hidden garden (the perfect setting) if the grid lines don't align with it.
• The New Robot (Adaptive Grid): This robot is smart.
  • Filtering: It looks at the map and says, "These areas look like dead ends. I'll stop checking them." It throws away the bad guesses.
  • Densifying: It looks at the areas that almost sound good and says, "Let's zoom in here! Let's check 100 tiny spots right next to this promising area."
  • The Result: Instead of mowing the whole lawn, it focuses all its energy on the tiny patch of flowers that actually blooms.

Why This Matters: The "CityCOVID" Example
The authors tested this on a massive simulation of the COVID-19 pandemic in Chicago (called CityCOVID). This simulation involves 2.7 million virtual people.

• The Challenge: You can't run this simulation millions of times because it takes too long.
• The Success: Their new method found specific "scenarios" (trajectories) that matched real hospital data much faster and more accurately than the old methods.
• The Benefit: It's not just about finding the right numbers for the virus. It's about finding the specific stories of how the virus spread that make sense. This helps public health officials say, "If we do X, here is the likely outcome," rather than just "On average, it might be okay."

The Takeaway
This paper is about stopping the practice of "averaging out" the chaos of the real world. Instead, it teaches computers how to hunt down the specific, chaotic, real-life scenarios that mimic what actually happened, using a smart, adaptive search strategy that saves time and money.

In a nutshell:

• Old Way: "Let's find the average weather."
• New Way: "Let's find the exact days it rained exactly like it did last Tuesday, so we can better plan our picnic."

By treating randomness as a feature rather than a bug, and by using a smart, zooming-in search strategy, they can find the 'perfect matches' much faster.

Technical Summary

1. Problem Statement

The paper addresses the challenge of calibrating stochastic simulation models, particularly in epidemiology (e.g., epidemic spread), where model outputs are random realizations (trajectories) rather than deterministic values.

• The Limitation of Current Methods: Traditional Bayesian Optimization (BO) and calibration approaches typically rely on summary statistics (e.g., mean, median, or quantiles) of multiple simulation replicates to match observed data. This aggregation often leads to a loss of identifiability and obscures specific stochastic realizations that might be consistent with empirical data.
• The Core Issue: In stochastic models, the same parameter set ($x$) can yield vastly different outcomes depending on the random seed ($r$). Relying solely on parameter estimation ($x$) ignores the specific "path" (trajectory) the system took. The authors argue that for complex systems like epidemic models, identifying the specific tuple of parameters and random seeds $(x, r)$ that generates a data-consistent trajectory is crucial for actionable insights (e.g., for data assimilation or intervention planning).
• Computational Bottleneck: Finding these specific $(x, r)$ pairs is computationally expensive because the search space is augmented (parameters + seeds), and standard grid-based or continuous optimization methods struggle to efficiently explore this high-dimensional, noisy space.

2. Methodology

The authors propose a Trajectory-Oriented Optimization (TOO) framework that combines three key components:

A. Common Random Number Gaussian Process (CRNGP)

Instead of modeling the mean behavior, the authors treat the random seed $r$ as an explicit input to the surrogate model.

• Augmented Input Space: The simulator is modeled as a deterministic function $f(x, r)$ over the joint space of parameters $x$ and seeds $r$.
• Kernel Structure: They employ a separable kernel $k((x_i, r), (x_j, r')) = k_x(x_i, x_j) \times k_r(r, r')$.
  • $k_x$ models similarity between parameter sets.
  • $k_r$ models similarity between seeds (assumed to be a constant correlation $\rho$ between different seeds).
• Benefit: This allows the surrogate to make predictions at the individual trajectory level, enabling the optimization algorithm to distinguish between "good" and "bad" realizations for a given parameter set.

B. Thompson Sampling (TS)

The optimization strategy uses Thompson Sampling, a probabilistic acquisition function.

• Mechanism: At each iteration, a realization of the CRNGP posterior is sampled. The algorithm selects the $(x, r)$ pair that minimizes the discrepancy (e.g., RMSE) between the sampled trajectory and the observed data.
• Advantage: TS naturally balances exploration (searching new regions) and exploitation (refining known good regions) without requiring complex gradient calculations.

C. Adaptive Grid Sampling (aCRN)

To overcome the inefficiency of fixed grids in high-dimensional spaces, the authors introduce an Adaptive Grid Refinement Strategy that operates in two stages at each iteration:

1. Filtering: Candidates are generated via Latin Hypercube Sampling (LHS). Based on the current CRNGP posterior, a likelihood function is computed. Points with low likelihood of producing a good trajectory are discarded.
2. Densification: To maintain a fixed grid size, new points are generated around the surviving high-likelihood candidates using a Metropolis-Hastings (MH) inspired approach. New parameter values are proposed and accepted based on the surrogate likelihood, while the set of random seeds remains fixed (or is managed within a fixed set).

• Result: The search space dynamically concentrates computational resources on regions of the $(x, r)$ space that are statistically promising, rather than wasting budget on low-probability areas.

3. Key Contributions

1. Trajectory-Level Inference: A novel shift from calibrating only parameters to jointly estimating parameters and specific stochastic realizations $(x, r)$, treating stochasticity as a feature rather than noise to be averaged out.
2. CRNGP Surrogate: The formalization and application of a Gaussian Process that explicitly models the dependence between random seeds and parameters, enabling direct inference on individual trajectories.
3. Adaptive Batch Sampling Algorithm: The development of an adaptive grid strategy (Filtering + Densification) that significantly improves the efficiency of Thompson Sampling for stochastic simulators, allowing for faster convergence to data-consistent trajectories.
4. Scalability: Demonstration of the method's ability to handle computationally expensive, high-resolution Agent-Based Models (ABMs) via High-Performance Computing (HPC) workflows.

4. Results

The method was evaluated on two scales: a simple stochastic SIR compartmental model and a complex, city-scale Agent-Based Model (CityCOVID).

• SIR Model Experiments:

  • Comparison: The proposed aCRN (Adaptive CRNGP) was compared against fixed-grid CRNGP, flexible-seed CRNGP, and heteroskedastic GP (hetGP) methods.
  • Quality: aCRN consistently identified a higher proportion of trajectories with low Root Mean Squared Error (RMSE) across various simulation budgets.
  • Speed-to-Solution: Using a Relative Area Under the Curve (rAUC) metric, aCRN was shown to find high-quality trajectories significantly earlier in the optimization process than competitors. This is critical for time-sensitive decision-making.
  • Exploration: Unlike methods that get stuck in local parameter optima, aCRN explored a broader range of the parameter space while maintaining focus on high-likelihood seeds.

• CityCOVID (Agent-Based Model) Experiments:

  • Context: Applied to a model of 2.7 million agents simulating COVID-19 in Chicago, calibrated against hospitalization and death data.
  • Performance: aCRN outperformed the next best method (fHet) in finding trajectories that strictly matched the dual-objective (deaths + hospitalizations).
  • Diversity: aCRN discovered diverse parameter-seed combinations that fit the data, whereas fixed-grid methods tended to converge on a single point and repeatedly sample the same region.
  • Efficiency: The adaptive strategy allowed the method to identify actionable trajectories within a limited budget (3000 simulations), which is a fraction of what traditional likelihood-free methods (like IMABC) require.

5. Significance

• Decision Support: By recovering specific trajectories rather than just average behaviors, the method provides "snapshots" of epidemic evolution that are consistent with reality. These can be used as initial conditions for sequential data assimilation or for testing specific intervention strategies.
• Efficiency: The adaptive grid strategy reduces the computational cost of calibrating expensive stochastic simulators, making it feasible to use high-fidelity ABMs in real-time or near-real-time public health scenarios.
• Generalizability: While motivated by epidemiology, the framework is applicable to any domain using stochastic simulators (e.g., climate modeling, finance, engineering) where understanding specific realizations is more valuable than aggregate statistics.

In summary, the paper presents a robust, efficient framework for "staying on track" with stochastic simulations, moving beyond the limitations of summary statistics to directly discover the specific parameter and random seed combinations that explain observed data.