# Electronic Correlations and Absence of Superconductivity in the Collapsed
Phase of LaFe2As2

Jianzhou Zhao<EMAIL_ADDRESS>Co-Innovation Center for New Energetic
Materials, Southwest University of Science and Technology, Mianyang, 621010,
People’s Republic of China  Research Laboratory for Quantum Materials,
Singapore University of Technology and Design, Singapore 487372, Singapore
Yilin Wang<EMAIL_ADDRESS>Hefei National Laboratory for Physical
Sciences at Microscale, University of Science and Technology of China, Hefei,
Anhui 230026, China  Xiaolong Feng  Research Laboratory for Quantum Materials,
Singapore University of Technology and Design, Singapore 487372, Singapore
Shengyuan A. Yang  Research Laboratory for Quantum Materials, Singapore
University of Technology and Design, Singapore 487372, Singapore  Center for
Quantum Transport and Thermal Energy Science, School of Physics and
Technology, Nanjing Normal University, Nanjing 210023, China

###### Abstract

The interplay between structural phase and electronic correlations has been an
intriguing topic of research. An prominent example is the pressure-induced
uncollapsed to collapsed tetragonal phase transition observed in CaFe2As2,
which is accompanied with the emergence of superconductivity in the collapsed
phase. Recently, a very similar structural phase transition was discovered in
LaFe2As2, but in contrast to CaFe2As2, superconductivity was only observed in
the uncollapsed phase, not the collapsed phase. Previous studies have
attributed this puzzling observation to the differences in the two materials’
band coherence, orbital occupation, and Fermi surface topology. Here, we
present a comparative study of LaFe2As2 and CaFe2As2 using the DFT+DMFT
method. Surprisingly, we find that although La appears to have a valence
higher than Ca, the doped one electron actually primarily resides on the La
site. This leads to almost the same total Fe-$3d$ occupancy and electronic
correlation strength as well as similar Lifshiftz transition in the Fermi
surface topology for the two materials. In addition, we show that the two
materials in both structural phases belong to the category of Hund’s metals.
Our results indicate that the electronic structures of LaFe2As2 and CaFe2As2
are not too different, which further suggest that superconductivity might also
be induced in the collapsed phase of LaFe2As2 under similar non-hydrostatic
conditions as for CaFe2As2.

††preprint: Preprint: , .

Introduction. Electronic correlations, Fermi-surface topology and magnetism
play important roles in unconventional iron-based superconductor materials [1,
2, 3, 4, 5]. It is extremely intriguing that all three aspects can be
dramatically changed during certain structural phase transitions. A prominent
example is CaFe2As2 (Ca122), which exhibits a pressure-induced “uncollapsed”
tetragonal (UT) to “collapsed” tetragonal (CT) structural phase transition [6,
7, 8, 9, 10, 11]. At ambient conditions, Ca122 is in the UT phase with the
ThCr2Si2-type lattice structure, and directly lowering the temperature below
170 K drives it into an orthorhombic phase, with a concomitant transition from
paramagnetic to antiferromagnetic phase [12, 13]. Interestingly, upon applying
a pressure, the magnetism can be suppressed, and the UT phase transits into
the iso-structural CT phase which maintains the same crystal symmetry $I4/mmm$
but has a strong reduction of the $c$-axis by $\sim 8$% [14, 15, 7]. More
importantly, superconductivity was observed in the CT phase with $T_{c}$ of 12
K at 0.3 GPa [16, 13, 17, 18].

Recently, a similar UT to CT phase transition was reported in LaFe2As2 (La122)
[19]. However, unlike Ca122, superconductivity (with $T_{c}\sim 12$ K) was
observed in the UT phase, but not in the CT phase. The puzzle of the absence
of superconductivity in CT-La122 has attracted significant attention [20, 21].
Formally, La3+ has a higher valence than Ca2+, hence one may tend to attribute
the difference to the different doping levels of Fe $d$-orbitals, namely,
La122 may be viewed as 0.5 electron doping from Ca122 (Fe2+), resulting a
nominal oxidation state with Fe1.5+ [19]. Particularly, in Ref. [22], via a
comparative study of La122 and Ca122 using _ab initio_ QSGW+DMFT techniques,
it was found that CT-La122 exhibits much enhanced coherence of its electronic
structure than CT-Ca122, and this distinction was argued to underly the
absence/presence of superconductivity in CT-La122/Ca122. However, it is noted
that the $c$ lattice parameter used for CT-Ca122 in Ref. [22] is only 1%
smaller than UT-Ca122, larger than the experimental value by $7\sim 8\%$ [8,
18]. This is a sizable difference. Hence, when taking experimental lattice
parameters, whether the comparison result still holds remains an open
question.

In this work, we address this question by a comparative study of the
electronic structures in both UT and CT phases of La122 and Ca122, using the
DFT+DMFT method with an “exact” double-counting scheme. The key finding is
that despite the different nominal valences of La and Ca, La122 and Ca122 are
_not_ too different in terms of their normal-state electronic structures.
Particularly, the distinction between CT-La122 and CT-Ca122 regarding the
coherence is not observed here. We show that both material systems have
similar electronic correlation strengths and similar Lifshitz transitions of
Fermi surface topology. These observations can be understood by noting that
the extra electron doped by La in fact is largely retained at the La site
rather than moving to the Fe-As layers that dominate the low-energy physics
and (possible) superconductivity. In addition, we show that both materials
belong to the category of Hund’s metals where the Hund’s coupling dominates
the correlation effects [23, 24, 25, 26, 27, 28, 29]. Since it was reported
that superconductivity in CT-Ca122 is very sensitive to pressure and occurs
only under non-hydrostatic conditions [30], our result implies that
superconductivity may also be induced in CT-La122 under similar conditions.

Table 1:  Lattice parameters $a$, $c$ and the As-Fe-As angles $\alpha_{\text{As-Fe-As}}$ used in this work [19, 8].  | Lattice Parameters
---|---
$a$ (Å) |  $c$ (Å) | $\alpha_{\text{As-Fe-As}}$
UT-La122 | 3.938 | 11.732 | 110.8$\degree$
CT-La122 | 4.004 | 11.014 | 118.1$\degree$
UT-Ca122 | 3.903 | 11.591 | 110.6$\degree$
CT-Ca122 | 3.982 | 10.684 | 115.9$\degree$

Methods. We perform fully charge self-consistent DFT+DMFT calculations using
the EDMFTF software package [31], based on the full-potential linear augmented
plane-wave method implemented in the WIEN2k code [32, 33]. A rotationally
invariant form of local on-site Coulomb interaction Hamiltonian parameterized
by Hubbard $U$ and Hund’s coupling $J_{H}$ is applied on all five Fe-$3d$
orbitals. We choose $U=5.0$ eV and $J_{H}=0.8$ eV in this work, which are
typical values for iron based superconductors [26, 34, 35]. The impurity
problem is solved by the hybridization expansion version of the continuous-
time quantum Monte Carlo solver [36]. We choose an “exact” double counting
scheme developed by Haule [37], which eliminates the double counting issue in
correlated materials. The self-energy on real frequency is obtained by the
analytical continuation method of maximum entropy [31]. We use the
experimental lattice parameters for both of the UT and CT phases of La122 [19]
and Ca122 [8], which are listed in Table 1 (also see Supplemental Material
(SM) [38]). The effective mass enhancement by correlations is defined by
$m^{*}/m_{\text{DFT}}=1/\mathcal{Z}$, where $\mathcal{Z}$ is the quasi-
particle weight. We directly obtain
$\mathcal{Z}=1-\frac{\partial\text{Im}\Sigma(\mathrm{i}\omega_{n})}{\partial\omega_{n}}\big{|}_{\omega_{n}\rightarrow
0^{+}}$ from the self-energy on Matsubara frequencies, to avoid large error-
bar in analytic continuations. More computation details are presented in SM
[38].

Figure 1: The DFT band structures for UT and CT phases of La122 and Ca122. Fat bands with Fe-$3d_{xy}$ and La-$5d$/Ca-$3d$ characters are shown in blue and red colors, respectively. The black arrows indicate the hybridization gaps between La-$5d$ and Fe-$3d$ bands. Table 2:  (left) Total occupancy numbers of Fe-$3d$, La-$5d$/Ca-$3d$ and As-$4p$ orbitals, obtained from DFT and DFT+DMFT calculations at $T=300$ K. (middle) Orbital-resolved effective mass-enhancement $m^{*}/m_{\text{DFT}}=1/\mathcal{Z}$, and (right) the imaginary part of self-energy at zero frequency $-\text{Im}\Sigma(\mathrm{i}0^{+})$ obtained from DFT+DMFT calculations at $T=300$ K. | Occupancy by DFT | Occupancy by DMFT | $m^{*}/m_{\text{DFT}}=1/\mathcal{Z}$ |  $-\text{Im}\Sigma(\mathrm{i}0^{+})$ (meV)
---|---|---|---|---
Fe-$3d$ | La/Ca-$d$ | As-$4p$ | Fe-$3d$ | La/Ca-$d$ | As-$4p$ | $d_{z^{2}}$ | $d_{x^{2}-y^{2}}$ | $d_{xz}/d_{yz}$ | $d_{xy}$ | $d_{z^{2}}$ | $d_{x^{2}-y^{2}}$ | $d_{xz}/d_{yz}$ | $d_{xy}$
UT-La122 | 6.516 | 0.968 | 3.066 | 6.524 | 0.983 | 3.131 | 1.92 | 1.94 | 2.02 | 2.09 | 9.47 | 3.35 | 18.32 | 31.68
CT-La122 | 6.489 | 0.973 | 3.102 | 6.520 | 0.977 | 3.089 | 1.64 | 1.75 | 1.60 | 1.49 | 4.29 | 2.95 | 11.55 | 11.74
UT-Ca122 | 6.507 | 0.383 | 3.056 | 6.546 | 0.385 | 3.078 | 1.83 | 1.80 | 2.00 | 1.94 | 4.00 | 2.93 | 15.33 | 17.18
CT-Ca122 | 6.515 | 0.394 | 2.993 | 6.559 | 0.396 | 3.016 | 1.65 | 1.75 | 1.79 | 1.60 | 0.86 | 0.18 | 5.24 | 4.47

Orbital-resolved occupancy numbers. Table 2 shows the normalized total
occupancy numbers of Fe-$3d$, La-$5d$/Ca-$3d$ and As-$4p$ orbitals, which
include the states both inside atomic sphere and the interstitial region [39],
obtained from our DFT and DFT+DMFT calculations. The most significant
observation is that although La122 was previously considered as one electron
doping from Ca122 [21], with nominal oxidation state of Fe1.5+, our
calculations (both DFT and DFT+DMFT) show almost the same Fe-$3d$ occupancy
$\sim 6.5$ for both La122 and Ca122. (The extra 0.5 electrons on iron are due
to strong Fe-$3d$ and As-$4p$ hybridizations.) Meanwhile, the La-$5d$
occupancy is about 1.0 in both UT and CT phases. The As-$4p$ occupancies are
also found to be very close in both materials ($\sim 3.1$), which indicates
their similar strengths of $3d$-$4p$ hybridizations. This is also reflected in
the partial DOS presented in SM [38].

These results clearly show that, from Ca122 to La122, the extra one electron
brought in by La does not go to Fe-As layers (especially the Fe-$3d$
orbitals), instead it resides on the La site, which leads to an unusual
oxidation state close to $+2$, not $+3$ for La. This is related to the strong
hybridizations between the very broad La-$5d$/$6s$ states and the Fe-$3d$
states, similar to the situation in the newly discovered infinite-layer
nickelate superconductor, Nd1-xSrxNiO2, where the Nd-$5d$/$6s$ states also
strongly hybridize with Ni-$3d$ states and contribute to the Fermi surface
[40, 41, 42, 43]. In Fig. 1, we plot the DFT fat bands projected to the
Fe-$3d_{xy}$ orbital (blue) and the La-$5d$/Ca-$3d$ orbitals (red). Indeed, we
find substantial La-$5d$ states below the Fermi level for La122, whereas
negligible Ca-$3d$ component is found in the similar range for Ca122. Notably,
such hybridization between Fe-$3d$ and La-$5d$/$6s$ states induces gaps in the
$3d_{xy}$ band near the $\Gamma$ point in La122 (marked by the arrows in Fig.
1), which is a major difference between band structures of La122 and Ca122.
The occupancy results for Ca122 obtained here agree well with previous
calculations [44].

Since the low-energy physics of these materials are mainly determined by the
Fe-As layers, the similar total Fe-$3d$ occupancies should lead to their
similar properties. Particularly, the two materials are expected to have
similar electronic correlation strength and similar Lifshitz transition of
Fermi surface across the UT-CT phase transition, as we demonstrate below.

Figure 2: The variations of (a,b) orbital-resolved effective mass enhancement
$m^{*}/m_{\text{DFT}}=1/\mathcal{Z}$ and (c,d) the effective quasi-particle
scattering rate $-\text{Im}\Sigma(\mathrm{i}0^{+})$ for Fe-$3d$ orbitals as
functions of temperature. The left (right) column is for La122 (Ca122).

Electronic correlations. We then investigate the correlation effects in both
UT and CT phases of La122, and make a comparison with Ca122.

The calculated orbital-resolved enhancement of effective mass
$m^{*}/m_{\text{DFT}}$ and the quasi-particle scattering rate
$-\text{Im}\Sigma(i0^{+})$ due to correlation effects at room temperature
($300$ K) are presented in Table 2(the imaginary part of Matsubara self-energy
at room temperature are detailed in SM [38]). Let’s first look at La122. In
the UT phase, the mass enhancement ranges from 2.09 to 1.92, with $d_{xy}$
($d_{z^{2}}$) as the most (least) correlated orbital. Especially, the $d_{xy}$
orbital shows much larger effective quasi-particle scattering rate than other
orbitals. When transiting into the CT phase, the correlation effects are
clearly weakened, which is reflected in the increased band width from UT to CT
phase. In CT-La122, The mass enhancement ranges from 1.75 to 1.49, but now
with the $d_{x^{2}-y^{2}}$ ($d_{xy}$) as the most (least) correlated orbital.
The scattering rate decreases significantly across the phase transition. For
the $d_{xy}$ orbital, the value decreases from 31.68 to 18.32. This is
primarily due to that the $3d_{xy}$ bands are pushed below the Fermi level in
the CT phase (see Fig. 1), with substantial increase of the occupancy away
from half filling.

In comparison, the results for Ca122 exhibit the same trend of variation in
the UT-CT phase transition as expected from their similar total occupancy of
Fe-$3d$ orbitals, which is also observed in the optical conductivity results
[45]. The mass enhancement of Ca122 in Table 2 are consistent with
experimental [46] and previous DFT+DMFT results [44, 35]. Particularly, the
bands of CT-La122 (UT-La122) show similar coherence as CT-Ca122 (UT-Ca122).
This differs from the result in Ref. [22], in which a much less coherence was
observed for CT-Ca122 (compared to CT-La122) as a larger $c$ lattice parameter
was used.

We further study the temperature dependence of the mass enhancement and the
quasi-particle scattering rate. The results are shown in Fig. 2. One observes
that the mass enhancement decreases and the scattering rate increases
significantly with increasing temperature, for both UT and CT phases of the
two materials. Especially, very large scattering rates appear at high
temperatures. These are typical Hund’s metal behaviors [26, 27, 28, 29]. We
have also verified that these values would be strongly suppressed in the
absence of Hund’s coupling [38].

Figure 3: The ratio of SW, $K(T,\Omega)/K(400K,\Omega)$, as a function of the
cutoff frequency $\Omega$ for (a) La122 and (b) Ca122.

The spectral weight (SW) transfer in the optical conductivity
$\sigma_{1}(T,\omega)$ is considered as a hallmark of correlated multi-orbital
physics [47, 48, 49, 28, 50]. Here, we calculate the optical conductivity (see
SM [38]) and investigate the integrated SW:

$K(T,\Omega)=\frac{Z_{0}}{\pi^{2}}\int^{\Omega}_{0}{\sigma_{1}(T,\omega)\mathrm{d}\omega},$
(1)

where $Z_{0}=376.73\ \Omega$ is the impedance of free space. In Fig. 3, we
plot the ratio $K(T,\Omega)/K(400\ \text{K},\Omega)$ as a function of the
cutoff frequency $\Omega$. One observes that in each phase, the results for
the two materials are similar and exhibit the same trend of SW transfer: With
decreasing temperature, the SW transfers from $\sim$1000 cm-1 of the mid
infrared (MIR) range to both far infrared (FIR) and near infrared (NIR)
ranges. The transfer to FIR represents the Drude component narrowing when
temperature decreases, which is a common phenomenon for metals. On the other
hand, the transfer from MIR to NIR could be explained by the correlation
effects due to Hund’s coupling [49], which suppresses the spin fluctuation to
form high spin state on the Fe ion (the atomic configuration is detailed in SM
[38]). In addition, the SW ratio recovers to $\approx 1$ at about 12000 cm-1,
at which the energy scale of spectral-weight transfer is defined. This scale
is about the same size of Hund’s coupling $J$ in our calculation. As a
comparison, we note that in the metallic phase of Mott-Hubbard systems, the SW
transfer direction is from high to low energies into the Drude component with
decreasing temperature [47]. And in that case, the energy scale is associated
with the Hubbard $U$, which is much larger than the Hund’s coupling.

Figure 4: The correlated electronic structures obtained from the DFT+DMFT
calculations at $T=300$ K. The columns from left to right are for UT-La122,
CT-La122, UT-Ca122 and CT-Ca122, respectively. (a)-(d) Momentum resolved
spectral function $\mathcal{A}(k,\omega)$. (e)-(h) Fermi surface in the (001)
plane. (i)-(l) Fermi surface in the (110) plane. The corresponding DFT results
are shown in dashed lines for reference.

Fermi surfaces. The momentum-resolved spectral functions
$\mathcal{A}(k,\omega)$ from our DFT+DMFT calculations at $T=300$ K are
plotted in Fig. 4(a-d). The quasi-particle bands near the Fermi level are
obviously renormalized compared with the LDA bands (black dashed lines). The
quasi-particle bandwidth is narrowed to about half of the LDA bands, which is
consistent with renormalization factor $\mathcal{Z}$ discussed above. The
correlated electronic structure undergoes obvious incoherence to coherence
crossover during the UT-to-CT phase transition in both La122 and Ca122. Again,
one can observe that CT-La122 and CT-Ca122 have similar coherence. In
addition, we find that the all spectra become much more coherent when ignoring
the Hund’s coupling [38], manifesting features of Hund’s metals [27].

In Fig. 4(e-l), we plot the corresponding Fermi surfaces in the $k_{z}=0$
plane and in the $k_{x}=k_{y}$ plane. Overall, compared to the DFT results,
the Fermi surface topology is not significantly modified by the correlation
effects. For UT-La122, there is one hole Fermi surface with $d_{xy}$ character
in the $k_{z}=0$ plane around the $\Gamma$ point (Fig. 4(e)), while for UT-
Ca122, there are three hole Fermi surfaces (Fig.4(g), the outer one is of
$d_{xy}$ character and the inner two are of $d_{xz}$/$d_{yz}$ character [44]).
In the $k_{x}=k_{y}$ plane, UT-La122 has another Fermi surface with La-$5d$
character around $\Gamma$ (Fig. 4(i)), while it disappears in UT-Ca122 because
of the negligible Ca-$3d$ occupancy (Fig. 4(k)). Across the UT-CT phase
transition, the most obvious change is that the conduction band at $\Gamma$
point in UT phase is pushed down below the Fermi level in the CT phase, which
occurs for both materials (see also Fig. 1). As a result, the hole Fermi
surfaces around $\Gamma$ in UT phase of both materials disappear in the CT
phase, corresponding to a Lifshitz transition of Fermi surface topology.

Discussion. Our result shows that in either UT or CT phase, La122 and Ca122
are similar in terms of qualitative band features, correlation strength, band
coherence, SW transfer, and Fermi surface topology. All these are connected
with the unusual oxidation state of La which results in their similar Fe (and
As) orbital occupation numbers. Particularly, we do not observe a marked
difference in band coherence and Fermi surface topology between CT-La122 and
CT-Ca122, indicating that these should not be primary factors for their
possibly different superconducting behaviors. In fact, it was noted that
superconductivity in CT-Ca122 is very sensitive to pressure and only occurs
under non-hydrostatic conditions [16, 30]. Based on their similarity, the same
may also be true for CT-La122. We suggest further experiments to test out
possible superconductivity in CT-La122 with different pressure conditions.

Our result also indicates that both phases of La122 (and Ca122) are Hund’s
metals in the normal state. Recently, the discovery of superconductivity in
Nd1-xSrxNiO2 attracted great interest [40]. There, it was shown that doped
NdNiO2 is a multi-orbital Hund’s metal, and Ni-$3d$ occupancies are very
similar in both LaNiO2 and SrNiO2 compounds [42, 43]. Those results are very
similar to our findings here. Therefore, this study may also help to shed
light on the common features of Hund’s superconductors.

## References

  * Kamihara  _et al._ [2008] Y. Kamihara, T. Watanabe, M. Hirano, and H. Hosono, J. Am. Chem. Soc. 130, 3296 (2008).
  * Paglione and Greene [2010] J. Paglione and R. L. Greene, Nat. Phys. 6, 645 (2010).
  * Johnston [2010] D. C. Johnston, Adv. Phys. 59, 803 1061 (2010).
  * Stewart [2011] G. R. Stewart, Rev. Mod. Phys. 83, 1589 (2011).
  * Si  _et al._ [2016] Q. Si, R. Yu, and E. Abrahams, Nat. Rev. Mater. 1, 16017 (2016).
  * Ni  _et al._ [2008] N. Ni, S. Nandi, A. Kreyssig, A. I. Goldman, E. D. Mun, S. L. Bud’ko, and P. C. Canfield, Phys. Rev. B 78, 014523 (2008).
  * Kreyssig  _et al._ [2008] A. Kreyssig, M. A. Green, Y. Lee, G. D. Samolyuk, P. Zajdel, J. W. Lynn, S. L. Bud’ko, M. S. Torikachvili, N. Ni, S. Nandi, J. B. Leão, S. J. Poulton, D. N. Argyriou, B. N. Harmon, R. J. McQueeney, P. C. Canfield, and A. I. Goldman, Phys. Rev. B 78, 184517 (2008).
  * Saha  _et al._ [2012] S. R. Saha, N. P. Butch, T. Drye, J. Magill, S. Ziemak, K. Kirshenbaum, P. Y. Zavalij, J. W. Lynn, and J. Paglione, Phys. Rev. B 85, 024525 (2012).
  * Bud’ko  _et al._ [2016] S. L. Bud’ko, X. Ma, M. Tomić, S. Ran, R. Valentí, and P. C. Canfield, Phys. Rev. B 93, 024516 (2016).
  * Yildirim [2008] T. Yildirim, Phys. Rev. Lett. 102, 037003 (2008).
  * Colonna  _et al._ [2011] N. Colonna, G. Profeta, A. Continenza, and S. Massidda, Phys. Rev. B 83, 094529 (2011).
  * Goldman  _et al._ [2008] A. I. Goldman, D. N. Argyriou, B. Ouladdiaf, T. Chatterji, A. Kreyssig, S. Nandi, N. Ni, S. L. Bud’ko, P. C. Canfield, and R. J. McQueeney, Phys. Rev. B 78, 100506 (2008).
  * Jeffries  _et al._ [2012] J. R. Jeffries, N. P. Butch, K. Kirshenbaum, S. R. Saha, G. Samudrala, S. T. Weir, Y. K. Vohra, and J. Paglione, Phys. Rev. B 85, 184501 (2012).
  * Goldman  _et al._ [2009] A. I. Goldman, A. Kreyssig, K. Prokeš, D. K. Pratt, D. N. Argyriou, J. W. Lynn, S. Nandi, S. A. J. Kimber, Y. Chen, Y. B. Lee, G. Samolyuk, J. B. Leão, S. J. Poulton, S. L. Bud’ko, N. Ni, P. C. Canfield, B. N. Harmon, and R. J. McQueeney, Phys. Rev. B 79, 024513 (2009).
  * Pratt  _et al._ [2009] D. K. Pratt, Y. Zhao, S. A. J. Kimber, A. Hiess, D. N. Argyriou, C. Broholm, A. Kreyssig, S. Nandi, S. L. Bud’ko, N. Ni, P. C. Canfield, R. J. McQueeney, and A. I. Goldman, Phys. Rev. B 79, 060510 (2009).
  * Torikachvili  _et al._ [2008] M. S. Torikachvili, S. L. Bud’ko, N. Ni, and P. C. Canfield, Phys. Rev. Lett. 101, 057006 (2008).
  * Zhao  _et al._ [2015] K. Zhao, C. Stingl, R. S. Manna, C. Q. Jin, and P. Gegenwart, Phys. Rev. B 92, 235132 (2015).
  * Chen  _et al._ [2016] D.-Y. Chen, J. Yu, B.-B. Ruan, Q. Guo, L. Zhang, Q.-G. Mu, X.-C. Wang, B.-J. Pan, G.-F. Chen, and Z.-A. Ren, Chin. Phys. Lett. 33, 067402 (2016).
  * Iyo  _et al._ [2019] A. Iyo, S. Ishida, H. Fujihisa, Y. Gotoh, I. Hase, Y. Yoshida, H. Eisaki, and K. Kawashima, J. Phys. Chem. Lett. 10, 1018 (2019).
  * Mazin  _et al._ [2019] I. I. Mazin, M. Shimizu, N. Takemori, and H. O. Jeschke, Phys. Rev. Lett. 123, 267001 (2019).
  * Usui and Kuroki [2019] H. Usui and K. Kuroki, Phys. Rev. Research 1, 033025 (2019).
  * Acharya  _et al._ [2020] S. Acharya, D. Pashov, F. Jamet, and M. van Schilfgaarde, Phys. Rev. Lett. 124, 237001 (2020).
  * Werner  _et al._ [2008] P. Werner, E. Gull, M. Troyer, and A. J. Millis, Phys. Rev. Lett. 101, 30 (2008).
  * Medici [2011] L. d. Medici, Phys. Rev. B 83, 205112 (2011).
  * Medici  _et al._ [2011] L. d. Medici, J. Mravlje, and A. Georges, Phys. Rev. Lett. 107, 256401 4 (2011).
  * Yin  _et al._ [2011] Z. P. Yin, K. Haule, and G. Kotliar, Nat. Mater. 10, 932 (2011).
  * Georges  _et al._ [2013] A. Georges, L. d. Medici, and J. Mravlje, Annu. Rev. Condens. Matter Phys. 4, 137 (2013).
  * Deng  _et al._ [2016] X. Deng, K. Haule, and G. Kotliar, Phys. Rev. Lett. 116, 256401 (2016).
  * Haule and Kotliar [2009] K. Haule and G. Kotliar, New J. Phys. 11, 025021 (2009).
  * Yu  _et al._ [2009] W. Yu, A. A. Aczel, T. J. Williams, S. L. Bud’ko, N. Ni, P. C. Canfield, and G. M. Luke, Phys. Rev. B 79, 020511 (2009).
  * Haule  _et al._ [2010] K. Haule, C.-H. Yee, and K. Kim, Phys. Rev. B 81, 195107 (2010).
  * Blaha  _et al._ [2018] P. Blaha, K. Schwarz, G. K. H. Madsen, D. Kvasnicka, J. Luitz, R. Laskowski, F. Tran, and L. D. Marks,  _WIEN2k, An Augmented Plane Wave + Local Orbitals Program for Calculating Crystal Properties_ (Karlheinz Schwarz, Techn. Universität Wien, Austria, 2018).
  * Blaha  _et al._ [2020] P. Blaha, K. Schwarz, F. Tran, R. Laskowski, G. K. H. Madsen, and L. D. Marks, J. Chem. Phys. 152, 074101 (2020).
  * Yin  _et al._ [2012] Z. P. Yin, K. Haule, and G. Kotliar, Phys. Rev. B 86, 195141 (2012).
  * Mandal  _et al._ [2014] S. Mandal, R. E. Cohen, and K. Haule, Phys. Rev. B 90, 060501 (2014).
  * Gull  _et al._ [2011] E. Gull, A. J. Millis, A. I. Lichtenstein, A. N. Rubtsov, M. Troyer, and P. Werner, Rev. Mod. Phys. 83, 349 (2011).
  * Haule [2015] K. Haule, Phys. Rev. Lett. 115, 196403 (2015).
  * [38] See Supplemental Material at [url] for: The Supplemental Material includes Refs. .
  * Haule  _et al._ [2008] K. Haule, J. H. Shim, and G. Kotliar, Phys. Rev. Lett. 100, 226402 (2008).
  * Li  _et al._ [2019] D. Li, K. Lee, B. Y. Wang, M. Osada, S. Crossley, H. R. Lee, Y. Cui, Y. Hikita, and H. Y. Hwang, Nature 572, 624 (2019).
  * Hepting  _et al._ [2020] M. Hepting, D. Li, C. J. Jia, H. Lu, E. Paris, Y. Tseng, X. Feng, M. Osada, E. Been, Y. Hikita, Y.-D. Chuang, Z. Hussain, K. J. Zhou, A. Nag, M. Garcia-Fernandez, M. Rossi, H. Y. Huang, D. J. Huang, Z. X. Shen, T. Schmitt, H. Y. Hwang, B. Moritz, J. Zaanen, T. P. Devereaux, and W. S. Lee, Nat. Mater. 19, 381 (2020).
  * Wang  _et al._ [2020] Y. Wang, C.-J. Kang, H. Miao, and G. Kotliar, Phys. Rev. B 102, 161118 (2020).
  * Petocchi  _et al._ [2020] F. Petocchi, V. Christiansson, F. Nilsson, F. Aryasetiawan, and P. Werner, Phys. Rev. X 10, 041047 (2020).
  * Diehl  _et al._ [2014] J. Diehl, S. Backes, D. Guterding, H. O. Jeschke, and R. Valentí, Phys. Rev. B 90, 085110 (2014).
  * Xing  _et al._ [2016] Z. Xing, T. J. Huffman, P. Xu, A. J. Hollingshad, D. J. Brooker, N. E. Penthorn, M. M. Qazilbash, S. R. Saha, T. Drye, C. Roncaioli, and J. Paglione, Phys. Rev. B 94, 064514 (2016).
  * Coldea  _et al._ [2009] A. I. Coldea, C. M. J. Andrew, J. G. Analytis, R. D. McDonald, A. F. Bangura, J.-H. Chu, I. R. Fisher, and A. Carrington, Phys. Rev. Lett. 103, 026404 (2009).
  * Qazilbash  _et al._ [2008] M. M. Qazilbash, A. A. Schafgans, K. S. Burch, S. J. Yun, B. G. Chae, B. J. Kim, H. T. Kim, and D. N. Basov, Phys. Rev. B 77, 115121 (2008).
  * Wang  _et al._ [2012] N. L. Wang, W. Z. Hu, Z. G. Chen, R. H. Yuan, G. Li, G. F. Chen, and T. Xiang, J. Phys. Condens. Matter 24, 294202 (2012).
  * Schafgans  _et al._ [2012] A. A. Schafgans, S. J. Moon, B. C. Pursley, A. D. LaForge, M. M. Qazilbash, A. S. Sefat, D. Mandrus, K. Haule, G. Kotliar, and D. N. Basov, Phys. Rev. Lett. 108, 147002 (2012).
  * Corasaniti  _et al._ [2020] M. Corasaniti, R. Yang, K. Sen, K. Willa, M. Merz, A. A. Haghighirad, M. L. Tacon, and L. Degiorgi, Phys. Rev. B 102, 161109 (2020).


# Diverse Imitation Learning via Self-Organizing Generative Models

Author Name Affiliation<EMAIL_ADDRESS>Arash Vahabpour1 Tianyi Wang1
Qiujing Lu1 Omead Pooladzandi1 Vwani Roychowdhury1 1University of California,
Los Angeles (UCLA) {vahabpour, tianyiw, qiujing, opooladz<EMAIL_ADDRESS>
###### Abstract

Imitation learning is the task of replicating expert policy from
demonstrations, without access to a reward function. This task becomes
particularly challenging when the expert exhibits a mixture of behaviors.
Prior work has introduced latent variables to model variations of the expert
policy. However, our experiments show that the existing works do not exhibit
appropriate imitation of individual modes. To tackle this problem, we adopt an
encoder-free generative model for behavior cloning (BC) to accurately
distinguish and imitate different modes. Then, we integrate it with GAIL to
make the learning robust towards compounding errors at unseen states. We show
that our method significantly outperforms the state of the art across multiple
experiments.

##  1 Introduction

The goal of imitation learning is to learn to perform a task from expert
trajectories without a reward signal. Towards this end, two approaches are
studied in the literature. The first approach is behavior cloning (BC)
(Pomerleau, 1991), which learns the mapping between individual expert states-
action pairs in a supervised manner. It is well known that BC neglects long-
range dynamics within a trajectory. This leads to problem of compounding error
(Ross and Bagnell, 2010; Ross et al., 2011), that is, when a trained model
visits states that are slightly different from those visited in expert
demonstrations, the model takes erroneous actions, and the errors quickly
accumulate as the policy unrolls. Therefore, BC can only succeed when training
data is abundant.

An alternative approach, known as inverse reinforcement learning (IRL)
(Ziebart et al., 2008), recovers a reward function that makes the expert
policy optimal. Such an approach often requires iterative calculation of the
optimal policy with reinforcement learning (RL) in an inner loop, making it
computationally expensive. The computational cost can be reduced by jointly
optimizing a policy—via RL, and a cost function—via maximum causal entropy IRL
(Ziebart et al., 2008, 2010). The resulting algorithm is called generative
adversarial imitation learning (GAIL) (Ho and Ermon, 2016), and is closely
connected to generative adversarial networks (Goodfellow et al., 2014). GAIL
has shown better promise than BC, both in sample efficiency in the number of
expert trajectories, and in robustness towards unseen states.

Expert demonstrations of a particular task can display variations, even given
the same initial state. This diversity can often stem from multiple human
experts performing the same task in different ways. Alternatively, it might
arise from pursuing multiple unlabeled objectives, e.g. a robotic arm moving
towards different goals. In such scenarios, neither BC nor GAIL enjoys any
explicit mechanism to capture different modes of behavior. Particularly, BC
averages out actions from different modes. On the other hand, GAIL learns a
policy that captures only a subset of control behaviors, because adversarial
training is prone to mode collapse (Li et al., 2017; Wang et al., 2017).

Imitation of diverse behaviors is addressed in several prior works. Fei et al.
(2020); Merel et al. (2017) use labels of expert modes in training. These
works differ from ours in that we assume expert modes are unlabeled. InfoGAIL
(Li et al., 2017) and Intention-GAN (Hausman et al., 2017) augment the
objective of GAIL with the mutual information between generated trajectories
and the corresponding latent codes. This mutual information is evaluated with
the help of a posterior network. Wang et al. (2017) use a variational
autoencoder (VAE) module to encode expert trajectories into a continuous
latent variable. These latent codes are supplied as an additional input to
both the policy and the discriminator of GAIL. Nevertheless, as we show in our
experiments, both InfoGAIL and VAE-GAIL perform poorly in practice.

Multimodal imitation learning requires finding the correspondence between
expert trajectories and a latent code. Wang et al. (2017) suggest learning
this correspondence using the encoder in a VAE. However, this imposes the
difficulty of designing an isolated (as opposed to end-to-end) recurrent
module to encode sequences of state-action pairs. Moreover, VAEs are mainly
well suited for continuous latent variables. We propose to address these
challenges by removing the recurrent encoder module, and instead directly
search over samples of the latent variable.Vondrick et al. (2016); Bojanowski
et al. (2017); Hoshen et al. (2019) have recently considered similar encoder-
free models.

Prior work on encoder-free models has mainly focused on non-sequential data.
Our work is the first to consider encoder-free models for sequential decision
making, and in particular, multimodal imitation learning. We choose the name
self-organizing generative model (SOG) for our core algorithm, indicating that
it learns semantically smooth latent spaces. We present a unified framework
for both discrete and continuous latent variables. In both cases, we analyse
the relationship between our SOG algorithm and algorithms based on
maximization of the marginal data likelihood, e.g. Kingma and Welling (2013).
Moreover, we propose an extension to our SOG algorithm that enables training
with high dimensional latent spaces. Finally, we combine SOG—as a BC
method—with GAIL, to get the best of both worlds. This attempt is inspired by
the empirical study of Jena et al. (2020), where in a unimodal setting, a
combination of BC with GAIL is shown to yield better performance in fewer
training iterations.

In a nutshell, this work mainly focuses on imitation of diverse behaviors
while distinguishing different modes. Our contributions are summarized as
follows:

  1. 1.

We adopt SOG as an encoder-free generative model that can be used for
multimodal BC. We show that SOG procedure is closely connected to maximizing
the marginal likelihood of data. Besides, we introduce an extension of SOG
that enables learning with high-dimensional latent spaces.

  2. 2.

We integrate SOG with GAIL to benefit from accuracy of SOG and robustness of
GAIL at the same time. Thus, the resulting model (a) can generate trajectories
that are faithful to the expert, (b) is robust to unseen states, (c) learns a
latent variable that models variations of the expert policy, and (d) captures
all modes successfully without mode collapse.

  3. 3.

Our empirical results show that our model considerably outperforms the
baselines.

##  2 Background

###  2.1 Preliminaries

We consider an infinite-horizon discounted Markov decision process (MDP) as a
tuple $(\mathcal{S},\mathcal{A},P,r,\rho_{0},\gamma)$. Namely, $\mathcal{S}$
denotes the state space, $\mathcal{A}$ denotes the action space,
$P:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow\mathbb{R}$
represents the state transition probability distribution,
$r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ is the reward function,
$\rho_{0}:\mathcal{S}\rightarrow\mathbb{R}_{\geq 0}$ represents the
distribution over initial states, and $\gamma\in(0,1)$ is the reward discount
factor.

Let $\pi:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}_{\geq 0}$ denote a
stochastic policy function. Demonstrations of the policy $\pi$ are sequences
of state-action pairs, sampled as follows: ${\bm{s}}_{0}\sim\rho_{0}$,
${\bm{a}}_{t}\sim\pi({\bm{a}}_{t}|{\bm{s}}_{t})$, ${\bm{s}}_{t+1}\sim
P({\bm{s}}_{t+1}|{\bm{s}}_{t},{\bm{a}}_{t})$. When states and actions
generated by policy $\pi$, we define expected return as
$\mathbb{E}_{\pi}\left[r({\bm{s}},{\bm{a}})\right]\coloneqq\mathbb{E}_{\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}r({\bm{s}}_{t},{\bm{a}}_{t})\right]$.

###  2.2 Imitation Learning

The goal of imitation learning is to replicate the expert policy from
demonstrations without access to the underlying reward function. Two
approaches exist for this problem. The first approach is BC, which treats the
state-action pairs as inputs and outputs of a policy function, and optimizes
for the $L_{2}$ loss between the expert actions and the predicted actions. The
second approach is GAIL, which optimizes the following objective:

$\begin{split}\min_{\begin{subarray}{c}\pi_{\bm{\theta}}\end{subarray}}\max_{\begin{subarray}{c}D\end{subarray}}\,\\{&\mathbb{E}_{\pi_{\bm{\theta}}}\left[\log
D({\bm{s}},{\bm{a}})\right]\\\
+&\mathbb{E}_{\pi_{E}}\left[\log\left(1-D({\bm{s}},{\bm{a}})\right)\right]\\\
-&\lambda H(\pi_{\bm{\theta}})\\}.\end{split}$ (1)

In this equation, $\pi_{\bm{\theta}}$ denotes the policy function realized by
a neural network with weights $\bm{\theta}$, and $\pi_{E}$ denotes the expert
policy. Moreover, $D$ denotes the discriminative classifier that distinguishes
between state-action pairs from the trajectories generated by
$\pi_{\bm{\theta}}$ and $\pi_{E}$. Lastly,
$H(\pi_{\bm{\theta}})\coloneqq\mathbb{E}_{\pi_{\bm{\theta}}}[-\log\pi_{\bm{\theta}}({\bm{a}}|{\bm{s}})]$
represents the discounted causal entropy of the policy $\pi_{\bm{\theta}}$
(Bloem and Bambos, 2014). We provide additional details about optimization of
Equation 1 in Appendix C.

###  2.3 Imitation of Diverse Expert Trajectories

We model the variations in the expert policy with a latent variable
${\bm{z}}$. In particular, ${\bm{z}}$ is sampled from a discrete or continuous
prior distribution $p({\bm{z}})$ before each rollout, and specifies the mode
of behavior. Therefore, we formalize the generative process of a multimodal
expert policy as follows: ${\bm{z}}\sim
p({\bm{z}}),{\bm{s}}_{0}\sim\rho_{0},{\bm{a}}_{t}\sim\pi_{E}({\bm{a}}_{t}|{\bm{s}}_{t},{\bm{z}}),{\bm{s}}_{t+1}\sim
P({\bm{s}}_{t+1}|{\bm{a}}_{t},{\bm{s}}_{t})$.

##  3 Method

In this section, we aim to propose an algorithm for imitation of diverse
behaviors. We first introduce a generative model for learning the distribution
of arbitrary datasets. Namely, let ${\bm{x}}$ and ${\bm{y}}$ respectively
denote an input data point and a corresponding output data point. We assume
that ${\bm{y}}$ is generated through a two-stage random process:

  1. 1.

A latent variable ${\bm{z}}$ (independent of ${\bm{x}}$) is sampled from a
prior distribution $p({\bm{z}})$. This prior can be a given discrete
distribution over $K$ categories with probability masses
$\pi_{1},\ldots,\pi_{K}$, or an arbitrary continuous distribution, e.g. a
multivariate Gaussian $\displaystyle
p({\bm{z}})=\mathcal{N}({\bm{z}};\mathbf{0},{\bm{I}})$.

  2. 2.

A value ${\bm{y}}$ is sampled from a conditional distribution
$p({\bm{y}}|{\bm{z}},{\bm{x}};f)$ of the following form:

$p({\bm{y}}|{\bm{z}},{\bm{x}};f)=\mathcal{N}({\bm{y}};f({\bm{z}},{\bm{x}}),\sigma^{2}{\bm{I}}),$
(2)

where $f$ maps ${\bm{z}}$ and ${\bm{x}}$ to the mean of the distribution, and
$\sigma^{2}$ is the isotropic noise variance in the space of ${\bm{y}}$.

###  3.1 The SOG Algorithm in General

Given a dataset $\mathcal{D}=\\{({\bm{x}}_{i},{\bm{y}}_{i})\\}_{i=1}^{N}$,
consisting of $N$ i.i.d samples generated by the process above, we design an
algorithm to estimate $\bm{\theta}$, such that the samples generated by the
inferred model match the given data in distribution. First, one needs to
restrict the estimation of function $f$ within a family of parameterized
functions, e.g. neural networks $f_{\bm{\theta}}$ with weight parameters
$\bm{\theta}$. Next, starting with a random initialization of $\bm{\theta}$,
an iterative algorithm is adopted:  (1) Given current $\bm{\theta}$, one
computes estimates of the latent variables ${\bm{z}}_{i}$, $i=1,\cdots,N$,
corresponding to each data point. In a marginal likelihood framework, a
posterior distribution for ${\bm{z}}_{i}$ is estimated. In our SOG algorithm,
we instead obtain point estimates ${\bm{z}}_{i}={\bm{z}}_{i}^{*}$, where
${\bm{z}}_{i}^{*}=\operatorname*{arg\,max}_{{\bm{z}}_{i}}p({\bm{y}}_{i}|{\bm{z}}_{i},{\bm{x}}_{i};\,\bm{\theta})$.
If the latent variable is discrete, ${\bm{z}}_{i}^{*}$ is found by an
exhaustive search over ${\bm{z}}\in\\{1,\ldots,K\\}$. In the continuous case,
the optimal ${\bm{z}}$ is searched over samples of the prior distribution
$p({\bm{z}})$.  (2) Given the best latent codes, one takes a gradient step on
$\bm{\theta}$ to increase the likelihood
$p({\bm{y}}_{i}|{\bm{z}}_{i}^{*},{\bm{x}}_{i};\,\bm{\theta})$.  Both steps
decrease the same loss function, guaranteeing convergence.

Algorithm 1 Self-Organizing Generative Model (SOG)

1:Define: Loss function
$\mathcal{L}\left(\hat{{\bm{y}}},{\bm{y}}\right)\coloneqq||\hat{{\bm{y}}}-{\bm{y}}||^{2}$

2:Input: Initial parameters of policy network $\bm{\theta}_{0}$

3:for $epoch=1,2,\ldots$ do

4: for $iteration=1,2,\ldots$ do

5: Sample a minibatch of $N_{\text{data}}$ data points
$({\bm{x}}_{i},{\bm{y}}_{i})$ from the dataset.

6: for $i=1,2,\ldots,N_{\text{data}}$ do

7: Sample $N_{z}$ latent codes ${\bm{z}}_{j}\sim p({\bm{z}})$.

8: Calculate
${\bm{z}}_{i}\coloneqq\arg\min_{{\bm{z}}_{j}}\mathcal{L}\left(f_{\bm{\theta}}({\bm{z}}_{j},{\bm{x}}_{i}),{\bm{y}}_{i}\right)$.

9: end for

10: Calculate
$\mathcal{L}_{\mathrm{SOG}}=\sum_{i=1}^{N_{\text{data}}}\mathcal{L}\left(f_{\bm{\theta}}({\bm{z}}_{i},{\bm{x}}_{i}),{\bm{y}}_{i}\right)$
and its gradients w.r.t. ${\bm{\theta}}$.

11: Update $f_{\bm{\theta}}$ by stochastic gradient descent on ${\bm{\theta}}$
to minimize $\mathcal{L}_{\mathrm{SOG}}$.

12: end for

13:end for

Formal relationships and equivalences between the SOG algorithm – that
iteratively computes a Maximum Likelihood Estimate (MLE) of ${\bm{z}}$ for
given $\bm{\theta}$ and then increases the likelihood of data computed at
these MLE estimates of ${\bm{z}}$ by updating the model parameters
$\bm{\theta}$ – and algorithms based on maximizing the marginalized likelihood
– that iteratively compute a posterior distribution of ${\bm{z}}$ for given
$\bm{\theta}$, and then increase the likelihood of data marginalized over the
distribution of ${\bm{z}}$ by updating the model parameters $\bm{\theta}$ –
are included in Appendices D and E.

###  3.2 Adopting SOG for Multimodal Imitation Learning

In the following, we introduce two approaches for multimodal imitation
learning:  (1) adopting SOG for multimodal BC, and  (2) combining multimodal
BC with GAIL using SOG.  Due to space limits, the corresponding algorithms are
presented in Appendix A.

#### Multimodal Behavior Cloning with SOG

We adopt Algorithm 1 for learning the relationship between state-action pairs
in a multimodal BC setting. We enforce an additional constraint that the
latent code is shared across each trajectory. Hence, we propose Algorithm 2.

#### Multimodal Combination of BC and GAIL using SOG

Next, we introduce Algorithm 3, where we combine SOG, as a means for BC, with
GAIL to ensure robustness towards unseen states. This algorithm is inspired by
Jena et al. (2020), which in a unimodal setting optimizes a weighted sum of BC
loss and the GAIL “surrogate” loss (see Appendix C).

###  3.3 Further Properties of SOG

#### Sample Complexity

Training procedure of Algorithm 1, given a high-dimensional continuous latent
variable, requires a prohibitively large number of latent code samples (7 of
the algorithm). In Appendix G, we propose a computationally efficient
extension of this algorithm suited for high dimensions of the latent space.

#### Self-Organization in the Latent Space

In Appendix F, we show that latent codes corresponding to nearby data points
get organized close to each other. In addition, visual results in Sections B.3
and G show that different regions of the latent space organize towards
generating different modes of data. Similar properties have been studied in
“self-organizing maps” Bishop (2006), without a deep network. Hence, we chose
the phrase “self-organization” in naming of Algorithm 1.

##  4 Experiments

In our experiments, we demonstrate that our method can recover and distinguish
all modes of expert behavior, and replicate each robustly and with high
fidelity. We evaluate our models, SOG-BC and SOG-GAIL, against two multimodal
imitation learning baselines: InfoGAIL (Li et al., 2017) and VAE-GAIL (Wang et
al., 2017).

###  4.1 Experiment Setup

We adopt an experiment from Li et al. (2017) in which an agent moves freely at
limited velocities in a 2D plane. In this experiment, the expert produces
three distinct circle-like trajectories. In another series of experiments, we
evaluate our model on several complex robotic locomotion tasks, simulated via
the MuJoCo simulator (Todorov et al., 2012). These experiments, borrowed from
Rakelly et al. (2019), include multimodal tasks with discrete or continuous
modes. Discrete tasks include Ant-Fwd-Back, Ant-Dir-6, HalfCheetah-Fwd-Back,
Humanoid-Dir-6, Walker2d-Vel-6, and Hopper-Vel-6. Additionally, we conduct two
experiments with continuous modes, namely FetchReach and HalfCheetah-Vel. All
these experiments are named after standard MuJoCo environments. The suffix
“Fwd-Back” indicates that the modes correspond to forward or backward moving
directions. Besides, the suffix “Dir-6” implies six moving directions, namely
$k\cdot 2\pi/6$ angles. On the other hand, “Vel” and “Vel-6” indicates that
the expert selects velocities uniformly at random, resulting in different
modes of behavior. The sets of expert velocities are listed in Section B.1.
Later, we refer to the environments Ant, HalfCheetah, Humanoid, Walker2d, and
Hopper, as “locomotion” tasks.

The FetchReach experiment involves a simulated 7-DoF robotic arm (Plappert et
al., 2018). At the beginning of each episode, a target point is uniformly
sampled from a cubic region. The expert controls the arm to reach the desired
target within a tolerance range.

We explain more details about our experiment setup in Section B.1.

###  4.2 Evaluation

Table 1: Mean and standard deviation of the rewards for locomotion tasks.

Data set | SOG-BC | SOG-GAIL ($\lambda_{S}=1$) | InfoGAIL | VAE-GAIL | Expert
---|---|---|---|---|---
Circles | $992.1\pm 1.0$ | $985.9\pm 10.8$ | $766.0\pm 67.9$ | $912.3\pm 8.6$ | $998.3\pm 0.1$
Ant-Fwd-Back | $1165.2\pm 32.1$ | $1101.0\pm 61.7$ | $220.6\pm 296.3$ | $-385.3\pm 67.0$ | $1068.7\pm 109.7$
Ant-Dir-6 | $1073.2\pm 206.6$ | $1023.2\pm 166.1$ | $-14.5\pm 87.6$ | $-572.9\pm 62.9$ | $1031.7\pm 253.7$
HalfCheetah-Fwd-Back | $221.6\pm 957.3$ | $1532.6\pm 148.5$ | $484.2\pm 919.4$ | $84.0\pm 152.2$ | $1686.0\pm 135.4$
Humanoid-Dir-6 | $5996.0\pm 326.4$ | $5457.8\pm 640.9$ | $1333.9\pm 461.4$ | $2285.5\pm 946.1$ | $6206.6\pm 292.6$
Walker2d-Vel-6 | $1915.3\pm 402.1$ | $1698.7\pm 650.0$ | $947.3\pm 105.3$ | $1183.6\pm 68.0$ | $1964.6\pm 394.6$
Hopper-Vel-6 | $2222.3\pm 431.1$ | $2015.3\pm 547.3$ | $1216.8\pm 70.8$ | $1065.8\pm 30.1$ | $2229.7\pm 438.6$

#### Visualization

In Figure 1, we visualize the generated trajectories of four locomotion tasks
with discrete modes. Throughout this figure, we can see that both SOG-BC and
SOG-GAIL successfully learn the different modes. Moreover, we observe that
SOG-BC performs slightly more accurately than SOG-GAIL. We discuss the
differences between the two algorithms in a later paragraph, where we compare
them in terms of their robustness. In Figure 4, we visualize the velocities of
Walker2d-Vel-6 and Hopper-Vel-6 experiments under different policies. We
observe that while SOG-BC successfully distinguishes and reconstructs the
desired velocities, the baseline models fail.

Table 2: Metrics for reached targets in the FetchReach experiment: average hit
rate and estimated entropy of achieved targets (higher is better).

Metric | SOG-BC | SOG-GAIL | InfoGAIL | VAE-GAIL | Expert
---|---|---|---|---|---
Entropy (nats) | $2.05$ | $1.89$ | $0.27$ | $0.85$ | $2.18$
Hit Rate | $100\%$ | $97.0\%$ | N/A | $18.6\%$ | $100\%$

Table 3: Mutual information between the latent variable (target velocity of the embedded trajectory in the case of VAE-GAIL) and generated velocities in HalfCheetah-Vel. SOG-BC | SOG-GAIL | InfoGAIL | VAE-GAIL
---|---|---|---
$1.584$ | $1.431$ | $0.145$ | $0.750$

Circles |  |  |  |  |
---|---|---|---|---|---
Ant-Fwd-Back |  |  |  |  |
Ant-Dir-6 |  |  |  |  |
Humanoid-Dir-6 |  |  |  |  |
| Expert | VAE-GAIL | InfoGAIL | SOG-BC | SOG-GAIL ($\lambda_{S}=1$)

Figure 1:  Discrete latent variable: locomotion towards multiple directions.
Trajectories of the imitated policies for four tasks. Different latent codes
at inference time are color coded. From top to bottom: number of modes are 3,
2, 6, 6; number of trajectories per latent code are 1, 3, 3, 3. Both SOG-BC
and SOG-GAIL precisely separate and imitate the modes in all the experiments.
However, VAE-GAIL and InfoGAIL fail to separate the modes and to imitate the
expert.

(a) VAE-GAIL

(b) InfoGAIL

(c) SOG-BC

Figure 2:  Continuous latent variable: HalfCheetah-Vel. In (a), since VAE-GAIL
uses a high latent dimension, in the horizontal axis we consider the desired
velocity of the expert trajectory corresponding to the latent embedding.
However, in (b), (c) the horizontal axis corresponds to (the CDF of) the 1-D
Gaussian latent variable. The CDF is applied for better visualization. In SOG-
BC, different velocities in range $[1.5,3]$ are replicated with a higher
certainty.

Walker2d-Vel-6 |  |  |  |  |
---|---|---|---|---|---
Hopper-Vel-6 |  |  |  |  |
| Expert | VAE-GAIL | InfoGAIL | SOG-BC | SOG-GAIL ($\lambda_{S}=1$)

Figure 3:  Discrete latent variable: locomotion at six velocities. Velocities
generated by the policies vs rollout time steps. Three rollouts per mode are
visualized. Each mode is marked with a distinct color. Both SOG-BC and SOG-
GAIL separate and imitate most of the modes, producing the desired velocities.
VAE-GAIL and InfoGAIL exhibit an inferior performance.

(a) Expert

(b) VAE-GAIL

(c) InfoGAIL

(d) SOG-BC

Figure 4:  Continuous latent variable: FetchReach. Sample trajectories of the
robotic arm. In (b), (d) expert trajectories are embedded and fed to the
learned model for reconstruction. In (c), since InfoGAIL lacks any module for
encoding expert trajectories, no target point is considered, and random codes
are used to generate trajectories. Blue circles mark different targets. Colors
of the trajectories are chosen at random for better visual distinction.
Compared to the baselines, SOG-BC better reaches different targets.

In Figures 4, 8 and 2, we visualize the results of continuous experiments
FetchReach and HalfCheetah-Vel. In these figures, SOG-BC and SOG-GAIL produce
similar plots, hence we drop the SOG-GAIL plot.

In Figure 4, we plot the trajectories of the Fetch robotic arm for different
samples of the latent variable. We observe that SOG-BC is able to control the
robotic arm towards any point in the 3D space. In Figure 8, we plot the
reached targets for different choices of the latent code, and observe a
semantically meaningful interpolation in the latent space. This video contains
animations of test rollouts.

In Figure 2, we plot the spread of instantaneous horizontal velocities of the
HalfCheetah agent as the latent variable varies. In SOG-BC and InfoGAIL, the
horizontal axis corresponds to the 1-D latent code used to train the model.
However, since VAE-GAIL requires a high dimensional latent variable for best
performance, the correspondence between latent codes and the model generated
horizontal velocities cannot be directly visualized. Therefore, we utilize the
following fact: if the latent embedding of an expert trajectories is fed to
the learned policy, the generated trajectory produces similar velocities as
the expert trajectory. Thus, for VAE-GAIL we plot the variations in the
generated velocities as a function of the target velocity of the corresponding
expert trajectory. We observe that compared to the baselines, SOG-BC better
associates the latent variable to goal velocities in the range $[1.5,3]$.
Also, SOG-BC produces a narrower error band, which indicates better certainty
given fixed latent codes. This video illustrates the imitated behavior of SOG-
BC at different velocities as the latent code varies.

#### Metrics

In Table 1, we use ground-truth reward functions to evaluate the collected
rewards under different policies. We explain these ground-truth rewards in
Section B.1. In each entry of the table, we report the best iteration in terms
of the mean reward. Particularly, we calculate mean and std of the rewards
across 100 rollouts. To find the mode-specific reward for each latent code, we
select a correspondence of latent codes to actual modes that gives the best
expected sum of rewards. This metric measures whether all modes are separated
and properly imitated, because each mode has a distinct reward function.
Across all experiments, SOG-BC performs significantly better than the
baselines. The only exception is HalfCheetah-Dir, where SOG-GAIL performs
better. We attribute this latter observation to the difficulty of the task,
which is alleviated by the explorations of GAIL.

In Table 2, we evaluate the FetchReach experiment through two metrics. First,
we collect the achieved targets under different policies, and clip them within
the cubic region from which the targets are sampled. Then, we quantify the
desired uniformity of the distribution of the achieved targets over the cube.
Since uniform distribution is the maximum entropy distribution over
rectangular supports, we alternatively estimate the entropy of the achieved
targets via the method of Kozachenko and Leonenko (1987). As a second metric,
we embed expert trajectories, and measure the hit rate of the reconstructed
trajectories given the same targets as the expert.

In the HalfCheetah-Vel experiment, to measure the correspondence between the
latent variable and generated instantaneous horizontal velocities, we estimate
their mutual information through the method of Kraskov et al. (2004). The
results are presented in Table 3.

#### Robustness

We design additional experiments to show that integration of SOG with GAIL in
Algorithm 3 makes the learned policies robust towards unseen states. Observing
space limits, we present the results in Section B.2.

##  5 Concluding Remarks

Our empirical results show that SOG-BC and SOG-GAIL distinguish and learn
different modes of behavior better than the baseline models, particularly in
the discrete case. This advantage can be attributed to the differences between
the methods in inducing mode separation in GAIL. Specifically, VAE-GAIL and
InfoGAIL directly modify the GAIL optimization to encourage multimodality.
However, our method learns multiple modes in GAIL indirectly and through
integration with multimodal BC.

## References

  * Andrychowicz et al. [2017] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba.  Hindsight experience replay.  arXiv preprint arXiv:1707.01495, 2017.
  * Bishop [2006] Christopher M Bishop.  Pattern Recognition and Machine Learning.  springer, 2006.
  * Bloem and Bambos [2014] Michael Bloem and Nicholas Bambos.  Infinite Time Horizon Maximum Causal Entropy Inverse Reinforcement Learning.  In 53rd IEEE Conference on Decision and Control, pages 4911–4916. IEEE, 2014.
  * Bojanowski et al. [2017] Piotr Bojanowski, Armand Joulin, David Lopez-Paz, and Arthur Szlam.  Optimizing the Latent Space of Generative Networks.  arXiv preprint arXiv:1707.05776, 2017.
  * Fei et al. [2020] Cong Fei, Bin Wang, Yuzheng Zhuang, Zongzhang Zhang, Jianye Hao, Hongbo Zhang, Xuewu Ji, and Wulong Liu.  Triple-gail: a Multi-Modal Imitation Learning Framework with Generative Adversarial Nets.  arXiv preprint arXiv:2005.10622, 2020.
  * Goodfellow et al. [2014] Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative Adversarial Networks.  arXiv preprint arXiv:1406.2661, 2014.
  * Hausman et al. [2017] Karol Hausman, Yevgen Chebotar, Stefan Schaal, Gaurav Sukhatme, and Joseph Lim.  Multi-Modal Imitation Learning from Unstructured Demonstrations using Generative Adversarial Nets.  arXiv preprint arXiv:1705.10479, 2017.
  * Ho and Ermon [2016] Jonathan Ho and Stefano Ermon.  Generative Adversarial Imitation Learning.  arXiv preprint arXiv:1606.03476, 2016.
  * Hoshen et al. [2019] Yedid Hoshen, Ke Li, and Jitendra Malik.  Non-Adversarial Image Synthesis with Generative Latent Nearest Neighbors.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5811–5819, 2019.
  * Jena et al. [2020] Rohit Jena, Changliu Liu, and Katia Sycara.  Augmenting GAIL with BC for sample efficient imitation learning.  arXiv preprint arXiv:2001.07798, 2020.
  * Kingma and Welling [2013] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  arXiv preprint arXiv:1312.6114, 2013.
  * Kozachenko and Leonenko [1987] LF Kozachenko and Nikolai N Leonenko.  Sample estimate of the entropy of a random vector.  Problemy Peredachi Informatsii, 23(2):9–16, 1987.
  * Kraskov et al. [2004] Alexander Kraskov, Harald Stögbauer, and Peter Grassberger.  Estimating Mutual Information.  Physical review E, 69(6):066138, 2004.
  * Li et al. [2017] Yunzhu Li, Jiaming Song, and Stefano Ermon.  Infogail: Interpretable Imitation Learning from Visual Demonstrations.  arXiv preprint arXiv:1703.08840, 2017.
  * Lillicrap et al. [2015] Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra.  Continuous control with deep reinforcement learning.  arXiv preprint arXiv:1509.02971, 2015.
  * Merel et al. [2017] Josh Merel, Yuval Tassa, Dhruva TB, Sriram Srinivasan, Jay Lemmon, Ziyu Wang, Greg Wayne, and Nicolas Heess.  Learning human behaviors from motion capture by adversarial imitation.  arXiv preprint arXiv:1707.02201, 2017.
  * Plappert et al. [2018] Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, et al.  Multi-goal reinforcement learning: Challenging robotics environments and request for research.  arXiv preprint arXiv:1802.09464, 2018.
  * Pomerleau [1991] Dean A Pomerleau.  Efficient Training of Artificial Neural Networks for Autonomous Navigation.  Neural computation, 3(1):88–97, 1991.
  * Radford et al. [2015] Alec Radford, Luke Metz, and Soumith Chintala.  Unsupervised representation learning with deep convolutional generative adversarial networks.  arXiv preprint arXiv:1511.06434, 2015.
  * Rakelly et al. [2019] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen.  Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables.  In International conference on machine learning, pages 5331–5340. PMLR, 2019.
  * Ross and Bagnell [2010] Stéphane Ross and Drew Bagnell.  Efficient Reductions for Imitation Learning.  In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 661–668. JMLR Workshop and Conference Proceedings, 2010.
  * Ross et al. [2011] Stéphane Ross, Geoffrey Gordon, and Drew 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, pages 627–635. JMLR Workshop and Conference Proceedings, 2011.
  * Schulman et al. [2015a] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz.  Trust Region Policy Optimization.  In International conference on machine learning, pages 1889–1897. PMLR, 2015.
  * Schulman et al. [2015b] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel.  High-Dimensional Continuous Control Using Generalized Advantage Estimation.  arXiv preprint arXiv:1506.02438, 2015.
  * Schulman et al. [2017] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal Policy Optimization Algorithms.  arXiv preprint arXiv:1707.06347, 2017.
  * Todorov et al. [2012] Emanuel Todorov, Tom Erez, and Yuval Tassa.  MuJoCo: A physics engine for model-based control.  In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026–5033. IEEE, 2012.
  * Vondrick et al. [2016] Carl Vondrick, Hamed Pirsiavash, and Antonio Torralba.  Anticipating visual representations from unlabeled video.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 98–106, 2016.
  * Wang et al. [2017] Ziyu Wang, Josh Merel, Scott Reed, Greg Wayne, Nando de Freitas, and Nicolas Heess.  Robust Imitation of Diverse Behaviors.  arXiv preprint arXiv:1707.02747, 2017.
  * Ziebart et al. [2008] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, and Anind K Dey.  Maximum Entropy Inverse Reinforcement Learning.  In Aaai, volume 8, pages 1433–1438. Chicago, IL, USA, 2008.
  * Ziebart et al. [2010] Brian D Ziebart, J Andrew Bagnell, and Anind K Dey.  Modeling Interaction via the Principle of Maximum Causal Entropy.  In ICML, 2010.

##  Appendix A Imitation Learning Algorithms

In this section, we present the imitation learning algorithms described in
Section 3.2.

Algorithm 2 SOG-BC: Multimodal Behavior Cloning with SOG

1:Define: Loss function
$\mathcal{L}\left(\hat{{\bm{a}}},{\bm{a}}\right)\coloneqq||\hat{{\bm{a}}}-{\bm{a}}||^{2}$

2:Input: Initial parameters of policy network $\bm{\theta}_{0}$; expert
trajectories $\tau_{E}\sim\pi_{E}$

3:for $iteration=0,1,2,\ldots$ do

4: Sample state-action pairs $\chi_{E}\sim\tau_{E}$ with the same batch size.

5: Sample $N_{z}$ latent codes $\displaystyle{\bm{z}}_{i}\sim p({\bm{z}})$ for
each trajectory.

6: Calculate
${\bm{z}}_{E}\coloneqq\arg\min_{{\bm{z}}_{i}}\hat{\mathbb{E}}_{\tau_{E}}\left[\mathcal{L}\left(f_{\bm{\theta}}({\bm{z}}_{i},{\bm{s}}),{\bm{a}}\right)\right]$.

7: Calculate
$\mathcal{L}_{\mathrm{SOG}}=\hat{\mathbb{E}}_{\tau_{E}}\left[\mathcal{L}\left(f_{\bm{\theta}}({\bm{z}}_{E},{\bm{s}}),{\bm{a}}\right)\right]$
and its gradients w.r.t. $\bm{\theta}$.

8: Update $f_{\bm{\theta}}$ by stochastic gradient descent on $\bm{\theta}$ to
minimize $\mathcal{L}_{\mathrm{SOG}}$.

9:end for

Algorithm 3 SOG-GAIL: Multimodal Combination of BC and GAIL

1:Input: Initial parameters of policy and discriminator networks,
$\bm{\theta}_{0},{\bm{w}}_{0}$; expert trajectories $\tau_{E}\sim\pi_{E}$

2:for $i=0,1,2,\ldots$ do

3: Sample a latent code ${\bm{z}}_{i}\sim p({\bm{z}})$, and subsequently a
trajectory $\tau_{i}\sim\pi_{\bm{\theta}_{i}}(\cdot|{\bm{z}}_{i})$.

4: Sample state-action pairs $\chi_{i}\sim\tau_{i}$ and $\chi_{E}\sim\tau_{E}$
with the same batch size.

5: Update the discriminator parameters from ${\bm{w}}_{i}$ to ${\bm{w}}_{i+1}$
with the gradient
$\begin{split}&\mathbb{\hat{E}}_{\tau_{i}}\left[\nabla_{{\bm{w}}}\log
D_{{\bm{w}}}({\bm{s}},{\bm{a}})\right]+\\\
&\mathbb{\hat{E}}_{\tau_{E}}\left[\nabla_{{\bm{w}}}\log\left(1-D_{{\bm{w}}}({\bm{s}},{\bm{a}})\right)\right].\end{split}$

6: Calculate the surrogate loss $\mathcal{L}_{\mathrm{PPO}}$ using the PPO
rule with the following objective
$\mathbb{\hat{E}}_{\tau_{i}}\left[\nabla_{{\bm{w}}}\log
D_{{\bm{w}}}({\bm{s}},{\bm{a}})\right]-\lambda_{H}H(\pi_{\bm{\theta}_{i}}).$

7: Calculate the SOG loss $\mathcal{L}_{\mathrm{SOG}}$ per Algorithm 2.

8: Take a policy step from $\bm{\theta}_{i}$ to $\bm{\theta}_{i+1}$ w.r.t. the
objective $\mathcal{L}_{\mathrm{PPO}}+\lambda_{S}\mathcal{L}_{\mathrm{SOG}}$.

9:end for

##  Appendix B Further Details of the Experiments

In this appendix, we first present details about the setup of the experiments.
Then, we demonstrate the robustness of our method against noise. Lastly, we
show that a meaningful correspondence holds between the latent space and
target locations in the Fetch experiment.

###  B.1 Experiment Setup

In the Circles environment, the observed state at time $t$ is a concatenation
of positions from time $t-4$ to $t$. Expert demonstrations appear in three
modes, each trying to produce a distinct circle-like trajectory. The expert
tries to maintain a $(2\pi/100)$ rad/s angular velocity along the perimeter of
a circle. However, its displacement at each step incurs a 10$\%$-magnitude 2D
Gaussian noise from the environment.

To train an expert policy for MuJoCo locomotion tasks, we used Pearl [Rakelly
et al., 2019], a few-shot reinforcement learning algorithm. For the FetchReach
experiment, we considered a pretrained DDPG+HER model as expert [Lillicrap et
al., 2015; Andrychowicz et al., 2017].

In Section 4, we measured the rewards of different experiments as follows. For
the Circles experiment, we applied a Gaussian kernel to the distance of the
agent from the perimeter of each of the three reference trajectories. Hence,
each step yields a reward between 0 and 1. For all MuJoCo locomotion
environments, i.e. Ant, HalfCheetah, Walker2d, Hopper, and Humanoid, we used
the original rewards introduced in Rakelly et al. [2019]. We ran all
experiments with four random seeds and presented the results for the best
training iteration in terms of expected sum of rewards.

In experiment with varied velocities, having a suffix “Vel” in their name,
target velocities of the expert are sampled uniformly from the values listed
in Table 4. As a reference, HalfCheetah has a torso length of around 1 m.

Table 4: Set of possible expert velocities in different environments. Environment | Walker2d-Vel | Hopper-Vel | HalfCheetah-Vel
---|---|---|---
Possible Velocities | $\\{-1.5,-1,-0.5,0.5,1,1.5\\}$ | $\\{-2,-1.5,0.5,1,1.5,2\\}$ |  Interval of $[1.5,3]$

Episode lengths in different environments are listed in Table 5.

Table 5: Episode lengths in different environments. Environment | Circles | Ant | HalfCheetah | Humanoid | Walker2d | Hopper | FetchReach
---|---|---|---|---|---|---|---
Episode Length | $1000$ | $200$ | $200$ | $1000$ | $1000$ | $1000$ | $50$

We used the same network architectures across all algorithms. The policy
network passes the states and latent codes through separate fully connected
layers, and then adds them and feeds them to the final layer. The
discriminator receives state-action pairs and produces scores for the GAN
objective. The posterior network for InfoGAIL has the same architecture as the
discriminator, except it has a final softmax layer in the discrete case. For
continuous latent codes, InfoGAIL adopts a Gaussian posterior. The
architectures are illustrated in Figure 5.

(a) Policy network

(b) The discriminator

Figure 5: Neural network architectures. (a) the policy network, (b) the
discriminator in GAIL.

For each of the algorithms SOG-GAIL, InfoGAIL, and VAE-GAIL, we pretrain the
policy network with behavior cloning for better performance. Since the policy
network has an input head for the latent code, in SOG-GAIL we provide the
latent codes according to Algorithm 2. Also, because InfoGAIL is not equipped
with a module to extract the latent codes prior to training, we feed random
samples as the input latent codes during pretraining.

Across all experiments, we varied the coefficients $\lambda_{I}$,
corresponding to the mutual information term in InfoGAIL, and $\lambda_{S}$
from SOG-GAIL (Algorithm 3), at values of $\lambda=i*10^{j};\ i=1,2,5;\
0.01\leq\lambda\leq 1$.

###  B.2 Robustness Analysis

We design two experiments to show that SOG-GAIL learns policies that are
robust towards unseen states. In the first experiment, upon generating test
trajectories for Ant-Dir-6, we switch the movement angles to all six possible
directions. Since the motion angles within each expert trajectory are
consistent, such changes create unseen circumstances. We observe that
maneuvers with acute angles causes the agent trained with SOG-BC to topple
over and make the episode terminate. In contrast, SOG-GAIL manages to switch
between all six directions.

Secondly, we vary the coefficient $\lambda_{S}$ in SOG-GAIL in several
experiments. High values of $\lambda_{S}$ in limit lead to SOG-BC. We observe
that increasing $\lambda_{S}$ improves scores in Table 1 across all
experiments (except for HalfCheetah-Dir). We now perturb each learned policy
by taking random actions with a 20% probability. We observe that under this
perturbation, several of the experiments perform worse at higher values of
$\lambda_{S}$. This indicates that combination of SOG with GAIL is more robust
to perturbations. The results of our two tests are respectively depicted in
Figures 7 and 7.

(a) Goal

(b) SOG-GAIL

(c) SOG-BC

Figure 6:  Robustness of SOG-BC vs SOG-GAIL: Ant-Dir-6. To create unseen
situations, we switch the latent codes to those corresponding to the
directions shown in (a). In (b), (c) we show three trajectories generated in
different colors. We observe that the ant agent trained with SOG-GAIL performs
the desired task flawlessly, while the one trained with SOG-BC often topples
over and fails the task.

(d) Unperturbed policy

(e) Perturbed policy

Figure 7: Optimal choice of $\bm{\lambda_{S}}$. Parameter $\lambda_{S}$ is the
coefficient of the SOG loss in Algorithm 3. Performance of Algorithm 3 is
illustrated as $\lambda_{S}$ varies. We consider two cases of with/without
perturbations. These perturbations are imposed by taking random actions with a
chance of 20%. The vertical axis is normalized such that the best performance
in unperturbed settings achieves a score of 1, and a random policy achieves 0.
In Figure (b), we observe that the policy trained with higher values of
$\lambda_{S}$ shows less robustness to perturbations in several experiments .

###  B.3 Latent Space Interpretability

In Figure 8 we demonstrate that in the Fetch experiment, different areas of
the latent space smoothly correspond to various target points.

Figure 8:  Self-organization in the latent space: FetchReach. In each plot,
two dimensions among the three dimensions of the latent space are fixed,
whereas the third is varied. Each line corresponds to the location of the
robotic arm at the terminal state of different trajectories. In each line, two
latent dimensions take a fixed value, while the third varies. Lines are
randomly colored for better visual distinction.

##  Appendix C Optimization of GAIL

GAIL optimizes the objective of Equation 1 in two alternating steps: (1) a
gradient ascent step to increase it with respect to the discriminator
parameters, and (2) a Trust Region Policy Optimization (TRPO) [Schulman et
al., 2015a] or Proximal Policy Optimization (PPO) [Schulman et al., 2017] step
to decrease it with respect to the policy parameters.

TRPO and PPO are policy optimization algorithms that take the biggest possible
policy improvement step without stepping so far that performance accidentally
deteriorates. In particular, PPO achieves this by optimizing the following
“surrogate” objective at $i^{th}$ iteration:

$\begin{split}&\mathcal{L}_{\mathrm{PPO}}({\bm{s}},{\bm{a}},\bm{\theta}_{i},\bm{\theta})=\\\
&\min\left(\frac{\pi_{\bm{\theta}}({\bm{a}}|{\bm{s}})}{\pi_{\bm{\theta}_{i}}({\bm{a}}|{\bm{s}})}A^{\pi_{\bm{\theta}_{i}}}({\bm{s}},{\bm{a}}),\;\;g(\epsilon,A^{\pi_{\bm{\theta}_{i}}}({\bm{s}},{\bm{a}}))\right),\end{split}$
(3)

where $A$ is the advantage function [Schulman et al., 2015b], and $g$ is a
“clipping” function for some small parameter $\epsilon>0$:

$g(\epsilon,A)=\left\\{\begin{array}[]{ll}(1+\epsilon)A&A\geq 0\\\
(1-\epsilon)A&A<0.\end{array}\right.$ (4)

##  Appendix D Discrete Latent Variables: Maximum Marginal Likelihood Approach
vs SOG

Expectation-maximization (EM) is a well-known method for optimizing the
maximum likelihood objective –marginalized over the latent variables – in
models involving latent variable. In this section, we demonstrate how the
iterative procedure of Algorithm 1 can be derived as a special case of the EM
algorithm when the latent variable has a discrete distribution.

###  D.1 The EM Algorithm in General

First, we introduce the EM method in general. Let
${\bm{X}}=\\{{\bm{x}}_{i}\\}_{i=1}^{N}$,
${\bm{Y}}=\\{{\bm{y}}_{i}\\}_{i=1}^{N}$,
${\bm{Z}}=\\{{\bm{z}}_{i}\\}_{i=1}^{N}$ constitute a dataset generated by the
procedure described in Section 3, and let $\bm{\Pi}=\\{\pi_{k}\\}_{k=1}^{K}$
denote the set of probability masses for the discrete latent codes. The EM
framework defines a total likelihood function for any given arbitrary
distribution $q({\bm{Z}})$ as follows:

$\displaystyle\begin{split}\mathcal{L}(&q({\bm{Z}});\bm{\theta},\bm{\Pi})\\\
=&\mathbb{E}_{{\bm{Z}}\sim
q({\bm{Z}})}\log\left\\{p({\bm{Y}},{\bm{Z}}|{\bm{X}};\>\bm{\theta},\bm{\Pi})\right\\}+H(q({\bm{Z}}))\end{split}$
(5a) $\displaystyle\begin{split}=&\mathbb{E}_{{\bm{Z}}\sim
q({\bm{Z}})}\log\left\\{p({\bm{Y}}|{\bm{X}};\>\bm{\theta},\bm{\Pi})\right\\}\\\
&+\mathbb{E}_{{\bm{Z}}\sim
q({\bm{Z}})}\log\left\\{\frac{p({\bm{Z}}|{\bm{X}},{\bm{Y}};\>\bm{\theta},\bm{\Pi})}{q({\bm{Z}})}\right\\}\end{split}$
(5b) $\displaystyle\begin{split}\equiv&\log
p(\text{{data}}|\bm{\theta},\bm{\Pi})\\\
&-D_{\mathrm{KL}}\big{(}{q({\bm{Z}})}\>||\>{p({\bm{Z}}|{\bm{X}},{\bm{Y}};\>\bm{\theta},\bm{\Pi})}\big{)}\end{split}$
(5c)

After initialization of model parameters, the optimization proceeds by
alternating between two steps:

  1. 1.

Expectation Step (E Step).

$\displaystyle q^{t+1}({\bm{Z}})$
$\displaystyle\leftarrow\displaystyle\operatorname*{arg\,max}_{q}\mathcal{L}(q;\>\bm{\theta}^{t},\bm{\Pi}^{t})$
(6a)
$\displaystyle=p({\bm{Z}}|{\bm{X}},{\bm{Y}};\>\bm{\theta}^{t},\bm{\Pi}^{t})$
(6b) $\displaystyle\equiv\prod_{i=1}^{N}r_{i,{\bm{z}}_{i}}^{t}$ (6c)

  2. 2.

Maximization Step (M Step).

$\displaystyle\begin{split}\bm{\theta}^{t+1},&\bm{\Pi}^{t+1}\\\
\leftarrow&\operatorname*{arg\,max}_{\bm{\theta},\bm{\Pi}}\mathcal{L}(q^{t+1};\>\bm{\theta},\bm{\Pi})\end{split}$
(7a) $\displaystyle=\;$
$\displaystyle{\operatorname*{arg\,max}_{\bm{\theta},\bm{\Pi}}\>\mathbb{E}_{{\bm{Z}}\sim
q^{t+1}}[\log P({\bm{Y}},{\bm{Z}}|{\bm{X}};\>\bm{\theta},\bm{\Pi})]}$ (7b)
$\displaystyle\begin{split}=\;&\operatorname*{arg\,max}_{\bm{\theta},\bm{\Pi}}\sum_{i=1}^{N}\sum_{k=1}^{K}r_{ik}^{t}(\log\pi_{k}\\\
&+\log p({\bm{y}}_{i}|{\bm{x}}_{i},{\bm{z}}_{i}=k;\>\bm{\theta}))\end{split}$
(7c)

where

$\begin{split}r_{ik}^{t}&=p({\bm{z}}_{i}=k\>\lvert\>{\bm{x}}_{i},{\bm{y}}_{i};\>\bm{\theta}^{t},\bm{\Pi}^{t})\\\
&=\frac{p({\bm{y}}_{i}|{\bm{x}}_{i},{\bm{z}}_{i}=k;\>\bm{\theta}^{t})\>\pi_{k}^{t}}{\sum_{l=1}^{K}p({\bm{y}}_{i}|{\bm{x}}_{i},{\bm{z}}_{i}=l;\>\bm{\theta}^{t})\>\pi_{l}^{t}}.\end{split}$
(8)

Finally, we introduce the update rules in the M step for
$\bm{\theta},\bm{\Pi}$. Since the M step does not have a closed-form solution
w.r.t. $\bm{\theta}$, on can instead perform a batch gradient ascent step.
This approach for an intractable M step is known as generalized EM [Bishop,
2006]. Also, it is straightforward to derive the following update rule for
$\bm{\Pi}$ by maximizing the objective of Equation 7c subject to the
constraint $\sum_{k=1}^{K}\pi_{k}=1$:

$\pi_{k}^{t}\leftarrow\frac{\sum_{i}r_{ik}^{t}}{\sum_{il}r_{il}^{t}}.$ (9)

The E step can be carried out in two ways:

  1. 1.

Soft. using all the posterior probabilities in Equation 8, so that each data
point has a soft-distribution over the latent codes at each iteration step.
This corresponds to the original EM algorithm above.

  2. 2.

Hard. approximating the posterior probabilities by a one-hot distribution,
driven by the best latent code for the M step. In the following subsection, we
derive when the regular EM reduces to the Hard EM method.

###  D.2 Hard Assignment EM and SOG (Algorithm 1)

Hard assignment EM has a long history in the ML community and the assumptions
made to derive this algorithm are well known and are used to motivate and
justify the $K$-means algorithm from an EM perspective [Bishop, 2006]. As in
SOG, one replaces the probabilities $r_{ik}$ in Equation 8 with binary values
denoted as $\gamma_{ik}$. This hard assignment is performed by allocating all
the probability mass to the categorical variable for which the posterior
probability is the maximum.

Let
$k_{i}^{*}=\operatorname*{arg\,max}_{k}p({\bm{z}}_{i}=k\>\lvert\>{\bm{x}}_{i},{\bm{y}}_{i};\>\bm{\theta}^{t},\bm{\Pi}^{t})$.
Then in the Hard EM case, the E step becomes

$\gamma_{ik}^{t}=\begin{cases}1,&k=k_{i}^{*}\\\
0,&k\not=k_{i}^{*}\end{cases},$ (10)

and the M step of Equation 7c (w.r.t. $\bm{\theta}$) simplifies to:

$\operatorname*{arg\,min}_{\bm{\theta}}\sum_{i=1}^{N}||f({\bm{z}}_{i}=k_{i}^{*},{\bm{x}}_{i};\bm{\theta})-{\bm{y}}_{i}||^{2}.$
(11)

It follows from Equation 8 that if $\pi_{k}=1/K;\>k=1,\ldots,K$, then

$\begin{split}&\operatorname*{arg\,max}_{k\in\\{1,\ldots,K\\}}p({\bm{z}}_{i}=k\>\lvert\>{\bm{x}}_{i},{\bm{y}}_{i};\>\bm{\theta}^{t})\\\
&\hskip
21.68121pt=\operatorname*{arg\,max}_{k\in\\{1,\ldots,K\\}}p({\bm{y}}_{i}\>\lvert\>{\bm{z}}_{i}=k,{\bm{x}}_{i},;\>\bm{\theta}^{t}).\end{split}$

Hence, the Hard EM is equivalent to the update step in Algorithm 1 in this
special case.

Next we provide a sufficient condition under which both the Hard EM and SOG
(for any values of $\pi_{k}$’s) become a special case of the general EM.
First, recall that

$\begin{split}&p({\bm{y}}_{i}|{\bm{z}}_{i}=k,{\bm{x}}_{i};\,\bm{\theta})\\\
&\quad\propto\exp[-||f({\bm{z}}_{i}=k,{\bm{x}}_{i};\,\bm{\theta})-{\bm{y}}_{i}||^{2}/(2\sigma^{2})],\end{split}$
(12)

where $\sigma$ determines the spread around the manifold of data. Let

$\displaystyle k_{i}^{\dagger}$
$\displaystyle=\operatorname*{arg\,max}_{k_{i}\in\\{1,\ldots,K\\}}p({\bm{y}}_{i}|{\bm{z}}_{i}=k,{\bm{x}}_{i};\,\bm{\theta})$
(13a)
$\displaystyle=\operatorname*{arg\,min}_{k_{i}\in\\{1,\ldots,K\\}}||f({\bm{z}}_{i}=k,{\bm{x}}_{i};\bm{\theta})-{\bm{y}}_{i}||^{2}.$
(13b)

Assuming a very small spread $\sigma$, we get for all $k\not=k_{i}^{\dagger}$

$\lim_{\sigma\rightarrow
0}\frac{p({\bm{y}}_{i}|{\bm{z}}_{i}=k,{\bm{x}}_{i};\,\bm{\theta})}{p({\bm{y}}_{i}|{\bm{z}}_{i}=k_{i}^{\dagger},{\bm{x}}_{i};\,\bm{\theta})}=0\;,$
(14)

and substituting these probabilities in the equation for $r_{ik}^{t}$ in
Equation 8, we obtain that the posteriors follow the same pattern:

$\lim_{\sigma\rightarrow
0}p({\bm{z}}_{i}=k|{\bm{x}}_{i},{\bm{y}}_{i};\,\bm{\theta})=\begin{cases}1,&k=k_{i}^{\dagger}\\\
0,&k\not=k_{i}^{\dagger}\end{cases},$ (15)

Thus, the SOG algorithm that computes an MLE estimate of ${\bm{z}}$ at every
iteration (i.e. the value of the latent variable for which data likelihood is
maximized given $\bm{\theta}$) is equivalent to the standard EM when
$\sigma\approx 0$.

##  Appendix E Continuous Latent Variables: Maximum Marginal Likelihood
Approach vs SOG

In this section, we consider the case of continuous latent variables, and
demonstrate that in the asymptotic limit where $\sigma$ is sufficiently small,
the same relationship derived in the preceding section for the discrete case
still holds: The marginalized data log-likelihood
$p({\bm{y}}|{\bm{x}};\,\bm{\theta})$ is well approximated by the maximum of
the data log-likelihood computed over the latent variables as done in
Algorithm 1. In this limit, we can use Laplace approximation for integrals.
Laplace approximation provides a general way to approach marginalization
problems. The basic setting for Laplace approximation is a multivariate
integral of the form:

$I=\int e^{\,-h({\bm{z}})/t}\>d{\bm{z}},$ (16)

where ${\bm{z}}$ is a $d$-dimensional variable, and $h({\bm{z}})$ is a scalar
function with a unique global minimum at ${\bm{z}}={\bm{z}}^{*}$. Laplace
approximation considers a Taylor series expansion of $h({\bm{z}})$ around
${\bm{z}}^{*}$. In this way, as $t\rightarrow 0\,$, the integral of Equation
16 is approximated as:

$I=e^{\,-h({\bm{z}}^{*})/t}\>(2\pi)^{d/2}\>|{\bm{A}}|^{1/2}\>t^{\,d/2}\
+O(t),$ (17)

where ${\bm{A}}=({\bm{D}}^{2}h({\bm{z}}^{*}))^{-1}$ is the inverse of the
Hessian of $h({\bm{z}})$ evaluated at ${\bm{z}}^{*}$.

Now, consider the expression for the data likelihood of the model in Section 3

$\displaystyle I$ $\displaystyle=p({\bm{y}}|{\bm{x}};\,\bm{\theta})$ (18a)
$\displaystyle=\int
p({\bm{y}}|{\bm{z}},{\bm{x}};\,\bm{\theta})p({\bm{z}})d{\bm{z}}$ (18b)
$\displaystyle\begin{split}&\propto\int\exp\left[-||f({\bm{z}},{\bm{x}})-{\bm{y}}||^{2}/(2\sigma^{2})\right]\\\
&\hskip 21.68121pt\times\exp\left[-||{\bm{z}}||^{2}/2\right]d{\bm{z}}.\\\
\end{split}$ (18c)

For $\sigma$ sufficiently small, we can use Laplace approximation for this
integral. Specifically, comparing with Equation 16 we identify:

$\displaystyle h({\bm{z}})$
$\displaystyle\coloneqq||f({\bm{z}},{\bm{x}})-{\bm{y}}||^{2}+\sigma^{2}||{\bm{z}}||^{2},$
(19a) $\displaystyle t$ $\displaystyle\coloneqq 2\sigma^{2},$ (19b)

thus deriving:

$\displaystyle{\bm{z}}^{*}$
$\displaystyle=\operatorname*{arg\,min}h({\bm{z}})$ (20a)
$\displaystyle=\operatorname*{arg\,min}||f({\bm{z}},{\bm{x}})-{\bm{y}}||^{2}+\sigma^{2}||{\bm{z}}||^{2}$
(20b)
$\displaystyle\approx\operatorname*{arg\,min}||f({\bm{z}},{\bm{x}})-{\bm{y}}||^{2},$
(20c)

which is equivalent to 8 of Algorithm 1. The only nuance is that Algorithm 1
finds an approximation to ${\bm{z}}^{*}$ by searching among a batch of sampled
candidate ${\bm{z}}$’s. However, given large batches of ${\bm{z}}$’s, the best
candidate stays close to ${\bm{z}}^{*}$ with a high probability. Sampling of
${\bm{z}}$ from a Gaussian prior ensures that $||{\bm{z}}||$ remains small,
making the approximation of Equation 20c valid.

Plugging the terms in Equation 19 into Equation 17, we obtain the following
approximation to the data log-likelihood, which is later maximized w.r.t.
model parameters $\bm{\theta}$,

$\begin{split}\log
p({\bm{y}}|{\bm{x}};\,\bm{\theta})\approx-\,&\frac{1}{2\sigma^{2}}\left(||f({\bm{z}}^{*},{\bm{x}})-{\bm{y}}||^{2}+\sigma^{2}||{\bm{z}}^{*}||\right)\\\
+&\,\frac{1}{2}\log|{\bm{A}}|\\\ +&\,d\log\sigma\\\
+&\,\text{constant}.\end{split}$ (21)

We notice that in the asymptote of $\sigma\rightarrow 0$, the term
$||f({\bm{z}}^{*},{\bm{x}})-{\bm{y}}||^{2}\,/(2\sigma^{2})$ dominates the
other terms. Particularly, $||{\bm{z}}^{*}||$ and $\log|{\bm{A}}|$ do not
depend on $\sigma$, and the growth of $\log\sigma$ is dominated by that of
$1/\sigma^{2}$. Therefore, in the limit we can approximate Equation 21 with
only the first summand. That is,

$\log
p({\bm{y}}|{\bm{x}};\,\bm{\theta})\approx-\,\frac{1}{2\sigma^{2}}\,||f({\bm{z}}^{*},{\bm{x}})-{\bm{y}}||^{2},$
(22)

which yields 10 of Algorithm 1 since $\sigma$ is constant w.r.t.
$\bm{\theta}$. We assume Lipchitz continuity for $f$. Thus, the approximate
choice of ${\bm{z}}^{*}$, given large batches over the latent codes,
introduces an insignificant error to Equation 22.

It is worth mentioning that similar approximation results are well studied in
the framework of Bayesian model selection. In that context, the data
likelihood is marginalized instead over the model parameters, and also the
asymptotic variable is rather the number of observed data samples becoming
large Bishop [2006].

##  Appendix F The Self-Organization Effect

In this section, we discuss how the self-organization phenomenon, that is the
formation of similar output data in different regions of the latent space,
naturally emerges in Algorithm 1. In the following, we present three
propositions demonstrating different aspects of the self-organization
phenomenon in the latent space, under the assumption that the function
$f(\cdot)$ is a sufficiently smooth function of its inputs
$({\bm{x}},{\bm{z}})$, e.g. it is Lipschitz continuous:

  1. 1.

Continuity in the data space: We discuss that for every data point, there is a
neighborhood in the data space such that for the same latent code ${\bm{z}}$,
the variations in the $L_{2}$-loss in Algorithm 1 over all points in the
neighborhood is small.

  2. 2.

Two sufficiently close data points will have nearby winning latent codes (or
the same winning latent codes in the discrete case): We show that if for a
data point, a particular latent code results in a significantly better
$L_{2}$-loss than another latent code, then the same will also hold within a
neighborhood of that data point. Thus, if ${\bm{z}}$ is a winning code for a
data point, then the points in its neighborhood will also have winning codes
close to ${\bm{z}}$ (or the same winning code ${\bm{z}}$ in the discrete
case).

  3. 3.

Network training can be a collaborative process: We demonstrate that if we
update the network parameters $\bm{\theta}$ to improve the $L_{2}$-loss for a
particular data point (at a latent code $z$), then the same update will also
collaboratively improve the loss for other data points within the neighborhood
of the main data point (w.r.t. the same latent code $z$).

All three propositions provide insights on why the SOG training algorithm
converges fast. In particular, in the mini-batch optimization in any
implementation of the SOG algorithm, these propositions imply that the
parameter updates for a sample batch of data points, reduce the loss for
points neighboring to the sampled data points as well.

Formally, we shall make the assumption that the network $f$ is a Lipschitz
continuous function with the constant $M$. We also define the
$\delta_{x}$-$\delta_{y}$-neighborhood of a data point $({\bm{x}},{\bm{y}})$
as
$B({\bm{x}}_{1},{\bm{y}}_{1};\delta_{x},\delta_{y})=\\{({\bm{x}},{\bm{y}})\>\lvert\>\left\lVert{\bm{x}}-{\bm{x}}_{1}\right\rVert<\delta_{x}/M,\>\left\lVert{\bm{y}}-{\bm{y}}_{1}\right\rVert<\delta_{y}\\}$.
For later convenience, we define $\delta=\delta_{x}+\delta_{y}$. Finally, we
denote $d({\bm{x}},{\bm{y}},{\bm{z}};\bm{\theta})\coloneqq\left\lVert
f({\bm{x}},{\bm{z}};\bm{\theta})-{\bm{y}})\right\rVert$.

######  Proposition F.1.

For any data point $({\bm{x}}_{1},{\bm{y}}_{1})$, if
$({\bm{x}}_{2},{\bm{y}}_{2})\in
B({\bm{x}}_{1},{\bm{y}}_{1};\delta_{x},\delta_{y})$, then
$|d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}};\theta)-d({\bm{x}}_{2},{\bm{y}}_{2},{\bm{z}};\theta)|<\delta_{x}+\delta_{y}=\delta$

###### Proof.

From $({\bm{x}}_{2},{\bm{y}}_{2})\in
B({\bm{x}}_{1},{\bm{y}}_{1};\delta_{x},\delta_{y})$, follows that
$\left\lVert{\bm{x}}_{1}-{\bm{x}}_{2}\right\rVert<\delta_{x}/M$,
$\left\lVert{\bm{y}}_{1}-{\bm{y}}_{2}\right\rVert<\delta_{y}$. Therefore:

$\displaystyle|\left\lVert
f({\bm{x}}_{1},{\bm{z}};\bm{\theta})-{\bm{y}}_{1}\right\rVert-\left\lVert
f({\bm{x}}_{2},{\bm{z}};\bm{\theta})-{\bm{y}}_{2}\right\rVert|$ (23a)
$\displaystyle\leq\left\lVert
f({\bm{x}}_{1},{\bm{z}};\bm{\theta})-{\bm{y}}_{1}-(f({\bm{x}}_{2},{\bm{z}};\bm{\theta})-{\bm{y}}_{2})\right\rVert$
(23b) $\displaystyle\leq\left\lVert
f({\bm{x}}_{1},{\bm{z}};\bm{\theta})-f({\bm{x}}_{2},{\bm{z}};\bm{\theta})\right\rVert+\left\lVert{\bm{y}}_{1}-{\bm{y}}_{2}\right\rVert$
(23c) $\displaystyle\leq
M\left\lVert{\bm{x}}_{1}-{\bm{x}}_{2}\right\rVert+\left\lVert{\bm{y}}_{1}-{\bm{y}}_{2}\right\rVert$
(23d) $\displaystyle<\delta_{x}+\delta_{y}=\delta.$ (23e)

where Equations 23b and 23c follow reverse triangle inequality and triangle
inequality, respectively. Besides, Equation 23d follows from the Lipschitz
continuity assumption. ∎

######  Proposition F.2.

Let $({\bm{x}}_{1},{\bm{y}}_{1})$ be a data point, and
${\bm{z}}_{i},{\bm{z}}_{j}$ be two choices of the latent code.
If
$\dfrac{d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}}_{i};\bm{\theta})}{d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}}_{j};\bm{\theta})}>C>1$
for a constant $C$, and
$d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}}_{j};\bm{\theta})$ is finite, then
$\exists\delta_{x},\delta_{y}$ such that

$\forall({\bm{x}}_{2},{\bm{y}}_{2})\in
B({\bm{x}}_{1},{\bm{y}}_{1};\delta_{x},\delta_{y}),\quad\frac{d({\bm{x}}_{2},{\bm{y}}_{2},{\bm{z}}_{i};\bm{\theta})}{d({\bm{x}}_{2},{\bm{y}}_{2},{\bm{z}}_{j};\bm{\theta})}>1$

###### Proof.

$\forall\eta>0$, $\exists\delta_{x},\delta_{y}$ such that
$\delta_{x}+\delta_{y}=\delta<\eta
d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}}_{j};\bm{\theta})$. Then using proposition
F.1, the following sequence of inequalities holds

Δd(x_2, y_2, z_i) &≥d(x_1, y_1, z_i) - δ
¿ Cd(x_1,y_1, z_j) - δ
¿ Cd(x_1,y_1, z_j) - ηd(x_1,y_1, z_j)
= (C-η)d(x_1,y_1, z_j)
¿ (C-η)(d(x_2,y_2, z_j) - δ)
¿ (C-η)(1-η1-η)d(x_2,y_2, z_j),

where the step from (24) to (24) is done by realizing the fact that

$\eta
d({\bm{x}}_{2},{\bm{y}}_{2},{\bm{z}}_{j})>\eta\big{(}d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}}_{j})-\delta\big{)}>(1-\eta)\delta$

and therefore
$\delta<\frac{\eta}{1-\eta}d({\bm{x}}_{2},{\bm{y}}_{2},{\bm{z}}_{j})$.

It can be easily verified that there always exists a small enough value of
$\eta\rightarrow 0^{+}$ such that

$(C-\eta)\left(1-\frac{\eta}{1-\eta}\right)>1$

and therefore

$d({\bm{x}}_{2},{\bm{y}}_{2},{\bm{z}}_{i})>d({\bm{x}}_{2},{\bm{y}}_{2},{\bm{z}}_{j}).$

∎

######  Proposition F.3.

Let $({\bm{x}}_{1},{\bm{y}}_{1})$ be a data point and let ${\bm{z}}$ be a
choice of the latent code, respectively. Also assume an improvement update of
the network parameters $\bm{\theta}_{1}$ to $\bm{\theta}_{2}$ that ensures
$d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}};\bm{\theta}_{1})>d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}};\bm{\theta}_{2})$.
Then, there exist $\delta_{x}$ and $\delta_{y}$ such that
$\forall({\bm{x}}_{2},{\bm{y}}_{2})\in
B({\bm{x}}_{1},{\bm{y}}_{1};\delta_{x},\delta_{y}),\;\;d({\bm{x}}_{2},{\bm{y}}_{2},{\bm{z}};\bm{\theta}_{2})<d_{(\delta_{x},\delta_{y})}^{\;\max}(\bm{\theta}_{1})$
, where

$d_{(\delta_{x},\delta_{y})}^{\;\max}(\bm{\theta})=\max_{({\bm{x}},{\bm{y}})\in
B({\bm{x}}_{1},{\bm{y}}_{1};\delta_{x},\delta_{y})}d({\bm{x}},{\bm{y}},{\bm{z}};\bm{\theta})$

is the largest reconstruction distance for data points in the neighborhood of
$({\bm{x}}_{1},{\bm{y}}_{1})$

###### Proof.

Let
$a=d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}};\bm{\theta}_{1})-d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}};\bm{\theta}_{2})>0$,
and again $\delta=\delta_{x}+\delta_{y}$, then $\forall\delta_{x},\delta_{y}$,
and $\forall({\bm{x}}_{2},{\bm{y}}_{2})\in
B({\bm{x}}_{1},{\bm{y}}_{1};\delta_{x},\delta_{y})$

$\displaystyle d({\bm{x}}_{2},{\bm{y}}_{2},{\bm{z}};\bm{\theta}_{2})$
$\displaystyle\leq d_{(\delta_{x},\delta_{y})}^{\;\max}(\bm{\theta}_{2})$
(25a) $\displaystyle\leq
d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}};\bm{\theta}_{2})+\delta$ (25b)
$\displaystyle\leq
d({\bm{x}}_{1},{\bm{y}}_{1},{\bm{z}};\bm{\theta}_{1})-a+\delta$ (25c)
$\displaystyle\leq
d_{(\delta_{x},\delta_{y})}^{\;\max}(\bm{\theta}_{1})-a+\delta$ (25d)

where Equation 25b follows from Proposition F.1.

Since $a$ is finite, $\exists\delta_{x},\delta_{y}$ such that $\delta<a$, and
hence

$d({\bm{x}}_{2},{\bm{y}}_{2},{\bm{z}};\bm{\theta}_{2})<d_{(\delta_{x},\delta_{y})}^{\;\max}(\bm{\theta}_{1}).$

∎

##  Appendix G Extension of Algorithm 1 to Large Latent Spaces

Algorithm 1 suffers from an exponential growth of the number of samples
required, as the latent space dimension increases. To address this, one can
consider a coordinate-wise search of the latent code ${\bm{z}}$, which results
in Algorithm 4. This modification reduces the computational cost from
exponential to linear.

Algorithm 4 Coordinate-Wise Search in SOG

1: Define: Loss function
$\mathcal{L}(\hat{{\bm{y}}},{\bm{y}})\coloneqq||\hat{{\bm{y}}}-{\bm{y}}||^{2}$,
block size $\Delta$.

2:for $epoch=1,2,\ldots$ do

3: for $iteration=1,2,\ldots$ do

4: Sample a minibatch of $N_{\text{data}}$ data points
$({\bm{x}}_{i},{\bm{y}}_{i})$ from the dataset.

5: for $i=1,2,\ldots,N_{\text{data}}$ do

6: Initialize latent code ${\bm{z}}_{i}=\bm{0}\in\mathbb{R}^{d}$

7: for $b=1,2,\ldots,(d/\Delta)$ do

8: Initialize candidate latent codes ${\bm{z}}_{j}\coloneqq{\bm{z}}_{i};\
j=1,\ldots,N_{z}$.

9: Replace the $b$’th $\Delta$-element block in ${\bm{z}}_{j}$ with a sample
from $\mathcal{N}(\bm{0},{\bm{I}}^{\Delta\times\Delta})$;
$\>j=1,\ldots,N_{z}$.

10: Calculate
${\bm{z}}_{i}\coloneqq\arg\min_{{\bm{z}}_{j}}\mathcal{L}(f_{\bm{\theta}}({\bm{z}}_{j},{\bm{x}}_{i}),{\bm{y}}_{i})$.

11: end for

12: end for

13: Calculate
$\mathcal{L}_{\mathrm{SOG}}=\sum_{i=1}^{N_{\text{data}}}\mathcal{L}(f_{\bm{\theta}}({\bm{z}}_{j},{\bm{x}}_{i}),{\bm{y}}_{i})$
and its gradient w.r.t. $\bm{\theta}$.

14: Update $f_{\bm{\theta}}$ by stochastic gradient descent on $\bm{\theta}$
to minimize $\mathcal{L}_{\mathrm{SOG}}$.

15: end for

16:end for

In the FetchReach experiment, the latent dimension is 3 and we adopted
Algorithm 4 that resulted in superior results (see Figure 4). Besides, in
Figures 9, 10 and 11, we present visual results of Algorithm 4 demonstrating
the capability of our method in fitting complex distributions. The axes in
Figures 9 and 10 are the CDFs of the Gaussian latent codes for better
visualization. We used the network architecture of the generator in Radford et
al. [2015].

(a) Synthesized MNIST digits

(b) Embedding of MNIST test data in 2 dimensions

Figure 9: MNIST dataset. Output results of Algorithm 4 for the latent
dimension of 2. (a) Synthesized images (b) Embedding of MNIST test data in 2
dimensions by the search mechanism of the SOG algorithm. Color coding
represents ground truth digit labels. It is evident that despite the
unsupervised training of the SOG algorithm, data samples corresponding to each
digits are evenly organized across the latent space. Figure 10: Fashion MNIST.
Output results of Algorithm 4 for the latent dimension of 8. Each $6\times 6$
block (or block of blocks, etc.) sweeps over one dimension of the latent
space. Due to limited space, only part of the entire latent space is plotted.
Full results are provided in the link (see the summplementary material zip
file in the review stage). Different latent dimensions have smoothly captured
semantically meaningful aspects of the data, i.e. type of clothing, color
intensity, thickness, etc.

(a) CelebA ground-truth data

(b) Reconstruction by SOG algorithm

Figure 11: CelebA. Output results of Algorithm 4 for the latent dimension of
16. (a) Ground truth samples. (b) Reconstructed samples synthesized by SOG.
Note that the reconstructed images learned many of the salient features of the
original faces, including the pose of the faces, as well as the hairstyle in
most cases.