Accelerated Magnonic Motional Cooling with Deep Reinforcement Learning

Bijita Sarma Sangkha Borah A Kani Jason Twamley

Supporting Information

The workflow of a typical RL is shown in Fig. S1. Here, the block named as the RL-agent is depicted as comprising a set of neural network layers which act as function approximators. This type of RL is known as deep reinforcement learning (or DRL for short). On the other hand, the block in the right depicts the RL-environment, which is the world wherein the RL-agent lives and can affect change via actions on the dynamics of the environment. In RL terminology, the environment signifies the real/virtual space in which the rules of the physics of the problem/system are encoded along with a reward estimation function returning some real scalar values that determine whether the applied action has led to desirable changes or not. In practice, the RL-agent makes a certain number of interactions by first randomly applying some actions on the RL-environment based on which the weights of the nonlinear neural network function estimators are updated. The procedure is repeated several hundreds/thousands of times, depending on the complexity of the problem. Such a learning process is analogous to the learning process of a human, where the agent learns by conducting trial and error experiments on the RL-environment. Depending on the type of the RL algorithm used, the neural network based agent can explore even newer possibilities not usually explored, thus outperforming the analogous supervised learning agent.

The algorithm that the RL-agent uses to determine the actions is called its policy, which in the case of DRL, represents the neural network itself that takes the observations as the input, and outputs the actions. A policy can be deterministic in which case the action of the RL-agent is determined by the policy parameters $\theta$ for a given state $s_{t}$ at time $t$: $a_{t}=\mu_{\theta}(s_{t})$, or stochastic in which the actions are sampled from a probability distribution conditioned on $s_{t}$: $a_{t}\sim\pi_{\theta}(\cdot|s_{t})$. The task of the RL-agent is to optimize the policy parameters ($\theta$) to maximize the net discounted rewards $R(\tau)=\sum_{t=0}^{\infty}\gamma^{t}r_{t}$ over a trajectory $\tau=(s_{0},a_{0},s_{1},a_{2},...)$, for the discount factor $\gamma\in(0,1)$. To achieve that, the expected return over the discounted rewards $J(\pi)=\mathbb{E}[R(\tau)]$ is optimized to obtain the optimal policy $\pi^{\star}=\mathrm{argmax}\,J(\pi)$.

As a fundamental concept in RL, the value function gives a prediction of the expected, cumulative, discounted, future reward, and provides a measure of how well a given state $s$ or state-action pair $(s,a)$ behaves to generate a higher net return. The state value $V_{\pi}(s)=\mathbb{E}[R_{t}|s_{t}=s]$ is defined as the expected return for following policy $\pi$ from state $s$. On the other hand, the action value $Q_{\pi}(s,a)=\mathbb{E}[R_{t}|s_{t}=s,a_{t}=a]$ is the expected return for selecting action $a$ in state $s$ and then following policy $\pi$. The value functions obey the so-called Bellman equations Sutton and Barto (2018), which can be solved self-consistently. For example, the action-value function obeys,

There are several approaches to optimize the policy, that can be broadly categorized into three types, namely (a) policy-gradient-based, (b) value-based and (c) actor-critic based methods. In value-based approaches, such as the Q-learning methods, the policy is optimized to get the net maximum value functions solving the Bellman equations as discussed above. On the other hand, in policy gradient methods the policy parameters are optimized by using gradient descent algorithms:

where $\mathbb{E}$ represents the expectation value over the trajectory $\tau$. This basic approach can be improved by using a baseline function, $b(s_{t})$, to reduce the variance of gradient estimation and forms the basis of the more advanced state-of-the-art DRL actor-critic algorithms. In the most generalized form, the policy gradient methods work by optimizing the following objective (loss) function:

where, $\pi_{\theta}$ is a stochastic policy, and $\hat{A}_{t}=Q(s_{t},a_{t})-V(s_{t})$ is an estimator of the advantage function at timestep t, since $R_{t}$ is an estimate of $Q(a_{t},s_{t})$.

An actor-critic algorithm learns both a policy and a state-value function, and the value function is used for bootstrapping, i.e., updating a state from subsequent estimates, to reduce variance and accelerate learning Sutton and Barto (2018), while the critic updates action-value function parameters, and the actor updates policy parameters, in the direction suggested by the critic.

The Soft Actor-Critic (SAC) algorithm is a recently proposed actor-critic algorithm Haarnoja et al. (2018) in the field of reinforcement learning. The main difference of SAC compared to other actor-critic methods is that its policy is optimized in an entropy-regularized manner and is inherently stochastic. The policy is trained to maximize a tradeoff between expected return and entropy, with an increase in entropy leading to more exploration and, in addition, preventing the policy from converging prematurely. The following discussion closely follows the theory presented in Achiam (2021).

In entropy-regularized RL, the agent gets a bonus reward at each time step proportional to the entropy of the policy,

where $H(P)=\underset{x\sim P}{\mathbb{E}}[{-\log P(x)}]$ denotes the entropy computed from the probability distribution $P$, and $\alpha>0$ is the trade-off coefficient. The corresponding value functions in this setting, $V^{\pi}$ and $Q^{\pi}$ are changed to,

With these definitions, $V^{\pi}$ and $Q^{\pi}$ are connected by,

The Bellman equation for $Q^{\pi}$ becomes,

where, the expectation over the next states, $r$ and the states, $s^{\prime}$ come from the replay buffer, while the next actions $\tilde{a}^{\prime}$ are sampled fresh from the policy. SAC concurrently learns a policy $\pi_{\theta}$ and two Q-functions $Q_{\phi_{1}},Q_{\phi_{2}}$ and set up the loss functions for each Q-function, and takes the minimum Q-value between the two Q approximators,

where the target is given by

For the studies reported in this article, we design the RL-environments using the framework template of OpenAI-Gym Brockman et al. (2016). The RL-agents are modelled using SAC policy, with two critic models following the implementation available in the open-source software package Stable-Baselines3 Raffin et al. (2019). For modelling the neural networks and optimization of the network parameters we have used the PyTorch Paszke et al. (2019) ML module.

The effective two-mode magno-mechanical model shown in Fig. 1(b) in the main text consists of a YIG sphere put in a microwave (MW) cavity, where a uniform magnetic field, $B_{0}$ is applied in the $z$-direction that induces spin wave magnons precessing in the $xy$-plane. The magnetic field component of the cavity mode couples to the magnons through magnetic dipole coupling Tabuchi et al. (2014); Zhang et al. (2014). The YIG sphere also acts as an excellent mechanical resonator because of the magnetostrictive effect, that gives rise to vibrational acoustic phonon modes Zhang et al. (2016); Li et al. (2018). We consider the magnon modes to be driven by a MW field with frequency $\omega_{p}$ and amplitude $\epsilon_{p}$ Wang et al. (2018); Li et al. (2018). The Hamiltonian of the system in a frame rotating with the magnon drive frequency is given by,

where $a(a^{\dagger})$, $m(m^{\dagger})$ and $b(b^{\dagger})$ are the annihilation (creation) operators for the cavity, magnon and phonon modes respectively, $\delta_{i}=\omega_{i}-\omega_{p}$ are the detunings, $J$ is the magnon-photon coupling strength and $G$ is the magno-mechanical coupling rate. The Rabi frequency $\epsilon_{p}$ is given by $\sqrt{5N}\gamma B_{0}/4$, where $\gamma/2\pi=28\ \rm{GHz/T}$ is the gyromagnetic ratio, $N=\rho_{n}V$ is the total number of spins, where $\rho_{n}=4.22\times 10^{27}\ \rm{m}^{-3}$ is the spin density of YIG and $V$ is the volume of the sphere.

The quantum Langevin equations for the average dynamics of the fluctuation operators are given by

where $\tilde{G}$ is the driving-enhanced magno-mechanical coupling. When the cavity has a low Q-factor, with a damping rate of the order of $\kappa_{a}\gg\{J,\ \tilde{G},\ \kappa_{m},\kappa_{b}\}$, the cavity field can be adiabatically eliminated, that gives rise to an effective magno-mechanical Hamiltonian of the form,

where the effective parameters are given by, $\Delta_{m}=\delta_{m}-\frac{|J|^{2}}{\delta_{a}^{2}+\kappa_{a}^{2}}\delta_{a}$, $\tilde{\kappa}_{m}=\kappa_{m}+\frac{|J|^{2}}{\delta_{a}^{2}+\kappa_{a}^{2}}\kappa_{a}$. Hence, when the sphere is put at a node of the cavity magnetic field, or the cavity damping is very high, effectively the system reduces to a two-mode coupled system with a magno-mechanical interaction.

The magnon and mechanical modes have largely different frequencies, with $\omega_{m}\sim 2\pi\times 10\ {\rm{GHz}}$ and $\omega_{b}\sim 2\pi\times 10\ {\rm{MHz}}$, therefore the magnon mode behaves as a ‘cold’ sink at zero temperature for the ‘hot’ mechanical mode. At the red sideband, $\Delta_{m}=\omega_{b}$, by virtue of the resonant interaction in the weak coupling regime, between the mechanical mode and magnon anti-Stokes sideband, it gives rise to transfer of thermal quanta from the mechanical mode to the magnon mode, which is the working principle of the well-known sideband cooling technique. In Fig. S3 we show the sideband cooling dynamics with several values of $\tilde{G}$, in the regime $\tilde{G}<\omega_{b}$. The time limit for such irreversible cooling is denoted by $\tau_{\rm{SB}}$. From the figure, one can get an estimate of the steady-state cooling time limit and the optimum cooling quotient, $\tilde{n}_{b}$. Considering an optimum cooling quotient of, $\tilde{n}_{b}\approx 10^{-4}$, the best steady-state cooling from such sideband cooling mechanism is found to occur at a time limit of the order of $\tau_{\rm{SB}}\omega_{b}/(2\pi)\approx 40$. As can be seen from the figure, such sideband cooling is not possible with higher values of the magno-mechanical coupling, $\tilde{G}$ because of the effect of the counter-rotating terms beyond RWA. In our DRL-scheme, we show that cooling can be obtained even in the strong coupling regime with proper choice of coupling modulations. And more interestingly, it leads to lower time limit for cooling.

For the two-mode magno-mechanical system, the control parameter is the coupling strength $\tilde{G}(t)$, which can be complex in nature (see Eq. 1 of the main text of the article). In the DRL model, the real and imaginary parts of $\tilde{G}$ are considered as the two actions that are learned by the DRL-agent through trial and error. In a given iteration of the RL workflow at time $t$, the RL-agent chooses two actions based on which the system (RL-environment) makes the dynamics to $t=t+\delta t$, where $\delta t=0.1\omega_{b}/(2\pi)$ during which $\tilde{G}(t)$ remains fixed. The observations that the RL-agent gets from the RL-environment comprise the instantaneous values of the moments obtained by solving the second order coupled differential equations. The actions of the RL-agent are also added to the list of observations, which we found to be advantageous for learning. The SAC controller is made based on the following settings of hyperparameters,

The goal of the problem is to optimize the neural network parameters such that it learns the non-trivial combinations of the controls, $\tilde{G}$, so that the phonon occupancy of the YIG sphere is reduced thereby cooling it significantly. For that the agent is given the following reward function,

where $n_{T}$ is the thermal phonon number, $b^{\dagger}b$ is the number operator of the mechanical mode, and $\langle A\rangle$ denotes the expectation value of the operator $A$. The task of the RL-agent would be to maximize the net reward, $R=\sum_{i}^{{\mathcal{N}}_{t}}r_{i}$ over an episode. Note that this, in theory, represents a highly challenging task as for each control (action), the number of possible combinations in a given episode is $\mathcal{N}_{C}^{\mathcal{N}_{t}}$, where $\mathcal{N}_{t}$ is the number of total time steps allowed within an episode, and ${\mathcal{N}_{C}}$ is the number of possible discretization of the control strength bounded by $\tilde{G}_{\rm max}$. Since in our case, the actions are continuous the problem is even harder.

The RL-agent is trained for the following set of cases: $\tilde{G}_{\rm max}/(\sqrt{2}\omega_{b})=\left\{0.1,0.2,0.3,0.5,1,2,3,4,5\right\}$. The learning curve for $\tilde{G}_{\rm max}/(\sqrt{2}\omega_{b})=5$ is shown in Fig. S4. In these plots, we have shown how the RL-agent gradually learns the non-intuitive control sequences by optimizing the neural network weights to achieve the goal of attaining larger net returns determined by Eq. S17. We have also shown the performance of the RL-agent at different stages to demonstrate the learning process, as explained in the caption of Fig. S4 in more detail.

For the tripartite system model, we consider the cooling of a levitated spherical YIG ferrimagnet with high spin density, trapped in a harmonic potential wherein spin-wave magnon modes are excited by applying an external bias magnetic field. In a large external homogeneous magnetic field, $B_{0}$, applied in the $\hat{z}$-direction, the YIG sphere is magnetized to its saturation magnetization, $M_{s}=5.87\times 10^{5}\mathrm{Am}^{-1}$, whereas the small magnetization fluctuations, $\bm{m}(\bm{r},t)$ around the fully magnetized state is given by, $\bm{M}(\bm{r},t)=M_{S}\hat{z}+\bm{m}(\bm{r},t)$. For the homogeneously magnetized fundamental magnon mode (Kittel mode), precessing around the $\hat{z}$-axis with an eigen-frequency, $\omega_{m}=|\gamma|B_{0}$, the magnetization fluctuation is given by $\bm{m}(\bm{r},t)=\mathcal{M}_{K}(\hat{x}+i\hat{y})m+h.c.$ Fletcher and Bell (1959), which induces a magnetic field, ${\bm{b}}(\bm{r},t)=\mu_{0}\frac{\mathcal{M}_{K}}{3}\frac{R^{3}}{r^{3}}e^{i\phi}\left(2\sin\theta\ \hat{r}-\cos\theta\ \hat{\theta}-i\hat{\phi}\right)m+\textit{h.c.},$ outside the sphere, where, $\mathcal{M}_{K}=\sqrt{\frac{\hbar\gamma M_{s}}{2V}}$ is the zero-point magnetization, $\mu_{0}$ is the vacuum permeability, $m$ is the bosonic magnon operator, $R$ and $V$ are the radius and the volume of the sphere. As schematically shown in the main text, a driven WGM optical microsphere with a magnetostrictive rod (MR) attached to it, is placed near the YIG magnet. The magnon mode field modulates the axial length of the MR, which modulates the WGM optical mode frequency that is depicted by a coupling of the form $\Omega_{S}a^{\dagger}a(m+m^{\dagger})$ between the WGM optical mode $(a)$ and the magnon bosonic mode $(m)$, where $\Omega_{S}=\Delta\omega$ is the optical frequency shift i.e. the single photon magnon-cavity coupling.

The magnon mode can also be coupled to the COM motion of the YIG sphere, by applying a spatially inhomogeneous external time-dependent magnetic field, $\mathbf{H}_{g}(y,t)$ Hoang et al. (2016); Delord et al. (2018), which satisfies the weak driving, $\left|\mathbf{H}_{g}(y,t)\right|\ll H_{0}$, and small-curl, $\left|\nabla\times\mathbf{H}_{g}(y,t)\|\mathbf{r}\right|\ll\left|\mathbf{H}_{g}(y,t)\right|$ conditions. The COM motion of the magnet with frequencies $\omega_{x},\omega_{y}$, and $\omega_{z}$, in $\hat{x},\hat{y}$ and $\hat{z}$ directions, is depicted by the Hamiltonian, ${H}_{b}=\sum_{j=x,y,z}\omega_{j}{b}_{j}^{\dagger}{b}_{j}$, where the bosonic ladder operators describe annihilation and creation of a motional quantum along the direction $j$. The quantum Hamiltonian describing the interaction between the COM motion, and the spin-wave magnetization due to the gradient magnetic field is given by, $\mathcal{H}_{mb}(t)=-\frac{\mu_{0}}{2}\int dV\left[M_{S}\hat{z}+{\mathbf{m}}(\mathbf{r})\right]\cdot\mathbf{H}_{g}\left(\mathbf{y},t\right)$. Considering a time-varying gradient magnetic field of the form, $\mathbf{H}_{g}(\mathbf{y},t)=\frac{b_{g}(t)}{\mu_{0}}y\hat{y}$, ($b_{g}$ in units of $[\mathrm{T}/\mathrm{m}]$), the interaction Hamiltonian for the COM motion in the $\hat{y}$-direction (frequency $\omega_{y}\equiv\omega_{b}$ for $b_{y}\equiv b$) and the magnon mode is given by, $\mathcal{H}_{mb}(t)=\tilde{\Omega}_{P}(t)\left(\hat{b}+\hat{b}^{\dagger}\right)\left(\hat{m}+\hat{m}^{\dagger}\right),$ with $\tilde{\Omega}_{P}(t)=\frac{b_{g}(t)}{4}\sqrt{\frac{|\gamma|M_{S}}{\rho\omega_{b}}}$, where $\rho=5170\ \rm{kg/m^{3}}$ is the mass density of YIG. Taking into account all these interactions, in the rotating frame of the cavity drive and the displacement picture of the average field in each mode Wang and Clerk (2012), the Hamiltonian is given by

where ${\Delta}_{a}$ is the cavity detuning, and $\tilde{\Omega}_{S}$ is the driven optomagnonic coupling.

In this tripartite system, the optical and COM modes are not directly coupled, however both are coupled to the magnon mode. In the large detuning condition, $\omega_{m}\gg\sqrt{\tilde{\Omega}_{P}^{2}+\tilde{\Omega}_{S}^{2}}$, and under RWA, the model can be reduced to a bipartite effective system with the form of Hamiltonian,

where $\Delta_{\rm{eff}}=(\tilde{\Omega}_{P}^{2}-\tilde{\Omega}_{S}^{2})/(2\omega_{m})$ and $\Omega_{\rm{eff}}=\tilde{\Omega}_{S}\tilde{\Omega}_{P}/\omega_{m}$. This results in a transfer time limit of $\tilde{\tau}_{\rm{lim}}=\pi\omega_{m}/(2\tilde{\Omega}_{S}\tilde{\Omega}_{P})$. In the three-mode system given by the full tripartite model under RWA, a widely applicable scheme of population transfer between the uncoupled modes in an irreversible manner is to apply STIRAP-like ideal pulses Bergmann et al. (1998). However, only when the total operation time of pulse sequences, $T\gg\tilde{\tau}_{\rm{lim}}$, adiabatic effective population transfer can occur between the optical and COM modes in this system. In Fig. S5, we show the cooling dynamics with ideal STIRAP-like Gaussian pulses. This is obtained with counter-intuitive sequence of the coupling parameters for transfer of quanta from the COM mode to the optical mode. The cooling time is of the order of the limit given from the above expression, $\tilde{\tau}_{\rm{lim}}$, which shows that these pulses are not adiabatic, but rather gives the shortest time for excitation transfer. The performance of the cooling quotient can be seen from the figure. As the time limit is the shortest possible time for cooling, one has to compromise on the effective cooling achieved. For the coupling parameters in the range, $\tilde{\Omega}_{i}^{\rm{max}}/\omega_{b}=\{6,8,10\}$, the RWA is valid, whereas for higher values of $\tilde{\Omega}_{i}^{\rm{max}}$ beyond $12\omega_{b}$, the counter-rotating terms have significant effect. Therefore, even if the increase of coupling parameters shows sign of the lowering of transfer time, it is not possible to obtain effective cooling in these cases. On the contrary, our proposed DRL-scheme finds finely-tuned pulses to go beyond such limits to obtain effective and much faster cooling.

The tripartite system is modelled similarly to the two-mode model described earlier, except the fact that an additional Gaussian noise layer is added to the final layer of the network to facilitate more exploration to uncover the non-triviality. Here the controls are, $\tilde{\Omega}_{\rm P/S}$, which are taken to be real and positive. As earlier, the observables are the instantaneous solutions of the coupled differential equations of the second order moments of the quantum master equation of the tripartite system, along with the controls, $\tilde{\Omega}_{\rm P/S}$. In a given episode, the RL-agent makes a total of 150 interactions with the RL-environment in timesteps of $dt=0.1$, each time applying its actions, maintaining a generous exploration-exploitation trade-off. The reward function, for this particular case is,

where $n_{T}$ is the thermal phonon number, $b^{\dagger}b~{}(m^{\dagger}m)$ is the number operator of the phonon (magnon) modes, and $\lambda=10$ is chosen which guarantees to shuffle population out of the phonon mode and into the optical cavity mode leading to cooling of the phonon mode, without transferring to the magnon modes.

Given the large computational overhead to solve the coupled second order moment equations for the tripartite opto-magno-mechanical system with state-of-the-art experimental parameter choice of $\omega_{m}=10^{5}\,\omega_{b}$, the problem is solved in two steps. In the first step, we train an auxiliary system with $\omega_{m}=10^{3}\omega_{b}$ and $\tilde{\Omega}_{\rm P/S}=10\ \omega_{b}$, for which solution of the set of equations can be obtained without much computational effort. The trained auxiliary model is then used as a supervisor/teacher for the actual system, that we call primary, with $\omega_{m}=10^{5}\omega_{b}$ and $\tilde{\Omega}_{\rm P/S}=100\ \omega_{b}$. Training the actual system for a few hundred episodes with periodic evaluation of the RL-agent yields the best trained model. In the literature of RL, this can be considered as a form of imitation learning.

The (auxiliary) RL-agent is trained for the following set of cases: $\tilde{\Omega}_{\rm P/S}^{\rm max}=10,12,15$. The learning curve for the first case is shown in Fig. S6. In these plots, we have shown how the RL-agent gradually learns the non-intuitive control sequences by optimizing the neural network weights to achieve the goal of attaining larger net returns determined by Eq. S20. We have also shown the performance of the RL-agent at different stages of the learning process, as explained in the caption of Fig. S6 in more detail.