Robust Optimization for Quantum Reinforcement Learning Control using Partial Observations

Consider the Hilbert space $\mathcal{H}=\mathbb{C}^{2^{N}}$. The $N$-qubit pure quantum state is defined as a complex-valued unit vector in $\mathcal{H}$, where $N$ is the number of qubits. Starting from an initial state $\ket{\psi_{i}}\in\mathcal{H}$, we apply the following control sequence

where $H_{0}$ and $H_{1}$ are noncommuting control Hamiltonians and $p$ is the control depth. For quantum state transfer problem, the objective is to choose the correct control actions $(\vec{\gamma},\vec{\beta})$ to steer the state from $\ket{\psi_{i}}$ to a target $\ket{\psi_{f}}$. The existing quantum RL control schemes use the following state fidelity

to define the reward at each intermediate step $t$, which requires a huge amount of quantum measurements or accurate simulation of the system dynamics. In contrast, the intermediate states will only be partially observed for the implementation of RL in this paper.

For QAOA, the control scheme is similar as shown in Fig. 1, except that the objective function may be slightly different.

Since the pure quantum state is represented by a complex-valued vector $\ket{\psi}$ in $\mathcal{H}$, the space of quantum states can be visualized as a generalized unit sphere called Bloch sphere. Consequently, the coordinates for each state can be calculated by projecting the state vector onto an orthogonal basis. Denote the single-qubit Pauli operators as

and the basis states of a single qubit as

It is clear that the Pauli operators span the space of one-qubit states, and thus the coordinates of a state vector $\vec{x}$ on the Bloch sphere can be derived by

with $|x_{1}|^{2}+|x_{2}|^{2}+|x_{3}|^{3}=1$. It should be noted that Eq. (5) is the mathematical formulation of quantum measurements, with the Pauli operators being the measurement operations and $\ket{\psi}$ being the state to be measured. $x_{1},x_{2},x_{3}$ are the corresponding measurement results which are obtained by averaging the measured outcomes on a large number of copies of $\ket{\psi}$. The quantum state is fully observable if all three observations in Eq. (5) can be made.

The coordinates of multi-qubit states can be defined in a similar way. Denote $\sigma_{0}=I$. The coordinates of a multi-qubit state are given by

except for $k_{1}=k_{2}=...=k_{N}=0$. Hence, the dimension of an $N$-qubit pure state vector is $4^{N}-1$. In this paper, partial observations are made on a subset of these generalized Pauli operators.

First, we introduce the quantum RL control scheme of this paper.

• •

  State. The quantum state is defined by the coordinate vector $\vec{x}$ on the Bloch sphere. Due to the noise and uncertain environments, the agents may only receive partial observations of the current state.

  To be more specific, according to the control model (1), the quantum state evolves as

  $$\ket{\psi_{t}}=\left\{\begin{array}[]{ll}e^{-i\gamma_{t}(H_{0}+\delta H_{0})}\ket{\psi_{t-1}}&\mbox{if}\ t\ \mbox{is\ odd},\\ e^{-i\beta_{t}(H_{1}+\delta H_{1})}\ket{\psi_{t-1}}&\mbox{if}\ t\ \mbox{is\ even}.\end{array}\right.$$

  (7)

  Due to the unknown perturbations $\delta H_{0}$ and $\delta H_{1}$, the quantum state at each time step $t$ is hidden to the agent unless a full observation is conducted.

• •

  Partial Observation. The quantum observation is made by applying quantum measurements to the states with respect to a subset of generalized Pauli operators, which will give the estimates $\vec{o}_{t}$ for only a part of the coordinates of the state. Alternatively, the full set of coordinates may be observed, with a fraction of the observations $\vec{o}_{t}$ being single-shot or extremely noisy.

• •

  Reward. As depicted in Fig. 2, the reward for optimizing the quantum control action is defined as the reduced distance between the partially observed coordinates of the current state and target state on the Bloch sphere, which is calculate by $r_{t}=\sigma(|\vec{o}_{t}-\vec{o}_{target}|-|\vec{o}_{t+1}-\vec{o}_{target}|),\sigma>0$.

• •

  Action. The action $a_{t}$ is chosen from a continuous interval, representing the values of the control durations (or angles) $\gamma_{j},\beta_{j}$. Similar to the classical RL, the actor generates the action $a_{t}$ through an actor-critic framework given the observations as input. After that, the action is used to parameterize the next unitary control.

The general framework for the quantum RL control is depicted in Fig. 3. At the time step $t$, the agent receives the partial observation $\vec{o}_{t+1}$ which allows us to calculate the reward $r_{t}$. Here we use a done flag to mark the end of each trajectory. If $t=2p$, the done flag $f_{t}$ is set to 1, otherwise it is set to 0. After taking the action $a_{t}$, the experience $(\vec{o}_{t},a_{t},r_{t},\vec{o}_{t+1},f_{t})$ is sent into the replay buffer which records the historical data. The parameters of the agent are optimized by sampling mini-batches from the replay buffer. More details on POMDP can be found in Appendix. A.

Any established RL-based frameworks can be used to train the agent as defined in Fig. 3. Here we adopt the TD3 deep neural networks from [33]. TD3 is based on the commonly-used actor-critic structure as depicted in Fig. 4, which involves a total of two actor networks and four critic networks. TD3 is said to be able to solve the problem of overestimating the Q-value. During the training process, each critic in the predict network is optimized to minimize the mean-square-error between the predicted $Q_{j\in\{1,2\}}$ and the estimated target $\hat{Q}$. The pseudo-code that implements the quantum robust control using TD3 is shown in Algorithm III.2.

Since any RL-based frameworks can be readily used to implement the proposed algorithm of this paper as an alternative, we do not intend to expand on the design of the deep neural network structures of TD3. Further details on the exact implementation of the neural networks can be found in Appendix. B and [33].

In particular, step 6 in Algorithm III.2 will be replaced by real quantum control in practical implementation on quantum hardware. In that case, Algorithm III.2, which is solely based on partial measurements, can still work with unknown perturbations and noises. However, if prior knowledge on the source of the perturbation and noise is available, the RL control model can be pretrained with classical simulations, which will provide significant cost reduction since experimenting on quantum hardware is very expensive at present.

For single-qubit robust control, the number of control steps is chosen to be 20. We adopt the experimental setup from [16], in which the initial and target states are given by

and the control Hamiltonians are given by

The perturbations $\delta\tilde{H}_{0}$ and $\delta\tilde{H}_{1}$ to $H_{0}$ and $H_{1}$ are randomly generated for each trajectory of training. That is, at the beginning of each trajectory of training, we randomly generate the diagonal elements of $\delta H_{0}$ and $\delta H_{1}$ from a uniform distribution $U(0,1)$. The real and imaginary parts of the elements of the upper triangle of $\delta H_{0}$ and $\delta H_{1}$ are sampled from $U(0,1)$ as well. The lower triangle is then generated as the complex conjugate of the upper triangle to guarantee that the perturbed Hamiltonian is Hermitian. $\delta$ is a positive parameter that controls the magnitude of perturbation. The action values are confined within $\mathcal{A}=[-5,5]$, and the initial values of actions for each trajectory is uniformly sampled from $\mathcal{A}=[-5,5]$. The mini-batch size is set as $100$ and the maximal length of history is set as $20$ for the simulations. The parameters of the agent are optimized using the Adam optimizer with a learning rate of $0.00001$ for actor and critics. Also, we set $\sigma=10$ in the reward function. We assume that the partial observations are made on the observables $\sigma_{x}$ and $\sigma_{z}$. Therefore, the observations obtained at time step $t$ are calculated by $\mathcal{O}=\{\bra{\psi_{t}}\sigma_{x}\ket{\psi_{t}},\;\{\bra{\psi_{t}}\sigma_{z}\ket{\psi_{t}}\}$.

The average fidelities between the initial and target states during the training are shown in Fig. 5(a). The model is trained for 20000 times, corresponding to 20000 complete trajectories. As $\delta$ increases, the average

Next, we perform additional tests to demonstrate the generalizability of our method. In Fig. 6(a), the agents are trained with $\delta_{train}=0,0.1,0.5,1$, respectively. Then the parameters of the agents are fixed, and the performances of the models are tested with different $\delta$ on the same initial and target states. An interesting observation is that models trained with higher level of noise ($\delta_{train}=0.5,1$) show better generalizability than models trained with little or no noise ($\delta_{train}=0,0.1$). For example, when $\delta_{train}=1$, the test fidelity can still reach above 0.9 with any $\delta_{test}$ that is larger than 0.3. In contrast, when $\delta_{train}=0.1$, the test fidelity quickly drops below 0.9 before $\delta_{test}=0.5$. This observation suggests that adding a certain amount of noise during training is beneficial for improving the robustness and generalizability of the model.

We have also tested the performance of the trained model on different initial and target states. In Fig. 6(b), the agent is trained with $\delta=0.1$, and $100000$ pairs of initial and target states are randomly generated for test by

with

The test experiments are still conducted with $\delta=0.1$. It is clearly that the trained model works well with unseen initial and target states, with more than 32500 pairs of states having achieved a final fidelity above 0.95.

Our method is then compared to other commonly-used algorithms in Table 1 under the same experimental setup with full observability. These algorithms consist of two gradient-based methods: stochastic gradient descent (SGD) [34], Krotov algorithm [35] and two commonly-used RL algorithms: deep Q-learning [4] and deep policy gradient [12, 19]. SGD is one of the simplest gradient-based optimization algorithms, in which the control is updated using the approximated gradient of the cost function based on random batches. In Krotov algorithm, the initial state is firstly propagated forward to obtain the evolved state. Then the evolved state is projected to the target state, defining a co-state that encapsulates the mismatch between the two. Finally, the co-state is propagated backward to the initial state, during which the controls are updated. As this process converges, the co-state becomes identical to the target state. For comparison, each algorithm has been run 100 times to calculate the average fidelity without assuming any noise or Hamiltonian uncertainty, with codes adopted from [19]. The achieved fidelity of our method ($\delta=0$) is higher than those of other algorithms including deep Q-Learning when noise is not taken into consideration. If noise is added to the control Hamiltonian of our method, e.g. $\delta=0.1$ or $\delta=0.5$, the performance of our method is still better than the conventional SGD and Krotov methods and similar to deep Q-learning and policy gradient method, which did not consider any noise effect. These numerical results demonstrate that our method can at least achieve the same level of performance in comparison with the state-of-the-art algorithms using only partial observations while being adaptive to noise.

It should be noted that we have $|\vec{x}|^{2}=1$, or that the control landscape is a unit sphere. As a result, the estimates of $\langle\sigma_{x}\rangle$ and $\langle\sigma_{z}\rangle$ can be used to jointly determine the value of $|\langle\sigma_{y}\rangle|$ which is unobserved. The RL model may have learned to infer this hidden relation with partial observations to promote the control precision.

Consider the commonly-used spin-chain example for multi-qubit control [24], in which the control Hamiltonians are defined as

and

where $J$ is the coupling strength. Then initial state is given by

Without loss of generality, we let $N=2$ and $J=0.5$ for the demonstration. In this case, the control Hamiltonians and initial state are written as

and

Here the target state is chosen as the ground state of

The observations at time step $t$ are made as follows

The numerical results in Fig. 7 show that the RL algorithm with 6 observations can achieve high-fidelity control against a low level of perturbations to the Hamiltonian. As the amplitude of the perturbations increases, the average fidelities will decrease. In particular, the fluctuations in the ultimate fidelities have become more severe with $\delta=0.2$, which can be seen from Fig. 7(b). This may be due to the fact that the number of randomly perturbed elements of the control Hamiltonians is growing exponentially with the number of qubits, which indicates an exponential growth in the uncertainties of the Hamiltonians. In addition, the number of partial observations for two qubits is only 6, while the dimension of the state space is $4^{2}-1=15$. In contrast, 3 observations are enough for the complete determination of a single-qubit state, and thus partial observation with 2 observables will make the single-qubit algorithm more tolerant of a large $\delta$.

We also test the trained agents with different noise levels, as shown in Fig. 8. Since the number of observations is too few to counteract the increased effects of noise in the two-qubit system, the agents suffer from an abrupt performance deterioration if $\delta_{test}$ is much larger than $\delta_{train}$. Similar to the single-qubit example, the trained model with the largest $\delta_{train}=0.05$ has the most stable performance during the test under a small variation of the noise amplitude.

As can be seen from Table 2, our method does not always outperform the state-of-the-art reinforcement learning algorithms. However, with partial observations, our method can at least achieve the same level of performance in comparison with the state-of-the-art algorithms. SGD is not studied in the two-qubit case since it has been proved to be an inefficient method for solving many-body problems [19].

Next, we train the agents with 4 different configurations of observations as follows

Here the noise level is set as $\delta=0.01$. As shown in Fig. 9, the performance of the algorithm continues to improve as the number of observations increases. However, the achieved fidelity with 8 observations is worse than that with 6 observations. This may be because the reward becomes overly large with too many observations. In particular, $\sigma$ is set as 3.5 in the reward function for all the experiments. Reducing $\sigma$ may improve the performance of the algorithm with more observations.

We consider a 6-qubit example which uses QAOA to solve the combinatorial optimization problem [36]. The target problem-based Hamiltonian is defined as

whose ground states are the solutions to a specific combinatorial optimization problem. The control Hamiltonian is given by

Then initial state is set as

The goal of QAOA is to minimize $\langle H_{0}\rangle$. It has been argued that even if the optimization process may not be able to reach the ground states, at least an approximate solution can be found using the quantum-classical algorithm. We let $N=6$, and randomly generate the weights $\{\omega_{jk}\}$ from a uniform distribution $U(0,1)$ for illustration.

The full set of observations are given by

which can be used for the complete determination of the expectation $\langle H_{0}\rangle$ with respect to $\ket{\psi_{t}}$. In the numerical experiment, only $60\%$ of the observations are made at each time step to calculate the reward, while the other observations are just generated as random noise from $U([-1,1])$. This configuration is proposed to demonstrate the model performance under the extreme condition that $40\%$ of the observations are pure noise, which is also an instance of partially observable process as aforementioned. The real observations are made as follows

Although the other noisy observations are not used to calculate the reward, they are still contained in $\vec{o}_{t}$ to train the agent. It should be noted that only in the last step, the reward is calculated using the full set of real observations. To be more precise, the rewards for the intermediate steps are calculated as

while the reward for the last step is calculated as

This definition of reward function is slightly different from that of the state transfer problem, as the target state is unknown in the current case. Here, the purpose of the reward function is to encourage the decrease of $\langle H_{0}\rangle$ on the partially observed subsystem.

The optimization performance is shown in Fig. 10, where we have used the following approximation ratio[24]

as the metric to evaluate the accuracy of the QAOA algorithm with partial observation. Here $E_{min}$ and $E_{max}$ are the lowest and highest values of $\langle H_{0}\rangle$. The control depth is set as $p=5$, and we let $\sigma=10$. The numerical result demonstrates that the algorithm with only $60\%$ accurate observations can still achieve a good-enough approximation ratio when $\delta\leq 0.05$. In particular, the approximation ratio can reach above $0.9$, which may yield an acceptable solution to the combinatorial optimization problem, as optimizing the approximation ratio beyond a desired value has been proven to be NP-hard with classical algorithms [37].

We have compared our method with Trotterized quantum annealing control protocol by increasing $p$ as $p=5,7,9,11,13,15$. Part of the numerical results with optimized coefficients are shown in Table 3, and the complete results are given in Supplementary Materials. The Trotterized quantum annealing algorithm is executed using PyQPanda without assuming any Hamiltonian uncertainties, and the total durations of control are initialized with the same value. Under the same control depth, the accuracy of the Trotterized quantum annealing algorithm is higher than our method. Considering that our method additionally deals with the robust control problem with Hamiltonian noise and only uses $60\%$ of full observations, such accuracy is acceptable for QAOA. It should also be noted that our strategy achieves a relatively high accuracy in the first few control steps in order to reduce the cumulative effect of Hamiltonian noise, which greatly reduces the total duration of control as compared to Trotterized quantum annealing. In practice, the control depth $p$ cannot be increased to infinity due to robustness considerations.