††thanks: These authors contributed equally to this work††thanks: These
authors contributed equally to this work††thanks: Current address: Institut
f$\ddot{u}$r Experimentalphysik und Zentrum f$\ddot{u}$r Quantenphysik,
Universit$\ddot{a}$t Innsbruck, 6020 Innsbruck, Austria

# Aharonov-Bohm Caging and Inverse Anderson transition in Ultracold Atoms

Hang Li  Zhaoli Dong  Interdisciplinary Center of Quantum Information, State
Key Laboratory of Modern Optical Instrumentation, Zhejiang Province Key
Laboratory of Quantum Technology and Device, Department of Physics, Zhejiang
University, Hangzhou 310027, China  Stefano Longhi  Dipartimento di Fisica,
Politecnico di Milano, Piazza L. da Vinci 32, I-20133 Milano, Italy  IFISC
(UIB-CSIC), Instituto de Fisica Interdisciplinar y Sistemas Complejos, Palma
de Mallorca, Spain  Qian Liang  Interdisciplinary Center of Quantum
Information, State Key Laboratory of Modern Optical Instrumentation, Zhejiang
Province Key Laboratory of Quantum Technology and Device, Department of
Physics, Zhejiang University, Hangzhou 310027, China  Dizhou Xie
Interdisciplinary Center of Quantum Information, State Key Laboratory of
Modern Optical Instrumentation, Zhejiang Province Key Laboratory of Quantum
Technology and Device, Department of Physics, Zhejiang University, Hangzhou
310027, China  Bo Yan<EMAIL_ADDRESS>Interdisciplinary Center of
Quantum Information, State Key Laboratory of Modern Optical Instrumentation,
Zhejiang Province Key Laboratory of Quantum Technology and Device, Department
of Physics, Zhejiang University, Hangzhou 310027, China

###### Abstract

Aharonov-Bohm (AB) caging, a special flat-band localization mechanism, has
spurred great interest in different areas of physics. AB caging can be
harnessed to explore the rich and exotic physics of quantum transport in
flatband systems, where geometric frustration, disorder and correlations act
in a synergetic and distinct way than in ordinary dispersive band systems. In
contrast to the ordinary Anderson localization, where disorder induces
localization and prevents transport, in flat band systems disorder can induce
mobility, a phenomenon dubbed inverse Anderson transition. Here, we report on
the experimental realization of the AB cage using a synthehtic lattice in the
momentum space of ultracold atoms with tailored gauge fields, demonstrate the
geometric localization due to the flat band and the inverse Anderson
transition when correlated binary disorder is added to the system. Our
experimental platform in a many-body environment provides a fashiinating
quantum simulator where the interplay between engineered gauge fields,
localization, and topological properties of flat band systems can be finely
explored.

Figure 1: The experimental scheme and energy spectrum of Aharonov-Bohm cage.
a. Schematic diagram of a rhombic lattice with three sites ($A_{n}$, $B_{n}$
and $C_{n}$) per unit cell. The sites $A_{n}(n=0,1,...)$ are put on the $F=2$
energy level, while the $B_{n}$ and $C_{n}(n=0,1,...)$ sites are rooted in the
$F=1$ energy level. The complex hopping coefficients between nearest-
neighboring sites are denoted as $Je^{i\theta_{m}}(m=1,2,3,4)$. b. The energy
dispersion $E(k)$ of the chain as a function of the flux $\phi$ (middle). The
left and right insets display the dispersion relation for $\phi=\pi$ and
$\phi=0$, respectively.

Introduction. Exploring the deep insights into localization, disorder and
transport is of upmost relevance in different areas of physics and in modern
quantum technologies, ranging from condensed matter physics [1, 2, 3] to a
variety of artificial quantum and classical systems [4, 5, 6, 7, 8, 9, 10, 11,
12]. When sufficiently strong static uncorrelated disorder is added to a
regular lattice structure, complete localization of wave functions and the
absence of transport is usually observed for non-interacting particles as a
result of Anderson localization [1]. A different approach to create
localization is inducing a flat-band spectrum: here destructive interference
among different propagation paths is realized via the special lattice
geometry, with the formation of perfectly localized compact modes [13]. The
Aharonov-Bohm (AB) caging [14] provides a paradigmatic example of the latter
type of localization, induced by an artificial gauge field. With the
development of modern quantum technologies, a variety of platforms, mainly
based on ultracold quantum gas [15, 16, 17, 18, 19], superconducting circuits
[20, 21] and photonic crystals [22, 23, 10] can be applied for creating
artificial synthetic magnetic flux, thus studying the flat-band localization.
Recently, the flat-band localization in the AB cage has been demonstrated in
photonic lattice systems with curved or auxiliary waveguides [24, 25]. Owing
to the diverging effective mass in a flat band, the system becomes very
sensitive to disorder and interactions, and the emerging phenomena, such as
criticalilty and multifractality, can significantly deviate from conventional
Anderson localization [26, 27, 28, 29, 30, 31, 32, 33, 34, 35]. An important
example is the inverse Anderson transition, i.e. a transition from an
insulating to a metallic phase induced by disorder in a lattice where all
bands are flat. This phenomenon was predicted more than ten years ago for a
three-dimensional diamond lattice with flat bands [27], however its
experimental observation has remained elusive so far, mainly because of the
impracticality of engineering and controlling special three-dimensional
lattice geometries. In low-dimensional flat band systems, inverse Anderson
transition is rather generally prevented, and a competition between geometric
frustration and Anderson localization is typically observed [36]. However, it
has been recently predicted that under certain correlated disorder inverse
Anderson transition could be observed in quasi one-dimensional (1D) AB caging
systems [37].

In this Letter, we experimentally realize a quasi 1D rhombic chain with
synthetic magnetic flux and demonstrate flat-band localization in the
momentum-space lattice of ultracold atoms [38, 39, 40, 41]. Thanks to the
flexible engineering abilities offered by the artificial lattice in momentum
space, we can control on-site correlated disorder in the system and explore
the interplay between disorder and localization, with the observation of the
inverse Anderson transition under anti-symmetric binary disorder. Since our
artificial lattice platform is accomplished under the many-body environment
[42, 43, 44, 45], it could offer the possibility to experimentally explore
other exotic properties of flat band systems with strongly-correlated
particles [46, 13, 47, 48, 49, 50].

Theoretical model and experimental realization. Theoretically, we consider a
quasi-1D rhombic lattice with three coupled sublattices (denoted as A, B, and
C) as shown in Fig. 1(a). The system is described by the following effective
Hamiltonian

$\displaystyle
H_{\text{eff}}=\sum_{n}[-J(e^{i\theta_{3}}\hat{a}_{n}^{\dagger}\hat{b}_{n-1}+e^{-i\theta_{4}}\hat{a}_{n}^{\dagger}\hat{b}_{n}+e^{-i\theta_{2}}\hat{a}_{n}^{\dagger}\hat{c}_{n-1}+$
(1) $\displaystyle\ \ \
e^{i\theta_{1}}\hat{a}_{n}^{\dagger}\hat{c}_{n}+\text{H.c.})-(\Delta_{n}^{(a)}\hat{a}_{n}^{\dagger}\hat{a}_{n}+\Delta_{n}^{(b)}\hat{b}_{n}^{\dagger}\hat{b}_{n}+\Delta_{n}^{(c)}\hat{c}_{n}^{\dagger}\hat{c}_{n})]$

where $J$ is the coupling amplitude between neighboring sites,
$\hat{a}_{n},\hat{b}_{n}$ and $\hat{c}_{n}$
($\hat{a}_{n}^{\dagger},\hat{a}_{n}^{\dagger}$ and $\hat{a}_{n}^{\dagger}$)
are the annihilation (creation) operators for particles at the A, B and C
sites of the $n$-th unit cell, $\phi$ is the synthetic magnetic flux of each
plaquette defined as $\phi=\sum_{m=i}^{4}\theta_{m}$, and
$\Delta^{(i)}_{n}\;(i=a,b,c)$ are the on-site energy potentials in the three
sites at the $i$-th unit cell.

In a clean lattice system, where the on-site energy shifts are set to be zero,
$i.e.~{}\Delta_{n}^{(i)}=0$, there are three bands with dispersion relations:
$E_{0}=0$ and $E_{\pm}=\pm 2J\sqrt{1+\cos(\phi/2)\cos(k+\phi/2)}$, where
$k\;(-\pi\leq k<\pi)$ is the Bloch wavenumber; see the middle panel of Fig.
1(b). Notably, for $\phi=\pi$ the lattice displays three flat bulk bands [left
inset of Fig. 1(b)] with compact localized states. In this case transport is
forbidden owing to geometric localization (AB caging effect). Conversely. for
$\phi\neq\pi$ the system supports only one flat bulk band and two dispersive
bulk bands, as shown in the right inset of Fig. 1(b) for $\phi=0$.

Besides the clean lattice configuration, we can add the on-site static
disorder terms to explore the insulating and metallic phases in the rhombic
chain. For $\phi\neq\pi$, disorder can induce Anderson localization of
dispersive bands and prohibit transport. However, for $\phi=\pi$, i.e. in the
fully flat band case, an inverse Anderson transition can arise: adding
disorder breaks the geometric localization of compact states and enables
transport in the lattice [37].

Experimentally, we engineer a chain of Aharonov-Bohm rhombic rings along the
momentum lattice in a ${}^{87}\textrm{Rb}$ Bose-Einstein condensate (BEC) with
number $~{}\sim 2\times 10^{5}$ [19, 51, 52]. As illustrated in Fig. 1(a), the
three sublattice sites are experimentally realized in the synthetic dimensions
of atomic internal states, $i.e.$ the ground-state hyperfine state
$\left|F=1,m_{F}=0\right>$ for sites A and B, and $\left|F=2,m_{F}=0\right>$
for site C. The adjacent sites between different internal states are coupled
by Raman transitions, and the adjacent sites in $F=1$ are coupled via Bragg
transitions [53, 54]. The tunneling amplitude is set to be $J=0.95(5)$ kHz for
all the experiments, and the synthetic magnetic flux (denoted as $\phi$) is
controlled by the relative phases of coupling lasers [54]. In total, we
implement a chain of 10 unit cells with open boundary condition (OBC) and
mainly focus on the bulk state dynamics.

Figure 2: Aharonov-Bohm cage dynamics. a. The Aharonov-Bohm cage evolution
with flux $\phi=\pi$ in each plaquette. b. Measuring the population
distribution of sites A, B, and C in 1 $ms$. The AB cage effect appears in the
middle adjacent unit cells. The small population leakage of the middle cage
can be mainly ascribed to the imperfect experimental parameters. The solid
lines represent the fitting of experimental data. c. The evolution dynamics
with flux $\phi=0$ in each plaquette. d. The population distribution of sites
A, B, and C in the middle unit cell with flux $\phi=0$. All the fitted solid
lines show a fast exponential decay feature . In (a) and (c), the population
of all lattice sites has been normalized in each displayed time instant.

Aharonov-Bohm cage dynamics. As Fig. 1(b) shows, the energy spectrum of AB
caging limit ($\phi=\pi$ and without on-site energy shift) has three non-
dispersive bulk bands. The zero eigenvalue ($E_{0}=0$) has the localized
eigenstate
$\left|\psi_{n,0}^{bulk}\right>=(b^{\dagger}_{n-1}+c^{\dagger}_{n-1}+b^{\dagger}_{n}-c^{\dagger}_{n})\left|0\right>$,
in which the $A$ state is not involved. And the other two eigenstates of flat
bands are
$\left|\psi_{n,\pm}^{bulk}\right>=(b^{\dagger}_{n-1}+c^{\dagger}_{n-1}\mp
2a^{\dagger}_{n}-b^{\dagger}_{n}-c^{\dagger}_{n})\left|0\right>$ [24, 37].
Therefore, the initial excitation at the site $A_{n}$ can be viewed as the
superposition of $\left|\psi_{n,\pm}^{bulk}\right>$, whose dynamics shows the
feature of breathing motion and the oscillation frequency is determined by
their eigenfrequency difference. In the experiment, we prepare the initial BEC
wave packet at site A of the middle unit cell (labeled as $0_{th}$) and
observe the evolution process of its population. Figure 2(a) shows the AB
caging evolution dynamics in a period of 1 $ms$, which displays the
localization and breathing features. As Fig. 2(b) shows, the oscillated
population of site A (or B and C) gives the frequency $\omega=$3.72(5) kHz,
which is consistent with the energy spectrum of Fig. 1(b) ($\omega=4J$ in
theory). Due to the decoherence of our Raman coupling lattice, the fitted
decay lifetime of cage dynamics is 0.70(2) $ms$, which is depicted by the
cosine decay of Fig. 2(b). The ideal AB cage sustains compact localized
states, while the actual experimental result suffers for a small population
leakage as depicted in Fig. 2(b), which is mainly ascribable to the imperfect
symmetric neighboring coupling and decoherence [54]. By comparison, we also
show the evolution dynamics with $\phi=0$ in Figs. 2(c) and 2(d). The
geometric localization feature disappears and each site population of the
$0_{th}$ unit cell displays a fast near-exponential decay behavior.

Inverse Anderson transition. As mentioned before, the system with non-$\pi$
synthetic magnetic flux possesses one flat band ($E_{0}$) and two dispersive
bands ($E_{\pm}$). For $\phi\neq\pi$, the addition of static disorder will
generally lift the flat band eigenvalue degeneracy and cause mixing inside
this band, which results in diffusion as those states with dispersive bands.
For $\phi=\pi$, i.e. in the fully flat band case, the addition of uncorrelated
on-site static disorder is not able to induce an inverse Anderson transition
[36]. In fact, for strong uncorrelated disorder ($i.e.$ larger than $2J$)
localization features as in the ordinary Anderson localization are observed:
the energy bands are mixed and the localization length decreases as the
disorder strength increases [36]. On the other hand, for weak disorder ($i.e.$
smaller than $2J$) the Anderson localization effect is overwhelmed: the
eigenstates from each band are separated by gaps and their localization
lengths do not depend on the disorder strength [36].

Figure 3: The localization and transport with symmetric vs. anti-symmetric
correlated disorder. a. Schematic diagram of adding on-site disorders. b. The
momentum lattices evolution of ballistic transport with the anti-symmetric on-
site disorder. c. The momentum lattices evolution of AB cage localization with
the symmetric on-site disorder. d. The compared evolutions of distance
$\mathcal{D}(t)$ with a different type of on-site disorder. All above results
are set with flux $\phi=\pi$. In (b) and (c), the population has also been
normalized in each displayed moment. The above experimental results refer to a
single realization of correlated Bernoulli disorder.

In contrast, the correlated on-site static disorder in sublattices B and C can
induce transport, i.e. an inverse Anderson transition can be observed for
correlated disorder in the system [37]. There are two typical cases of
correlated disorder, $i.e.$ the symmetric
($\Delta_{n}^{(b)}=\Delta_{n}^{(c)}$) and anti-symmetric
($\Delta_{n}^{(b)}=-\Delta_{n}^{(c)}$) correlated disorder, where
$\Delta_{n}^{(b)}$ and $\Delta_{n}^{(c)}$ are independent stochastic variables
with the same probability density function of zero mean. As shown in [37] and
briefly reviewed in [54], the symmetric case shares similar localization
dynamics as the AB caging: band degeneracy is lifted by the disorder but the
eigenstates are still compact localized states. More striking behaviors appear
in the anti-symmetric correlated scenario
$\Delta_{n}^{(b)}=-\Delta_{n}^{(c)}$, where the existence of extended states
can be observed, indicating the disorder-induced transport regime (inverse
Anderson transition). The localization or mobility features dynamics for anti-
symmetric correlated disorder depends on the forms of the probability density
function of the stochastic variables $\Delta_{n}^{(b)}=-\Delta_{n}^{(c)}$
[37]. Specifically. the uniform distribution of probability density function
prevents mobility, while the Bernoulli distribution, where $\Delta^{(b)}_{n}$
can take randomly only the two values $\pm\Delta$, induces ballistic
transport.

Figure 4: The energy bands and evolution with anti-symmetric Bernoulli
disorder. a. The quasienergy spectrum for different on-site disorder
strengths, where we have calculated eigenvalues of a chain with 100 unit
cells. The left and right inset correspond to the smallest and largest
bandwidth, respectively. b. The measured distance with various disorder
strengths in $0.6~{}ms$. The inset line corresponds to the theoretical
calculated distance. The above measurements refer to a single realization of
Bernoulli disorder.

To characterize the transport dynamics, we excite the system in a single site
and monitor the spreading dynamics using the second-order moment of position
operator $\mathcal{D}^{2}(t)$, given by

$\mathcal{D}^{2}(t)=\sum_{n}n^{2}(\left|a_{n}\right|^{2}+\left|b_{n}\right|^{2}+\left|c_{n}\right|^{2}).$
(2)

A distance $\mathcal{D}(t)$ that linearly grows with time, i.e.
$\mathcal{D}(t)\sim t$ , is the clear signature of ballistic transport regime.
Figure 3(a) shows a schematic of the setup where on-site disorder energy shift
is added in our momentum lattice system [54], whereas the experimental results
of wave packet spreading are illustrated in Figs.3(b-d). In the experiment,
the initial excitation of BEC wavepacket is in the sublattice A of unit cell
at site $n=0$, $i.e.~{}a_{n}(0)=\delta_{n,0}$ and $b_{n}(0)=c_{n}(0)=0$. The
figures clearly show that the symmetric and anti-symmetric correlated disorder
display completely different evolution dynamics: while ballistic transport is
clearly observed in the antisymmetric case, corresponding to an inverse
Anderson transition, wave spreading is prevented in the symmetric case
[Fig.3(d)].

To further unveil the features of the anti-symmetric disorder case, in
Fig.4(a) we show the energy spectrum of the finite rhombic chain with the flux
$\phi=\pi$ under antisymmetric disorder with a Bernoulli distribution,
parametrized in the strength $\Delta$ of disorder. We can see that the energy
spectrum changes from the flat bands at $\Delta=0$ to dispersive bands as the
disorder strength increase from zero to a large value. The theoretical
analysis [54] indicates that for $\Delta>0$ the Hamiltonian displays
absolutely continuous spectrum with three dispersive Bloch bands, and their
gaps vanish at $\Delta=\Delta_{0}=\sqrt{2}J$. At this turning point, the
transport is fastest due to the largest bandwidth. Experimentally, we fix the
evolution time to be $t=0.6\;ms$ and measure the distance $\mathcal{D}(t)$ for
different disorder strengths. As shown in Fig. 4(b), our experimental results
support this prediction. However, due to the finite size effect, our
experimentally measured turning amplitude is $\Delta_{0}=1.7J$, which is
larger than the theoretically predicted value. If the rhombic chain is long
enough or the measured evolution time is much longer, our numerical results
show that the turning point is approaching the prediction of the ideal
infinite chain [54]. The aberration of the dynamics at larger disorder
strength is mainly due to the decoherence of our system, which slower the
quantum transport and makes the distance smaller. This effect is more
important for fast transport, thus the distance is reduced more for larger
disorder strength.

Conclusion. In conclusion, we reported on the experimental realization of a
quasi-1D rhombic chain in a momentum lattice of ultracold atoms, and
demonstrated Aharonov-Bohm caging dynamics when a synthetic gauge field is
applied. Thanks to the flexible engineering abilities provided by our
platform, we could finely tune the tunneling process and introduce
controllable disorder into the system, unraveling the interplay between
disorder and transport in a system with geometric localization. In particular,
we reported on the first experimental observation of an inverse Anderson
localization in a flat band system [27], i.e. the delocalization of the wave
functions and ballistic transport induced by suitable correlated static
disorder in the AB system. The AB photonic cage and inverse Anderson
transition reported in our experiments are observed in a regime where particle
interaction is negligible, however thanks to the many-body environment offered
by the ultracold gas platform our synthetic lattice setup could be feasible
for the experimental demonstrations of other exotic effects arising in
strongly-correlated flat bands systems.

###### Acknowledgements.

Acknowledgement:- We acknowledge the support from the National Key Research
and Development Program of China under Grant No. 2018YFA0307200, the National
Natural Science Foundation of China under Grants No. U21A20437 and No.
12074337, Natural Science Foundation of Zhejiang Province under Grant No.
LR21A040002, Zhejiang Province Plan for Science and Technology Grant No.
2020C01019, and the Fundamental Research Funds for the Central Universities
under Grant No. 2021FZZX001-02.

## References

  * Anderson [1958] P. W. Anderson, Phys. Rev. 109, 1492 (1958).
  * Hasan and Kane [2010] M. Z. Hasan and C. L. Kane, Rev. Mod. Phys. 82, 3045 (2010).
  * Qi and Zhang [2011] X.-L. Qi and S.-C. Zhang, Rev. Mod. Phys. 83, 1057 (2011).
  * Houck  _et al._ [2012] A. A. Houck, H. E. Tureci, and J. Koch, Nature Physics 8, 292 (2012).
  * Blatt and Roos [2012] R. Blatt and C. F. Roos, Nature Physics 8, 277 (2012).
  * Bloch  _et al._ [2012] I. Bloch, J. Dalibard, and S. Nascimbene, Nature Physics 8, 267 (2012).
  * Bloch  _et al._ [2008] I. Bloch, J. Dalibard, and W. Zwerger, Reviews of Modern Physics 80, 885 (2008).
  * Cooper  _et al._ [2019] N. R. Cooper, J. Dalibard, and I. B. Spielman, Reviews of Modern Physics 91, 015005 (2019).
  * Browaeys and Lahaye [2020] A. Browaeys and T. Lahaye, Nature Physics 16, 132 (2020).
  * Ozawa  _et al._ [2019] T. Ozawa, H. M. Price, A. Amo, N. Goldman, M. Hafezi, L. Lu, M. C. Rechtsman, D. Schuster, J. Simon, O. Zilberberg, and I. Carusotto, Reviews of Modern Physics 91, 015006 (2019).
  * Cai  _et al._ [2013] J. M. Cai, A. Retzker, F. Jelezko, and M. B. Plenio, Nature Physics 9, 168 (2013).
  * Segev  _et al._ [2013] M. Segev, Y. Silberberg, and D. N. Christodoulides, Nature Photonics 7, 197 (2013).
  * Leykam  _et al._ [2018] D. Leykam, A. Andreanov, and S. Flach, Advances in Physics: X 3, 1473052 (2018).
  * Vidal  _et al._ [1998] J. Vidal, R. Mosseri, and B. Douçot, Phys. Rev. Lett. 81, 5888 (1998).
  * Lin  _et al._ [2009] Y. J. Lin, R. L. Compton, K. Jimenez-Garcia, J. V. Porto, and I. B. Spielman, Nature 462, 628 (2009).
  * Aidelsburger  _et al._ [2013] M. Aidelsburger, M. Atala, M. Lohse, J. T. Barreiro, B. Paredes, and I. Bloch, Physical Review Letters 111, 185301 (2013).
  * Miyake  _et al._ [2013] H. Miyake, G. A. Siviloglou, C. J. Kennedy, W. C. Burton, and W. Ketterle, Phys. Rev. Lett. 111, 185302 (2013).
  * Atala  _et al._ [2014] M. Atala, M. Aidelsburger, M. Lohse, J. T. Barreiro, B. Paredes, and I. Bloch, Nature Physics 10, 588 (2014).
  * Gou  _et al._ [2020] W. Gou, T. Chen, D. Xie, T. Xiao, T.-S. Deng, B. Gadway, W. Yi, and B. Yan, Phys. Rev. Lett. 124, 070402 (2020).
  * Abilio  _et al._ [1999] C. C. Abilio, P. Butaud, T. Fournier, B. Pannetier, J. Vidal, S. Tedesco, and B. Dalzotto, Phys. Rev. Lett. 83, 5102 (1999).
  * Kou  _et al._ [2017] A. Kou, W. C. Smith, U. Vool, R. T. Brierley, H. Meier, L. Frunzio, S. M. Girvin, L. I. Glazman, and M. H. Devoret, Phys. Rev. X 7, 031037 (2017).
  * Lu  _et al._ [2014] L. Lu, J. D. Joannopoulos, and M. Soljacic, Nature Photonics 8, 821 (2014).
  * Khanikaev and Shvets [2017] A. B. Khanikaev and G. Shvets, Nature Photonics 11, 763 (2017).
  * Mukherjee  _et al._ [2018] S. Mukherjee, M. Di Liberto, P. Öhberg, R. R. Thomson, and N. Goldman, Phys. Rev. Lett. 121, 075502 (2018).
  * Kremer  _et al._ [2020] M. Kremer, I. Petrides, E. Meyer, M. Heinrich, O. Zilberberg, and A. Szameit, Nature Communications 11, 907 (2020).
  * Vidal  _et al._ [2001] J. Vidal, P. Butaud, B. Douçot, and R. Mosseri, Phys. Rev. B 64, 155306 (2001).
  * Goda  _et al._ [2006] M. Goda, S. Nishino, and H. Matsuda, Phys. Rev. Lett. 96, 126401 (2006).
  * Chalker  _et al._ [2010] J. T. Chalker, T. S. Pickles, and P. Shukla, Phys. Rev. B 82, 104209 (2010).
  * Bodyfelt  _et al._ [2014] J. D. Bodyfelt, D. Leykam, C. Danieli, X. Yu, and S. Flach, Phys. Rev. Lett. 113, 236403 (2014).
  * Flach  _et al._ [2014] S. Flach, D. Leykam, J. D. Bodyfelt, P. Matthies, and A. S. Desyatnikov, Epl 105, 30001 (2014).
  * Leykam  _et al._ [2017] D. Leykam, J. D. Bodyfelt, A. S. Desyatnikov, and S. Flach, European Physical Journal B 90, 1 (2017).
  * Shukla [2018] P. Shukla, Phys. Rev. B 98, 054206 (2018).
  * Gneiting  _et al._ [2018] C. Gneiting, Z. Li, and F. Nori, Phys. Rev. B 98, 134203 (2018).
  * Baboux  _et al._ [2016] F. Baboux, L. Ge, T. Jacqmin, M. Biondi, E. Galopin, A. Lemaître, L. Le Gratiet, I. Sagnes, S. Schmidt, H. E. Türeci, A. Amo, and J. Bloch, Phys. Rev. Lett. 116, 066402 (2016).
  * Liu  _et al._ [2022] F. Liu, Z.-C. Yang, P. Bienias, T. Iadecola, and A. V. Gorshkov, Phys. Rev. Lett. 128, 013603 (2022).
  * Gligorić  _et al._ [2020] G. Gligorić, D. Leykam, and A. Maluckov, Phys. Rev. A 101, 023839 (2020).
  * Longhi [2021] S. Longhi, Optics Letters 46, 2872 (2021).
  * Meier  _et al._ [2016] E. J. Meier, F. A. An, and B. Gadway, Phys. Rev. A 93, 051602 (2016).
  * Meier  _et al._ [2018] E. J. Meier, F. A. An, A. Dauphin, M. Maffei, P. Massignan, T. L. Hughes, and B. Gadway, Science 362, 929 (2018).
  * Lapp  _et al._ [2019] S. Lapp, J. Ang’ong’a, F. A. An, and B. Gadway, New Journal of Physics 21, 045006 (2019).
  * Xie  _et al._ [2019] D. Xie, W. Gou, T. Xiao, B. Gadway, and B. Yan, npj Quan. Inf. 5, 55 (2019).
  * Xie  _et al._ [2020] D. Xie, T.-S. Deng, T. Xiao, W. Gou, T. Chen, W. Yi, and B. Yan, Phys. Rev. Lett. 124, 050502 (2020).
  * An  _et al._ [2021] F. A. An, B. Sundar, J. Hou, X.-W. Luo, E. J. Meier, C. Zhang, K. R. A. Hazzard, and B. Gadway, Phys. Rev. Lett. 127, 130401 (2021).
  * Xiao  _et al._ [2021] T. Xiao, D. Z. Xie, Z. L. Dong, T. Chen, W. Yi, and B. Yan, Science Bulletin 66, 2175 (2021).
  * Wang  _et al._ [2022] Y. Wang, J.-H. Zhang, Y. Li, J. Wu, W. Liu, F. Mei, Y. Hu, L. Xiao, J. Ma, C. Chin, and S. Jia, arXiv:2204.12730 (2022).
  * Vidal  _et al._ [2000] J. Vidal, B. Douçot, R. Mosseri, and P. Butaud, Phys. Rev. Lett. 85, 3906 (2000).
  * Jünemann  _et al._ [2017] J. Jünemann, A. Piga, S.-J. Ran, M. Lewenstein, M. Rizzi, and A. Bermudez, Phys. Rev. X 7, 031057 (2017).
  * Danieli  _et al._ [2020] C. Danieli, A. Andreanov, and S. Flach, Phys. Rev. B 102, 041116 (2020).
  * Zurita  _et al._ [2020] J. Zurita, C. E. Creffield, and G. Platero, Advanced Quantum Technologies 3, 1900105 (2020).
  * Kuno  _et al._ [2020] Y. Kuno, T. Orito, and I. Ichinose, New Journal of Physics 22, 013032 (2020).
  * Xie  _et al._ [2018] D. Xie, D. Wang, W. Gou, W. Bu, and B. Yan, J Opt. Soc. Am. B 35, 500 (2018).
  * Chen  _et al._ [2021] T. Chen, W. Gou, D. Z. Xie, T. Xiao, W. Yi, J. Jing, and B. Yan, Npj Quantum Information 7, 78 (2021).
  * Liang  _et al._ [2022] Q. Liang, D. Xie, Z. Dong, H. Li, H. Li, B. Gadway, W. Yi, and B. Yan, arXiv:2201.09478 (2022).
  * [54] See Supplemental Information for details. .


# An Improved Data Augmentation Scheme for Model Predictive Control Policy
Approximation

Dinesh Krishnamoorthy The author is with the Department of Mechanical
Engineering, Eindhoven University of Technology, 5600 MB, Eindhoven, The
Netherlands<EMAIL_ADDRESS>
###### Abstract

This paper considers the problem of data generation for MPC policy
approximation. Learning an approximate MPC policy from expert demonstrations
requires a large data set consisting of optimal state-action pairs, sampled
across the feasible state space. Yet, the key challenge of efficiently
generating the training samples has not been studied widely. Recently, a
sensitivity-based data augmentation framework for MPC policy approximation was
proposed, where the parametric sensitivities are exploited to cheaply generate
several additional samples from a single offline MPC computation. The error
due to augmenting the training data set with inexact samples was shown to
increase with the size of the neighborhood around each sample used for data
augmentation. Building upon this work, this letter paper presents an improved
data augmentation scheme based on predictor-corrector steps that enforces a
user-defined level of accuracy, and shows that the error bound of the
augmented samples are independent of the size of the neighborhood used for
data augmentation.

###### Index Terms:

Data Augmentation, direct policy approximation, imitation learning, parametric
sensitivities

##  I Background

###  I-A Pre-computed control policies

Implementing optimization-based controllers such as model predictive control
(MPC) for fast dynamic systems with limited computing and memory capacity has
motivated research on pre-computing and storing the optimal control policy
offline such that it can be used online without recomputing the optimization
problem. This was first studied under the framework of explicit MPC for linear
time-invariant systems [1]. However, the synthesis of the explicit MPC law
does not scale well, since the number of piecewise affine polyhedral regions
required to capture the MPC control law can increase exponentially with
problem size. Extension to nonlinear systems and economic objectives are also
not trivial.

An alternative approach is to approximate the optimal policy using parametric
function approximators, such as deep neural networks, which are broadly
studied under the context of _learning from demonstrations_. The idea of
approximating an MPC policy using neural networks was first proposed in [2],
but this idea remained more or less dormant (due to the high cost of offline
training). However, with the recent promises of deep learning, there has been
an unprecedented surge of interest in approximating control policies, see e.g.
[3, 4, 5, 6, 7, 8, 9, 10, 11, 12]. This has also inspired a number of
applications ranging from energy[6], automotive [13, 14], chemical processing
[10, 15], robotics [16], spacecraft [9], and healthcare [17] to name a few.
Research developments in this direction have been predominantly devoted to
understanding the safety and performance of approximate policies [4, 13, 8,
18]. Although, these are very important developments in the direction of MPC
policy approximation, a major challenge of this approach that hinders
practical implementation is the cost of training the policy, which is not well
studied in the literature, as also noted in [19].

Figure 1: MPC policy approximation pipeline with data augmentation scheme for
efficient data generation.

###  I-B High cost of offline learning

Consider the “expert” policy which is given by solving the MPC problem

$\displaystyle V^{*}(x(t))=\min_{x_{k},u_{k}}$
$\displaystyle\;\sum_{k=0}^{N-1}\ell_{Q}(x_{k},u_{k})+\ell_{P}(x_{N})$ (1a)
s.t. $\displaystyle x_{k+1}=f(x_{k},u_{k},d)\quad\forall
k\in\mathbb{I}_{0:N-1}$ (1b)
$\displaystyle{\color[rgb]{0,0,0}x\in\mathcal{X},\quad u\in\mathcal{U}},\quad
x_{N}\in\mathcal{X}_{f}$ (1c) $\displaystyle x_{0}=x(t)$ (1d)

where $x\in{X}\subseteq\mathbb{R}^{n_{x}}$,
$u\in\mathcal{U}\subseteq\mathbb{R}^{n_{u}}$, and $d\in\mathbb{R}^{n_{d}}$
denote the states, inputs, and the model parameters, respectively.
$f:\mathbb{R}^{n_{x}}\times\mathbb{R}^{n_{u}}\times\mathbb{R}^{n_{d}}\rightarrow\mathbb{R}^{n_{x}}$
denotes the system model,
$\ell_{Q}:\mathbb{R}^{n_{x}}\times\mathbb{R}^{n_{u}}\rightarrow\mathbb{R}$ and
$\ell_{P}:\mathbb{R}^{n_{x}}\rightarrow\mathbb{R}$ denote the stage and
terminal cost, which are parameterized by the tuning weights $Q$ and $P$,
respectively, which are treated very generically in this letter to include
different MPC formulations, including economic stage and terminal costs. $N$
is the length of the prediction horizon,  (1c) denotes the path and terminal
constraints, and (1d) denotes the initial condition constraint. Solving (1) at
each time step and implementing the first control input gives the implicit MPC
policy $u^{*}=\pi^{*}(x(t))$.

Approximating the expert policy $\pi^{*}(x)$ using supervised learning
requires a data set consisting of the expert demonstrations in the form of
optimal state-action pairs
$\mathcal{D}:=\\{(x_{i},u_{i}^{*})\\}_{i=1}^{N_{s}}$, which tells us what
action $u_{i}^{*}$ the expert (i.e., MPC) would take at state $x_{i}$. Using
this data set, we can learn a parameterized policy $\pi(x;\theta)$ that tells
what actions to take as a function of the current state $x$, such that
$\pi(x;\theta)$ mimics $\pi^{*}(x)$. The data set $\mathcal{D}$ is generated
by solving the MPC problem offline for different realizations of $x_{i}$,
ideally covering the entire feasible state space $\mathcal{X}_{feas}$.

The availability of a sufficiently rich training data set covering the entire
state space is a key stipulation for satisfactory learning [19, 4]. In fact,
it has also been shown that the learned policy can result in instability when
insufficient demonstrations are used for policy fitting [20, 21]. This means
the MPC problem (1) must be solved for several different state realizations
$x_{i}$, the size of which increases exponentially as the problem dimension
increases, (much like the scalability issues with explicit MPC, albeit
offline). For example, if we want to include at least $s$ samples along each
dimension, then we need to generate at least $N_{s}=s^{n_{x}}$ samples. To
this end, the first challenge is the cost of generating the data set
$\mathcal{D}:=\\{(x_{i},u_{i}^{*})\\}_{i=1}^{N_{s}}$.

Assume now that a sufficiently rich data set $\mathcal{D}$ has been generated
offline covering $\mathcal{X}$, and an approximate policy $\pi(x;\theta)$ is
learned and has passed the necessary statistical validations and is certified
for online use. Now, if any of the MPC tuning parameters $Q$, $P$ change, or
if the model parameter $d$ is updated online (which is not uncommon), then
this renders the learned policy $\pi(x;\theta)$ useless. A new training data
set $\mathcal{D}$ now has to be generated using the modified MPC formulation
and the approximate policy must be re-trained and re-validated from scratch.
This is perhaps one of the biggest drawbacks of MPC policy approximation. In
theory, this can of course be circumvented by training a parametric function
that is also a function of the weights and model parameters $\pi(p;\theta)$,
where for example $p:=[x,\text{vec}(Q),\text{vec}(P),d]^{\mathsf{T}}$ and
$\pi:\mathcal{P}\rightarrow\mathcal{U}$ with $\mathcal{P}$ being the the
combined state and parameter space. However, this makes the offline training
problem even more expensive since the dimension of $p$ can be extremely large.
Despite the surge of interest in MPC policy approximation, its practical
impact will depend on addressing the pivotal challenge of offline data
generation.

###  I-C Data Augmentation

In machine learning literature, and more particularly, computer vision and
image classification, the issue of data-efficiency is addressed using a simple
yet powerful technique known as “Data Augmentation”, broadly defined as a
strategy to artificially increase the number of training samples using
computationally inexpensive transformations of the existing samples [22, 23,
24]. When the data set is a collection of images (i.e. pixel-based data), data
augmentation techniques include geometric transformations (such as rotate,
translate, crop, flip, etc.), photometric transformations (such as colorize,
saturation, contrast, brightness, etc.), and elastic transformation (such as
deformation, shear, grid distortion, etc.) of existing images to artificially
augment several additional training samples.

Unfortunately, when it comes to data sets consisting of optimal state-action
pairs sampled from an optimal policy such as $\pi^{*}(x)$, the existing data
augmentation methods are not applicable. For example, what do
geometric/photometric/elastomeric transformations even mean for an optimal
state-action pair $(x_{i},u_{i}^{*})$?! By data augmentation for optimal
policy observations, we mean the following: Given an optimal state-action pair
$(x_{i},u_{i}^{*})$ observed by querying the expert policy $\pi^{*}(x_{i})$,
we would like to augment additional data points using computationally
inexpensive transformations that would tell us what the optimal action
${u}_{j}^{*}$ would be for states $x_{j}:=x_{i}+\Delta x$ in the vicinity of
$x_{i}$. If we can get a reasonable approximation
$\hat{u}_{j}\approx{u}_{j}^{*}$ without actually querying the expert, then the
inexact optimal state-action pair $(x_{j},\hat{u}_{j})$ can be augmented to
the set of demonstrations used for learning the policy. In other words,
instead of simply learning from a given set of expert demonstrations, we would
like to augment additional samples that are _inferred_ from the expert
demonstrations, thereby enabling us to generalize to states not included in
the original set of expert queries.

Noting that the MPC problem (1) is parametric in the initial condition, a
sensitivity-based data augmentation scheme amenable for optimal state-action
pairs was recently proposed in [25], where it was shown that augmenting
additional samples using parametric sensitivities only require the solution to
a system of linear equations, which is computationally much cheaper than
solving the optimization problem. In [25], the error bound between the expert
policy and approximate policy learned from the augmented data set was shown to
depend on the size of the region $\Delta x$ used for data augmentation.
Therefore, if we want a certain desired accuracy, depending on the problem
Lipschitz properties, this limits the size of $\Delta x$ that can be used for
data augmentation. Building on our recent work [25], the main contribution of
this letter paper is an improved data augmentation scheme, where the augmented
data samples are enforced to be within an user-specified accuracy by using
additional corrector steps. This would make the error bound independent of the
size of $\Delta x$, thereby enabling us to augment samples from a larger
neighborhood $\Delta x$ without jeopardizing accuracy.

##  II Sensitivity-based data augmentation

###  II-A Preliminaries

For the sake of generality, we rewrite the MPC problem (1) in the standard
parametric NLP form

$\displaystyle\Pi(p):\quad\min_{\mathbf{w}}$ $\displaystyle\;J(\mathbf{w},p)$
(2a) s.t. $\displaystyle\;c(\mathbf{w},p)=0,\quad g(\mathbf{w},p)\leq 0$ (2b)

where $p\in\mathbb{R}^{n_{p}}$ includes current state of the system $x_{i}$,
vectorized tuning parameters $Q,P$, system model parameters $d$, or any other
parameter that can change during online deployment.
$\mathbf{w}\in\mathbb{R}^{n_{w}}$ denotes the primal decision variables (with
the optimal action $u^{*}_{i}\in\mathbb{R}^{n_{u}}$ contained within
$\mathbf{w}^{*}$),
$J:\mathbb{R}^{n_{p}}\times\mathbb{R}^{n_{w}}\rightarrow\mathbb{R}$ denotes
the objective function,
$c:\mathbb{R}^{n_{p}}\times\mathbb{R}^{n_{w}}\rightarrow\mathbb{R}^{n_{c}}$
and
$g:\mathbb{R}^{n_{p}}\times\mathbb{R}^{n_{w}}\rightarrow\mathbb{R}^{n_{g}}$
denotes the set of equality and inequality constraints, respectively.

The Lagrangian of the optimization problem (2) is given by
$\mathcal{L}(\mathbf{w},\lambda,\mu,p)=J(\mathbf{w},p)+\lambda^{\mathsf{T}}c(\mathbf{w},p)+\mu^{\mathsf{T}}g(\mathbf{w},p)$
where $\lambda\in\mathbb{R}^{n_{c}}$ and $\mu\in\mathbb{R}^{n_{g}}$ are the
dual variables. For the inequality constraints, we define
$g_{\mathbb{A}}\subseteq g$ as the set of active inequality constraints (i.e.
$g_{\mathbb{A}}(\mathbf{w},p)=0$) and assume strict complimentarity (i.e.,
$\mu_{\mathbb{A}}>0$). The first order necessary conditions of optimality can
be denoted compactly as,

$\displaystyle\varphi(\mathbf{s}(p),p)$
$\displaystyle:=\begin{bmatrix}\nabla\mathcal{L}(\mathbf{w},\lambda,\mu,p)\\\
c(\mathbf{w},p)\\\ g_{\mathbb{A}}(\mathbf{w},p)\end{bmatrix}=0$ (3)
$\displaystyle\mu_{\mathbb{A}}$
$\displaystyle>0,\mu_{\bar{\mathbb{A}}}=0,g_{\bar{\mathbb{A}}}(\mathbf{w},p)<0$

Any primal-dual vector
$\mathbf{s}^{*}(p):=[\mathbf{w}^{*},\lambda^{*},\mu^{*}]^{\mathsf{T}}$ is
called a KKT-point if it satisfies the first order necessary conditions of
optimality (3).

###### Assumption 1

The cost and constraints $J(\cdot,\cdot)$, $c(\cdot,\cdot)$, and
$g(\cdot,\cdot)$ of the NLP problem $\Pi(p)$ are twice continuously
differentiable in the neighborhood of the KKT point $\mathbf{s}^{*}(p)$.

###### Assumption 2

Linear independence constraint qualification, second order sufficient
conditions and strict complementarity holds for $\Pi(p)$.

The above assumption implies that the primal-dual solution vector
$\mathbf{s}^{*}(p_{i})$ is a unique local minimizer of $\Pi(p)$. Given
Assumptions 1 and 2, it was shown in [26] that for parametric perturbations
$\Delta p$ in the neighborhood of $p_{i}$, there exists a unique, continuous,
and differentiable vector function $\mathbf{s}(p_{i}+\Delta p)$ which is a KKT
point satisfying LICQ and SOSC for $\Pi(p_{i}+\Delta p)$, and that the
solution vector satisfies $\|\mathbf{s}^{*}(p_{i}+\Delta
p)-\mathbf{s}^{*}(p_{i})\|\leq L_{s}\|\Delta p\|$, where the notation
$\|\cdot\|$ by default denotes Euclidean norm.

######  Theorem 1 ([25])

Given Assumptions 1 and 2, let $u_{i}^{*}$ be the optimal policy obtained by
querying $\Pi(p_{i})$ for some $p_{i}\in\mathcal{P}$. Then additional samples
$\\{(p_{i}+\Delta p_{j},u_{i}^{*}+\Delta u_{j})\\}_{j}$ in the neighborhood of
$p_{i}$ with the same active constraint set $g_{\mathbb{A}}$ can be augmented
to the data set without querying $\Pi(p_{i}+\Delta p_{j})$, where $\Delta
u_{j}\subset\Delta\mathbf{s}_{j}$ is given by the linear predictor
$\Delta\mathbf{s}_{j}=\mathcal{H}_{i}\Delta p_{j}$.

###### Proof:

Taylor series expansion of $\mathbf{s}^{*}(p)$ around $p_{i}$, gives

$\displaystyle\mathbf{s}^{*}(p_{i}+\Delta
p_{j})=\mathbf{s}^{*}(p_{i})+\frac{\partial\mathbf{s}^{*}}{\partial p}\Delta
p_{j}+\mathcal{O}(\|\Delta p_{j}\|^{2})$ (4)

Since $\mathbf{s}^{*}(p)$ satisfies (3) for any $p$ in the neighborhood of
$p_{i}$, using the implicit function theorem, we have

$\displaystyle\frac{\partial}{\partial
p}\varphi(\mathbf{s}^{*}(p),p)\bigg{|}_{p=p_{i}}=\frac{\partial\varphi}{\partial\mathbf{s}}\frac{\partial\mathbf{s}^{*}}{\partial
p}+\frac{\partial\varphi}{\partial p}=0$ $\displaystyle\Rightarrow$
$\displaystyle\frac{\partial\mathbf{s}^{*}}{\partial
p}=-\left[\frac{\partial\varphi}{\partial\mathbf{s}}\right]^{-1}\frac{\partial\varphi}{\partial
p}=:\mathcal{H}_{i}$

where,

$\frac{\partial\varphi}{\partial\mathbf{s}}:=\begin{bmatrix}\nabla^{2}_{\mathbf{w}\mathbf{w}}\mathcal{L}&\nabla_{\mathbf{w}}{c}&\nabla_{\mathbf{w}}{g}_{\mathbb{A}}\\\
\nabla_{\mathbf{w}}{c}^{\mathsf{T}}&0&0\\\
\nabla_{\mathbf{w}}{g}_{\mathbb{A}}^{\mathsf{T}}&0&0\end{bmatrix},\frac{\partial\varphi}{\partial
p}:=\begin{bmatrix}\nabla^{2}_{\mathbf{w}p}\mathcal{L}\\\
\nabla_{p}{c}^{\mathsf{T}}\\\
\nabla_{p}{g}_{\mathbb{A}}^{\mathsf{T}}\end{bmatrix}$

For the sake of notational simplicity we will drop the subscripts such that
$\nabla(\cdot)$ indicates $\nabla_{\mathbf{w}}(\cdot)$ by default, unless
otherwise specified explicitly. Substituting this in (4) and ignoring the
higher order terms, we get the linear predictor

${\color[rgb]{0,0,0}\mathbf{s}^{*}(p_{i}+\Delta
p_{j})\approx\hat{\mathbf{s}}(p_{i}+\Delta
p_{j})}=\mathbf{s}^{*}(p_{i})+\underbrace{\mathcal{H}_{i}\Delta
p_{j}}_{:=\Delta\mathbf{s}_{j}}$ (5)

which contains the inexact optimal action $\hat{u}_{j}=u_{i}^{*}+\Delta
u_{j}$. ∎

######  Remark 1 (Matrix factorization and inverse)

The main computation in the linear predictor involves getting
$\mathcal{H}_{i}$. Notice that the first term
$\frac{\partial\varphi}{\partial\mathbf{s}}$ is nothing but the KKT matrix,
which is already factorized when solving the NLP $\Pi(p_{i})$ using e.g. IPOPT
[27]. More importantly, augmenting any number of data samples using a single
data sample $(p_{i},u_{i}^{*})$ requires computing $\mathcal{H}_{i}$ only
once! That is, in the simplest case, using only a single NLP solve at $x_{i}$,
and a single linear solve (to compute $\mathcal{H}_{i}$), we can efficiently
augment several data samples $\Delta s_{j}=\mathcal{H}_{i}\Delta p_{j}$.

Figure 2: Schematic representation of the inexact solution manifold with
$N_{s}=3$ points where the NLP is solved exactly (denoted by red circle), and
parameter space $\mathcal{P}$ divided into the regions $\Delta\mathcal{P}_{i}$
around each $p_{i}$, which are used for data augmentation. Figure 3: Schematic
illustration of the proposed sensitivity-based data augmentation. Blue line
indicate the linear predictor step, and the orange arrows indicate the
corrector steps. The set $\mathcal{S}(p)$ is shown in gray.

To understand the effect of augmenting the data set with inexact samples,
consider the case where the true solution manifold of $\Pi(p)$ is given by
$\mathbf{s}^{*}(p)$ for all $p\in\mathcal{P}$, which is sampled at $N_{s}$
discrete points $\\{p_{i}\\}_{i=1}^{N_{s}}\in\mathcal{P}$. Assume that there
exists a piece-wise affine inexact solution manifold, denoted by
$\hat{\mathbf{s}}(p)$ for all $p\in\mathcal{P}$ given by the linear predictor
(5) within a neighborhood $\Delta\mathcal{P}_{i}$ around each $p_{i}$ for all
$i=1,\dots,N_{s}$. We make the following assumption regarding the neighborhood
$\Delta\mathcal{P}_{i}$.

###### Assumption 3

For all $i=1,\dots,N_{s}$, the neighborhood $\Delta\mathcal{P}_{i}$ is chosen
around each $p_{i}$ such that $p_{i}\in\langle\Delta\mathcal{P}_{i}\rangle$,
$\bigcup_{i=1}^{N_{s}}\Delta\mathcal{P}_{i}=\mathcal{P}$, and that the active
constraint set remains the same within each neighborhood
$\Delta\mathcal{P}_{i}$ , and the solution vector satisfies
$\|\mathbf{s}^{*}(p_{i}+\Delta p)-\mathbf{s}^{*}(p_{i})\|\leq L_{s}\|\Delta
p\|$.

We construct the idea of such a piece-wise affine inexact manifold
$\hat{\mathbf{s}}(p)$ given by the linear predictor (5), so that the augmented
data points are considered to be sampled from this inexact manifold. Learning
the policy then involves fitting a parametric function $\pi(x,\theta)$ to the
augmented training data set $\hat{u}(p)\subset\hat{\mathbf{s}}(p)$, such that

$\theta_{1}=\arg\min_{\theta}\sum_{i=1}^{N_{s}+M}\|\pi(p_{i};\theta)-\hat{u}(p_{i})\|^{2}$
(6)

###### Assumption 4

The functional form of $\pi(p;\theta)$ has sufficiently rich parametrization
and $\exists\;\hat{\theta}$ such that $\hat{u}(p)=\pi(p;\hat{\theta})$.

######  Theorem 2 ([25])

Consider a problem with the same setup as in Theorem 1, where the base data
set $\mathcal{D}^{0}$ with $N_{s}$ samples is obtained by querying $\Pi(p)$,
which is augmented with $M$ inexact samples from $\hat{u}(p)$ using (5). Under
Assumptions 3 and 4,

$\|\pi(p;\theta_{1})-\pi^{*}(p)\|\leq\max\left(\left\\{L_{p_{i}}R_{i}^{2}\right\\}_{i=1}^{N_{s}}\right)$
(7)

in probability as $M\rightarrow\infty$ holds if $\theta_{1}$ is a consistent
estimator of (6), with $R_{i}:=\sup_{\Delta p\in\Delta\mathcal{P}_{i}}\|\Delta
p\|$.

###### Proof:

Due to the continuity and differentiability of $\mathbf{s}^{*}(p_{i})$[26],
the following holds for all $\Delta p\in\Delta\mathcal{P}_{i}$,

$\displaystyle\|\hat{\mathbf{s}}(p_{i}+\Delta p)-\mathbf{s}^{*}(p_{i}+\Delta
p)\|\leq L_{p_{i}}\|\Delta p\|^{2}\leq L_{p_{i}}R_{i}^{2}$

Aggregating over all the regions results in the inequality

$\|\hat{\mathbf{s}}(p)-\mathbf{s}^{*}(p)\|\leq\max\left(\left\\{L_{p_{i}}R_{i}^{2}\right\\}_{i=1}^{N_{s}}\right)\quad\forall
p\in\mathcal{P}$

Since ${u}\subset{\mathbf{s}}(p)$, we have
$\|\hat{u}(p)-u^{*}(p)\|\leq\|\hat{\mathbf{s}}(p)-\mathbf{s}^{*}(p)\|$, which
quantifies the maximum deviation between the optimal policy and the augmented
inexact samples. If $\theta_{1}$ is a consistent estimator of (6), then
inequality (7) follows under Assumption 4. ∎

From this we can see that the error due to learning a policy based on
augmented data samples depends on the problem Lipschitz properties
$L_{p_{i}}$, and the size of the neighborhood $\Delta p_{i}$ used to augment
the data samples(cf. Fig. 2). If one were to use a large neighborhood $\Delta
p_{i}$ in order to reduce the number of offline NLP computations, then this
can potentially affect the performance of the learned policy.

Algorithm 1 Improved predictor-corrector sensitivity-based data augmentation.

1:$\Pi(p)$, $\mathcal{P}$, $\mathcal{D}=\emptyset$, $\epsilon_{tol}$

2:for $i=1,\dots,N_{s}$ do

3: Sample $p_{i}\in\mathcal{P}$

4: $\mathbf{s}^{*}(p_{i})\leftarrow$ Solve $\Pi(p_{i})$

5: Extract $u_{i}^{*}$ from the solution vector $\mathbf{s}^{*}(p_{i})$

6: $\mathcal{D}\leftarrow\mathcal{D}\cup\\{(p_{i},u_{i}^{*})\\}$

7:
$\mathcal{H}_{i}\leftarrow-\left[\frac{\partial\varphi(\mathbf{s}^{*}(p_{i}),p_{i})}{\partial\mathbf{s}}\right]^{-1}\frac{\partial\varphi(\mathbf{s}^{*}(p_{i}),p_{i})}{\partial
p}$

8: for $j=1,\dots,m$ do

9: Sample $\Delta p_{j}\in\Delta\mathcal{P}_{i}$ in the neighborhood of
$p_{i}$

10: $\hat{\mathbf{s}}(p_{i}+\Delta
p_{j})\leftarrow\mathbf{s}^{*}(p_{i})+\mathcal{H}_{i}\Delta{p}_{j}$
$\triangleright$ predictor

11: while $\|\nabla\mathcal{L}(\hat{\mathbf{s}}(p_{i}+\Delta
p_{j}))\|>\epsilon_{tol}$ do $\triangleright$ corrector

12: $\hat{\mathbf{s}}(p_{i}+\Delta
p_{j})\leftarrow\begin{bmatrix}\hat{\mathbf{w}}_{ij}\\\ 0\\\
0\end{bmatrix}-\mathcal{M}_{ij}^{-1}\mathcal{Q}_{ij}$

13: end while

14: Extract $\hat{u}_{j}^{*}$ from the solution vector
$\mathbf{\hat{s}}^{*}(p_{i}+\Delta{p}_{j})$

15: end for

16:end for

17:$\mathcal{D}$

###  II-B Data augmentation with user-enforced accuracy

In order to address this issue, we now propose an improved data augmentation
scheme, where in addition to the linear predictor, corrector steps are taken
to reduce the approximation error. Notice that the KKT conditions can
equivalently be expressed in the following matrix notation

$\begin{bmatrix}\nabla^{2}\mathcal{L}&\nabla{c}&\nabla{g}_{\mathbb{A}}\\\
\nabla{c}^{\mathsf{T}}&0&0\\\
\nabla{g}_{\mathbb{A}}^{\mathsf{T}}&0&0\end{bmatrix}\begin{bmatrix}0\\\
\lambda^{*}\\\ \mu_{\mathbb{A}}^{*}\end{bmatrix}+\begin{bmatrix}\nabla J\\\
c\\\ g_{\mathbb{A}}\end{bmatrix}$

Adding this with the predictor step results in

$\begin{bmatrix}\nabla^{2}\mathcal{L}&\nabla{c}&\nabla{g}_{\mathbb{A}}\\\
\nabla{c}^{\mathsf{T}}&0&0\\\
\nabla{g}_{\mathbb{A}}^{\mathsf{T}}&0&0\end{bmatrix}\begin{bmatrix}\Delta w\\\
\Delta\lambda+\lambda^{*}\\\
\Delta\mu_{\mathbb{A}}+\mu_{\mathbb{A}}^{*}\end{bmatrix}+\begin{bmatrix}\nabla
J\\\ c\\\
g_{\mathbb{A}}\end{bmatrix}+\begin{bmatrix}\nabla^{2}_{\mathbf{w}p}\mathcal{L}\\\
\nabla_{p}{c}^{\mathsf{T}}\\\
\nabla_{p}{g}_{\mathbb{A}}^{\mathsf{T}}\end{bmatrix}\Delta p=0$

Setting $\Delta p=0$ gives us the SQP correction step

$\begin{bmatrix}H&\nabla{c}&\nabla{g}_{\mathbb{A}}\\\
\nabla{c}^{\mathsf{T}}&0&0\\\
\nabla{g}_{\mathbb{A}}^{\mathsf{T}}&0&0\end{bmatrix}\begin{bmatrix}\Delta w\\\
\Delta\lambda+\lambda^{*}\\\
\Delta\mu_{\mathbb{A}}+\mu_{\mathbb{A}}^{*}\end{bmatrix}+\begin{bmatrix}\nabla
J\\\ c\\\ g_{\mathbb{A}}\end{bmatrix}=0$ (8)

Expanding $\hat{\mathbf{s}}(p_{i}+\Delta
p_{j}):=[\hat{\mathbf{w}}_{ij},\hat{\lambda}_{ij},\hat{\mu}_{ij}]^{\mathsf{T}}$,
the approximate solution is updated as

$\begin{bmatrix}\hat{\mathbf{w}}_{ij}\\\ \hat{\lambda}_{ij}\\\
\hat{\mu}_{ij}\end{bmatrix}=\begin{bmatrix}\hat{\mathbf{w}}_{ij}\\\ 0\\\
0\end{bmatrix}-{\underbrace{\begin{bmatrix}H&\nabla c&\nabla g_{\mathbb{A}}\\\
\nabla c^{\mathsf{T}}&0&0\\\ \nabla
g_{\mathbb{A}}^{\mathsf{T}}&0&0\end{bmatrix}}_{:=\mathcal{M}_{ij}}}^{-1}\underbrace{\begin{bmatrix}\nabla
J\\\ c\\\ g_{\mathbb{A}}\end{bmatrix}}_{:=\mathcal{Q}_{ij}}$ (9)

where $H$ (Hessian of the Lagrangian), $\nabla c$, $\nabla
g_{\mathbb{A}}$,$\nabla J$, $c$, and $g_{\mathbb{A}}$ that make up
$\mathcal{M}_{ij}$ and $\mathcal{Q}_{ij}$ are all evaluated at
$\hat{\mathbf{s}}(p_{i}+\Delta p_{j})$. The corrector steps are taken, until
the optimality residual defined by
$\|\nabla\mathcal{L}(\hat{\mathbf{s}}(p_{i}+\Delta p_{j}))\|$ is less than
some user defined tolerance $\epsilon_{tol}$.  If the active constraint set
$g_{\mathbb{A}}$ remains the same (cf. Assumption 3), then the corrector step
corresponds to a Newton direction on the KKT condition (cf. (8)). See [28] and
the references therein for detailed description of the corrector step. The
proposed data augmentation scheme with predictor-corrector steps is summarized
in Algorithm 1.

###### Theorem 3

Consider a problem with the same setup as in Theorem 1, where the base data
set $\mathcal{D}^{0}$ with $N_{s}$ samples is obtained by querying $\Pi(p)$,
which is augmented with $M=N_{s}m$ inexact samples from $\hat{\mathbf{s}}(p)$
using Algorithm 1. Under Assumption 4,

$\|\pi(p;\theta_{1})-\pi^{*}(p)\|\leq r(p)$ (10)

holds in probability as $M\rightarrow\infty$ if $\theta_{1}$ is a consistent
estimator of (6), where
$r(p):=\sup_{\mathbf{s}(p)\in\mathcal{S}(p)}\|\mathbf{s}(p)-\mathbf{s}^{*}(p)\|$,
with
$\mathcal{S}(p):=\\{\mathbf{s}(p)|\|\nabla\mathcal{L}(\hat{\mathbf{s}}(p))\|\leq\epsilon_{tol}\\}$.

###### Proof:

At the solution manifold $\mathbf{s}^{*}(p)$, we have
$\nabla\mathcal{L}(\mathbf{s}^{*}(p))=0$. From the continuity of the cost and
constraints,
$\exists\;\mathcal{S}(p):=\\{\mathbf{s}(p)\;|\;\|\nabla\mathcal{L}({\mathbf{s}}(p))\|\leq\epsilon_{tol}\\}$
containing the solution manifold $\mathbf{s}^{*}(p)$ in its interior (cf. Fig.
3) for all $p$. Solving the corrector steps until the bound on the optimality
residual satisfies
$\|\nabla\mathcal{L}(\hat{\mathbf{s}}(p))\|\leq\epsilon_{tol}$ (cf. line 10 in
Algorithm 1) implies that the $\hat{\mathbf{s}}(p)\in\mathcal{S}(p)$. Defining
the maximum distance from the exact solution manifold $\mathbf{s}^{*}(p)$ to
the boundary of the set $\mathcal{S}(p)$ for any $p$ as

$r(p):=\sup_{\mathbf{s}(p)\in\mathcal{S}(p)}\|\mathbf{s}(p)-\mathbf{s}^{*}(p)\|$

results in the inequality $\|\hat{\mathbf{s}}(p)-\mathbf{s}^{*}(p)\|\leq
r(p)$. Since ${u}\subset{\mathbf{s}}(p)$, we have
$\|\hat{u}(p)-u^{*}(p)\|\leq\|\hat{\mathbf{s}}(p)-\mathbf{s}^{*}(p)\|$.
Furthermore, from Assumption 4, we have that

$\displaystyle\|\hat{u}(p)-u^{*}(p)\|=$
$\displaystyle\|\pi(p;\hat{\theta})-\pi^{*}(p)\|\leq r(p)$

Our result follows, if $\theta_{1}$ is a consistent estimator of (6). ∎

The bound $r(p)$ clearly depends on $\epsilon_{tol}$, which can be controlled
by the user, as opposed to the bound in Theorem 2, which depends on the
distance from the original sample.

######  Remark 2 (Augmented samples)

The augmented samples can be approximated from the NLP sample, or another
previously augmented sample (similar to a path-following scheme). In either
case, corrector steps are taken until the optimality residual is less than
$\epsilon_{tol}$. Therefore, Theorem 3 is valid in both cases.

######  Remark 3 (Active constraint set)

Note that if a parturbation in $p$ changes the active constraint set
$g_{\mathbb{A}}$, this would induce non-smooth points in the solution manifold
$\mathbf{s}^{*}(p)$.  Capturing the non-smooth points in $\mathbf{s}^{*}(p)$
using piecewise-linear prediction manifolds requires solving a predictor-
corrector QP [29, 28]. However, in the context of data augmentation, a simpler
alternative could be to discard any samples $p_{i}+\Delta p$ that induce a
change in the active constraint set. In other words, Assumption 3 limits the
size of the neighborhood around which additional samples can be augmented.

TABLE I: Example 1 - Comparison of the total CPU time required to generate the different data sets, and the maximum error.  | full NLP | predictor only | predictor-corrector
---|---|---|---
CPU time [s] | CPU time [s] | Max Error | CPU time [s] | Max Error
Case 1 | 466.42 | 4.92 | 1.34 | 11.95 | 0.0074
Case 2 | 456.21 | 0.45 | 12.51 | 8.67 | 0.028
Case 3 | 542.73 | 0.15 | 79.88 | 12.96 | 0.018

(a)

(b)

Figure 4: Example 1 - (a) Data set with optimal state-action data pairs for
the inverted pendulum example. Red circles indicate samples where the MPC
problem is solved exactly, and gray dots indicate samples that are augmented
using only the linear predictor [25] (left subplots), and the proposed
approach using additional corrector steps (right subplots). (b) Closed-loop
trajectories of the approximate policy trained with augmented data set using
only the linear predictor [25] (red dotted lines in left subplots), and with
augmented data set using proposed approach (red dashed lines in right
subplots), corresponding to the data sets shown in Fig. 4(a), which are
benchmarked against the true MPC policy (solid black lines) for various
starting points.

##  III Numerical experiments

###  III-A Example 1: Inverted Pendulum

We demonstrate the performance of the data augmentation framework using the
benchmark inverted pendulum problem with the nonlinear dynamics
$ml^{2}\ddot{\omega}=u-b\dot{\omega}-mgl\sin\omega$ described by the two
states, angle $\omega$, and angular velocity $\dot{\omega}$. We consider the
case where the expert policy is given by a nonlinear model predictive
controller with an objective to drive the pendulum to its inverted position.
The input $u$ is the torque, which is used to maintain the pendulum at its
inverted position of $\omega=3.14$ rad, $\dot{\omega}=0$ . In this example, we
have $p=[\omega,\dot{\omega}]^{\mathsf{T}}$ and
$\mathcal{P}=[0,2\pi]\times[-5,5]$. Note that we deliberately choose this
simple example to facilitate visualization of the exact and inexact solution
manifolds to illustrate the data augmentation schemes.

To approximate the MPC policy, we wish to generate $100\times 100$ state-
action data pairs using a grid-based sampling strategy to evenly cover the
parameter space $\mathcal{P}$, which would normally take $10^{4}$ offline NLP
computations. In each of the following cases, we also solve the full NLP to
compute $\pi^{*}(p)$ at each sample to serve as a benchmark. The NLP problems
$\Pi(p)$ are solved using IPOPT [30]with MUMPS linear solver. The optimization
problem and the NLP sensitivities were formulated using CasADi v3.5.1 [31].
All augmented samples are based on the sensitivity updates from the exact
sample obtained by solving the NLP. All computations were performed on a 2.6
GHz processor with 16GB memory.

#### Case 1: Solve 100 offline NLP problems

In this case, we queried the NMPC expert to generate $N_{s}=100$ samples using
a $10\times 10$ grid as shown in Fig. 4(a) in red circles. Using each data
sample, we further augment additional 100 data samples around each $p_{i}$,
leading to a total of $M=$ 104 augmented data samples using only the linear
predictor as done in [25], as well as the proposed improved predictor-
corrector data augmentation scheme. These are shown in Fig. 4(a) in gray dots,
where it can be seen that both the data augmentation schemes are able to
sufficiently capture the solution manifold. The maximum error due to
approximation when using only the linear predictor is 1.34, whereas the
maximum error was 0.0074 with the improved data augmentation.

#### Case 2: Solve 9 offline NLP problems

We then consider the case where we queried the NMPC expert only $N_{s}=9$
times using a $3\times 3$ grid as shown in Fig. 4(a) in red circles. Using
each data sample, we further augment additional 1089 data samples around each
sample, leading to a total of 9081 augmented data samples using only $9$
offline NLP computations. Since the size of the neighborhood increases, the
inexact solution manifold using only the linear predictor is now visbily
different from the exact solution manifold (with 9 piecewise affine
manifolds), with the maximum error of 12.5. On the other hand, when using the
improved data augmentation scheme, the maximum error is only 0.028. This shows
that despite the large neighborhood used for data augmentation, we can get the
generate data samples with a desired accuracy at only marginally increased
computational cost (cf. Table I).

#### Case 3: Solve only 1 offline NLP problem

We then consider the extreme case, where we only solve $N_{s}=1$ offline NLP
problem as shown in Fig. 4(a). Using this single NLP computation, we augment
10000 data samples using a $100\times 100$ grid. Using only the linear
predictor, naturally we get a single affine manifold, leading to large
approximation errors. However, using the improved data augmentation scheme, we
are still able to augment data samples with a desired accuracy, where the
maximum error is only 0.018, as opposed to 79.88.

The total CPU time to generate the data samples using the two data
augmentation schemes, as well as solving all the data points exactly along
with the maximum error are summarized in Table I, which also shows the trade-
off between computation cost and accuracy. We then learn the approximate
policy using the data sets from the six different cases shown in Fig. 4(a).
For the approximate policies in each case, we choose a generalized regression
neural network (a variant of the radial basis function network) that was
trained using the newgrnn function in  MATLAB v2020b. The hyperparameters of
$\pi(p,\theta)$ were the same across all the cases, and the only variation was
the data set used for training. The closed-loop trajectories of the
approximate policies compared to the true MPC policy starting from different
initial conditions are shown in Fig. 4(b). Here, red dotted lines denote the
closed-loop trajectories of the the approximate policy learned using augmented
data set with only the linear predictor, red dashed lines denote the closed-
loop trajectories of the approximate policy learned using augmented data set
using the proposed predictor-corrector based data augmentation, summarized in
Algorithm 1, and solid black lines denote the true MPC policy that serves as a
benchmark. These results clearly show that the control policy learned using
only the augmented data set with linear predictor steps results in reasonable
performance if the neighborhood used for data augmentation is small. As the
size of the neighborhood used for data augmentation increases, the
approximation error increases (cf. Theorem 2), resulting in poorer closed-loop
performance. However, using the proposed data augmentation scheme with user
enforced accuracy, we are able to learn a satisfactory control policy, even
when a large neighborhood is used for data augmentation. This clearly
demonstrates the benefit of using the improved data augmentation scheme, which
enables us to use a larger neighborhood for data augmentation at the cost of a
few additional linear solves (corrector steps).

###  III-B Example 2: Quadcopter control

The proposed data augmentation method is also applied to a quadcopter control
problem, results from which are briefly summarized here. The quadcopter is
modelled as shown in [32]. The system has $x\in\mathbb{R}^{12}$ states and
$u\in\mathbb{R}^{4}$ inputs. A nonlinear model predictive controller is
designed to control the quadcopter to a desired target $x_{target}$ described
by an xyz-coordinate. The NMPc problem has a stage cost
$\ell(x,u))\|x-x_{target}\|_{Q}^{2}+\|u\|_{R}^{2}+\|\Delta u\|_{S}^{2}$. The
goal is to train a control policy $\pi(p_{i};\theta)$ that takes as input the
current states, control action, and the desired target, to compute the next
action. That is, in this case $p\in\mathbb{R}^{19}$. For this high dimensional
system, the goal is to generate training data points using the proposed data
augmentation scheme, which can be used to train a function approximator
$\pi:\mathbb{R}^{19}\rightarrow\mathbb{R}^{4}$.

$N_{s}=20$ random samples are generated, where the MPC problem is solved
exactly. We further randomly sample m=500 data points with ±20% variation
around each exact MPC sample, where the data is augmented from the MPC
solution using the proposed method. By doing so, we generated a total of 10020
samples. For the predictor-corrector data augmentation, the corrector steps
were used until the optimality residual was less than a tolerance of
$\epsilon_{tol}=1e-2$.

Fig. 5(a) shows the approximation error $\|\hat{u}(p_{i})-u^{*}(p_{i})\|$ when
using predictor-only data augmentation [25] (shown in blue), and the
predictor-corrector data augmentation (shown in red). Their corresponding
computation times are shown in Fig. 5(b), compared against generating all the
samples exactly using the full NLP (yellow). To generate all the 10020 samples
exactly required a total CPU time of 850.3s, whereas predictor-corrector data
augmentation took 100.3s, and predictor-only data augmentation took 2.3s. This
clearly shows the trade-off between the computation cost and the accuracy of
the augmented samples.

(a)

(b)

Figure 5: Example 2 - (a) Approximation error
$\|\hat{u}(p_{i})-u^{*}(p_{i})\|$ when using predictor-only data augmentation
(blue), and predictor-corrector data augmentation (red). (b) CPU times for
generating the full NLP data set (yellow), predictor-only augmented data set
(blue), and predictor-corrector augmented data set (red). These two plots
clearly show the trade-off between accuracy and the computation cost when
generating training data to learn the MPC policy.

##  IV Discussion

In this section, we provide some discussions on the placing our work on data
augmentation in the context of other related literature.

###  IV-A Data augmentation and sampling techniques

As mentioned in the introduction, generating large data sets that can be used
for training and testing is critical for scaling MPC policy approximations.
This challenge, if at all considered in the current works, is mainly addressed
via efficient sampling techniques. This warrants a discussion on how our
proposed data augmentation scheme fits in the context of the different
sampling strategies that have been proposed in the MPC literature.

#### Adaptive sampling using resolution criteria

A recursive sampling algorithm was proposed in [33], where new sample
locations are chosen with higher resolution as long as neighboring values are
still differing. At each sample location, the optimization problem is solved
exactly to get the corresponding optimal action. One can use the same
algorithm as in [33] to decide where to sample the state-space, and if the
newly generated sample locations are within a neighborhood of a previously
computed sample satisfying the assumptions of Theorem 1, then one can use the
proposed data augmentation instead of solving the exact NLP at the newly
sampled states.

#### Sampling based on geometric random walks

The authors in [19] recently proposed an algorithm to efficiently generate
large data sets based on geometric random walks instead of independent
sampling, where starting from a feasible state, small steps are taken
iteratively along a random line until each step is feasible. Such an algorithm
naturally lends itself to our proposed data augmentation scheme, since this
algorithm is based on the fact that new samples are generated based on
previous successful samples. At every new step taken along the random walk,
the corresponding optimal action can can be approximated using the previous
sample, instead of solving it exactly.

#### Control-oriented sampling

Based on the premise that it is unlikely that all states from the feasible
region are equally likely to be observed during closed-loop operation, control
oriented sampling strategies have been proposed in [34]. In this approach,
starting from a feasible initial state, the consecutive samples are selected
based on noisy closed-loop simulations, such that the samples are selected
from a tube of closed loop trajectories likely to be observed in operation.
Again, such an algorithm naturally lends itself to our proposed data
augmentation scheme, since the new samples are generated based on previous
successful samples.

In general, the proposed data augmentation technique is complementary to the
different sampling strategies, in the sense that, once the sample location is
chosen using any sampling technique, instead of solving the optimization
problem exactly at the chosen sample location, the proposed data augmentation
technique can be used to efficiently approximate the corresponding optimal
action with a desired accuracy at the given sample location. Given that the
main motivation for developing different sampling strategies is to scale MPC
policy approximation to larger systems, combining these developments with our
proposed data augmentation scheme would enhance this goal. As such, the data
augmentation scheme proposed in this paper complements the developments on
sampling strategies for offline training.

###  IV-B Data augmentation for a broader class of learning from demonstration
problems

The cost of generating sufficiently rich training data set is also an open
challenge in a broad class of learning from demonstrations, such as imitation
learning with interactive expert, and inverse optimal control. In this
subsection, we briefly describe how the proposed data augmentation method can
be used with a broad class of learning from demonstrations.

#### Imitation learning with interactive expert

Imitation learning with interactive expert is similar to the direct policy
approximation approach described above, but the main difference is that we
have access to an interactive expert in-the-loop, who we can query in order to
improve our policy online [35, 36, 20, 37]. For example, consider a system
where we have an approximate policy running online, but we have an MPC
controller running in the background that provides feedback on the optimal
action for the state visited by the system.

That is, we first train an initial policy $\pi_{k}$ using the labeled data set
$\mathcal{D}_{k}:=\\{(x_{i},u_{i}^{*})\\}$. The current policy $\pi_{k}$ is
then rolled out to collect the trajectory
$\\{x_{t},\pi_{k}(x_{t})\\}_{t=1}^{T}$. For each trajectory, an interactive
expert also provides the optimal policy for the states visited by the current
roll-out (called interactive feedback), which are aggregated such that
$\mathcal{D}_{k+1}=\mathcal{D}_{k}\cup\\{(x_{t},u^{*}_{t})\\}_{t=1}^{T}$ for
$k=1,2,\dots$, and the policy function is updated on-the-fly by learning from
the aggregated data set $\mathcal{D}_{k+1}$.

In the case of an interactive expert that provides feedback on the rolled out
policy (i.e.what the expert would have done in the states visited by the
current policy), the proposed data-augmentation framework can be used to
augment additional samples around each state where the expert provides
feedback. The “augmented” interactive feedback can then be used to improve the
policy online. By doing so, we posit that this would enable us to generalize
the feedback provided by the interactive expert to other nearby states not
visited by the current policy (i.e. using the expert feedback, infer what the
expert would have done in other nearby states not visited by the current
policy).

The use of sensitivity-based data augmentation for learning to control an
inverted pendulum using imitation learning with interactive expert can be
found in Appendix A.

#### Inverse Optimal Control

An alternative approach to direct policy approximation is to use the noisy
observations of the optimal policy to compute the cost function of a simpler
optimization problem that is consistent with the expert demonstrations
$\mathcal{D}:=\\{(x_{i},u_{i}^{*})\\}_{i=1}^{N}$. That is, we assume that
$\mathcal{D}:=\\{(x_{i},u_{i}^{*})\\}_{i=1}^{N}$ comes from a simpler
optimization problem, and we want to learn the cost function of this simpler
optimization problem that is consistent with the demonstrations. By doing so,
the policy is given by the simpler optimization problem, instead of the
original complex optimization problem. For example, a framework for
approximating a complex MPC problem with a simple optimizing controller using
inverse optimization was shown in [38]. The proposed data augmentation scheme
can also be used in such frameworks.

#### Value-function approximation

Instead of approximating the policy function directly, one can obtain the
optimal value function $V^{*}(x_{i})$ from the expert queries, and approximate
the optimal value function using the optimal state-value pairs
$\\{(x_{i},V^{*}(x_{i}))\\}$. The function approximator can then be used as
the optimal cost-to-go function in simple myopic optimizing controllers in the
context of approximate dynamic programming (ADP), much like in [38](where
instead of using a quadratic cost function, any function approximator can be
used). The proposed data augmentation framework can not only augment the
optimal actions, but can also be used to augment the optimal state-value pairs
based on a few expert queries. To do this, (1) is solved for various state
realizations $\\{x_{i}\\}$, and the corresponding optimal value function
$V^{*}(x_{i})$ is collected as the training data set. After each offline MPC
solve, additional samples at $x_{i}+\Delta x_{j}$ can be augmented using our
proposed data augmentation scheme. The augmented data set
$\\{(x_{i},V^{*}(x_{i}))\\}$ can then be used to fit a parametric function
approximator that predicts the optimal value function $V^{*}(x)$ for any given
$x$. This would be a better approach than learning the policy directly, since
in this approach, we solve a much simpler and easier optimization problem, and
by doing so, we can have feasibility guarantees by explicitly enforcing the
constraints in the (simpler) optimization problem.

##  V Conclusion

To summarize, this paper presented an improved data augmentation scheme for
cheaply generating the training data samples for MPC policy approximation
using only a fraction of the computation effort. We showed that by using
additional corrector steps we can enforce a desired level of accuracy in the
augmented samples (cf. Algorithm 1), which amounts to a few additional linear
solves. By doing so the approximation error does not depend on the size of the
neighborhood used for data augmentation, but on the user-defined optimality
residual (cf. Theorem 3), which allows us to augment samples from larger
neighborhoods while without jeopardizing accuracy. This was demonstrated on a
inverted pendulum control problem, which clearly shows the trade-off between
accuracy and computational effort. The proposed data augmentation approach can
be used with any sampling strategy, as well as for a broad class of learning
from demonstration problems including imitation learning with interactive
expert, value function approximation, and inverse optimal control.

## References

  * [1] A. Bemporad, M. Morari, V. Dua, and E. N. Pistikopoulos, “The explicit linear quadratic regulator for constrained systems,” _Automatica_ , vol. 38, no. 1, pp. 3–20, 2002.
  * [2] T. Parisini and R. Zoppoli, “A receding-horizon regulator for nonlinear systems and a neural approximation,” _Automatica_ , vol. 31, no. 10, pp. 1443–1451, 1995.
  * [3] S. Chen, K. Saulnier, N. Atanasov, D. D. Lee, V. Kumar, G. J. Pappas, and M. Morari, “Approximating explicit model predictive control using constrained neural networks,” in _2018 Annual American control conference (ACC)_. IEEE, 2018, pp. 1520–1527.
  * [4] M. Hertneck, J. Köhler, S. Trimpe, and F. Allgöwer, “Learning an approximate model predictive controller with guarantees,” _IEEE Control Systems Letters_ , vol. 2, no. 3, pp. 543–548, 2018.
  * [5] X. Zhang, M. Bujarbaruah, and F. Borrelli, “Safe and near-optimal policy learning for model predictive control using primal-dual neural networks,” in _2019 American Control Conference (ACC)_. IEEE, 2019, pp. 354–359.
  * [6] J. Drgoňa, D. Picard, M. Kvasnica, and L. Helsen, “Approximate model predictive building control via machine learning,” _Applied Energy_ , vol. 218, pp. 199–216, 2018.
  * [7] B. Karg and S. Lucia, “Efficient representation and approximation of model predictive control laws via deep learning,” _IEEE Transactions on Cybernetics_ , 2020.
  * [8] J. A. Paulson and A. Mesbah, “Approximate closed-loop robust model predictive control with guaranteed stability and constraint satisfaction,” _IEEE Control Systems Letters_ , 2020.
  * [9] C. Sánchez-Sánchez, D. Izzo, and D. Hennes, “Learning the optimal state-feedback using deep networks,” in _2016 IEEE Symposium Series on Computational Intelligence (SSCI)_. IEEE, 2016, pp. 1–8.
  * [10] P. Kumar, J. B. Rawlings, and S. J. Wright, “Industrial, large-scale model predictive control with structured neural networks,” _Computers & Chemical Engineering_, vol. 150, p. 107291, 2021.
  * [11] B. M. Åkesson and H. T. Toivonen, “A neural network model predictive controller,” _Journal of Process Control_ , vol. 16, no. 9, pp. 937–946, 2006.
  * [12] Y. Cao and R. B. Gopaluni, “Deep neural network approximation of nonlinear model predictive control,” _IFAC-PapersOnLine_ , vol. 53, no. 2, pp. 11 319–11 324, 2020.
  * [13] X. Zhang, M. Bujarbaruah, and F. Borrelli, “Near-optimal rapid mpc using neural networks: A primal-dual policy learning framework,” _IEEE Transactions on Control Systems Technology_ , vol. 29, no. 5, pp. 2102–2114, 2020\.
  * [14] Y. S. Quan and C. C. Chung, “Approximate model predictive control with recurrent neural network for autonomous driving vehicles,” in _2019 58th Annual Conference of the Society of Instrument and Control Engineers of Japan (SICE)_. IEEE, 2019, pp. 1076–1081.
  * [15] B. Karg and S. Lucia, “Approximate moving horizon estimation and robust nonlinear model predictive control via deep learning,” _Computers & Chemical Engineering_, p. 107266, 2021.
  * [16] J. Nubert, J. Köhler, V. Berenz, F. Allgöwer, and S. Trimpe, “Safe and fast tracking on a robot manipulator: Robust mpc and neural network control,” _IEEE Robotics and Automation Letters_ , vol. 5, no. 2, pp. 3050–3057, 2020.
  * [17] A. D. Bonzanini, J. A. Paulson, D. B. Graves, and A. Mesbah, “Toward safe dose delivery in plasma medicine using projected neural network-based fast approximate NMPC,” _IFAC-PapersOnLine_ , vol. 53, no. 2, pp. 5279–5285, 2020.
  * [18] B. Karg, T. Alamo, and S. Lucia, “Probabilistic performance validation of deep learning-based robust NMPC controllers,” _International Journal of Robust and Nonlinear Control_ , vol. 31, no. 18, pp. 8855–8876, 2021.
  * [19] S. W. Chen, T. Wang, N. Atanasov, V. Kumar, and M. Morari, “Large scale model predictive control with neural networks and primal active sets,” _Automatica_ , vol. 135, p. 109947, 2022.
  * [20] S. Ross, G. Gordon, and D. Bagnell, “A reduction of imitation learning and structured prediction to no-regret online learning,” in _Proceedings of the fourteenth international conference on artificial intelligence and statistics_. JMLR Workshop and Conference Proceedings, 2011, pp. 627–635.
  * [21] M. Palan, S. Barratt, A. McCauley, D. Sadigh, V. Sindhwani, and S. Boyd, “Fitting a linear control policy to demonstrations with a kalman constraint,” in _Learning for Dynamics and Control_. PMLR, 2020, pp. 374–383.
  * [22] T. Tran, T. Pham, G. Carneiro, L. Palmer, and I. Reid, “A bayesian data augmentation approach for learning deep models,” in _Advances in neural information processing systems_ , 2017, pp. 2797–2806.
  * [23] L. Taylor and G. Nitschke, “Improving deep learning with generic data augmentation,” in _2018 IEEE Symposium Series on Computational Intelligence (SSCI)_. IEEE, 2018, pp. 1542–1547.
  * [24] C. Shorten and T. M. Khoshgoftaar, “A survey on image data augmentation for deep learning,” _Journal of Big Data_ , vol. 6, no. 1, pp. 1–48, 2019.
  * [25] D. Krishnamoorthy, “A sensitivity-based data augmentation framework for model predictive control policy approximation,” _IEEE Transactions on Automatic Control_ , vol. In-Press, 2021.
  * [26] A. V. Fiacco, “Sensitivity analysis for nonlinear programming using penalty methods,” _Mathematical programming_ , vol. 10, no. 1, pp. 287–311, 1976\.
  * [27] H. Pirnay, R. López-Negrete, and L. T. Biegler, “Optimal sensitivity based on IPOPT,” _Mathematical Programming Computation_ , vol. 4, no. 4, pp. 307–331, 2012.
  * [28] V. Kungurtsev and J. Jäschke, “A predictor-corrector path-following algorithm for dual-degenerate parametric optimization problems,” _SIAM Journal on Optimization_ , vol. 27, no. 1, pp. 538–564, 2017.
  * [29] J. F. Bonnans and A. Shapiro, “Optimization problems with perturbations: A guided tour,” _SIAM review_ , vol. 40, no. 2, pp. 228–264, 1998.
  * [30] A. Wächter and L. T. Biegler, “On the implementation of an interior-point filter line-search algorithm for large-scale nonlinear programming,” _Mathematical programming_ , vol. 106, no. 1, pp. 25–57, 2006.
  * [31] J. A. E. Andersson, J. Gillis, G. Horn, J. B. Rawlings, and M. Diehl, “CasADi – A software framework for nonlinear optimization and optimal control,” _Mathematical Programming Computation_ , vol. 11, no. 1, pp. 1–36, 2019.
  * [32] A. Bemporad, C. A. Pascucci, and C. Rocchi, “Hierarchical and hybrid model predictive control of quadcopter air vehicles,” _IFAC Proceedings Volumes_ , vol. 42, no. 17, pp. 14–19, 2009.
  * [33] J. Nubert, “Learning-based approximate model predictive control with guarantees - joining neural networks with recent robust MPC,” Master’s thesis, ETH Zürich, 2019.
  * [34] D. Krishnamoorthy, A. Mesbah, and J. A. Paulson, “An adaptive correction scheme for offset-free asymptotic performance in deep learning-based economic MPC,” _IFAC-PapersOnLine_ , vol. 54, no. 3, pp. 584–589, 2021.
  * [35] B. D. Argall, B. Browning, and M. Veloso, “Learning robot motion control with demonstration and advice-operators,” in _2008 IEEE/RSJ International Conference on Intelligent Robots and Systems_. IEEE, 2008, pp. 399–404.
  * [36] S. Chernova and M. Veloso, “Interactive policy learning through confidence-based autonomy,” _Journal of Artificial Intelligence Research_ , vol. 34, pp. 1–25, 2009.
  * [37] S. Jauhri, C. Celemin, and J. Kober, “Interactive imitation learning in state-space,” _arXiv preprint arXiv:2008.00524_ , 2020.
  * [38] A. Keshavarz, Y. Wang, and S. Boyd, “Imputing a convex objective function,” in _2011 IEEE international symposium on intelligent control_. IEEE, 2011, pp. 613–619.

###  -A Direct policy approximation with interactive expert

The proposed data augmentation framework is applied in the context of
imitation learning with interactive expert on the same inverted pendulum
example. The objective here is to learn a policy function that would bring the
pendulum from rest position ($\omega=\dot{\omega}=0$) to its inverted position
($\omega=3.14,\dot{\omega}=0$).

#### Task

Learn to invert a pendulum from rest.

#### Initial policy

Linear policy $u=-Kx$, with a$K$ that is not able to accomplish the task.

#### Learning framework

We first start with no prior expert demonstrations, but with an initial policy
that is rolled out on the system. For this, we chose a linear policy
$u=-11\omega-7\dot{\omega}+35$. For each state visited in the rollout, the
expert provides feedback on what the expert would have done in these states.
Around each expert feedback, we further augment 25 samples using the proposed
data augmentation framework.

We first learn a policy function using only the expert feedback. To
demonstrate the advantage of augmenting additional samples inferred from the
expert feedback, we then learn a policy function based on the expert feedback
and the augmented samples. For both these policies, we choose a generalized
regression neural network that was trained using the newgrnn function in
MATLAB v2020b. The newly obtained approximate policy is then rolled out and
this procedure is repeated for 25 times.

Fig. 6 left subplots shows the states that are visited in each rollout, where
the expert provides feedback (shown in red circles). As it can be seen, even
after 25 rollouts the number of demonstrations in the state-space is not much,
which can affect the quality of the policy learned using only the expert
feedback. Whereas, by augmenting additional samples around each expert
feedback, we can cheaply generate several data samples covering a wider state-
space than just what was visited in the rollouts. This is shown in black dots
in Fig. 6 right subplots. (From the 3D plot, one can already visually
recognize the shape of the optimal solution manifold when we augment
additional data samples).

Figure 6: Closed loop performance over 15 rollouts when using only the expert
feedback (shown in red), and using botht he expert and augmented feedback
(shown in black). The bold lines shows the performance at the 15th rollout.
Figure 7: Closed-loop performance of all the rollouts using only expert
feedback (left subplot) and using expert and augmented feedback (right
subplot). Table shows whether the task in a particular rollout was successful
or not.

The performance of the policy learned using only the expert feedback or using
the expert and the augmented feedback for all the 25 rollouts are shown in
Fig. 7. Here it can be seen that by augmenting additional samples, we are able
to better learn the policy, and hence learn to control the inverted pendulum
better than if we only use the expert feedback. This is because, after each
rollout, we not only learn from the expert feedback, but also infer what the
expert would have done in other states not visited in the current rollout.