Skip to content

Environment: rlquantopt/jx/env.py

env.py is a functional re-implementation of v1 ZCQPEE with the semantics of the paper run (05-12-24_201634). There is no environment object: state is an explicit pytree and every function is pure. Its defaults are the paper's MDP, checked step by step against v1; everything added since (named-gate objective, measured observations, finite shots, calibration context, clipping, carrier actions, the realistic device) is a switch in EnvConfig that is off by default.

cfg = EnvConfig()                                   # static: sets shapes and episode length
obs, state = reset(key, cfg)
obs, state, reward, terminated, truncated, info = step(state, action, cfg)

The MDP in one table

Value (paper run)
Pulse 1000 samples of 50 ps (T = 50 ns), piecewise constant
Step K = 3 samples; 333 steps per episode
Action \(a \in [-1, 1]^3\): amplitude deltas, × 20 rad/ns, cumulatively summed
Amplitude bound (
Observation 28 numbers: 12 complex amplitudes in polar form, the 3 segment amplitudes / 20, time, all × 0.9
Reward (-\log_{10} J_T - 0.0044 - 10^{-3}\sum

The newer MDP

Ideas i07-i10 replace this MDP with a calibration policy: it observes a measured device calibration instead of the state and steers the knobs of a carrier instead of individual samples. See The calibration policy (new MDP).

One step

def step(state: EnvState, action, cfg: EnvConfig):
    """One env step. ``action`` in [-1, 1]^K (clipped, as SB3 does) are amplitude deltas."""
    if cfg.physics_model == "device":
        from rlquantopt.jx import env_device
        return env_device.step(state, action, cfg)
    K = cfg.n_time_steps
    knobs = state.knobs
    if cfg.action_mode == "carrier":
        action = jnp.clip(jnp.reshape(action, (3,)), -1, 1).astype(fdtype())
        knobs = knobs + action * jnp.asarray(cfg.carrier_steps, fdtype())
        knobs = knobs.at[0].set(jnp.clip(knobs[0], 0.0, cfg.a_scale)).at[2].set(jnp.clip(knobs[2], -cfg.a_scale, cfg.a_scale))
        f = cfg.carrier_ghz
        if "context" in cfg.obs_mode:       # the measured detuning, not the true one
            f = f + jnp.sign(cfg.model.omega_s[1] - cfg.model.omega_s[0]) * (state.context[1] - state.context[0]) * 1e-3
        t = (state.cur_idx + jnp.arange(K) + 0.5) * cfg.dt
        amps = knobs[2] + knobs[0] * jnp.cos(2 * jnp.pi * f * t + knobs[1])
    else:
        action = jnp.clip(jnp.reshape(action, (K,)), -1, 1).astype(fdtype())
        deltas = action * (cfg.delta_scale or cfg.a_scale)
        prev = jnp.where(state.cur_idx == 0, 0.0, state.amps_cur[-1])
        amps = prev + jnp.cumsum(deltas)
    amps_norm = amps / cfg.a_scale
    oob = jnp.max(jnp.abs(amps_norm)) > cfg.a_norm_max
    excess = jnp.sum(jnp.maximum(jnp.abs(amps_norm) - cfg.a_norm_max, 0.0))
    amps = jnp.clip(amps_norm, -cfg.a_norm_max, cfg.a_norm_max) * cfg.a_scale

    sector = physics.propagate(state.ham, state.sector, amps, cfg.dt)
    G = physics.realised_gate(sector)
    JT, C, U = metrics.cost(G, cfg.objective, cfg.concurrence_weight, cfg.unitarity_weight)
    reward = -jnp.log10(jnp.maximum(JT, cfg.jt_floor)) * cfg.rew_scale - cfg.rew_thresh
    tv = jnp.sum(jnp.abs(jnp.diff(amps))) * cfg.tv_penalty_scale
    reward = reward - tv
    if cfg.oob_mode == "clip":          # keep the pulse, penalise pushing against the bound
        reward = (reward - cfg.oob_penalty * excess).astype(fdtype())
    else:                               # v1: leaving the bound ends the episode
        oob_reward = cfg.terminal_penalty * (1 - state.cur_idx / cfg.pulse_length)
        reward = jnp.where(oob, oob_reward, reward).astype(fdtype())

    k_obs, k_next = jax.random.split(state.key)
    new_state = state._replace(sector=sector, amps_cur=amps, key=k_next, knobs=knobs)
    obs = observation(new_state, state.cur_idx, cfg, k_obs)
    cur_idx = state.cur_idx + K
    new_state = new_state._replace(cur_idx=cur_idx)
    terminated = cur_idx + K - 1 >= cfg.pulse_length
    truncated = oob if cfg.oob_mode == "terminate" else jnp.array(False)
    info = dict(JT=JT, concurrence=C, unitarity=U, tv_penalty=tv, oob=oob,
                t=cur_idx * cfg.dt)
    return obs, new_state, reward, terminated, truncated, info

Read it in four blocks:

  1. Action → amplitudes. The deltas are summed onto the last amplitude of the previous segment (0 at the first step). If any amplitude leaves the bound, it is clipped and oob is set.
  2. Physics. propagate applies the three samples to the sector state; realised_gate and cost_JT give \(J_T\), \(C\), \(U\) at the end of this segment.
  3. Reward. \(-\log_{10} J_T\) minus the v1 offset rew_thresh (\(-\log_{10} 0.99\)) minus the TV penalty on the segment. On oob the reward is replaced by the truncation penalty.
  4. Bookkeeping. The observation uses the time before the index advances (as v1 does); the episode terminates when the next segment would run past the pulse.

Nothing here is a Python branch on data: jnp.where replaces if, so the step can be compiled and vectorised.

Observation

def observation(state: EnvState, idx_for_time, cfg: EnvConfig, key=None):
    parts = []
    if cfg.obs_mode.startswith("measured"):
        pops, paulis = physics.measured_observables(state.sector)
        if cfg.shots:
            pops, paulis = shot_estimates(key, pops, paulis, cfg.shots)
        parts.append(jnp.concatenate([2 * pops - 1, paulis]))          # all in [-1, 1]
    elif cfg.obs_mode == "amplitudes":
        z = physics.sector_amplitudes(state.sector)
        parts.append(jnp.stack([2 * jnp.abs(z) - 1, jnp.angle(z) / jnp.pi], axis=-1).reshape(-1))
    if "context" in cfg.obs_mode:
        parts.append(state.context / jnp.asarray(cfg.context_scale_mhz, fdtype()))
    if cfg.action_mode == "carrier":
        A, phi, off = state.knobs
        amps = jnp.stack([A / cfg.a_scale, off / cfg.a_scale, jnp.cos(phi), jnp.sin(phi)])
    else:
        amps = state.amps_cur / cfg.a_scale / cfg.a_norm_max
    t = idx_for_time * 2 / cfg.pulse_length - 1
    obs = jnp.concatenate(parts + [amps, jnp.atleast_1d(t)]) * cfg.obs_scale
    return obs.astype(fdtype())

sector_amplitudes returns the N=1 images of |010⟩ and |100⟩ (3 amplitudes each) and the N=2 image of |110⟩ (6 amplitudes): exactly the 12 entries v1 extracted from the 27-dim state vector. Each complex amplitude \(z\) becomes \((2|z| - 1,\ \arg z / \pi)\).

The observation is a simulator privilege

State amplitudes are not measurable on hardware. That is fine for replicating the paper; for a sim-to-real version use obs_mode="measured" below.

Measured observations (obs_mode="measured")

With EnvConfig(obs_mode="measured") (CLI --obs-mode measured, idea i07) the agent sees only what a lab can read out after preparing an input state and playing the pulse so far:

# Experimentally measurable observables (idea i06): what a lab reads out after preparing an input state
# and playing the pulse so far. The coupler is not read out (traced out); each transmon is read out in
# three levels, so leakage to |2> is visible.
_A = np.zeros((3, 3)); _A[0, 0] = _A[1, 1] = 1
_PX = np.zeros((3, 3)); _PX[0, 1] = _PX[1, 0] = 1
_PY = np.zeros((3, 3), complex); _PY[0, 1], _PY[1, 0] = -1j, 1j
_PZ = np.zeros((3, 3)); _PZ[0, 0], _PZ[1, 1] = 1, -1
PAULIS_2Q = np.stack([np.kron(a, b) for a in (_A, _PX, _PY, _PZ) for b in (_A, _PX, _PY, _PZ)])   # (16, 9, 9)
COMP_IDX = np.array([0, 1, 3, 4])            # |00>, |01>, |10>, |11> among the 9 qubit states |q0 q1>
N_MEASURED = 3 * 5 + 3 * 16


def _full(idx, v):
    return jnp.zeros(FULL_DIM, v.dtype).at[idx].set(v)


def _qubit_rho(psi):
    """Reduced density matrix of the two transmons (9x9), coupler traced out."""
    p = psi.reshape(3, 3, 3)
    return jnp.einsum("abc,dec->abde", p, p.conj()).reshape(9, 9)


def measured_observables(state: SectorState):
    """63 numbers: for inputs |01>, |10>, |11> the readout probabilities of 00, 01, 10, 11 and of any
    transmon in |2>; for inputs |0+>, |+0>, |+1> the 16 two-qubit Pauli expectations (I on the
    computational levels only), which carry the relative phases that populations cannot see."""
    U1, psi2 = state
    psi01 = _full(SECTOR1_IDX, U1[:, S1_POS_010])
    psi10 = _full(SECTOR1_IDX, U1[:, S1_POS_100])
    psi11 = _full(SECTOR2_IDX, psi2)
    e000 = jnp.zeros(FULL_DIM, U1.dtype).at[0].set(1.0)
    pops = []
    for psi in (psi01, psi10, psi11):
        p = jnp.real(jnp.diag(_qubit_rho(psi)))[COMP_IDX]
        pops.append(jnp.concatenate([p, jnp.atleast_1d(1 - p.sum())]))
    paulis = []
    for psi in ((e000 + psi01) / np.sqrt(2), (e000 + psi10) / np.sqrt(2), (psi01 + psi11) / np.sqrt(2)):
        rho = _qubit_rho(psi)
        paulis.append(jnp.real(jnp.einsum("kij,ji->k", PAULIS_2Q, rho)))
    return jnp.concatenate(pops), jnp.concatenate(paulis)

63 numbers: for the inputs |01⟩, |10⟩, |11⟩ the readout probabilities of 00, 01, 10, 11 and of "a transmon in |2⟩" (15), and for the superposition inputs |0+⟩, |+0⟩, |+1⟩ the 16 two-qubit Pauli expectation values (48), which carry the relative phases that populations cannot see. The coupler is traced out, since it is not read out; an excitation left in it reads as qubits in |00⟩. Values are exact expectations for now; finite-shot estimates are the next step. On hardware each value needs its own experiment per step, which is the price of observing mid-pulse.

Objective (objective)

objective="pe" (default) is the paper's perfect-entangler \(J_T\). objective="sqrt_iswap" (CLI --objective sqrt_iswap) uses \(1 - F\), the average gate fidelity to √iSWAP after free virtual-Z corrections (Gate metrics); the reward is \(-\log_{10}(1 - F)\) with the same shaping. info["JT"] always holds the active objective's cost, so training, evaluation and GRAPE code work unchanged (GrapeConfig(objective=...)).

Finite shots, calibration context, clipping (idea i09)

  • shots=N (--shots N): the measured observables become estimates from N shots per measurement setting, drawn with the key carried in the environment state; one step costs 30 settings.
def shot_estimates(key, pops, paulis, shots):
    """Finite-shot estimates of the measured observables, ``shots`` per measurement setting.

    Populations: a multinomial over the 5 readout outcomes of each input. Pauli expectations: each
    value from a multinomial over +1, -1 and "a transmon left the computational levels" (probability
    1 - <II>), so the estimate is (n+ - n-)/N. Values sharing a setting are sampled independently (their
    correlations are ignored). One step costs 3 + 3 x 9 = 30 settings, i.e. 30 N shots.
    """
    k1, k2 = jax.random.split(key)
    p = jnp.clip(pops.reshape(3, 5), 0, 1)
    n_pop = jax.random.multinomial(k1, shots, p / p.sum(-1, keepdims=True))
    ii = jnp.repeat(paulis.reshape(3, 16)[:, :1], 16, axis=1).reshape(-1)
    q = jnp.clip(jnp.stack([(ii + paulis) / 2, (ii - paulis) / 2, 1 - ii], -1), 0, 1)
    n_pauli = jax.random.multinomial(k2, shots, q / q.sum(-1, keepdims=True))
    return (n_pop / shots).reshape(-1), (n_pauli[:, 0] - n_pauli[:, 1]) / shots
  • obs_mode="context" / "measured+context": the policy also sees the qubit and coupler frequency offsets, measured once per episode with Gaussian error context_noise_mhz (spectroscopy, before the pulse); "context" alone is an open-loop policy that needs no measurement during the pulse.
  • oob_mode="clip" (--oob-mode clip --oob-penalty λ): an amplitude beyond the bound is clipped and costs λ per unit of normalised excess instead of ending the episode.
  • delta_scale (--delta-scale): rad/ns per sample for a unit action; scale it down for longer steps (n_time_steps), or one step sweeps the whole amplitude range.

See ibis for what these do in practice.

Carrier actions (action_mode="carrier", idea i09)

The step above has a second branch. With action_mode="carrier" (--action-mode carrier) an action is three numbers in [−1, 1], scaled by carrier_steps = (1 rad/ns, 0.2 rad, 1 rad/ns) and added to three slow knobs held in state.knobs: the amplitude A (kept in [0, 20] rad/ns), the phase φ and the offset b (kept within ±20 rad/ns). The next K samples are then

\[ u(t_k) = b + A\cos(2\pi f_d t_k + \phi), \qquad t_k \text{ the sample midpoints, in absolute time,} \]

so the carrier runs on continuously across steps. \(f_d\) is the nominal qubit-qubit detuning (cfg.carrier_ghz, 0.86 GHz); in the context modes it is corrected by the measured frequency offsets, so the policy uses what spectroscopy says, not the true detuning. Everything after the amplitudes (clipping, physics, reward) is shared with the delta branch. act_dim is 3 for any K, and with obs_mode="context" the observation is 8 numbers: the calibration (3), the knobs (A, b, cos φ, sin φ) and the time. context_bias_mhz adds a fixed error to the calibration, a stale calibration for stress tests. Why this is a new MDP, with a diagram of how a pulse is built: The calibration policy.

The realistic device (physics_model="device", idea i10)

With physics_model="device" (--physics device), reset and step hand over to env_device.py, with the same interface and reward. What changes:

Paper model ("rwa") Device ("device")
Physics physics.py: RWA, excitation-number sectors device.py: no RWA, parity blocks (7 + 10 states), direct coupling g₁₂, dressed computational states (jackal)
Control Coupler frequency shift u, rad/ns Coupler flux offset from idle, in flux quanta (bound ±0.15 Φ₀ as train.py sets it), through an asymmetric-SQUID curve
Samples 50 ps 0.1 ns, held for one AWG sample (awg_dt, 1 ns), then a first-order flux-line filter (filter_tau, 0.5 ns) whose state carries across steps (state.filt)
Carrier At the qubit detuning, 0.86 GHz At the measured dressed detuning, about 0.16 GHz
Drift Frequencies, couplings (fractions or MHz) Physical: qubit MHz, coupler flux offset in mΦ₀, couplings and anharmonicities (not in the calibration)
At the bound oob_mode Always clipped with a penalty

device_simplified=True swaps in the model one would write down first (RWA, no g₁₂, a linear flux curve, ideal AWG and flux line), with the same calibration and actions, so a policy trained on it can be deployed on the full model: the sim-to-sim transfer test.

Auto-reset for training

Training runs thousands of environments in lock-step, so an environment that finishes must restart inside the compiled loop:

def step_autoreset(key, state: EnvState, action, cfg: EnvConfig, resample=True):
    """Step, and reset when the episode ends. Returns the pre-reset obs as ``info['final_obs']``.

    With ``resample=False`` the episode restarts on the same frequencies, which
    mimics v1 ZCQPEEWRD (drift drawn once per env instance).
    """
    obs, new_state, reward, terminated, truncated, info = step(state, action, cfg)
    done = terminated | truncated
    obs_reset, state_reset = reset(key, cfg) if resample else reset_to(state.omega_s, cfg)
    state_out = jax.tree_util.tree_map(lambda r, s: jnp.where(done, r, s), state_reset, new_state)
    obs_out = jnp.where(done, obs_reset, obs)
    info = dict(info, final_obs=obs)
    return obs_out, state_out, reward, terminated, truncated, info
  • Both branches are computed and tree_map(jnp.where, ...) picks one per environment.
  • info["final_obs"] keeps the last observation of the finished episode, so the agent can bootstrap the value of a truncated episode (as Stable-Baselines3 does).
  • resample=True draws new qubit frequencies at every reset (domain randomisation as the paper describes it). resample=False keeps the frequencies of the environment, which is what v1 ZCQPEEWRD actually did: its drift was drawn once per env instance.

Domain randomisation

cfg = EnvConfig(max_drift=1e-3)       # omega_s *= 1 + U(-0.1 %, +0.1 %) at every reset

sample_omega_s draws the frequencies and reset_to(omega, cfg) builds the episode's sector Hamiltonian, which lives in the state (state.ham). A vmap over environments therefore gives each environment its own physics at no extra cost.

Evaluating a fixed pulse

rollout_pulse(amps, cfg, omega_s) applies a stored pulse sample by sample and returns \(J_T\), \(C\), \(U\) after every sample. The convention matches v1: amps[k] is held on \((t_k, t_{k+1}]\), which is amplist[k + 1] in the pulse CSVs.

Verified against v1

tests/test_jx_env.py replays three random action sequences (small, medium and saturating, the last one hitting the amplitude bound) through v1 ZCQPEE and the JAX env, comparing observations, rewards, termination and truncation at every step. The JAX env is also checked against full-space expm to 1e-12; v1's adaptive ODE drifts ~1e-8 from exact over an episode.