qenvs 0.3.0a1

Physically validated quantum control environments for reinforcement learning and optimal control.

Reusable physics. Reproducible control.

QEnvS (Quantum Environment Suite) is a model zoo of open-quantum-system control problems (transmon gates, quantum batteries, qubit and three-level state transfer; thermal machines later), packaged so that you can plug them into an RL agent or a gradient-based optimiser in a few lines, instead of re-deriving the physics for every paper.

Every model ships with analytic reference results that run as tests (pi-pulse inversion, exponential decay, ergotropy bounds, collective-charging scaling, ...). If a model is in QEnvS, its physics has been checked against closed-form solutions and against an independent simulator (QuTiP).

Built on dynamiqs (JAX): environments are jit-able, vmap-able across thousands of trajectories, and differentiable, so the same model serves RL (PPO/SAC), GRAPE-style optimal control and MPC.

Install

Full instructions, GPU setup and troubleshooting: Getting started. Supported versions and the API stability policy: Stability and versioning.

pip install qenvs            # core: jax + dynamiqs
pip install 'qenvs[gym]'     # + gymnasium adapter (stable-baselines3, CleanRL, ...)

Quickstart

Functional JAX API (batch 1024 episodes on one GPU):

import jax, jax.numpy as jnp
import qenvs

env = qenvs.make("qubit/StateTransfer-v0", horizon=10, dt=0.2)

keys = jax.random.split(jax.random.PRNGKey(0), 1024)
states, ts = jax.vmap(env.reset)(keys)
actions = jnp.ones((1024, env.action_dim))
states, ts = jax.jit(jax.vmap(env.step))(states, actions)
print(ts.reward.shape, ts.info["fidelity"].shape)  # (1024,) (1024,)

Gradient-based control through the same environment:

def loss(actions):
    _, rewards = env.rollout(actions)
    return -rewards.sum()  # = -(F_final - F_initial)


grad = jax.grad(loss)(jnp.zeros((env.horizon, env.action_dim)))

Gymnasium for classic RL libraries:

from qenvs.wrappers import GymnasiumEnv
from stable_baselines3 import PPO

gym_env = GymnasiumEnv(env)
PPO("MlpPolicy", gym_env).learn(20_000)

Command line, for humans and coding agents (every command has --json; uvx qenvs ... works without installing):

qenvs list --json
qenvs info qubit/StateTransfer-v0 --json       # physics, kwargs, references, usage
qenvs run qubit/StateTransfer-v0 --policy pi-pulse --set gamma1=0.1 --json
qenvs validate                                  # per-check error vs tolerance + versions
qenvs doctor                                    # install / backend / optional extras check

Repository conventions for coding agents are in AGENTS.md; an LLM index of the docs is at /llms.txt.

Architecture

flowchart LR
    M["Model<br/>drift, controls, jump ops,<br/>initial state, observables"] --> E["Environment<br/>reset / step / rollout<br/>jit, vmap, grad"]
    E --> G["GymnasiumEnv<br/>SB3, CleanRL"]
    E --> R["registry.make()<br/>namespace/Name-vN"]
    M --> V["validation.CHECKS<br/>analytic, constraints, numerics"]
    V --> C["qenvs validate / tests"]
    E --> O["GRAPE via jax.grad<br/>PPO via vmap"]

The physics lives only in Model; everything above it is an adapter. A model is merged only with checks in qenvs.validation that compare it against closed-form results, physical constraints and an independent simulator (QuTiP).

Models and tasks

Each model has a page under docs/models/ with its Hamiltonian, assumptions and validation table.

Bring your own system: HamiltonianModel.from_operators(drift, controls, jump_ops, initial_state) accepts NumPy, JAX, dynamiqs or QuTiP operators, plugs into every environment, and qenvs validate --model my_module:MODEL runs generic physics and numerics checks on it (guide). For RL libraries, GymnasiumEnv(env, normalize_actions=True) and the batched GymnasiumVectorEnv (guide). For hardware toolchains, qenvs.pulses stores pulses with units and provenance (JSON, NPZ), resamples them exactly to an AWG sample period and converts rad/ns to MHz (guide). See the roadmap.

Baselines

examples/grape_jax.py (open-loop Adam, GRAPE-style, through env.rollout), examples/ppo_jax.py (vectorised PPO, equinox + optax) and examples/benchmark.py (table in docs/benchmarks.md). Closed qubit/StateTransfer-v0, 3 seeds, about a minute on CPU: GRAPE 1.0000, PPO 0.9997 +- 0.0002 (optimum 1.0).

examples/benchmark_battery.py charges the battery tasks and compares GRAPE and PPO with a provable optimum (closed resonant Dicke and qubit batteries reduce to a single pulse-area parameter) and a one-pulse physical baseline. GRAPE reaches the optimum on every closed task ($N$ = 2, 4, 8); on the open Dicke batteries no method found more than the best single rectangular pulse, while time-shaped GRAPE pulses win on the open qubit battery. PPO at default settings stays within about 2 percent of the closed optima. Tables in docs/benchmarks.md.

Why not X?

• qutip-qoc / GRAPE / CRAB optimise a pulse for one fixed model; QEnvS gives you the model as a reusable, tested object plus an environment interface.
• qiskit-gym, qgym are circuit-level (synthesis, routing); QEnvS is Hamiltonian/pulse-level with open-system dynamics.
• Paper repositories (e.g. ICML 2025 RL-for-quantum-control code) are reproductions of one paper; QEnvS is the shared, maintained layer under them.

Contributing

See CONTRIBUTING.md. New models must come with analytic validation tests.

Citation

Software DOI (all versions): 10.5281/zenodo.23359969. Version 0.2.0: 10.5281/zenodo.23374561 (0.1.0: 10.5281/zenodo.23359970). Machine-readable metadata is in CITATION.cff (GitHub shows a "Cite this repository" button).

@software{park_qenvs,
  author  = {Park, Beomdo},
  title   = {{QEnvS}: physically validated quantum control environments for
             reinforcement learning and optimal control},
  year    = {2026},
  version = {0.2.0},
  doi     = {10.5281/zenodo.23359969},
  url     = {https://github.com/beomdo-park/QEnvS}
}

License

Apache-2.0. Third-party attributions are listed in NOTICE.