Reinforcement Learning for Quantum State Preparation

For the Spring 2026 offering of Caltech’s Ph 22, Computational Physics Laboratory III, I was one of the two teaching assistants who designed and built the final project, in which students train reinforcement learning agents to control a quantum system.

The physics setting is the Tavis–Cummings model, the canonical cavity quantum electrodynamics system of atoms coupled to an optical cavity. The agent chooses laser pulses, each with a duration, an amplitude, and a phase. The environment integrates the Schrödinger equation for each pulse and rewards the agent for steering the system toward a target quantum state such as a Fock state, a coherent state, a Bell state, or a Schrödinger cat state.

I built the Gymnasium environment wrapping the physics simulator, the shaped reward design, the graded ladder of target states, and the baseline policies students must beat. Students implement REINFORCE from scratch and can compare it against PPO, SAC, and TQC through Stable-Baselines3, across MLP, LSTM, and Transformer policies. The project also covers experimental rigor through multi-seed evaluation, reward ablations, and a transfer learning recipe that pretrains on an easier target before finetuning on the hardest non-Gaussian cat state.

The course materials are private to preserve the assignment for future offerings. If you would like to know more, feel free to reach out.