`, `train.py`, `sb3_import.py`¶

The agents re-implement what the paper used from Stable-Baselines3 (SB3) and sb3-contrib, in pure JAX, so that rollout and update run as one compiled function next to the environment. TRPO is the paper's algorithm and is kept for the replication; PPO has been the workhorse since v2.12 (same best gate, about twice as fast, training results), and every idea on the idea farm starts from it. The research agents built on top of PPO are walked through in Research agents.

Algorithm (`--algo`)	File	Used for
`trpo`	`agents/trpo.py`	The paper replication (results)
`ppo`	`agents/ppo.py`	Everything since v2.12: seeds, drift, large coupler drift, gecko, ibis, jackal
`ppo_judge`, `chameleon`, `dragonfly`, `ppo_plus`	`agents/ppo_judge.py`, `chameleon.py`, `dragonfly.py`, `ppo_plus.py`	Ideas i01-i05 (Research agents)
`ppo_plus` on toy environments, `fox.FoxAdapter`	`toybench.py`, `foxbench.py`, `hippogriff.py`	Ideas i05, i06, i08 (own entry points)

Networks¶

ActorCritic in agents/common.py mirrors SB3's ActorCriticPolicy as used in the paper:

separate policy and value MLPs, [128, 128] by default (--hidden; the calibration policies of ibis and jackal use [64, 64]);
activation: ReLU for TRPO (what the paper run used; SB3's default, and the v1 script since June 2025, is tanh), tanh for PPO by default; --activation also offers leaky ReLU and GELU (ibis and jackal);
orthogonal initialisation: gain √2 for hidden layers, 0.01 for the policy head, 1 for the value head;
a state-independent Gaussian log-std, initialised at 0;
actions are sampled unclipped (for the log-probabilities) and clipped to [−1, 1] only in the environment.

Rollouts, PPO and evaluation talk to the policy only through sample, mode, logp, entropy and value, so another policy class can plug in unchanged: axolotl's autoregressive policy does (agents/autoregressive.py).

Collecting a rollout¶

def collect(model, runner: RunnerState, cfg: jenv.EnvConfig, n_steps: int, gamma: float,
            resample_drift: bool = True):
    """Roll out ``n_steps`` in every env. Arrays in the returned Transition are (n_steps, n_envs, ...)."""
    v_step = jax.vmap(jenv.step_autoreset, in_axes=(0, 0, 0, None, None))

    def body(runner, _):
        params, env_state, obs, key = runner
        key, k_act, k_env = jax.random.split(key, 3)
        action, logp = model.sample(params["actor"], obs, k_act)
        value = model.value(params["critic"], obs)
        n = obs.shape[0]
        next_obs, env_state, reward, term, trunc, info = v_step(
            jax.random.split(k_env, n), env_state, action, cfg, resample_drift)
        bootstrap = model.value(params["critic"], info["final_obs"])
        reward = reward + gamma * bootstrap * (trunc & ~term)
        tr = Transition(obs, action, logp, value, reward.astype(jnp.float32), term | trunc,
                        info["JT"], info["concurrence"], info["unitarity"], info["t"])
        return RunnerState(params, env_state, next_obs, key), tr

    return jax.lax.scan(body, runner, None, n_steps)

The line to understand is the truncation bootstrap: reward + gamma * V(final_obs) * (trunc & ~term). When an episode is cut short by the amplitude bound (truncation), SB3 treats the lost future as \(\gamma V(s_{\text{final}})\) instead of 0, and we do the same. Normal episode ends (termination at 50 ns) get no bootstrap.

The truncation trap

This convention is harmless when training from scratch, but after a drift an over-optimistic critic makes overshooting pay, and fine-tuning collapses into it (found in fox). With --oob-mode clip (ibis on) the episode never truncates, so the question does not arise.

Advantages¶

def gae(traj: Transition, last_value, gamma, lam):
    def body(carry, tr):
        next_adv, next_value = carry
        not_done = 1.0 - tr.done
        delta = tr.reward + gamma * next_value * not_done - tr.value
        adv = delta + gamma * lam * not_done * next_adv
        return (adv, tr.value), adv

    _, adv = jax.lax.scan(body, (jnp.zeros_like(last_value), last_value), traj, reverse=True)
    return adv, adv + traj.value

Generalised advantage estimation [schulman2015gae], as a reverse scan over time. done cuts the recursion at episode boundaries, which matters because auto-reset puts the next episode's first observation in the same array.

PPO: the loss¶

agents/ppo.py is SB3's PPO with the clipped surrogate [schulman2017]:

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

The ratio \(r = \pi_\theta(a\mid s)/\pi_{\text{old}}(a\mid s)\) is formed from the log-probabilities stored at collection time, so the rollout policy is never re-evaluated.
Advantages are normalised per minibatch, as SB3 does, before the clipped surrogate \(-\min(r\hat A,\ \operatorname{clip}(r, 1\pm\epsilon)\hat A)\) with \(\epsilon = 0.2\).
The value loss is a plain squared error to the λ-return (SB3's default, no value clipping), weighted by vf_coef = 0.5. Actor and critic are separate networks, so the weight only matters through the shared optimiser's gradient clipping.
Entropy of a Gaussian with a state-independent log-std depends only on the log-std; its coefficient is 0, as in the paper's setup.
Diagnostics: approx_kl is the low-variance estimator \(\mathbb E[(r - 1) - \log r]\), clip_frac the share of samples outside the clip range. On the gate environment they settle near 0.1 and 0.5, far above the usual 0.01-0.02 and 0.1-0.2; braking them stalls learning (badger).

PPO: one update¶

@jax.jit
def update(runner: RunnerState, opt_state):
    runner, traj = collect(model, runner, env_cfg, cfg.n_steps, cfg.gamma, cfg.resample_drift)
    last_value = model.value(runner.params["critic"], runner.obs)
    adv, ret = gae(traj, last_value, cfg.gamma, cfg.gae_lambda)
    flat = lambda x: x.reshape((cfg.batch_size,) + x.shape[2:])
    data = (flat(traj.obs), flat(traj.action), flat(traj.logp), flat(adv), flat(ret))

    def epoch(carry, key):
        params, opt_state = carry
        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 = carry
            (loss, stats), grads = jax.value_and_grad(loss_fn, has_aux=True)(params, mb)
            updates, opt_state = tx.update(grads, opt_state, params)
            return (optax.apply_updates(params, updates), opt_state), stats

        return jax.lax.scan(minibatch, (params, opt_state), mbs)

    key, k_ep = jax.random.split(runner.key)
    (params, opt_state), stats = jax.lax.scan(
        epoch, (runner.params, opt_state), jax.random.split(k_ep, cfg.n_epochs))
    runner = runner._replace(params=params, key=key)
    stats = jax.tree_util.tree_map(jnp.mean, stats)
    stats.update(rollout_stats(traj))
    return runner, opt_state, stats

Everything in one jax.jit:

Collect n_steps × n_envs transitions (64 × 128 = 8192 by default) with collect, then bootstrap the last value and compute GAE (γ = 0.99, λ = 0.95).
Flatten the time and environment axes into one batch.
Epochs as a scan: each of the 10 epochs draws a fresh permutation, splits the batch into 32 minibatches of 256 and runs a scan of gradient steps over them.
The optimiser clips the global gradient norm at 0.5 and runs Adam (ε = 1e-5) with the paper's harmonic decay, applied per gradient step: \(\eta_i = 3\cdot10^{-4} / (1 + 10^{0.4}\, i / n_{\text{grad steps}})\).

Because the rollout, GAE and all 320 gradient steps are a single XLA program, one PPO update on the gate environment takes about 1.5 s on the laptop CPU; 20M steps (2441 updates) take about an hour.

	PPO (`PPOConfig`)	TRPO (`TRPOConfig`)
Rollout per update	64 envs × 128 steps	4 envs × 2048 steps (the paper)
Actor step	10 epochs × 32 minibatches of the clipped surrogate	One natural-gradient step with a KL limit of 0.01
Critic	Same loss, same optimiser as the actor	10 epochs of Adam, minibatches of 128
Default activation	tanh (SB3); ReLU in all our PPO runs up to gecko, GELU in ibis and jackal	ReLU (the paper run)
Learning rate	3e-4, harmonic decay	3e-4, harmonic decay (critic)

TRPO: the actor step¶

This follows sb3_contrib.TRPO.train step by step [schulman2015]: the natural-gradient direction by conjugate gradient on the KL Hessian (15 iterations, damping 0.1), a maximal step that makes the KL equal 2 * target_kl along it, and a backtracking line search that must keep KL < 0.01 and improve the surrogate.

def actor_step(actor_params, obs, action, logp_old, adv):
    flat0, unravel = ravel_pytree(actor_params)
    mean_old, log_std_old = jax.lax.stop_gradient(model.dist(actor_params, obs))

    def surrogate(flat):
        mean, log_std = model.dist(unravel(flat), obs)
        return jnp.mean(adv * jnp.exp(gaussian_logp(mean, log_std, action) - logp_old))

    def kl(flat):
        mean, log_std = model.dist(unravel(flat), obs)
        return jnp.mean(gaussian_kl(mean, log_std, mean_old, log_std_old))

    obj0, g = jax.value_and_grad(surrogate)(flat0)
    grad_kl = jax.grad(kl)

    def fvp(v):
        return jax.grad(lambda f: grad_kl(f) @ v)(flat0) + cfg.cg_damping * v

    step_dir = conjugate_gradient(fvp, g, cfg.cg_max_steps)
    max_step = jnp.sqrt(2 * cfg.target_kl / (step_dir @ fvp(step_dir)))
    coeffs = cfg.line_search_shrinking_factor ** jnp.arange(cfg.line_search_max_iter)
    cands = flat0[None] + (coeffs * max_step)[:, None] * step_dir[None]
    objs = jax.vmap(surrogate)(cands)
    kls = jax.vmap(kl)(cands)
    ok = (kls < cfg.target_kl) & (objs > obj0)
    first = jnp.argmax(ok)
    success = ok.any()
    new_flat = jnp.where(success, cands[first], flat0)
    stats = dict(policy_objective=jnp.where(success, objs[first], obj0),
                 kl=jnp.where(success, kls[first], 0.0), line_search_success=success.astype(jnp.float32))
    return unravel(new_flat), stats

ravel_pytree flattens the actor parameters into one vector, so the CG and the line search are plain linear algebra.
fvp is a Hessian-vector product of the KL: the gradient of \(\nabla \mathrm{KL}\cdot v\). No Hessian is ever formed.
The line search evaluates all 10 candidate step sizes at once with vmap, then keeps the first acceptable one; SB3 evaluates them one by one, with the same result.

TRPO: one full update¶

@jax.jit
def update(runner: RunnerState, opt_state):
    runner, traj = collect(model, runner, env_cfg, cfg.n_steps, cfg.gamma, cfg.resample_drift)
    params = runner.params
    last_value = model.value(params["critic"], runner.obs)
    adv, ret = gae(traj, last_value, cfg.gamma, cfg.gae_lambda)
    flat = lambda x: x.reshape((cfg.rollout_size,) + x.shape[2:])
    obs, action, logp_old, adv, ret = map(flat, (traj.obs, traj.action, traj.logp, adv, ret))
    adv_n = (adv - adv.mean()) / (adv.std() + 1e-8)

    actor, stats = actor_step(params["actor"], obs, action, logp_old, adv_n)

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

        def mb_step(carry, mb):
            critic, opt_state = carry
            loss, grads = jax.value_and_grad(value_loss)(critic, *mb)
            updates, opt_state = tx.update(grads, opt_state, critic)
            return (optax.apply_updates(critic, updates), opt_state), loss

        return jax.lax.scan(mb_step, (critic, opt_state), mbs)

    key, k_ep = jax.random.split(runner.key)
    (critic, opt_state), v_losses = jax.lax.scan(
        epoch, (params["critic"], opt_state), jax.random.split(k_ep, cfg.n_critic_updates))
    runner = runner._replace(params=dict(actor=actor, critic=critic), key=key)
    stats.update(value_loss=v_losses.mean(), mean_reward=traj.reward.mean(), min_JT=traj.JT.min(),
                 frac_done=traj.done.mean(), std=jnp.exp(actor["log_std"]).mean())
    return runner, opt_state, stats

After the actor step the critic gets 10 epochs of Adam on the value loss (minibatches of 128), with the paper's harmonic learning-rate decay \(\eta = 3\cdot10^{-4} / (1 + 10^{0.4}\,p)\), where \(p\) is the fraction of training done. With the paper's 4 environments × 2048 steps, one update is 8192 environment steps.

The training loop¶

model, runner, opt_state = algo.init(jax.random.PRNGKey(args.seed), env_cfg, algo_cfg)
update = algo.make_update(model, env_cfg, algo_cfg)
eval_fn = jax.jit(lambda params: getattr(algo, "evaluate", evaluate)(model, params, env_cfg))
steps_per_update = algo_cfg.n_envs * algo_cfg.n_steps
eval_every = max(1, args.eval_every // steps_per_update)

rows, evals = [], {k: [] for k in ("step", "reward", "JT", "concurrence", "unitarity", "amps", "alive")}
t0 = time.perf_counter()
for i in range(algo_cfg.n_updates):
    runner, opt_state, stats = update(runner, opt_state)
    step = (i + 1) * steps_per_update
    row = dict(step=step, wall=time.perf_counter() - t0, **{k: float(v) for k, v in stats.items()})
    if (i + 1) % eval_every == 0 or i == algo_cfg.n_updates - 1:
        ep = jax.device_get(eval_fn(runner.params))
        row.update(eval_summary(ep, env_cfg))
        evals["step"].append(step)
        for k in ("reward", "JT", "concurrence", "unitarity", "amps", "alive"):
            evals[k].append(np.asarray(ep[k]))
        np.savez_compressed(os.path.join(out, "evals.npz"), **{k: np.stack(v) for k, v in evals.items()})
        with open(os.path.join(out, f"params_{step}.pkl"), "wb") as f:
            pickle.dump(jax.device_get(runner.params), f)
        print(f"[{step:>10d}] {row['wall']:7.0f}s  R={row['eval_return']:8.1f}  "
              f"minJT={row['eval_JT_min']:.2e}@{row['eval_t_JT_min']:.1f}ns  "
              f"firstPE={row['eval_t_first_PE']:.2f}ns  1-U={row['eval_1mU_at_best']:.1e}", flush=True)
    rows.append(row)
    if (i + 1) % 10 == 0 or i == algo_cfg.n_updates - 1:
        _write_csv(os.path.join(out, "progress.csv"), rows)

Every algorithm exposes the same init(key, env_cfg, algo_cfg) and make_update(model, env_cfg, algo_cfg), so the loop does not know which one it runs. Every --eval-every steps it runs one deterministic episode (mean actions) on the nominal device, and saves:

File in `runs/<name>/` (or `runs/iNN_<idea>/<name>/` with `--idea`)	Content
`config.json`	All arguments, the environment and algorithm configs, backend, precision
`progress.csv`	One row per update: losses, KL, clip fraction, rollout statistics, eval summary
`evals.npz`	Every evaluation episode: per-step reward, J_T (or 1 − F), C, U, the pulse, alive mask
`params_<step>.pkl`	Parameters at each evaluation

evals.npz is what scripts/replicate_training_figs.py turns into the paper's Figs. 5-8. The evaluation here is on the nominal device only; the drifted-device evaluations are separate scripts (coupler_drift_experiment.py, ibis_eval.py, jackal_eval.py, Scripts and CLI).

The run directory name is <algo>_<timestamp>_s<seed>_<tag>, and os.makedirs fails rather than overwrite an existing run. With --idea <animal> it lands in the idea's own folder (rlquantopt/jx/ideas.py).

Loading v1 agents¶

def load_sb3_policy(zip_path):
    import torch
    with zipfile.ZipFile(zip_path) as z:
        sd = torch.load(io.BytesIO(z.read("policy.pth")), map_location="cpu", weights_only=True)
    sd = {k: v.numpy() for k, v in sd.items()}
    hidden = []
    i = 0
    while f"mlp_extractor.policy_net.{i}.weight" in sd:
        hidden.append(sd[f"mlp_extractor.policy_net.{i}.weight"].shape[0])
        i += 2      # Linear, Tanh, Linear, Tanh, ...
    n_layers = len(hidden)

    def mlp(net, head):
        layers = {f"Dense_{j}": _dense(sd, f"mlp_extractor.{net}.{2 * j}") for j in range(n_layers)}
        layers[f"Dense_{n_layers}"] = _dense(sd, head)
        return {"params": layers}

    act_dim = sd["action_net.weight"].shape[0]
    model = ActorCritic.make(act_dim, tuple(hidden), sb3_activation(zip_path))
    params = dict(actor=dict(net=mlp("policy_net", "action_net"), log_std=jnp.asarray(sd["log_std"], jnp.float32)),
                  critic=mlp("value_net", "value_net"))
    return model, params

SB3 stores torch Linear weights as (out, in); Flax Dense kernels are (in, out), hence the transpose in _dense. The activation is read from the checkpoint's policy_kwargs: the paper policy says ReLU. Loading it with tanh gives a completely different first action (+0.81 instead of −0.17), which is how the activation mismatch was found.

from rlquantopt.jx.env import EnvConfig
from rlquantopt.jx.sb3_import import load_sb3_policy
from rlquantopt.jx.agents.common import evaluate
model, params = load_sb3_policy("rlquantopt/rl_agents/ZCQPEE_pl-1000_T-50ns_delta_mode-TRPO/05-12-24_201634/rl_model_12566528_steps.zip")
episode = evaluate(model, params, EnvConfig())     # regenerates the paper's RL pulse

Agents and training: rlquantopt/jx/agents/, train.py, sb3_import.py¶