# Tricolour: an optimized SumThreshold flagger for MeerKAT

Benjamin V. Hugo1,2, Simon Perkins1, Bruce Merry1, Tom Mauch1 and Oleg M.
Smirnov2,1

###### Abstract

We present Tricolour, a package for Radio Frequency Interference mitigation of
wideband finely channelized MeerKAT correlation data. The MeerKAT passband is
heavily affected by interference from satellite, mobile, aircraft and
terrestrial transponders. Coupled with typical data rates in excess of 100
GiB/hr at 208kHz channelization resolution, mitigation poses a significant
processing challenge. Our flagger is highly configurable, parallel and
optimized, employing Dask and Numba technologies to implement the widely used
SumThreshold and MAD interference detection algorithms. We find that typical
208kHz channelized datasets can be processed at rates in excess of 400 GiB/hr
for a typical L-band flagging strategy on a modern dual-socket Intel Xeon
server.

Preprint
Astronomical Data Analysis Software and Systems XXX
J.E. Ruiz, F. Pierfederici, and P. Teuben, eds.
ASP Conference Series
©2021 Astronomical Society of the Pacific 1 0.3cm

1South African Radio Astronomy Observatory, Cape Town, South Africa;
<EMAIL_ADDRESS>
2Rhodes University, Makhanda (Grahamstown), Eastern Cape, South Africa

##  1 Radio Frequency Interference impacts on MeerKAT observation

Radio Frequency Interference (RFI) can be defined as any interfering signal
that negatively affects the observable instrumental bandwidth of a radio
telescope. The most prominent RFI impacting cm wavelength observations using
MeerKAT (Jonas et al. 2018) are spatially coherent man-made telecommunications
and navigational satellite transmissions. These substantially impacts all
observation, particularly that of L-band, shown in Fig. 1 for a typical
observation.

Figure 1.: Mean visibility amplitude of a typical unflagged and uncalibrated
L-band MeerKAT observation. Darkened areas show subbands where persistent
interference is most likely, and serves as a “static” background mask during
application of mitigation routines.

RFI affecting large portions of the bandwidth presents a substantial risk for,
especially, sensitive continuum radio science projects, because the
sensitivity of such observations scale as the square root of the observable
bandwidth. As shown in Fig. 1, man-made interference power is orders of
magnitude above the relatively dominating sub-Jy Active Galactic Nuclei (AGN)
population at cm wavelengths — even on a bright calibrator field! A large
number of science cases require very long observations to target imaging
sensitivities at the level of only a few $\mu\text{Jy}/\text{beam}$ RMS noise,
e.g. Mauch et al. (2020). Multi-hour observations using all available
bandwidth is only one factor determining image observation sensitivity —
another important factor is the number of unique interferometer spacings
(baselines) contributing to the final synthesized image. This scales, roughly,
quadratically with the number of antennae in the interferometric array.

The dump rate of MeerKAT’s raw correlated visibilities (excluding metadata) is
0.14 and 1.10 TiB / hr for coarse and fine channelization respectively, at a
typical dump rate of 8s. RFI mitigation on a large wideband sensitive array
therefore poses a significant processing challenge, requiring a flexible and
parallel implementation of mitigation routines.

##  2 RFI detection and mitigation

Outlier detection routines are commonly used to detect and mitigate RFI in
interferometry. The results are stored in a spectro-polarimetric flag array
with the same shape as the data column, which is used in further calibration
and imaging of the radio product.

Mitigation of man-made RFI is substantially simplified by two properties of
the measured signal. Firstly, transmission signals are ordinarily highly
polarized. Much of the bright cm dominating radio AGN population is linearly
polarized to a couple of percent. Secondly, a radio interferometer measures
phase information of the spatial locality of emitting sources. The complex
vector rotates quickly on long spacings for all declinations away from the
celestial poles. Through complex averaging, a substantial portion of the
detectable RFI is washed out.

Various simple and effective outlier detection techniques that can be used to
detect the remaining RFI not washed out in the fringes of the interferometer,
are discussed in Offringa et al. (2010). The SumThreshold algorithm is a
simple and fast windowed procedure widely used for RFI detection that can be
applied to two dimensions — that of time and frequency per spacing. The method
clips values where the sum of the values within a time and channel window
exceeds an interatively adjusted average threshold. The method is very
sensitive to RFI while maintaining good false-positive classification rates
for large enough window sizes. The algorithm is trivially parallel over the
unique interferometer spacings (baselines), but requires a reordering step to
select the time and frequency subsets per baseline from data that is
traditionally stored in time-sorted row order. For a large array such as
MeerKAT this step is amortized by the dominating run time of the sum-
thresholding for various window sizes and number of iterations. To aid in
background estimation on short spacings, we provide options to apply “static”
masks to flag large bands of persistent RFI (such as the one shown in Fig. 1)
on adjustable ranges of baseline length, as well as options to form residual
products based on predicted sky models.

The sizes of the windows, as well as the sensitivity adjustment parameter,
have to be fine tuned by hand for the particular science case and
channelization regime of the observation, to minimize false positives. To
achieve this we provide a YAML-based interface to the user to allow users to
customize strategies to their requirements. Typical flagging statistics and
results for MeerKAT are shown in Figure 2.

Figure 2.: Here we show bandwidth and spatial flagging statistics after
applying typical multi-pass flagging strategies for L (first row) and UHF
(second row) bands of calibrator PKS B1934-638. UHF band is generally very
clean compared to L-band, however we note that the large spread in visibility
amplitudes at the bottom of the band is due to the substantial cumulative
contribution from the off-axis AGN population, due to the wide field of view
of the instrument. The contribution from navigational systems is substantially
reduced in comparison to what was seen in Fig. 1

##  3 A parallel, scalable implementation

Our flagger is distributed as an Open Source Python package, available on the
Python Package Index as “tricolour". We have implemented the SumThreshold
method written in Numpy (Harris et al. 2020) and cache-optimized Numba (Lam et
al. 2015) kernels. The user-available version of the flagger ingests data from
Measurement Set v2.0 columns exposed as chunked Dask Arrays by Dask-ms . This
architecture is more fully described in Perkins et al. (2021a). A trimmed
version of the flagger is used in the online MeerKAT Science Data Processor
which flags most of the substantial RFI in buffered raw data at line rates,
before storage to the MeerKAT archive.

We profiled our implementation with a dual-socket Intel Xeon 8160 system with
550 GiB of memory applying a typical multi-pass L-band flagging strategy to a
sufficiently large ($>100~{}\text{GiB}$) coarsely channelized (208kHz) dataset
captured at a dump rate of 1s. After tuning thread affinity and chunk size to
optimize cache performance (through non-intrusive L3 cache profiling with the
Linux kernel utility perf) we show that our flagger scales well to more than
20 physical cores (Figure 3) and flags data at over 400 GiB/hr.

Figure 3.: Scaling performance of our flagging software after fine-tuning
cache performance through the chunk size parameters in the user interface.
Dashed line represents perfect linear scaling for reference. Overall core
occupancy is good although scaling is ultimately limited by the memory
intensity of the underlying algorithm.

##  4 Conclusions

We have built a highly configurable and scalable RFI flagger for the MeerKAT
radio telescope, that is capable of processing coarsely channelized wideband
data well above line rate with typical flagging strategies. Our flagging
software has substantially cut down on the processing time needed for
telescope commissioning, as well as recent continuum science observation
processing through the CARACal (Jozsa et al. 2021) pipeline, as well as
processing of the Mightee continuum survey (Delhaize et al. 2020) through the
Oxkat (Heywood 2020) pipeline.

While Tricolour is currently restricted to operation on a single node, Dask
offers the potential for horizontal scaling on a cluster, as described in
Perkins et al. (2021b). As we support data ingest through the generic, defacto
standard, Measurement Set v2.0 interface, modification of the flagger to
support instruments other than MeerKAT should not be difficult to accomplish.

## Acknowledgements

All figures in this paper were generated using Matplotlib - a 2D graphics
package used for Python for application development, interactive scripting,and
publication-quality image generation across user interfaces and operating
systems. The MeerKAT radio telescope is a facility operated by the South
African Radio Astronomy Observatory. A business unit of the National Research
Foundation of South Africa.

## References

  * Delhaize et al. (2020) Delhaize, J., et al. 2020, MIGHTEE: Are giant radio galaxies more common than we thought? MNRAS,accepted
  * Harris et al. (2020) Harris, C. R., et al. 2020, Nature, 585, 357
  * Heywood (2020) Heywood, I. 2020, oxkat: Semi-automated imaging of MeerKAT observations. 2009.003
  * Jonas et al. (2018) Jonas, J., et al. 2018, in MeerKAT Science: On the Pathway to the SKA (SISSA Medialab), vol. 277, 001
  * Jozsa et al. (2021) Jozsa, G., et al. 2021, in ADASS XXX, edited by J.-E. Ruiz, & F. Pierfederici (San Francisco: ASP), vol. TBD of ASP Conf. Ser., 999 TBD
  * Lam et al. (2015) Lam, S. K., et al. 2015, in Proceedings of the Second Workshop on the LLVM Compiler Infrastructure in HPC, 1
  * Mauch et al. (2020) Mauch, T., et al. 2020, The Astrophysical Journal, 888, 61
  * Offringa et al. (2010) Offringa, A., et al. 2010, Monthly Notices of the Royal Astronomical Society, 405, 155
  * Perkins et al. (2021a) Perkins, S., et al. 2021a, in ADASS XXX, edited by J.-E. Ruiz, & F. Pierfederici (San Francisco: ASP), vol. TBD of ASP Conf. Ser., 999 TBD
  * Perkins et al. (2021b) — 2021b, in ADASS XXX, edited by J.-E. Ruiz, & F. Pierfederici (San Francisco: ASP), vol. TBD of ASP Conf. Ser., 999 TBD


# Extra invariant and plasma inhomogeneity to improve zonal flow

Alexander M. Balk  missing Department of Mathematics, University of Utah, 155
South 1400 East, Salt Lake City, Utah 84112

###### Abstract

Zonal flows are known to diminish turbulent transport in magnetic fusion.
Interstingly, there is an adiabatic invariant that implies the emergence of
zonal flow. The paper shows that if this invariant is decreasing (due to some
external factors) then the emerging zonal flow is better. It is also shown
that the plasma inhomogeneity can lead to the decrease of the adiabatic
invariant. A simple condition for such decrease is found.

Fusion plasma, Transport barriers, drift/Rossby waves, Additional
conservation.

##  I Introduction

One of the problems of nuclear fusion is to decrease transport of heat and
particles from the core of a fusion device to its walls. In devices with
magnetic confinement, the decreased transport could be achieved by creating
strong alternating poloidal flow (often called “zonal” flow in analogy with
hydrodynamics). This is reviewed in Horton (1999); Diamond  _et al._ (2005).
The possibility to decrease transport is based on the idea Biglari  _et al._
(1990) that alternating zonal/poloidal flow interrupts ballistic transport
(decreases mean free path) across the flow. At the same time, the poloidal
flow does not contribute to the radial transport.

We consider zonal/poloidal flow not as the flow with exactly zero poloidal
wave number ($q=0$), but as the turbulence whose poloidal wave number is much
less than the radial wave number ($|q|\ll|p|$). The zonal flow turbulence is a
continuous continuation of other plasma turbulence with the entire spectrum of
wave vectors ${\bf k}=(p,q)$. So, zonal flow that efficiently reduces
transport, should have a large amount of large-scale energy tightly
concentrated around the $p$-axis. This energy will eventually lead to the
coherent zonal flow with exactly zero poloidal number (like the emergence of
large vortex in 2D hydrodynamics), but this process is not considered in the
present paper. The problem of interaction between the drift turbulence and the
coherent zonal flow was studied for decades and is intensively studied now,
which is reviewed in Marston and Tobias (2023).

We suppose that a significant mode of plasma dynamics represent drift waves
with dispersion relation

$\displaystyle\omega({\bf k})=\frac{\nu+\mu k^{2}}{1+k^{2}}\,q,$ (1)

where the wave vector ${\bf k}=(p,q)$ is non-dimensionalized by the Larmor
(Rossby) radius $\rho$, and $k^{2}=p^{2}+q^{2}$. The interaction of these
waves can have different forms, and we do not need to specify it. The
frequencies $\mu$ and $\nu$ can depend on the coherent part of the
poloidal/zonal flow Parker (2016); Ruiz  _et al._ (2019); Zhu and Dodin
(2021). More about this is in Section III.

The waves similar to the plasma drift waves occur in different problems. As
well known, the waves with the same dispersion relation are the Rossby waves
in atmosphere and ocean Vallis (2006); then often $\nu=\beta\rho$ ($\beta$ is
the Coriolis parameter) and $\mu=0$. The dynamics of Rossby waves leads to the
generation of zonal flows in atmospheres and oceans on several planets. An
example of this phenomenon are Jupiter’s stripes. Much slower waves (slower
than the relatively slow Rossby waves) also appear in astrophysical
magnethydrodynamics Hide (1966); Zaqarashvili  _et al._ (2007); Balk (2014,
2022). Such waves are often called the magnetic Rossby waves. In this case,
$\nu=0$ and $\mu\rho^{3}=B_{0}^{2}/\beta$ ($B_{0}$ is the basic zonal magnetic
field, measured in the velocity units). These waves occur in the upper layer
of Earth’s liquid iron outer core (“ocean of the core” Braginsky (1998)). They
also take place in Sun’s tachocline.

It is interesting that the emergence of zonal/poloidal flow follows from the
extra conservation. To describe this extra invariant, let us first recall the
question considered by L. Boltzmann Boltzmann (1875, 1876) about binary
collisions

$\displaystyle{\bf p}_{1}+{\bf p}_{2}$ $\displaystyle=$ $\displaystyle{\bf
p}_{3}+{\bf p}_{4},$ (2a) $\displaystyle{\bf p}_{1}^{2}+{\bf p}_{2}^{2}$
$\displaystyle=$ $\displaystyle{\bf p}_{3}^{2}+{\bf p}_{4}^{2}$ (2b) (${\bf
p}_{i}$ are the momentum vectors of the two particles before and after the
collision), see also Cercignani (1990); Balk and Ferapontov (1998). Does there
exist an independent quantity $\phi({\bf p})$ which is conserved in the
collisions (2ab) $\displaystyle\phi({\bf p}_{1})+\phi({\bf p}_{2})=\phi({\bf
p}_{3})+\phi({\bf p}_{4})?$ (2c) In other words, does there exist a function
$\phi({\bf p})$ such that the equation (2c) holds for any vectors ${\bf
p}_{i}\;(i=1,2,3,4)$ bound by the relations (2ab)? And the function $\phi({\bf
p})$ is not a mere linear combination of the momentum ${\bf p}$ and energy
${\bf p}^{2}$.

Boltzmann had shown that such quantity does not exist.

Similar, one can consider resonance interactions of drift/Rossby waves

$\displaystyle{\bf k}_{1}$ $\displaystyle=$ $\displaystyle{\bf k}_{2}+{\bf
k}_{3},$ (3a) $\displaystyle\omega({\bf k}_{1})$ $\displaystyle=$
$\displaystyle\omega({\bf k}_{2})+\omega({\bf k}_{3})$ (3b) (${\bf k}_{i}$ are
the wave vectors) and pose the question: Does there exist an independent
quantity $\eta({\bf k})$ which is conserved in the interactions (3ab)
$\displaystyle\eta({\bf k}_{1})=\eta({\bf k}_{2})+\eta({\bf k}_{3})?$ (3c) In
other words, does there exist a function $\eta({\bf k})$ such that the
equation (3c) holds for any vectors ${\bf k}_{i}\;(i=1,2,3)$ bound by the
relations (3ab)? And the function $\eta({\bf k})$ is not a mere linear
combination of the momentum ${\bf k}$ and energy $\omega({\bf k})$.

For the dispersion relation (1) such function does exist, namely Balk  _et
al._ (1991); Balk (1991)

$\displaystyle\eta({\bf
k})=\arctan{\frac{p+q\sqrt{3}}{k^{2}}}-\arctan{\frac{p-q\sqrt{3}}{k^{2}}}\,.$
(4)

Thus, in addition to the energy and momentum, a system of drift/Rossby waves
has an extra invariant ${\mathcal{I}}=\int\eta_{\bf k}{N}_{\bf k}\,d{\bf k}$.
The function ${N}_{\bf k}=\varepsilon_{\bf k}/\omega_{\bf k}$ is the wave
action spectrum (or the phase space density of quasi-particles),
$\varepsilon_{\bf k}$ is the energy spectrum. (To have positive ${N}_{\bf k}$,
one needs to restrict wave vectors ${\bf k}$ to half-plane $q>0$.)

The extra invariant stems from the integrability work by V. Zakharov and E.
Shulman, published in a series of papers throughout 1980s (in particular,
Zakharov and Schulman (1980, 1988)). They found that conservation in resonance
interactions implies cancellation of small divisors, and therefore, leads to
an additional approximate invariant of the corresponding Hamiltonian system.
An important feature of their theory is that one can make physical predictions
without knowing the exact form of nonlinearity. It is only required that the
wave dynamics is Hamiltonian (which is usually the case for a physical
system). Some aspects of this theory turned out Ferapontov (1992) to be
related to the web geometry Blaschke (1932). This allowed to show Balk and
Ferapontov (1998) that the Rossby wave extra invariant $\mathcal{I}$ (or the
extra conserved quantity $\eta$) is unique: All invariants are linear
combinations of the energy, the enstrophy, and the extra invariant
$\mathcal{I}$.

The extra conservation implies the emergence of zonal flow Balk (2005), see
also Section II. Here we will see even more: If the extra invariant is
decreasing with time (due to some external factors), then the emerging zonal
flow is better; it diminishes transport even more efficiently.

The extra invariant and its physical implications were established through the
work of several people, in particular, Galperin  _et al._ (2006); Weichman
(2006); Balk and van Heerden (2006); Sukoriansky  _et al._ (2007); Lee and
Smith (2007); Read  _et al._ (2007); Yoshi-Yuki  _et al._ (2007); Balk and
Yoshikawa (2009); Nazarenko and Quinn (2009); Balk and Zakharov (2009); Balk
_et al._ (2011); Weichman (2012); Saito and Ishioka (2013, 2016). In Nazarenko
and Quinn (2009), the adiabatic invariant $I$ was called “zonostrophy”
(“zonostrophic instability” is unrelated).

##  II Decreasing extra invariant leads to efficient zonal flow.

First, let us note the energy accumulation at large scales. This is obvious
for Rossby waves when $\mu=0$ and $\nu\neq 0$; then the energy follows the
inverse cascade. But even in general, the energy accumulation at large scales
appears to be true. For instance, in the classic 3-dimensional hydrodynamic
turbulence, the energy follows the direct cascade, but the large scales carry
most of the energy (the Kolmogorov spectrum has infrared divergence). Similar
situation occurs for the gravity waves on the ocean surface Zakharov (1985).
This system conserves two integrals: the energy and the wave action. The wave
action follows the inverse cascade, while the energy follows the direct
cascade. Still, most of the energy is carried by large scales (the Kolmogorov-
Zakharov energy spectrum, determined by the wave action flux, has infrared
divergence). In both of these situations, most of the energy flows towards
large $k$ and dissipates there. However, a small fraction of the energy
reaches large scales, and the energy accumulates there. Similar phenomenon
seems to hold Balk (2014, 2022) for the turbulence of magnetic Rossby waves,
when $\nu=0$ and $\mu\neq 0$. Then the energy and the enstrophy trade places;
and so, there is direct cascade of energy, but the energy accumulates at small
$k$. When both parameters $\mu$ and $\nu$ in (1) are non-zero, then there is
no clear direction for the energy cascade, but presumably, the energy
accumulates at large scales.

Now recall the well known conserved quantities for the system of interacting
drift waves. These are the energy $\mathcal{E}$ and the zonal momentum
$\mathcal{M}$

$\displaystyle\mathcal{E}=\int\omega_{\bf k}\;N_{\bf k}\;d{\bf
k},\qquad\mathcal{M}=\int q\;N_{\bf k}\;d{\bf k}.$ (5)

The enstrophy is a linear combination of $\mathcal{E}$ and $\mathcal{M}$. The
$p$-momentum is not a real physical invariant because the function
$p/\omega_{\bf k}$ has singularity (when $q\rightarrow 0$), see Balk and
Yoshikawa (2009) for details.

To show the emergence of zonal flow, we need to consider the following linear
combination of the extra invariant $\mathcal{I}$ with $\mathcal{E}$ and
$\mathcal{M}$
$\displaystyle\tilde{\mathcal{I}}=\mathcal{I}-2\sqrt{3}\,\frac{\mathcal{E}-\mu\,\mathcal{M}}{\nu-\mu}=\int\tilde{\eta}_{\bf
k}\,N_{\bf k}\,d{\bf k},$ (6a) where $\displaystyle\tilde{\eta}_{\bf
k}=\eta_{\bf k}-2\sqrt{3}\,\frac{\omega_{\bf k}-\mu\,q}{\nu-\mu}=\eta_{\bf
k}-2\sqrt{3}\,\frac{p}{1+k^{2}}.$ (6b) Like $\eta_{\bf k}$, the function
$\tilde{\eta}_{\bf k}$ is independent of $\mu$ and $\nu$.

Let us write the modified extra invariant (6) in terms of the energy spectrum:
$\tilde{\mathcal{I}}=\int R_{\bf k}\,\varepsilon_{\bf k}\;d{\bf k}$, where the
ratio $R_{\bf k}=\tilde{\eta}_{\bf k}/\omega_{\bf k}$ shows how much extra
invariant $\tilde{\mathcal{I}}$ is carried by the unit amount of energy with
wave vector ${\bf k}=(p,q)$. The contour plot of $R_{\bf k}$ is presented in
figure 1.

Figure 1: Contour plot of the ratio $R_{\bf k}=\tilde{\eta}_{\bf
k}/\omega_{\bf k}$ in logarithmic scale: The color represents the values of
$\ln(R_{\bf k})$. The figure is qualitatively the same for various value of
the parameters $\mu$ and $\nu$ in (1); for this particular plot,
$\mu=1,\nu=2$. The graph is shown only for positive $p,q$, due to the
symmetries $R(p,q)=R(-p,q)=R(p,-q)$.

Figure 1 shows that the accumulated large scale energy should concentrate near
the $p$-axis (poloidal/zonal flow). Indeed, if during the energy transfer
towards large scales, the energy accumulates away from $p$-axis, then the
extra invariant $\tilde{\mathcal{I}}$ would significantly increase (in
contradiction to the extra conservation). This is because — according to the
figure — $R_{\bf k}$ is small at large $k$ but large at small $k$, unless $q$
is close to zero. This is also supported by the asymptotic behavior:

$\displaystyle\frac{(\nu+\mu k^{2})R_{\bf
k}}{8\sqrt{3}}=\left\\{\begin{array}[]{ll}\frac{q^{2}(5p^{2}+q^{2})}{5k^{8}}+O(k^{-6})&(k\rightarrow\infty),\\\
\\\ \frac{q^{2}}{p^{2}(1+p^{2})^{2}}+O(q^{4})&(q\rightarrow
0).\end{array}\right.\quad$ (10)

Now suppose, the extra invariant $\tilde{I}$ is decreasing because of some
external factors, but the energy $\mathcal{E}$ is not decreasing. Then a
lesser amount of the extra invariant is available for the spreading of the
energy away from the $p$-axis, and the energy has to concentrate more tightly
around $p$-axis. This means that the emerging zonal flow is more
unidirectional, and it contributes less to the radial transport. In addition,
the zonal flow interrupts the radial transport, decreasing the mean free path.
It is worthwhile to note that the transport decrease occurs for transport by
any plasma mode (not just by drift waves).

Actually, one needs to decrease the extra invariant as much as possible,
without decreasing the energy. This is relatively easy because the energy
kernel (1) and the extra invariant kernel (4) are quite different.

There are at least two ways to decrease the extra invariant without decreasing
the energy.

First, one can introduce some dissipation and pumping in different areas of
the ${\bf k}$-plane, so that the extra invariant $\tilde{\mathcal{I}}$ is
dissipated, while the energy $\mathcal{E}$ is pumped. This approach was
studied in Balk (2019).

The present paper reports another opportunity which is due to the
inhomogeneity of fusion plasma.

##  III Evolution of the invariant by the inhomogeneous wave kinetic equation

Let us describe the turbulence of the drift/Rossby waves by the wave action
spectrum $N(p,q,x,y,t)$ obeying the wave kinetic equation

$\displaystyle\frac{\partial{N}}{\partial
t}+\frac{\partial(\omega,N)}{\partial({\bf k},{\bf x})}=St[{N}]\,,$ (11)

where we use the notation

$\displaystyle\frac{\partial({\mathcal{F}},{\mathcal{G}})}{\partial(u,v)}=\frac{\partial{\mathcal{F}}}{\partial
u}\cdot\frac{\partial{\mathcal{G}}}{\partial
v}\;-\;\frac{\partial{\mathcal{F}}}{\partial
v}\cdot\frac{\partial{\mathcal{G}}}{\partial u}$

for arbitrary functions ${\mathcal{F}}$, ${\mathcal{G}}$ of arbitrary
variables $u$, $v$ (vector or scalar). The l.h.s. in equation (11) describes
the slow refraction of waves, while the r.h.s. (stoss-term or collision
integral) describes changes in $N$ due to the 3-wave resonance interactions
(3). We assume that the interactions occur locally in physical space, and so,
the three interacting waves have the same plasma parameters (the inhomogeneity
does not enter $St[N]$).

Let us consider the evolution by the equation (11) of the integral

$\displaystyle\tilde{\mathcal{I}}=\int\tilde{\eta}(p,q)\;N(p,q,x,y,t)\;dp\,dq.$
(12)

We multiply equation (11) by $\tilde{\eta}$ and integrate over $d{\bf k}$.
Since the quantity $\tilde{\eta}$ is conserved in the 3-wave resonance
interactions, the term $St[N]\tilde{\eta}$ integrates to zero. The l.h.s. in
(11) leads to the flowing of $\tilde{I}$ in the physical plane (the divergence
of the flux) and to a source/sink term

$\displaystyle\frac{\partial{\tilde{\mathcal{I}}}}{\partial
t}+\frac{\partial}{\partial{\bf x}}\cdot\int\frac{\partial\omega}{\partial{\bf
k}}\,{\tilde{\eta}}\,N\,d{\bf k}=\int S\,N\,d{\bf k}$ (13)

with the kernel

$\displaystyle S=\frac{\partial(\omega,{\tilde{\eta}})}{\partial({\bf k},{\bf
x})}=-\frac{\partial\omega}{\partial x}\frac{\partial{\tilde{\eta}}}{\partial
p}=$
$\displaystyle\frac{p\,(\nu^{\prime}+k^{2}\mu^{\prime})\;\;16\sqrt{3}\,(1+4p^{2})\,q^{4}}{(1+k^{2})^{3}[(p+\sqrt{3}q)^{2}+k^{4}][(p-\sqrt{3}q)^{2}+k^{4}]}\,,$
(14)

where the frequencies $\mu$ and $\nu$ in the dispersion relation (1) slowly
depend on the radial coordinate $x$; prime denotes the $x$-derivative.

If the function $\tilde{\eta}$ were replaced by $\omega$, then the last
integral in (13) would vanish, and we would have the energy conservation
(energy flows throughout the $xy$-plane). So, the inhomogeneous part of the
kinetic equation can decrease the extra invariant $\tilde{\mathcal{I}}$, while
the energy $\mathcal{E}$ is always conserved by the entire kinetic equation.

The parameters of the dispersion relation (1) depend on the zonal flow
velocity $U(x)$ Parker (2016); Ruiz  _et al._ (2019); Zhu and Dodin (2021)

$\displaystyle\mu=U(x)/\rho,\qquad\nu=\rho[\kappa+U^{\prime\prime}(x)]+U(x)/\rho\,.\quad$
(15)

The radial $x$ and poloidal $y$ coordinates are considered in generalized
sense: Radial direction is the direction of the anti-gradient of some quantity
$Q$, the poloidal direction is orthogonal to the radial direction and to the
confining magnetic field (the contour lines $Q=$const are not necessarily
circles). The constant $\kappa$ is determined by the local value of the
gradient $\nabla Q$ ($\kappa$ is analogous to the $\beta$-parameter in
hydrodynamics).

The $U^{\prime\prime}$ term in the second expression (15) is often neglected,
but recently its significance is re-considered Parker (2016); Ruiz  _et al._
(2019); Zhu and Dodin (2021).

We see from (13)-(III) that the invariant $\tilde{I}$ would decrease if the
drift wave turbulence spectrum had significant asymmetric part

$\displaystyle N_{a}(p,q,x,y,t)=N(p,q,x.y,t)-N(-p,q,x,y,t)$

and

$\displaystyle[\rho^{2}U^{\prime\prime\prime}+(1+k^{2})U^{\prime}]N_{a}<0$
(16)

(most of the factors in (III) are positive and cancel out). The condition (16)
is especially simple if the $U^{\prime\prime}$ term in (15) is neglected. Then
$\tilde{I}$ decreases when $U^{\prime}N_{a}<0$. This means: If zonal flow
velocity $U(x)$ is decreasing, the turbulence spectrum should be skewed
towards positive $p$ (the spectrum $N$ should be bigger for positive $p$ than
for negative $p$). If zonal flow velocity $U(x)$ is increasing, the turbulence
spectrum should be skewed towards negative $p$. In general, the condition (16)
involves the length scale. The condition (16) can be used to create a
transport barrier at a certain location.

##  IV Conclusion

The dynamics of interacting drift waves has an adiabatic invariant; Section I.
The presence of this extra invariant implies the emergence of zonal flow,
which diminishes the transport of heat and particles (by any plasma mode, not
just by drift waves), and so, serves as transport barrier. The present paper
makes two points:

1). If the extra invariant is decreasing (due to some external factors), while
the energy is not decreasing, then the emerging zonal flow is a better
transport barrier; Section II. The extra invariant kernel so much differs from
the energy kernel, that it is relatively easy to decrease the extra invariant
without decreasing the energy.

2). The plasma inhomogeneity can decrease the extra invariant, while
preserving the energy; Section III. To see this, we have applied the
inhomogeneous wave kinetic equation (the three waves involved in the collision
integral having the same plasma parameters) and derived a simple condition
(16) for the decrease of the extra invariant.

## References

  * Horton (1999) W. Horton, Reviews of Modern Physics 71, 735 (1999).
  * Diamond  _et al._ (2005) P. H. Diamond, S.-I. Itoh, K. Itoh, and T. S. Hahm, Plasma Phys. Control. Fusion 47, R35 (2005).
  * Biglari  _et al._ (1990) H. Biglari, P. Diamond, and P. Terry, Physics of Fluids B: Plasma Physics 2, 1 (1990).
  * Marston and Tobias (2023) J. Marston and S. Tobias, Annual Review of Fluid Mechanics 55, 351 (2023), https://doi.org/10.1146/annurev-fluid-120720-031006 .
  * Parker (2016) J. B. Parker, Journal of Plasma Physics 82, 595820602 (2016).
  * Ruiz  _et al._ (2019) D. Ruiz, M. Glinsky, and I. Dodin, Journal of Plasma Physics 85, 1 (2019).
  * Zhu and Dodin (2021) H. Zhu and I. Y. Dodin, Physics of Plasmas 28, 032303 (2021).
  * Vallis (2006) G. K. Vallis,  _Atmospheric and Oceanic Fluid Dynamics. Fundamentals and Large-scale Circulation_ (Cambridge, 2006).
  * Hide (1966) R. Hide, Proc. Roy. Soc. London 259, 615 (1966).
  * Zaqarashvili  _et al._ (2007) T. V. Zaqarashvili, R. Oliver, J. L. Ballester, and B. M. Shergelashvili, A&A 470, 815 (2007).
  * Balk (2014) A. M. Balk, ApJ 796, 143 (8pp) (2014).
  * Balk (2022) A. M. Balk, The Astrophysical Journal 926, 2 (2022).
  * Braginsky (1998) S. I. Braginsky, Earth Planets Space 50, 641 (1998).
  * Boltzmann (1875) L. Boltzmann, Wien. Ber. 72, 427 (1875).
  * Boltzmann (1876) L. Boltzmann, Sitzungsberichte der academie der wissenschaften wien 74, 503 (1876).
  * Cercignani (1990) C. Cercignani, J. Stat. Phys. 58, 817 (1990).
  * Balk and Ferapontov (1998) A. M. Balk and E. V. Ferapontov, in  _Nonlinear waves and weak turbulence_ , edited by V. E. Zakharov (Amer. Math. Soc. Trans. Ser. 2, vol. 182, 1998) pp. 1–30.
  * Balk  _et al._ (1991) A. M. Balk, S. V. Nazarenko, and V. E. Zakharov, Phys. Lett. A 152, 276 (1991).
  * Balk (1991) A. M. Balk, Phys. Lett. A 155, 20 (1991).
  * Zakharov and Schulman (1980) V. E. Zakharov and E. I. Schulman, Physica D 1, 192 (1980).
  * Zakharov and Schulman (1988) V. E. Zakharov and E. I. Schulman, Physica D 29, 283 (1988).
  * Ferapontov (1992) E. V. Ferapontov, “Web geometry and mathematical physics,” (Kluwer Academic Publishers, Norwell, MA, 1992) pp. 310–323, appendix in the book “Geometry and Algebra of Multidimensional Three-Webs” by M. A. Akivis and A. M. Shelekhov.
  * Blaschke (1932) W. Blaschke,  _Topological differential geometry_ (Chicago University Press, 1932).
  * Balk (2005) A. M. Balk, Phys. Lett. A 345, 154 (2005).
  * Galperin  _et al._ (2006) B. Galperin, S. Sukoriansky, N. Dikovskaya, P. L. Read, Y. H. Yamazaki, and R. Wordsworth, Nonlin. Processes Geophys. 13, 89 (2006).
  * Weichman (2006) P. Weichman, Physical review. E, Statistical, nonlinear, and soft matter physics 73, 036313 (2006).
  * Balk and van Heerden (2006) A. M. Balk and F. van Heerden, Physica D 223, 109 (2006).
  * Sukoriansky  _et al._ (2007) S. Sukoriansky, N. Dikovskaya, and B. Galperin, Journal of the Atmospheric Sciences 64, 3312 (2007).
  * Lee and Smith (2007) Y. Lee and L. M. Smith, J. Fluid Mech. 576, 405 (2007).
  * Read  _et al._ (2007) P. L. Read, Y. H. Yamazaki, S. R. Lewis, P. D. Williams, R. D. Wordsworth, K. Miki-Yamazaki, J. Sommeria, and H. Didelle, Journal of the Atmospheric Sciences 64, 4031 (2007).
  * Yoshi-Yuki  _et al._ (2007) H. Yoshi-Yuki, S. Nishizawa, S.-I. Takehiro, M. Yamada, K. Ishioka, and S. Yoden, Journal of The Atmospheric Sciences - J ATMOS SCI 64, 4246 (2007).
  * Balk and Yoshikawa (2009) A. M. Balk and T. Yoshikawa, Physica D 238, 384 (2009).
  * Nazarenko and Quinn (2009) S. Nazarenko and B. Quinn, Phys. Rev. Lett. 103, 118501 (2009).
  * Balk and Zakharov (2009) A. M. Balk and V. E. Zakharov, Phys. Lett. A 373, 4049 (2009).
  * Balk  _et al._ (2011) A. M. Balk, F. van Heerden, and P. B. Weichman, Phys. Rev. E 83, 046320 (2011).
  * Weichman (2012) P. Weichman, Physical review letters 109, 235002 (2012).
  * Saito and Ishioka (2013) I. Saito and K. Ishioka, Phys. Fluids 25, 076602 (2013).
  * Saito and Ishioka (2016) I. Saito and K. Ishioka, J. Met. Soc. Japan. Ser. II 94, 25 (2016).
  * Zakharov (1985) V. E. Zakharov, in  _Basic plasma physics_ , Vol. 2, edited by A. A. Galeev and R. N. Sudan (Elsevier (North-Holland Publishing Company), Amsterdam, 1985) pp. 3–36.
  * Balk (2019) A. M. Balk, Phys. Rev. Res. 1, 033180 (2019).


# Learning Perception and Planning with Deep Active Inference

###### Abstract

Active inference is a process theory of the brain that states that all living
organisms infer actions in order to minimize their (expected) free energy.
However, current experiments are limited to predefined, often discrete, state
spaces. In this paper we use recent advances in deep learning to learn the
state space and approximate the necessary probability distributions to engage
in active inference.

Index Terms—  active inference, deep learning, perception, planning

##  1 Introduction

Active inference postulates that action selection in biological systems, in
particular the human brain, is actually an inference problem where agents are
attracted to a preferred prior state distribution in a hidden state space [1].
To do so, each living organism builds an internal generative model of the
world, by minimizing the so-called free energy. The idea of active inference
stems from neuroscience [2, 3] and has already been adopted to solve different
control and learning tasks [4, 5, 1]. These experiments however typically make
use of manually engineered state transition models and predefined, often
discrete, state spaces.

In this paper we show that we can also learn the state space and state
transition model, by using deep neural networks as probability density
estimators. By sampling from the learnt state transition model, we can plan
ahead minimizing the expected free energy, trading off goal-directed behavior
and uncertainty-resolving behavior.

##  2 Active Inference

Active inference states that every organism or agent entertains an internal
model of the world, and implicitly tries to minimize the difference between
what it believes about the world and what it perceives, hence minimizing its
own variational free energy [2]. Moreover, the agent believes that it will
minimize its expected free energy in the future, in effect turning action
selection into an inference problem. This boils down to optimizing for two
distinct objectives. On the one hand the agent actively samples the world to
update its internal model of the world and better explain observations. On the
other hand the agent is driven to visit preferred states that it believes a
priori it will visit—a kind of global prior— which carry little expected free
energy.

Formally, we assume an agent entertains a generative model
$P(\tilde{{\bm{o}}},\tilde{{\bm{s}}},\tilde{{\bm{a}}},\pi)$ of the
environment, which specifies the joint probability of observations, actions
and their hidden causes, where actions are determined by some policy $\pi$. If
the environment is modelled as a Markov Decision Process (MDP) this generative
model factorizes as:

$\small\begin{split}P(\tilde{{\bm{o}}},\tilde{{\bm{s}}},\tilde{{\bm{a}}},\pi)&=P(\pi)P({\bm{s}}_{0})\\\
&\quad\prod_{t=1}^{T}P({\bm{o}}_{t}|{\bm{s}}_{t})P({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})P({\bm{a}}_{t-1}|\pi)\end{split}$
(1)

The free energy is then defined as:

$\begin{split}F&=\mathbb{E}_{Q}[\log Q(\tilde{{\bm{s}}})-\log
P(\tilde{{\bm{o}}},\tilde{{\bm{s}}},\tilde{{\bm{a}}})]\\\
&=D_{\mathrm{KL}}(Q(\tilde{{\bm{s}}})\|P(\tilde{{\bm{s}}},\tilde{{\bm{a}}}|\tilde{{\bm{o}}}))-\log
P(\tilde{{\bm{o}}})\\\
&=D_{\mathrm{KL}}(Q(\tilde{{\bm{s}}})\|P(\tilde{{\bm{s}}},\tilde{{\bm{a}}}))-\mathbb{E}_{Q}[\log
P(\tilde{{\bm{o}}}|\tilde{{\bm{s}}})]\end{split}$ (2)

where $Q(\tilde{{\bm{s}}})$ is an approximate posterior distribution. The
second equality shows that the free energy is minimized when the KL divergence
term becomes zero, meaning that the approximate posterior becomes the true
posterior, in which case the free energy becomes the negative log evidence.
This can also be rewritten as the third equality, which is the negative
evidence lower bound (ELBO) that also appears in the variational autoencoders
(VAE) framework [6, 7].

In active inference, agents infer the actions that will result in visiting
states of low expected free energy. They do this by sampling actions from a
prior belief about policies according to how much expected free energy that
policy will induce. Formally, this means that the probability of picking a
policy is given by [8]:

$\begin{split}P(\pi)&=\sigma(-\gamma G(\pi))\\\
G(\pi)&=\sum_{\tau}G(\pi,\tau)\end{split}$ (3)

where $\sigma$ is the softmax function with precision parameter $\gamma$,
which governs the agent’s goal-directedness and randomness in its behavior.
Here $G$ is the expected free energy at future timestep $\tau$ when following
policy $\pi$ [8]:

$\displaystyle\begin{split}G(\pi,\tau)&=\mathbb{E}_{Q({\bm{o}}_{\tau},{\bm{s}}_{\tau}|\pi)}[\log
Q({\bm{s}}_{\tau}|\pi)-\log P({\bm{o}}_{\tau},{\bm{s}}_{\tau}|\pi)]\\\
&=\mathbb{E}_{Q({\bm{o}}_{\tau},{\bm{s}}_{\tau}|\pi)}[\log
Q({\bm{s}}_{\tau}|\pi)-\log P({\bm{o}}_{\tau}|{\bm{s}}_{\tau},\pi)\\\
&\quad-\log P({\bm{s}}_{\tau}|\pi)]\\\
&=D_{\mathrm{KL}}(Q({\bm{s}}_{\tau}|\pi)\|P({\bm{s}}_{\tau}))+\mathbb{E}_{Q({\bm{s}}_{\tau})}[H({\bm{o}}_{\tau}|{\bm{s}}_{\tau})]\end{split}$
(4)

We used
$Q({\bm{o}}_{\tau},{\bm{s}}_{\tau}|\pi)=P({\bm{o}}_{\tau}|{\bm{s}}_{\tau})Q({\bm{s}}_{\tau}|\pi)$
and that the prior probability $P({\bm{s}}_{\tau}|\pi)$ is given by a
preferred state distribution $P({\bm{s}}_{\tau})$. This results into two
terms: a KL divergence term between the predicted states and the prior
preferred states, and an entropy term reflecting the expected ambiguity under
predicted states. Action selection in active inference thus entails:

  1. 1.

Evaluate $G(\pi)$ for each policy $\pi$

  2. 2.

Calculate the belief $P(\pi)$ over policies

  3. 3.

Infer the next action using $P(\pi)P({\bm{a}}_{t+1}|\pi)$

##  3 Deep Active Inference

For the agent’s model, we use deep neural networks to parameterize the various
factors of equation (1): i.e. the transition model
$p_{\theta}({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t})$ and the likelihood
distribution $p_{\xi}({\bm{o}}_{t}|{\bm{s}}_{t})$. Also the approximate
posterior is parameterized by a neural network
$p_{\phi}({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t},{\bm{o}}_{t})$. All
distributions are parameterized as i.i.d multivariate Gaussian distributions,
i.e. the outputs of the neural networks are the means $\mu$ and standard
deviations $\sigma$ of each Gaussian. Sampling is done using the
reparameterization trick, computing $\mu+\epsilon\sigma$ with $\epsilon\sim
N(0,1)$, which allows for backpropagation of the gradients. Minimizing the
free energy then boils down to minimizing the following loss function:

$\displaystyle\begin{split}\forall
t&:\underset{\phi,\theta,\xi}{\text{minimize}}:-\log
p_{\xi}({\bm{o}}_{t}|{\bm{s}}_{t})\\\
&+D_{\mathrm{KL}}(p_{\phi}({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1},{\bm{o}}_{t})\|p_{\theta}({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1}))\end{split}$
(5)

Figure 1 shows an overview of the information flow between the transition
model, approximate posterior and likelihood neural networks. To engage in
active inference using these models, we need to estimate $G(\pi,\tau)$, which
involves estimating $Q({\bm{s}}_{\tau}|\pi)$. As our model takes a state
sample as input, and only estimates the state distribution of the next
timestep, the only way to get an estimate of the state distribution at a
future timestep $\tau>t+1$ is by Monte Carlo sampling. Concretely, to infer
$P(\pi)$, we sample for each policy $N$ trajectories following $\pi$ using
$p_{\theta}$. This results in $N$ state samples $\hat{{\bm{s}}}_{\tau}$, for
which we can get $N$ observation estimates $\hat{{\bm{o}}}_{\tau}$ via
$p_{\xi}$. To be able to calculate the KL divergence and entropy, we use a
Gaussian distribution fitted on the samples’ mean and variance. We then
estimate the expected free energy for each policy from current timestep $t$
onward as follows:

$\displaystyle\hat{G_{t}}(\pi)$
$\displaystyle=\\!\sum_{\tau=t+1}^{t+K}\\!D_{\mathrm{KL}}(\mathcal{N}(\mu_{\hat{{\bm{s}}}_{\tau}},\sigma_{\hat{{\bm{s}}}_{\tau}})\|P({\bm{s}}_{\tau}))+\frac{1}{\rho}H(\mathcal{N}(\mu_{\hat{{\bm{o}}}_{\tau}},\sigma_{\hat{{\bm{o}}}_{\tau}}))$
$\displaystyle\;\;+\sum_{\pi^{\prime}}\sigma(-\gamma\hat{G}_{t+K}(\pi^{\prime}))\hat{G}_{t+K}(\pi^{\prime})$
(6)

The first summation term looks $K$ timesteps ahead, calculating the KL
divergence between expected and preferred states and the entropy on the
expected observations. We also introduce an additional hyperparameter $\rho$,
which allows for a trade-off between reaching preferred states on the one hand
and resolving uncertainty on the other hand. The second summation term implies
that after $K$ timesteps, we continue to select policies according to their
expected free energy, hence recursively re-evaluating the expected free energy
of each policy at timestep $t+K$. In practice, we unroll this $D$ times,
resulting into a search tree with an effective planning horizon of $T=K\times
D$.

Fig. 1: We train simultaneously a transition model
$p_{\theta}({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t})$, an approximate
posterior distribution model
$p_{\phi}({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t},{\bm{o}}_{t})$, and a
likelihood distribution $p_{\xi}({\bm{o}}_{t}|{\bm{s}}_{t})$ model, by
minimizing the variational free energy.

##  4 Experiments

We experiment with the Mountain Car problem, where an agent needs to drive the
car up the mountain in 1D by throttling left or right, as shown on Figure 2.
The top of the mountain can only be reached by first building up momentum
before throttling right. The agent spawns at a random position, and only
observes a noisy position sensor and has no access to its current velocity. At
each timestep, the agent can choose between two policies: $\pi_{l}$ to
throttle to the left, $\pi_{r}$ to throttle to the right. We experiment with
two flavors of this environment: one where the agent starts with fixed zero
velocity, and one where the agent starts with a random initial velocity.

Fig. 2: The mountain car environment. The shown position of the car at $-0.5$
is the starting position in our evaluations.

For our generative model, we instantiate
$p_{\theta}({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t})$,
$p_{\phi}({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t},{\bm{o}}_{t})$ and
$p_{\xi}({\bm{o}}_{t}|{\bm{s}}_{t})$ as fully connected neural networks with
20 hidden neurons, and a state space with 4 dimensions. To bootstrap the
model, we train on actions and observation of a random agent minimizing the
loss function in Equation (5) using stochastic gradient descent. Next, we
instantiate an active inference agent that uses Equation (3) to plan ahead and
select the best policy. As preferred state distribution $P({\bm{s}}_{\tau})$,
we manually drive the car up the mountain and evaluate the model’s posterior
state at the end of the sequence $\hat{{\bm{s}}}_{end}$, and set
$P({\bm{s}}_{\tau})=\mathcal{N}(\hat{{\bm{s}}}_{end},1)$. To limit the
computations, the active inference agent plans ahead for 90 timesteps,
allowing to switch policy every 30 timesteps, effectively evaluating a search
tree with depth 3, using 100 samples for each policy ($K=30,D=3,N=100$).

Figure 3 shows the sampled trajectories for all branches of the search tree,
in the case the model is bootstrapped with only a single observation at
position $-0.5$. This is a challenging starting position as the car needs
enough momentum in order to reach up the hill from there. In the case of a
random starting velocity, the generative model is not sure about the velocity
after only the first observation. This is reflected by the entropy (i.e. the
expected ambiguity) of the sampled trajectories. Now following $\pi_{r}$ from
the start will sometimes reach the preferred state, depending on the initial
velocity. In this case the active inference agent’s behavior is determined by
the parameter $\rho$. For $\rho>1$, the agent will act greedily, preferring
the policy that has a chance of reaching the top early, cf. Figure 3e. When
setting $\rho<<1$, the entropy term will play a bigger role, and the agent
will select the policy that is less uncertain about the outcomes, rendering a
more cautious agent that prefers a more precise and careful policy, moving to
the left first – see Fig. 3a. We found setting $\rho=0.1$ yields a good trade-
off for the mountain car agent.

In the environment with no initial velocity, the transition model learnt by
the agent is quite accurate and the entropy terms are an order of magnitude
lower, as shown in Figure 4. However, in terms of preferred state the lowest
KL is still achieved by following $\pi_{r}$. This is due to the fact that the
KL term is evaluated each timestep, and moving to the left, away from the
preferred state in the sequence outweighs the benefit of reaching the
preferred state in the end. Choosing $\rho=0.1$ again forces the agent to put
more weight on resolving uncertainty, preferring the policy in Figure 4a.

(a) left-right-right

(b) left-right-left

(c) left-left-right

(d) left-left-left

(e) right-right-right

(f) right-right-left

(g) right-left-right

(h) right-left-left

Fig. 3: Depending on the random initial velocity, the car will reach the hill
fast using the right policy only part of the cases (e), however starting with
the left policy first also reaches the hill top and with lower entropy on the
trajectories (a). A greedy agent ($\rho>1$) will pick (e) whereas a cautious
agent ($\rho<<1$) will favor (a). For each policy we report the values of KL,
H and G for $\rho=0.1$.

(a) left-right-right

(b) left-right-left

(c) left-left-right

(d) left-left-left

(e) right-right-right

(f) right-right-left

(g) right-left-right

(h) right-left-left

Fig. 4: When the environment starts the car with a fixed zero velocity, the
model is on average much more certain on the predicted trajectories, resulting
in lower entropy terms. However, policy (e) still achieves the lowest KL
value, as this term is evaluated each timestep, and moving away from the
preferred state yields a high KL penalty. When choosing $\rho=0.1$, the agent
again favors (a). For each policy we report the values of KL, H and G for
$\rho=0.1$.

##  5 Discussion

Using deep neural networks to instantiate the generative model and to
approximate both prior and posterior distributions, has the advantage that the
generative model is independent of any state representation. The model can
learn the best state representation for the observed data. Employing deep
neural networks also opens up the possibility of using high-dimensional sensor
inputs, e.g. images. The downside of our model, however, is the required
sampling step, which means that a distribution is only calculated for the next
timestep, and distributions for timesteps $\tau$ further in the future can
only be approximated by sampling.

Another point of discussion is the definition of the preferred state
distribution. In our case we opted for a Gaussian state distribution, centered
around the state visited by an expert demonstration, similar to our earlier
work [9]. However, the standard deviation of this distribution will determine
the absolute value of the KL term in Equation (3). A small standard deviation
will blow up the KL term, completely ignoring the entropy term. A large
standard deviation will assign probability mass to neighboring states,
possibly introducing local optima that don’t reach the actual goal state. To
mitigate this, we introduced an additional $\rho$ parameter that balances risk
and ambiguity.

Finally, planning by generating and evaluating trajectories of the complete
search tree is computationally expensive. In this paper, we intentionally
pursued this approach in order to directly investigate the effect of the KL
term versus the entropy term. To mitigate the computational load, one might
amortize the resulting policy by training a policy neural network
$p_{\pi}({\bm{a}}_{t}|{\bm{s}}_{t})$ based on the visited states and the
planned actions by the agent, similar to [10]. In other approaches, the policy
is learned directly through end-to-end training. For example, K. Ueltzhöffer
uses evolution strategies to learn both a model and a policy, that requires
part of the state space to be fixed, to contain information of the preferred
state [11]. B. Millidge on the other hand amortizes the expected free energy
$G$ as function of the state, similar to value function estimation in
reinforcement learning [12]. Again, however, the perception part is omitted
and the state space is fixed upfront.

##  6 Conclusion

In this paper, we have shown how generative models parameterized by neural
networks are trained by minimizing the free energy, and how these can be
exploited by an active inference agent to select the optimal policy. We will
further extend our models to work in more complex environments, in particular
towards more complex sensory inputs such as camera, lidar or radar data.

## References

  * [1] Karl Friston, Thomas FitzGerald, Francesco Rigoli, Philipp Schwartenbeck, and Giovanni Pezzulo,  “Active inference: A Process Theory,”  Neural Computation, vol. 29, pp. 1–49, 2017.
  * [2] Karl Friston,  “The free-energy principle: A unified brain theory?,”  Nature Reviews Neuroscience, vol. 11, no. 2, pp. 127–138, 2010\.
  * [3] Karl Friston, James Kilner, and Lee Harrison,  “A free energy principle for the brain,”  Journal of Physiology Paris, vol. 100, no. 1-3, pp. 70–87, 2006\.
  * [4] Karl Friston, Spyridon Samothrakis, and Read Montague,  “Active inference and agency: Optimal control without cost functions,”  Biological Cybernetics, vol. 106, no. 8-9, pp. 523–541, 2012.
  * [5] Karl Friston, Philipp Schwartenbeck, Thomas FitzGerald, Michael Moutoussis, Timothy Behrens, and Raymond J. Dolan,  “The anatomy of choice: active inference and agency,”  Frontiers in Human Neuroscience, vol. 7, no. September, pp. 1–18, 2013.
  * [6] Diederik P. Kingma and Max Welling,  “Auto-encoding variational bayes,”  CoRR, vol. abs/1312.6114, 2013.
  * [7] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra,  “Stochastic backpropagation and approximate inference in deep generative models,”  in Proceedings of the 31st International Conference on Machine Learning, 2014.
  * [8] Philipp Schwartenbeck, Johannes Passecker, Tobias Hauser, Thomas H B FitzGerald, Martin Kronbichler, and Karl J Friston,  “Computational mechanisms of curiosity and goal-directed exploration,”  bioRxiv, p. 411272, 2018.
  * [9] Cedric De Boom, Tim Verbelen, and Bart Dhoedt,  “Deep active inference for state estimation by learning from demonstration,”  in Conference of Complex Systems (CCS), Satellite Symposium on Complexity from Cells to Consciousness: Free Energy, Integrated Information, and Epsilon Machines, 2018.
  * [10] Ozan Catal, Johannes Nauta, Tim Verbelen, Pieter Simoens, and Bart Dhoedt,  “Bayesian policy selection using active inference,”  in Workshop on “Structure & Priors in Reinforcement Learning” at ICLR 2019 : proceedings, 2019, p. 9.
  * [11] Kai Ueltzhöffer,  “Deep active inference,”  Biological Cybernetics, vol. 112, no. 6, pp. 547–573, Dec 2018\.
  * [12] Beren Millidge,  “Deep active inference as variational policy gradients,”  CoRR, vol. abs/1907.03876, 2019.