# Imitation Learning from Observations under Transition Model Disparity

Tanmay Gangwani1, Yuan Zhou2 & Jian Peng134
1University of Illinois at Urbana-Champaign
2BIMSA 3AIR, Tsinghua University 4HeliXon Limited
<EMAIL_ADDRESS>

###### Abstract

Learning to perform tasks by leveraging a dataset of expert observations, also
known as imitation learning from observations (ILO), is an important paradigm
for learning skills without access to the expert reward function or the expert
actions. We consider ILO in the setting where the expert and the learner
agents operate in different environments, with the source of the discrepancy
being the transition dynamics model. Recent methods for scalable ILO utilize
adversarial learning to match the state-transition distributions of the expert
and the learner, an approach that becomes challenging when the dynamics are
dissimilar. In this work, we propose an algorithm that trains an intermediary
policy in the learner environment and uses it as a surrogate expert for the
learner. The intermediary policy is learned such that the state transitions
generated by it are close to the state transitions in the expert dataset. To
derive a practical and scalable algorithm, we employ concepts from prior work
on estimating the support of a probability distribution. Experiments using
MuJoCo locomotion tasks highlight that our method compares favorably to the
baselines for ILO with transition dynamics mismatch 111Code:
https://github.com/tgangwani/AILO.

##  1 Introduction

Imitation Learning (IL) is a framework that trains an agent to perform desired
skills by leveraging expert demonstrations of those skills. Compared to the
standard Reinforcement Learning (RL) approach, IL offers the benefit of not
requiring a reward function, that can be difficult to specify for complicated
objectives. Recent IL methods that integrate efficiency with deep-RL and are
performant in high-dimensional state-action spaces include behavioral-cloning-
based algorithms (Ross et al., 2011; Brantley et al., 2019), and adversarial
IL algorithms inspired by maximum entropy inverse-RL (Ho & Ermon, 2016; Finn
et al., 2016). Imitation Learning from Observations (ILO) refers to the
setting where the expert demonstrations consist of only observations (or
states), while the expert actions are unavailable. ILO is beneficial when the
measurement of the expert action is difficult, e.g., in kinesthetic teaching
in robotics or when learning with motion capture datasets.

Adversarial IL methods frame the problem as the minimization of an
f-divergence between the state-action visitation distributions of the expert
and the learner (Ke et al., 2019). Since expert actions are absent in ILO, the
analogous methodology here is to minimize the f-divergence between the state-
transition distributions of the expert and the learner (Arumugam et al.,
2020). A state-transition distribution for a policy is the joint distribution
over the current state and the next state, and is defined formally in Section
2. Choosing the f-divergence to be the Jensen–Shannon (JS) divergence has
enabled successful imitation using algorithms such as GAIL (Ho & Ermon, 2016)
and GAIfO (Torabi et al., 2018b), for IL and ILO, respectively.

The transition dynamics model of an environment governs the distribution over
the next state, given the current state and action. In this paper, we focus on
ILO in the scenario of a transition dynamics mismatch between the expert and
the learner environments. Such discrepancy could manifest in real-world
applications of imitation learning when there are subtle differences in the
physical attributes of the system used to collect the demonstrations and the
system where the learner policy is run. Adversarial ILO methods such as GAIfO,
that attempt to train the learner by matching its state-transition
distribution with that of the expert, perform very well when the expert and
learner operate in a shared environment, under the same dynamics. When the
dynamics differ, however, matching the state-transition distributions becomes
challenging since the state transitions provided in the expert demonstrations
could be infeasible under the dynamics in the learner’s environment. To form
some intuition, consider ILO in a 2D grid with discrete states, where the path
demonstrated by the expert follows the main diagonal from the bottom-left
grid-cell to the top-right grid-cell. Suppose that the transition dynamics for
the learner environment are such that only horizontal and vertical moves are
permitted on the grid, i.e., no diagonal motion for the learner. In this case,
an optimal learner’s movement is still in the same general direction as the
expert, and it covers all the expert states (along with some adjacent states).
However, matching the state-transition distributions is an ineffective
strategy since the learner cannot reproduce the (one-step) state transitions
of the expert.

To alleviate the challenges of ILO under dynamics mismatch, we propose an
algorithm that trains an intermediary policy in the learner environment, and
hope to use it as a surrogate expert for training the learner (imitator). We
refer to this policy as the advisor. For the advisor to be effective, the
state transitions generated by it in the learner environment should be as
close as possible to the state transitions in the expert dataset. We formalize
this concept in terms of the cross-entropy distance between state-conditional
next-state distributions of the expert and the advisor. To convert this into a
practical and scalable algorithm for training the advisor, we incorporate
ideas from distribution support estimation (Wang et al., 2019). Simultaneous
to the advisor training, the learner agent is updated to imitate the advisor.
Crucially, the advisor operates in the same environment as the learner, making
the distribution-matching IL objective amenable.

We evaluate the efficacy of our ILO algorithm using five locomotion control
tasks from OpenAI Gym where we introduce a mismatch between the dynamics of
the expert and the learner by changing different configuration parameters. We
demonstrate that our approach compares favorably to the baseline ILO
algorithms in many of the considered scenarios.

##  2 Preliminaries

We model the RL environment as a discounted, infinite horizon Markov Decision
Process (MDP). At every discrete timestep, the agent observes a state
$(s\in\mathcal{S})$, generates an action $(a\in\mathcal{A})$ from a stochastic
policy $\pi(a|s)$, receives a scalar reward $r(s,a)$, and transitions to the
next state $(s^{\prime})$ sampled from the transition dynamics model
$p(s^{\prime}|s,a)$. In the infinite horizon setting, the future rewards are
discounted by a factor of $\gamma\in[0,1)$. Let $d_{\pi}^{t}(s)$ denote the
distribution induced by $\pi$ over the state-space at a particular timestep
$t$. The stationary discounted state distribution of $\pi$ is then defined as
$\rho_{\pi}(s)=(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}d_{\pi}^{t}(s)$. The RL
objective of maximizing the cumulative discounted sum of rewards can be framed
as $\max_{\pi}\mathbb{E}_{\rho_{\pi}(s,a)}[r(s,a)]$, where
$\rho_{\pi}(s,a)=\rho_{\pi}(s)\pi(a|s)$ is the state-action distribution (also
known as the occupancy measure). Lastly, we define the state-transition
distribution for a policy as
$\rho_{\pi}(s,s^{\prime})=\sum_{a}p(s^{\prime}|s,a)\rho_{\pi}(s,a)$ 222With
slight abuse of notation, we use the symbol $\rho_{\pi}$ for state, state-
action, and state-transition distributions of a policy $\pi$. We would provide
context around their usage to avoid any confusion..

###  2.1 GAIL and GAIfO

Generative Adversarial Imitation Learning (Ho & Ermon, 2016) is a widely
popular model-free IL method that builds on the Maximum Causal Entropy
Inverse-RL (MaxEnt-IRL) framework (Ziebart, 2010). MaxEnt-IRL models the
expert behavior with a policy that maximizes its $\gamma$-discounted causal
entropy, $\mathcal{H}(\pi)=\mathbb{E}_{\pi}[-\log\pi(a|s)]$, while satisfying
a feature matching constraint. GAIL considers a regularized version of the
dual to this primal problem. It shows that RL with the reward function
recovered as the solution of the regularized dual is equivalent to directly
learning a policy whose state-action distribution is similar to that of the
expert. For a specific choice of the regularizer, this similarity is
quantified by the JS divergence between the two state-action distributions,
$D_{JS}[\rho_{\pi}(s,a)\ ||\ \rho_{\pi_{e}}(s,a)]$. Based on these ideas, GAIL
seeks to learn a policy with the objective:

$\min_{\pi}\max_{D}\>\mathbb{E}_{\rho_{\pi}}\big{[}\log
D(s,a)\big{]}+\mathbb{E}_{\rho_{\pi_{e}}}\big{[}\log(1-D(s,a))\big{]}-\lambda\mathcal{H}(\pi)$

where $D:\mathcal{S}\times\mathcal{A}\rightarrow(0,1)$ is the discriminator
that provides the rewards for training the learner policy $\pi$, and the inner
maximization over $D$ approximates $D_{JS}(\cdot)$ similar to GANs (Goodfellow
et al., 2014). To empirically estimate the expectation under
$\rho_{\pi_{e}}(s,a)$, state-action pairs are sampled from the available
expert demonstrations. In the ILO setting, however, expert actions are not
included in the demonstrations. To mitigate this challenge, Torabi et al.
(2018b) propose GAIfO, which adapts to ILO by modifying the GAIL objective to
match the state-transition distributions of the expert and the learner, i.e.,
$D_{JS}[\rho_{\pi}(s,s^{\prime})\ ||\ \rho_{\pi_{e}}(s,s^{\prime})]$.
Correspondingly, the discriminator in GAIfO is a function of the state
transitions $D(s,s^{\prime})$.

###  2.2 Support Estimation via RED

We summarize a recently proposed method for estimating the support of a
distribution in high dimensions (RED; Wang et al. (2019)), since it forms a
core ingredient of our final algorithm. Let $\mathcal{X}$ denote a set and $p$
be a probability distribution on $\mathcal{X}$. Denote by
$\text{supp}(p)=\\{x\in\mathcal{X}\mid p(x)\neq 0\\}$, the support of the
distribution $p$. Given any $x\in\mathcal{X}$, the task is to know if
$x\in\text{supp}(p)$. Towards this goal, Wang et al. (2019) combine ideas from
kernel-based support estimation (De Vito et al., 2014) and RND (Burda et al.,
2018), and consider the following objective:

$\theta^{*}=\operatorname*{arg\,min}_{\theta}\mathbb{E}_{x\sim
p(x)}\|f_{\theta}(x)-f_{\tilde{\theta}}(x)\|^{2}_{2}$

where $f_{\theta}:\mathcal{X}\rightarrow\mathbb{R}^{K}$ is a trainable
function parameterized by $\theta$, while $f_{\tilde{\theta}}$ is a fixed
function with randomly initialized parameters $\tilde{\theta}$. Define the
score function as (with constant positive scalar $\lambda$):

$r_{\text{RED}}(x)=\exp({-\lambda\|f_{\theta^{*}}(x)-f_{\tilde{\theta}}(x)\|^{2}_{2}})$

Wang et al. (2019) conclude that the score $r_{\text{RED}}(x)$ is high if
$x\in\text{supp}(p)$, and is low otherwise. With neural network function
approximators, we thus obtain a smooth metric whose value decreases
(increases) as we move farther from (closer to) the support of the
distribution $p$.

##  3 Methods

We begin by defining some notations for our setup. We denote the MDP in which
the expert policy $(\pi_{e})$ operates as the $e$-MDP, while the learner (or
imitator) policy is run in the $l$-MDP. The two MDPs share all the attributes,
except for the transition dynamics function, which we symbolize with
$p_{e}(s^{\prime}|s,a)$ and $p_{l}(s^{\prime}|s,a)$, for the $e$-MDP and
$l$-MDP, respectively. $\rho_{e}(s)$, $\rho_{e}(s,a)$, and
$\rho_{e}(s,s^{\prime})$ are the state, state-action, and state-transition
distribution for the expert policy. These distributions depend on the $e$-MDP
dynamics $p_{e}(s^{\prime}|s,a)$. The corresponding distributions for the
learner are denoted by $\rho_{\pi}(\cdot)$ and they depend on the $l$-MDP
dynamics $p_{l}(s^{\prime}|s,a)$.

Adversarial ILO methods, such as GAIfO, that learn the imitator policy by
matching the state-transition distributions of the expert and learner aim to
solve the following primal problem:

$\max_{\pi}\>\mathcal{H(\pi)}\quad\text{s.t.}\quad\rho_{e}(s,s^{\prime})=\rho_{\pi}(s,s^{\prime})\quad\forall\>(s,s^{\prime})\in\mathcal{S}\times\mathcal{S}$
(1)

If the expert and the learner operate under different transition dynamics,
depending on the extent of the mismatch, it is possible that some (or all) of
the one-step state transitions of the expert are infeasible under the dynamics
function in $l$-MDP. Said differently, given a $(s_{e},s^{\prime}_{e})$ pair
sampled from the expert demonstrations, there could be no action in the
$l$-MDP from the state $s_{e}$ that results in $s^{\prime}_{e}$ as the next
state, as per the dynamics $p_{l}(s^{\prime}|s,a)$. This renders the state-
transition matching objective hard to optimize in practice. For instance, in
the practical implementation of GAIfO, the rewards for the imitator policy are
computed from a discriminator that is trained with binary classification on
$(s,s^{\prime})$ pairs from the expert and the learner. If the expert
transitions can’t be generated in $l$-MDP by the learner, then a high-capacity
discriminator could achieve perfect accuracy and thus fail to provide
informative rewards for imitation.

Consider a sequence of states $\\{s_{1},s_{2},\dotsc\\}$ generated by the
expert policy in the $e$-MDP, as shown in Figure 1(a). Given an expert state
$s_{i}$, it may not be possible to reach $s_{i+1}$ in a single timestep in the
$l$-MDP. Instead, we would like to find an alternative state $\tilde{s}_{i+1}$
that is reachable from $s_{i}$ in one step (i.e., feasible under the dynamics
$p_{l}$) and is close to the desired destination $s_{i+1}$. For instance, in
Figure 1(a), starting from the expert state $s_{1}$, $\tilde{s}_{2}$ is a more
desirable next state compared to $\hat{s}_{2}$.

(a) State transitions

(b) Comparison between GAIfO (left) and our approach (right)

Figure 1: (a) A sequence of states $s_{i}$ from the expert dataset. The states
$\tilde{s}_{i}$ are reached by sampling an action from the advisor policy
$\pi_{a}$ from every expert state. State $\hat{s}_{2}$ is less desirable than
$\tilde{s}_{2}$ since the latter is closer to the expert state $s_{2}$. The
squiggly lines show the path that a learner, that is optimized to match the
state-action distribution of the advisor, may take. (b) A high-level overview
of our approach. While GAIfO directly matches the state-transition
distributions of the expert and the learner, we learn an intermediary policy
(advisor) in the $l$-MDP that acts as the surrogate expert for the learner.
$\mathcal{D}$ is used to denote the corresponding distance metric.

###  3.1 Guidance via an advisor policy

To discover such feasible states $\tilde{s}_{i}$, we introduce an advisor
policy $\pi_{a}$ that operates in the $l$-MDP. $\pi_{a}$ is invoked only for
action selection on the expert states, rather than being run in a closed
feedback loop in the learner environment. In this way, $\pi_{a}$ is akin to a
contextual bandit policy. The goal with $\pi_{a}$ is to produce an action
$a_{i}\sim\pi_{a}(\cdot|s_{i})$ from the expert state $s_{i}$ such that the
next state in the $l$-MDP, $\tilde{s}_{i+1}\sim p_{l}(\cdot|s_{i},a_{i})$, is
close to the next state $s_{i+1}$ in the $e$-MDP under the expert policy. We
expand on the measure of closeness and the methodology to train the advisor in
the next subsection.

When the expert and learner dynamics are the same, the optimal advisor would
take the same actions as the expert policy. In the presence of a dynamics
mismatch, however, the advisor actions provide reasonably good guidance on how
to stay close to the expert’s state trajectory (Figure 1(a)). Therefore, the
advisor $\pi_{a}$ could act as a surrogate expert for the learner and the
learner’s imitation learning objective could be suitably modified. The main
advantage of the learner and the advisor operating in the same $l$-MDP is that
we could now attempt to match their visitation distributions, which is much
more structured than matching distributions across dynamics (Eq. 1). We
consider the IL objective:

$\max_{\pi}\>\mathcal{H(\pi)}\quad\text{s.t.}\quad\rho_{\pi_{a}}(s,a)=\rho_{\pi}(s,a)\quad\forall\>(s,a)\in\mathcal{S}\times\mathcal{A}$

where $\rho_{\pi_{a}}(s,a)$ is the state-action distribution of $\pi_{a}$.
Since the advisor is invoked only on the expert states, effectively,
$\rho_{\pi_{a}}(s,a)=\rho_{e}(s)\pi_{a}(a|s)$. We can convert the above
objective to an unconstrained optimization by using a parametric function
$f_{\omega}(s,a)$ as the Lagrange multiplier:

$\min_{\omega}\max_{\pi}\>\mathcal{J}(\pi,\omega)\coloneqq\mathcal{H(\pi)}+\mathbb{E}_{\rho_{e}(s)\pi_{a}(a|s)}[f_{\omega}(s,a)]-\mathbb{E}_{\rho_{\pi}(s,a)}[f_{\omega}(s,a)]$
(2)

We note that this objective bears resemblance to entropy-regularized
apprenticeship learning (Abbeel & Ng, 2004; Syed et al., 2008; Syed &
Schapire, 2008), modulo the use of an advisor policy that contributes the
favorable actions, instead of the expert policy.

###  3.2 Training the advisor

The advisor policy generates an action (deterministic $\pi_{a}$), or a
distribution over actions (stochastic $\pi_{a}$), given an expert state
sampled from the demonstration data. One suitable objective for training the
advisor is to minimize the dissimilarity between the destination state
achieved with the expert policy in $e$-MDP and that achieved with the advisor
in $l$-MDP. For instance, consider the scenario where the dynamics functions
are deterministic and known, denoted by $f_{e}$ and $f_{l}$ for the $e$-MDP
and $l$-MDP, respectively. Also, the policies $\pi_{e}$ and $\pi_{a}$ are
deterministic. Then, the advisor could be optimized with the following loss:

$\min_{\pi_{a}}\>\mathcal{J}(\pi_{a})\coloneqq\mathbb{E}_{s\sim\rho_{e}(s)}\Big{[}\mathcal{D}\big{[}f_{e}(s,\pi_{e}(s)),f_{l}(s,\pi_{a}(s))\big{]}\Big{]}$

where $\mathcal{D}$ is a distance measure, e.g., the $L_{2}$-norm in the state
space. However, we are interested in the setup with stochastic, unknown
dynamics functions and stochastic policies. To compute a distance metric in
this setting, we first define the state-conditional next-state distributions
for the expert and the advisor, by marginalizing over the actions:

$p_{e,\pi_{e}}(s^{\prime}|s)=\sum_{a}p_{e}(s^{\prime}|s,a)\pi_{e}(a|s)\quad;\quad
p_{l,\pi_{a}}(s^{\prime}|s)=\sum_{a}p_{l}(s^{\prime}|s,a)\pi_{a}(a|s)$

We then use the cross-entropy distance ($\mathbb{H}$) between these
distributions as a measure of closeness between the expert and the advisor:

$\displaystyle\pi_{a}^{*}$
$\displaystyle=\operatorname*{arg\,min}_{\pi_{a}}\>\mathcal{J}(\pi_{a})\coloneqq\mathbb{E}_{s\sim\rho_{e}(s)}\Big{[}\mathbb{H}\big{[}p_{l,\pi_{a}}(s^{\prime}|s),\>p_{e,\pi_{e}}(s^{\prime}|s)\big{]}\Big{]}$
(3)
$\displaystyle=\operatorname*{arg\,max}_{\pi_{a}}\>\mathcal{J}(\pi_{a})\coloneqq\mathbb{E}_{s\sim\rho_{e}(s)}\mathbb{E}_{a\sim\pi_{a}(a|s)}\mathbb{E}_{s^{\prime}\sim
p_{l}(s^{\prime}|s,a)}\big{[}\log p_{e,\pi_{e}}(s^{\prime}|s)\big{]}$
$\displaystyle=\operatorname*{arg\,max}_{\pi_{a}}\>\mathcal{J}(\pi_{a})\coloneqq\mathbb{E}_{s\sim\rho_{e}(s)}\mathbb{E}_{a\sim\pi_{a}(a|s)}\mathbb{E}_{s^{\prime}\sim
p_{l}(s^{\prime}|s,a)}\big{[}\log\rho_{e}(s,s^{\prime})\big{]}$

where, in the last equation, we have multiplied the term inside the log with
$\rho_{e}(s)$, a quantity independent of $\pi_{a}$. Figure 1(b) provides a
high-level overview of our approach.

####  3.2.1 An approximation based on RED

The objective in Eq. 3 presents a couple of challenges. Firstly, it is
infeasible to evaluate the density of any state transition $(s,s^{\prime})$
under the expert’s state-transition distribution $\rho_{e}(s,s^{\prime})$,
since this distribution is unknown and only a few samples from this
distribution are available to us in the form of the expert demonstrations.
Secondly, and more importantly, note that the state transition that ought to
be evaluated under $\rho_{e}(s,s^{\prime})$ is generated by the advisor policy
in the $l$-MDP. If no advisor policy can replicate the state transition
behavior of the expert, as is likely when there is a dynamics mismatch, then
the objective becomes degenerate with an optimal value of $-\infty$.

Our approach to mitigate these issues is to replace the log-density term in
Eq. 3 with an estimated value that quantifies the proximity of a given state
transition $(s,s^{\prime})$ to the manifold of the expert’s support in this
space, i.e., $\rho_{e}(s,s^{\prime})$. To get this value, we leverage RED
(Wang et al., 2019), which pre-trains a deep neural network using data samples
from the distribution of interest. Then, given a test sample, the network
outputs a continuous value that provides an estimate of how far the test
sample is from the distribution’s support. Section 2.2 provides a short
background on RED.

In our instantiation, we pre-train a RED network
$r^{\phi}_{\text{RED}}(s,s^{\prime})$ with state transitions from the expert
demonstrations. The network is then frozen and utilized in the objective to
train the advisor:

$\max_{\pi_{a}}\>\mathcal{J}(\pi_{a})\coloneqq\mathbb{E}_{s\sim\rho_{e}(s)}\mathbb{E}_{a\sim\pi_{a}(a|s)}\mathbb{E}_{s^{\prime}\sim
p_{l}(s^{\prime}|s,a)}\big{[}r^{\phi}_{\text{RED}}(s,s^{\prime})\big{]}$ (4)

For any advisor-generated state transition,
$r^{\phi}_{\text{RED}}(s,s^{\prime})$ provides a smooth metric in the range
$(0,1]$, whose value increases (decreases) as the transition moves closer to
(farther from) the support of the distribution $\rho_{e}(s,s^{\prime})$. Thus,
training $\pi_{a}$ to maximize this value yields an advisor that provides
guidance on how to stay close to the expert’s state trajectory when operating
in $l$-MDP (Figure 1(b)).

To provide further intuition on the use of RED, Table 1 shows the
$r^{\phi}_{\text{RED}}(s,s^{\prime})$ values obtained from a trained RED
network in two environments—Walker2d and HalfCheetah. The RED network

$\Big{\downarrow}$ |  Increasing noise | Walker2d | HalfCheetah
---|---|---|---
| N/S = 0 | 0.93 | 0.96
| N/S = 0.2 | 0.80 | 0.71
| N/S = 0.5 | 0.44 | 0.23
| N/S = 1 | 0.12 | 0.04
| N/S = 2 | 0.01 | 0.001

Table 1: Averaged $r^{\phi}_{\text{RED}}(s,s^{\prime})$ values for different
noise levels.

is trained with 50 state transitions $\\{s_{e},s^{\prime}_{e}\\}$ sampled from
$\rho_{e}(s,s^{\prime})$ and then evaluated under several noisy transitions
$\\{s_{e},s^{\prime}_{e}+\eta\\}$, where $\eta$ denotes zero-mean Gaussian
noise. We show the $r^{\phi}_{\text{RED}}$ values averaged across the
transitions for different settings of the noise-to-signal ratio (N/S), i.e.,
the ratio of the standard deviation of the noise to the standard deviation of
the states. We observe that $r^{\phi}_{\text{RED}}(s,s^{\prime})$ is close to
1.0 when the transitions are on the support of the expert, and it gradually
decreases as the transitions drift away from the support as a result of a
larger amount of added noise.

###  3.3 Algorithm and Implementation

It is possible to use a two-stage training procedure for the overall
algorithm—first, learn the advisor in $l$-MDP with Eq. 4, and then use it in
the IL objective laid out in Eq. 2 to optimize for the reward network
$f_{\omega}$ and the learner policy $\pi$. In the first stage of advisor
learning, although the RED network is trained offline, computing the gradients
for optimizing $\pi_{a}$ still requires environment interaction. More
importantly, since the advisor is only trained on the expert states, the
objective in Eq. 4 requires the capability to reset the environment to the
expert states in the $l$-MDP.

To alleviate this problem, we propose to train the advisor $(\pi_{a})$, the
reward network $(f_{\omega})$ and the learner $(\pi)$ jointly, in an iterative
manner, and reuse the environment interaction data generated with $\pi$ for
training $\pi_{a}$ as well, using the importance sampling trick. Specifically,
let $\beta$ denote the parameters of $\pi_{a}$. Then the gradient of the
objective function $\mathcal{J}(\pi_{a})$ is:

$\displaystyle\nabla={}$
$\displaystyle\nabla_{\beta}\Big{(}\mathbb{E}_{s\sim\rho_{e}(s)}\mathbb{E}_{a\sim\pi_{a}(a|s)}\mathbb{E}_{s^{\prime}\sim
p_{l}(s^{\prime}|s,a)}\big{[}r^{\phi}_{\text{RED}}(s,s^{\prime})\big{]}\Big{)}$
(5) $\displaystyle={}$
$\displaystyle\mathbb{E}_{s\sim\rho_{e}(s)}\mathbb{E}_{a\sim\pi_{a}(a|s)}\Big{(}\mathbb{E}_{s^{\prime}\sim
p_{l}(s^{\prime}|s,a)}[r^{\phi}_{\text{RED}}(s,s^{\prime})].\nabla_{\beta}\log\pi_{a}(a|s)\Big{)}$
$\displaystyle={}$
$\displaystyle\mathbb{E}_{\rho_{\pi}(s,a)}\underbrace{\frac{\rho_{e}(s)}{\rho_{\pi}(s)}\frac{\pi_{a}(a|s)}{\pi(a|s)}}_{\text{IS
ratios}}\Big{(}\mathbb{E}_{s^{\prime}\sim
p_{l}(s^{\prime}|s,a)}[r^{\phi}_{\text{RED}}(s,s^{\prime})].\nabla_{\beta}\log\pi_{a}(a|s)\Big{)}$

where the second equation uses the score function estimator (Kleijnen &
Rubinstein, 1996), and the third employs two importance sampling ratios. These
ratios are easy to estimate for our setting (please see Appendix A.1 for
details). Crucially, since the gradient is now computed with state-action data
from $\rho_{\pi}(s,a)$, we no longer require environment resets. Further, we
observe that approximating the inner expectation with a single sample is
sufficient for our tasks.

Input : A dataset of expert state transitions
$\mathcal{D}_{e}=\\{s_{e},s^{\prime}_{e}\\}$ collected in $e$-MDP

/* initialization and offline pre-training */

Initialize: advisor $\pi_{a}$, learner $\pi$, reward network
$f_{\omega}(s,a)$, RED network $r^{\phi}_{\text{RED}}(s,s^{\prime})$

Pre-train $r^{\phi}_{\text{RED}}(s,s^{\prime})$ with $\mathcal{D}_{e}$ using
the algorithm in Wang et al. (2019)

for  _iter in $\\{1,\dotsc,N\\}$_ do

/* data collection */

Roll out trajectories $\tau$ using $\pi$

Update reward network $f_{\omega}$ using $\tau$, $\pi_{a}$ and
$\mathcal{D}_{e}$ (Eq. 2)

/* update learner */

Compute reward $f_{\omega}(s,a)$ for each transition in $\tau$. Use $\tau$ to
update $\pi$ with MaxEnt RL

/* update advisor */

Compute reward $r^{\phi}_{\text{RED}}(s,s^{\prime})$ for each transition in
$\tau$. Use $\tau$ to update $\pi_{a}$ with Eq. 5

end for

Algorithm 1 AILO (Advisor-augmented Imitation Learning from Observations)

We abbreviate our method as AILO, short for Advisor-augmented Imitation
Learning from Observations. Algorithm 1 provides an outline. We start with the
offline pre-training of the RED network using a dataset of expert state
transitions collected in the $e$-MDP. Then, we perform iterative optimization
in the $l$-MDP where each iteration involves generating trajectories with the
current learner $\pi$, followed by gradient updates to the reward network
$f_{\omega}$, and to the advisor and learner policies. Following Eq. 2, the
learner policy is trained with a MaxEnt RL algorithm with per-timestep
entropy-regularized reward given by
$f_{\omega}(s_{t},a_{t})-\alpha\log\pi(a_{t}|s_{t})$, where $\alpha$ is the
entropy coefficient. In our experiments, we use the clipped-ratio PPO
algorithm (Schulman et al., 2017) and adaptively tune $\alpha$ as suggested in
prior work (Haarnoja et al., 2018).

##  4 Related work

There is a vast amount of literature on IL since it is a powerful framework to
train agents to perform complex behaviors without a reward specification. ILO,
where no expert action labels are available, presents several benefits as well
as some unique challenges, and thus, has garnered significant attention from
the community in recent times (Torabi et al., 2018a; b; Liu et al., 2018;
Edwards et al., 2019; Sun et al., 2019; Yang et al., 2019; Zhu et al., 2021).
ILO methods adapted from GAIL have been proposed for one-shot imitation of
diverse behaviors (Wang et al., 2017), training policies to generate human-
like movement patterns using motion-capture data (Peng et al., 2018; Merel et
al., 2017), and for locomotion control from raw visual data (Torabi et al.,
2018b). Arumugam et al. (2020) introduce a framework that casts adversarial
ILO as $f$-divergence minimization and provide insights on the design
decisions that impact performance.

Several approaches have been proposed to handle the differences between the
expert and the learner environments in terms of viewpoints, visual
appearances, presence of distractors, and morphology changes (Stadie et al.,
2017; Gupta et al., 2017; Liu et al., 2018; Sermanet et al., 2018). They
typically proceed by learning a domain-invariant representation and matching
features in that space. Learning such a representation is not required in our
setup since the state-space is shared between the $l$-MDP and the $e$-MDP.
Methods for robustness to shifts in the action-space and the dynamics model
have also been researched. Zolna et al. (2019) propose to match state-pair
distributions of the expert and the learner, where the states in a pair are
sampled with random time gaps, rather than being consecutive. Liu et al.
(2019) learn an inverse action model to predict deterministic actions in the
$l$-MDP that could generate expert-like state transitions and use it to
regularize policy updates. Gangwani & Peng (2020) filter the trajectories
generated in $l$-MDP based on their similarity to the states in the expert
dataset and perform (self-) imitation on these. The key methodological
difference between these methods and our work is that we learn an intermediary
stochastic policy (advisor) in the $l$-MDP by bringing its state-conditional
next-state distribution closer to that of the expert, and propose an
instantiation of this idea using an approximation based on support estimation.

##  5 Experiments

We evaluate the efficacy of AILO using continuous-control locomotion
environments from OpenAI Gym (Brockman et al., 2016), modeled using the MuJoCo
physics simulator (Todorov et al., 2012). We include a description of the
tasks, the baselines, and the learning curves. Further details on the
architectures and the hyperparameters are provided in Appendix A.2.

Environments. We consider five tasks - {Half-Cheetah, Walker, Hopper, Ant,
Humanoid}. To create a discrepancy between the expert and the learner
transition dynamics, we modify one physical property in the learner
environment from the set - {density, gravity, joint-friction}. Specifically,
for a task $\mathcal{T}$ from the task-set, we denote:

  * •

$\mathcal{T}$ (heavy) $\rightarrow$ learner agent has 2$\times$ the mass of
the expert agent

  * •

$\mathcal{T}$ (light) $\rightarrow$ gravity in the learner’s environment is
half the value in the expert’s

  * •

$\mathcal{T}$ (drag) $\rightarrow$ the friction coefficient for all the joints
in the learner is 2$\times$ the value in the expert

(a) Performance on $\mathcal{T}$ (heavy) environments

(b) Performance on $\mathcal{T}$ (drag) environments

(c) Performance on $\mathcal{T}$ (light) environments

Figure 2: Learning curves for AILO and the baselines for different
environments with discrepancy in dynamics

Baselines. We contrast AILO with two baselines – a.) GAIfO strives to minimize
the JS divergence between the state-transition distributions and is briefly
described in section 2.1; b.) VAIL (Peng et al., 2018) applies the concept of
variational information bottleneck (Alemi et al., 2016) to the GAIL
discriminator for improved regularization and has been successfully used for
ILO with motion-capture data. To limit the effect of confounding factors
during comparison, we share modules across AILO and the two baselines to the
best of our ability. Concretely, all discriminator/reward networks use the
same architecture and the gradient-penalty regularization (Mescheder et al.,
2018) and thus exhibit the same reward biases. Furthermore, the MaxEnt-RL PPO
module is the same for all algorithms.

Performance. Figures 2(a)–2(c) plot the learning curves for all the algorithms
across the different tasks (heavy, drag, light). We show the average episodic
returns achieved by the learner in the $l$-MDP, normalized to the returns
achieved by the expert in $e$-MDP. The plots include the mean and the standard
deviation of returns over 6 independent runs with random seeds.

We observe that AILO provides a noticeable improvement in learning efficiency
in several situations, such as Half-Cheetah (heavy, drag, light), Walker
(heavy), Hopper (heavy), Ant (drag), Humanoid (drag, light); while being
comparable to the best baseline in other cases. For GAIfO, we find that it
does not learn any useful skill for Walker (heavy) and exhibits training
instability for Half-Cheetah (heavy, drag). We attribute this to the
difficulty of matching the state-transition distributions across different
dynamics. VAIL, which matches state distributions, instead of state-transition
distributions, is a stronger baseline and works well in several cases. Lastly,
we highlight a few failure modes of AILO – in Ant (light) and Humanoid
(heavy), we note that AILO (and baselines) do not make much progress towards
imitating the demonstrated behavior within our time budget, motivating the
need for future enhancements to enable efficient skill transfer in these
challenging setups.

Ablation on the degree of dynamics mismatch. For the empirical results in
Figure 2, the variation in mass, gravity, or friction, between the $e$-MDP and
the $l$-MDP was kept at a constant factor.

Figure 3: Results with different amount of dynamics mismatch

In Figure 3, we consider the Walker task and plot the systematic degradation
in the performance of AILO and the baseline GAIfO, as the $l$-MDP parameters
drift further from the $e$-MDP parameters. We show the final average episodic
returns achieved by the learner in the $l$-MDP, normalized to the returns
achieved by the expert in $e$-MDP, as a function of the degree of dynamics
mismatch. In the plots, the expert parameters (mass, friction) are kept
constant and the learner parameters are varied such that the ratio increases
from a starting value of 1 (no mismatch). We observe that although imitation
naturally becomes more challenging as the dynamics become more different, the
degradation with AILO is more graceful compared to GAIfO.

##  6 Conclusion

In this paper, we present AILO, our algorithm for imitation learning from
observations under transition model disparity between the expert and the
learner environments. Rather than directly matching the state-transition
distributions across environments, we train an intermediary policy (advisor)
in the learner environment and use it as a surrogate expert for the learner.
Towards learning an advisor that acts as an effective surrogate, we propose to
minimize the cross-entropy distance between the state-conditional next-state
distributions of the advisor and the expert. To realize this idea into a
scalable ILO algorithm, we leverage prior work on support estimation (RED).
Our experiments on five MuJoCo locomotion tasks with different types of
dynamics discrepancies show that AILO compares favorably to the baseline ILO
methods in many cases.

## References

  * Abbeel & Ng (2004) Pieter Abbeel and Andrew Y Ng.  Apprenticeship learning via inverse reinforcement learning.  In _Proceedings of the twenty-first international conference on Machine learning_ , pp. 1, 2004.
  * Alemi et al. (2016) Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy.  Deep variational information bottleneck.  _arXiv preprint arXiv:1612.00410_ , 2016.
  * Arumugam et al. (2020) Dilip Arumugam, Debadeepta Dey, Alekh Agarwal, Asli Celikyilmaz, Elnaz Nouri, and Bill Dolan.  Reparameterized variational divergence minimization for stable imitation.  _arXiv preprint arXiv:2006.10810_ , 2020.
  * Brantley et al. (2019) Kiante Brantley, Wen Sun, and Mikael Henaff.  Disagreement-regularized imitation learning.  In _International Conference on Learning Representations_ , 2019.
  * Brockman et al. (2016) Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym.  _arXiv preprint arXiv:1606.01540_ , 2016.
  * Burda et al. (2018) Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov.  Exploration by random network distillation.  _arXiv preprint arXiv:1810.12894_ , 2018.
  * De Vito et al. (2014) Ernesto De Vito, Lorenzo Rosasco, and Alessandro Toigo.  A universally consistent spectral estimator for the support of a distribution.  _Appl Comput Harmonic Anal_ , 37:185–217, 2014.
  * Edwards et al. (2019) Ashley Edwards, Himanshu Sahni, Yannick Schroecker, and Charles Isbell.  Imitating latent policies from observation.  In _International Conference on Machine Learning_ , pp. 1755–1763. PMLR, 2019.
  * Finn et al. (2016) Chelsea Finn, Paul Christiano, Pieter Abbeel, and Sergey Levine.  A connection between generative adversarial networks, inverse reinforcement learning, and energy-based models.  _arXiv preprint arXiv:1611.03852_ , 2016.
  * Gangwani & Peng (2020) Tanmay Gangwani and Jian Peng.  State-only imitation with transition dynamics mismatch.  _arXiv preprint arXiv:2002.11879_ , 2020.
  * Goodfellow et al. (2014) Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  In _Advances in neural information processing systems_ , pp. 2672–2680, 2014.
  * Gupta et al. (2017) Abhishek Gupta, Coline Devin, YuXuan Liu, Pieter Abbeel, and Sergey Levine.  Learning invariant feature spaces to transfer skills with reinforcement learning.  _arXiv preprint arXiv:1703.02949_ , 2017.
  * Haarnoja et al. (2018) Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _International Conference on Machine Learning_ , pp. 1861–1870. PMLR, 2018.
  * Ho & Ermon (2016) Jonathan Ho and Stefano Ermon.  Generative adversarial imitation learning.  In _Advances in Neural Information Processing Systems_ , pp. 4565–4573, 2016.
  * Ke et al. (2019) Liyiming Ke, Matt Barnes, Wen Sun, Gilwoo Lee, Sanjiban Choudhury, and Siddhartha Srinivasa.  Imitation learning as $f$-divergence minimization.  _arXiv preprint arXiv:1905.12888_ , 2019.
  * Kleijnen & Rubinstein (1996) Jack PC Kleijnen and Reuven Y Rubinstein.  Optimization and sensitivity analysis of computer simulation models by the score function method.  _European Journal of Operational Research_ , 88(3):413–427, 1996.
  * Liu et al. (2019) Fangchen Liu, Zhan Ling, Tongzhou Mu, and Hao Su.  State alignment-based imitation learning.  _arXiv preprint arXiv:1911.10947_ , 2019.
  * Liu et al. (2018) YuXuan Liu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine.  Imitation from observation: Learning to imitate behaviors from raw video via context translation.  In _2018 IEEE International Conference on Robotics and Automation (ICRA)_ , pp. 1118–1125. IEEE, 2018.
  * 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.
  * Mescheder et al. (2018) Lars Mescheder, Andreas Geiger, and Sebastian Nowozin.  Which training methods for gans do actually converge?  In _International conference on machine learning_ , pp. 3481–3490. PMLR, 2018.
  * Peng et al. (2018) Xue Bin Peng, Angjoo Kanazawa, Sam Toyer, Pieter Abbeel, and Sergey Levine.  Variational discriminator bottleneck: Improving imitation learning, inverse rl, and gans by constraining information flow.  _arXiv preprint arXiv:1810.00821_ , 2018.
  * 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_ , pp. 627–635, 2011.
  * Schulman et al. (2015) 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.
  * Sermanet et al. (2018) Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain.  Time-contrastive networks: Self-supervised learning from video.  In _2018 IEEE International Conference on Robotics and Automation (ICRA)_ , pp. 1134–1141. IEEE, 2018.
  * Stadie et al. (2017) Bradly C Stadie, Pieter Abbeel, and Ilya Sutskever.  Third-person imitation learning.  _arXiv preprint arXiv:1703.01703_ , 2017.
  * Sun et al. (2019) Wen Sun, Anirudh Vemula, Byron Boots, and Drew Bagnell.  Provably efficient imitation learning from observation alone.  In _International Conference on Machine Learning_ , pp. 6036–6045. PMLR, 2019.
  * Syed & Schapire (2008) Umar Syed and Robert E Schapire.  A game-theoretic approach to apprenticeship learning.  In _Advances in neural information processing systems_ , pp. 1449–1456, 2008.
  * Syed et al. (2008) Umar Syed, Michael Bowling, and Robert E Schapire.  Apprenticeship learning using linear programming.  In _Proceedings of the 25th international conference on Machine learning_ , pp. 1032–1039, 2008.
  * 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_ , pp. 5026–5033. IEEE, 2012.
  * Torabi et al. (2018a) Faraz Torabi, Garrett Warnell, and Peter Stone.  Behavioral cloning from observation.  _arXiv preprint arXiv:1805.01954_ , 2018a.
  * Torabi et al. (2018b) Faraz Torabi, Garrett Warnell, and Peter Stone.  Generative adversarial imitation from observation.  _arXiv preprint arXiv:1807.06158_ , 2018b.
  * Wang et al. (2019) Ruohan Wang, Carlo Ciliberto, Pierluigi Vito Amadori, and Yiannis Demiris.  Random expert distillation: Imitation learning via expert policy support estimation.  In _International Conference on Machine Learning_ , pp. 6536–6544. PMLR, 2019.
  * 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.
  * Yang et al. (2019) Chao Yang, Xiaojian Ma, Wenbing Huang, Fuchun Sun, Huaping Liu, Junzhou Huang, and Chuang Gan.  Imitation learning from observations by minimizing inverse dynamics disagreement.  _arXiv preprint arXiv:1910.04417_ , 2019.
  * Zhu et al. (2021) Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou.  Off-policy imitation learning from observations.  _arXiv preprint arXiv:2102.13185_ , 2021.
  * Ziebart (2010) Brian D Ziebart.  Modeling purposeful adaptive behavior with the principle of maximum causal entropy.  2010\.
  * Zolna et al. (2019) Konrad Zolna, Negar Rostamzadeh, Yoshua Bengio, Sungjin Ahn, and Pedro O Pinheiro.  Reinforced imitation in heterogeneous action space.  _arXiv preprint arXiv:1904.03438_ , 2019.

##  Appendix A Appendix

###  A.1 Importance sampling for advisor training

Following Eq. 5, the importance sampling factor to be estimated is
$\frac{\rho_{e}(s)}{\rho_{\pi}(s)}\frac{\pi_{a}(a|s)}{\pi(a|s)}$, which is a
product of two ratios. The second ratio $\frac{\pi_{a}(a|s)}{\pi(a|s)}$ is
easily computable since we have both the advisor policy $(\pi_{a})$ and the
learner policy $(\pi)$ as parameterized Gaussian distributions. The first
ratio $\frac{\rho_{e}(s)}{\rho_{\pi}(s)}$ can be approximated by training a
binary classifier to distinguish the states sampled from the distributions
$\rho_{e}(s)$ and $\rho_{\pi}(s)$. Concretely, consider the objective:

$D^{\prime}=\operatorname*{arg\,max}_{D:\mathcal{S}\rightarrow(0,1)}\>\mathbb{E}_{\rho_{e}(s)}\big{[}\log
D(s)\big{]}+\mathbb{E}_{\rho_{\pi}(s)}\big{[}\log(1-D(s))\big{]}$

Then, we can estimate the ratio of the state distributions as
$\frac{\rho_{e}(s)}{\rho_{\pi}(s)}\approx\frac{D^{\prime}(s)}{1-D^{\prime}(s)}$.

###  A.2 Hyperparameters and implementation details

Hyperparameter | Value
---|---
Discriminator / Reward network ($f_{\omega}$) |
architecture | 3 layers, 128 hidden-dim, tanh
# updates per iteration | 5
optimizer | Adam (lr = $1\times 10^{-4}$)
RL agent (shared) |
policy arch. | 3 layers, 64 hidden-dim, tanh
critic arch. | 3 layers, 64 hidden-dim, tanh
PPO clipping | 0.2
PPO epochs per iteration | 5
optimizer | Adam (lr = $1\times 10^{-4}$)
discount factor $\left(\gamma\right)$ | $0.99$
GAE factor (Schulman et al., 2015) | $0.95$
entropy target | $-|\mathcal{A}|$
Table 2: Hyperparameters for AILO and the baselines

The hyperparameters used for our experiments are mentioned in Table 2. The
discriminator network used by the baselines and the reward network employed in
AILO share a common structure. The same RL agent is used across AILO and all
the baselines. Additionally, for the baselines, we tested the hyperparameters
mentioned in the respective papers and tuned them with grid-search.

Expert data collection. We train the expert policy in $e$-MDP using
environmental rewards with the PPO algorithm. We then generate a few rollouts
with the trained policy and sub-sample state transitions from these rollouts.
We create the expert dataset with 50 such transitions for {Half-Cheetah,
Walker, Hopper, Ant}, and 1000 transitions for Humanoid.

Absolute value for the expert scores. In Figure 2, we plot the average
episodic returns achieved by the learner in the $l$-MDP, normalized to the
returns achieved by the expert in the $e$-MDP. For completeness, Table 3
reports the absolute value of the expected returns obtained by the expert per
episode (max-episode-length = 1000 timesteps) in the $e$-MDP.

Task | Returns
---|---
Half-Cheetah | 6441
Walker | 4829
Hopper | 3675
Ant | 5765
Humanoid | 9678
Table 3: Expert’s performance in $e$-MDP

###  A.3 Illustration of the advisor and the learner in grid-world

We discuss the differences between the advisor policy $\pi_{a}$ and the
learner policy with the help of an illustration in the deterministic grid-
world environment (Figure 4). The leftmost plot shows the demonstrated expert
states $\\{s_{1},s_{2},s_{3},s_{4},s_{5}\\}$. The $e$-MDP and the $l$-MDP are
the same grid-world but with the following transition dynamics mismatch –
while the $e$-MDP allows diagonal hops such that the shown expert transitions
are consecutive, in the $l$-MDP, the agent is restricted to only horizontal
and vertical movements to the nearby cells.

As discussed in Section 3.2, we seek to train an advisor such that the
dissimilarity between the destination state achieved with the expert policy in
$e$-MDP and that achieved with the advisor in $l$-MDP is minimized. The middle
plot in Figure 4 shows the actions (marked with red arrows) that an optimal
deterministic advisor would take. Note that the advisor is trained only on the
expert states to optimize for the next state; it is not trained for long-term
behavior and thus struggles with decision-making on the non-expert states. At
the state $s_{3}$, the green states are the favorable states (minimum distance
to $s_{4}$), while the red states are the unfavorable states (maximum distance
to $s_{4}$). In the $l$-MDP, the dataset $\\{s_{i},\pi_{a}(s_{i})\\}$ is then
used to train a learner that runs in a closed feedback loop with the
environment.

The learner is trained with the standard adversarial imitation learning
objective to match its stationary discounted state-action distribution with
that of the generated dataset $\\{s_{i},\pi_{a}(s_{i})\\}$ in the $l$-MDP.
With adversarial imitation learning, the learner reproduces the demonstrated
action at the demonstrated state, while also being trained to navigate back to
the manifold of the demonstrated states when it encounters non-demonstrated
states, thus enabling long-horizon imitation. The advisor, that is trained to
optimize just the immediate action at each expert state (similar in principle
to behavior cloning), is incapable of long-horizon imitation. The rightmost
plot in Figure 4 shows the trajectory that an optimal learner may produce.

Figure 4: Illustration of the expert states, the actions by the advisor
policy, and the path taken by the learner policy in a grid-world environment.

###  A.4 Comparing the advisor and the learner performance

In Table 4, we report the average episodic returns achieved by the advisor and
the learner in the $l$-MDP at the end of the training, normalized to the
returns achieved by the expert in $e$-MDP. We show data for three environments
and all the types of transition dynamics mismatch considered in the paper –
Heavy (H), Drag (D), Light (L). We note that the learner outperforms the
advisor by a substantial margin for all the scenarios. This is because, unlike
the learner, the advisor is not optimized for long-horizon performance (cf.
Appendix A.3).

|  | AILO-advisor | AILO-learner
---|---|---|---
HalfCheetah | (H) | 0.19 | 0.57
(D) | 0.27 | 0.82
(L) | 0.18 | 0.63
Walker | (H) | 0.07 | 0.50
(D) | 0.14 | 0.78
(L) | 0.17 | 0.81
Hopper | (H) | 0.53 | 0.67
(D) | 0.48 | 0.77
(L) | 0.66 | 0.84
Table 4: Normalized average episodic returns achieved by the advisor and the
learner in the $l$-MDP.

###  A.5 Comparison to Behavioral Cloning and RL baselines

In this section, we compare AILO with BCO (Torabi et al., 2018a). BCO is an
ILO algorithm that first learns an inverse dynamics model $p(a|s,s^{\prime})$
in the $l$-MDP, and then uses this model to infer the expert actions from the
given expert state-transition dataset. These inferred actions are used to
train the learner agent using the standard behavioral cloning strategy.

Table 5 reports the average episodic returns achieved by BCO and AILO,
normalized to the returns achieved by the expert in $e$-MDP. We show data for
all the environments and all the types of transition dynamics mismatch
considered in the paper – Heavy (H), Drag (D), Light (L). We note that BCO
performs much worse than AILO in all scenarios, with the exception of Ant (H).
This can be attributed to two factors: 1.) the compounding errors problem in
the behavioral cloning approach, and 2.) BCO does not explicitly account for
the transition dynamics mismatch while inferring actions from the learned
inverse dynamics model.

Although the agent does not receive rewards in the $l$-MDP in the imitation
learning setup, we consider an oracle RL baseline that trains the agent to
maximize the rewards in the $l$-MDP using the clipped-ratio PPO algorithm
(Schulman et al., 2017). This baseline provides us with an estimate on the
upper bound of the performance of AILO. Similar to the AILO numbers, the RL
baseline scores are normalized to the returns achieved by the expert in
$e$-MDP. The data in Table 5 shows that there is a lot of scope to optimize
AILO further, especially for environments such as Walker and Ant.

|  | BCO | AILO | RL in $l$-MDP
---|---|---|---|---
HalfCheetah | (H) | 0.28 $\pm$ 0.02 | 0.57 $\pm$ 0.08 | 0.67
(D) | 0.47 $\pm$ 0.03 | 0.82 $\pm$ 0.04 | 1.06
(L) | 0.09 $\pm$ 0.01 | 0.63 $\pm$ 0.05 | 0.78
Walker | (H) | 0.03 $\pm$ 0.02 | 0.50 $\pm$ 0.14 | 0.72
(D) | 0.04 $\pm$ 0.01 | 0.78 $\pm$ 0.03 | 1.13
(L) | 0.04 $\pm$ 0.03 | 0.81 $\pm$ 0.08 | 0.93
Hopper | (H) | 0.10 $\pm$ 0.01 | 0.67 $\pm$ 0.05 | 0.65
(D) | 0.09 $\pm$ 0.01 | 0.77 $\pm$ 0.04 | 0.76
(L) | 0.11 $\pm$ 0.05 | 0.84 $\pm$ 0.03 | 0.86
Ant | (H) | 0.73 $\pm$ 0.01 | 0.65 $\pm$ 0.04 | 1.04
(D) | 0.02 $\pm$ 0.00 | 0.19 $\pm$ 0.13 | 0.83
(L) | 0.06 $\pm$ 0.02 | 0.02 $\pm$ 0.14 | 0.75
Humanoid | (H) | 0.02 $\pm$ 0.00 | 0.07 $\pm$ 0.00 | 0.48
(D) | 0.01 $\pm$ 0.00 | 0.78 $\pm$ 0.05 | 0.71
(L) | 0.01 $\pm$ 0.00 | 0.26 $\pm$ 0.02 | 0.32
Table 5: Normalized average episodic returns for BCO and RL with rewards in
$l$-MDP.

###  A.6 Comparison with I2L

In this section, we compare AILO with I2L (Gangwani & Peng, 2020). I2L is an
algorithm that builds on top of adversarial imitation learning and introduces
specific components to handle the dynamics mismatch between the $e$-MDP and
the $l$-MDP. AILO and I2L share the principle of curating an intermediate
dataset in the $l$-MDP, which is then used to train the learner with
adversarial imitation learning. In the case of I2L, the intermediate dataset
is created by applying a filtration step on the trajectory generated in the
$l$-MDP, by using a learned Wasserstein critic to rank the trajectories.
Differently, in AILO, the intermediate dataset is obtained by the application
of the advisor policy on the demonstrated expert states. These are two quite
distinct and orthogonal approaches that could potentially be blended to create
an effective intermediate dataset for the learner agent.

Table 6 reports the average episodic returns achieved by I2L, along with the
performance of the other algorithms considered in this paper. We normalize all
the scores to the returns achieved by the expert in $e$-MDP and highlight the
best algorithm in each row with bold. The table verifies that I2L works very
well for ILO under dynamics model disparity. Among the 15 environment-mismatch
scenarios, GAIfO is the best method for 3, VAIL for 1, I2L for 5, and AILO for
7. This clearly indicates that a single approach may be insufficient to solve
a large variety of ILO problems and opens up an interesting future direction
of understanding the unique characteristics of these algorithms (e.g., the
contrasting approaches to create the intermediate dataset in I2L and AILO)
towards designing efficient ILO algorithms.

|  | GAIfO | VAIL | I2L | AILO
---|---|---|---|---|---
HalfCheetah | (H) | 0.34 $\pm$ 0.17 | 0.49 $\pm$ 0.09 | 0.59 $\pm$ 0.04 | 0.57 $\pm$ 0.08
(D) | 0.41 $\pm$ 0.19 | 0.65 $\pm$ 0.07 | 0.80 $\pm$ 0.22 | 0.82 $\pm$ 0.04
(L) | 0.52 $\pm$ 0.03 | 0.31 $\pm$ 0.20 | 0.48 $\pm$ 0.19 | 0.63 $\pm$ 0.05
Walker | (H) | 0.13 $\pm$ 0.07 | 0.34 $\pm$ 0.23 | 0.56 $\pm$ 0.09 | 0.50 $\pm$ 0.14
(D) | 0.50 $\pm$ 0.29 | 0.78 $\pm$ 0.03 | 0.64 $\pm$ 0.24 | 0.78 $\pm$ 0.03
(L) | 0.86 $\pm$ 0.07 | 0.81 $\pm$ 0.05 | 0.76 $\pm$ 0.05 | 0.81 $\pm$ 0.08
Hopper | (H) | 0.44 $\pm$ 0.04 | 0.50 $\pm$ 0.10 | 0.67 $\pm$ 0.07 | 0.67 $\pm$ 0.05
(D) | 0.79 $\pm$ 0.03 | 0.78 $\pm$ 0.04 | 0.78 $\pm$ 0.05 | 0.77 $\pm$ 0.04
(L) | 0.70 $\pm$ 0.06 | 0.75 $\pm$ 0.07 | 0.70 $\pm$ 0.08 | 0.84 $\pm$ 0.03
Ant | (H) | 0.69 $\pm$ 0.05 | 0.55 $\pm$ 0.06 | 0.48 $\pm$ 0.22 | 0.65 $\pm$ 0.04
(D) | 0.09 $\pm$ 0.12 | 0.14 $\pm$ 0.06 | 0.17 $\pm$ 0.10 | 0.19 $\pm$ 0.13
(L) | 0.03 $\pm$ 0.14 | 0.02 $\pm$ 0.14 | 0.13 $\pm$ 0.11 | 0.02 $\pm$ 0.14
Humanoid | (H) | 0.06 $\pm$ 0.01 | 0.06 $\pm$ 0.00 | 0.07 $\pm$ 0.01 | 0.07 $\pm$ 0.00
(D) | 0.57 $\pm$ 0.08 | 0.65 $\pm$ 0.10 | 0.33 $\pm$ 0.22 | 0.78 $\pm$ 0.05
(L) | 0.18 $\pm$ 0.02 | 0.15 $\pm$ 0.02 | 0.45 $\pm$ 0.21 | 0.26 $\pm$ 0.02
Table 6: Normalized average episodic returns for different ILO algorithms.


# Input-Aware Dynamic Timestep Spiking Neural Networks for Efficient In-Memory
Computing ††thanks: This work was supported in part by CoCoSys, a JUMP2.0
center sponsored by DARPA and SRC, Google Research Scholar Award, the NSF
CAREER Award, TII (Abu Dhabi), the DARPA AI Exploration (AIE) program, and the
DoE MMICC center SEA-CROGS (Award #DE-SC0023198).

Yuhang Li, Abhishek Moitra, Tamar Geller, Priyadarshini Panda  Department of
Electrical Engineering, Yale University
New Haven, CT 06511, USA

###### Abstract

Spiking Neural Networks (SNNs) have recently attracted widespread research
interest as an efficient alternative to traditional Artificial Neural Networks
(ANNs) because of their capability to process sparse and binary spike
information and avoid expensive multiplication operations. Although the
efficiency of SNNs can be realized on the In-Memory Computing (IMC)
architecture, we show that the energy cost and latency of SNNs scale linearly
with the number of timesteps used on IMC hardware. Therefore, in order to
maximize the efficiency of SNNs, we propose input-aware Dynamic Timestep SNN
(DT-SNN), a novel algorithmic solution to dynamically determine the number of
timesteps during inference on an input-dependent basis. By calculating the
entropy of the accumulated output after each timestep, we can compare it to a
predefined threshold and decide if the information processed at the current
timestep is sufficient for a confident prediction. We deploy DT-SNN on an IMC
architecture and show that it incurs negligible computational overhead. We
demonstrate that our method only uses 1.46 average timesteps to achieve the
accuracy of a 4-timestep static SNN while reducing the energy-delay-product by
80%.

###### Index Terms:

Spiking neural networks, in-memory computing, dynamic inference

##  I Introduction

Deep learning has revolutionized many challenging computational tasks such as
computer vision and natural language processing [10] using Artificial Neural
Networks (ANNs). These successes, however, have come at the cost of tremendous
computing resources and high latency [6]. Over the past decade, Spiking Neural
Networks (SNNs) have gained popularity as an energy-efficient alternative to
ANNs [14, 17]. SNNs are different from ANNs in that they process inputs over a
series of timesteps, whereas ANNs infer over what can be considered a single
timestep. The biologically-plausible neurons in SNNs maintain a variable
called membrane potential, which controls the behavior of the SNN over a
series of timesteps. When the membrane potential exceeds a certain threshold,
the neuron fires, creating a spike, and otherwise, the neuron remains inactive
(neuron outputs a 0 or 1). Such spike-based computing creates sparsity in the
computations and replaces multiplications with additions.

Figure 1: Energy estimation on our IMC architecture using VGG-16 on CIFAR-10
dataset. (A) energy ratio of each unit, (B) energy/latency vs. timesteps.

Although the binary spike nature of SNNs eliminates the need for
multiplications, compared to ANNs, SNNs require significantly more memory
access due to multi-timestep computations on traditional von-Neumann
architectures (called the “memory wall problem”) [21]. To alleviate this
problem, In-Memory Computing (IMC) hardware is used to perform analog dot-
product operations to achieve high memory bandwidth and compute parallelism
[15]. In this work, we mainly focus on achieving lower energy and latency in
the case of IMC-implemented SNNs while maintaining iso-accuracy.

Fig. 1(A) shows a component-wise energy distribution for the CIFAR10-trained
VGG-16 network on 64$\times$64 4-bit RRAM IMC architecture. Among these, the
digital peripherals (containing crossbar input switching circuits, buffers,
and accumulators) entail the highest energy cost (45%). The IMC-crossbar and
analog-to-digital converter (ADC) consumes the second highest energy (25%).
Efforts to lower energy consumption have been made by previous works. As an
example, prior IMC-aware algorithm-hardware co-design techniques [13, 22],
have used pruning, quantization to reduce the ADC and crossbar energy, and
area cost. However, in the case of SNNs, the improvement is rather limited
because the crossbar and ADC only occupy 25% of the overall energy cost.

Unlike ANNs, the number of timesteps in SNNs plays an important role in
hardware performance, which is orthogonal to data precision or sparsity. In
Fig. 1(B) we investigate how timesteps affect the energy consumption and
latency of an SNN. Note that both metrics are normalized to the performance of
a 1-timestep SNN. We find that both energy consumption and latency scale
linearly with the number of timesteps, up to $4.9\times$ more energy and
$8\times$ more latency when changing the number of timesteps from 1 to 8. More
importantly, if one can reduce the number of timesteps in SNNs, then all parts
in Fig. 1(A) can benefit from the energy and latency savings. These findings
highlight the tremendous potential to optimize SNNs’ performance on IMC
hardware.

In fact, [8, 12, 3] have explored ways to reduce the number of timesteps from
an algorithmic perspective. They all train an SNN with a high number of
timesteps first and then finetune the model with fewer timesteps later.
However, their method decreases the number of timesteps for all input samples,
thereby inevitably leading to an accuracy-timestep trade-off. In this paper,
we tackle this problem with another solution. _We view the number of timesteps
during inference as a variable conditional to each input sample._ We call our
method Dynamic Timestep Spiking Neural Network (DT-SNN) as it varies the
number of timesteps based on each input sample. In particular, we use entropy
thresholding to determine the appropriate number of timesteps for each sample.
To further optimize our algorithm in practice, we design a new training loss
function and implement our algorithm on an IMC architecture.

The main contributions of our work are summarized below:

  1. 1.

Based on what we have seen thus far, this is the first work that changes the
number of timesteps in SNNs based on the input, reducing computational
overhead and increasing inference efficiency without compromising task
performance.

  2. 2.

To achieve that goal, we propose using entropy thresholding to distinguish the
number of timesteps required. Meanwhile, we also provide a new training loss
function and an IMC implementation of our method.

  3. 3.

Extensive experiments are carried out to demonstrate the efficacy and
efficiency of DT-SNN. For example, the DT-SNN ResNet-19 achieves the same
accuracy as the 4-timestep SNN ResNet-19 with only an average of 1.27-timestep
on the CIFAR-10 dataset, reducing 84% energy-delay-product.

##  II Preliminaries

We start by introducing the basic background of SNNs. We denote the overall
spiking neural network as a function $f_{T}({\bm{x}})$ (${\bm{x}}$ is the
input image), its forward propagation can be formulated as

${\bm{y}}=f_{T}({\bm{x}})=\frac{1}{T}\sum_{t=1}^{T}h\circ g^{L}\circ
g^{L-1}\circ g^{L-2}\circ\cdots g^{1}({\bm{x}}),$ (1)

where $g^{\ell}({\bm{x}})=\mathrm{LIF}({\bm{W}}^{\ell}{\bm{x}})$ denotes the
$\ell$-th block. A block contains a convolutional layer, a leaky integrate-
and-fire (LIF) layer, and an optional normalization layer placed in between
the former two layers [23]. $L$ represents the total number of blocks in the
network and $h(\cdot)$ denotes the final linear classifier. In this work, we
use the direct encoding method, i.e., using $g^{1}({\bm{x}})$ to encode the
input tensor into spike trains, as done in recent SNN works [20]. To get the
final prediction, we repeat the inference process $T$ times and average the
output from the classifier.

SNNs emulate the biological neurons using LIF layers. For each timestep, the
input current charges the membrane potential ${\bm{u}}$ in the LIF neurons.
When the membrane potential exceeds a certain threshold, a spike ${\bm{s}}$
will be fired to the next layer, given by

${{\bm{u}}}^{\ell}[t+1]=\tau{\bm{u}}^{\ell}[t]+{\bm{W}}^{\ell}{\bm{s}}^{\ell}[t],$
(2) ${\bm{s}}^{\ell+1}[t+1]=\begin{cases}1&\text{if
}\bm{u}^{\ell}[t+1]>V_{th}\\\ 0&\text{otherwise}\end{cases},\ \ $ (3)

where $\tau\in(0,1]$ is the leaky factor, mimicking the potential decay. If a
spike is fired, the membrane potential will be reset to 0, _i.e._
$({\bm{u}}[t+1]={\bm{u}}[t+1]*(1-{\bm{s}}[t+1]))$.

In the spiking neurons, all functions except the spike firing function (Eq.
(3)) can be normally differentiated. The firing function generates a 0
gradient if ${\bm{u}}^{\ell}[t+1]\neq V_{th}$, otherwise, it generates a
gradient with infinity. This impedes the gradient-based optimization in SNNs.
To overcome this problem, we leverage the surrogate gradient training method
[19]. Specifically, in the forward propagation, we keep the original LIF
neuron dynamics, while in the backward propagation, we use another function:

$\frac{\partial{\bm{s}}^{\ell}[t]}{\partial{\bm{u}}^{\ell}[t]}=\max(0,V_{th}-|{\bm{u}}^{\ell}[t]-V_{th}|)$
(4)

##  III Methodology

In this section, we first introduce the algorithm for our work. Then we
demonstrate the hardware implementation of DT-SNN.

Figure 2: The impact of the number of timesteps on the accuracy. We test the
spiking VGG-16 on three datasets (CIFAR10, CIFAR100, TinyImageNet).

###  III-A Dynamic Timestep Spiking Neural Network

Because spikes in an SNN are sparse and binary, the number of timesteps,
therefore, controls the density of information inside the SNNs. Generally,
more timesteps help SNNs explore more temporal information and thus achieve
higher task performance. Fig. 2 demonstrates that when the number of timesteps
of an SNN VGG-16 is increased from 1 to 4 during inference, the accuracy
increases as well. Together with the hardware performance as shown in Fig.
1(B), the number of timesteps $T$ controls a trade-off between hardware
performance and task performance on an SNN model $f_{T}(\cdot)$.

Unlike the conventional approach where $T$ is selected and fixed for all
images, we propose a conditional probability of selecting $T$ for different
input $\bm{x}$. We call our method the Dynamic Timestep Spiking Neural Network
(DT-SNN). More concretely, denote $\mathbb{P}(T|\bm{x})$ as the conditional
probability of $T$ with respect to $\bm{x}$, DT-SNN is given by

$f_{\hat{T}\sim\mathbb{P}(T|\bm{x})}=\frac{1}{\hat{T}}\sum_{t=1}^{\hat{T}}h\circ
g^{L}\circ g^{L-1}\circ g^{L-2}\circ\cdots g^{1}({\bm{x}}).$ (5)

DT-SNN allows allocating a different number of timesteps for each input
sample. As seen in Fig. 2, we find that the majority of samples can be
correctly classified with fewer timesteps. For example, on the CIFAR-100
dataset, 69.39% of overall test data can be correctly predicted using only 2
timesteps. Yet only 2.9% of test data needs full timesteps ($T=4$) to get the
right prediction. If we compare the hardware performance, the 4-timestep model
brings 86% more energy consumption and 100% more latency than the 2-timestep
model. This observation is also applicable to other datasets like CIFAR10 and
TinyImageNet.

##### Choosing the Right $T$

Our objective in DT-SNN is to reduce the unnecessary timesteps as much as
possible while not compromising accuracy. However, finding the appropriate
timestep for each input data is non-trivial. In this work, we use entropy to
determine $T$. Formally, given a dataset that has $K$ classes, the prediction
probability $\pi(\bm{y}|\bm{x})$ is calculated by the Softmax function
$(\sigma(\cdot))$, given by

$\pi(\bm{y}_{i}|\bm{x})=\sigma_{i}(f(\bm{x}))=\frac{\exp(f(\bm{x})_{i})}{\sum_{j=1}^{K}\exp(f(\bm{x})_{j})},$
(6)

where $\pi(\bm{y}_{i}|\bm{x})$ is the probability of predicting $i$-th class.
The entropy can be further calculated by

$E_{f}(\bm{x})=-\frac{1}{\log
K}\sum_{i=1}^{K}\pi(\bm{y}_{i}|\bm{x})\log\pi(\bm{y}_{i}|\bm{x}).$ (7)

Here, $\log K$ ensures the final entropy is normalized to $(0,1]$. The entropy
measures the state of uncertainty. For instance, if all classes have an equal
probability of $\frac{1}{K}$, the entropy will become 1, meaning the current
state is completely random and uncertain. Instead, if one class’s probability
is approaching 1 while others are approaching 0, the entropy moves towards 0,
indicating the state is becoming certain.

Generally, the prediction accuracy is highly correlated with entropy. If the
model is certain about some inputs (low entropy), the prediction would be
highly probable to be correct, and vice versa [5]. Therefore, we select the
$T$ if the entropy is lower than some pre-defined threshold $\theta$, given by

$\hat{T}(\bm{x})=\operatorname*{arg\,min}_{\hat{T}}\\{E_{f_{\hat{T}}}(\bm{x})<\theta|1\leq\hat{T}<T\\}\cup\\{T\\}.$
(8)

Here, the $\hat{T}$ is selected based on the lowest timestep that can have
lower-than-$\theta$ entropy. If none of them can have confident output, the
SNN will use maximum timesteps, _i.e.,_ $T$.

##### Training DT-SNN

Originally, the loss function for training an SNN is the cross-entropy
function, given by:

$\mathcal{L}(\bm{x},\bm{z})=-\frac{1}{B}\sum_{i=1}^{K}\bm{z}_{i}\log\pi(f_{T}(\bm{x})_{i}|\bm{x}),$
(9)

where ${\bm{z}}$ is the label vector and $B$ is the batch size. Although the
output from lower timesteps implicitly contributes to $f_{T}({\bm{x}})$, there
lacks some explicit guidance to them. As shown in Fig. 2, the accuracy in the
first timestep is always low. Here, we propose to explicitly add a loss
function to each timestep output. The new loss function is defined as:

$\mathcal{L}(\bm{x},\bm{z})=-\frac{1}{TB}\sum_{t=1}^{T}\sum_{i=1}^{K}\bm{z}_{i}\log\pi(f_{t}(\bm{x})_{i}|\bm{x}),$
(10)

In practice, we find this loss function does not change much training time on
GPUs. Also, we will demonstrate that adding this loss function can benefit the
accuracy of the outputs from all timesteps, further improving our DT-SNN.

##### Relation to Early Exit in ANN

Conceptually our DT-SNN is similar to the Early-Exit technique in ANN [18, 1]
which adds multiple exits to different layers. Here, we want to clarify the
relation between our DT-SNN and early exit: (1) DT-SNN is operated in the time
dimension, it naturally fits SNNs and does not require any additional layers,
while early exit has to add classifier layers in each branch; (2) DT-SNN has a
higher potential than early exit: in experiments section, we will show that
the majority of the examples can only use the first timestep, while the first
exit in ANNs outputs marginal examples. Furthermore, DT-SNN is fully
complementary to the early exit, that is, we can further add the early exit
technique to SNN to fulfill even higher efficiency.

###  III-B Hardware Implementation

Figure 3: Figure showing (a) monolithic-tiled IMC architecture implementation of an SNN (b) architecture of the $\sigma-E$ module for softmax and entropy value computation. TABLE I: Hardware Implementation Parameters. Technology | 32nm CMOS
---|---
Crossbar Size & Crossbars/Tile | 64 & 64
Device & Weight Precision | 4-bit RRAM ($\sigma$/$\mu$=20%) & 8-bit
$R_{off}$/$R_{on}$ | 10 at $R_{on}$=20k$\Omega$
GB, Tile & PE Buffer Size | 20KB, 10KB & 5KB
$V_{DD}$ & $V_{read}$ | 0.9V & 0.1V
$\sigma$ & $E$ LUT size | 3KB & 3KB
TABLE II: Comparison between static SNN and DT-SNN in terms of timesteps, accuracy, and normalized energy cost. Model | Method | CIFAR-10 | CIFAR-100 | TinyImageNet | CIFAR10-DVS
---|---|---|---|---|---
$T$ | Acc. | Energy | $T$ | Acc. | Energy | $T$ | Acc. | Energy | $T$ | Acc. | Energy
VGG-16 | SNN | 4 | 93.17 | 1.00$\times$ | 4 | 72.39 | 1.00$\times$ | 4 | 58.48 | 1.00$\times$ | 10 | 74.4 | 1.00$\times$
DT-SNN | 1.46 | 93.58 | 0.46$\times$ | 2.03 | 72.43 | 0.56$\times$ | 2.14 | 58.27 | 0.60$\times$ | 5.25 | 74.4 | 0.54$\times$
ResNet-19 | SNN | 4 | 94.03 | 1.00$\times$ | 4 | 73.31 | 1.00$\times$ | 4 | 56.70 | 1.00$\times$ | 10 | 75.0 | 1.00$\times$
DT-SNN | 1.27 | 93.87 | 0.41$\times$ | 1.90 | 73.48 | 0.53$\times$ | 2.01 | 56.96 | 0.56$\times$ | 5.02 | 74.8 | 0.52$\times$

We implement DT-SNN on a tiled-monolithic chip architecture [2] as shown in
Fig. 3a. First, the individual layers of an SNN are mapped onto tiles. The
number of tiles occupied by an SNN layer depends on factors such as the
crossbar size, the number of input and output channels in a layer, kernel
size, and the number of crossbars per tile. To implement DT-SNN-specific
functionality, we incorporate the following modifications in conventionally
used architectures 1) A digital $\sigma-{E}$ module to jointly compute the
SoftMax ($\sigma(\cdot)$) and entropy ($E$) followed by threshold ($\theta$)
comparison to detect whether to exit or not. 2) Timesteps are processed
sequentially without pipelining. This eliminates the delay and hardware
overhead (energy and area cost) required to empty the pipeline in case of
dynamic timestep inference. The tiles additionally contain global accumulators
(GA) and global buffers (GB) for accumulating partial sums and storing the
intermediate outputs, respectively from different processing elements (PE). At
the tile level, all modules are connected via a Network-on-Chip (NoC)
interconnect. Each tile consists of several PEs, accumulators, and buffers.
Each PE contains several crossbars, accumulators and buffers. The PEs and
crossbars are connected by an H-Tree interconnect. We use a standard 2D-IMC
crossbar connected to peripherals such as switch matrix, multiplexers, analog-
to-digital converters (ADCs), and Shift-$\&$-Add circuits. The switch matrix
provides input voltages at the source lines (SL) while simultaneously
activating the word lines (WL). The voltages accumulate over the bit lines
(BL). The analog MAC value (partial sum output) is converted to a digital
value using the ADC. Multiplexers enable resource sharing of ADCs and
Shift-&-Add circuits among multiple crossbar columns to reduce the area
overheads. The digital partial sum outputs from different crossbars, PEs and
tiles are accumulated using the PE, tile, and global accumulators,
respectively. For all layers except the last, the final MAC outputs from the
GA are transferred to the LIF module for the non-linear activation
functionality. The spike outputs are relayed to the tiles mapping the
subsequent layers.

For the last layer (a fully connected layer) the GA accumulated MAC output is
directed to the $\sigma-E$ module as shown in Fig. 3a (using red arrows).
Inside the $\sigma-E$ module MAC outputs are stored in the $y$-FIFO buffer.
The depth of $y$-FIFO depends on the dataset. For example, in CIFAR10, the
FIFO depth is 10. Data from the $y$-FIFO is passed to the address lines of the
$\sigma$-LUT to compute $\sigma$ which are pushed into the $\sigma$-FIFO. The
$\sigma$-FIFO outputs are sent as inputs to the Entropy Module that contains
LUT for $\log(\sigma)$ computation. The Entropy Module additionally contains a
multiplier and accumulator circuit (comprised of an adder and register) to
implement the entropy computation using Eq. (7). If the computed entropy is
less than the threshold $\theta$, the inference is terminated and new input
data is loaded into the GB. DT-SNN is implemented on the IMC architecture
using parameters shown in Table I.

Energy Consumption of the $\sigma-E$ module: Based on 32nm CMOS
implementations, we find that the energy consumed by the $\sigma-E$ module for
one timestep is merely $2e^{-5}\times$ of the 1 timestep inference energy
consumed by the IMC architecture, which is negligible.

##  IV Experimental Results

In this section, we present the evaluation of both the task performance and
the hardware performance of DT-SNN, highlighting its extreme efficacy and
efficiency.

Figure 4: Comparison between static SNN and DT-SNN in terms of Energy-Delay-
Product (EDP) (normalized to the static SNN).  Figure 5: Accuracy vs. EDP
curve for both static SNN (drawn by pink line) and DT-SNN (drawn by blue
line). For the static SNN, we report the accuracy and EDP when
$T=\\{1,2,3,4\\}$. For the DT-SNN, we draw the pie charts illustrating the
distribution of $\hat{T}({\bm{x}})$.

###  IV-A Comparison with Static SNN

We select 4 popular visual benchmarks for evaluation, CIFAR-10 [9], CIFAR-100
[9], TinyImageNet [4], and CIFAR10-DVS [11] dataset. For the architecture, we
choose to use VGG-16 [16] and ResNet-19 [7]. We compare our DT-SNN with the
static SNN, _i.e.,_ an SNN that uses a fixed number of timesteps for all
inputs. The training method for the static SNN and the DT-SNN are kept the
same, except that the static SNN uses Eq. (9) as training loss while DT-SNN
uses Eq. (10). We train them with a batch size of 256, a learning rate of 0.1
followed by a cosine decay. The L2 regularization is set to 0.0005. The number
of timesteps is set to 4 as done in existing work [23]. For task metrics, we
report the top-1 accuracy on the visual benchmarks. As for hardware metrics,
we measure them based on the parameters shown in Table I and further normalize
them w.r.t. static SNNs. We report the number of timesteps ($T$), energy, and
energy-delay-product (EDP). Note that the cost of DT-SNN is varied across
input data, thus we average the hardware metrics in the test dataset.

####  IV-A1 Comparison of Accuracy, Energy Cost, and $T$

We summarize the results of 4 datasets in Table II. Here, we test static SNN
with the full number of timesteps, _i.e.,_ $T=4$, and compare the hardware
performances with DT-SNN under a similar accuracy level. We find that DT-SNN
only needs 1.46 average timesteps on the CIFAR-10 dataset, bringing more than
50% energy saving. For the other three datasets, DT-SNN requires roughly half
the number of timesteps used in a static SNN model. Nevertheless, DT-SNN
reduces at least 40% of energy cost when compared to static SNN.

####  IV-A2 Comparison of EDP

We next compare the EDP between static SNNs and DT-SNNs. EDP is more suitable
for measuring a holistic performance in hardware because it considers both
time and energy efficiency. Fig. 4 shows the EDP comparison normalized by the
EDP of the static SNN. We can find that DT-SNN is extremely efficient as it
reduces 61.2%$\sim$80.9% EDP of static SNNs. These results highlight the
efficiency brought by our method, leading to both energy cost and latency
reduction.

####  IV-A3 Accuracy vs. EDP curve

The static SNN can adjust the number of timesteps for all inputs to balance
its accuracy and efficiency. Our DT-SNN can also adjust threshold $\theta$ to
balance such a trade-off. Here, we draw an accuracy-EDP curve in Fig. 5. _Note
that here the EDP is normalized to the EDP of the 1-timestep static SNN._ We
evaluate static SNN at 1,2,3, and 4 timesteps, and evaluate DT-SNN using three
different thresholds. It can be seen that our DT-SNN is placed in the top-left
corner, indicating a better accuracy-EDP trade-off than the static SNN.
Remarkably, DT-SNN can bring significant improvement in low-timestep
scenarios. For instance, DT-SNN VGG-16 increases the accuracy by 17% when
compared to the 1-timestep static counterpart on the CIFAR-10 dataset, while
it only has $\sim$10% higher EDP.

In order to further visualize the dynamic timesteps in our method, Fig. 5
provides three pie charts in each case, which show the percentage of input
examples that are inferred with 1, 2, 3, or 4 timesteps. Notably, $T=1$ is
usually the most selected case for the input due to the fact that most input
examples can be correctly predicted using only 1 timestep. As the threshold
decreases, more images start to use higher timesteps. Overall, we find $T=3$
and $T=4$ are rarely used in DT-SNN, demonstrating the effectiveness of our
method to reduce redundant timesteps.

Figure 6: Accuracy vs. the number of timesteps. (A) Comparison with prior
works, (B) Comparison under non-ideal (NI) device variation in the IMC.

####  IV-A4 Comparison with Prior Work

Here, we also compare our method with prior work on SNNs. We compare tdBN [23]
and Dspike [12] with our static-SNN and DT-SNN trained with Eq. (10). Fig.
6(A) shows the accuracy under different $T$. Our DT-SNN reaches a new state of
the art.

####  IV-A5 Non-Ideal Accuracy with Device Variations

So far, all the experiments are performed without considering device
conductance variation. In Fig. 6(B), we compare DT-SNN and static SNN under
20% device conductance variation. We simulate this by adding noise to the
weights post-training. We see that DT-SNN still maintains higher accuracy
while eliminating the redundant timesteps compared to static-SNN .

###  IV-B Acceleration in General Processors

TABLE III: Throughput comparison on the GPU.

Dataset&Model CIFAR10 VGG-16 CIFAR10 ResNet-19 Method $T$ Acc. Throughput $T$
Acc. Throughput SNN 1 76.30 199.3 1 89.58 185.3 2 91.34 121.8 2 92.51 105.9 3
92.54 85.19 3 93.40 77.14 4 93.01 64.34 4 94.03 62.65 DT-SNN 1.10 93.01 176.7
1.07 92.95 169.3 1.46 93.58 142.0 1.27 93.87 147.5 2.11 93.71 105.9 1.46 94.07
135.7

In the previous section, we demonstrated that our DT-SNN can accelerate the
inference on an IMC architecture. Here, we show that apart from the IMC
architecture we simulated, our method can also be applicable to other types of
hardware like digital processors. To this end, we measure the inference
throughput (_i.e.,_ the number of images inferred per second) on an RTX 2080Ti
GPU simulated by PyTorch using batch size 1. Table III lists the accuracy,
(averaged) timesteps, and throughput. As can be seen from the table, the
throughput on GPU significantly reduces as the timesteps are increased.
Compared to static SNNs, our DT-SNN substantially improves the throughput
while not sacrificing accuracy. For example, DT-SNN ResNet-19 with 1.07
averaged timesteps can infer 169.3 images per second, quite close to the
1-timestep static SNN (185.3 images per second), yet still bringing 3.4%
accuracy improvement.

###  IV-C Ablation Study

In this section, we ablate the training loss function choice. To verify this,
we train static an SNN VGG-16 on CIFAR-10 either with Eq. (9) or Eq. (10) and
test the corresponding accuracies of both static SNN and DT-SNN. Fig. 7
demonstrates the comparison between these two loss functions. We find that our
training loss boosts the accuracy of all timesteps in the static SNN. In
particular, the first timestep of VGG-16 changes from 76.3% to 91.5%. This
result proves that explicit guidance from label annotations should be added to
the lower timesteps. Meanwhile, it also increases the full timestep
performance, resulting in a 0.6% accuracy increase on VGG-16.

This improvement is extremely beneficial to DT-SNN. According to the pie
charts describing the distribution of $\hat{T}$, we find our training loss
function enables a smaller number of timesteps to be required to classify the
test data, thus reducing the EDP significantly.

###  IV-D Visualization

Figure 7: Comparison between two loss functions using accuracy-EDP curves on
the CIFAR-10 dataset. Annotations are aligned with Fig. 5.

In this section, we visualize the input images that are differentiated by our
DT-SNN. Ideally, we anticipate DT-SNN can identify whether an image is easy or
hard to infer (corresponding to 4 or 1 timestep). To maximize the
differentiation, we use a low threshold to filter out the high timesteps, so
that only the easiest images can be classified in the first timestep and vice
versa. Fig. 8 presents the results on the TinyImageNet dataset. Generally, we
find the images inferred with 1 timestep exhibit a simple structure: a clear
object placed in the center of a clean background. In contrast, hard images
require 4 timesteps and usually mix the background and the object together,
making the object imperceptible.

##  V Conclusion

In this work, we introduce the Dynamic Timestep Spiking Neural Network, a
simple yet significantly effective method that selects different timesteps for
different input samples. DT-SNN determines the suitable timestep based on the
confidence of the model output, seamlessly fitting the nature of sequential
processing over time in SNNs. Moreover, DT-SNN is practical. It can be
deployed to IMC architecture and even general digital processors. Extensive
experiments prove that DT-SNN establishes a new state-of-the-art trade-off
between hardware efficiency and task performance.

Figure 8: Visualization of input samples from $T=1$ (left) and $T=4$ (right).

## References

  * [1] Tolga Bolukbasi et al.  Adaptive neural networks for efficient inference.  In ICML, pages 527–536. PMLR, 2017.
  * [2] Pai-Yu Chen et al.  Neurosim: A circuit-level macro model for benchmarking neuro-inspired architectures in online learning.  IEEE TCAD, 37(12):3067–3080, 2018.
  * [3] Sayeed Shafayet Chowdhury et al.  Spatio-temporal pruning and quantization for low-latency spiking neural networks.  In 2021 IJCNN, pages 1–9. IEEE, 2021.
  * [4] Jia Deng et al.  Imagenet: A large-scale hierarchical image database.  In 2009 CVPR, pages 248–255. Ieee, 2009.
  * [5] Chuan Guo et al.  On calibration of modern neural networks.  In ICML, pages 1321–1330. PMLR, 2017.
  * [6] Song Han et al.  Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding.  arXiv preprint arXiv:1510.00149, 2015.
  * [7] Kaiming He et al.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on CVPR, pages 770–778, 2016\.
  * [8] Youngeun Kim et al.  Revisiting batch normalization for training low-latency deep spiking neural networks from scratch.  Frontiers in neuroscience, page 1638, 2021.
  * [9] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * [10] Yann LeCun et al.  Deep learning.  nature, 521(7553):436–444, 2015.
  * [11] Hongmin Li et al.  Cifar10-dvs: an event-stream dataset for object classification.  Frontiers in neuroscience, 11:309, 2017.
  * [12] Yuhang Li et al.  Differentiable spike: Rethinking gradient-descent for training spiking neural networks.  NeurIPS, 34, 2021.
  * [13] Jie Peng et al.  Cmq: Crossbar-aware neural network mixed-precision quantization via differentiable architecture search.  IEEE TCAD, 2022.
  * [14] Kaushik Roy et al.  Towards spike-based machine intelligence with neuromorphic computing.  Nature, 575(7784):607–617, 2019.
  * [15] Abu Sebastian et al.  Memory devices and applications for in-memory computing.  Nature nanotechnology, 15(7):529–544, 2020.
  * [16] Karen Simonyan et al.  Very deep convolutional networks for large-scale image recognition.  arXiv preprint arXiv:1409.1556, 2014.
  * [17] Amirhossein Tavanaei et al.  Deep learning in spiking neural networks.  Neural Networks, 111:47–63, 2019.
  * [18] Surat Teerapittayanon et al.  Branchynet: Fast inference via early exiting from deep neural networks.  In 2016 23rd ICPR, pages 2464–2469. IEEE, 2016.
  * [19] Yujie Wu et al.  Spatio-temporal backpropagation for training high-performance spiking neural networks.  Frontiers in neuroscience, 12:331, 2018.
  * [20] Yujie Wu et al.  Direct training for spiking neural networks: Faster, larger, better.  In AAAI, volume 33, pages 1311–1318, 2019.
  * [21] Ruokai Yin et al.  Sata: Sparsity-aware training accelerator for spiking neural networks.  arXiv preprint arXiv:2204.05422, 2022.
  * [22] Geng Yuan et al.  Tinyadc: Peripheral circuit-aware weight pruning framework for mixed-signal dnn accelerators.  In 2021 DATE, pages 926–931. IEEE, 2021.
  * [23] Hanle Zheng et al.  Going deeper with directly-trained larger spiking neural networks.  In AAAI, volume 35, pages 11062–11070, 2021.


# The Particle in a Box
Warrants an Examination

John R<EMAIL_ADDRESS>
Department of Physics and Department of Mathematics
University of Florida, Gainesville, FL 32611-8440

###### Abstract

The particle in a box is a simple model that has a classical Hamiltonian
$H=p^{2}$ (using $2m=1$), with a limited coordinate space, $-b<q<b$, where
$0<b<\infty$. Using canonical quantization, this example has been fully
studied thanks to its simplicity, and it is a common example for beginners to
understand. Despite its repeated analysis, there is a feature that puts the
past results into question. In addition to pointing out the quantization
issue, the procedures of affine quantization can lead to a proper quantization
that nesaeccsrily points toward more complicated eigenfunctions and
eigenvalues, which deserve to be solved.

##  1 A Standard Quantization of
the Particle in a Box

The basic operators of canonical quantization (CQ), namely $P$ and $Q$, have
appropriate representations that range over entire real values, which leads
both operators to being self-adjoint, i.e., $P^{\dagger}=P$ and
$Q^{\dagger}=Q$. These operators admit a Schrödinger representation, e.g..,
$P\rightarrow-i\hbar(d/dx)$ and $Q\rightarrow x$, with $-\infty<x<\infty$.

Since our story takes place entirely in the interval $-b<x<b$, we seek the
solutions of

$\displaystyle-\hbar^{2}(d^{2}\phi(x)/dx^{2})\;=E\;\phi(x)$ (1)

with the requirement $\phi(-b)=\phi(b)=0$ to match the fact that $\phi(x)=0$
for all of $x\leq-b$ and $x\geq b$, thanks to ‘implicit infinite walls’. This
connection at $x=\pm b$ ensures that $\phi(x)$ will be continuous. Clearly,
the $\cos$ and $\sin$ functions fit the requirement, like
$\phi_{n}(x)=\cos(n\pi x/2b)$, for $n=1,3,5,...$, and $\phi_{n}(x)=\sin(n\pi
x/2b)$, for$n=2,4,6,...$. The eigenvalues then become
$E_{n}=n^{2}\pi^{2}/4b^{2}$, now for $n=1,2,3,4,5...$. Evidently, two
derivatives of either $\sin$ or $\cos$ just leads to a simple factor times the
original term. The norm of the eigenfunctions becomes
${\textstyle\int}_{-b}^{b}|\cos(n\pi x/2b)|^{2}\;dx<\infty$ for $n=1,3,5,...$,
and $\int_{-b}^{b}|\sin(n\pi x/2b)|^{2}\;dx<\infty$ for $n=2,4,6,...$. Hence,
the Hilbert space is composed of arbitrary sums of these eigenfunctions that
lead to finite normalizations.

Clearly, this first section points toward a bona-fide quantization of a
particle in a box that has been uniformly accepted; for a standard analysis,
see [1].

##  2 A Common Review of Derivatives

###  2.1 A toy example of the problem

The feature that is of concern deals with the proper derivatives of the
prospective eigenfunctions. This feature can be seen in a simple example.

Consider the continuous function $f(x)$, where $-1<x<1$, which is defined as
$f(x)=0$ for $-1\leq x\leq 0$, and $f(x)=x$ for $0\leq x\leq 1$.222Picture a
flat floor under a flat door that leans upward from the floor at 45 degrees. A
simple review of derivatives shows that
$f^{\prime}(x)=\lim_{\epsilon\rightarrow
0}[f(x+\epsilon)-f(x-\epsilon)]/2\epsilon$, which, for our example, leads to
$f^{\prime}(x)=0$ for $x<0$ and $f^{\prime}(x)=1$ for $x>0$. In addition, we
find that $f^{\prime}(0)=1/2$. This result has led to $f^{\prime}(x)$ becoming
a discontinuous function. The next step is defining
$f^{\prime\prime}(x)=\lim_{\epsilon\rightarrow
0}[f^{\prime}(x+\epsilon)-f^{\prime}(x-\epsilon)]/2\epsilon$. The result is
that $f^{\prime\prime}(x)=0$ for $x<0$ and $f^{\prime\prime}(x)=0$ for $x>0$.
In addition, it follows that $f^{\prime\prime}(0)=\lim_{\epsilon\rightarrow
0}[f^{\prime}(\epsilon)-f^{\prime}(-\epsilon)]/2\epsilon]=\infty$. If we
imagine that $2\epsilon\rightarrow dx$, then $f^{\prime\prime}(x)=\delta(x)$,
which is Dirac’s delta function, formally chosen so that
${\textstyle\int}_{-1}^{1}\delta(x)\,dx=1$. Summarizing, we have found that
$f^{\prime\prime}(x)=\delta(x)$. This result has been because $f^{\prime}(x)$
is a discontinuous function. Let us add that while
${\textstyle\int}_{-1}^{1}|f^{\prime\prime}(x)|\;dx<\infty$, it follows that
${\textstyle\int}_{-1}^{1}|f^{\prime\prime}(x)|^{2}\;dx=\infty$, which is the
reason to reject $f^{\prime\prime}(x)$ in any Hilbert space.

###  2.2 Application to a particle in a box

The proposed eigenfunction ground state of the particle in a box starts with

$\displaystyle\hskip-25.00003pt\phi_{1}(x)=0\;\;\;(for\;x\leq-b)$
$\displaystyle=\cos(\pi x/2b)\;\;\;(for\;-b\leq x\leq b)$
$\displaystyle=0\;\;\;(for\;b\leq x).$ (2)

However, two derivatives of the eigenfunction ground state, i.e.,
$\phi_{1}(x)=\cos(\pi x/2b)$, lead to

$\displaystyle\hskip-26.00009pt\phi_{1}^{\prime\prime}(x)=0\;\;\;(for\;x<-b)$
$\displaystyle=(\pi^{2}/8b^{2})\,\delta(x+b)$
$\displaystyle=-(\pi^{2}/4b^{2})\cos(\pi x/2b)\;\;\;(for\;-b<x<b)$
$\displaystyle=-(\pi^{2}/8b^{2})\,\delta(x-b)$
$\displaystyle=0\;\;\;(for\;b<x).$ (3)

This equation does not match the desired form of
$-\hbar^{2}\phi_{1}^{\prime\prime}(x)=E_{1}\,\phi_{1}(x)$, and moreover,
${\textstyle\int}_{-\infty}^{\infty}|\phi_{1}^{\prime\prime}(x)|^{2}\;dx=\infty$,
which can not be accepted by any Hilbert space.

While we have only considered the standard ground state, every proposed
eigenfunction would have a similar story and none of them could be a proper
eigenfunction with a finite eigenvalue, e.g.,
$-\hbar^{2}\phi_{n}^{\prime\prime}(x)=E_{n}\;\phi_{n}(x)$ with
$|E_{n}|<\infty$. Nor would they belong to the usual Hilbert space due to
having $\delta(x\pm b)$ terms in their second derivative. In the author’s
opinion, this behavior fails the conventual quantization of a particle in a
box, as it was reviewed in Sec. 1.

In the next section, we will propose a valid quantization of the particle in a
box using an alternate quantization procedure.

###  2.3 Half of the expected eigenfunctions

The sentence, taken from Sec. 1, “This connection at $x\pm b$ ensures that
$\phi(x)$ will be continuous. Clearly, the $\cos$ and $\sin$ functions fit the
requirement, like $\phi_{n}(x)=\cos(n\pi x/2b)$, for $n=1,3,5,...$, and
$\phi_{n}(x)=\sin(n\pi x/2b)$, for $n=2,4,6,...$,” makes it clear that only
half of the sin and cos terms have been accepted. This means that the standard
box treatment leads to only half of the set of eigenfunctions that were
presumed.

##  3 A Valid Quantization of
the Particle in a Box

Instead of CQ, we use affine quantization (AQ) in order to produce a valid
quantization of the particle in a box. Fortunately, AQ has been designed to
deal with reduced coordinate spaces.

###  3.1 A brief review of affine quantization

Some classical problems require reduced coordinates, like the half-harmonic
oscillator with the classical Hamiltonian $H=(p^{2}+q^{2})/2$ provided that
$q>0$. In so doing, $P^{\dagger}\neq P$, and we seek substitutes for $p$ and
$P$. We choose the dilation variable $d=pq$ and $q$ as the new variables. To
ensure their independence, we eliminate $q=0$, and then discard all $q<0$,
keeping all $q>0$.333A real-life example, is cutting a long thread into two
parts and only keeping one of them. The basic affine quantum operators are now
$D=(P^{\dagger}Q+QP)/2$ and $Q>0$. Coherent states, used to connect classical
and quantum realms [2], have established that with special classical variables
we find that $H^{\prime}(d,q)\rightarrow H^{\prime}(D,Q)$. This implies that
the classical Hamiltonian can now be $H^{\prime}=(d^{2}/q^{2}+q^{2})/2$, along
with $q>0$, and the affine quantum Hamiltonian for this model then becomes

$\displaystyle
H^{\prime}=(DQ^{-2}D+Q^{2})/2=[P^{2}+(3/4)\hbar^{2}/Q^{2}+Q^{2}]/2\>.$ (4)

Happily, the ‘3/4’ term ensures that both $P^{\dagger}$ and
$(P^{\dagger})^{2}$ act like $P$ and $P^{2}$ in this equation.

This example has proved its validity in several papers that have exposed the
eigenfunctions and eigenvalues for the half-harmonic oscillator [3, 4, 5].
This favorable result for the eigenfunctions being of the form
$\phi_{n}(x)=x^{3/2}(polynomial)_{n}e^{-x^{2}/2\hbar}$, with $n=0,1,2,...$, a
fact that the first derivative remains a completely continuous function. To
see what the $x^{3/2}$ factor can do for the relevant two derivatives, just
examine the simple equation
$[-\hbar^{2}\,(d^{2}/dx^{2})+(3/4)\hbar^{2}/x^{2}]\;x^{3/2}=?$.

###  3.2 Affine quantization of the particle in a box

The example of the last section removed $q=0$ from the coordinate space. Now
we remove two points, namely $q=-b$ and $q=b$. That leaves three different
spaces, and we keep only the middle part, namely, where $-b<q<b$. We use a
suitable dilation variable, namely, $d^{\prime}=p(b^{2}-q^{2})$, along with
$q$, which is restricted by $-b<q<b$. Next, we offer
$D^{\prime}=[P^{\dagger}(b^{2}-Q^{2})+(b^{2}-Q^{2})P]/2$, and $-b<Q<b$.

Using the new affine variables, the classical Hamiltonian is
$p^{2}=d^{\prime}(b^{2}-q^{2})^{-2}d^{\prime}$, which is promoted to

$\displaystyle
H^{\prime}=D^{\prime}(b^{2}-Q^{2})^{-2}D^{\prime}=P^{2}+\hbar^{2}[2Q^{2}+b^{2}]/[(b^{2}-Q^{2})^{2}]\;.$
(5)

This $\hbar$-factor has been selected from a general study of various dilation
operators in [6].

It is noteworthy to examine the last equation when $Q$ accepts a Schrödinger
representation, i.e., $Q=x$, and $x$ is extremely close to either $\pm b$. In
that case, $[2x^{2}+b^{2}]/[(b^{2}-x^{2})^{2}]\simeq(3/4)/(b\pm x)^{2}$, a
relation that closely resembles the very strong properties from the present
model with those of the former model. Using this fact leads to the suggestion
that part of any eigenfunction of (5) is likely to be
$\psi(x)=(b^{2}-x^{2})^{3/2}(remainder)$.444A very different use of (5) is to
accept the outside space, $|x|>b$, and reject $|x|<b$, which then it becomes
an ‘anti-box’. Now this system has a similarity to a toy ‘black hole’. It may
happen that particles could pile up close to an ‘end of space’, while being
attracted by a simple, ‘gravity-like’, pull of a potential such as
$V(x)=W/|x|$. If you choose AQ, the $\hbar$-term in (5) could prevent the
particles from falling ‘out of space’.

Efforts to find eigenvalues and eigenfunctions for (5) are open to help shed
further information on this effort to find a proper quantization of the
particle in a box.

##  4 Conclusion

Affine quantization has also been able to help with other problems as well.
Recent efforts have been focused on quantum field theory [7], and the intro-
duction of an affine path integral quantization of gravity [8]. Another
article features a direct comparison of CQ and AQ regarding the half-harmonic
oscillator with its important implications for gravity and field theory [9].

Thanks are given to B.-S. Skagerstam, S. Shabanov, J. Govaerts, and A. Kempf
for helpful comments of earlier versions of this paper.

Note: The author has no conflicts to disclose.

## References

  * [1] Wikipedia: ‘Particle in a box’.
  * [2] J. Klauder,‘The Favored Classical Variables to Promote to Quantum Operators’, Journal of High Energy Physics, Gravitation and Cosmology 6, 828-832 (2020); DOI:10.4236/jhepgc.2020.64055.
  * [3] J. Klauder, ‘Quantum Gravity Made Easy’, Journal of High Energy Physics, Gravitation and Cosmology 6, 90-102 (2020); DOI:10.4236/jhepgc.2020.61009, Sec, 1.5.
  * [4] L. Gouba, ‘Affine Quantization on the Half Line’. Journal of High Energy Physics, Gravitation and Cosmology 7, 352-365 (2021); DOI:10.4236/jhepgc.2021.71019.
  * [5] C. Handy, ‘Affine Quantization of the Harmonic Oscillator on the Semi-bounded domain $(-b,\infty)$ for $b:0\rightarrow\infty$’; arXiv:2111:10700.
  * [6] R. Fantoni and J. Klauder, ‘Kinetic Factors in Affine Quantization and Their Role in Field Theory Monte Carlo’, International-Journal-of-Modern-Physics-A-0217-751X (accepted, but nor published yet).
  * [7] R. Fantoni and J. Klauder, ”Eliminating Nonrenormalizability Helps Prove Scaled Affine Quantization of $\varphi^{4}_{4}$ is Nontrivial” International Journal of Modern Physics A, Volume 37, Issue 6, id. 2250029; arXiv:2109.13447.
  * [8] J. Klauder, ‘A Straight Forward Path to a Path Integration of Einstein’s Gravity’; arXiv:2203.15141.
  * [9] J. Klauder, ‘Using a Toy Model to Improve the Quantization of Gravity and Field Theories’, Journal of High Energy Physics, Gravitation and Cosmology 8, 303-308 (2022); DOI:10.4236/jhepgc.2022.82024.