# Policy Contrastive Imitation Learning

Jialei Huang  Zhaoheng Yin  Yingdong Hu  Yang Gao

###### Abstract

Adversarial imitation learning (AIL) is a popular method that has recently
achieved much success. However, the performance of AIL is still unsatisfactory
on the more challenging tasks. We find that one of the major reasons is due to
the low quality of AIL discriminator representation. Since the AIL
discriminator is trained via binary classification that does not necessarily
discriminate the policy from the expert in a meaningful way, the resulting
reward might not be meaningful either. We propose a new method called Policy
Contrastive Imitation Learning (PCIL) to resolve this issue. PCIL learns a
contrastive representation space by anchoring on different policies and
generates a smooth cosine-similarity-based reward. Our proposed representation
learning objective can be viewed as a stronger version of the AIL objective
and provide a more meaningful comparison between the agent and the policy.
From a theoretical perspective, we show the validity of our method using the
apprenticeship learning framework. Furthermore, our empirical evaluation on
the DeepMind Control suite demonstrates that PCIL can achieve state-of-the-art
performance. Finally, qualitative results suggest that PCIL builds a smoother
and more meaningful representation space for imitation learning.

Imitation Learning, ICML

##  1 Introduction

Imitation is one of the fundamental capabilities of an intelligent agent
(Hussein et al., 2017). Animals and humans can acquire many skills by
mimicking each other (Byrne, 2009). In engineering, imitation learning also
enables many robotics applications. One mainstream class of imitation learning
algorithms is the adversarial imitation learning (AIL) (Ho & Ermon, 2016). AIL
converts the imitation task into a distribution matching problem and proposes
to imitate it by training a policy against an adversarial discriminator. AIL
has enjoyed great success on many imitation tasks: it achieves superior
performance (Ho & Ermon, 2016; Kostrikov et al., 2018), and has been
experimentally proven to alleviate some of the distributional drift issue, and
can work even without expert actions (Torabi et al., 2018b).However, AIL is
hard to train in practice, usually involving careful tuning of discriminator
neural network sizes and learning rates (Wang et al., 2017; Kim & Park, 2018;
Orsini et al., 2021). The fragility of the discriminator (Peng et al., 2018)
not only leads to poor performance but also severely limits the applicability
of AIL to a broader range of tasks.

Numerous techniques have been proposed to improve the performance of AIL, such
as using regularization and gradient penalties (Fu et al., 2017; Kostrikov et
al., 2018; Gulrajani et al., 2017). Some works also propose to use different
distribution metrics (e.g., KL divergence, Wasserstein distance) (Xiao et al.,
2019) for distribution matching and show some improvements. Though these
methods show encouraging results, we notice that they ignore one crucial
aspect of the problem: the representation of AIL’s discriminator. To be
specific, the discriminator in AIL is usually trained with binary
classification loss that distinguishes expert transitions from agent
transitions. This discriminator is then used to define rewards. However, since
the only goal of the discriminator is to distinguish the expert from the
agent, it does not necessarily learn a good, smooth representation space that
can provide a reasonable comparison between the behavior of two agents.Ideal
representations should be able to provide semantically meaningful signals to
compare the expert policy and the agent policy.

Figure 1: Comparison between the representation space of AIL method and our
method. Since the AIL methods use a binary classification objective to
distinguish expert and non-expert transitions, the representation space is
only required to separate two classes in two disjoint subspaces. So the
embedding space is not required to be semantically meaningful enough, e.g.
(Left) the distance between 2 expert data points may be even longer than the
distance between expert data point and sub-optimal non-expert data point.
(Right) We overcome this limitation by proposing PCIL. Our method enforces the
compactness of the expert’s representation. This ensures that the learned
representation can capture common, robust features of the expert’s
transitions, which leads to a more meaningful representation space.

In this paper, we propose a new algorithm called Policy Contrastive Imitation
Learning (PCIL) to achieve this goal. Instead of training with a binary-
classification objective, we propose to train a discriminator representation
space with the contrastive learning loss. Our method differs from the prior
representation learning approach in AIL in that we perform contrastive
learning between different policies. More specifically, we push the expert’s
representation together and pull the agent policy’s representation away from
them. We define the imitation learning reward via cosine similarity between
the policy’s and expert’s transition.

As is shown in Figure 1, the discriminator (binary classifier) might not have
a good representation space: the distance between two expert transitions can
be even larger than the distance between an expert transition and an agent
transition. This implies that the discriminator may not encode some common
features of the expert’s behavior and may use non-robust features to compare
the behavior. The lack of proper structure in the representation space of
traditional discriminators may yield low-quality AIL rewards. To alleviate
this issue, we explicitly define a constraint on the representation space,
requiring that the distance between expert transitions be smaller than their
distance to the agent’s transitions. This is a stronger constraint on the
discriminator’s representation: it is easy to derive a binary classifier from
our learned representation. However, the binary classifier’s representation
space does not necessarily satisfy our objective.

We validate PCIL by first showing that PCIL is performing distribution
matching between the expert and the agent policy. We also carry out empirical
evaluation by benchmarking our method on the DeepMind Control Suite (Tassa et
al., 2018). Experimental results show that our method is able to achieve
state-of-the-art results. Through ablation study and qualitative
visualization, we find that our method is more effective than prior
representation learning methods and able to provide a better representation
space for imitation learning.

In summary, our contributions in this paper are as follows. (1). We point out
a new direction to improve the performance of the AIL methods, i.e., going
beyond naive binary classification and leveraging more stable and meaningful
representation learning algorithms for imitation. (2). We propose an algorithm
called Policy Contrast Imitation Learning (PCIL) method that instantiates such
an improvement and establishes its connection to apprenticeship learning from
a theoretical perspective. (3). We evaluate our method on the DeepMind Control
Suite and achieve state-of-the-art performance. Through ablation studies, we
highlight its essential difference from previous contrastive learning methods
in AIL.

##  2 Preliminaries

###  2.1 Notations

In this paper, we model the imitation leanring problem as a markov decision
process
$\mathcal{M}=\left(\mathcal{S},\mathcal{A},p_{0}(s),p\left(s^{\prime}\mid
s,a\right),r\left(s,a,s^{\prime}\right),\gamma\right)$. Here, $\mathcal{S}$ is
the state space. $\mathcal{A}$ is the action space. $p_{0}(s)$ defines the
initial state distribution. $p(s^{\prime}|s,a)$ defines the transition
dynamics. $r(s,a,s^{\prime})$ is the reward function. $\gamma$ is the discount
factor. The goal is to maximize the expected return of the learned policy
$\pi$, which is defined by

$\displaystyle\mathcal{J}(\pi)=\mathbb{E}_{\begin{subarray}{c}s_{0}\sim
p_{0}(s),\\\ a_{i}\sim\pi(\cdot|s_{i}),\\\ s_{i+1}\sim
p(\cdot|s_{i},a_{i})\end{subarray}}\left[\sum_{k=0}^{\infty}\gamma^{k}r\left(s_{k},a_{k},s_{k+1}\right)\right].$

For the imitation learning problem, the algorithm does not have access to the
reward function and the transition dynamics. Instead, it is provided with an
expert demonstration dataset $\mathcal{D}$ sampled from an expert policy
$\pi_{E}$, which can perform well in $\mathcal{M}$. Here, $\mathcal{D}$ takes
the form of $\\{(s_{i}^{E},a_{i}^{E})\\}$, where $(s_{i}^{E},a_{i}^{E})$ is
sampled from $\rho_{E}$, the stationary state-action visiting distribution of
$\pi_{E}$. The imitation learning algorithm is then required to reproduce the
expert’s behavior with $\mathcal{D}$.

###  2.2 Adversarial Imitation Learning

One popular class of the imitation learning algorithm is the adversarial
imitation learning (AIL), whose vanilla version is Generative Adversarial
Imitation Learning (GAIL) (Ho & Ermon, 2016). The idea of GAIL is to minimize
the divergence between $\rho_{E}$ and $\rho_{\pi}$. It uses a discriminator
$D(s,a)$ to distinguish expert’s transitions
$(s^{E}_{i},a^{E}_{i})\sim\mathcal{D}$ from the policy transitions
$(s,a)\sim\rho_{\pi}$, which is trained by maximizing the objective

$\displaystyle\mathcal{L}=\ $
$\displaystyle\mathbb{E}_{(s_{i},a_{i})\sim\pi}[\log(D(s_{i},a_{i}))]+$
$\displaystyle\mathbb{E}_{(s_{i}^{E},a_{i}^{E})\sim\mathcal{D}}[\log(1-D(s_{i}^{E},a_{i}^{E}))].$

To achieve imitation, the agent policy is then required to fool the
discriminator, which can only be possible when the policy $\pi$ resembles the
expert $\pi_{E}$. Specifically, GAIL defines an adversarial reward
$r(s_{t},a_{t})=-\log(1-D(s_{t},a_{t}))$, and trains $\pi$ to maximize the
expected return with respect to this reward using on-policy RL algorithms.

##  3 Policy Contrastive Imitaiton Learning

###  3.1 Overview

We propose a novel representation-learning-based approach called Policy
Contrastive Imitation Learning to improve the AIL reward. The overview of PCIL
is illustrated in Figure 2. Our key insight is to learn a policy-contrastive
representation space. Unlike the contrastive learning studied in previous AIL
literature, the policy-contrastive representation here is learned by anchoring
on the rollout of different policies, based on which we can compare the
behavior of different policies in a meaningful way. We will discuss the
training of this representation in Section 3.2 and the reward design in
Section 3.3. Then, we will show the convergence of our algorithm in Section
3.4.

Figure 2: Illustration of our contrastive learning approach. We first select
an anchor state (the orange) from the expert trajectory. Then, we select a
positive state sample (the red) from another expert trajectory and a negative
state sample (the green) from the agent trajectory. We map these selected
states to the representation space. Finally, we push the representation of the
anchor state and the positive state together and pull the representation of
negative samples away from the representation of the anchor state.

###  3.2 Contrastive Policy Representation for Imitation

The vanilla AIL algorithms are based on unconstrained representations and can
be very non-robust. One possible approach to handle this problem is to learn a
more meaningful representation by contrastive learning. Researchers (He et
al., 2020) have found that it can usually learn semantically meaningful
representation, leading to better performance on downstream tasks, such as
classification. However, though it is effective in the field of supervised
learning, prior work (Chen et al., 2021) has found that it does not greatly
improve AIL much.

To understand the reason behind this, we first recall that the contrastive
learning method learns by drawing the representation of one training sample
${\mathbf{x}}$ towards a similar positive sample ${\mathbf{x}}_{p}$, and
pulling it away from a dissimilar, negative sample ${\mathbf{x}}_{n}$. Some
heuristic rules determine the choice of positive and negative samples: the
positive sample is usually defined as a data augmentation of ${\mathbf{x}}$.
Therefore, this augmentation decides what should be similar to the
representation. Nevertheless, we argue that the definition of similarity in
the previous works is not strong enough for AIL. This is because, in AIL, the
representation should also be able to let us discern good behavior from the
bad behavior. However, the difference between good and bad behavior here can
sometimes be very faint. For example, consider a case where a robot misses the
exact point to execute a certain action and, as a result, fails to accomplish
the task. The good and bad states right before this point may look very
similar. More specifically, the difference can simply be minor in a particular
physical measurement, like the distance. In this case, the representation
should consider this as a semantic component and be sensitive to such a
difference to succeed. Unfortunately, the difference between the positive
sample and the anchor sample in the previous representation learning methods
is usually very large and overwhelms the difference between good and bad
states. As a result, the model may not distinguish between good and bad states
effectively. Though combining the representation learning objective and the
AIL objective may help combat this problem by enforcing a hard distinguishing
constraint over contrastive representation, in practice we find this does not
work well (Section 4.4).

These observations motivate us to learn a representation that is semantically
meaningful and able to distinguish between good and bad states. We find a
surprisingly simple yet very effective approach: we can consider the samples
drawing from the same policy as the positive samples and the samples from all
the other policies as negative samples. In the case of imitation learning, our
samples can be naturally divided into two categories, namely expert and non-
expert samples. Then given the encoder $\Phi:\mathcal{S}\to\mathbb{S}$ that
maps the state to a representation vector in a high-dimensional sphere, we
define its infoNCE representation loss function as follows:

$\displaystyle\mathcal{L}=\mathbb{\mathbb{E}}_{\begin{subarray}{c}{\mathbf{x}}_{0}=({\mathbf{s}},{\mathbf{a}})\sim\mathcal{D},\\\
{\mathbf{x}}_{p}=({\mathbf{s}}_{p},{\mathbf{a}}_{p})\sim\mathcal{D},\\\
\tilde{{\mathbf{x}}}_{i}=(\tilde{{\mathbf{s}}}_{i},\tilde{{\mathbf{a}}}_{i})\sim\rho_{\pi}\end{subarray}}\left[\Phi({\mathbf{x}}_{0})^{T}\Phi({\mathbf{x}}_{p})+\right.$
(1)
$\displaystyle\left.{\log\left(\exp\Phi({\mathbf{x}}_{0})^{T}\Phi({\mathbf{x}}_{p})+\sum_{i=1}^{n}\exp\Phi({\mathbf{x}}_{0})^{T}\Phi(\tilde{{\mathbf{x}}}_{i})\right)}\right].$

Here, ${\mathbf{x}}_{0}$ is some state-action pair from the expert
transitions; ${\mathbf{x}}_{p}$ is some transition from the expert data acting
as positives; $\tilde{{\mathbf{x}}}_{i}$ is some agent transition working as
negatives. In other words, we require $\Phi$ to draw the expert samples
towards each other and pull all the policy samples away from the expert
samples.

Our proposed objective function is a strictly stronger constraint on the
discriminator. The binary classification discriminator can find any hyper-
plane that separates the two types of transitions, with no constraint on how
the transitions are embedded. However, our objective enforces the pair-wise
distance constraint between any triplets. One can derive a binary expert-
policy transition classifier from a trained $\Phi$ by computing
$\Phi({\mathbf{x}}_{0})^{T}\Phi({\mathbf{x}})>t$, where ${\mathbf{x}}_{0}$ is
any expert transition, ${\mathbf{x}}$ is the transition to be classified and
$t$ is some threshold. However, on the other hand, the binary classification
induced latent space might not satisfy our constraint. As illustrated in
Figure 1, the latent space for the binary classification discriminator might
have some expert-expert pairs that are even further away than some expert-
agent pairs.

Interestingly, our approach echoes the supervised contrastive learning (SCL)
(Khosla et al., 2020), which suggests that we consider all the samples in one
class as similar, positive samples. Here we also consider all the samples in
the expert demonstrations as similar. However, unlike SCL, we do not require
the samples from the agents to be similar to each other. This is because the
agent transitions are generated from different policies during the training
process.

###  3.3 Similarity-Based Imitation Reward

With a representation that can capture the difference between good and bad
states, we can then define a reward function to encourage imitation learning.
Though using an AIL-style reward with this representation is still possible,
we find that a better choice is to use a cosine similarity metric to define
the reward. It has several advantages: it is bounded and appears relatively
smooth in practice, leading to more stable learning. Concretely, we define:

$r({\mathbf{x}})=\Phi\left({\mathbf{x}}\right)^{T}\mathbb{E}_{{\mathbf{x}}_{E}\sim\mathcal{D}}\Phi\left({\mathbf{x}}_{E}\right).$
(2)

Nevertheless, in practice evaluating the latter expectation can be time-
consuming since $\Phi$ is frequently updated. Therefore, we use a random
expert sample for the reward calculation. From this reward, we can see that a
policy can only obtain high rewards when it frequently visits the expert’s
distribution. This naturally connects our method to the distribution matching,
and we provide the theoretical analysis of our algorithm in the following
subsection.

###  3.4 Theoretical Analysis

In this part, we first show that PCIL can be reduced to Apprenticeship
Learning (AL) (Abbeel & Ng, 2004). Based on this observation, we can further
show that PCIL is also perform distribution matching by minimizing a
divergence $\rho^{E}$ and $\rho^{\pi}$ that is equivalent to the total
variation divergence.

Recall that an AL problem takes the following form (Ho & Ermon, 2016):

$\underset{\pi}{\min}\max_{r\in\mathcal{R}}\displaystyle\mathop{\mathbb{E}}_{\begin{subarray}{c}{\mathbf{x}}=({\mathbf{s}},{\mathbf{a}})\sim\mathcal{D}\end{subarray}}\left[r({\mathbf{x}})\right]-\displaystyle\mathop{\mathbb{E}}_{\begin{subarray}{c}{\mathbf{x}}=({\mathbf{s}},{\mathbf{a}})\sim\rho_{\pi}\end{subarray}}\left[r({\mathbf{x}})\right],$
(3)

where $\mathcal{R}$ is a set of reward functions. AL plays a min-max game
between the policy $\pi$ and the reward function $r$. Intuitively, in the
inner loop we would like to find a cost function such that the expert data’s
cummulative return is higher than that of the agent’s, and their gap is
maximized. Meanwhile, the policy $\pi$ tries to minimize this gap.

Now, we reduce our objective to the AL formulation. For simplicity, we
consider the case that we only have one negative sample. We notice that
Equation 1 is then

$\mathcal{L}=\displaystyle\mathop{\mathbb{E}}\left[-\log\frac{\exp\Phi({\mathbf{x}}_{0})^{T}\Phi({\mathbf{x}}_{p})}{\left(\exp\Phi({\mathbf{x}}_{0})^{T}\Phi({\mathbf{x}}_{p})+\exp\Phi({\mathbf{x}}_{0})^{T}\Phi({{\mathbf{x}}}_{n}\right)}\right]$
(4)

As suggested by (Khosla et al., 2020), we can apply the Taylor expansion trick
to approximate this loss function with the following form:

$\mathcal{L}\approx\mathbb{E}~{}[\|\Phi({\mathbf{x}}_{0})-\Phi({\mathbf{x}}_{p})\|^{2}-\|\Phi({\mathbf{x}}_{0})-\Phi({{\mathbf{x}}}_{n})\|^{2}].$
(5)

Note that we drop the constant terms and the scaling constant since they do
not affect the optimization objective. Moreover, since $\Phi$ embeds the data
points to the sphere, we have
$\|\Phi({\mathbf{x}})\|^{2}=1,\forall{\mathbf{x}}$. As a result, we can
further expand each term above and have

$\mathcal{L}=\mathbb{E}\left[\Phi({\mathbf{x}}_{0})^{T}\Phi({\mathbf{x}}_{n})-\Phi({\mathbf{x}}_{0})^{T}\Phi({{\mathbf{x}}}_{p})\right].$
(6)

Since the variables in this equation are independent from each other, we are
minimizing

$\begin{split}\mathcal{L}=\mathbb{E}_{{\mathbf{x}}_{n}\sim\rho_{\pi}}[\mathbb{E}_{{\mathbf{x}}_{0}\sim\mathcal{D}}[\Phi({\mathbf{x}}_{0})]^{T}\Phi({\mathbf{x}}_{n})]-\\\
\mathbb{E}_{{\mathbf{x}}_{p}\sim\mathcal{D}}[\mathbb{E}_{{\mathbf{x}}_{0}\sim\mathcal{D}}[\Phi({\mathbf{x}}_{0})]^{T}\Phi({{\mathbf{x}}}_{p})].\end{split}$
(7)

Let the reward function
$r_{\theta}({\mathbf{x}})=\mathbb{E}_{{\mathbf{x}}_{0}\sim\mathcal{D}}[\Phi_{\theta}({\mathbf{x}}_{0})]^{T}\Phi_{\theta}({\mathbf{x}})$
as we defined in Equation 2, then minimizing the Equation 7 is exactly doing
the maximization of

$\displaystyle\mathbb{E}_{{\mathbf{x}}\sim\mathcal{D}}\left[r_{\theta}({\mathbf{x}})\right]-\displaystyle\mathbb{E}_{{\mathbf{x}}\sim\rho_{\pi}}\left[r_{\theta}({\mathbf{x}})\right],$
(8)

which is exactly the inner maximization loop of AL. Then optimizing the policy
with respect to this $r_{\theta}$ is exactly the outer loop. Hence, our
algorithm is reduced to AL.

Finally, we can show that we are minimizing the total variation divergence
between the expert and policy distribution as follows. First, let us define
the inner objective of Equation 3 as

$D_{cont}(\rho^{E}||\rho^{\pi})=\max_{r_{\theta}}[\mathbb{E}_{{\mathbf{x}}\sim\rho^{E}}r_{\theta}({\mathbf{x}})-\mathbb{E}_{{\mathbf{x}}\sim\rho^{\pi}}r_{\theta}({\mathbf{x}})].$

Here we assume that the expert dataset $\mathcal{D}$ is large enough, so
sampling from $\mathcal{D}$ can be regarded as sampling from $\rho^{E}$. Then
we can show that

Theorem 1 (Equivalence to TV divergence)

$0.25D_{TV}(\rho^{E}||\rho^{\pi})\leq D_{cont}(\rho^{E}||\rho^{\pi})\leq
2.0D_{TV}(\rho^{E}||\rho^{\pi}).$

Here,
$D_{TV}(\rho^{E}||\rho^{\pi})=\frac{1}{2}\int_{\mathcal{X}}|\rho^{E}(x)-\rho^{\pi}(x)|dx$
is the total variation divergence. Therefore, our objective is a proxy of
$D_{TV}(\rho^{E}||\rho^{\pi})$.

We leave the proof of the theorem to Appendix E.

##  4 Experiments

In this section, we empirically evaluate PCIL on an extensive set of tasks
from the DeepMind control suite (Tassa et al., 2018), a widely used benchmark
for continuous control. Our experiments are designed to answer the following
questions: (1) Can PCIL achieve expert performance, and how sample efficient
is PCIL compared to state-of-the-art imitation learning algorithms? (2) How
does the representation space of PCIL differ from that of the AIL methods? (3)
How does our method perform when we use different representation learning
methods and reward design?

###  4.1 Experimental Setup

#### Environments

We experiment with 10 MuJoCo (Todorov et al., 2012) tasks provided by DeepMind
Control Suite. The selected tasks cover various difficulty levels, ranging
from simple control problems, such as the single degree of freedom cart pole,
to complex high-dimensional tasks, such as the quadruped run. The episode
length for all tasks is 1000 steps, where a per-step ground truth environment
reward is in the unit interval $[0,1]$. For each task, we train an expert
policy using DrQ-v2 (Yarats et al., 2021) with the true environment reward
function and use it to collect 10 demonstrations. We refer readers to Appendix
A for the full task list and more details about the demonstrations.

#### Training Details

To update the encoder, we randomly sample 128 expert transitions and 128 agent
transitions from a replay buffer. For arbitrary expert transition, any other
expert transition is considered a positive sample, and all the agent
transitions constitute the set of negative samples. We update the encoder by
minimizing Equation 1 with respect to these samples. We use DrQ-v2 (Yarats et
al., 2021) as the underlying RL algorithm to train the agent with the cosine
similarity reward given in Equation 2. We use a budget of 2M environment steps
for all the experiments. Further implementation details can be found in
Appendix B.

#### Baselines

We compare PCIL to Behavioral Cloning (BC) and two major classes of imitation
learning algorithms:

  1. 1.

Adversarial IRL: We consider Discriminator-Actor-Critic (DAC) (Kostrikov et
al., 2018), a state-of-the-art AIL method that employs an unbiased AIL reward
function and performs off-policy training to reduce environmental
interactions.

  2. 2.

Trajectory-matching IRL: Primal Wasserstein Imitation Learning (PWIL) (Dadashi
et al., 2020) and Sinkhorn Imitation Learning (SIL) (Papagiannis & Li, 2020)
are two recently proposed trajectory-matching imitation learning methods. PWIL
computes the reward based on an upper bound of Wasserstein distance. SIL
computes the reward based on Sinkhorn distances (Cuturi, 2013).

To ensure a fair comparison, we implement all the baselines using the same RL
algorithm. The implementation details of these algorithms are in the Appendix.

Expert  BC  SIL  PWIL  DAC  PCIL (Ours)

Figure 3: Comparisons of algorithms on 6 selected tasks. See Appendix D for
more tasks. For every 20k environment steps, we perform 10-episode rollouts of
the policy without exploration noise and report average episode returns over
the 10 episodes. We plot the mean performance over 3 seeds together with the
shaded regions, which represent 95% confidence intervals.

###  4.2 Main Results

We show the performance curves of 6 tasks in Figure 3, which are averaged over
three random seeds. More results on DeepMind control tasks are provided in
Appendix D. We find that PCIL is able to outperform the existing methods on
all of these tasks. It achieves near-expert performance within our online
sample budget in all considered tasks except Hopper Hop. In terms of sample
efficiency, i.e., the number of environment interactions required to solve a
task, PCIL shows significant improvements over prior methods on five tasks:
Cheetah Run, Finger Spin, Hopper Hop, Hopper Stand, and Quadruped Run. For the
remaining tasks, PCIL achieves similar results compared with the state-of-the-
art adversarial imitation learning method DAC. In particular, we notice that
PCIL’s performance gain is larger on more difficult tasks (e.g., Cheetah Run,
Quadruped Run). On those easier tasks (e.g., Walker Stand, Walker Walk), the
baselines are also able to achieve strong results.

###  4.3 Analysis of Representation Space

Figure 4: t-SNE visualization results for DAC (top panels) and our PCIL method
(bottom panels). The blue color indicates the expert’s transition. The lighter
(yellow) color indicates agent transitions with higher real reward while the
darker (green) indicates lower real reward.

We visualize the representation space of PCIL and DAC using t-SNE (Van der
Maaten & Hinton, 2008) in Figure 4. For DAC, since there is no explicit
representation learning in the discriminator, we treat the last hidden layer
of its discriminator as the representation. We randomly sample 128 expert
transitions and 256 agent transitions for visualization. For a fair
comparison, all the transitions of PCIL and DAC are selected from the same
episode during the training process on the same task. We use color to indicate
the real environment reward of the agent’s transitions. The lighter (yellow)
color indicates a higher reward for agent transition while the darker (green)
indicates a lower reward agent transition. The blue color indicates the
expert’s transition.

We observe that in the representation space of PCIL, the expert transitions
are concentrated in a cluster. Moreover, the distance between each agent
transition to the cluster of expert transition is highly correlated with the
real reward of that agent transition. This fact suggests that the contrastive
objective of PCIL indeed induces a meaningful representation space here. On
the contrary, the representation space of DAC is much less structured. The
expert transitions are scattered throughout the representation space of DAC.
Moreover, we identify that in the DAC’s representation space, the agent’s real
ward does not correlate well with its distance to the expert transitions.
These facts show that our method indeed learns a better representation space.

###  4.4 Ablation Studies

As described in our method, our method has two components: a policy
contrastive representation for imitation (Section 3.2) and a similarity-based
imitation reward (Section 3.3). In this part, we carry out ablation studies to
analyze their effects. We first introduce our design choices as follows.

Table 1: Ablation studies on the policy contrastive representation and
similarity-based imitation reward. We report the average final returns on 4
selected tasks over 3 random seeds and standard deviations are given in error
bars.
Methods | Finger Spin | Walker Run | Hopper Stand | Hopper Hop
---|---|---|---|---
PCL + Sim-reward (Ours) | $\bf 964.2\pm 24.2$ | $\bf 778.4\pm 14.1$ | $\bf 771.2\pm 142.1$ | $\bf 43.2\pm 20.1$
TCN + Sim-reward | $\bf 975.2\pm 45.1$ | $180.5\pm 11.4$ | $2.0\pm 1.2$ | $6.8\pm 0.2$
PCL + GAIL-reward | $2.5\pm 2.4$ | $18.5\pm 6.5$ | $1.2\pm 1.8$ | $0.1\pm 0.1$
TCN + GAIL-reward | $35.4\pm 19.2$ | $19.8\pm 4.2$ | $1.5\pm 0.3$ | $17.3\pm 4.6$

#### PCIL Representation v.s. TCN Representation

We replace our proposed policy contrastive objective with other contrastive
learning methods. For this purpose, we adopt the popular self-supervised
representation learning method that leverages temporal information: Time-
Contrastive Networks (TCN) (Sermanet et al., 2018). In this case, the positive
samples are selected within a small window around the anchor sample, while the
negative samples are selected from distant time steps in the same rollout
trajectory. See Appendix C for implementation details.

#### Similarity-based Reward v.s. GAIL-like Reward

We also ablate the similarity-based reward in PCIL by replacing the
similarity-based imitation reward with a GAIL-like reward. Specifically, we
train a linear binary classifier on the policy contrastive embedding space to
distinguish expert or non-expert data. In this case, the embedding space is
still trained by PCIL contrastive loss, and the GAIL reward’s gradient is
detached from the embedding network. We use the same reward predictor as other
GAIL-style methods (Kostrikov et al., 2018), i.e. $\log(D(x))-\log(1-D(x))$.

#### Analysis

By comparing rows 1 and 2 in Table 1, we find that the approach with TCN
encoder does not work in three out of four environments. This is because the
optimization goal of TCN is not to distinguish between expert and non-expert
data. Thus the reward produced by comparing expert and non-expert data in the
learned representation space is not necessarily meaningful. Note that the case
in row 2 is no longer an adversarial IRL method. We also consider a case (row
4 in Table 1) where we use TCN and GAIL-like reward predictor, but the
performance of this method is poor. Moreover, we observe that in the absence
of the similarity-based imitation reward (compare rows 1 and 3 in Table 1),
our method does not work. This is because our representation space has metric-
space characteristics. As a result, we should use a distance-based metric to
compute the reward. In conclusion, the two components of our method are
necessary for achieving good performance.

##  5 Related Work

###  5.1 Imitation Learning

Imitation learning is a class of algorithms that enables a robot to acquire
behaviors from a demonstration dataset. There are two classes of imitation
learning algorithms: behavioral cloning (BC) and Inverse Reinforcement
Learning (IRL) (Ng et al., 2000). BC is a simple supervised learning algorithm
that directly fits the expert’s action. However, some work suggests that it
has some drawbacks: it suffers from covariate shift problem (Ross et al.,
2011), and it is hard to learn from a demonstration dataset without expert
actions (Torabi et al., 2018a). Instead, IRL (Abbeel & Ng, 2004) proposes to
recover the underlying policy by estimating the underlying reward function and
then maximizing the overall return with this reward. In particular, a recent
branch of IRL is the AIL, which proposes to match agents’ state-action
distribution with experts via adversarial training. GAIL (Ho & Ermon, 2016)
proposed a maximum entropy occupancy measure matching method which learns a
discriminator to bypass the need to recover the expert’s reward function.
Later, several works proposed an improved version of the GAIL methods
(Ghasemipour et al., 2020; Blondé et al., 2022; Baram et al., 2017; Kostrikov
et al., 2018; Fu et al., 2017). AIRL (Fu et al., 2017) replaced the Shannon-
Jensen divergence used in GAIL by KL divergence to measure similarity between
state-action pair distributions. (Baram et al., 2017) bridges the GAIL
framework to model-based reinforcement learning. DAC (Kostrikov et al., 2018)
improved the sample efficiency by leveraging a replay buffer without
importance sampling and dealing with the absorbing state problem. In contrast
to these works, we focus on the representation of AIL’s discriminator and
reformulating AIL in a contrastive embedding space.

###  5.2 Representation Learning For Policy Learning

In this work, we propose a representation-learning-based approach to improve
imitation learning. As imitation learning is a major class of policy learning
algorithms, we review works that use self-supervised learning to improve
policy learning in this part. Pioneer works (Mirowski et al., 2016; Jaderberg
et al., 2016; Shelhamer et al., 2016; Lample & Chaplot, 2017) explore using
auxiliary objectives (e.g., predict some property of the environment) as a
self-supervision signal. Recent works employ more general self-supervised
objectives (Oord et al., 2018). In particular, (Srinivas et al., 2020) are
based on contrastive representation learning. (Sermanet et al., 2018) learn
representation from multiview video using time contrastive learning. Some
other methods also explore the use of self-supervised representation
pretrained on environment data (Ha & Schmidhuber, 2018) or from real-world
images (Xiao et al., 2022; Parisi et al., 2022; Nair et al., 2022). In
imitation learning, (Mandi et al., 2022) proposes to use contrastive learning
for one-shot imitation learning in robotics. (Chen et al., 2021) investigate
the use of representation for imitation learning. However, their result
suggests that self-supervised representation learning only provides a small
improvement of imitation learning algorithms’ performance. Our method differs
from all these existing works by proposing to anchor on different policies and
learn a discriminative self-supervised representation for imitation learning.

##  6 Conclusion

In this paper, we suggested a new approach to improve adversarial imitation
learning algorithms: to learn a more meaningful, discriminative representation
space for imitation. To this end, we proposed a new algorithm called PCIL. We
conducted a theoretical analysis of our method and showed its connection to
apprenticeship learning. We also conducted experiments on the DeepMind Control
Suite and showed that PCIL could achieve state-of-the-art performance.
Moreover, we used an ablation study to highlight its difference from the
previous representation learning method. In the future, we will focus on
further improving our loss function design. For example, can we anchor on the
agent policies at different training stages? It will also be interesting to
extend the proposed representation learning method in the relaxed setting of
IL, like the scene where we can access both the reward and demonstration.

##  7 Reproducibility

We implement our algorithm according to parameters and details described in
Appendix B and Section 4.1. We will release our code and data.

##  8 Acknowledgement

This work is supported by the Ministry of Science and Technology of the
People´s Republic of China, the 2030 Innovation Megaprojects ”Program on New
Generation Artificial Intelligence” (Grant No. 2021AAA0150000). This work is
also supported by the National Key R&D Program of China (2022ZD0161700). We
also thanks Guanqi Zhan and Luke Melas-Kyriazi for proof reading.

## References

  * Abbeel & Ng (2004) Abbeel, P. and Ng, A. Y.  Apprenticeship learning via inverse reinforcement learning.  In _Proceedings of the twenty-first international conference on Machine learning_ , pp. 1, 2004.
  * Arjovsky et al. (2017) Arjovsky, M., Chintala, S., and Bottou, L.  Wasserstein generative adversarial networks.  In _International conference on machine learning_ , pp. 214–223. PMLR, 2017.
  * Baram et al. (2017) Baram, N., Anschel, O., Caspi, I., and Mannor, S.  End-to-end differentiable adversarial imitation learning.  In _International Conference on Machine Learning_ , pp. 390–399. PMLR, 2017.
  * Blondé et al. (2022) Blondé, L., Strasser, P., and Kalousis, A.  Lipschitzness is all you need to tame off-policy generative adversarial imitation learning.  _Machine Learning_ , 111(4):1431–1521, 2022.
  * Byrne (2009) Byrne, R. W.  Animal imitation.  _Current Biology_ , 19(3):111–114, 2009.
  * Chen et al. (2021) Chen, X., Toyer, S., Wild, C., Emmons, S., Fischer, I., Lee, K.-H., Alex, N., Wang, S. H., Luo, P., Russell, S., et al.  An empirical investigation of representation learning for imitation.  In _Neural Information Processing Systems Datasets and Benchmarks Track (NeurIPS)_ , 2021.
  * Cuturi (2013) Cuturi, M.  Sinkhorn distances: Lightspeed computation of optimal transport.  _Advances in neural information processing systems_ , 26, 2013.
  * Dadashi et al. (2020) Dadashi, R., Hussenot, L., Geist, M., and Pietquin, O.  Primal wasserstein imitation learning.  _arXiv preprint arXiv:2006.04678_ , 2020.
  * Fu et al. (2017) Fu, J., Luo, K., and Levine, S.  Learning robust rewards with adversarial inverse reinforcement learning.  _arXiv preprint arXiv:1710.11248_ , 2017.
  * Fujimoto et al. (2018) Fujimoto, S., Hoof, H., and Meger, D.  Addressing function approximation error in actor-critic methods.  In _International conference on machine learning_ , pp. 1587–1596. PMLR, 2018.
  * Ghasemipour et al. (2020) Ghasemipour, S. K. S., Zemel, R., and Gu, S.  A divergence minimization perspective on imitation learning methods.  In _Conference on Robot Learning_ , pp. 1259–1277. PMLR, 2020.
  * Gulrajani et al. (2017) Gulrajani, I., Ahmed, F., Arjovsky, M., Dumoulin, V., and Courville, A. C.  Improved training of wasserstein gans.  _Advances in neural information processing systems_ , 30, 2017.
  * Ha & Schmidhuber (2018) Ha, D. and Schmidhuber, J.  World models.  _arXiv preprint arXiv:1803.10122_ , 2018.
  * He et al. (2020) He, K., Fan, H., Wu, Y., Xie, S., and Girshick, R.  Momentum contrast for unsupervised visual representation learning.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 9729–9738, 2020.
  * Ho & Ermon (2016) Ho, J. and Ermon, S.  Generative adversarial imitation learning.  _Advances in neural information processing systems_ , 29, 2016.
  * Hussein et al. (2017) Hussein, A., Gaber, M. M., Elyan, E., and Jayne, C.  Imitation learning: A survey of learning methods.  _ACM Computing Surveys (CSUR)_ , 50(2):1–35, 2017\.
  * Jaderberg et al. (2016) Jaderberg, M., Mnih, V., Czarnecki, W. M., Schaul, T., Leibo, J. Z., Silver, D., and Kavukcuoglu, K.  Reinforcement learning with unsupervised auxiliary tasks.  _arXiv preprint arXiv:1611.05397_ , 2016.
  * Khosla et al. (2020) Khosla, P., Teterwak, P., Wang, C., Sarna, A., Tian, Y., Isola, P., Maschinot, A., Liu, C., and Krishnan, D.  Supervised contrastive learning.  _Advances in Neural Information Processing Systems_ , 33:18661–18673, 2020.
  * Kim & Park (2018) Kim, K.-E. and Park, H. S.  Imitation learning via kernel mean embedding.  In _Thirty-Second AAAI Conference on Artificial Intelligence_ , 2018\.
  * Kostrikov et al. (2018) Kostrikov, I., Agrawal, K. K., Dwibedi, D., Levine, S., and Tompson, J.  Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning.  _arXiv preprint arXiv:1809.02925_ , 2018.
  * Kostrikov et al. (2020) Kostrikov, I., Yarats, D., and Fergus, R.  Image augmentation is all you need: Regularizing deep reinforcement learning from pixels.  _arXiv preprint arXiv:2004.13649_ , 2020.
  * Lample & Chaplot (2017) Lample, G. and Chaplot, D. S.  Playing fps games with deep reinforcement learning.  In _Thirty-First AAAI Conference on Artificial Intelligence_ , 2017\.
  * Langley (2000) Langley, P.  Crafting papers on machine learning.  In Langley, P. (ed.), _Proceedings of the 17th International Conference on Machine Learning (ICML 2000)_ , pp. 1207–1216, Stanford, CA, 2000\. Morgan Kaufmann.
  * Lillicrap et al. (2015) Lillicrap, T. P., Hunt, J. J., Pritzel, A., Heess, N., Erez, T., Tassa, Y., Silver, D., and Wierstra, D.  Continuous control with deep reinforcement learning.  _arXiv preprint arXiv:1509.02971_ , 2015.
  * Mandi et al. (2022) Mandi, Z., Liu, F., Lee, K., and Abbeel, P.  Towards more generalizable one-shot visual imitation learning.  _2022 IEEE International Conference on Robotics and Automation (ICRA)_ , 2022.
  * Mirowski et al. (2016) Mirowski, P., Pascanu, R., Viola, F., Soyer, H., Ballard, A. J., Banino, A., Denil, M., Goroshin, R., Sifre, L., Kavukcuoglu, K., et al.  Learning to navigate in complex environments.  _arXiv preprint arXiv:1611.03673_ , 2016.
  * Nair et al. (2022) Nair, S., Rajeswaran, A., Kumar, V., Finn, C., and Gupta, A.  R3m: A universal visual representation for robot manipulation.  _arXiv preprint arXiv:2203.12601_ , 2022.
  * Ng et al. (2000) Ng, A. Y., Russell, S., et al.  Algorithms for inverse reinforcement learning.  In _Icml_ , volume 1, pp. 2, 2000.
  * Oord et al. (2018) Oord, A. v. d., Li, Y., and Vinyals, O.  Representation learning with contrastive predictive coding.  _arXiv preprint arXiv:1807.03748_ , 2018.
  * Orsini et al. (2021) Orsini, M., Raichuk, A., Hussenot, L., Vincent, D., Dadashi, R., Girgin, S., Geist, M., Bachem, O., Pietquin, O., and Andrychowicz, M.  What matters for adversarial imitation learning?  _Advances in Neural Information Processing Systems_ , 34:14656–14668, 2021.
  * Papagiannis & Li (2020) Papagiannis, G. and Li, Y.  Imitation learning with sinkhorn distances.  _arXiv preprint arXiv:2008.09167_ , 2020.
  * Parisi et al. (2022) Parisi, S., Rajeswaran, A., Purushwalkam, S., and Gupta, A.  The unsurprising effectiveness of pre-trained vision models for control.  _arXiv preprint arXiv:2203.03580_ , 2022.
  * Peng et al. (2018) Peng, X. B., Kanazawa, A., Toyer, S., Abbeel, P., and Levine, S.  Variational discriminator bottleneck: Improving imitation learning, inverse rl, and gans by constraining information flow.  _arXiv preprint arXiv:1810.00821_ , 2018.
  * Ross et al. (2011) Ross, S., Gordon, G., and Bagnell, D.  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. JMLR Workshop and Conference Proceedings, 2011.
  * Sermanet et al. (2018) Sermanet, P., Lynch, C., Chebotar, Y., Hsu, J., Jang, E., Schaal, S., Levine, S., and Brain, G.  Time-contrastive networks: Self-supervised learning from video.  In _2018 IEEE international conference on robotics and automation (ICRA)_ , pp. 1134–1141. IEEE, 2018.
  * Shelhamer et al. (2016) Shelhamer, E., Mahmoudieh, P., Argus, M., and Darrell, T.  Loss is its own reward: Self-supervision for reinforcement learning.  _arXiv preprint arXiv:1612.07307_ , 2016.
  * Silver et al. (2014) Silver, D., Lever, G., Heess, N., Degris, T., Wierstra, D., and Riedmiller, M.  Deterministic policy gradient algorithms.  In _International conference on machine learning_ , pp. 387–395. PMLR, 2014.
  * Srinivas et al. (2020) Srinivas, A., Laskin, M., and Abbeel, P.  Curl: Contrastive unsupervised representations for reinforcement learning.  _arXiv preprint arXiv:2004.04136_ , 2020.
  * Tassa et al. (2018) Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. L., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A., et al.  Deepmind control suite.  _arXiv preprint arXiv:1801.00690_ , 2018.
  * Todorov et al. (2012) Todorov, E., Erez, T., and Tassa, Y.  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) Torabi, F., Warnell, G., and Stone, P.  Behavioral cloning from observation.  _arXiv preprint arXiv:1805.01954_ , 2018a.
  * Torabi et al. (2018b) Torabi, F., Warnell, G., and Stone, P.  Generative adversarial imitation from observation.  _arXiv preprint arXiv:1807.06158_ , 2018b.
  * Van der Maaten & Hinton (2008) Van der Maaten, L. and Hinton, G.  Visualizing data using t-sne.  _Journal of machine learning research_ , 9(11), 2008.
  * Wang et al. (2017) Wang, Z., Merel, J. S., Reed, S. E., de Freitas, N., Wayne, G., and Heess, N.  Robust imitation of diverse behaviors.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Xiao et al. (2019) Xiao, H., Herman, M., Wagner, J., Ziesche, S., Etesami, J., and Linh, T. H.  Wasserstein adversarial imitation learning.  _arXiv preprint arXiv:1906.08113_ , 2019.
  * Xiao et al. (2022) Xiao, T., Radosavovic, I., Darrell, T., and Malik, J.  Masked visual pre-training for motor control.  _arXiv preprint arXiv:2203.06173_ , 2022.
  * Yarats et al. (2021) Yarats, D., Fergus, R., Lazaric, A., and Pinto, L.  Mastering visual continuous control: Improved data-augmented reinforcement learning.  _arXiv preprint arXiv:2107.09645_ , 2021.

##  Appendix A Environments

We use 10 continuous control tasks from the DeepMind control suite (Tassa et
al., 2018). The summary for each task is provided in Table 2.

Task | $\mathrm{dim}({\mathcal{S}})$ | $\mathrm{dim}({\mathcal{A}})$
---|---|---
Finger Spin | $6$ | $2$
Hopper Stand | $14$ | $4$
Pendulum Swingup | $2$ | $1$
Walker Stand | $18$ | $6$
Walker Walk | $18$ | $6$
Acrobot Swingup | $4$ | $1$
Cheetah Run | $18$ | $6$
Hopper Hop | $14$ | $4$
Quadruped Run | $56$ | $12$
Walker Run | $18$ | $6$
Table 2: A detailed description of each tasks used in our experiments.

#### Demonstrations

For each task, we train expert policies using DrQ-v2 (Yarats et al., 2021) on
the actual environment rewards. We run 3 seeds and pick the seed that achieves
the highest return. Then we use this expert policy to collect 10
demonstrations.

##  Appendix B Algorithm Details

###  B.1 Implementation

#### RL agent

We use DrQ-v2 as the underlying RL algorithm. DrQ-v2 is an off-policy actor-
critic algorithm for continuous control. The core of DrQ-v2 is Deep
Deterministic Policy Gradient (DDPG) (Lillicrap et al., 2015) augmented with
$n$-step returns. The critic is trained using clipped double Q-learning
(Fujimoto et al., 2018) to reduce the overestimation bias in the target value.
The deterministic actor is trained using deterministic policy gradients (DPG)
(Silver et al., 2014). We also follows the setting of actor’s and critic’s
neural network architectures in state-based DrQ-v2 (Yarats et al., 2021).

#### Contrastive encoder

The contrastive encoder is implemented as a 4 layer MLP with hidden size [256,
256, 256]. The output dimension is 64. Following the architecture in (Yarats
et al., 2021), the contrastive encoder, the critic and the actor share the
same encoder backbone. This shared encoder is trained with the gradient of the
critic alone, which is also following the suggestion of (Kostrikov et al.,
2020; Yarats et al., 2021). The input of this shared encoder is state $s$ of a
transition.

#### Reward predictor

Reward of the agent transition is computed according to Equation 2. Thus, we
randomly sample one expert transition from the expert replay buffer to compute
agent reward. Empirically, we find that using the mean embedding of the expert
data yields similar performance.

#### Gradient penalty

In order to make the algorithm more stable, we use the gradient penalty
technique (Gulrajani et al., 2017) widely used in Wasserstein-GANs (Arjovsky
et al., 2017). We make minor adjustments to accommodate our policy contrastive
loss. GAIL-like methods usually constrain the gradient norm of the
discriminator’s output with respect to its input. While for PCIL, the
contrastive encoder’s output needs one more step. Specifically, the output
embeddings are first used to calculate rewards following Equation 2, then we
compute and penalize the gradient norm of the rewards. We use $10$ as the
weighting for the gradient penalty.

###  B.2 Hyperparameters

Table 3 lists the hyperparameters that are used for all baseline methods and
our method. Expert data ratio in PCIL means the ratio between expert data and
batch size. A ratio of 0.5 means that half of the batch is expert data and the
other half is the agent data. The contrastive learning usually needs a
temperature scaling after computing the cos-similarity, before computing the
exponential. For simplicity, we ignored it in the main text. In the
experiment, we follow prior contrastive learning work (He et al., 2020) and
use a typical value of $0.07$ for the temperature.

Methods | Parameter | Value
---|---|---
all methods | Replay buffer size | 500k
| Agent update frequency | 2
| Optimizer | Adam
| Learning rate | 1e-4
| Critic soft-update rate | 0.01
| Random seed | 1,2,3
| RL batch size | 128
| Discriminator training batch size | 256
| Hidden dim | 256
PCIL | Expert data ratio | 0.5
| Contrastive temperature | 0.07
Table 3: The hyperparameters of baseline methods and our method.

##  Appendix C Ablation study implementation details

###  C.1 TCN representation

TCN encoder shares the same network architecture as the PCL encoder. During
the training process in TCN encoder, for a random sampled anchor, we use data
point adjacent to it as positive pair and another random sampled data point as
negative pair. We also set contrastive temperature to $0.07$ and batch size to
$256$, which is the same as PCL encoder.

###  C.2 GAIL-like reward

GAIL reward predictor can be seen as a simplified version of GAIL
discriminator which has only one linear classifier layer. The reward predictor
is trained independently to distinguish whether the input is from a expert
data or non-expert data with a binary classification loss.

##  Appendix D Additional Experimental Results

Figure 5 shows the performance of PCIL on the other 4 tasks from the DeepMind
Control suite. We notice that the performance of some relatively easy tasks
has saturated. All the baselines achieve expert performance on Pendulum
Swingup. On Walker Stand and Walker Walk, PCIL is competitive with DAC, which
already demonstrates impressive sample efficiency.

Expert  BC  SIL  PWIL  DAC  PCIL (Ours)

Figure 5: Comparisons of algorithms on the other 4 tasks.

##  Appendix E Proof

At the last of Section 3.4, we have shown that

$\min_{\pi}\max_{r_{\theta}}\mathbb{E}_{x\sim\mathcal{D}}r_{\theta}(x)-\mathbb{E}_{x\sim\rho^{\pi}}r_{\theta}(x),$

where
$r_{\theta}=(\mathbb{E}_{x^{\prime}\sim\mathcal{D}}\Phi_{\theta}(x^{\prime})dx^{\prime})^{T}\Phi_{\theta}(x)$,
and $\mathcal{D}$ is the expert dataset. We may as well regart $\mathcal{D}$
as $\rho^{E}$ when the dataset size is large.

Define the inner objective as

$D_{cont}(\rho^{E}||\rho^{\pi})=\max_{r_{\theta}}[\mathbb{E}_{x\sim\rho^{E}}r_{\theta}(x)-\mathbb{E}_{x\sim\rho^{\pi}}r_{\theta}(x)].$

In this part, we prove

Theorem 1

$0.25D_{TV}(\rho^{E}||\rho^{\pi})\leq D_{cont}(\rho^{E}||\rho^{\pi})\leq
2.0D_{TV}(\rho^{E}||\rho^{\pi}).$

Here,
$D_{TV}(\rho^{E}||\rho^{\pi})=\frac{1}{2}\int_{\mathcal{X}}|\rho^{E}(x)-\rho^{\pi}(x)|dx$
is the total variation divergence.

Proof of Theorem 1

Step 1 We first turn the inner optimization to a simple form. Let
$v_{\theta}=\mathbb{E}_{x^{\prime}\sim\mathcal{D}}\Phi_{\theta}(x^{\prime})$,
then we find a corresponding orthonormal matrix $M_{\theta}$ such that
$M_{\theta}v_{\theta}=[\|v_{\theta}\|,0,0,0,...,0]^{T}$. As a result,

$r_{\theta}=(\mathbb{E}_{x^{\prime}\sim\mathcal{D}}\Phi_{\theta}(x^{\prime}))^{T}\Phi_{\theta}(x)=(M_{\theta}\mathbb{E}_{x^{\prime}\sim\mathcal{D}}\Phi_{\theta}(x^{\prime}))^{T}M_{\theta}\Phi_{\theta}(x)=\|v_{\theta}\|[M_{\theta}\Phi_{\theta}(x)]_{1}.$

If we define
$g_{\theta}=[M_{\theta}\Phi_{\theta}(x)]_{1}:\mathcal{X}\to[-1,1]$, then the
inner objective can be written as:

$\mathsf{InnerOPT}$ Maximize (w.r.t. $\theta$, $\alpha$)

$\alpha\int_{\mathcal{X}}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx,$

subject to

$\alpha=\left|\int_{\mathcal{X}}g_{\theta}(x)\rho^{E}(x)dx\right|.$

Note that the constraint is derived from the following fact:

$\alpha=\|v_{\theta}\|=\|M_{\theta}v_{\theta}\|=\|\mathbb{E}_{x^{\prime}\sim\rho^{E}}M_{\theta}\Phi_{\theta}(x^{\prime})\|=\left|\int_{\mathcal{X}}g_{\theta}(x)\rho^{E}(x)dx\right|.$

Step 2 In this step, we show that $0.5D_{TV}(\rho^{E}||\rho^{\pi})\leq
D_{cont}(\rho^{E}||\rho^{\pi})$. We do this by constructing a $g_{\theta}$,
and evaluate the corresponding objective function. First, let
$S=\\{x\in\mathcal{X}:\rho^{E}(x)\geq\rho^{\pi}(x)\\}$. Then, we split the
discussion into two cases.

Case 1 If $\mu^{E}(S)\geq\mu^{E}(S^{c})$.

we define

$g_{\theta}(x)=\begin{cases}1&\text{$x\in S$},\\\ -\beta&\text{$x\in
S^{c}$}\end{cases}.$

where $\beta\in[0,\mu^{E}(S)]$ is a scalar. Then, we know that

$\alpha=|\mu^{E}(S)-\beta\mu^{E}(S^{c})|=|\mu^{E}(S)-\beta(1-\mu^{E}(S))|=|(1+\beta)\mu^{E}(S)-\beta|=(1+\beta)\mu^{E}(S)-\beta.$

Meanwhile, we know

$\int_{\mathcal{X}}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx=\int_{S}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx+\int_{S^{c}}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx$

and

$RHS=\int_{S}|\rho^{E}(x)-\rho^{\pi}(x)|dx+\beta\int_{S^{c}}|\rho^{E}(x)-\rho^{\pi}(x)|dx\geq\beta\int_{\mathcal{X}}|\rho^{E}(x)-\rho^{\pi}(x)|dx$

Therefore, we have

$\displaystyle D_{cont}(\rho^{E}||\rho^{\pi})$
$\displaystyle\geq\alpha\int_{\mathcal{X}}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx\geq((1+\beta)\mu^{E}(S)-\beta)\beta\int_{\mathcal{X}}|\rho^{E}(x)-\rho^{\pi}(x)|dx$
$\displaystyle=2((1+\beta)\mu^{E}(S)-\beta)\beta
D_{TV}(\rho^{E}||\rho^{\pi}).$

Since $\mu^{E}(S)\geq\mu^{E}(S^{c})$ and $\mu^{E}(S)+\mu^{E}(S^{c})=1$, we
have $\mu^{E}(S)\geq 0.5$. Then,

$RHS=2((1+\beta)\mu^{E}(S)-\beta)\beta
D_{TV}(\rho^{E}||\rho^{\pi})\geq(1-\beta)\beta D_{TV}(\rho^{E}||\rho^{\pi}).$

In particular, if we pick $\beta=0.5$, then the right hand side is
$0.25D_{TV}(\rho^{E}||\rho^{\pi})$, so in this case,

$D_{cont}(\rho^{E}||\rho^{\pi})\geq 0.25D_{TV}(\rho^{E}||\rho^{\pi}).$

Case 2 If $\mu^{E}(S)<\mu^{E}(S^{c})$. This can be done in a similar way.

We define

$g_{\theta}(x)=\begin{cases}\beta&\text{$x\in S$},\\\ -1&\text{$x\in
S^{c}$}\end{cases}.$

where $\beta\in[0,0.5]$ is a scalar. Then, we know that

$\alpha=|\beta\mu^{E}(S)-\mu^{E}(S^{c})|=|\beta\mu^{E}(S)-(1-\mu^{E}(S))|=|(1+\beta)\mu^{E}(S)-1|=1-(1+\beta)\mu^{E}(S).$

Meanwhile, we know

$\int_{\mathcal{X}}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx=\int_{S}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx+\int_{S^{c}}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx$

and

$RHS=\beta\int_{S}|\rho^{E}(x)-\rho^{\pi}(x)|dx+\int_{S^{c}}|\rho^{E}(x)-\rho^{\pi}(x)|dx\geq\beta\int_{\mathcal{X}}|\rho^{E}(x)-\rho^{\pi}(x)|dx$

Therefore, we have

$\displaystyle D_{cont}(\rho^{E}||\rho^{\pi})$
$\displaystyle\geq\alpha\int_{\mathcal{X}}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx\geq(1-(1+\beta)\mu^{E}(S))\beta\int_{\mathcal{X}}|\rho^{E}(x)-\rho^{\pi}(x)|dx$
$\displaystyle=2(1-(1+\beta)\mu^{E}(S))\beta D_{TV}(\rho^{E}||\rho^{\pi}).$

Since $\mu^{E}(S^{c})>\mu^{E}(S)$ and $\mu^{E}(S)+\mu^{E}(S^{c})=1$, we have
$\mu^{E}(S)<0.5$. Then,

$RHS=2(1-(1+\beta)\mu^{E}(S))\beta
D_{TV}(\rho^{E}||\rho^{\pi})\geq(1-\beta)\beta D_{TV}(\rho^{E}||\rho^{\pi}).$

In particular, if we pick $\beta=0.5$, then the right hand side is
$0.25D_{TV}(\rho^{E}||\rho^{\pi})$, so in this case,

$D_{cont}(\rho^{E}||\rho^{\pi})\geq 0.25D_{TV}(\rho^{E}||\rho^{\pi}).$

Putting Case 1 and 2 together, we can conclude that

$D_{cont}(\rho^{E}||\rho^{\pi})\geq 0.25D_{TV}(\rho^{E}||\rho^{\pi}).$

Step 3 Finally, we show that

$D_{cont}(\rho^{E}||\rho^{\pi})<2.0D_{TV}(\rho^{E}||\rho^{\pi}).$

This is actually quite straightforward. Noting that

$\alpha=\left|\int_{\mathcal{X}}g_{\theta}(x)\rho^{E}(x)dx\right|\leq\int_{\mathcal{X}}|g_{\theta}(x)\rho^{E}(x)|dx=\int_{\mathcal{X}}|g_{\theta}(x)|\rho^{E}(x)dx\leq\int_{\mathcal{X}}\rho^{E}(x)dx=1,$

we have

$\displaystyle\alpha\int_{\mathcal{X}}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx$
$\displaystyle\leq|\alpha||\int_{\mathcal{X}}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx|\leq|\int_{\mathcal{X}}g_{\theta}(x)[\rho^{E}(x)-\rho^{\pi}(x)]dx|$
$\displaystyle\leq\int_{\mathcal{X}}|g_{\theta}(x)||\rho^{E}(x)-\rho^{\pi}(x)|dx\leq\int_{\mathcal{X}}|\rho^{E}(x)-\rho^{\pi}(x)|dx=2D_{TV}(\rho^{E}||\rho^{\pi}).$

This then completes the proof.


# Statistical Dynamics of Wealth Inequality
in Stochastic Models of Growth

Jordan T. Kemp1 and Luís M. A. Bettencourt2,3 1Department of Physics,
University of Chicago, Chicago, Illinois 60637, USA  2Department of Ecology
and Evolution, University of Chicago, Chicago, Illinois 60637, USA  3Mansueto
Institute for Urban Innovation, University of Chicago, Chicago, Illinois
60637, USA

###### Abstract

Understanding the statistical dynamics of growth and inequality is a
fundamental challenge to ecology and society. Recent analyses of wealth and
income dynamics in contemporary societies show that economic inequality is
very dynamic and that individuals experience substantially different growth
rates over time. However, despite a fast growing body of evidence for the
importance of fluctuations, we still lack a general statistical theory for
understanding the dynamical effects of heterogeneneous growth across a
population. Here we derive the statistical dynamics of correlated growth rates
in heterogeneous populations. We show that correlations between growth rate
fluctuations at the individual level influence aggregate population growth,
while only driving inequality on short time scales. We also find that growth
rate fluctuations are a much stronger driver of long-term inequality than
earnings volatility. Our findings show that the dynamical effects of
statistical fluctuations in growth rates are critical for understanding the
emergence of inequality over time and motivate a greater focus on the
properties and endogenous origins of growth rates in stochastic environments.

††preprint: APS/123-QED

Exponential growth and broad inequality are general features of population
dynamics in biology and society [1]. In biology, the maintenance of diverse
populations is associated with greater biodiversity and with larger pools of
variability enabling faster processes of evolution by natural selection [2].
In human societies, rates of long-term (economic) growth are much higher than
in most ecosystems, often generating widening inequalities and leading to
familiar stresses of social justice and equity [3, 4].

Researchers have recently taken a special interest in the dynamics of economic
inequality [1], applying multidisciplinary approaches that include historical
data analyses [5] and searches for social, political and economic mechanisms
that generate, distribute, and rebalance wealth within populations [3, 6].
Research in economics has emphasized the importance of heterogeneity in
populations [4] as a source of widening wealth inequality. These
heterogeneities associated with a number of different population features such
as the divergence of incomes between capital and labor [3], between management
and workers within firms [7, 8], and between people with different educational
attainment [9]. These primarily empirical approaches point to the fundamental
importance of diverging income growth for different subpopulations and point
to the need for rebalancing mechanisms if inequality is to remain controlled.
However, we still do not have statistical theories applicable in general
circumstances and, as a consequence, we lack a clear picture of how to
concurrently manage growth and inequality.

To this end, research in physics has derived a number of results on the role
of stochasticity as a source of inequality, on inter-agent exchanges
(particularly in the mean-field limit), and proposed redistribution mechanisms
in the context of stochastic geometric growth models [10, 11, 12, 13, 14, 15].
The main focus of this work has been the design of wealth redistribution
schemes towards creating long time stationary limits [16, 17, 18] with
parametrically controlled levels of inequality and whether such stationary
solutions exist [19, 20].

While these models have been used to fit data on wealth distributions [18,
20], their specialized approaches have left open questions on the fundamental
statistical dynamics underpinning the generation and allocation of wealth. For
example, evidence suggests models following strict application of Gibrat’s law
(individual growth rates independent of scale) cannot characterize the rapid
emergence of inequality experienced in recent years [21, 11], motivating
explicit analysis of open-ended dynamical effects due to various sources of
fluctuations. These developments support multiplicative stochastic growth as a
starting point for modeling [22], but call for the consideration of
statistical effects due to fluctuations and correlations in model parameters,
especially growth rates.

To address these issues, we derive the time evolution of the resource (wealth)
distribution in a statistical population of heterogeneous agents experiencing
geometric random growth. We focus on the dynamical effects of growth rate
fluctuations and the effects of their correlations with resources across
agents and over time. We show that these two effects lead to two dynamical
time scales, which require different control measures so that inequality does
not explode in a population over time. Specifically, we show that natural
schemes to reduce inequality are subject to reversal over longer time scales
because of variability of growth rates in the population. We end by discussing
the mechanisms that may simultaneously lead to sustained exponential growth
and control of long term sources of inequality in terms of agent-based
learning in correlated stochastic environments.

The fundamental model for the dynamics of resources in populations relies on
stochastic exponential growth [21, 22]. In its simplest form, known as
geometric Brownian motion, agents generate wealth by (re)investing incomes net
of costs. Crucially these models, unlike additive stochastic growth, generate
population lognormal and power-law statistics, which characterize observed
wealth and income statistics [23, 24]

Specifically, the dynamics of stochastic multiplicative growth start by
tracking the resources (wealth), $r(t)$, of a specific agent at time $t$. This
quantity changes in time via the difference between an income, $y(t)$, minus
costs, $c(t)$. The difference, $y-c$ (net income), is then defined to be
proportional to resources, expressed as a (potentially $r$ dependent)
stochastic growth rate, $\eta(t)=(y(t)-c(t))/r(t)$, for $r>0$. It follows that
the time evolution of resources obeys a simple (but multiplicative) stochastic
differential equation $dr(t)=\eta rdt+\sigma rdW_{t}$, where the volatility,
$\sigma$, is the standard deviation of the average growth rate, $\eta$. When
these parameters are independent of $r$ and $t$, this stochastic equation can
be integrated via It$\bar{\textrm{o}}$ calculus [22] to give

$\ln\frac{r(t)}{r(0)}=\bigg{(}{\eta}-\frac{\sigma^{2}}{2}\bigg{)}t+\sigma
W(t),$ (1)

where $\gamma={\eta}-\sigma^{2}/2$ is the effective growth rate and $r_{0}$
are the agent’s resources at $t=0$, $r_{0}=r(0)$. The decrease in the mean
growth rate ${\eta}$ due to finite volatility is an important feature of
multiplicative growth. The quantity $\sigma W(t)$ is a Wiener process with
magnitude proportional to the volatility and units of $t^{-1/2}$.

Under these circumstances the population dynamics follow from the single agent
time evolution, but correlations between $\eta$ and $r_{0}$ introduce new
considerations for the average growth rate. When growth rates are identical
and statistically independent in the population, the average growth rate is
the average over the growth rates of all individuals. However, when initial
log wealth $\ln r$ and growth rates $\eta$ are correlated, the growth of the
average becomes (Appendix A) $(\gamma
r)_{N}=\big{(}1+\textrm{covar}_{N}(\frac{\gamma_{i}}{\bar{\gamma}},\frac{r_{i}}{\bar{r}})\big{)}\bar{\gamma}\bar{r}$
for population average resources $\bar{r}$, and covariance in agent parameters
over the population $\textrm{covar}_{N}$.

We denote the term,
$\gamma^{\prime}=\bar{\gamma}+\textrm{covar}_{N}\big{(}\gamma_{i},\frac{r_{i}}{\bar{r}}\big{)}$,
the population effective growth rate. This quantity can be calculated
analytically for normally distributed quantities $\gamma$ and $\ln
r/{\bar{r}}$. The result involves the variances of both quantities as well as
the Pearson correlation coefficient between them, as
$\textrm{covar}_{N}(\gamma_{i},\frac{r_{i}}{\bar{r}})=\rho\sigma_{G}\sigma_{r}e^{\sigma^{2}_{r}/2\bar{r}}$.
The effective growth rate becomes

$\gamma^{\prime}=\eta-\frac{\sigma^{2}}{2}+\rho\sigma_{G}\sigma_{r}e^{\sigma^{2}_{r}/2\bar{r}},$
(2)

where $y\equiv\ln r/\bar{r}$, and the lognormal distributed resources
introduce an exponential term. If instead resources were normally distributed,
this term becomes unity. This expression shows that a positive covariance
between initial log resources and growth rates results in a higher population
effective growth rate and vice versa. Of course, associating higher-growth
opportunities with wealthier individuals will exacerbate inequalities in the
population over time. We denote populations with this configuration as
regressive and with negative correlation as progressive. We visualize the
tradeoff between short-term change in inequality and growth through numerical
simulations in Fig. 1.

Figure 1:  Monte Carlo simulations of growth dynamics with correlated growth
rates and initial resources. Left: Aggregate resource trajectories, of
populations with $\rho$ ranging from $-1$ (blue) to $1$ (red). Inset: The
population averaged effective growth rate (blue) scales linearly in $\rho$,
following Eq. 2. The temporal averaged growth rate (gray) scales similarly via
$\gamma=\gamma^{\prime}+\sigma^{2}$ [10]. Right: Distributions of log wealth
at different time steps for a regressive, $\rho=.75$ (bottom), and
progressive, $\rho=-.75$ (top) population. Black curves represent normal
distribution fits. The regressive assignments broaden more quickly, indicating
that the populations are becoming more rapidly unequal. Inset: Variance of the
progressive (blue) and regressive (red) populations over many time steps. Data
are simulated from populations of $10^{5}$ agents with Gaussian distributed
initial resources of mean $\bar{r}=10$, mean growth rate $\eta=.04$, and
standard deviations $\sigma_{\eta}=.015,\sigma_{r}=.682$.

We see that these simple considerations present an apparent paradox for any
attempt to simultaneously maximize total wealth (social welfare) and reduce
inequality. Decreasing $\rho$, thereby making the society more progressive
results in a social opportunity cost in terms of a decline in average growth.
This effect may even lead to decline (negative exponential growth) in
societies starting out with low average growth and high volatility. One way
out of this dilemma is for a progressive policy assignment to create (in ways
to be specified) higher average growth rates than the regressive case.

To analyze this possibility, Eq. 2 introduces the threshold correlation below
which aggregate effective growth rate becomes negative. Computed under the
condition $\gamma^{\prime}=0$, it is
$\rho_{c}=r\big{(}\frac{\sigma^{2}}{2}-\eta\big{)}/2\sigma_{r}\sigma_{G}$.
Similarly, the critical volatility marking the crossover from positive to
negative average growth, denoted $\sigma_{c}$, can be determined by
rearranging the expression for $\sigma$. This relationship is plotted in Fig.
2. We can directly compare the efficacy of progressive and regressive
assignments by computing the ratio of growth rates between a population with
growth rate $\gamma$, and its hypothetical progressive counterpart,
$\gamma_{p}$ with distribution parameters with superscript ${(p)}$. In the
simplest case where we assume identical initial population conditions up to
the assignment of growth rates (sign of $\rho$), Eq. 2 produces the condition
where average growth is preserved in the progressive society

$\gamma_{p}=\gamma+\rho\sigma_{r}\big{(}\sigma_{G}+{\sigma_{G}}^{(p)}\big{)}e^{\sigma^{2}_{r}/2\bar{r}}.$
(3)

Thus, a progressive arrangement must have a larger average growth rate in
order to achieve a population effective growth rate equal to its regressive
counterpart. This difference is made worse in regimes with stronger starting
inequality and growth rate fluctuations. Fig. 2 demonstrates this equation’s
agreement with population Monte-Carlo simulations in which $r_{0}$ is normally
distributed.

Figure 2:  Parameter spaces for progressive/regressive growth rate ratios
across variances (top) and growth rate regimes across correlation and
volatility values (bottom) Upper Left: Contour plot of $\gamma_{p}/\gamma$ for
correlations and growth rate variances with condition
$\sigma_{G}=\sigma_{G}^{(p)}$. Upper Right: A Monte-Carlo simulation comparing
the correlated average growth rates of populations with identical parameters
up to a change in sign of $\rho$ agrees with theory. To stay competitive, the
progressive society requires a larger growth rate as either parameter
increases. Lower: The effective growth rate for
$\eta=.04,\sigma_{r}=1,\sigma_{G}=.025$ transitions from positive to negative
along the correlation axis at high volatility ($\sigma\approx.275$).

So far, we have shown that the population average growth rate collects a
correction from covariances between resources and growth rate variations. This
introduces a trade-off between managing inequality in the short run and
maximizing overall growth which can only be resolved, at each time interval,
if progressive assignments also result in higher average growth rates.

But what happens over the longer term, as growth rate fluctuations persist? To
answer this question, we consider the solution to the Fokker-Planck equation
(FPE) for the time dependent probability of resources $r$ for each agent. The
solution with initial conditions $r(0)=r_{0}$ (a delta function at the
individual level) is a lognormal distribution [25] with time dependent mean
and variance. This can be re-written as a Gaussian distribution of $\ln r$ as

$P(\ln r,t|\gamma,\ln r_{0})=\frac{d\ln
r}{\sqrt{2\pi\sigma^{2}t}}\exp\bigg{[}-\frac{(\ln r-\ln r_{0}-\gamma
t)^{2}}{2\sigma^{2}t}\bigg{]},$ (4)

where at early times, the expected resource value is given by the linear
equation $\langle\ln r(t)\rangle=\ln r_{0}+\gamma t$, and asymptotically
approaches $\gamma t$. The variance increases linearly in time
$\langle\delta\ln r\rangle^{2}=\sigma^{2}t$, which shows that individual
wealth also becomes more uncertain as a function of time. The critical
volatility at which growth ceases is, as before, given by
$\sigma_{c}=\sqrt{2{\eta}}$ . We consider Eqn. 4 the homogeneous population
solution.

Changing variables and working explicitly with time average growth rates
yields further insight into the dynamics. Dividing each factor of the
exponential by t, and performing the variable transformation
$\frac{1}{t}\ln\frac{r}{r_{0}}\rightarrow\phi,td\phi=d\ln r$ yields the
distribution of growth rates

$P(\phi,t|\gamma,\sigma)=\frac{d\phi}{\sqrt{2\pi\sigma^{2}/t}}\exp\bigg{[}-\frac{(\phi-\gamma)^{2}}{2\sigma^{2}/t}\bigg{]}.$
(5)

This shows explicitly that the statistics of the time averaged growth rates
are simpler than those of resources. The average growth rate becomes
stationary at long times as the mean remains constant, while growth rate
fluctuations decrease in time as $1/t$, converging in distribution to a delta
function on a time scale $t>>\sigma^{2}$.

This suggests a simple picture emerging at the population level; however
varying the growth rates among agents in the population will introduce
complications. The most direct route to assess these effects follows by
positing Gaussian distributions on growth rates, from the asymptotic behavior
of growth rate distributions [26, 27, 22], and initial log resources across
the population, from the solution to the FPE.

To obtain the dynamical solution for the distribution of resources in the
population, we then convolve Eq. 4 with a bivariate static Gaussian
distribution of mean initial log resources $\mu_{0}$, variance
$\sigma_{r}^{2}$, and mean growth rate $G$, variance $\sigma_{G}^{2}$ via
(Appendix B)

$P(\ln r,t)=\frac{d\ln r}{\sqrt{2\pi\Sigma^{2}}}\exp\bigg{[}-\frac{(\ln
r-\mu_{0}-Gt)^{2}}{2\Sigma^{2}}\bigg{]},$ (6)

where the variance in the population, $\Sigma^{2}$, is given by

$\Sigma^{2}=\sigma_{r}^{2}+(\sigma^{2}+2\rho\sigma_{r}\sigma_{G})t+\sigma_{G}^{2}t^{2}.$
(7)

We consider Eq. 6 the heterogeneous population solution. The quadratic
expression in time for variance is characterized by two different timescales,
$t_{c1}=\sigma_{r}^{2}/(\sigma^{2}+2\rho\sigma_{r}\sigma_{G})$ and $\
t_{c2}=(\sigma^{2}+2\rho\sigma_{r}\sigma_{G})/\sigma_{G}^{2}$. The natural
timescale in this analysis is years [3], such that early times are on the
order of a few years, intermediate over a few decades, and long timescales are
several decades to centuries. Effects over long timescales can thus be thought
of as generational, although we do not incorporate life cycle dynamics in this
analysis. Annual growth rates are typically on the order of a few percent a
year ($\gamma\simeq 10^{-2}t^{-1}$), whereas log wealth varies in the
population on a typical scale of $\delta\ln r\simeq 10^{1}$, so the magnitude
of $\sigma_{G}$ is naturally some orders smaller than $\sigma_{r}$. This
results in $t_{c1}<t_{c2}$, producing three distinct dynamical regimes for the
population variance and hence inequality dynamics. For early times $t<t_{c1}$,
the variance of resources in the population is given approximately by its
initial condition $\Sigma^{2}\simeq\sigma_{r}^{2}$. For intermediate times
$t_{c1}<t<t_{c2}$, $\Sigma^{2}\simeq(\sigma^{2}+2\rho\sigma_{r}\sigma_{G})t$.
The intermediate regime introduces the explicit decrease in inequality from a
progressive assignment, so long as $\sigma^{2}<2\rho\sigma_{r}\sigma_{G}$.
However, this benefit is short-lived, as for later times $t>t_{c2}$,
$\Sigma^{2}\simeq\sigma_{G}^{2}t^{2}$, and variance invariably explodes due to
fluctuations in growth rates. Note that the long time regime does not require
biases between resources and growth rates; it persists under the weaker
conditions of a finite variance in growth rates and appears at earlier times
for larger growth rate variances.

Finally, we observe that the asymptotic population growth rate distribution
simplifies to a Gaussian

$\lim_{t\rightarrow\infty}P(\phi,t)=\frac{d\phi}{\sqrt{2\pi\sigma_{G}^{2}}}\exp\bigg{[}-\frac{(\phi-G)^{2}}{2\sigma^{2}_{G}}\bigg{]},$
(8)

in contrast to the case of a single agent, reflecting the Gaussian
distribution of growth rates in the population. Over the long term, population
dynamics are thus dominated by fluctuations in growth rates, which lead to
widening inequality the larger the growth rate variance.

This effect introduces a general mechanism that can account for increases in
inequality that occur over intermediate and long timescales [21], and
incorporate findings that empirical growth rates are heterogeneous beyond
fluctuation noise, and correlate with wealth over time [28, 29, 30]. It also
asks that we shift our focus from the dynamics of wealth itself to the
statistics of its drivers in terms of variations of growth rates.

We now analyze the detailed dynamics of inequality using two standard metrics.
First, the Gini coefficient, $G_{\rm ini}$, is the most common metric
measuring inequality within a population. It ranges from $G_{\rm ini}=0$ for
equally distributed wealth, and trends towards 1 as a society becomes
maximally unequal. For lognormally distributed wealth in the infinite
population limit, $G_{ini}=\textrm{Erf}\big{[}\Sigma/2\big{]}$. Second, we
examine the social cost of high growth on inequality, a topic of interest and
debate in the social sciences [31, 32]. Using the coefficient of variation
(CV), denoted $c_{v}$, we compare how quickly a distribution’s mean increases
relative to its standard deviation [33]. The quantity
$c_{v}=\Sigma/(\mu_{0}+Gt)$ where, from Eq. 2,
$G=\eta-\frac{\sigma^{2}}{2}+\rho\sigma_{G}\sigma_{r}e^{\sigma_{r}^{2}/2\bar{r}}$.
A more equitable economy would decrease $c_{v}$; growing without increasing
inequality. In the egalitarian limit, $c_{v}=0$ as every agent has the same
resources, and variance is zero. Conversely, an exploding $c_{v}$ in the
positive or negative direction indicates increasing inequality, marked by
positive or negative aggregate growth respectively. From Eq. 7, without
dynamical redistribution measures, societies will invariably become maximally
unequal as multiplicative dynamics are dominated by fluctuations in earnings
[1].

|
---|---
|

Figure 3:  Trajectories of inequality metrics in various parameter regimes;
$\sigma=.15,\sigma_{r}=.341,\mu_{0}=1,\eta=.06,\sigma_{G}=.025$. See legend
for line type mapping. Left: Plots for low volatility, $\sigma=.1$,
populations with different correlation coefficient values. At early times,
$G_{ini}$ increases more quickly for $\rho<0$, and slightly decreases in a
society with low volatility. Variances in growth rates in both cases drive the
$G_{ini}$ to 1 more rapidly than in the homogeneous case. For $\rho>0$,
$c_{v}$ jumps higher than in an uncorrelated society, but rapidly approaches a
lower asymptotic value due to a higher average growth rate with comparable
inequality. Conversely, $c_{v}$ initially decreases for $\rho<0$, but reaches
a higher asymptotic value as long term inequality remains comparable at the
cost of a lower average population growth rate. Right: Regressive populations
($\rho=.75$) at different values of volatility. A higher volatility causes
more rapid increase in $G_{ini}$ in both population configurations, and
volatility positively influences the peak and asymptotic values of $c_{v}$.
Progressive assignments hurt long term growth while marginally affecting long-
term inequality dynamics, while heterogeneities play a dominant role in
accelerating the emergence of inequality.

Homogeneous

Regressive Progressive

Figure 4: Contour plots for $\partial_{t}c_{v}$ in identical populations to
Fig. 3, with $\eta=.035$. At low volatilities, homogeneous $c_{v}$ initially
increases, then decreases with magnitude increasing in $\sigma$. For $\rho>0$,
this behavior persists following initial times of high increase in $c_{v}$.
The initial decrease observed for $\rho<0$ only occurs at low volatility, and
$c_{v}$ does not decrease in later times, as demonstrated in Fig. 3. Past the
homogeneous critical volatility, given by $\sigma_{c}\approx.265$, $c_{v}$
blows up as the negative growth rate drives average log resources through zero
to a negative value, and is seen by the bright yellow bands. More negative
$\gamma$, caused by high $\sigma$ or negative $\rho$, reduces the time needed
to pass through $\ln r=0$. In this regime of instability, resources continue
to decrease as inequality grows. Introducing heterogeneity shifts this
crossover point based on the value of the covariance term in the growth rate.

We explore the dynamics of homogeneous and heterogeneous populations by
comparing time-series trajectories of $c_{v}$ and $G_{ini}$ across several
values of correlation coefficient and volatility. Fig.3 demonstrates the
impact of heterogeneity on inequality (Gini), and its relative effects on
growth (CV). In homogeneous populations, volatility spurs inequality from
uniform initial conditions causing $c_{v}$ to initially increase, then peak
and decrease as resources increase in $t^{1/2}$ faster than population
standard deviation. The quantity $G_{ini}$ asymptotically approaches 1 in all
parameter cases. Heterogeneity causes $c_{v}$ to increase across all cases
with asymptotically constant behavior as the time leading terms of the
numerator and denominator of $c_{v}$ cancel. Heterogeneity causes $G_{ini}$ to
generally increase more rapidly. The differences in progressive and regressive
$G_{ini}$ are negligible after intermediate times, suggesting that the initial
configuration has little effect on its long-term macroscopic dynamics.
Volatility strongly determines the value of $c_{v}$ for homogeneous and
heterogeneous populations, while only significantly affecting the behavior of
$G_{ini}$ in the homogeneous case.

The dynamics of $G_{ini}$ are thus straightforward, but the dependence of
$c_{v}$ on both growth rates and dynamical variance manifests interesting
dynamics. To probe the dynamics of $c_{v}$, we produce quasi-phase diagrams
displaying $\partial_{t}c_{v}$ for a continuous spectrum of volatilities over
time. Fig. 4 demonstrates these dynamical regimes for the homogeneous and
heterogeneous progressive and regressive cases. It shows the magnitude of the
initial increase in $c_{v}$ scales with $\sigma$, as does the magnitude of the
eventual decrease. The eventual decrease is weaker in regressive populations,
and is nonexistent for any volatility value of the progressive population due
to the lower progressive growth rate. Furthermore, they reveal a divergent
regime past the critical volatility, where resources decrease while inequality
continues to increase.

These results show that growth rate fluctuations pose a general challenge to
any proposed mechanism to address widening inequality that must be addressed
independently of volatility in single agent trajectories. Specifically,
dynamical balancing of growth rates is required to control long term
inequality dynamics. This raises the issue that any method for reducing growth
rate fluctuations while preserving growth must include knowledge of the
origins of wealth growth rate statistics. Extending the stochastic growth
models commonly used to include environmental opportunities and agent-based
learning, in ways analogous to gambling [34] and portfolio theory [35],
naturally produces situations where individual growth rates become distinct
and history dependent but also where average growth rates can be maximized and
variances minimized over time [34, 22, 36]. In such dynamical picture, initial
correlations will also take on dynamical properties in ways that we stated
anticipating here. In recognizing recent research on the effects of
heterogeneous and dynamical growth on distributions of wealth [25, 37, 38, 39,
40], we seek a theoretical framework for wealth dynamics that both complements
these phenomenological approaches and incorporates strategic agent behavior in
statistical environments. These topics will be presented in future work.

In summary, we explored the effects of fluctuations in growth rates and wealth
statistics in the temporal development of inequality in a population of
heterogeneous agents. We set up the general problem in the context of
multiplicative random growth and derived closed form expressions for how the
population inequality changes over time, thus identifying the most important
parameters at early and later times. We found that population variances in
growth rates, their correlations with resources, and individual temporal
volatility to be the primary drivers of inequality in general situations. The
sustained presence of variance in growth rates in a population produces a
dominant effect on long-term inequality. While redistribution methods, such as
those explored by mean-field or other more complex network exchange models
[16, 17, 11] reduce the dynamical impacts of volatility, we have shown that
directly addressing growth rate variance in a population is a more fundamental
and general requirement for arresting or reversing increases in wealth
inequality.

Acknowledgements: We thank Arvind Murugan, Thomas Bourany, and Akhil Ghanta
for discussions and comments on the manuscript. This work is supported by the
Mansueto Institute for Urban Innovation and the Department of Physics at the
University of Chicago and by a National Science Foundation Graduate Research
Fellowship (Grant No. DGE 1746045 to JTK).

## References

  * [1] M. Scheffer, B. Van Bavel, I. Leemput, and E. Nes, “Inequality in nature and society,” Proceedings of the National Academy of Sciences, vol. 114, p. 201706412, 11 2017.
  * [2] S. Frank, “Natural election. iv. the price equation,” J. Evo. Biol., vol. 25, no. 6, pp. 97–111, 2012.
  * [3] T. Piketty and E. Saez, “Inequality in the long run,” Science, vol. 343, no. 6186, pp. 838–843, 2014.
  * [4] F. Cunha and J. J. Heckman, “The evolution of inequality, heterogeneity and uncertainty in labor earnings in the u.s. economy,” October 2007.
  * [5] L. A. Keister, Wealth in America: Trends in wealth inequality.  Cambridge University Press, 2000.
  * [6] B. Düring, D. Matthes, and G. Toscani, “Kinetic equations modelling wealth redistribution: a comparison of approaches,” Physical Review E, vol. 78, no. 5, p. 056103, 2008.
  * [7] J. Bakija, A. Cole, and B. Heim, “Jobs and income growth of top earners and the causes of changing income inequality: Evidence from u.s. tax return data,” Mar 2012.
  * [8] A. B. Atkinson, T. Piketty, and E. Saez, “Top incomes in the long run of history,” Journal of Economic Literature, vol. 49, pp. 3–71, March 2011\.
  * [9] C. Goldin and L. F. Katz, “The race between education and technology: The evolution of u.s. educational wage differentials, 1890 to 2005,” March 2007.
  * [10] J.-P. Bouchaud, “Wealth condensation in a simple model of economy,” Physica A: Statistical Mechanics and its Applications, vol. 282, pp. 536–545, 2000.
  * [11] D. Garlaschelli and M. I. Loffredo, “Effects of network topology on wealth distributions,” Journal of Physics A, vol. 41, p. 224018, 2008.
  * [12] P. Degond, J. G. Liu, and C. Ringhofer, “Evolution of the distribution of wealth in an economic environment driven by local nash equilibria,” Journal of Statistical Physics, vol. 154, no. 3, pp. 751–780, 2014.
  * [13] M. Medo, “Breakdown of the mean-field approximation in a wealth distribution model,” arXiv.org, Quantitative Finance Papers, vol. 2009, 09 2008.
  * [14] A. Chakraborti and B. Chakrabarti, “Statistical mechanics of money: How saving propensity affects its distribution,” The European Physical Journal B - Condensed Matter and Complex Systems, vol. 17, 09 2000.
  * [15] M. Patriarca, E. Heinsalu, and A. Chakraborti, “Basic kinetic wealth-exchange models: common features and open problems,” The European Physical Journal B, vol. 73, pp. 145–153, 2006.
  * [16] J.-P. Bouchaud, “On growth-optimal tax rates and the issue of wealth inequalities,” Journal of Statistical Mechanics: Theory and Experiment, vol. 2015, 08 2015.
  * [17] Y. Berman, E. Ben-Jacob, and Y. Shapira, “The dynamics of wealth inequality and the effect of income distribution,” PloS one, vol. 11, p. e0154196, 04 2016.
  * [18] J. Li, B. Boghosian, and C. Li, “The affine wealth model: An agent-based model of asset exchange that allows for negative-wealth agents and its empirical validation,” Physica A: Statistical Mechanics and its Applications, vol. 516, 10 2018.
  * [19] O. Peters, “The ergodicity problem in economics,” Nature Physics, vol. 15, pp. 1216–1221, 12 2019.
  * [20] Y. Berman, O. Peters, and A. Adamou, “Wealth inequality and the ergodic hypothesis: Evidence from the united states,” Claremont McKenna College Robert Day School of Economics & Finance Research Paper Series, 2020.
  * [21] X. Gabaix, J.-M. Lasry, P.-L. Lions, and B. Moll, “The dynamics of inequality,” Econometrica, vol. 84, pp. 2071–2111, 11 2016.
  * [22] L. Bettencourt, “Urban growth and the emergent statistics of cities,” Science Advances, vol. 6, p. eaat8812, 08 2020.
  * [23] A. I. Saichev, Y. Malevergne, and D. Sornette, “Theory of zipf’s law and beyond,” Springer Science & Business Media, vol. 632, 2009.
  * [24] A. Clauset, C. R. Shalizi, and M. E. Newman, “Power-law distributions in empirical data,” SIAM review, vol. 51, no. 4, pp. 661–703, 2009.
  * [25] V. Stojkoski, T. Sandev, L. Basnarkov, L. Kocarev, and R. Metzler, “Generalised geometric brownian motion: Theory and applications to option pricing,” 10 2020.
  * [26] M. Takayasu, H. Watanabe, and H. Takayasu, “Generalised central limit theorems for growth rate distribution of complex systems,” Journal of Statistical Physics, vol. 155, pp. 47–71, 2013.
  * [27] M. Wyart and J.-P. Bouchaud, “Statistical models for company growth,” Physica A-statistical Mechanics and Its Applications, vol. 326, pp. 241–255, 2003.
  * [28] F. Guvenen, “Learning your earning: Are labor income shocks really very persistent?,” American Economic Review, vol. 97 (3), pp. 687–712, 2007\.
  * [29] A. Fagereng, L. Guiso, D. Malacrino, and L. Pistaferri, “Heterogeneity and persistence in returns to wealth,” Econometrica, vol. 88, no. 1, pp. 115–170, 2020.
  * [30] J. Hubmer, P. Krusell, and A. A. Smith Jr, “Sources of us wealth inequality: Past, present, and future,” Nber macroeconomics annual, vol. 35, no. 1, pp. 391–455, 2021.
  * [31] R. E. Lucas, “The industrial revolution: past and future,” 2003 Annual Report of the Federal Reserve Bank of Minneapolis, pp. 5,20, 2004.
  * [32] J. C. Córdoba and G. Verdier, “Lucas vs. lucas: On inequality and growth,” Economic Growth, 2007.
  * [33] L. G. Bellu and P. Liberati, “Policy impacts on inequality – simple inequality measures,” EASYPol, Analytical tools. Policy Support Service, Policy Assistance Division, FAO, 2006.
  * [34] J. L. Kelly, “A new interpretation of information rate,” IRE Trans. Inf. Theory, vol. 2, pp. 185–189, 1956.
  * [35] M. Hilbert, “The more you know, the more you can grow: an information theoretic approach to growth in the information age,” Entropy, vol. 19, no. 2, p. 82, 2017.
  * [36] T. M. Cover and J. A. Thomas, Elements of Information Theory 2nd Edition (Wiley Series in Telecommunications and Signal Processing).  Wiley-Interscience, July 2006.
  * [37] D. B. Nelson, “Conditional heteroskedasticity in asset returns: A new approach,” Econometrica: Journal of the Econometric Society, pp. 347–370, 1991.
  * [38] Y. Hsing, “Economic growth and income inequality: the case of the us,” International Journal of Social Economics, 2005.
  * [39] M. Cagetti and M. De Nardi, “Wealth inequality: Data and models,” Macroeconomic Dynamics, vol. 12, no. S2, pp. 285–313, 2008.
  * [40] Y. Achdou, J. Han, J.-M. Lasry, P.-L. Lions, and B. Moll, “Income and wealth distribution in macroeconomics: A continuous-time approach,” tech. rep., National Bureau of Economic Research, 2017.
  * [41] J. Tufto, “Covariance between a normally distributed variable and its exponent.” Cross Validated.  URL:https://stats.stackexchange.com/q/274853 (version: 2017-04-20).

##  Appendix A Parameter Covariances

The time dependent resources of an agent is given by Eq. 1

$r(t)=r_{0}\exp{\bigg{[}\bigg{(}\eta-\frac{\sigma^{2}}{2}\bigg{)}t+\sigma
W(t)}\bigg{]}$ (9)

We compute the population effective growth rate over population of size $N$ by
separating the expected growth rate into a mean and covariance term

$\displaystyle(\gamma r)_{N}$ $\displaystyle=$
$\displaystyle\frac{1}{N}\sum_{j=1}^{N}\gamma_{j}r_{j}$ (10) $\displaystyle=$
$\displaystyle\frac{1}{N}\sum_{j=1}^{N}(\gamma_{j}-\bar{\gamma})r_{j}+\frac{1}{N}\sum_{j=1}^{N}\bar{\gamma}r_{j}$
(11) $\displaystyle=$
$\displaystyle\bar{\gamma}\bar{r}+\textrm{covar}_{N}(\gamma_{i},r_{i})$ (12)

where the population effective growth rate can be factored out as
$\gamma^{\prime}=\bar{\gamma}+\textrm{covar}_{N}(\gamma_{i},\frac{r_{i}}{\bar{r}})$.

Assuming $r$ is a lognormal distributed quantity, then so is $r/\bar{r}$.
Define the normal variable $y\equiv\ln
r/\bar{r}\sim\mathcal{N}(0,\sigma_{r})$, then the covariance is equal to
$\textrm{covar}_{N}(\gamma,e^{y})$. The first moment is calculated

$\begin{split}\mathrm{E}[\gamma e^{y}]&=\int d\gamma dy\gamma
e^{y}P(\gamma,y)\\\ &=\int dye^{y}\mathrm{E}[\gamma|Y=y]=\int
dye^{y}(\bar{\gamma}+\rho\frac{\sigma_{\gamma}}{\sigma_{y}}y)\\\
&=\bar{\gamma}\mathrm{E}[e^{y}]+\rho\frac{\sigma_{\gamma}}{\sigma_{y}}\mathrm{E}[ye^{y}]=\bar{\gamma}e^{\sigma^{2}_{y}/2}+\rho\sigma_{\gamma}\sigma_{y}e^{\sigma^{2}_{y}/2}\end{split}$
(14)

Where the expectation value of $ye^{y}$ is calculated in [41]. With an
expectation value defined, the covariance becomes

$\begin{split}\textrm{Cov}[\gamma,e^{y}]&=\textrm{E}[\gamma
e^{\gamma}]-\textrm{E}[\gamma]\textrm{E}[e^{y}]\\\
&=(\bar{\gamma}+\rho\sigma_{\gamma}\sigma_{y})e^{\sigma^{2}_{y}/2}-\bar{\gamma}e^{\sigma^{2}_{y}/2}\\\
&=\rho\sigma_{\gamma}\sigma_{y}e^{\sigma^{2}_{y}/2}\end{split}$ (15)

##  Appendix B Convolution of distributions

We compute a joint probability distribution by convolving our dynamical growth
distribution from Eq. 4 with a distribution of initial resources and growth
rates.

$P(\ln r,t)=\int P(\ln r,t|\gamma,\ln r_{0})P(\ln r_{0},\gamma)d\ln
r_{0}d\gamma$ (16)

The log resources, $\ln{r_{0}}$, and growth rates follows a bivariate normal
distribution

$P(\ln r_{0},G)=\mathcal{N}\big{(}\\{\mu_{0},G\\},K_{\Sigma}\big{)}$ (17)

for mean initial resources $\mu_{0}$ and mean growth rate $G$. $\Sigma$ gives
the covariance matrix, and will be calculated through the convolution process.
The respective variances are given by $\sigma_{r},\sigma_{G}$ with correlation
coefficient $\rho$. The growth rate terms can be factored out of the resource
integral by completing the square. The integral is then evaluated using
Gaussian integral identities, leaving a second round of completing the square,
this time for $\mu_{0}$. The resulting coefficients reduce, leaving the
Gaussian equation

$P(\ln r,t)=\frac{d\ln r}{\sqrt{2\pi\Sigma^{2}}}\exp\bigg{[}-\frac{(\ln
r-\mu_{0}-Gt)^{2}}{2\Sigma^{2}}\bigg{]}$ (18)

Where
$\Sigma^{2}=\sigma_{r}+\sigma^{2}t+2\sigma_{r}\sigma_{G}\rho+\sigma^{2}t^{2}$.
The covariance matrix is thus deduced as

$K_{\Sigma}=\begin{pmatrix}\sigma_{r}^{2}&0&\sigma_{r}\sigma_{G}\rho t\\\
0&\sigma^{2}t&0\\\ \sigma_{G}\sigma_{r}\rho t&0&\sigma_{G}^{2}t^{2}\\\
\end{pmatrix}$ (19)