Skip to content

RLQuantOpt v2: JAX upgrade planΒΆ

Goal: a GPU-native, differentiable re-implementation of ZCQPEE and its training stack, used first to replicate Grech et al., QST 11 015030 (2026), and then as the base for the follow-up work (fair baselines, robustness, meta-RL).

v1.0.0 (PyTorch/SB3 + QuTiP) stays in the repo untouched as the reference implementation. The JAX code lives in rlquantopt/jx/.

Key design decision: exploit excitation-number conservationΒΆ

The RWA Hamiltonian (paper eq. 1) commutes with the total excitation number N = a1†a1 + a2†a2 + b†b. The four computational basis states live in fixed-N sectors:

State N Sector dimension (3 levels each)
|000⟩ 0 1
|010⟩, |100⟩ 1 3 (|001⟩, |010⟩, |100⟩)
|110⟩ 2 6 (|002⟩, |011⟩, |020⟩, |101⟩, |110⟩, |200⟩)

So instead of integrating a 27-dimensional ODE with QuTiP, each 50 ps piecewise-constant sample is an exact propagator exp(-i H_N dt) of a 3Γ—3 and a 6Γ—6 Hermitian matrix (via eigh). These are exactly the amplitudes the v1 observation already uses. The realised 4Γ—4 gate is block diagonal (1 βŠ• Vβ‚‚β‚“β‚‚ βŠ• e^{iΟ†}), which gives the Weyl-chamber coordinates in closed form (no general 4Γ—4 eig, which JAX does not support on GPU), so the full reward is jit-able, vmap-able and differentiable.

Milestones and versionsΒΆ

Each milestone is one or more commits on main, tagged v2.<minor>.<patch>.

Version Milestone Status
2.0.0 This plan; version bump; rlquantopt.jx package skeleton done
2.1.0 Physics + metrics (sector propagators, closed-form Weyl coordinates), cross-checked against QuTiP and weylchamber done
2.2.0 Functional ZCQPEE env (jit/vmap), step-by-step equivalent to v1 done
2.3.0 Robustness maps of the stored RL and Krotov pulses vs the paper's Julia data (Figs. 10-12) done
2.4.0 PPO and TRPO in pure JAX; SB3 checkpoint import; training CLI done
2.5.0 Policy-level generalisation with the paper's own policies (Figs. 13-17) done
2.6.0 TRPO training replication from scratch (Figs. 5-8) done, 1 seed
2.7.0 Differentiable GRAPE baseline and QSL curve (Fig. 3) done
2.8.0 Replication report, docs/replication.md done
2.9.0 RL β†’ GRAPE refinement and robust (ensemble) GRAPE (results) done
2.10.0 Documentation site (MkDocs, GitHub Pages) done
2.11.0 Docs by topic, landing page, bibliography, hardware drift ranges, PPO vs TRPO done
2.12.0 PPO seeds, losses, drift-range robustness, RL vs GRAPE cost per device, site redesign done
2.13.0 Hypothesis page (drift, dimension, learned policies); drift literature tied to experiments; coupler/coupling drift in the env done
2.14.0 Fabrication-range comparison; test T1; matched-gate-time control done
2.14.2 Compute in logical-core hours (benchmarked step prices) done
2.15.0 Large coupler drift (Β±140 MHz): first regime where one robust pulse fails; RL warm start reaches the target in 100 steps done
2.16.0 Measurable observations (readout populations, Pauli expectations) and √iSWAP gate fidelity with free Z (idea i07) done
2.17.0 The idea farm (i01-i06, i08): research agents, toy benchmarks, run registry; ibis (i09) started: clipping, finite shots, calibration context, refinement from measured data done
2.18.0 Carrier actions: an open-loop policy fed only a frequency calibration reaches 3.9e-3 on all drifted devices with ~1e4 shots (ibis) done
2.18.1 Hippogriff's belief fixed (exact grid belief; consistency check) done
2.19.0 Ibis concluded (4 seeds, calibration stress, unseen drift, K = 30 / 60, 32 Γ— 32); the calibration MDP page; jackal started; docs consolidated (status stamps, lessons, code walkthroughs) done
2.20.0 Jackal (i10), first results: a 35Γ— sim-to-sim gap; trained on the realistic device 7.2e-3 on drifted devices; an evaluation bug fixed done
2.21.0 Jackal: close the 5Γ— margin to gradient-optimised knob schedules; then hippogriff's identification on the device open
next Hippogriff on jackal's device (identification of the drift the calibration does not see); an improvement-equivalent tracker in knob space (the farm); decoherence in the device model; the hardware contract of a target platform; three-qubit model with a spectator (T5) open

Beyond replicationΒΆ

The list drawn up after the replication, and where each item went:

  1. ~~Tune PPO (KL per update ~0.1 is too high).~~ A KL brake stalls learning on this task: PPO needs its large steps (badger, v2.17.0).
  2. ~~Named-gate targets with average gate fidelity.~~ √iSWAP with free Z, leakage included (v2.16.0). CZ is open.
  3. Fair baselines: ensemble GRAPE and RL β†’ GRAPE refinement (v2.9.0), matched gate time (v2.14.0), core-hour accounting (v2.14.2) are done; still open: β‰₯5 seeds per method, RL from demonstrations.
  4. ~~Context-conditioned policies.~~ A policy conditioned on a measured calibration is the best policy we have (ibis, v2.18.0); a belief over the drift it cannot see is next (hippogriff on jackal).
  5. Open-system training and a flux-line transfer function: the filter and the AWG are in the jackal device model; decoherence is open.
  6. ~~Measurement-based observations (finite shots).~~ Done (gecko, ibis); model-based RL was tried in several forms on the idea farm, with identification (hippogriff) the one that paid.
  7. Meta-RL over transmon parameter distributions: the drift-conditioned and calibration-conditioned policies are the amortised form of it; tracking with improvement equivalence is the next step.

Replication targets (from the paper)ΒΆ

Target Paper value Source
Speed of PE gate found by RL ~10 ns Figs. 7-8
Best errors C > 0.9999, U > 0.999, J_T ~ 1e-4 Sec. 4.2, App. B
Dominant pulse frequency ~0.86 GHz (|Ο‰1 βˆ’ Ο‰2|/2Ο€) Fig. 5
RL pulse robustness Broad low-J_T region over Β±1 % Δω1, Δω2 Fig. 10
Krotov robustness Narrow band (good guess), two pockets (bad guess) Figs. 11-12
Policy generalisation island Δω1 ∈ [βˆ’1.4, 1.6] MHz, Δω2 ∈ [βˆ’1.4, 0.6] MHz at reward β‰₯ 3.8 Fig. 16
Domain randomisation Β±0.1 % Wider region, error floor > 1e-4 Fig. 17
Training budget 13.3M steps, ~9 h CPU App. B

Reference semantics copied from v1 (paper run 05-12-24_201634)ΒΆ

  • 1000 samples of Ξ”t = 50 ps (T = 50 ns); K = 3 samples per step; 333 steps per episode.
  • Action ∈ [βˆ’1, 1]Β³ β†’ deltas Γ— 20 rad/ns, cumulative sum on the last amplitude; |A| ≀ 20 rad/ns (10/Ο€ GHz). Leaving the bound clips the amplitude and truncates the episode with reward βˆ’20Β·(1 βˆ’ t/T).
  • Reward per step: βˆ’log10(1 βˆ’ (C + 3U)/4) βˆ’ log10-offset(0.99 β†’ 0.00436) βˆ’ 1e-3Β·Ξ£|Ξ”A| within the segment.
  • Observation (28): polar (2|z|βˆ’1, arg z/Ο€) of the N=1 sector for |010⟩, |100⟩ and the N=2 sector for |110⟩, then the segment amplitudes / 20, then time in [βˆ’1, 1] (taken before the index advances), all Γ— 0.9.
  • Realised gate G[i, j] = ⟨ψ_i(t)|b_j⟩ (as Qobj.overlap), Weyl coordinates from weylchamber.c1c2c3 rounded to 8 digits, concurrence from weylchamber.concurrence.