Wrappers

Bases: Env

Wrap a functional environment as a gymnasium.Env (single instance, numpy I/O).

step and reset are jit-compiled once; subsequent calls are cheap. For batched training use :class:qenvs.wrappers.vector.GymnasiumVectorEnv or jax.vmap over the functional API directly.

Parameters:

Name	Type	Description	Default
env	Environment	functional environment.	required
seed	int	initial PRNG seed (overridden by reset(seed=...)).	0
normalize_actions	bool	if True, action_space is Box(-1, 1) and actions are multiplied by env.action_bound before stepping (recommended for SB3 and CleanRL, which assume symmetric normalised action spaces). The observation space is unchanged.	False

Episode end: every episode ends with terminated=True after env.horizon steps. QEnvS tasks are fixed-duration control problems: the final time is part of the task (the reward is the fidelity gain, so the return is the final fidelity), so reaching the horizon is a true terminal state, not a time-limit cut-off, and value bootstrapping must not continue past it. truncated is therefore always False.

Episode keys: the run key is jax.random.PRNGKey(seed) and the i-th episode since the last seeding (i = 0 at the first reset) uses keys.episode_key(run_key, i) = fold_in(run_key, i), so any episode can be regenerated from the seed and its index. reset(seed=...) restarts the run (episode index 0). Deterministic tasks ignore the key. The info of a wrapped environment may hold a nested dictionary ("diagnostics"); it is returned as a nested dictionary of numpy arrays.

Bases: VectorEnv

num_envs copies of a functional environment stepped by one jitted vmap.

Implements the gymnasium>=1.0 VectorEnv API with batched numpy I/O.

Autoreset follows AutoresetMode.NEXT_STEP: the step on which an episode ends returns its true final observation with terminated=True; the following step call ignores the action of that sub-environment, returns the fresh initial observation, reward 0, terminated=False and truncated=False. All returned arrays (observations, rewards, flags and every info value) are writable numpy copies. step expects actions of shape (num_envs, action_dim); a 1-D array of shape (num_envs,) raises ValueError.

Episodes in QEnvS always end by termination at the horizon, so truncated is always False (see the RL libraries guide).

Episode keys follow the single-run rule of :class:GymnasiumEnv: the run key is jax.random.PRNGKey(seed) and the episode with index i uses keys.episode_key(run_key, i) = fold_in(run_key, i). The j-th episode of sub-environment b has index i = j * num_envs + b, so the vector env replays, in an interleaved order, the episodes of GymnasiumEnv(seed=seed), and any episode can be regenerated from the seed and its index (stochastic resets included). The info of a wrapped environment may hold a nested dictionary ("diagnostics"); it is returned as a nested dictionary of batched arrays with its own _key masks.

Parameters:

Name	Type	Description	Default
env	Environment	functional environment.	required
num_envs	int	number of parallel sub-environments.	required
seed	int	base seed (used until reset(seed=...) is called).	0
normalize_actions	bool	if True, action_space is Box(-1, 1) and actions are multiplied by env.action_bound before stepping.	False

Uncertainty and observation

Model-parameter randomisation: a different device (parameter draw) every episode.

:class:ParameterRandomization wraps any environment whose model declares randomisable parameters (Model.parameter_specs). At reset it draws one offset per randomised parameter from the episode key (leaf keys.parameter_key(episode_key, name)), keeps it fixed for the whole episode, and runs the inner environment on model.with_parameters(...) swapped in with equinox.tree_at.

The episode data are "offset/<name>" (the drawn offset delta, in the unit of the parameter) and "value/<name>" (nominal + delta). Samples never leave the support of a parameter: a distribution whose range could leave it is rejected at construction (never clipped).

ParameterRandomization

Bases: Wrapper

Randomise model parameters per episode (fixed within the episode).

Parameters:

Name	Type	Description	Default
env	Environment	environment to wrap; its model must declare the named parameters (:meth:~qenvs.models.base.Model.parameter_specs). Several wrappers may be nested; the model of the innermost environment is the one randomised.	required
distributions	dict[str, Distribution]	mapping parameter name to an offset distribution (:class:~qenvs.distributions.Normal, Uniform, TruncatedNormal). The sampled value is nominal + offset in the unit of the parameter. A distribution that states a unit must agree with the model's unit.	required
strength	float	finite factor >= 0 multiplying the configured widths (see :mod:qenvs.distributions); 0 gives the nominal model exactly.	1.0
oracle	bool	if true, append (value - nominal) / w_ref of every randomised parameter with w_ref > 0 to the observation, in the order of distributions. w_ref is the model reference width, independent of the configured widths and of strength, so the scale is the same for a training and an evaluation distribution. The oracle tells the policy the true parameters, which a real device would have to estimate with a separate calibration; with the "state" observation it makes the observation Markov again.	False

Raises:

Type	Description
ValueError	unknown or non-randomisable parameter (the message gives the reason for dissipation rates and structural sizes), a distribution that conflicts with the support or the unit, two names with the same key tag, or an invalid strength.
TypeError	env has no :class:~qenvs.models.base.Model.

Key use: the wrapper passes the episode key it receives in reset unchanged to the inner environment, and draws parameter name from keys.parameter_key(episode_key, name), so adding, removing or reconfiguring one parameter (or any other wrapper) never changes the samples of the others.

oracle_names property

Parameters of the oracle observation: randomised and with a reference width > 0.

parameter_names property

Names of the randomised parameters, in the order of distributions.

episode_model(state)

The model of this episode (nominal model with the sampled parameters).

Control errors: miscalibrated and noisy control amplitudes.

:class:ControlErrors applies, at step k and on channel c,

u = clip((a (1 + epsilon) + o) (1 + xi_k), -B, B)

where a is the commanded amplitude, epsilon the relative gain error (fixed for the episode), o an absolute offset in the units of the action (fixed for the episode), xi_k a relative noise redrawn every step and B the action_bound of the wrapped environment. This is an explicitly defined simulation model; it does not claim to be the order in which every device produces its errors. The relative errors have the whole real line as support (a gain below -1 inverts the pulse; the simulation is defined for any real gain), the clip at B keeps the amplitude in the allowed range, and a positive gain can be enforced with a :class:~qenvs.distributions.TruncatedNormal.

ControlErrors

Bases: Wrapper

Gain error, offset and per-step noise on the control amplitudes (see module docstring).

Parameters:

Name	Type	Description	Default
env	Environment	environment to wrap.	required
amplitude_scale	Distribution | None	distribution of the relative gain error epsilon (episode fixed).	None
offset	Distribution | None	distribution of the absolute offset o in action units (episode fixed).	None
step_noise	Distribution | None	distribution of the relative per-step noise xi_k.	None
correlated	bool	if true, one value is drawn and shared by all channels (for example a common AWG gain error); if false (default) every channel gets its own draw. It applies to every enabled error source.	False

A source that is None is off. With all three off the wrapper only reports info["clip_fraction"].

Info: every reset and step adds info["clip_fraction"], the share of channels clipped at +-action_bound after the errors were applied (0 at reset). Noise followed by clipping piles probability up at the bounds, so report it with results.

Key use: the wrapper passes the episode key unchanged to the inner environment; gain and offset come from keys.calibration_key(episode_key, keys.GAIN / keys.OFFSET) and the noise of step k from keys.step_noise_key(episode_key, k), so switching one source on or off, changing a distribution, or correlated never changes another source.

applied_action(state, action)

Amplitude u sent to the inner environment at the next step, and the clip fraction.

apply_control_errors(action, gain, offset, noise, bound)

u = clip((a (1 + gain) + offset) (1 + noise), -bound, bound) and the clipped fraction.

None switches a term off (no operation at all, so the output equals clip(a) bit for bit when everything is off). Returns (u, fraction) where fraction is the share of channels with |u| > bound before clipping.

Observation modes and finite-shot measurement: ObservationModel.

The tasks of QEnvS hand the agent the full quantum state (or the accumulated propagator). ObservationModel replaces that observation by the kind of information an experiment provides: expectation values of a fixed list of observables, optionally estimated from a finite number of shots, or only the elapsed time (open-loop, "blind" control).

Modes (observation=):

• "state": the observation of the wrapped environment, unchanged (default, the released behaviour). For gate tasks this is the accumulated propagator or superoperator, which an experiment can only get from process tomography, so it is an ideal-simulator baseline and not a measurement-statistics observation (see :attr:observation_kind).
• "expectations": <O_j> for the observables of Model.observables() of the nominal model, in insertion order, then the time fraction t / duration.
• "correlations": the "expectations" list followed by the observables of Model.correlation_observables() (second-order correlators), then the time fraction. Only DickeBattery and DickeBatteryLocal define correlators among the built-in models; for any other model the base method returns none and "correlations" gives the same observables as "expectations". The common observables get the same measurement noise as in "expectations" (same key leaves).
• "time": only t / duration.

Whatever the wrapped environments append after the base observation (for example the oracle parameters of a parameter randomiser) is kept at the end of the new observation.

Finite-shot model (shots=n, only for "expectations" and "correlations"): each observable is measured n times in its own eigenbasis, in a separate experiment, on the ideal evolved state. The outcome counts are drawn from the exact multinomial distribution of the distinct eigenvalues (degenerate eigenvalues are one outcome), implemented as a chain of conditional binomial draws, and the estimate is the sample mean of the outcome eigenvalues. Its variance is exactly (<O^2> - <O>^2) / n. What is assumed: no measurement back-action on the state, no readout error, no sharing of shots between commuting observables (the cost is counted conservatively), and the observables are those of the nominal model, so the estimate is of the nominal observable on the (possibly randomised) physical state. shots=None returns the exact simulator value.

Information boundary:

1. The policy sees the observation only.
2. In the partial modes (everything but "state") info carries no exact value (fidelity, leakage, energy, ergotropy, expectations): only the cost keys shots_used and clip_fraction (when a control-error wrapper provides it). The "state" mode keeps the released info.
3. diagnostics=True adds the exact values as info["diagnostics"], for evaluation and debugging.
4. :meth:~qenvs.wrappers.Wrapper.info and :meth:~qenvs.wrappers.Wrapper.diagnostics always return the exact values; the API cannot stop a caller from feeding them to a policy.
5. The reward is not changed by the mode: it is the exact simulator gain of the task quantity, so even in the partial modes the cumulative reward reveals the exact performance trajectory. This is the usual simulation-RL assumption and differs from an experiment, where the reward would itself be estimated from measurements.

Random numbers: observable j at observation time t (t = 0 after reset, t = k after step k) draws from keys.measurement_key(episode_key, t, j), nothing else.

COST_INFO_KEYS = ('shots_used', 'clip_fraction') module-attribute

Keys that stay in info in the partial observation modes (cost, not performance).

DEGENERACY_RTOL = 1e-10 module-attribute

Eigenvalues closer than DEGENERACY_RTOL * max(1, |O|_max) form one outcome.

DUPLICATE_TOL = 1e-12 module-attribute

Two observables closer than this in operator norm are duplicates (rejected).

HERMITIAN_RTOL = 1e-10 module-attribute

Relative tolerance of max|O - O^dag| / max(1, max|O|) for an observable.

MAX_SHOTS = 2 ** 24 module-attribute

Largest shots: outcome counts are floats, exact integers up to 2**24 in float32.

ObservationModel

Bases: Wrapper

Observation mode and finite-shot measurement of an environment (see module docs).

Parameters:

Name	Type	Description	Default
env	Environment	environment to wrap (a task, or another wrapper such as a parameter randomiser; the wrapper does not change dynamics or rewards).	required
observation	str	"state" (default), "expectations", "correlations" or "time".	'state'
shots	int | None	None (exact expectation values) or a positive integer number of measurements per observable and observation, at most MAX_SHOTS = 2**24 (outcome counts stay exact integers in float32). Only allowed with "expectations" and "correlations".	None
diagnostics	bool	if True, info["diagnostics"] holds the exact values (evaluation and debugging only; never feed them to a policy).	False
observables	Sequence[str] | None	names of the Model.observables() to measure, in the order given (default: all). Needed when the model lists two names for the same operator, for example energy and p1 of a qubit with omega0 = 1, which is rejected as a duplicate.	None

Raises:

Type	Description
ValueError	unknown mode, shots that is not a positive integer, is above 2**24 or is given for "state"/"time", a gate task in "expectations"/"correlations", no observables, an unknown name, a non-Hermitian operator, two observables that are the same operator, or a diagnostics that is not a bool.
TypeError	observables given as a single string instead of a sequence.

Attributes:

Name	Type	Description
observation	str	the mode.
shots	int | None	shots per observable and observation, or None.
observable_names	tuple[str, ...]	names of the measured observables, in observation order.
observation_kind	str	what the "state" mode carries for this task ("state", "propagator" or "superoperator").
shots_per_step	int	shots * len(observable_names), the measurements of one observation (0 without shots).

measure(qstate, key, t)

The values the policy sees for the measured observables, shape (K,).

Parameters:

Name	Type	Description	Default
qstate	QArray	ideal quantum state of the (possibly randomised) physical system.	required
key	Array	episode key; observable j uses keys.measurement_key(key, t, j).	required
t	int | Array	observation time (0 after reset, k after step k).	required

outcome_probabilities(qstate)

Outcome probabilities of every measured observable on qstate.

observation_kind(env)

Kind of the "state" observation of env.

"state" for state-transfer and battery tasks, "propagator" for a closed gate task and "superoperator" for an open one. The last two are accumulated maps that an experiment only obtains by process tomography, so the "state" mode of a gate task is an ideal-simulator baseline.

outcome_counts(key, probs, shots)

Multinomial outcome counts as an exact chain of conditional binomial draws.

Outcome m gets Binomial(shots - sum_{l<m} c_l, p_m / sum_{l>=m} p_l); the last outcome with positive probability gets all that is left, so the counts always sum to shots. A probability of exactly 0 gives count 0 and an outcome that has probability 1 gets all the shots, without calling the binomial sampler (whose inversion branch does not terminate for p = 0).

Parameters:

Name	Type	Description	Default
key	Array	PRNG key; the draws of outcome m use the m-th key of split(key, M).	required
probs	Array	outcome probabilities, shape (M,), non-negative, summing to 1.	required
shots	int | Array	number of measurements (a positive integer).	required

Returns:

Type	Description
Array	Counts of shape (M,) (a float array holding integers, dtype of probs).

Offset distributions for parameter randomisation and control errors.

Every distribution draws an offset delta from the nominal value (the parameter value is nominal + delta), in the unit of the parameter. All three are frozen dataclasses with a unit field (empty means "not stated"; a stated unit must equal the unit declared by the model), analytic mean and stddev, and sample(key).

• Normal(std): scales std with strength; allowed only if the support of the parameter is the whole real line.
• Uniform(half_width): scales half_width; [nominal - h, nominal + h] with the effective half-width h must lie inside the support.
• TruncatedNormal(std, low, high): scales std only (low and high are fixed offsets); low <= 0 <= high, low < high and [nominal + low, nominal + high] must lie inside the support.

Order of operations (fixed): (1) the configured width is multiplied by strength to give the effective width, (2) the range of values the effective distribution can produce is checked against the support of the parameter at construction of the wrapper, (3) samples are drawn. Samples are never clipped to the support, since that would change the declared distribution; the only operation after the draw is a guard against floating-point rounding (TruncatedNormal and Uniform offsets are kept inside their own interval). A zero effective width gives an offset of exactly 0 (also for TruncatedNormal).

TruncatedNormal.low and high are offsets from the nominal value and do not grow with strength; when they are not symmetric the mean offset is not zero, and :meth:mean and :meth:stddev give the exact values.

Distribution

Bases: ABC

Offset distribution (base class of Normal, Uniform and TruncatedNormal).

Subclasses are frozen dataclasses. All methods take strength, the factor that turns the configured width into the effective width (see the module docstring).

check_support(name, nominal, support, strength=1.0)

Raise ValueError if the effective distribution can leave support.

nominal is the nominal parameter value and support the closed interval declared by the model. Only the range of attainable values is checked; nothing is ever clipped.

mean(strength=1.0) abstractmethod

Exact mean of the offset.

offset_range(strength=1.0) abstractmethod

Closed interval (lo, hi) of the offsets the effective distribution can give.

sample(key, strength=1.0, shape=(), dtype=None) abstractmethod

Draw offsets of the given shape from key (works under jit and vmap).

stddev(strength=1.0) abstractmethod

Exact standard deviation of the offset.

Normal dataclass

Bases: Distribution

Gaussian offset N(0, (std * strength)^2).

Allowed only for parameters whose support is the whole real line (a Gaussian has unbounded tails, so it cannot honour a bounded support); use :class:Uniform or :class:TruncatedNormal otherwise.

Attributes:

Name	Type	Description
std	float	standard deviation in the unit of the parameter (>= 0).
unit	str	unit of std; empty means "not stated".

TruncatedNormal dataclass

Bases: Distribution

Gaussian offset truncated to [low, high]: N(0, (std * strength)^2) on that interval.

low and high are offsets from the nominal value (low <= 0 <= high, low < high, both finite) and are not scaled by strength; only std is. When low != -high the mean offset is not zero: :meth:mean and :meth:stddev are the exact moments of the truncated distribution.

Attributes:

Name	Type	Description
std	float	standard deviation of the underlying Gaussian, in the unit of the parameter (>= 0; 0 gives an offset of exactly 0).
low	float	lower end of the offset interval (<= 0).
high	float	upper end of the offset interval (>= 0).
unit	str	unit of std, low and high; empty means "not stated".

Uniform dataclass

Bases: Distribution

Uniform offset on [-h, h] with h = half_width * strength.

Attributes:

Name	Type	Description
half_width	float	configured half-width in the unit of the parameter (>= 0); the standard deviation is h / sqrt(3).
unit	str	unit of half_width; empty means "not stated".

check_strength(strength)

Return strength as a float; raise ValueError unless it is finite and >= 0.

Base class of functional environment wrappers (no third-party dependencies).

A wrapper is itself an :class:~qenvs.envs.base.Environment. Its state is a :class:WrappedState that holds the inner state, the episode key and the wrapper's own episode data (sampled parameters, calibration errors). Wrappers nest: every wrapper passes the key it receives in reset unchanged to the inner environment and derives its own random numbers from that episode key through :mod:qenvs.keys, so the samples of one wrapper never depend on which other wrappers are present.

WrappedState

Bases: Module

State of a wrapped environment.

Attributes:

Name	Type	Description
inner	Any	state of the wrapped environment (an EnvState or another WrappedState).
key	Array	episode key received in reset.
step	Array	number of steps taken in this episode.
data	dict[str, Array]	the wrapper's episode data (for example sampled parameters).

base property

The innermost EnvState.

qstate property

Quantum state of the innermost environment.

t property

Time of the innermost environment.

Wrapper

Bases: Environment

Environment that delegates to env; subclasses override what they change.

Attributes not defined on the wrapper (model, dt, target, ...) are read from the wrapped environment.

unwrapped property

The innermost (unwrapped) environment.

diagnostics(state)

Exact diagnostics of state; for evaluation and debugging, never a policy input.

info(state)

Exact diagnostics of state from the inner environment (not a policy input).

rollout(actions, key=None)

Open-loop rollout of actions (horizon, action_dim); returns final state, rewards.