Introduction

Imitation Learning (IL) or Learning from Demonstrations (LfD) [1,6,18] aims to learn a policy directly from expert demonstrations, without access to the environment for more data or any reward signal. One successful IL paradigm is Generative Adversarial Imitation Learning (GAIL) [18], which employs generative adversarial network (GAN) [15] to jointly learn a generator (as a stochastic policy) to mimic expert behaviors, and a discriminator (as a reward signal) to distinguish the generated vs expert behaviors. The learned policy produces behaviors similar to the expert, and the similarity is evaluated using the reward signal, in Jensen-Shannon (JS) divergence (with a constant shift of log 4 [24]) between the distributions of learner vs expert behaviors. Thus, GAIL can be viewed as a variational divergence minimization (VDM) [25] problem with JS-divergence as the objective.

Figure 1: f -divergences and policies from GAIL, RKL-VIM, and f -GAIL on Walker task [32].

Beyond JS-divergence (as originally employed in GAIL), variations of GAIL have been proposed [18,13,12,20,14], essentially using different divergence measures from the fdivergence family [24,25], for example, behavioral cloning (BC) [26] with Kullback-Leibler (KL) divergence [24], AIRL [13] and RKL-VIM [20] with reverse KL (RKL) divergence [24], and DAGGER [28] with the Total Variation (TV) [7]. Choosing the right divergence is crucial in order to recover the expert policy more accurately with high data efficiency (as observed in [20,14,18,13,25,33]).

Motivation. All the above literature works rely on a fixed divergence measure manually chosen a priori from a set of well-known divergence measures (with an explicit analytic form), e.g., KL, RKL, JS, ignoring the large space of all potential divergences. Thus, the resulting IL network likely learns a sub-optimal learner policy. For example, Fig. 1 shows the results from GAIL [18] and RKL-VIM [20], which employ JS and RKL divergences, respectively. The learned input density distributions (to the divergence functions) are quite dispersed (thus with large overall divergence) in Fig. 1(a), leading to learner policies with only 30%-70% expert return in Fig. 1(b). In this work, we are motivated to develop a learnable model to search and automatically find an appropriate discrepancy measure from the f -divergence family for GAIL.

Our f -GAIL. We propose f -GAIL -a new generative adversarial imitation learning model, with a learnable f -divergence from the underlying expert demonstrations. The model automatically learns an f -divergence between expert and learner behaviors, and a policy that produces expert-like behaviors. In particular, we propose a deep neural network structure to model the f -divergence space. Fig. 1 shows a quick view of our results: f -GAIL learns a new and unique f -divergence, with more concentrated input density distribution (thus smaller overall divergence) than JS and RKL in Fig. 1(a); and its learner policy has higher performance (80%-95% expert return) in Fig. 1(b) (See more details in Sec 4). The code for reproducing the experiments are available at https: //github.com/fGAIL3456/fGAIL. Our key contributions are summarized below:

• We are the first to model imitation learning with a learnable divergence measure from fdivergence space, which yields better learner policies, than pre-defined divergence choices (Sec 2). • We develop an f * -network structure, to model the space of f -divergence family, by enforcing two constraints, including i) convexity and ii) f (1) = 0 (Sec 3). • We present promising comparison results of learned f -divergences and the performances of learned policies with baselines in six different physics-based control tasks (Sec 4).

2 Problem Definition

Preliminaries

Markov Decision Processes (MDPs). In an MDP denoted as a 6-tuple S, A, P, r, ρ 0 , γ where S is a set of states, A is a set of actions, P : S × A × S → [0, 1] is the transition probability distribution, r : S × A → R is the reward function, ρ 0 : S → R is the distribution of the initial state s 0 , and γ ∈ [0, 1] is the discount factor. We denote the expert policy as π E , and the learner policy as π. In addition, we use an expectation with respect to a policy π to denote an expectation with respect to the trajectories it generates:

, with s 0 ∼ ρ 0 , a t ∼ π(a t |s t ), s t+1 ∼ P(s t+1 |s t , a t ) and h as any function.

f -Divergence. f -Divergence [24,23,11] is a broad class of divergences that measures the difference between two probability distributions. Different choices of f functions recover different divergences, e.g. the Kullback-Leibler (KL) divergence, Jensen-Shannon (JS) divergence, or total variation (TV) distance [22]. Given two distributions P and Q, an absolutely continuous density function p(x) and q(x) over a finite set of random variables x defined on the domain X , an f -divergence is defined as

with the generator function f : R + → R as a convex, lower-semicontinuous function satisfying f (1) = 0. The convex conjugate function f * also known as the Fenchel conjugate [16] is

]} (See more details in [25]). Common choices of f functions are summarized in Tab. 1 and the plots of corresponding f * are visualized in Fig. 4.

Imitation Learning as Variational f -Divergence Minimization (VDM). Imitation learning aims to learn a policy for performing a task directly from expert demonstrations. GAIL [18] is an IL solution employing GAN [15] structure, that jointly learns a generator (i.e., learner policy) and a discriminator (i.e., reward signal). In the training process of GAIL, the learner policy imitates the behaviors from the expert policy π E , to match the generated state-action distribution with that of the expert. The distance between these two distributions, measured by JS divergence, is minimized. Thus the GAIL objective is stated as follows:

where T is a binary classifier distinguishing state-action pairs generated by π vs π E , and it can be viewed as a reward signal used to guide the training of policy π. H(π) = E π [-log π(a|s)] is the γdiscounted causal entropy of the policy π [18]. Using the variational lower bound of an f -divergence, several studies [20,14,25,5] have extended GAIL to a general variational f -divergence minimization (VDM) problem for a fixed f -divergence (defined by a generator function f ), with an objective below,

However, all these works rely on manually choosing an f -divergence measure, i.e., f * , which is limited by those well-known f -divergence choices (ignoring the large space of all potential fdivergences), thus lead to a sub-optimal learner policy. Hence, we are motivated to develop a new and more general GAIL model, which automatically searches an f -divergence from the f -divergence space given expert demonstrations.

2.2 Problem Definition: Imitation Learning with Learnable f -Divergence.

Divergence Choice Matters! As observed in [20,14,13,25,33], given an imitation learning task, defined by a set of expert demonstrations, different divergence choices lead to different learner policies. Taking KL divergence and RKL divergence (defined in eq. ( 4) below) as an example, let p(x) be the true distribution, and q(x) be the approximate distribution learned by minimizing its divergence from p(x). With KL divergence, the difference between p(x) and q(x) is weighted by p(x). Thus, in the ranges of x with p(x) = 0, the discrepancy of q(x) > 0 from p(x) will be ignored.

On the other hand, with RKL divergence, q(x) becomes the weight. In the ranges of x with q(x) = 0, RKL divergence does not capture the discrepancy of q(x) from p(x) > 0. Hence, KL divergence can be used to better learn multiple modes from a true distribution p(x) (i.e., for mode-covering), while RKL divergence will perform better in learning a single mode (i.e., for mode-seeking).

Beyond KL and RKL divergences, there are infinitely many choices in the f -divergence family, where each divergence measures the discrepancy between expert vs learner distributions from a unique perspective. Hence, choosing the right divergence for an imitation learning task is crucial and can more accurately recover the expert policy with higher data efficiency.

f -GAIL: Imitation Learning with Learnable f -Divergence. Given a set of expert demonstrations to imitate and learn from, the f -divergence, that can highly evaluate the discrepancy between the learner and expert distributions (i.e., the largest f -divergence from the family), can better guide the learner to learn from the expert (as having larger improvement margin). As a result, in addition to the policy function π, the reward signal function T , we aim to learn a (convex conjugate) generator function f * as a regularization term to the objective. The f -GAIL objective is as follows,

where F * denotes the admissible function space of f * , namely, each function in F * represents a valid f -divergence. The conditions for a generator function f to represent an f -divergence include: i) convexity and ii) f (1) = 0. In other words, the corresponding convex conjugate f * needs to be i) convex (the convexity constraint), ii) inf u∈dom f * {f * (u) -u} = 0 (the zero gap constraint, namely, the minimum distance from f * (u) to u is 0)foot_0 . Functions satisfying these two conditions form the admissible space F * . Note that the zero gap constraint can be obtained by combining convex conjugate f (v) = sup u∈dom f * {uv -f * (u)} and f (1) = 0. Tab. 1foot_1 below shows a comparison of our proposed f -GAIL with the state-of-the-art GAIL models [18,13,14,20]. These models use pre-defined f -divergences, where f -GAIL can learn an f -divergence from f -divergence family. Table 1: f -Divergence and imitation learning (JS * is a constant shift of JS divergence by log 4). 5) IL Method FAIRL [14] RKL-VIM [20], AIRL [13] GAIL [18] f -GAIL (Ours) 3

Imitation Learning with Learnable f -Divergence

There are three functions to learn in the f -GAIL objective in eq. ( 5), including the policy π, the f * -function f * , and the reward signal T , where we model them with three deep neural networks parameterized by θ, ω and φ respectively. Following the generative-adversarial approach [15], f * φ and T ω networks together can be viewed as a discriminator. The policy network π θ is the generator. As a result, the goal is to find the saddle-point of the objective in eq. ( 5), where we minimize it with respect to θ and maximize it with respect to ω and φ. In this section, we will tackle two key challenges including i) how to design an algorithm to jointly learn all three networks to solve the f -GAIL problem in eq. ( 5)? (See Sec 3.1); and ii) how to design the f * φ network structure to enforce it to represent a valid f -divergence? (See Sec 3.2). Fig. 2 shows the overall f -GAIL model structure.

f -GAIL Algorithm

Our proposed f -GAIL algorithm is presented in Alg. 1. It uses the alternating gradient method (instead of one-step gradient method in f -GAN [25]) to first update the f * -function f * φ and the reward signal T ω in a single back-propagation, and then update the policy π θ . It utilizes Adam [21] gradient step on ω to increase the objective in eq. ( 5) with respect to both T ω and f * φ , followed by a shifting operation on f * φ to guarantee the zero gap constraint (See Sec 3.2 and eq. ( 7)). Then, it uses the Trust Region Policy Optimization (TRPO) [29] step on θ to decrease eq. ( 7) with respect to π θ .

Algorithm 1 f -GAIL

Require: Initialize parameters of policy π θ , reward signal T ω , and f * φ networks as θ 0 , ω 0 and φ 0 (with shifting operation eq. ( 7) required on φ 0 to enforce the zero gap constraint); expert trajectories τ E ∼ π E containing state-action pairs. Ensure: Learned policy π θ , f * -function f * φ and reward signal T ω . 1: for each epoch i = 0, 1, 2, ... do

2:

Sample trajectories τ i ∼ π θi .

3:

Sample state-action pairs: D E ∼ τ E and D i ∼ τ i with the same batch size.

4:

Update ω i to ω i+1 and φ i to φ i+1 by ascending with the gradients:

5:

Estimate the minimum gap δ with gradient descent in Alg. 2 and shift f * φi+1 (by eq. 7).

6:

Take a policy step from θ i to θ i+1 , using the TRPO update rule to decrease the objective:

-ÊDi [f * φi+1 (T ωi+1 (s, a))] -H(π θi ). The architecture of the f * φ network is crucial to obtain a family of convex conjugate generator functions f * that represents the entire f -divergence space. To achieve this goal, two constraints need to be guaranteed (as discussed in Sec 3.2), including i) the convexity constraint, i.e., f * (u) is convex, and ii) the zero gap constraint, i.e., inf u∈dom f * {f * (u) -u} = 0. To enforce the convex constraint, we implement the f * φ network with a neural network structure convex to its input. Moreover, in each epoch, we estimate the minimum gap of δ = inf u∈dom f * {f * (u) -u}, with which we shift it to enforce the zero gap constraint. Below, we detail the design of the f * φ network. 1. Convexity constraint on f * φ network. The f * φ network takes a scalar input u from the reward signal network T ω output, i.e., u = T ω (s, a), with (s, a) as a state-action pair generated by π θ . To ensure the convexity of the f * φ network, we employ the structure of a fully input convex neural network (FICNN) [3] with a composition of convex nonlinearites (e.g., ReLU) and linear mappings (See Fig. 2). The convex structure consists of multiple layer perceptrons. Differing from a fully connected feedforward structure, it includes shortcuts from the input layer u to all subsequent layers, i.e., for each layer i = 0,

where z i denotes the i-th layer activation, g i represents non-linear activation functions, with W (z) 0 ≡ 0. b s is a bias over both the input u and the last layer output z k , which is used to enforce the zero gap constraint (as detailed below). As a result, the parameters in

1:k-1 to be non-negative and g i 's to be convex non-decreasing activation functions (e.g. ReLU) guarantee the network output to be convex to the input u = T ω (s, a). The convexity follows the fact that a non-negative sum of convex functions is convex and that the composition of a convex and convex non-decreasing function is also convex [9]. To ensure the non-negativity on W

3:

, where δ = inf

This shift guarantees zero gap constraint, and we delegate the proof to Appendix A. In each epoch, the estimation process of δ is detailed in Alg. 2, and the shift operation is implemented by updating b s = b s -δ/2. Fig. 3 illustrates the operations of estimating δ and shifting f * φ . Note that δ represents the minimum gap in function value between f * φ (u) and u. Shifting δ/2 over both input and output space of f * φ (u) (i.e., Line 5 in Alg. 1) enforces the zero gap constraint. Note that this shifting operation is also performed, when initializing the parameters φ 0 for f * φ (u), to make sure the training starts from a valid f -divergencefoot_2 .

Experiments

We evaluate Alg. 1 by comparing it with baselines on six physical-based control tasks, including the CartPole [8] from the classic RL literature, and five complex tasks simulated with MuJoCo [32], such as HalfCheetah, Hopper, Reacher, Walker, and Humanoid. By conducting experiments on these tasks, we show that i) our f -GAIL algorithm can learn diverse f -divergences, comparing to the limited choices in the literature (See Sec 4.1); ii) f -GAIL algorithm always learn policies performing better than baselines (See Sec 4.2); iii) f -GAIL algorithm is robust in performance with respect to structure changes in the f * φ network (See Sec 4.3). Each task in the experiment comes with a true reward function, defined in the OpenAI Gym [10]. We first use these true reward functions to train expert policies with trust region policy optimization (TRPO) [29]. The trained expert policies are then utilized to generate expert demonstrations. To evaluate the data efficiency of f -GAIL algorithm, we sampled datasets of varying trajectory counts from the expert policies, while each trajectory consists of about 50 state-action pairs. Below are five IL baselines, we implemented to compare against f -GAIL.

• Behavior cloning (BC) [26]: A set of expert state-action pairs is split into 70% training data and 30% validation data. The policy is trained with supervised learning. BC can be viewed as minimizing KL divergence between expert's and learner's policies [20,14]. • Generative adversarial imitation learning (GAIL) [18]: GAIL is an IL method using GAN architecture [15], that minimizes JS divergence between expert's and learner's behavior distributions. • BC initialized GAIL (BC+GAIL): As discussed in GAIL [18], BC initialized GAIL will help boost GAIL performance. We pre-train a policy with BC and use it as initial parameters to train GAIL. • Adversarial inverse reinforcement learning (AIRL) [13]: AIRL applies the adversarial training approach to recover the reward function and its policy at the same time, which is equivalent to minimizing the reverse KL (RKL) divergence of state-action visitation frequencies between the expert and the learner [14]. • Reverse KL -variational imitation (RKL-VIM) [20]: the algorithm uses the RKL divergence instead of the JS divergence to quantify the divergence between expert and learner in GAIL architecture 4 .

For fair comparisons, the policy network structures π θ of all the baselines and f -GAIL are the same in all experiments, with two hidden layers of 100 units each, and tanh nonlinearlities in between. The implementations of reward signal networks and discriminators vary according to baseline architectures, and we delegate these implementation details to Appendix B. All networks were always initialized randomly at the start of each trial. For each task, we gave GAIL, BC+GAIL, AIRL, RKL-VIM and f -GAIL exactly the same amount of environment interactions for training.

4.1 f * φ Learned from f -GAIL Fig. 4 shows that f -GAIL learned unique f * φ (u) functions for all six tasks, and they are different from those well-known divergences, such as RKL and JS divergences. Clearly, the learned f * φ (u)'s are convex and with zero gap from u, thus represent valid f -divergences. Moreover, the learned f -divergences are similar, when the underlying tasks share commonalities. For example, the two f * φ (u) functions learned from CartPole and Reacher tasks (Fig. 4(a) and (d)) are similar, because the two tasks are similar, i.e., both aiming to keep a balanced distance from the controlling agent to a target. On the other hand, both Hopper and Walker tasks aim to train the agents (with one foot for Hopper and two feet for Walker) to proceed as fast as possible, thus their learned f * φ (u) are similar (Fig. 4(c) and(e)). (See Appendix B for descriptions and screenshots of tasks.) We also plot Fig. 5 to show that given a task, the learned f * functions are consistent (small variance) for different sample sizes. Similar observations are made for tasks CartPole, Reacher and Humanoid as well.

In state-of-the-art IL approaches and our f -GAIL (from eq. ( 3) and ( 5)), the f * -function takes the learner reward signal u = T ω (s, a) (over generated state-action pairs (s, a)'s) as input. By examining the distribution of u, two criteria can indicate that the learner policy π θ is close to the expert π E : i. u centers around zero gap, i.e., f * (u) -u ≈ 0. This corresponds to the generator function f centered around f (p(s, a)/q(s, a)) ≈ f (1) = 0, with p and q as the expert vs learner distributions; ii. u has small standard deviation. This means that u concentrates on the nearby range of zero gap, leading to a small f -divergence between learner and expert, since D f (p(s, a) q(s, a)) ≈ q(s, a)f (1)d(s, a) = 0. φ (u) -u ≈ 0, equivalently, f (p(s, a)/q(s, a)) ≈ f (1) = 0, with close expert vs learner behavior distributions (i.e., p vs q). The distributions of input u were estimated by kernel density estimation [31] with Gaussian kernel of bandwidth 0.3. φ (u) with different sample sizes. In Fig. 4, we empirically estimated and showed the distributions of input u for the state-of-theart IL methods (including GAIL and RKL-VIM 5 ) and our f -GAIL. Fig. 4 shows that overall u distributions from our f -GAIL match the two criteria (i.e., close to zero gap and small standard deviation) better than baselines (See more statistical analysis on the two criteria across different approaches in Appendix B). This indicates that learner policies learned from f -GAIL are with smaller divergence, i.e., higher quality. We will provide experimental results on the learned policies to further validate this in Sec 4.2 below.

f -GAIL Performance in Policy Recovery

Fig. 6 shows the performances of our f -GAIL and all baselines under different training data sizes, and the tables in Appendix B provide detailed performance scores. In all tasks, our f -GAIL outperforms all the baselines. Especially, in more complex tasks, such as Hopper, Reacher, Walker, and Humanoid, f -GAIL shows a larger winning margin over the baselines, with at least 80% of expert performances for all datasets. GAIL shows lower performances on complex tasks such as Hopper, Reacher, Walker, and Humanoid, comparing to simple tasks, i.e., CartPole and HalfCheetah (with much smaller state and action spaces). Overall, BC and BC initialized GAIL (BC+GAIL) have the lowest performances comparing to other baselines and our f -GAIL in all tasks. Moreover, they suffer from data efficiency problem, with extremely low performance when datasets are not sufficiently large. These results are consistent with that of [19], and the poor performances can be explained as a result of compounding error by covariate shift [27,28]. AIRL performs poorly for Walker, with only 20% of expert performance when 4 trajectories were used for training, which increased up to 80% when using 25 trajectories. RKL-VIM had reasonable performances on CartPole, Hopper, Reacher, and Humanoid when sufficient amount of data was used, but was not able to get more than 80% expert performance for HalfCheetah, where our f -GAIL achieved expert performance. (See Tab. 6 in Appendix B for more detailed return values.) In terms of the convergence of the proposed f -GAIL, Fig. 7 below shows the training curve of f -divergence (i.e., the objective in eq. ( 5)) with respect to training epochs where it converges to less than 0.02 for HalfCheetah after 450 epochs. Similar results were observed in other tasks.

Ablation Experiments

In this section, we investigate how structure choices of the proposed f * φ network, especially, the network expressiveness such as the number of layers and the number of nodes per layer, affect the model performance. In experiments, we took the CartPole, HalfCheetah and Reacher tasks as examples, and fixed the network structures of policy π θ and the reward signal T ω . We changed the number of layers to be 1, 2, 4, and 7 (with 100 nodes each layer) and changed the number of nodes per layer to be 25, 50, 100 and 200 (with 4 layers). The comparison results are presented in Tab. 2. In simpler tasks with smaller state and action space, e.g. the CartPole, we observed quick convergence with f -GAIL, achieving expert return of 200. In this case, the structure choices do not have impact on the performance. However, in more complex tasks such as HalfCheetah and Reacher, a simple linear transformation of input (with one convex transformation layer) is not sufficient to learn a good policy function π θ . This naturally explains the better performances with the number of layers increased to 4 and the number of nodes per layer increased to 100. However, further increasing the number of layers to 7 and the number of nodes per layer to 200 decreased the performance a little bit. As a result, for

Discussion and Future Work

Our work makes the first attempt to model imitation learning with a learnable f -divergence from the underlying expert demonstrations. The model automatically learns an f -divergence between expert and learner behaviors, and a policy that produces expert-like behaviors.

Meaning of the best f -divergence. As a minimax optimization problem in eq. ( 5), f -GAIL searches for the best f -divergence in the "max" inner-loop given the current learned policy π learned from the "min" outer-loop, eventually leading to a stable solution of (π, f * ). Here, given an expert demonstration dataset, a better divergence can measure the discrepancy more precisely than other divergences, thus enables training a learner with closer behaviors to the expert.

Future work. This work focuses on searching within the f -divergence space, where Wasserstein distance [17,4] is not included. However, the divergence search space can be further extended to c-Wasserstein distance family [2], which subsumes f -divergence family and Wasserstein distance as special cases. Designing a network structure to represent c-Wasserstein distance family is challenging (we leave it as part of our future work), while a naive way is to model it as a convex combination of the f -divergence family (using our f * φ network) and Wasserstein distance. Moreover, beyond imitation learning, our f * -network structure can be potentially "coupled" with f -GAN [25] and f -EBM [33] to learn an f -divergence between the generated vs real data distributions (e.g., image and audio files), which in turn trains a higher quality generator.

Broader Impact

This paper aims to advance the imitation learning techniques, by learning an optimal discrepancy measure from f -divergence family, which has a wide range of applications in robotic engineering, system automation and control, etc. The authors do not expect the work will address or introduce any societal or ethical issues.