Fast State Stabilization using Deep Reinforcement Learning for Measurement-based Quantum Feedback Control

The policy function $\pi_{\theta}(a|\rho)$ maps the quantum state’s vectorized form to a set of control actions $a$. These actions correspond to the control signals applied to the control Hamiltonians $H_{1},H_{2},\dots,H_{M}$. For example, if there are two control Hamiltonians, the actions at each time step are given by $a:=\begin{bmatrix}u_{1},u_{2}\end{bmatrix}^{T}$, where $u_{1}$ and $u_{2}$ are the control amplitudes.

The DRL algorithm trains the policy to select actions that stabilize the quantum state. At each time step $t$, the agent uses the policy to choose an action $a_{t}$ based on the state $\rho_{t}$. The environment then updates the quantum state to $\rho_{t+1}$ according to the system dynamics and provides a reward $r_{t}$. This reward encourages the agent to take actions that reduces the difference between $\rho_{t+1}$ and the target state $\rho_{\text{target}}$.

It is important to note that the resulting state $\rho$ may not be a valid quantum state due to the inherent randomness in measurements and cumulative errors in solving the SME using classical computers [31]. For example, the state matrix may contain non-physical (negative) eigenvalues. To address this issue, a suitable approach is to check the eigenvalues of the estimated matrix at each step. When negative values are encountered, the state should be projected back onto a valid physical state. This can be achieved by finding the closest density matrix under the 2-norm [32]. This approximation ensures that a non-physical density matrix is transformed into the most probable positive semi-definite quantum state with a trace equal to $1$.

As mentioned in Section III, the interaction between the agent and the environment forms an episode $\tau=\{\rho_{0},r_{0},a_{0},\rho_{1},r_{1},a_{1},\dots,\rho_{m}\},$ where $m$ is the final step of the episode, $\tau\in\mathbb{D}=\{\tau_{n}\}_{n=1,2,\dots,N}$, with $N$ denoting the total number of possible sequences. Without confusion, we use $\tau$ to represent a single episode. The primary goal of the agent is to maximize the expected cumulative reward across all possible episodes. The construction of the reward function will be discussed in detail in the following section.

In this work, we propose a Partitioned Nonlinear Reward (PNR) function based on the distance $D_{\rho_{t}}$ (5). A lower $D_{\rho_{t}}$ indicates better alignment with the target state. The reward at each time step $r_{t}$ is defined as:

where $\overline{D}$ and $\underline{D}$ are the upper and lower bounds of the distance $D_{\rho_{t}}$, $\overline{R}$ and $\underline{R}$ are the upper and lower bounds of the reward, and $\mathfrak{e}$ and $\mathfrak{f}$ are parameters that regulate the slope of the reward curve.

In general, this reward function is motivated from the inverse proportional function (the relation between $r_{t}$ and $D_{\rho_{t}}$ when other parameters in (7) are fixed), which ensures that the reward value increases when the distance $D_{\rho_{t}}$ decreases. More specifically, the bounds $\overline{D}$ and $\underline{D}$ are designed to ensure $D_{\rho_{t}}\in[\underline{D},\overline{D}]$. The distance is mapped to a reward value constrained within the range $[\underline{R},\overline{R}]$. The two tunable parameters, $\mathfrak{f}$ and $\mathfrak{e}$ are used to adjust the steepness when $D_{\rho_{t}}$ is approaching $\underline{D}$ and $\overline{D}$, respectively.

The detailed effects of different $\mathfrak{e}$ and $\mathfrak{f}$ are plotted in Fig. 2. For the given bounds $\underline{D},\overline{D},\underline{R}$, and $\overline{R}$, when $\mathfrak{e}<\mathfrak{f}$, the reward curve near $\underline{D}$ is steeper than that near $\overline{D}$. This indicates that the rate of increase in the reward value is more pronounced as $D_{\rho_{t}}$ approaches $\underline{D}$. Conversely, when $\mathfrak{e}>\mathfrak{f}$, the trend is reversed.

In most existing work, the upper and lower bounds of $D_{\rho_{t}}$ are fixed at $\overline{D}=1$ and $\underline{D}=0$, respectively. The reward function in Eq. (7), however, offers the flexibility to adjust these bounds, as well as the corresponding reward values $\overline{R}$ and $\underline{R}$. This flexibility enables the division of the state space into multiple regions, which is particularly useful for addressing complex problems. Moreover, by adjusting $\overline{D}$ and $\underline{D}$, it becomes possible to assign positive reward values in certain regions of the state space and negative reward values (penalties) in others. This feature allows for more efficient and effective optimization, resulting in improved performance.

In this paper, we divide the state space into two regions using a partition parameter $d$. These two regions are denoted as Proximity Zone ($D_{\rho_{t}}<d$) and Exploration Zone ($D_{\rho_{t}}\geq d$). The same reward function in the formula (7) will be used in both of these two zones, and the corresponding bounds, $\underline{D},\overline{D}$ are determined by the partition parameter $d$. This division provides more flexibility to balance the reward and the penalty in the whole state space. One possible curve of the reward function in these two zones are shown in Fig.  3.

In the Proximity Zone ($D_{\rho_{t}}<d$), the state is considered close to success. Positive rewards encourage the agent to converge rapidly to the target state. The distance bounds in this region are $\overline{D}=\overline{D}_{P}=d$ and $\underline{D}=\underline{D}_{P}=0$, while the reward bounds are $\overline{R}_{P}>\underline{R}_{P}\geq 0$. As shown in Fig. 3(a), the reward function ensures that the slope of the curve increases as $D_{\rho_{t}}$ approaches zero, emphasizing rapid stabilization near the target state.

In the Exploration Zone ($D_{\rho_{t}}\geq d$), the system is considered far from the target state, and penalties are applied to guide exploration. The boundaries of this zone are $\overline{D}=\overline{D}_{E}=1$ and $\underline{D}=\underline{D}_{E}=d$, with the reward range satisfying $\underline{R}_{E}<\overline{R}_{E}\leq 0$. As illustrated in Fig. 3(b), the penalty decreases as $D_{\rho_{t}}$ approaches $d$, with the slope of the curve ensuring a progressively faster reduction. This mechanism encourages the state to transition into the Proximity Zone.

Remark 1: In the Exploration Zone, when the distance $D_{\rho_{t}}$ approaches $d$, we reduce penalties rather than introduce large positive rewards. This approach avoids potential reward hacking [33], where the agent might exploit increasing rewards by staying near $D_{\rho_{t}}\approx d$ without fully reaching the target. By applying small penalties at each step, the agent is encouraged to explore more broadly while progressively improving its policy.

Additionally, a small penalty proportional to the number of steps taken before stabilization is introduced to encourage efficiency. For example, the first step incurs a penalty of $-1\times 10^{-6}$, the second step $-2\times 10^{-6}$, and so forth. This step-based penalty discourages unnecessary delays and motivates the agent to achieve stabilization promptly.

The proposed reward function adjusts based on the distance $D_{\rho_{t}}$, with steeper slopes in both regions as the system approaches the target state. This design encourages exploration in the Exploration Zone while driving efficient convergence in the Proximity Zone. We show the superiority of our design by simulating different reward functions in Section VI.

With this reward function, we are now in the position to maximize the cumulative expected reward of the sequence, so for a complete sequences $\tau$, its cumulative reward can be expressed as

$A^{\theta}(\rho_{t},a_{t})=Q(\rho_{t},a_{t})-V^{\phi}(\rho_{t})$ is known as the advantage function in the field of RL, which is utilized to assess the desirability of taking a specific action $a_{t}$ at state $\rho_{t}$. $Q(\rho_{t},a_{t})=\sum_{t^{\prime}=t}^{m}\gamma^{t^{\prime}-t}r_{t^{\prime}}$ represents the action-value function, indicating the expected discounted reward for choosing action $a_{t}$ in state $\rho_{t}$, i.e., the cumulative sum of rewards until the end of the episode after executing this action. $r_{t^{\prime}}$ is the reward function in (7). The value of $\gamma$ lies between $0$ and $1$, determining the emphasis on long-term rewards (close to $1$) or short-term rewards (close to $0$). It effectively introduces a discounting mechanism for future rewards, thereby shaping the agent’s preference for future reward consideration when making decisions. $V^{\phi}(\rho_{t})$ is referred to as the state-value function (or baseline) and is modeled by a neural network with the same structure as the policy network but with different parameters $\phi$. It is primarily employed to approximate the discounted rewards from state $\rho_{t}$ to the end of an episode. Specifically, if the current state is $\rho_{t}$, and for all possible actions $a_{t}^{(1)},a_{t}^{(2)},\dots,a_{t}^{(\dots)}$, they correspond to discounted rewards $Q(\rho_{t},a_{t}^{(1)}),Q(\rho_{t},a_{t}^{(2)}),\dots,Q(\rho_{t},a_{t}^{(\dots)})$. As $V^{\phi}(\rho_{t})$ represents the expected value of the discounted rewards at $\rho_{t}$, we can use $Q(\rho_{t},a_{t}^{(1)}),Q(\rho_{t},a_{t}^{(2)}),\dots,Q(\rho_{t},a_{t}^{(\dots)})$ as features to approximate the value of $V^{\phi}(\rho_{t})$, representing the expected value of rewards in state $\rho_{t}$. When $A^{\theta}(\rho_{t},a_{t})>0$, action $a_{t}$ is considered better than average and is worth increasing the probability of being chosen in subsequent iterations while decreasing the probability otherwise.

Our DRL algorithm is implemented using the open-source Python library Stable-Baselines3 [34], while the quantum dynamic environment is constructed within the Gymnasium framework [35]. All simulations in this study are conducted on a computer equipped with an Apple M1 Pro chip and $32$ GB of memory, utilizing Python $3.11.3$, stable-baselines3 $2.3.2$, and Gymnasium $0.29.1$. We design a reasonable reward function to guide the DRL agent through iterative learning, aiming to train an excellent DRL agent capable of generating control signals to achieve the stability of the target entangled state.

The upper and lower bounds of the reward in the Proximity Zone, $\overline{R}_{P}$ and $\underline{R}_{P}$ , are set to $100$ and $1$, respectively, encouraging the system to get as close as possible to the target. In the Exploration Zone, the upper and lower bounds of $\overline{R}_{E}$ and $\underline{R}_{E}$ are set to $0$ and $-0.1$, respectively, maintaining a penalty without overly punishing the agent.

We set $\mathfrak{e}=2$ and $\mathfrak{f}=10$, emphasizing a steeper reward escalation as $D_{\rho_{t}}$ approaches $\underline{D}$.

The pseudo-code for the PPO in quantum state stabilization is shown in Algorithm 1.

which implies that two control channels are applied to this quantum system. And the forms of $H_{1}$ and $H_{2}$ represent $\sigma_{y}$ rotations on the first and second qubits, respectively, enabling independent control over each qubit.

The observable operator is chosen as:

which corresponds to the measurement of $\sigma_{z}$-like observable for the two qubits.

Specify the target state as

which is a symmetric two-qubit Bell state.

We utilize the previously summarized PPO algorithm to train the DRL agent. For the training trajectories, we set a time interval of $\Delta t=0.001$ a.u. for each measurement step, with a maximum evolution time of $T=20$ a.u., corresponding to a maximum of $20,000$ steps. At each step, the DRL agent interacts with the environment, obtaining system information to generate control signals, which are then stored for iterative updates of the policy. The total number of training steps is set to $10^{7}$. On the computer with configuration in Section IV-C, the training process requires approximately $50$ minutes. However, the primary focus of this study is the performance of the trained agent rather than the optimization of training time. In practical implementations, employing a pre-trained agent is feasible. Once the target state is specified, the agent does not require repeated updates and can be directly utilized after training is completed.

In order to evaluate its performance, we test the proposed strategy on $50$ randomly selected distinct initial quantum states $\rho_{0}$ (corresponding to different initial distances $D_{\rho_{0}}$, as indicated by the blue line in Fig. 4). The light blue lines represent the evolution trajectory of a specific initial state with respect to the target state, averaged over $50$ different trajectories under varying environmental noise, while the dark blue line depicts the average evolutionary trajectory of all the different initial states, i.e., the average trajectory of $2500$ different evolutionary trajectories. Smaller values of the distance $D_{\rho_{t}}$ indicate that the system is closer to the target state. It is worth noting that a trained agent can successfully stabilize any initial state to the target state with high fidelity. Furthermore, for comparison with the Lyapunov method mentioned in [36], we retain the same $50$ sets of randomly selected initial states and obtain the orange trajectories in Fig. 4 using the Lyapunov method. It can be observed that the control signals generated by the DRL agent outperform the Lyapunov method. Assuming that the time when the distance $D_{\rho_{t}}$ is less than $0.001$ is the evolution time of the system, the average evolution time of the $2500$ trajectories under the guidance of the DRL agent is $4.59$ a.u., while the average time using the Lyapunov method is $5.86$ a.u.. The DRL’s average stabilization time is improved by $22\%$ over the Lyapunov method. This indicates that our DRL approach successfully stabilizes these quantum states to the target state faster than the Lyapunov method.

Due to the degenerate form of the observable $c$ in (14), the system may converge to other state in the corresponding eigenspace related to $\lambda_{d}$, two control channels $H_{1}$ and $H_{2}$ based on Lyapunov method have been applied in [36, 8, 40] to solve this problem. For subsequent performance comparisons, two control channels are also used in this paper.

For any training trajectory, we take a time interval $\Delta t=0.001$ a.u. for each measurement step. Given that we have set the maximum evolution time as $T=40$ a.u., it means the maximum number of evolution steps for any trajectory during training is $40,000$. The total number of training steps is $10^{8}$.

For all instances, in order to compare with the Lyapunov methods presented in [8], we choose the same system Hamiltonian as $H_{0}=\text{diag}[1,-1,-1,1,1,-1,-1,1]$. The diagonal form of which indicates that the eigenvalues correspond to the energy levels of the system in the computational basis $|q_{1}q_{2}q_{3}\rangle$ (where $q_{i}\in\{0,1\}$ for three qubits). The target state is $\rho_{d}\rm Triangleq\rho_{\rm{GHZ}}$ (13), and the observable $c$ is chosen to be:

to measure correlations between the z-components of different pairs of qubits in a three-qubit system and also ensure that the target state is an eigenstate. The control Hamiltonians $H_{1}$ and $H_{2}$ are chosen as

and

Similar to the two-qubit systems, two control channels represented by $H_{1}$ and $H_{2}$ together with their strengths $u_{1}$ and $u_{2}$ provide control over the system’s dynamics. The forms of control Hamiltonian represent how the control action is applied to the three-qubit system. For example, the form of $H_{1}$ generally represents an independent $x$-axis control to the third qubit and a correlated $x$-axis interaction between the first and the second qubits.

We test the trained DRL agent in various environments to evaluate its performance and robustness. The goal is to assess how well the agent generalizes its learned policies to different scenarios and how it copes with perturbations and variations in the environment. To achieve this, we expose the trained DRL agent to a set of diverse environments, each with unique characteristics and challenges. These environments are carefully designed to represent a wide range of scenarios and potential disturbances that the agent might encounter in real-world applications. During the testing phase, we measure the agent’s performance in terms of its ability to achieve the desired objectives and maintain stability in each environment. In particular, we examine its response to changes in the measurement efficiency $\eta_{c}$ and time delay disturbances to assess its robustness and adaptability.

We first investigate the “perfect case”. In this paper, the “perfect case” entails assuming that negligible delay in solving the SME (1) by classical computers, and perfect detection, that is, measurement efficiency $\eta_{c}=1$. In contrast, situations where there is delay or imperfect detection within the system are collectively referred to as “imperfect cases”. We then show some performance indications for “imperfect cases”.

When the parameter $\theta$ is given, the expected value of the total reward, denoted as $J(\theta)$, is evaluated as the weighted sum of each sampled $\tau$ sequence, expressed by

To maximize $J(\theta)$, which indicates that our chosen policy parameters $\theta$ can lead to higher average rewards, we adopt the well-known gradient descent method. Thus, we take the derivative of the expected reward $J(\theta)$ in (20), resulting in the expression shown in

We use $\nabla f(x)=f(x)\nabla\log f(x)$ to derive the second row of (21). The last approximate equation is a result of practical gradient computations, where instead of calculating the expected reward for an entire trajectory, rewards contributed by each individual state-action pair ($\rho,a$) are computed separately. These individual rewards are then summed up to obtain the total cumulative reward for the optimization process. The direction of the policy update $\pi_{\theta}(a|\rho)$ is biased towards favoring state-action pairs that contribute to higher cumulative rewards within the sequence. For instance, if an action $a$ executed in state $\rho$ leads to a positive cumulative discounted reward, the subsequent update will enhance the probability of choosing action $a$ in state $\rho$, while diminishing the likelihood of selecting other actions. The update equation for the parameters $\theta$ is as follows:

where $\eta$ is the learning rate.

Once the policy $\pi_{\theta}(a|\rho)$ is updated, it necessitates the reacquisition of training data prior to the subsequent policy update. This arises due to the alteration in the probability distribution $p_{\theta}(\tau)$ brought about by the modified policy. Following data sampling, the parameter $\theta$ undergoes refinement, leading to the discarding of all prior data. Subsequent parameter updates mandate the collection of fresh data, constituting the fundamental principle underlying the conventional PG algorithm. However, in the context of quantum systems, the process of sampling system information is often characterized by time-intensive and computationally demanding operations. For instance, after each measurement, a classical computer is requisitioned to solve the SME (1) to ascertain the system’s state in the subsequent moment. This inability to reutilize previously acquired data contributes to a protracted training process. To address this challenge, an additional strategy $\pi_{\theta^{\prime}}$ is introduced, mirroring the architecture of $\pi_{\theta}(a|\rho)$. Instead of directly engaging with the environment for data gathering, the primary agent $\pi_{\theta}$ employs the auxiliary agent $\pi_{\theta^{\prime}}$ to interact with the environment and accumulate data. The objective is to subsequently utilize this data to train $\pi_{\theta}$ multiple times, effectively reducing the computational and resource demands for data collection. Ensuring the consistency of data sampled by $\pi_{\theta^{\prime}}$ with that of $\pi_{\theta}$, importance sampling [42] is introduced to facilitate this synchronization process. This approach contributes to enhancing data reuse and the overall efficiency of the training procedure. (21) is updated as:

Here, all the state-action pairs $(\rho_{t},a_{t})$ (or alternatively, all trajectories $\tau\in\mathbb{D}$) are sampled from $\pi_{\theta^{\prime}}$. The term ${p_{\theta}(\rho_{t},a_{t})}/{p_{\theta^{\prime}}(\rho_{t},a_{t})}$ represents the importance weight, which dynamically adjusts the weighting of data sampled by $\pi_{\theta^{\prime}}$ in real time to more accurately estimate the expected value under the target policy $\pi_{\theta}$.

The corresponding objective function from (23) can be calculated as:

Nonetheless, in the absence of constraint, such as when $A^{\theta^{\prime}}(\rho_{t},a_{t})>0$, indicating the desirability of specific action-state combinations, the agent’s inclination would be to elevate their likelihood, effectively amplifying the ${p_{\theta}(a_{t}|\rho_{t})}/{p_{\theta^{\prime}}(a_{t}|\rho_{t})}$ value. This scenario can lead to policy learning inaccuracies and an erratic learning process, impeding convergence. To counteract this, the PPO introduces a pivotal mechanism, termed the “clip ratio”. This clip ratio imposition serves to confine the proportions between the new and preceding policies, thereby ensuring congruence and augmenting the algorithm’s dependability. The following equation demonstrates the PPO Clipping algorithm, incorporating the clipping term to bound the difference between $p_{\theta}$ and $p_{\theta^{\prime}}$ during the policy update.

where $\varrho={p_{\theta}(a_{t}|\rho_{t})}/{p_{\theta^{\prime}}(a_{t}|\rho_{t})}$. The last two terms $1-\varsigma$, and $1+\varsigma$ in the clip function limit the boundaries of the first term. $\varsigma$ is a hyperparameter, typically set to $0.1$ or $0.2$. Exhaustively considering all possible sequences is typically infeasible, and thus in practical training, the objective function (25) is often formulated in the following manner:

where $N$ and $m$ represent finite real numbers that respectively signify the count of collected sequences and the maximum number of steps within each sequence.

Based on the above, we can see that PPO has three sets of network parameters for its update strategy:

• •

  One set of main policy parameters $\theta$, which is updated every time.

• •

  One set of policy parameter copies $\theta^{\prime}$, which interact with the environment and collect data. They utilize importance sampling to assist in updating the main policy parameters $\theta$. Typically, $\theta^{\prime}$ is updated only after several updates of $\theta$ has been performed.

• •

  One set of value network parameters $\phi$, which are updated based on the collected data using supervised learning to update the evaluation of states. They are also updated every time.