From b8d8bd000ffe7668ef234188140f91365e6e5e90 Mon Sep 17 00:00:00 2001
From: Marcus <MarcusNoack@lbl.gov>
Date: Fri, 24 Jul 2026 15:35:10 -0700
Subject: [PATCH 1/2] Add knowledge-gradient and noisy-EI acquisition functions

Two lookahead optimization acquisitions, available as built-in strings for both
single-task (GPOptimizer) and multi-task (fvGPOptimizer) decision-making:

- "knowledge gradient": expected increase in the maximum of the posterior mean
  after a fantasized measurement (KGCP scheme; reference set = data points plus
  the candidate). The correlated-KG expectation E[max_i a_i + b_i Z] is computed
  exactly via the upper envelope of lines (Frazier, Powell & Dayanik 2009),
  validated against Monte Carlo.
- "noisy expected improvement": EI averaged over the posterior of the incumbent,
  making it robust to observation noise (Letham et al. 2019). Monte-Carlo over
  posterior samples of the objective at the observed points, using common random
  numbers so the score is smooth for the acquisition optimizer.

For multi-task GPs both act on the task-summed objective sum_t f(x, t), matching
how the existing multi-task expected improvement / ucb scalarize the outputs. The
scalarized cross-covariance is taken from posterior_covariance()["S_flat"] using
the task-major (k = point + Npts*task) product-space ordering.

Both are tunable through the constructor args dict (reference-set size, seed, and
MC sample count for NEI). Wired into evaluate_gp_acquisition_function on both
branches, the ask() docstring, tests (single and multi-task), and the
acquisition-functions skill.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 CLAUDE.md                             |   2 +
 gpcam/gp_optimizer_base.py            |  15 +-
 gpcam/surrogate_model.py              | 210 +++++++++++++++++++++++++-
 skills/acquisition-functions/SKILL.md |  51 +++++++
 tests/test_gpCAM.py                   |  18 ++-
 5 files changed, 291 insertions(+), 5 deletions(-)

diff --git a/CLAUDE.md b/CLAUDE.md
index 66c727e..0966887 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -138,6 +138,8 @@ Sign convention to keep straight: acquisition functions are written to be **maxi
 
 Built-in acquisition strings are dispatched in `evaluate_gp_acquisition_function`, with **separate single-task vs multi-task branches** (`x_out is not None`) that do *not* support the same set — e.g. `"maximum"`, `"minimum"`, `"gradient"`, `"probability of improvement"`, and `"target probability"` exist only on the single-task branch. `"target probability"` requires `args={'a': lower, 'b': upper}`.
 
+`"knowledge gradient"` (KGCP) and `"noisy expected improvement"` (Monte-Carlo NEI) are lookahead optimization acquisitions available on **both** branches; the multi-task form scalarizes to the task-summed objective `sum_t f(x,t)`, like the multi-task EI/UCB. Both live as standalone functions in `surrogate_model.py` (`knowledge_gradient`, `noisy_expected_improvement`) with helpers (`_expected_max_of_affine` — the exact correlated-KG line integral; `_scalarized_blocks` — reference/candidate posterior blocks from `"S_flat"`, task-major `k = point + Npts*task` ordering). They need the cross-covariance between candidates and reference points, so unlike the point-wise acquisitions they call the full (non-`variance_only`) `posterior_covariance`. Tunable via `args` keys `kg_reference_set_size`/`kg_seed` and `nei_samples`/`nei_reference_set_size`/`nei_seed`.
+
 `ask()` returns `{'x': ..., 'f_a(x)': ..., 'opt_obj': ...}`. Non-Euclidean / mixed spaces are supported by passing `input_set` as a **list of candidates** instead of a bounds array.
 
 ### Pickling
diff --git a/gpcam/gp_optimizer_base.py b/gpcam/gp_optimizer_base.py
index ef43cf5..c941072 100755
--- a/gpcam/gp_optimizer_base.py
+++ b/gpcam/gp_optimizer_base.py
@@ -406,12 +406,14 @@ def ask(self,
             `ucb`,`lcb`,`maximum`,
             `minimum`, `variance`,`expected improvement`,
             `relative information entropy`,`relative information entropy set`,
-            `probability of improvement`, `gradient`,`total correlation`,`target probability`.
+            `probability of improvement`, `gradient`,`total correlation`,`target probability`,
+            `knowledge gradient`, and `noisy expected improvement`.
             In the multi-task case (using :py:meth:`gpcam.fvGPOptimizer)
             the following built-in acquisition functions can be used:
             `variance`, `relative information entropy`,
             `relative information entropy set`, `total correlation`, `ucb`, `lcb`,
-            and `expected improvement`.
+            `expected improvement`, `knowledge gradient`, and
+            `noisy expected improvement`.
             In the multi-task case, it is highly recommended to
             deploy a user-defined acquisition function due to the intricate relationship
             of posterior distributions at different points in the output space.
@@ -435,6 +437,15 @@ def ask(self,
             total correlation: extension of mutual information to more than 2 random variables;
             target probability: probability of a target. This needs a dictionary
             args = {'a': lower bound, 'b': upper bound} to be defined.
+            knowledge gradient: the expected increase in the maximum of the posterior
+            mean after a fantasized measurement (KGCP scheme over the data points plus
+            the candidate). For multi-task GPs it acts on the task-summed objective.
+            Tunable via args = {'kg_reference_set_size': int, 'kg_seed': int}.
+            noisy expected improvement: expected improvement averaged over the
+            posterior of the incumbent (best observed value), which makes it robust to
+            observation noise (Letham et al. 2019). For multi-task GPs it acts on the
+            task-summed objective. Tunable via
+            args = {'nei_samples': int, 'nei_reference_set_size': int, 'nei_seed': int}.
         method : str, optional
             A string defining the method used to find the maximum
             of the acquisition function. Choose from `global`,
diff --git a/gpcam/surrogate_model.py b/gpcam/surrogate_model.py
index 7a0d7f3..5f9a2b4 100755
--- a/gpcam/surrogate_model.py
+++ b/gpcam/surrogate_model.py
@@ -215,7 +215,8 @@ def evaluate_gp_acquisition_function(x, acquisition_function, gpo, x_out):
     if x_out is None:
         all_acq_func = ["variance", "relative information entropy", "relative information entropy set",
                         "ucb", "lcb", "maximum", "minimum", "gradient", "expected improvement",
-                        "probability of improvement", "target probability", "total correlation"]
+                        "probability of improvement", "target probability", "total correlation",
+                        "knowledge gradient", "noisy expected improvement"]
         if acquisition_function == "variance":
             res = np.sqrt(gpo.posterior_covariance(x, variance_only=True)["v(x)"])
             return res
@@ -261,6 +262,10 @@ def evaluate_gp_acquisition_function(x, acquisition_function, gpo, x_out):
             pdf = norm.pdf(gamma)
             cdf = norm.cdf(gamma)
             return std * (gamma * cdf + pdf)
+        elif acquisition_function == "knowledge gradient":
+            return knowledge_gradient(x, gpo, x_out=None)
+        elif acquisition_function == "noisy expected improvement":
+            return noisy_expected_improvement(x, gpo, x_out=None)
         elif acquisition_function == "target probability":
             try:
                 a = gpo.args["a"]
@@ -279,7 +284,8 @@ def evaluate_gp_acquisition_function(x, acquisition_function, gpo, x_out):
 
     else:
         all_acq_func = ["variance", "relative information entropy", "relative information entropy set",
-                        "ucb", "lcb", "expected improvement", "total correlation"]
+                        "ucb", "lcb", "expected improvement", "total correlation",
+                        "knowledge gradient", "noisy expected improvement"]
         if acquisition_function == "variance":
             res = gpo.posterior_covariance(x, x_out=x_out, variance_only=True)["v(x)"]
             return np.sum(res, axis=1)
@@ -315,10 +321,210 @@ def evaluate_gp_acquisition_function(x, acquisition_function, gpo, x_out):
             pdf = norm.pdf(gamma)
             cdf = norm.cdf(gamma)
             return std * (gamma * cdf + pdf)
+        elif acquisition_function == "knowledge gradient":
+            return knowledge_gradient(x, gpo, x_out=x_out)
+        elif acquisition_function == "noisy expected improvement":
+            return noisy_expected_improvement(x, gpo, x_out=x_out)
         else:
             raise Exception("No valid acquisition function string provided. Choose from ", all_acq_func)
 
 
+##########################################################################
+# Knowledge gradient and noisy expected improvement
+#
+# Both are lookahead acquisition functions that reason about the posterior of the
+# *objective* rather than the raw (noisy) observations, which is why they need the
+# cross-covariance between candidates and reference points, not just point-wise
+# variance. For multi-task GPs (``x_out is not None``) both operate on the
+# task-summed objective g(x) = sum_t f(x, t), matching how the built-in multi-task
+# ``expected improvement`` / ``ucb`` scalarize the outputs.
+##########################################################################
+def _acq_arg(gpo, key, default):
+    """Read an optional acquisition setting from ``gpo.args`` (a dict or None)."""
+    a = getattr(gpo, "args", None)
+    if isinstance(a, dict) and key in a and a[key] is not None:
+        return a[key]
+    return default
+
+
+def _reference_points(gpo, x_out, cap, rng):
+    """Input-space data points used as the reference/fantasy set.
+
+    For multi-task GPs ``gpo.x_data`` lives in the product (input x task) space, so
+    the unique input points come from ``get_data()['x data']``. Capped (subsampled)
+    for scalability; the subsample is stable within one ``ask()`` because the rng is
+    seeded.
+    """
+    if x_out is None:
+        xr = np.asarray(gpo.x_data, dtype=float)
+    else:
+        xr = np.asarray(gpo.get_data()["x data"], dtype=float)
+    if len(xr) > cap:
+        xr = xr[rng.choice(len(xr), size=cap, replace=False)]
+    return xr
+
+
+def _assumed_observation_noise(gpo, x_out):
+    """Assumed observation-noise variance for the (scalarized) objective.
+
+    Uses the mean of the measurement variances. For the task-summed multi-task
+    objective the noise of the sum is approximated as ``len(x_out)`` times that mean.
+    """
+    try:
+        v = np.asarray(gpo.get_data()["measurement variances"], dtype=float)
+        base = float(np.mean(v)) if v.size else 1e-6
+    except Exception:
+        base = 1e-6
+    if not np.isfinite(base) or base <= 0.0:
+        base = 1e-6
+    if x_out is not None:
+        base *= len(x_out)
+    return base
+
+
+def _scalarized_blocks(gpo, x_ref, x_cand, x_out):
+    """Posterior mean and covariance of the (task-summed) objective on [x_ref; x_cand].
+
+    Returns
+    -------
+    mu_ref   : (M,)   posterior mean at reference points
+    cov_ref  : (M, M) posterior covariance among reference points
+    mu_cand  : (N,)   posterior mean at candidates
+    var_cand : (N,)   posterior variance at candidates
+    cross    : (M, N) posterior covariance between reference points and candidates
+    """
+    M, N = len(x_ref), len(x_cand)
+    stack = np.vstack([x_ref, x_cand])
+    mean = np.asarray(gpo.posterior_mean(stack, x_out=x_out)["m(x)"])
+    S_flat = np.asarray(gpo.posterior_covariance(stack, x_out=x_out)["S_flat"])
+    if x_out is None:
+        mu = mean.reshape(-1)
+        cov = S_flat
+    else:
+        # Flat product-space ordering is task-major (k = point + Npts*task), so the
+        # (Npts*T, Npts*T) covariance is a T x T grid of (Npts, Npts) blocks; summing
+        # the task blocks gives Cov(sum_t f(.,t), sum_t' f(.,t')). See the v.reshape(
+        # ..., order='F') convention in fvgp.posterior_covariance.
+        T, P = len(x_out), M + N
+        mu = mean.reshape(P, T).sum(axis=1)
+        cov = S_flat.reshape(T, P, T, P).sum(axis=(0, 2))
+    cov = 0.5 * (cov + cov.T)
+    mu_ref, mu_cand = mu[:M], mu[M:]
+    cov_ref = cov[:M, :M]
+    cross = cov[:M, M:]
+    var_cand = np.clip(np.diag(cov)[M:], 0.0, None)
+    return mu_ref, cov_ref, mu_cand, var_cand, cross
+
+
+def _expected_max_of_affine(a, b):
+    """Exact ``E[max_i (a_i + b_i Z)]`` for a standard normal ``Z``.
+
+    Computed by intersecting the lines ``a_i + b_i z`` to find the upper envelope,
+    then integrating each dominant segment against the normal density. This is the
+    core of the correlated knowledge-gradient computation
+    (Frazier, Powell & Dayanik 2009).
+    """
+    a = np.asarray(a, dtype=float)
+    b = np.asarray(b, dtype=float)
+    order = np.lexsort((a, b))          # by slope b, ties broken by intercept a
+    a, b = a[order], b[order]
+    keep = np.concatenate([b[1:] != b[:-1], [True]])   # dedupe equal slopes, keep max a
+    a, b = a[keep], b[keep]
+    n = len(a)
+    if n == 1:
+        return float(a[0])
+    idx = [0]
+    z = [-np.inf]                       # z[k] = left breakpoint of dominant line idx[k]
+    for i in range(1, n):
+        while True:
+            j = idx[-1]
+            zi = (a[j] - a[i]) / (b[i] - b[j])   # crossing of line j and line i (b[i]>b[j])
+            if len(idx) > 1 and zi <= z[-1]:
+                idx.pop(); z.pop()               # line j never reaches the envelope
+            else:
+                break
+        idx.append(i); z.append(zi)
+    a, b = a[idx], b[idx]
+    zl = np.array(z)
+    zr = np.concatenate([zl[1:], [np.inf]])
+    # On (zl_k, zr_k) line k is the max; E[(a_k+b_k Z)1{zl<Z<zr}] uses
+    # E[Z 1{alpha<Z<beta}] = phi(alpha) - phi(beta).
+    return float(np.sum(a * (norm.cdf(zr) - norm.cdf(zl))
+                        + b * (norm.pdf(zl) - norm.pdf(zr))))
+
+
+def _jitter_cholesky(cov):
+    """Cholesky factor of ``cov`` with adaptive jitter; eigen-floor as last resort."""
+    cov = 0.5 * (cov + cov.T)
+    n = len(cov)
+    jit = 1e-9 * (np.trace(cov) / max(n, 1) + 1e-12)
+    for _ in range(8):
+        try:
+            return np.linalg.cholesky(cov + jit * np.eye(n))
+        except np.linalg.LinAlgError:
+            jit *= 10.0
+    w, Q = np.linalg.eigh(cov)
+    return Q @ np.diag(np.sqrt(np.clip(w, 1e-12, None)))
+
+
+def knowledge_gradient(x, gpo, x_out=None):
+    """Knowledge-gradient acquisition (maximized).
+
+    KG(x) is the expected increase in the maximum of the posterior mean after a
+    (fantasized) measurement at ``x``, over the reference set of data points plus
+    ``x`` itself (the KGCP scheme, Scott, Frazier & Powell 2011). Returns a 1d array
+    of one non-negative score per row of ``x``.
+
+    Optional ``gpo.args`` keys: ``kg_reference_set_size`` (default 100),
+    ``kg_seed`` (default 0).
+    """
+    x = np.atleast_2d(np.asarray(x, dtype=float))
+    cap = int(_acq_arg(gpo, "kg_reference_set_size", 100))
+    rng = np.random.default_rng(int(_acq_arg(gpo, "kg_seed", 0)))
+    x_ref = _reference_points(gpo, x_out, cap, rng)
+    mu_ref, _cov_ref, mu_cand, var_cand, cross = _scalarized_blocks(gpo, x_ref, x, x_out)
+    noise = _assumed_observation_noise(gpo, x_out)
+    base_ref = np.max(mu_ref)
+    out = np.empty(len(x))
+    for j in range(len(x)):
+        denom = np.sqrt(max(var_cand[j] + noise, 1e-12))
+        a = np.append(mu_ref, mu_cand[j])                    # intercepts: current mean
+        b = np.append(cross[:, j], var_cand[j]) / denom      # slopes: predictive update
+        out[j] = _expected_max_of_affine(a, b) - max(base_ref, mu_cand[j])
+    out[out < 0.0] = 0.0                                      # KG is non-negative in theory
+    return out
+
+
+def noisy_expected_improvement(x, gpo, x_out=None):
+    """Noisy expected-improvement acquisition (maximized).
+
+    Standard EI treats the incumbent (best value so far) as known, which is wrong
+    under observation noise. Noisy-EI (Letham et al. 2019) averages EI over the
+    posterior distribution of the incumbent: the best value across Monte-Carlo
+    samples of the objective at the observed points. Returns a 1d array of one
+    non-negative score per row of ``x``. Common random numbers (a seeded rng) keep
+    the score smooth across candidate evaluations.
+
+    Optional ``gpo.args`` keys: ``nei_samples`` (default 128),
+    ``nei_reference_set_size`` (default 100), ``nei_seed`` (default 0).
+    """
+    x = np.atleast_2d(np.asarray(x, dtype=float))
+    n_samples = int(_acq_arg(gpo, "nei_samples", 128))
+    cap = int(_acq_arg(gpo, "nei_reference_set_size", 100))
+    rng = np.random.default_rng(int(_acq_arg(gpo, "nei_seed", 0)))
+    x_ref = _reference_points(gpo, x_out, cap, rng)
+    mu_ref, cov_ref, mu_cand, var_cand, _cross = _scalarized_blocks(gpo, x_ref, x, x_out)
+    L = _jitter_cholesky(cov_ref)
+    Z = rng.standard_normal((len(mu_ref), n_samples))
+    f_ref = mu_ref[:, None] + L @ Z                          # (M, K) incumbent samples
+    y_star = np.max(f_ref, axis=0)                           # (K,) sampled incumbents
+    std = np.sqrt(np.maximum(var_cand, 1e-12))               # (N,)
+    d = mu_cand[:, None] - y_star[None, :]                   # (N, K)
+    gamma = d / std[:, None]
+    ei = std[:, None] * (gamma * norm.cdf(gamma) + norm.pdf(gamma))
+    return np.mean(ei, axis=1)
+
+
 def differential_evolution(func,
                            bounds,
                            tol,
diff --git a/skills/acquisition-functions/SKILL.md b/skills/acquisition-functions/SKILL.md
index d01d4a4..1eb4429 100644
--- a/skills/acquisition-functions/SKILL.md
+++ b/skills/acquisition-functions/SKILL.md
@@ -83,12 +83,63 @@ Pass these as strings to `gpo.ask(acquisition_function=...)`:
 | Relative Information Entropy | `"relative information entropy"` | Information-theoretic exploration |
 | Relative Information Entropy Set | `"relative information entropy set"` | Batch acquisition |
 | Total Correlation | `"total correlation"` | Batch acquisition |
+| Knowledge Gradient | `"knowledge gradient"` | Optimization; lookahead, one-step-optimal, robust to noise |
+| Noisy Expected Improvement | `"noisy expected improvement"` | Optimization when observations are noisy |
 
 To sanity-check any built-in or custom acquisition on a grid of candidates without calling `ask()`:
 ```python
 scores = gpo.evaluate_acquisition_function(x_grid, acquisition_function="ucb")
 ```
 
+### Knowledge gradient and noisy expected improvement
+
+Both are optimization acquisitions (find the maximum) that reason about the posterior
+of the underlying **function** rather than the raw noisy observations. They are the
+right choice over plain `"expected improvement"` when measurements are noisy, and
+they work for both single-task and multi-task GPs.
+
+```python
+# Single-task
+gpo.ask(bounds, acquisition_function="knowledge gradient")
+gpo.ask(bounds, acquisition_function="noisy expected improvement")
+
+# Multi-task — both act on the task-summed objective sum_t f(x, t),
+# matching the built-in multi-task "expected improvement" / "ucb"
+gpo.ask(bounds, x_out=np.array([0, 1, 2]),
+        acquisition_function="knowledge gradient")
+```
+
+- **Knowledge gradient (KG)** — the expected increase in the maximum of the posterior
+  *mean* after a hypothetical measurement at the candidate (KGCP scheme: the reference
+  set is the data points plus the candidate). It is a one-step-optimal, lookahead
+  criterion: it values a measurement by how much it is expected to improve your best
+  decision, so it keeps exploring even when simple EI has gone to zero. KG is ~0 on
+  top of existing data and largest where a measurement would most raise the achievable
+  maximum.
+- **Noisy expected improvement (NEI)** — expected improvement averaged over the
+  posterior of the incumbent (the best value so far). Plain EI treats the incumbent as
+  known exactly, which is wrong under observation noise; NEI integrates over it
+  (Monte-Carlo over posterior samples of the objective at the observed points, with
+  common random numbers so the score is smooth for the optimizer).
+
+Both accept optional tuning through the constructor `args` dict:
+
+```python
+gpo = GPOptimizer(x_data, y_data, args={
+    "kg_reference_set_size": 100,   # cap on reference points (subsampled if larger)
+    "kg_seed": 0,                   # seed for the reference-set subsample
+    "nei_samples": 128,             # Monte-Carlo incumbent samples for NEI
+    "nei_reference_set_size": 100,
+    "nei_seed": 0,                  # seed -> common random numbers, smooth objective
+})
+```
+
+Cost note: both build a joint posterior over the reference set plus candidates on each
+evaluation, so they are heavier than `"variance"` or `"ucb"`. Keep
+`*_reference_set_size` modest (the default 100 subsamples the data) for large datasets.
+Use `method="global"` (the default); the `"local"` finite-difference optimizer is weak
+on their flat score surfaces, exactly as with plain EI.
+
 ## Custom Acquisition Recipes
 
 ### Upper Confidence Bound (UCB)
diff --git a/tests/test_gpCAM.py b/tests/test_gpCAM.py
index de672b5..33599bf 100755
--- a/tests/test_gpCAM.py
+++ b/tests/test_gpCAM.py
@@ -104,9 +104,11 @@ def test_basic_multi_task(client):
         gp.update_hyperparameters(opt_obj)
         time.sleep(0.1)
     gp.stop_training(opt_obj)
-    acquisition_functions = ["variance","relative information entropy","relative information entropy set","total correlation", "ucb", "expected improvement"]
+    acquisition_functions = ["variance","relative information entropy","relative information entropy set","total correlation", "ucb", "expected improvement", "knowledge gradient", "noisy expected improvement"]
     for acq_func in acquisition_functions:
         gp.evaluate_acquisition_function(np.array([[0.0,0.6],[0.1,0.2]]), np.array([0,1]), acquisition_function = acq_func)
+    gp.ask(index_set_bounds, np.array([0,1]), acquisition_function="knowledge gradient", max_iter = 2)
+    gp.ask(index_set_bounds, np.array([0,1]), acquisition_function="noisy expected improvement", max_iter = 2)
     gp.ask(index_set_bounds,np.array([0.,1.]), max_iter = 2)
     gp.ask(index_set_bounds, max_iter = 2)
 
@@ -158,11 +160,25 @@ def test_acq_funcs(client):
     r = my_gpo.ask(np.array([[0.,1.],[0.,1.],[0.,1.]]),n = 1, acquisition_function="gradient", method = "local")
     r = my_gpo.ask(np.array([[0.,1.],[0.,1.],[0.,1.]]),n = 5, acquisition_function="variance", method = "hgdl", dask_client=client)
     r = my_gpo.ask(np.array([[0.,1.],[0.,1.],[0.,1.]]),n = 1, acquisition_function="target probability", method = "local")
+    r = my_gpo.ask(np.array([[0.,1.],[0.,1.],[0.,1.]]),n = 1, acquisition_function="knowledge gradient")
+    r = my_gpo.ask(np.array([[0.,1.],[0.,1.],[0.,1.]]),n = 1, acquisition_function="noisy expected improvement")
 
     r = my_gpo.ask([np.array([0.,1.,.5])], n = 1, acquisition_function="target probability", vectorized = False)
     r = my_gpo.ask([np.array([0.,1.,.5])], n = 1, acquisition_function="variance", vectorized = False)
     r = my_gpo.ask([np.array([0.,1.,.5])], n = 1, acquisition_function="ucb", vectorized = False)
 
+    # knowledge gradient and noisy expected improvement return one non-negative
+    # score per candidate (like expected improvement)
+    grid = np.random.uniform(size=(12, 3))
+    for acq in ["knowledge gradient", "noisy expected improvement"]:
+        v = my_gpo.evaluate_acquisition_function(grid, acquisition_function=acq)
+        assert v.shape == (12,)
+        assert np.all(np.isfinite(v)) and np.all(v >= -1e-9)
+    r = my_gpo.ask([np.array([0.,1.,.5]), np.array([0.5,0.5,0.5])], n = 1,
+                   acquisition_function="knowledge gradient", vectorized = False)
+    r = my_gpo.ask([np.array([0.,1.,.5]), np.array([0.5,0.5,0.5])], n = 1,
+                   acquisition_function="noisy expected improvement", vectorized = False)
+
 def test_pickle():
     import numpy as np
     from gpcam.gp_optimizer import GPOptimizer

From 78f4d3c8822902119e39601f25004b68b8a43ce7 Mon Sep 17 00:00:00 2001
From: Marcus <MarcusNoack@lbl.gov>
Date: Fri, 24 Jul 2026 15:55:33 -0700
Subject: [PATCH 2/2] Route to knowledge-gradient / noisy-EI in
 experiment-designer and multi-task skills

The acquisition-functions skill already documents both new acquisitions; this
surfaces them where a scientist actually picks one:

- experiment-designer: acquisition decision row + template comment now point to
  "noisy expected improvement" and "knowledge gradient" for noisy/lookahead
  optimization
- multi-task-advanced: note that EI/KG/noisy-EI work multi-task on the
  task-summed objective

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 skills/experiment-designer/SKILL.md | 4 +++-
 skills/multi-task-advanced/SKILL.md | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/skills/experiment-designer/SKILL.md b/skills/experiment-designer/SKILL.md
index 89b9fc3..7589656 100644
--- a/skills/experiment-designer/SKILL.md
+++ b/skills/experiment-designer/SKILL.md
@@ -47,7 +47,7 @@ Based on their answers, decide:
 |--------|----------|
 | **Optimizer class** | Pick by the support of the observations: `GPOptimizer` for unconstrained or negative-allowed `y`; `LogGPOptimizer` if `y > 0` (intensities, rates, concentrations); `LogitGPOptimizer` if `y ∈ [0, 1]` (yields, fractions, probabilities). A plain GP on positive-only or bounded data can predict invalid values — see `transformed-optimizers-advanced` skill. |
 | **Kernel** | Default Matérn-3/2 ARD is good for most cases. Use periodic kernel if periodicity is known. Use Matérn-1/2 for rough/discontinuous data, Matérn-5/2 or SE for very smooth. See `kernel-designer` skill for custom kernels. |
-| **Acquisition function** | `'variance'` for exploration/mapping. `'expected improvement'` or `'ucb'` for optimization (UCB exposes a tunable exploration/exploitation tradeoff via `beta`). Custom callable for multi-objective or constraints. See `acquisition-functions` skill. |
+| **Acquisition function** | `'variance'` for exploration/mapping. `'expected improvement'` or `'ucb'` for optimization (UCB exposes a tunable exploration/exploitation tradeoff via `beta`). For **noisy** measurements prefer `'noisy expected improvement'` (EI robust to observation noise) or `'knowledge gradient'` (lookahead, keeps exploring when EI stalls); both also work multi-task. Custom callable for multi-objective or constraints. See `acquisition-functions` skill. |
 | **Prior mean** | Default is a **constant equal to `mean(y_data)`** — not zero. Away from data the posterior reverts to that constant. Override with `prior_mean_function=` only if they have a physical model. See `prior-mean-functions` skill. |
 | **Noise model** | Use `noise_variances` if noise is known and uniform. Use `noise_function` if noise varies. See `noise-functions` skill. |
 | **Training strategy** | `method='global'` for first training, `method='local'` for re-training during the loop. Other options: `"mcmc"` (Bayesian — returns posterior samples over hyperparameters), `"adam"` (stochastic-gradient, fast, works well for high-dimensional hyperparameter vectors like deep kernels), `"hgdl"` (distributed local+global hybrid — needs a `dask_client`). |
@@ -218,6 +218,8 @@ hp_bounds = np.array(
 # ============================================================
 acquisition_function = "variance"  # Options: "variance", "expected improvement",
                                    #          "ucb", "relative information entropy",
+                                   #          "knowledge gradient" (lookahead, noise-robust),
+                                   #          "noisy expected improvement" (EI under noise),
                                    #          or a callable
 
 # ============================================================
diff --git a/skills/multi-task-advanced/SKILL.md b/skills/multi-task-advanced/SKILL.md
index 7c46673..cc00ebb 100644
--- a/skills/multi-task-advanced/SKILL.md
+++ b/skills/multi-task-advanced/SKILL.md
@@ -172,7 +172,7 @@ while reading a length scale as a task correlation.
 
 ## Important Notes
 
-1. **Multi-task acquisition**: Use `"relative information entropy set"` / `"relative information entropy"` / `"variance"` / `"total correlation"` for batch acquisition across tasks. Pass `x_out=...` to the `ask()` call. Custom callables are supported and often advisable when you care about a specific task or combination.
+1. **Multi-task acquisition**: Use `"relative information entropy set"` / `"relative information entropy"` / `"variance"` / `"total correlation"` for batch acquisition across tasks. For optimization, `"expected improvement"`, `"knowledge gradient"`, and `"noisy expected improvement"` all work multi-task — they act on the **task-summed** objective `sum_t f(x, t)` (`"noisy expected improvement"` is the one to reach for when measurements are noisy). Pass `x_out=...` to the `ask()` call. Custom callables are supported and often advisable when you care about a specific task or combination rather than the plain sum.
 2. **Missing task observations**: `y_data` can have `np.nan` entries (e.g., task 1 wasn't measured at some x); the corresponding `noise_variances` entry **must also** be `np.nan`. The GP just ignores those entries — no imputation needed.
 3. **Default kernel**: `fvGPOptimizer(x, y)` with no kernel uses a built-in deep kernel that learns its own hyperparameters and doesn't require bounds. If you supply a custom `kernel_function`, you become responsible for the full `init_hyperparameters` + `hyperparameter_bounds` layout.
 4. **Deep kernel via NN warping**: For harder multi-task structure, `from gpcam.deep_kernel_network import Network` gives you an MLP you parametrize from `hps` — see the `kernel-designer` skill for the pattern.