Research agents: the code of the idea farm

Each idea adds code next to PPO rather than changing it, so the baseline stays bit-for-bit what it was and every idea can be switched off. This page walks through that code, in the order of the animals; the results are on each animal's page. Stamps as on the farm: 🟢 active 🏁 concluded 📦 archived.

Module	Idea	Entry point	Status
agents/autoregressive.py, agents/ppo_judge.py	🐸 axolotl, 🦡 badger	train.py --algo ppo_judge	🏁
agents/chameleon.py	🐲 chameleon	train.py --algo chameleon	📦
agents/dragonfly.py	🦋 dragonfly	train.py --algo dragonfly	📦
agents/ppo_plus.py, toy_envs.py, toybench.py	🦔 echidna	python -m rlquantopt.jx.toybench, train.py --algo ppo_plus	🏁
agents/fox.py, foxbench.py	🦊 fox	python -m rlquantopt.jx.foxbench	🏁
env.py (objective, obs_mode), physics.measured_observables, metrics.fidelity_free_z	🦎 gecko	train.py --objective sqrt_iswap --obs-mode measured	🏁
hippogriff.py, toy_envs.DriftTracker	🦄 hippogriff	python -m rlquantopt.jx.hippogriff	🟢
env.py (shots, oob_mode, action_mode="carrier", context), metrics.fidelity_from_observables, refine.py	🦩 ibis	train.py --action-mode carrier --obs-mode context ..., scripts/ibis_eval.py	🏁
device.py, env_device.py	🐺 jackal	train.py --physics device, scripts/jackal_eval.py	🟢

🐸 🦡 A sample judge, and a KL brake (ppo_judge.py)

The judge's target. For a set of probe samples, how well does each one's PPO gradient agree with the gradient of the rest of the batch? Per-sample gradients are one vmap of grad away:

def help_targets(model, actor_params, obs, action, logp_old, adv, probe, ref):
    """cos(per-sample PPO gradient of each probe sample, mean PPO gradient of the ref samples)."""
    def sample_loss(p, o, a, lo, ad):
        ratio = jnp.exp(model.logp(p, o[None], a[None])[0] - lo)
        return -ratio * ad          # the clipped surrogate at ratio = 1, where the update starts

    per_sample = jax.vmap(sample_loss, in_axes=(None, 0, 0, 0, 0))
    G = jax.grad(lambda p: per_sample(p, obs[ref], action[ref], logp_old[ref], adv[ref]).mean())(actor_params)
    g = jax.vmap(jax.grad(sample_loss), in_axes=(None, 0, 0, 0, 0))(
        actor_params, obs[probe], action[probe], logp_old[probe], adv[probe])
    dots = jax.vmap(lambda gi: _tree_dot(gi, G))(g)
    norms = jnp.sqrt(jax.vmap(lambda gi: _tree_dot(gi, gi))(g))
    return dots / (norms * jnp.sqrt(_tree_dot(G, G)) + 1e-12)

A small network of (state, action, reward) learns to predict this cosine with a mean and a standard deviation; only confident verdicts reweight the PPO loss, renormalised so that the step size is unchanged:

def loss_fn(params, batch, coef):
    obs, action, logp_old, adv, ret, mark = batch
    logp = model.logp(params["actor"], obs, action)
    ratio = jnp.exp(logp - logp_old)
    adv = (adv - adv.mean()) / (adv.std() + 1e-8)
    l = -jnp.minimum(ratio * adv, jnp.clip(ratio, 1 - cfg.clip_eps, 1 + cfg.clip_eps) * adv)
    pg = l.mean()
    guide = coef * (mark * l).mean()
    pg_total = (pg + guide) / (1 + coef * mark.mean())
    v = model.value(params["critic"], obs)
    v_loss = ((ret - v) ** 2).mean()
    loss = pg_total + cfg.vf_coef * v_loss - cfg.ent_coef * model.entropy(params["actor"])
    return loss, dict(pg_loss=pg, guide_loss=guide, v_loss=v_loss,
                      approx_kl=((ratio - 1) - jnp.log(ratio)).mean(),
                      clip_frac=(jnp.abs(ratio - 1) > cfg.clip_eps).mean())

The target is the advantage's sign times a gradient alignment; a judge that sees only the reward cannot predict it, which is what the runs found.

Badger's brake is SB3's target_kl: once a minibatch starts above 1.5 × the target KL, the remaining steps of the update are skipped. Inside a scan nothing can break out of the loop, so every step is computed and a flag keeps the old parameters:

def epoch(carry, key):
    perm = jax.random.permutation(key, cfg.batch_size)
    mbs = jax.tree_util.tree_map(
        lambda x: x[perm].reshape((cfg.n_minibatches, -1) + x.shape[1:]), data)

    def minibatch(carry, mb):
        params, opt_state, stop = carry
        (loss, stats), grads = jax.value_and_grad(loss_fn, has_aux=True)(params, mb, coef)
        stop = stop | (stats["approx_kl"] > 1.5 * cfg.target_kl)     # SB3: checked before the step
        updates, new_opt = tx.update(grads, opt_state, params)
        keep = lambda new, old: jax.tree_util.tree_map(lambda n, o: jnp.where(stop, o, n), new, old)
        carry = (keep(optax.apply_updates(params, updates), params), keep(new_opt, opt_state), stop)
        return carry, dict(stats, taken=~stop)

    return jax.lax.scan(minibatch, carry, mbs)

(params, opt, _), stats = jax.lax.scan(
    epoch, (runner.params, state.opt, jnp.array(False)), jax.random.split(k_ep, cfg.n_epochs))
taken = stats.pop("taken").astype(jnp.float32)
stats = jax.tree_util.tree_map(lambda v: (v * taken).sum() / jnp.maximum(taken.sum(), 1), stats)
stats["frac_steps_taken"] = taken.mean()

The autoregressive policy (autoregressive.py) emits the three amplitude deltas of a step one at a time from one head, fed its own previous delta and the running amplitude. It exposes the same sample / mode / logp / entropy / value interface as ActorCritic, and logp teacher-forces the stored actions, so collect, evaluate and the PPO loss use it unchanged.

🐲 Latent tasks from a jumpy world model (chameleon.py)

A planner picks a latent task every 8 steps: the direction of steepest rise of the model's value, mixed with a random direction orthogonal to it, so exploring never undoes exploiting:

def exploit_direction(model: Chameleon, params, z, c):
    """Unit latent direction of steepest rise of the model's value at plan_gamma."""
    gam = jnp.full(z.shape[:-1], model.cfg.plan_gamma)
    u = jax.grad(lambda z_: model._wm(params, WorldModel.v, z_, c, gam).sum())(z)
    return u / (jnp.linalg.norm(u, axis=-1, keepdims=True) + 1e-8)


def orthogonal_explore(key, u):
    """Uniformly random unit vector orthogonal to each (unit) u."""
    xi = jax.random.normal(key, u.shape, u.dtype)
    xi = xi - jnp.sum(xi * u, -1, keepdims=True) * u
    return xi / (jnp.linalg.norm(xi, axis=-1, keepdims=True) + 1e-8)


def mix_goal(u, xi, eps):
    return jnp.sqrt(1 - eps)[..., None] * u + jnp.sqrt(eps)[..., None] * xi

The world model's loss: γ-discounted latent jumps, λ-returns, VICReg

def gamma_targets(z_next, done, gamma):
    """y_t = (1-gamma) z_{t+1} + gamma y_{t+1}, with the latent held after an episode ends. (T, B, d)."""
    def body(y_next, x):
        zn, d = x
        y = (1 - gamma)[:, None] * zn + gamma[:, None] * jnp.where(d[:, None], zn, y_next)
        return y, y

    _, y = jax.lax.scan(body, z_next[-1], (z_next, done), reverse=True)
    return y


def lambda_returns(r, v, done, gamma, lam):
    """TD(lambda) return bootstrapped with v (the rollout's last step bootstraps on its own v)."""
    v_next = jnp.concatenate([v[1:], v[-1:]], 0)

    def body(g_next, x):
        r_t, vn, d = x
        g = r_t + gamma * (1 - d) * ((1 - lam) * vn + lam * g_next)
        return g, g

    _, G = jax.lax.scan(body, v[-1], (r, v_next, done.astype(jnp.float32)), reverse=True)
    return G


def model_loss(model: Chameleon, mparams, seq, h0, key):
    """seq: ChamTransition fields, (T, B, ...) for B env sequences; h0: their history state at t = 0."""
    cfg = model.cfg
    params = dict(model=mparams)
    z = model._wm(params, WorldModel.encode, seq.obs)
    z_next = jax.lax.stop_gradient(model._wm(params, WorldModel.encode, seq.next_obs))

    def hist(h, x):
        z_t, a_prev, r_prev, first = x
        keep = (~first)[:, None]
        inp = jnp.concatenate([z_t, a_prev * keep, (r_prev * ~first * cfg.model_reward_scale)[:, None]], -1)
        h = model._wm(params, WorldModel.history, h * keep, inp)
        return h, h

    _, c = jax.lax.scan(hist, h0, (z, seq.a_prev, seq.r_prev, seq.first))
    T, B = seq.done.shape
    gamma = jax.random.choice(key, jnp.asarray(cfg.gammas, jnp.float32), (B,))
    gam_tb = jnp.broadcast_to(gamma, (T, B))

    y = gamma_targets(z_next, seq.done, gamma)
    jump = model._wm(params, WorldModel.jump, z, c, seq.action, gam_tb)
    jump_loss = jnp.mean(jnp.sum((jump - jax.lax.stop_gradient(y - z)) ** 2, -1))

    q = model._wm(params, WorldModel.q, z, c, seq.action, gam_tb)
    v = model._wm(params, WorldModel.v, z, c, gam_tb)
    G = jax.lax.stop_gradient(lambda_returns(seq.r_env * cfg.model_reward_scale, v, seq.done, gamma, cfg.model_lambda))
    value_loss = jnp.mean((q - G) ** 2 + (v - G) ** 2)

    zf = z.reshape(-1, z.shape[-1])
    zc = zf - zf.mean(0)
    std = jnp.sqrt(zc.var(0) + 1e-4)
    var_loss = jnp.mean(jax.nn.relu(1 - std))
    cov = zc.T @ zc / (zf.shape[0] - 1)
    cov_loss = (jnp.sum(cov ** 2) - jnp.sum(jnp.diag(cov) ** 2)) / zf.shape[-1]
    loss = jump_loss + value_loss + cfg.vicreg_coef * (var_loss + cov_loss)
    stats = dict(m_jump=jump_loss, m_value=value_loss, m_var=var_loss, m_cov=cov_loss,
                 m_q_minus_v=jnp.mean(jnp.abs(q - v)))
    return loss, {k: v.astype(jnp.float32) for k, v in stats.items()}


MODEL_STATS = ("m_jump", "m_value", "m_var", "m_cov", "m_q_minus_v")

🦋 A guided random walk with a pessimistic search (dragonfly.py)

One act call holds the whole idea: a goal step (exploit direction, pessimistic value, the position in the golden cache, the guide's decision on persistence τ, spread σ and exploration share ε), an Ornstein-Uhlenbeck walk in latent space, and a search over candidate actions scored by alignment with the task plus an ensemble's lower bound:

def act(model: Dragonfly, params, walk: Walk, obs, entropy, cache: Cache, key, deterministic=False):
    cfg = model.cfg
    k_guide, k_walk, k_cand = jax.random.split(key, 3)
    z, traces, c = model.observe(params, walk, obs)
    new_goal = walk.first | (walk.since >= cfg.goal_every)

    # Goal step: exploit direction, pessimistic value, cache position, guide decision
    u = jnp.where(new_goal[..., None], exploit_direction(model, params, z, c), walk.u)
    v_heads = model.m(params, DFModel.v, z, c, jnp.full(z.shape[:-1], cfg.gamma))
    v_mean, v_std = v_heads.mean(-1), v_heads.std(-1)
    p = cache_position(cache, walk.r_prev * cfg.model_reward_scale * ~walk.first, v_mean)
    gin = jnp.concatenate([c, p[..., None], jnp.broadcast_to(entropy, z.shape),
                           (v_mean - cfg.lcb_kappa * v_std)[..., None], v_std[..., None]], -1)
    gmean = model.guide_mean(params, gin)
    glog_std = params["guide"]["log_std"]
    raw = gmean if deterministic else gmean + jnp.exp(glog_std) * jax.random.normal(k_guide, gmean.shape, jnp.float32)
    tau_n, sig_n, eps_n = model.decode_guide(raw)
    if deterministic:
        eps_n = jnp.zeros_like(eps_n)
    sel = lambda new, old: jnp.where(new_goal.reshape(new_goal.shape + (1,) * (new.ndim - new_goal.ndim)), new, old)
    tau, sig, eps = sel(tau_n, walk.tau), sel(sig_n, walk.sig), sel(eps_n, walk.eps)
    graw, glogp, gin = sel(raw, walk.graw), sel(gaussian_logp(gmean, glog_std, raw), walk.glogp), sel(gin, walk.gin)

    # Exploration walk (every step): Ornstein-Uhlenbeck with per-dimension persistence and spread
    eta = jax.random.normal(k_walk, z.shape, jnp.float32)
    rho = jnp.exp(-1.0 / tau)
    xi = jnp.where(walk.first[..., None], sig * eta, rho * walk.xi + jnp.sqrt(1 - rho ** 2) * sig * eta)
    xi_perp = xi - jnp.sum(xi * u, -1, keepdims=True) * u
    xi_perp = xi_perp / (jnp.linalg.norm(xi_perp, axis=-1, keepdims=True) + 1e-8)
    g = jnp.sqrt(1 - eps)[..., None] * u + jnp.sqrt(eps)[..., None] * xi_perp

    # Candidate search: alignment with the task plus the pessimistic Q
    mean = model.pi_mean(params, obs, z, g)
    log_std = params["actor"]["log_std"]
    noise = jax.random.normal(k_cand, (cfg.n_candidates,) + mean.shape, jnp.float32)
    if deterministic:
        noise = noise.at[0].set(0.0)
    cands = mean + jnp.exp(log_std) * noise
    N = cfg.n_candidates
    zb, cb = jnp.broadcast_to(z, (N,) + z.shape), jnp.broadcast_to(c, (N,) + c.shape)
    gam = jnp.full((N,) + z.shape[:-1], cfg.plan_gamma)
    align = _cos(model.m(params, DFModel.jump, zb, cb, cands, gam), g)
    q_lcb = lower_bound(model.m(params, DFModel.q, zb, cb, cands, gam), cfg.lcb_kappa)
    score = align + cfg.search_value_coef * (q_lcb - q_lcb.max(0))
    best = jnp.argmax(score, 0)
    action = jnp.take_along_axis(cands, best[None, :, None], 0)[0]
    logp = gaussian_logp(mean, log_std, action)

    walk = walk._replace(traces=traces, z_prev=z, u=u, xi=xi, tau=tau, sig=sig, eps=eps, graw=graw, glogp=glogp,
                         gin=gin, since=jnp.where(new_goal, 1, walk.since + 1))
    info = dict(z=z, c=c, g=g, new_goal=new_goal, p=p, v_mean=v_mean, v_std=v_std,
                search_gain=score.max(0) - score.mean(0), vetoed=(q_lcb - q_lcb.max(0) < -0.5).mean(0))
    return action, logp, walk, info

History as leaky-integrator traces, trained in parallel with associative_scan

def model_loss(model: Dragonfly, mparams, seq: DFTransition, traces0, z_prev0, key):
    """seq: (T, B, ...) env sequences; traces0 / z_prev0: their trace state before t = 0."""
    cfg = model.cfg
    params = dict(model=mparams)
    z = model.m(params, DFModel.encode, seq.obs)
    z_next = jax.lax.stop_gradient(model.m(params, DFModel.encode, seq.next_obs))
    keep = (~seq.first)[..., None]
    z_prev = jnp.concatenate([z_prev0[None], z[:-1]], 0)
    z_prev = jnp.where(keep, z_prev, z)
    x = jnp.concatenate([z, z - z_prev, seq.a_prev * keep, (seq.r_prev * cfg.model_reward_scale)[..., None] * keep], -1)
    dec = jnp.asarray(cfg.trace_decays, jnp.float32)
    a = dec[None, None, :, None] * keep[..., None]                      # (T, B, K, 1)
    b = (1 - dec)[None, None, :, None] * x[:, :, None, :]               # (T, B, K, D)
    A, Bc = jax.lax.associative_scan(lambda e1, e2: (e1[0] * e2[0], e2[0] * e1[1] + e2[1]), (a, b), axis=0)
    traces = A * jax.lax.stop_gradient(traces0)[None] + Bc
    c = model.m(params, DFModel.context, traces)

    T, B = seq.done.shape
    k_gam, k_mask = jax.random.split(key)
    gamma = jax.random.choice(k_gam, jnp.asarray(cfg.gammas, jnp.float32), (B,))
    gam_tb = jnp.broadcast_to(gamma, (T, B))
    y = gamma_targets(z_next, seq.done, gamma)
    jump = model.m(params, DFModel.jump, z, c, seq.action, gam_tb)
    jump_loss = jnp.mean(jnp.sum((jump - jax.lax.stop_gradient(y - z)) ** 2, -1))
    jump_rel = jump_loss / (jnp.mean(jnp.sum(jax.lax.stop_gradient(y - z) ** 2, -1)) + 1e-8)

    q = model.m(params, DFModel.q, z, c, seq.action, gam_tb)            # (T, B, H)
    v = model.m(params, DFModel.v, z, c, gam_tb)
    G = jax.lax.stop_gradient(lambda_returns(seq.r_env * cfg.model_reward_scale, v.mean(-1), seq.done, gamma,
                                             cfg.model_lambda))[..., None]
    mask = jax.random.bernoulli(k_mask, cfg.head_keep, q.shape).astype(jnp.float32)
    value_loss = jnp.sum(mask * ((q - G) ** 2 + (v - G) ** 2)) / jnp.maximum(mask.sum(), 1.0)

    zf = z.reshape(-1, z.shape[-1])
    zc = zf - zf.mean(0)
    var_loss = jnp.mean(jax.nn.relu(1 - jnp.sqrt(zc.var(0) + 1e-4)))
    cov = zc.T @ zc / (zf.shape[0] - 1)
    cov_loss = (jnp.sum(cov ** 2) - jnp.sum(jnp.diag(cov) ** 2)) / zf.shape[-1]
    loss = jump_loss + value_loss + cfg.vicreg_coef * (var_loss + cov_loss)
    stats = dict(m_jump=jump_loss, m_jump_rel=jump_rel, m_value=value_loss, m_var=var_loss, m_cov=cov_loss,
                 m_head_std=jnp.mean(q.std(-1)))
    return loss, {k: v_.astype(jnp.float32) for k, v_ in stats.items()}

🦔 One ingredient at a time (ppo_plus.py, toy_envs.py)

ppo_plus is PPO with three switches, each off by default: Ornstein-Uhlenbeck action noise with the policy conditioned on the previous noise (so the PPO ratio stays an exact likelihood ratio), a pessimistic Q-ensemble veto, and golden-cache self-imitation [oh2018]. Acting:

def act(model: Plus, params, obs, eps_prev, first, key, veto_on, ret_std):
    cfg = model.cfg
    mean, log_std = model.ac.dist(params["actor"], obs)
    cmean, clog_std = cond_dist(mean, log_std, eps_prev, first, cfg.ou_rho)
    N = cfg.n_candidates if cfg.veto else 1
    cands = cmean + jnp.exp(clog_std) * jax.random.normal(key, (N,) + mean.shape, jnp.float32)
    choice = jnp.zeros(mean.shape[:-1], jnp.int32)
    vetoed = jnp.zeros(mean.shape[:-1], bool)
    if cfg.veto:
        qs = model.q(params["q"], jnp.broadcast_to(obs, (N,) + obs.shape), cands)
        lcb = qs.mean(-1) - cfg.lcb_kappa * qs.std(-1)                     # (N, B)
        ok = lcb >= lcb.max(0) - cfg.veto_delta * ret_std
        vetoed = veto_on & ~ok[0]
        choice = jnp.where(vetoed, jnp.argmax(ok, 0), 0)
    action = jnp.take_along_axis(cands, choice[None, :, None], 0)[0]
    logp = cond_logp(cmean, clog_std, action)
    eps = (action - mean) / jnp.exp(log_std)                              # the noise actually used
    return action, logp, eps, vetoed

The golden cache keeps the best transitions by return and by reward, at most one per ball in standardised observation space (canonical), rebuilt every rollout inside jit with a fixed-size fori_loop:

def update_cache(cfg: PlusConfig, cache: Cache, obs, action, eps_prev, first, G, r):
    M, K = cfg.cache_candidates, cfg.cache_size
    idx = jnp.concatenate([jnp.argsort(-G)[:M], jnp.argsort(-r)[:M]])
    pool = Cache(*(jnp.concatenate([a, b[idx]]) for a, b in zip(
        cache, (obs, action, eps_prev, first, G, r, jnp.ones(obs.shape[0], bool)))))
    scale = obs.std(0) + 1e-6
    big = 1e9
    norm = lambda x: (x - jnp.min(jnp.where(pool.valid, x, big))) / (
        jnp.max(jnp.where(pool.valid, x, -big)) - jnp.min(jnp.where(pool.valid, x, big)) + 1e-6)
    score = jnp.where(pool.valid, 0.5 * (norm(pool.G) + norm(pool.r)), -jnp.inf)
    order = jnp.argsort(-score)
    zs = pool.obs / scale

    def body(i, carry):
        acc, count = carry
        j = order[i]
        dist = jnp.linalg.norm(zs - zs[j], axis=-1) / jnp.sqrt(zs.shape[-1])
        ok = pool.valid[j] & ~jnp.any(acc & (dist < cfg.cache_radius)) & (count < K)
        return acc.at[j].set(ok), count + ok

    acc, count = jax.lax.fori_loop(0, score.shape[0], body, (jnp.zeros(score.shape[0], bool), 0))
    new = Cache(*(a[jnp.nonzero(acc, size=K, fill_value=0)[0]] for a in pool))
    return new._replace(valid=jnp.arange(K) < count)

toy_envs.py gives Pendulum, MountainCar and a Tracker (a miniature gate problem: delta actions on a bounded amplitude, \(-\log_{10}\) error reward) the same functional interface as the gate environment, and wraps the gate environment itself (Ham) in it, so one agent runs on all of them. DriftTracker adds a hidden device gain and offset for fox and hippogriff.

🦊 An improvement-equivalent model (fox.py)

Fox does not model the environment. After every PPO update it records the step taken in actor-parameter space and the return change it caused, and fits a Bayesian linear regression of the gain per unit step on the step's direction, in the basis of the last 8 PPO steps; the prior and noise precisions are learned by evidence maximisation [mackay1992]:

def fit(self):
    obs, w = self._data()
    if self.D is None or len(obs) < 2:
        self.post = None
        return
    Phi = np.stack([self.features(o.step) for o in obs])
    y = np.array([o.dJ / (np.linalg.norm(o.step) + 1e-12) for o in obs])
    m = Phi.shape[1]
    # Evidence maximisation (MacKay): learn the prior precision alpha and the noise precision b from the
    # data, so the scale of "return per unit step" is not fixed a priori
    var_y = max(float(np.mean(y ** 2)), self.cfg.noise_floor)
    alpha, b = self.cfg.prior / var_y, 1.0 / var_y
    n_eff = float(w.sum())
    for _ in range(30):
        A = alpha * np.eye(m) + b * (Phi.T * w) @ Phi
        A_inv = np.linalg.inv(A)
        mu = b * A_inv @ (Phi.T * w) @ y
        gamma = float(np.clip(m - alpha * np.trace(A_inv), 1e-6, m))
        alpha = gamma / max(float(mu @ mu), 1e-12)
        rss = float(np.sum(w * (y - Phi @ mu) ** 2))
        b = max(n_eff - gamma, 1e-6) / max(rss, self.cfg.noise_floor * n_eff)
    A_inv = np.linalg.inv(alpha * np.eye(m) + b * (Phi.T * w) @ Phi)
    mu = b * A_inv @ (Phi.T * w) @ y
    self.post = (mu, A_inv, 1.0 / b)

def predict(self, step):
    """Predicted return gain per unit step length along ``step``, and its standard deviation."""
    mu, A_inv, s2 = self.post
    phi = self.features(step)
    return float(phi @ mu), float(np.sqrt(phi @ A_inv @ phi))

def direction(self):
    """Posterior mean improvement direction (unit, parameter space) and gain per unit step along it."""
    mu, A_inv, _ = self.post
    g = self.D.T @ mu
    n = np.linalg.norm(g)
    return (g / n if n > 0 else g), float(n)

FoxAdapter wraps any PPO update: it adds half a PPO step along the posterior direction when the predicted gain is confident, keeps unexplained updates in an outlier cache, and stops when the optimistic predicted gain of another update falls below its sample cost. The principle behind it, and where it goes next, is on the farm page.

🦎 Measurable observations and a named gate (physics.py, metrics.py)

Gecko changed the environment, not the agent: objective="sqrt_iswap" rewards the average gate fidelity to √iSWAP after free virtual-Z corrections (Gate metrics), and obs_mode="measured" shows readout populations and Pauli expectations instead of state amplitudes (Environment).

🦄 A belief over the drift (hippogriff.py)

Three beliefs over the device's drift z, all pure functions, so the whole deployment episode is one scan: the linearised Kalman update (Joseph form), the same with a consistency check (inflate the covariance when the normalised innovation is too large), and exact Bayes on a grid; plus the expected information gain that drives probing:

def prior(d=2):
    return jnp.zeros(d, jnp.float32), jnp.eye(d, dtype=jnp.float32) / 3.0      # U(-1, 1): variance 1/3


def rho_of(P, P0):
    """Geometric-mean standard deviation of the belief relative to the prior's."""
    d = P.shape[0]
    return jnp.exp((jnp.linalg.slogdet(P)[1] - jnp.linalg.slogdet(P0)[1]) / (2 * d))


def ekf_update(m, P, y, yhat, J, R):
    S = J @ P @ J.T + R
    K = P @ J.T @ jnp.linalg.inv(S)
    A = jnp.eye(P.shape[0], dtype=P.dtype) - K @ J
    P = A @ P @ A.T + K @ R @ K.T                                   # Joseph form: stays symmetric PSD
    return m + K @ (y - yhat), 0.5 * (P + P.T)


def info_gain(J, P, R):
    return 0.5 * jnp.linalg.slogdet(jnp.eye(R.shape[0], dtype=R.dtype) + jnp.linalg.solve(R, J @ P @ J.T))[1]


def ekf_nis_update(m, P, y, yhat, J, R, P0, threshold):
    """EKF update after a consistency check: if the normalised innovation squared exceeds ``threshold``, the
    belief was too sure; P is scaled up by NIS / dim(y) (at most back to the prior's volume) first."""
    nu = y - yhat
    nis = nu @ jnp.linalg.solve(J @ P @ J.T + R, nu)
    d = P.shape[0]
    cap = jnp.exp((jnp.linalg.slogdet(P0)[1] - jnp.linalg.slogdet(P)[1]) / d)       # to the prior's volume
    alpha = jnp.where(nis > threshold, jnp.clip(nis / y.shape[0], 1.0, jnp.maximum(cap, 1.0)), 1.0)
    m, P = ekf_update(m, alpha * P, y, yhat, J, R)
    return m, P, nis


def make_grid(n):
    lin = jnp.linspace(-1, 1, n, dtype=jnp.float32)
    g1, g2 = jnp.meshgrid(lin, lin, indexing="ij")
    return jnp.stack([g1.ravel(), g2.ravel()], -1)


def grid_moments(logw, Z):
    w = jax.nn.softmax(logw)
    m = w @ Z
    d = Z - m
    return m, (d.T * w) @ d + 1e-8 * jnp.eye(Z.shape[1], dtype=Z.dtype)


def grid_update(logw, y, H, R):
    """Exact Bayes on the grid: add the Gaussian log-likelihood of y under each grid point's prediction H."""
    logw = logw - 0.5 * jnp.sum((y - H) ** 2 / jnp.diag(R), -1)
    return logw - jax.scipy.special.logsumexp(logw)


def predictive_info_gain(H, w, R):
    """1/2 log det(I + R^-1 C), C the spread of the predictions H over the belief w: exact for a linear
    model, and it sees the nonlinearity a Jacobian at the mean misses."""
    d = H - w @ H
    C = (d.T * w) @ d
    return 0.5 * jnp.linalg.slogdet(jnp.eye(R.shape[0], dtype=R.dtype) + jnp.linalg.solve(R, C))[1]


def sigma_points(m, P):
    d = m.shape[0]
    L = jnp.linalg.cholesky(d * P + 1e-9 * jnp.eye(d, dtype=P.dtype))
    return jnp.concatenate([m[None], m + L.T, m - L.T])

One deployment episode: belief update, trust region, sigma-point safety, probing

def make_episode(env: DriftTracker, model, actor_params, meas, gcfg: GriffinConfig, mode: str):
    """One deployment episode on a device, carrying the belief (m, P) in and out. vmapped over devices."""
    A = env.act_dim
    uses_filter = mode.startswith(("filter", "hippogriff"))
    R = meas.R if meas is not None else None
    grid = gcfg.filter == "grid"
    Z = make_grid(gcfg.grid_n) if grid else None
    on_grid = lambda s, a: jax.vmap(lambda z: meas.measure(s, a, z))(Z)                # (G, 2) predictions

    def episode(key, gain, offset, m, P, logw, P0):
        z_true = env.z_of(gain, offset)
        k0, key = jax.random.split(key)
        _, s = env.reset_with(k0, gain, offset, zhat=m)

        def body(carry, k):
            s, m, P, logw, alive, ret = carry
            zhat = {"oracle": z_true, "robust": jnp.zeros(2), "blind": jnp.zeros(2)}.get(mode, m)
            s = s._replace(zhat=zhat.astype(jnp.float32))
            mu = model.ac.mode(actor_params, env.obs(s))
            rho = rho_of(P, P0)
            probing = jnp.array(False)
            if mode.startswith("hippogriff"):
                c = gcfg.c_min + (1 - gcfg.c_min) * (1 - rho)
                noise = gcfg.probe_sigma * jax.random.normal(k, (gcfg.n_candidates, A), jnp.float32)
                cands = jnp.clip(jnp.concatenate([mu[None], jnp.zeros((1, A), jnp.float32), mu + noise]), -c, c)
                sp = sigma_points(m, P)
                worst = jax.vmap(lambda a: jax.vmap(lambda z: meas.max_out(s, a, z))(sp).max())(cands)
                safe = worst <= 1 - gcfg.safety_margin
                if grid:
                    w = jax.nn.softmax(logw)
                    ig = jax.vmap(lambda a: predictive_info_gain(on_grid(s, a), w, R))(cands)
                else:
                    Js = jax.vmap(lambda a: jax.jacfwd(lambda z: meas.measure(s, a, z))(m))(cands)
                    ig = jax.vmap(lambda J: info_gain(J, P, R))(Js)
                close = -jnp.sum((cands - jnp.clip(mu, -c, c)) ** 2, -1)
                probing = rho > gcfg.rho_stop
                score = jnp.where(probing, ig + gcfg.tie_beta * close, close)
                score = jnp.where(safe, score, -jnp.inf)
                a = cands[jnp.where(safe.any(), jnp.argmax(score), jnp.argmin(worst))]
            else:
                a = mu
            _, s2, r, term, trunc, info = env.step(s, a)
            nis = jnp.zeros((), jnp.float32)
            if uses_filter:
                noise = jnp.array([gcfg.noise_du, gcfg.noise_score], jnp.float32)
                y = jnp.stack([s2.u - s.u, info["score"]]) + noise * jax.random.normal(jax.random.fold_in(k, 1), (2,), jnp.float32)
                if grid:
                    logw2 = grid_update(logw, y, on_grid(s, a), R)
                    m2, P2 = grid_moments(logw2, Z)
                    logw = jnp.where(alive, logw2, logw)
                else:
                    J = jax.jacfwd(lambda z: meas.measure(s, a, z))(m)
                    yhat = meas.measure(s, a, m)
                    if gcfg.filter == "ekf_nis":
                        m2, P2, nis = ekf_nis_update(m, P, y, yhat, J, R, P0, gcfg.nis_threshold)
                    else:
                        m2, P2 = ekf_update(m, P, y, yhat, J, R)
                        nu = y - yhat
                        nis = nu @ jnp.linalg.solve(J @ P @ J.T + R, nu)
                m = jnp.where(alive, jnp.clip(m2, -gcfg.z_clip, gcfg.z_clip), m)
                P = jnp.where(alive, P2, P)
            out = dict(zerr=jnp.linalg.norm(m - z_true), rho=rho_of(P, P0), probing=probing & alive,
                       a_max=jnp.max(jnp.abs(a)), alive=alive, nis=nis)
            ret = ret + r * alive
            alive = alive & ~(term | trunc)
            return (s2, m, P, logw, alive, ret), out

        (_, m, P, logw, _, ret), out = jax.lax.scan(
            body, (s, m, P, logw, jnp.array(True), jnp.zeros((), jnp.float32)), jax.random.split(key, env.max_steps))
        return ret, m, P, logw, out

    return jax.jit(jax.vmap(episode, in_axes=(0, 0, 0, 0, 0, 0, None)))


def initial_belief(gcfg: GriffinConfig, n_devices):
    """(m, P, logw, P0) for n_devices at the prior; for the grid filter P0 is the uniform grid's own spread."""
    if gcfg.filter == "grid":
        Z = make_grid(gcfg.grid_n)
        logw = jnp.full((n_devices, Z.shape[0]), -jnp.log(Z.shape[0]), jnp.float32)
        m0, P0 = grid_moments(logw[0], Z)
    else:
        (m0, P0), logw = prior(), jnp.zeros((n_devices, 1), jnp.float32)
    return jnp.tile(m0, (n_devices, 1)), jnp.tile(P0, (n_devices, 1, 1)), logw, P0

🦩 Data-lean control (env.py, metrics.py, refine.py)

The environment options (finite shots, clipping with a penalty, calibration context, the carrier action) are in Environment and on The calibration policy (new MDP). Two pieces work from measured data only.

The fidelity from the 30 measurement settings, with no Hamiltonian: magnitudes from the readout populations, phases from the Pauli expectations:

def gate_from_observables(pops, paulis):
    """Reconstruct the physical gate <b_i|U|b_j> from the 63 measured observables (physics.measured_observables
    or their finite-shot estimates), assuming only that excitation number is conserved (block-diagonal gate,
    |00> untouched). Magnitudes from the readout populations, phases from the Pauli expectations:
    input |0+> gives <IX> + i<IY> = a = <01|U|01> and <XI> + i<YI> = c = <10|U|01>; |+0> gives b and d the same
    way; |+1> gives a* e and c* e (e = <11|U|11>) from the q1 = 1 and q0 = 1 projected coherences.
    Coherences with the coupler excited are neglected, a bias of the order of the coupler population."""
    p = pops.reshape(3, 5)
    s = paulis.reshape(3, 16)
    ph = lambda z: z / jnp.maximum(jnp.abs(z), 1e-12)
    mag = lambda x: jnp.sqrt(jnp.clip(x, 0, 1))
    a = mag(p[0, 1]) * ph(s[0, 1] + 1j * s[0, 2])
    c = mag(p[0, 2]) * ph(s[0, 4] + 1j * s[0, 8])
    b = mag(p[1, 1]) * ph(s[1, 1] + 1j * s[1, 2])
    d = mag(p[1, 2]) * ph(s[1, 4] + 1j * s[1, 8])
    w1 = ((s[2, 4] - s[2, 7]) + 1j * (s[2, 8] - s[2, 11])) / 2      # a* e, from X/Y on q0 with q1 in |1>
    w2 = ((s[2, 1] - s[2, 13]) + 1j * (s[2, 2] - s[2, 14])) / 2     # c* e, from X/Y on q1 with q0 in |1>
    e = mag(p[2, 3]) * ph(a * w1 + c * w2)
    return jnp.array([[1, 0, 0, 0], [0, a, b, 0], [0, c, d, 0], [0, 0, 0, e]], dtype=jnp.complex128)


def fidelity_from_observables(pops, paulis, V=SQRT_ISWAP):
    """Average gate fidelity to V (free Z) estimated from measured data only: no Hamiltonian involved."""
    P = gate_from_observables(pops, paulis)
    return fidelity_free_z(P.conj().T, V)[0]         # fidelity_free_z takes the realised_gate convention

Model-free refinement. SPSA [spall1992] estimates a gradient in all coordinates of a smooth pulse correction from two noisy fidelity estimates per iteration; every shot is counted. CMA-ES [hansen2016] is the other optimiser in refine.py.

def spsa(measure, exact, n_basis, shots, max_shots, seed=0, c=0.05, first_step=0.05, A=20, calib=10):
    rng = np.random.default_rng(seed)
    key = jax.random.PRNGKey(seed)
    x = np.zeros(n_basis)
    used, tr_s, tr_j = 0, [], []

    def f(x):
        nonlocal key, used
        key, k = jax.random.split(key)
        used += SETTINGS * shots
        return float(measure(jnp.asarray(x), k))

    g = [abs(f(x + c * d) - f(x - c * d)) / (2 * c) for d in rng.choice([-1.0, 1.0], (calib, n_basis))]
    a = first_step * (A + 1) ** 0.602 / max(np.mean(g), 1e-9)
    k = 0
    while used + 2 * SETTINGS * shots <= max_shots:
        ak, ck = a / (k + 1 + A) ** 0.602, c / (k + 1) ** 0.101
        d = rng.choice([-1.0, 1.0], n_basis)
        grad = (f(x + ck * d) - f(x - ck * d)) / (2 * ck) * d
        x = x - ak * grad
        tr_s.append(used)
        tr_j.append(float(exact(jnp.asarray(x))))
        k += 1
    return x, Trace(np.asarray(tr_s), np.asarray(tr_j))

🐺 A realistic device (device.py, env_device.py)

device.py is the tunable-coupler device of Sung et al. [sung2021] without the rotating-wave approximation. Without it, excitation number is no longer conserved, but parity is, so the Hamiltonian splits into an even block (7 states) and an odd one (10 states), propagated exactly per sample as in physics.py. The coupler frequency follows an asymmetric-SQUID flux curve, the couplings scale with it, and the computational states are the dressed idle eigenstates. env_device.py puts the carrier knobs on the coupler flux, holds each value for one AWG sample, low-pass filters it, draws physical drift and builds the spectroscopy calibration; env.py hands over to it when physics_model == "device". The model is described, with its sources, on the jackal page.