Measurement Based Feedback Quantum Control With Deep Reinforcement Learning for a Double-well Non-linear Potential

Sangkha Borah Bijita Sarma Michael Kewming Gerard J. Milburn Jason Twamley

Measurement Based Feedback Quantum Control With Deep Reinforcement Learning for Double-well Non-linear Potential: Supporting Information

Essentially there are three kinds of machine learning, namely supervised learning, unsupervised learning and reinforcement learning Goodfellow et al. (2016). In supervised learning, the ML model is provided with sufficient input and output data (labeled data), which are used to train the model to discover the hidden pattern/features in the datasets. For example, the dataset may consist of RGB images of different fruits, e.g., of apples, oranges, and bananas, with proper labels assigned by their names. The ML model can read the pixels of those images and correlate to the labels in the datasets by training on them. After training on enough datasets, the ML model is expected to acquire the necessary wisdom to predict from a new image (on which no training was done) whether it is an apple or not. On the other hand, in unsupervised learning, we do not provide labels to the datasets. After training, the model is able to make a classification of those fruits based on the hidden features in the pixel data of the images, and eventually acquire the expertise to classify new test images of fruits again, whether it is an apple, orange or banana or not.

Reinforcement learning (RL) Sutton and Barto (2018) is the third kind of ML, the basic workflow of which is shown in Fig. S1. It comprises of two main characters - the agent and the environment. The environment is the agent’s world in which it lives and interacts with. The agent is the brain of the RL and comprises of routines to function approximators, such as the artificial neural networks. In the case when the agent uses deep neural networks as function approximators (neural networks are non-linear function approximators), the RL is known as deep reinforcement learning (DRL). The RL-agent periodically interacts with the RL-environment in real time, in which it applies some changes to the dynamics of the system by changing some dependent variables. The changes it applies to the environment are called the actions. In return, the agent gets back full/partial information of the changes in the dynamics of the environment (known as the observation or the state of the system) along with a scalar value known as the reward signal. The reward function returns a scalar value that gives an estimation of the likelihood of the applied action bringing to desirable changes in the dynamics or not. The idea behind choosing a reward function is that rewarding or punishing an agent for its behaviour (actions applied) makes it more likely to repeat/forgo that behaviour in future. For that, the agent is set the goal to maximize the cumulative rewards (known as return) by changing its actions within a certain number of interactions in real time, known as an episode (a play in games). At the beginning of such a RL experiment, the agent applies the actions randomly, and later updates the parameters of the agents function approximators (the weights and biases of the neural networks in the case of DRL) based on which the next actions are chosen. In reality, the agent-environment loop has to be computed for several hundreds/thousands of episodes depending on the complexity of the problem of interest. For instance, for stochastic environments with noisy data, the learning process of the RL-agent is slower than compared to its deterministic counterpart.

To understand it better, we here try to demonstrate the working of DRL with a fundamental classical mechanics problem, shown in Fig. S1. The environment is an upside-down harmonic oscillator with a completely frictionless surface, such that a ball can not be held fixed on top of it without any external forces. The task of the agent would be to apply engineered amount of forces back and forth to keep it on top, namely in the blue region shown in the figure. The reward function is chosen such that it gets the maximum reward at $x=0$ and gradually smaller ones as we go away from the origin. This must be programmed together with the solver of the equations of motion in the RL environment. The agent, composed of a few layers of neural networks (hence known as deep reinforcement learning, DRL) can be trained for few thousand steps before it acquires superhuman or human-like skills to adapt the external forces to keep it stable within the blue region. A video demonstration of the same can be obtained in the supporting media or the GitHub link  Quantum Machines Unit (2021a).

The rule based on which the agent decides the action to be applied on the environment is known as its policy. In the case of DRL the rule roughly represents the neural network function approximators, as the rule is implicitly determined by some non-intuitive choices of the parameters (weights and biases) of the neural networks. A policy can be deterministic, in which case it is usually denoted by $a_{t}=\mu_{\theta}(s_{t})$ or stochastic, which is denoted by $a_{t}\sim\pi_{\theta}(\cdot|s_{t})$, where $a_{t}$ is the action chosen by the agent as a function of the state $s_{t}$ at time $t$. The dot ($\cdot$) in the case of stochastic policy ($a_{t}\sim\pi_{\theta}(\cdot|s_{t})$) represents the conditioning on $s_{t}$ (probabilistic), unlike the case of deterministic policy $a_{t}=\mu_{\theta}(s_{t})$, for which the action is a deterministic function that maps the state $s_{t}$ to the the policy $\mu$ at time $t$. A deterministic policy is usually chosen only for evaluating the performance of an agent after being trained.

The reward function $R$ is one of the most critical parts of RL, which is crucial for the agent’s learning. In its most general form, it is a function of the current state, the action applied and the next state of the environment, $r_{t}=R(s_{t},a_{t},s_{t+1}$, where $r_{t}$ is the scalar reward signal at time $t$. The cumulative reward is defined as,

where $\tau=(s_{0},a_{0},s_{1},s_{2}...s_{t},a_{t}..)$ represents the sequences of state and action over an episode known as a trajectory, and $T$ represents the horizon, giving the extent to which the summation is done. In this case, the return is known as finite-horizon un-discounted return. Often the choice is to use infinite horizon limit with a discount factor $\gamma\in(0,1)$, known as infinite-horizon discounted return:

which essentially signifies the fact that rewards now is better than rewards later. The choice of the discounted rewards also makes it mathematically convenient.

The task of the agent is to maximize the expected return $J(\pi)={\hat{\mathbb{E}}_{t}}\left[R(\tau)\right]$, where ${\hat{\mathbb{E}}_{t}}$ is the empirical average over a batch of samples, for which the policy is optimized for the optimal policy $\pi^{*}$: $\pi^{*}=\text{arg}\max_{\pi}J(\pi)$. The expected return for choosing a given state or state-action pair is known as the value function. The expected return for the starting in a state $s_{0}=s$ and following the policy afterwards gives the on-policy value function, given by,

On the other hand, the expected return for the starting in a state $s_{0}=s$ and taking an action $a_{0}=a$ and following the policy afterwards gives the on-policy action-value function, given by,

The value functions obey the Bellman equations, which can be solved self-consistently, for example, for action-value function, it is given by,

Sometimes, yet another function that combines both the types of values function is defined, given by,

This is called the advantage function and gives a measure of advantage of taking an specific action $a$ in state $s$ over the random selection of it following $\pi(\cdot|s)$.

In recent years, several different RL algorithms have been proposed, among which the leading contenders are the deep Q-learning Mnih et al. (2015), vanilla policy gradient methods Sutton and Barto (2018), and trust region natural policy gradient methods Schulman et al. (2015, 2017). In the modern scenario, a successful RL method is expected to be scalable to large models, parallel implementations, data efficiency, and robustness, and should be successful on a variety of problems without much of hyper-parameter tuning. While Q-learning methods are known for their efficiency in numerous discrete action-based environments, such as various Atari games Silver et al. (2016, 2017), they fail on many simple problems, and may suffer from the lack of guarantees for an accurate value function, the intractability problem resulting from uncertain (stochastic) state information as well as the complexity arising from continuous states and actions Schulman et al. (2017). Vanilla policy gradient methods Sutton and Barto (2018) are known for better convergence properties, are effective in high-dimensional or continuous action spaces, and can learn stochastic policies. However, they often converge to a local minimum rather than a global one, and have poor data efficiency and robustness, often leading to high variance in policy updates. The trust region policy optimization (TRPO) Schulman et al. (2015) is relatively complicated, and is not compatible with architectures that include noise such as dropout or parameter sharing between the policy and value function, or with auxiliary tasks. Proximal policy optimization (PPO) Schulman et al. (2017) is the latest inclusion in the list, which is shown to yield similar or better performance and robustness as TRPO, with much easier implementation and parallelization, and faster optimization of policy parameters with stochastic gradient methods, and is naturally compatible with noisy environments.

Policy gradient methods work by optimizing the following objective (loss) function:

where, $\pi_{\theta}$ is a stochastic policy, and $\hat{A}_{t}$ is an estimator of the advantage function at timestep t. $\hat{\mathbb{E}}_{t}$ is the empirical average over a finite batch of samples. In TRPO methods, the objective function (the “surrogate” objective) contains the probability ratio $r(t)$ of the new ($\pi_{\theta}$) and old policies ($\pi_{\theta_{old}}$) instead of $\log\pi_{\theta}$:

under the constraints that the average KL-divergence between new and old policies are less than a small cutoff, $\delta$:

Without the constraint, maximizing the objective would lead to an excessively large policy update Schulman et al. (2015). PPO modifies the above constraint by clipping the policy update, with a hyper-parameter $\epsilon\sim 0.2$, and optimizing the objective function with stochastic gradient method Schulman et al. (2017):

Q-learning, on the contrary, optimizes value function $Q_{\theta}(s,a)$ using the Bellman equation, as given in Eq. S11. Lately, different algorithms that combines policy gradient methods with Q-learning methods have been proposed and have been found to sidestep many of the drawbacks of the traditionally used approaches. These are known as the actor-critic methods, in which there are two policy networks, actor and critic. While the actor, based on a policy network, takes an action on the environment, the critic network receives the reward from the environment and makes useful suggestion to the actor by updating the parameters of the action value function. In this paper, we have used a DRL agent that combines PPO with advantage actor-critic (A2C) approach Mnih et al. (2016) , where the advantage function is is calculated as the difference between the Q-value for action $a_{t}$ at the state $s_{t}$ and the average value of the state: $\hat{A_{t}}(s_{t}|a_{t})=Q(s_{t}|a_{t})-V(s_{t})$, discussed above. The A2C framework allows synchronous training of multiple parallel worker environments simultaneously, which enables faster training. The basic workflow of A2C is shown in Fig. S2.

We have used a weak continuous measurement of $\hat{x}^{2}$ described using a master equation for the system density operator. In this section we derive this equation using a continuous limit for a sequence of weak measurements at random times.

A single weak measurement of a Hermitian operator $\hat{A}$ is described using an operation on the state of the system that gives the post-measurement state conditioned on a measurement result $z\in{\mathrm{R}}$

where $P(z)={\rm Tr}[\hat{\Upsilon}^{\dagger}(z)\hat{\Upsilon}(z)\rho]$ is the probability density of the measurement result and

with $[\hat{A},\hat{\Upsilon}(z)]=0$, and $g$ is the gain in inverse units to $\hat{A}$. The measurement result, $z$, is a classical random variable conditioned on the state of the measured system. The mean and variance of $z$ are given by

where $\Delta\hat{A}=\hat{A}-\langle\hat{A}\rangle$ with the quantum averages defined by $\langle\hat{A}\rangle={\rm tr}[\hat{A}\rho]$. It is clear that $\sigma$ can be interpreted as the measurement error over and above any intrinsic quantum uncertainty. The unconditional state of the system post measurement is then given by

We now move to a weak continuous measurement by assuming each of these single weak measurements occur at random, Poisson distributed times. At each time, the device that performs the measurement is discarded after the measurement result is recorded, and a new measurement device is supplied - ready for the subsequent measurement. This builds in the Markov condition for the irreversible measurement process. The master equation describing the unconditional dynamics of the entire system is given by Milburn (1987),

where $\gamma$ is the rate of measurements.

We now take the limit in which the rate of measurements satisfies $\gamma\delta t\gg 1$ where $\delta t$ is the time scale of the free dynamics of the system and the uncertainty of the measurements $\sigma\gg g^{2}$ such that the ratio $\gamma/\sigma$ is constant. This is the continuous weak measurement limit. The unconditional master equation then becomes

with $\Gamma=\gamma g^{2}/\sigma$. The effect of the double commutator term in this master equation is to destroy coherence in the diagonal basis of $\hat{A}$ and adds diffusive noise to incommensurate variables $\hat{B}$ where $[\hat{A},\hat{B}]\neq 0$. A good measurement requires that $\Gamma$ is large.

We define an observed classical stochastic process in the continuum limit as follows. Denote a complete record of the measurement results, indexed by the random times of the measurement times, as $z(t)$. Now define the observed process by the stochastic differential equation

where $dN(t)^{2}=dN(t)$ and $\overline{dN(t)}=\gamma dt$ while $z(t)$ is a Gaussian distribution. In the continuum limit this point process can be approximated by a diffusion process as

where the subscript $c$ indicates that this average must be computed from the state conditioned on the entire history of $y(t)$ to time $t$ and $dW(t)$ is the Wiener process. Note that for a good measurement in which $\Gamma$ is large, the signal-to-noise ratio also becomes large.

In order to evaluate the conditional averages in the observed process we need a continuous limit of the conditional state at each measurement. The result is the stochastic master equation Wiseman and Milburn (2009),

where

and

This equation is non linear; not surprising as for a conditional process the future must depend on the entire past history of measurement results.

Let us consider the case when we remove the DRL agent altogether and set the feedback Hamiltonian to be proportional to the measurement current, $dy(t)/dt$ (Eq. S25). The effect of feedback on $\rho_{c}$ can be introduced as,

where $\mathcal{K}$ is a super-operator such that $\mathcal{K}\rho_{c}=-i[F,\rho_{c}]$ for some Hermitian feedback action operator $F$. With Eq. (S26), we get the modified closed-loop stochastic master equation under feedback $F$,

Averaging over the noise we get the the Wiseman-Milburn unconditional master equation with Markovian feedback as

where $\rho=E(\rho_{f})$, the average over $\rho_{f}$. Looking at the unconditional master equation (S31), one can identify the feedback term, $-i[F,\hat{A}\rho+\rho\hat{A}]$, which modifies the base Hamiltonian evolution generated by $H$. The term $\Gamma{\mathcal{D}}[\hat{A}]\rho$, is the decoherence introduced by the continuous measurement of $\hat{A}$. However we now observe a new source of decoherence ${\mathcal{D}}[F]\rho$, introduced by the feedback itself Zhang et al. (2017). There is often this competition between the two sources of decoherence that decide an optimal value of $\Gamma$. For small values of $\Gamma$, the measurement current $dy/dt$, is very noisy and the feedback decoherence is high, while the continuous measurement decoherence is low. For large values of $\Gamma$, the feedback decoherence is low while the measurement decoherence is high. This competition carries over to the case studied in the main manuscript where the feedback action is determined by a DRL agent. We thus observe an optimal measurement rate, $\Gamma$ for the learning of the DRL agent, as shown in Figs. 2(a) and S7(a).

Using the model given in Eq. (2) of the paper, we created a DRL environment (shown in Fig.1(a) using the open-source platform OpenAI-Gym Brockman et al. (2016), in which QuTIP’s Johansson et al. (2012) stochastic master equation (SME) solver is used to compute the dynamics, see Eq. (6). The environment primarily includes two routines– the quantum dynamics/state of particle moving in a DW and modulation function to alter the quantum dynamics and a reward estimation function based on the observables (measurement results) in a step-wise manner for every interaction made by the agent. Our PPO agent Schulman et al. (2017) is constructed following the implementations of stable-baselines3  Raffin et al. (2019) in the A2C settings Mnih et al. (2016), where both the actor and critic are modelled with a set of fully connected feed-forward networks dimensions $512\times 256\times 128$, with the first layer as a shared network between the two. For modeling the neural networks we have used the open-source ML platform PyTorch Paszke et al. (2019). The input of the network is provided as the mean of the measurement currents, $I(t)$ from the last four time-steps (which helps learning faster rather than using the instantaneous current), along with the action returned by the network. In each episode of the training process the DRL interacts with the environment $1000$ times, in intervals of $\delta t=0.01$, and each time applies an action to the environment. The agent gathers trajectories as far out as the horizon limits (in our case 4000 steps, i.e., data from 4 episode for each environment) for each environment running in parallel (in our case 8 environments in total), then performs a stochastic gradient descent update of mini-batch size (in our case it is 100) on all the gathered trajectories for the specified number of epochs (in our case 10). We have found that a larger network along with a small learning rate ($10^{-5}$) helps the agent keep the learning process stable during longer simulations.

These simulations require a high truncation of the Hilbert-space dimension (we took $N_{\mathrm{trunc}}=60$), since the process of measurement often drives the state to high Fock numbers (see supporting media for demonstration or the GitHub link Quantum Machines Unit (2021b)). This is especially important for the DRL in the initial stages of learning when it is trying to gather information about the system by performing random actions. In our numerical experiments we chose $N_{\mathrm{trunc}}=60$, $\Gamma\leq 0.5$, and used 8 parallel vector environments to speed up the processing. After training within a few hundred such parallel environments the DRL learns to apply the actions in such a way that it limits the occupation to low Fock states, avoiding and errors with truncation limits. In addition, we observe that the efficiency of learning depends on the type of agent. Although traditional DQN methods based on discrete actions works well with simple conditional mean data, we find that they perform very poorly in the presence of noise. Actor-critic methods, especially the newer PPO agent (an evolution of A2C and TRPO), is better suited for dealing with a stochastic environment (stochastic data) than a purely value-based algorithm like DQN.

In the main manuscript we have chosen the agent controlled action to be $F(t)\propto(xp+px)$. This action augments the default DW quantum dynamics. It is natural to enquire if there are other, more effective actions that will improve the learning of the DRL process to achieve the goal of cooling the DW to its quantum ground state. Interestingly, we find that this form of action appears to optimally effective and choosing other forms of action results in weaker DRL efficiency. The action operator, $F(t)=A(t)f$, where $f=xp+px$, is the Hermitian squeezing operator. The reason why this action is so effective in learning and controlling the quantum dynamics of the particle in the double-well potential can be understood merely from classical perspectives.

Consider the classical streamlines of the feedback action Hamiltonian in the phase space of the double-well potential. The feedback, $f(x,p)$ generates a flow vector $(V_{x},V_{p})$ in classical phase space defined by, $V_{x}=\{x,f\}$ and $V_{p}=\{p,f\}$, where $\{,\}$ is the Poisson bracket. This flow is always tangential to the curves of $f=\text{constant}=f(x_{k},p_{k})$, where $(x_{k},p_{k})$ is a phase space point at time $t_{k}$ just before feedback, and moves the particle along curves of constant $f$. Fig. S3 shows the phase space streamlines for three different possible feedback action Hamiltonians. Feedback in the form of $f=x^{2}$ is likely to be very ineffective, as indicated by the streamlines in Fig. S3(a), as the flow is always orthogonal to the $x$-axis. Similarly, the streamlines for $f=p^{2}-x^{2}$ would push the state starting near the $p$ axis further away from the $x$ axis. On the other hand, for the case of $xp+px$, the streamlines push the particle in the direction of the x-axis, which is appropriate for approaching the double-well minima.

Numerically, we can demonstrate the best of the three controls by using these choices of the feedback with the Bayesian control with conditional data, discussed in details down below.

It is found that the choice of initial state for the time evolution via the stochastic master equation has a profound effect on how well the DRL agent can achieve its goal to bring the system to the ground state. The maximum (blue) and mean (orange) fidelity of the trained agents for different initial states are shown in Fig. S4. When the initial state, $\rho(0)$ is a thermal or coherent state, the DRL converges to an average fidelity of about 60%. However, if we use a small cat state or the ground state of the DW itself, the agent is able to achieve a mean trained fidelity of over 80% with noisy measurement data. We achieve similar high performance if we start with a thermal state projected on to even parity.

We benchmark the mean fidelity achieved as a function of the measurement efficiency, $\eta$ in the Fig. S5(a). We find that the DRL reveals robust control for $\eta\geq 50\%$. In Fig. S5(b), the learning of the DRL agent is compared with the one without environment damping (blue),in the presence of different collapse operators – (i) $\sqrt{\gamma}a$, and (ii) $\sqrt{\gamma}a^{\dagger}a$. Further improvement of the results could be achieved by implementing the suggestions mentioned at the end of the main manuscript.

Since the DRL is trained explicitly on the measurement data that varies with the double-well parameters (for example, the position of the well minima), measurement strength ($\Gamma$) and measurement time interval ($\delta t$), the generalizability of a trained model is poor across any significant changes of those parameters (see Fig. S6). However, this is not the case for the initial state $\rho(0)$. Although the choice of the type of $\rho(0)$ has a significant effect on the mean fidelity achieved by the agent, we found that a trained model with a given $\rho(0)$ (say coherent state) yields exactly the same behavior for others (say even-thermal state) and vice versa (see Fig. S4), which means that choosing a better prediction (or a similar basis with respect to the target state) would lead to better net fidelity, without the need to retrain the model for different initial states.

It is possible to choose fidelity as the reward function instead of measurement current discussed in the main text of the paper, shown in Fig. S7. However, fidelity is not directly accessible in real time in a laboratory experiment, and the following analysis is included for the sake of clarity and helpful comparison. When the DRL agent uses fidelity as a reward function we see that the learning behaviour, shown in (Fig. S7(a)), as a function of the measurement rate $\Gamma$, essentially mirrors the behaviour seen when the measurement current is used instead as a reward function (shown in Fig. 2(a) in the main text of the paper). In Fig. S7(b) the episodic mean reward $\tilde{F}(t)$ evolution in time during the training of the agent when $\Gamma=0.1$ [light(dark) blue includes(averages) noise]. In Fig. S7(c) the fidelity variation of a given random episode of the training agent is shown.

Here we benchmark our result against the state-based Bayesian feedback protocol (where the feedback is based on an estimate of the state) as proposed by Doherty et al., Doherty and Jacobs (1999); Stockton et al. (2004). In the context of the present work, the protocol is simplified to provide feedback of the form $\mathcal{F}(t)=-(\langle{x}_{c}^{2}\rangle(t)-3^{2})\times(xp+px)$ where $\langle x_{c}^{2}\rangle(t)$ denotes the conditional mean of the observable $x^{2}$. However, $\langle x_{c}^{2}\rangle(t)$ is not a quantity directly accessible in real experiments, and the above feedback condition should be replaced by the measurement current, $I(t)$. We find that when Bayesian feedback is driven by $I(t)$, it shows very little control over the dynamics. We performed a numerical simulation with 1000 copies of the system (ensemble), all evolving under feedback based on the mean of the measurement currents during each time step. The performances of these two approaches is compared in Table S1

We see from S1 that in order that the Bayesian control performs well only when it is provided the control with conditional mean data. When used with measurement currents, it, however, shows no control at all. Even when it is used with an ensemble of 1000 identical systems, it does not find any control beyond 50%. We found from our analysis that it is very important to provide the Bayesian controller with very accurate estimation of the mean currents that does not deviate significantly from the conditional means. The presence of a bad data point often throws the system into a non-recoverable state resulting in much lower fidelities. DRL, on the other hand, has a brain, and hence can recover from such bad data points.

It is also seen that when the Bayesian as well as the DRL controller is trained with conditional mean data, both yield comparable performances. It is interesting to observe the significant drop in achievable fidelity in presence of the environment damping $\sqrt{\gamma}a$, but not in presence of $\sqrt{\gamma}a^{\dagger}a$, environment dephasing. Thus, the drop of the performance of the DRL agent in presence of an environment damping $\sqrt{\gamma}a$ can be attributed to the reduction of the indistinguishably of the target state of the DW potential in the presence of $\sqrt{\gamma}a$ decoherence.

Finally, we found that the choice of the $xp+px$ feedback over $x^{2}$ and $p^{2}-x^{2}$, discussed above, can be demonstrated using Bayesian feedback with conditional mean data. We found a feedback of the form $x^{2}$ could never control the dynamics, while with $p^{2}-x^{2}$ achieves a mean fidelity of $\sim 60\%$. The mean fidelity of 85% shown shown in table S1 could only be obtained when used with the feedback $xp+px$.