Title: GENERATIVE PREDECESSOR MODELS FOR SAMPLE-EFFICIENT IMITATION LEARNING

Abstract: We propose Generative Predecessor Models for Imitation Learning (GPRIL), a novel imitation learning algorithm that matches the state-action distribution to the distribution observed in expert demonstrations, using generative models to reason probabilistically about alternative histories of demonstrated states. We show that this approach allows an agent to learn robust policies using only a small number of expert demonstrations and self-supervised interactions with the environment. We derive this approach from first principles and compare it empirically to a stateof-the-art imitation learning method, showing that it outperforms or matches its performance on two simulated robot manipulation tasks and demonstrate significantly higher sample efficiency by applying the algorithm on a real robot.
Recent advances in generative modeling, such as Goodfellow et al. (2014) ; Kingma and Welling (2013) ; Van Den Oord et al. (2016b;a); Dinh et al. (2016) , have shown great promise at modeling complex distributions and can be used to reason probabilistically about such state-action pairs.

1. INTRODUCTION:

Training or programming agents to act intelligently in unstructured and sequential environments is a difficult and central challenge in the field of artificial intelligence. Imitation learning provides an avenue to tackle this challenge by allowing agents to learn from human teachers, which constitutes a natural way for experts to describe the desired behavior and provides an efficient learning signal for the agent. It is thus no surprise that imitation learning has enabled great successes on robotic (Chernova and Thomaz, 2014) as well as software domains (e.g. Aytar et al. (2018) ). Yet, key challenges in the field are diverse and include questions such as how to learn from observations alone (e.g. Aytar et al. (2018)) , learning the correspondence between the expert's demonstrations and the agent's observations (e.g. Sermanet et al. (2017) ) as well as the question of how to integrate imitation learning with other approaches such as reinforcement learning (e.g. Vecerik et al. (2017) ). However, at the core of the imitation learning problems lies the challenge of utilizing a given set of demonstrations to match the expert's behavior as closely as possible. In this paper, we approach this problem considering the setting where the set of expert demonstrations is given up-front and the dynamics of the environment can only be observed through interaction.

In principle, imitation learning could be seen as a supervised learning problem, where the demonstrations are used to learn a mapping from observed states to actions. This solution approach is known as behavioral cloning. However, it has long been known that the sequential structure of the task admits more effective solutions. In particular, the assumptions made in supervised learning are restrictive and don't allow the agent to reason about the effect of its actions on it's future inputs. As a result, errors and deviations from demonstrated behavior tend to accumulate over time as small mistakes lead the agent to parts of the observation space that the expert has not explored (Ross and Bagnell, 2010) . In this work, we propose a novel imitation learning algorithm, Generative Predecessor Models for Imitation Learning (GPRIL), based on a simple core insight: Augmenting the training set with state-action pairs that are likely to eventually lead the agent to states demonstrated by the expert is an effective way to train corrective behavior and to prevent accumulating errors. Specifically, we propose to utilize Masked Autoregressive Flows (Papamakarios et al., 2017) to model long-term predecessor distributions, i.e. distributions over state-action pairs which are conditioned on a state that the agent will see in the future. Predecessor models have a long history in reinforcement learning (e.g. Peng and Williams (1993) ) with recent approaches using deep networks to generate off-policy transitions (Edwards et al., 2018; Pan et al., 2018) or to reinforce behavior leading to high-value states (Goyal et al., 2018) . Here, we use predecessor models to derive a principled approach to state-distribution matching and propose the following imitation learning loop:

1. Interact with the environment and observe state, action as well as a future state. To encode long-term corrective behavior, these states should be multiple steps apart. 2. Train a conditional generative model to produce samples like the observed state-action pair when conditioned on the observed future state. 3. Train the agent in a supervised way, augmenting the training set using data drawn from the model conditioned on demonstrated states. The additional training data shows the agent how to reach demonstrated states, enabling it to recover after deviating from expert behavior.

In the above, we laid out the sketch of an algorithm that intuitively learns to reason about the states it will observe in the future. In section 3, we derive this algorithm from first principles as a maximum likelihood approach to matching the state-action distribution of the agent to the expert's distribution. In section 4, we compare our approach to a state-of-the-art imitation learning method (Ho and Ermon, 2016) and show that it matches or outperforms this baseline on our domains while being significantly more sample efficient. Furthermore, we show that GPRIL can learn using demonstrated states alone, allowing for a wider variety of methods to be used to record demonstrations. Together these properties are sufficient to allow GPRIL to be applied in real-world settings, which we demonstrate in section 4.3. To our knowledge this is the first instance of dynamic, contact-rich and adaptive behavior being taught solely using the kinesthetic-teaching interface of a collaborative robot, without resorting to tele-operation, auxiliary reward signals, or manual task-decomposition.

2. BACKGROUND:



2.1. MARKOV DECISION PROCESSES WITHOUT REWARDS:

As is usual, we model the problem as a Markov decision process without reward. That is, given state and action sets S, A, the agent is observing states s ∈ S and taking actions a ∈ A. In this work, we use s and a to refer to states and actions observed during self-supervision and s ∈ S and a ∈ A to refer to target and demonstration states and actions. We furthermore use superscripts s (i) , a (i) to refer to specific instances, e.g. specific demonstrated state-action pairs, and subscripts, e.g. s t , a t , to indicate temporal sequences. The observed transitions are guided by the Markovian dynamics of the environment and the probability of transitioning from state s to state s by taking action a is denoted as p(s t+1 = s |s t = s, a t = a). The agent's behavior is defined by a stationary parametric policy π θ (a|s) while the expert's behavior is modeled by a stationary distribution π * (a|s). We denote as d π t (s) the probability of observing state s at time-step t when following policy π. Under the usual ergodicity assumptions, each such policy induces a unique stationary distribution of observed states d π (s) = lim t→∞ d π t (s) as well as a stationary joint state-action distribution ρ π (s, a) := π(a|s)d π (s). Furthermore, we use q π t to refer to the dynamics of the time reversed Markov chain induced by a particular policy π at time-step t

q π t (s t = s, a t = a|s t+1 = s ) = d t+1 (s ) -1 d t (s)π(a t |s t )p(s t+1 = s |s t = s, a t = a)

(1) and define q π (s t = s, a t = a|s t+1 = s ) := lim t→∞ q π t (s t = s, a t = a|s t+1 = s ). For the purposes of this work, we handle the episodic case with clear termination conditions by adding artificial transitions from terminal states to initial states. This creates a modified, ergodic MDP with identical state-distribution and allows us to assume arbitrarily large t such that q π t = q π . Finally, we extend this notation to multi-step transitions by writing q π (s t = s, a t = a|s t+j = s ).

2.2. IMITATION LEARNING:

In this work, we are considering two settings of imitation learning. In the first setting, the agent is given a set of observed states s (1) , s (2) , • • • , s (N ) and observed corresponding actions a (1) , a (2) , • • • , a (N ) as expert demonstrations. The goal in this setting is to learn a policy π θ that matches the expert's behavior as closely as possible. In the second setting, the agent is given the states observed by the expert but is not aware of the actions the expert has taken. Recent years have seen heightened interest in a related setting where to goal is to track expert state trajectories (Zhu et al., 2018; Peng et al., 2018; Pathak et al., 2018) . These approaches do not learn general policies that can adapt to unseen situations. A straightforward approach to train general policies in the first setting, usually referred to as behavioral cloning, is to treat the task as a supervised learning problem (e.g. Pomerleau (1989) ). However, as outlined in section 1, predictions made by π θ (a|s) influence future observations thus violating a key assumption of supervised learning, which states that inputs are drawn from an i.i.d. distribution. This has formally been analyzed by Ross et al. who introduce a family of algorithms (e.g. Ross and Bagnell (2010) ; Ross et al. (2011) ) that provably avoid this issue. However, these approaches require the expert to continuously provide demonstrations and thus are not applicable when the set of demonstrations is assumed to be fixed. A popular avenue of research that considers this setting is inverse reinforcement learning (IRL). Inverse reinforcement learning (Ng and Russell, 2000) aims to learn a reward function for which π * is optimal and thus captures the intent of the expert. The arguably most successful approach to IRL aims to match the state-action distribution to that of the demonstrations (Ziebart et al., 2008) with recent approaches extending these ideas to the model free case (Boularias et al., 2011; Finn et al., 2016; Fu et al., 2018) .

However, inverse reinforcement learning is indirect and ill-defined as many reward functions induce the same behavior. Recently, methods have been proposed that aim to match state-action distributions directly and achieve state-of-the-art result without learning a reward function first. Generative Adversarial Imitation Learning (GAIL) (Ho and Ermon, 2016) uses an adversarial objective, training a discriminator to identify demonstrations and using TRPO to train a policy that fools the discriminator. While GAIL is able to achieve impressive results, the adversarial objective can make the learning procedure unstable and unpredictable. This is especially true when parts of the state-space are not under the agent's control, yielding a setting resembling a conditional GAN (Mirza and Osindero, 2014) , which are prone to issues such mode collapse (Odena et al., 2017) . We compare our approach with GAIL in section 4. State Aware Imitation Learning (SAIL) (Schroecker and Isbell, 2017) is an alternative method that aims to learn the gradient of the state-action distribution using a temporal-difference update rule. This approach is able to avoid instabilities prompted by the adversarial learning rule but is only applicable to policies with a small number of parameters where learning a representation of the gradient is feasible. In this work, we follow a gradient descent approach similar to SAIL but estimate the gradient without representing it explicitly by a neural network. These methods can also be used to match state-distributions without actions. While the order of states is unspecified and the induced policy is therefore not unique, including transition information such as velocities can allow the agent to learn solely from expert state trajectories nonetheless.

2.3. GENERATIVE MODELS:

Recent years have seen great advances in deep generative models. A variety of approaches such as Generative Adversarial Networks (Goodfellow et al., 2014) , Variational Auto Encoders (Kingma and Welling, 2013) , autoregressive networks (e.g. Germain et al. (2015) ; Van Den Oord et al. (2016b; a) and normalizing flows (Dinh et al., 2016) have been proposed which enable us to learn complex distributions and efficiently generate samples. In this work, we use generative models to model the distribution of long-term predecessor state-action pairs. While the approach we propose is model agnostic, we choose to model this distribution using masked autoregressive flows (Papamakarios et al., 2017) (MAF) . MAFs are trained using a maximum likelihood objective, which allows for a stable and straight-forward training procedure. Autoregressive models are capable of representing complex distributions p(x); x ∈ R n by factoring the distribution p

(x) = p 1 (x 1 ) N -1 i=1 p i+1 (x i+1 |x 1 , . . . , x i

) and learning a model for each p i . In this paper, we model each x i to be distributed by x i ∼ N (•|µ i , σ i ) where each µ i and σ i is a function of x 1:i-1 . Masked autoencoders (Germain et al., 2015) provide a straight-forward approach to parameter sharing and allow representing these functions using a single network. MAFs stack multiple autoregressive models with different orderings and thus avoid the strong inductive bias imposed by the order of variables. Using the reparameterization trick, the autoregressive model can be seen as a deterministic and invertible transformation of a random variable: x = f (z) := µ + σz; z ∼ N (•|0, 1). The change of variable formula then allows us to calculate the density of x:

log p(x) = log p N (f -1 (x)) + log det(|J(f -1 (x))|) (2)

Algorithm 1 Generative Predecessor Models for Imitation Learning (GPRIL)

1: function GPRIL(NB, Nπ, B (batch size)) 2: for i ← 0..#Iterations do 3:

for k ← 0..NB do 4:

for n ← 0..B do 5:

Sample s

(n) t , a (n) t from replay buffer 6: Sample j ∼ Geom(1 -γ) 7: Sample s (n)

t+j from replay buffer 8:

Update ωs using gradient B n=0 ∇ω s log B s ωs (s

(n) t |s (n) t+j ) 9:

Update ωa using gradient B n=0 ∇ω a log B a ωa (a

(n) t |s (n) t , s (n) t+j ) 10:

for k ← 0..Nπ do 11:

for n ← 0..B do 12:

Sample s (n) , a (n) from expert demonstrations 13:

Sample

s (n) ∼ B s ωs (•|s (n) ), a (n) ∼ B a ωa (•|s (n) , s (n) ) 14: Update θ using gradient B n=0 βπ∇ θ log π θ (a (n) |s (n) ) + β d ∇ θ log π θ (a (n) |s (n) )

Where the autoregressive nature of f in eq. 2 allows for tractable computation of the second term.

MAFs chain multiple such transformations to derive highly expressive explicit density models able to model complex dynamics between target states and long-term predecessor state-action pairs.

3. GPRIL:

In section 1, we provided an intuitive framework for using predecessor models to augment our training set and achieve robust imitation learning from few samples. In this section, we will derive this algorithm based on state-action distribution matching. To this end, we first derive the gradient of the logarithmic state distribution based on samples from a long-term predecessor distribution that we will define below. In section 3.2, we describe how to train a generative model of the predecessor distribution, which will allow us to evaluate this gradient. Ascending on this gradient evaluated at demonstrated states leads the agent to stick to those states and thus provides a corrective measure (Schroecker and Isbell, 2017) . Furthermore, reproducing the states of the expert can be sufficient to achieve the correct behavior if the state-space is chosen appropriately as we will show in section 4. We will show how to use this gradient to match state-action-distributions in section 3.3.

3.1. ESTIMATING THE GRADIENT OF THE STATE DISTRIBUTION:

Here, we will show that the samples drawn from a long-term predecessor distribution conditioned on s enable us to estimate the gradient of the logarithmic state distribution ∇ θ log d π θ (s) and, later, to match the agent's state-action-distribution to that of the expert. To achieve this goal, we can utilize the fact that the stationary state distribution of a policy can be defined recursively in terms of the state distribution at the previous time step, similar to Morimura et al. (2010) :

EQUATION

Taking the derivative shows that this notion extends to the gradient as well as its logarithm:

EQUATION

The recursive nature of this gradient then allows us to unroll the gradient indefinitely. However, this process is cumbersome and will be left for appendix A. We arrive at the following equality:

EQUATION

The derivation of our approach now rests on two key insights: First, in ergodic Markov chains, such as the ones considered in our setting, decisions that are made at time t affect the probability of seeing state s at time t+j more strongly if j is small. In the limit, as j → ∞, the expectation of the gradient ∇ θ log π θ (a t |s t ) vanishes and the decision at time t only adds variance to the gradient estimate. Introducing a discount factor γ similar to common practice in reinforcement learning (Sutton and Barto, 1998) places more emphasis on decisions that are closer in time and can thus greatly reduce variance. We explore this interpretation further in appendix C. Second, by introducing a discount factor, the effective time-horizon is now finite. This allows us to replace the sum over all states and actions in each trajectory with a scaled expectation over state-action pairs. Formally, we can write this as follows and arrive at our main result:

EQUATION

where B π θ corresponds to the long-term predecessor distribution modeling the distribution of states and actions that, under the current policy π θ , will eventually lead to the given target state s:

EQUATION

3.2. LONG-TERM GENERATIVE PREDECESSOR MODELS:

In the previous section, we derived the gradient of the logarithm of the stationary state distribution as approximately proportional to the expected gradient of the log policy, evaluated at samples obtained from the long-term predecessor distribution B π θ . In this work, we propose to train a model B π θ ω to represent B π θ and use its samples to estimate ∇ θ log d π θ (s). However, rather than unrolling a time-reversed Markov model in time, which is prone to accumulated errors, we propose to use a generative model to directly generate jumpy predictions. We have furthermore found that imposing a sensible order on autoregressive models achieves good results and thus propose to use two conditional MAFs (Papamakarios et 

To train this model, we collect training data using self-supervised roll-outs: We sample states, actions and target-states where the separation in time between the state and target-state is selected randomly based on the geometric distribution parameterized by γ as a training set for B π θ ω . Training data for B s ωs , B a ωa are obtained by executing the current policy to obtain a sequence s 0 , a 0 , s 1 , a 1 , • • • , which we store in a replay buffer. In practice, we store data from multiple iterations in this buffer in order to decrease the variance of the gradient. While our algorithm does not explicitly account for off-policy samples, we found empirically that a short replay buffer does not degrade final performance while significantly improving sample efficiency. To obtain a training sample, we first pick s = s t and a = a t for a random t. We now select a future state s = s t+j+1 from that sequence. For any particular s t+j+1 we now have s, a ∼ q

π θ t (s t = •, a t = •|s t+j+1 = s) ≈ q π θ (s t = •, a t = •|s t+j+1 = s).

Note that in the episodic case, we can add transitions from terminal to initial states and pick t to be arbitrarily large such that the approximate equality becomes exact (as outlined in section 2.1). In non-episodic domains, we find the approximation error to be small for most t. Finally, we choose j at random according to a geometric distribution j ∼ Geom(1γ) and have a training triple s, a, s that can be used to train B a ωa and B s ωs as it obeys

EQUATION

)

3.3. MATCHING STATE-ACTION DISTRIBUTIONS WITH GPRIL:

State-action distribution matching has been a promising approach to sample-efficient and robust imitation learning (see section 2.2). While each policy induces a unique distribution of states and behavioral cloning would therefore be sufficient in the limit, it is sub-optimal the case of limited data. Matching the joint-distribution directly ensures that we minimize discrepancies between everything we observed from the expert and the behavior the agent exhibits. In this work, we propose a maximum likelihood based approach, ascending on the estimated gradient of the joint distribution: ∇ θ log ρ π θ (s, a) = ∇ θ log π θ (a|s) + ∇ θ log d π θ (s).

(11) where ∇ θ log π θ (a|s) can be computed directly by taking the gradient of the policy using the demonstrated state-action pairs and ∇ θ log d π θ (s) can be evaluated using samples drawn from B π θ ω (•, •|s) according to equation 7. We introduce scaling factors β π and β d to allow for finer control, interpolating between matching states only (β π = 0) and behavioral cloning (β d = 0) and have:

∇ θ log ρ π θ (s, a) ≈ β π log π θ (a|s) + β d E s,a∼B π θ (•,•|s) [∇ θ log π θ (a|s)] .

(12) Here, higher values of β π provide more supervised guidance while lower values aim to prevent accumulating errors. This gives rise to the full-algorithm: We fill the replay buffer by asynchronously collecting experience using the current policy. Simultaneously, we repeatedly draw samples from the replay buffer to update the predecessor models and use expert samples in combination with an equal number of artificial samples to update the policy. This procedure is described fully in algorithm 1.

4. EXPERIMENTS:

To evaluate our approach, we use a range of robotic insertion tasks similar to the domains introduced by Vecerik et al. (2017) but without access to a reward signal or, in some cases, expert actions. We choose these domains both for their practical use, and because they highlight challenges faced when applying imitation learning to the real world. Specifically, collecting experience using a robot arm is costly and demands efficient use of both demonstrations and autonomously gathered data. Furthermore, insertion tasks typically require complex searching behavior and cannot be solved by open-loop tracking of a given demonstration trajectory when the socket position is variable. We first compare against state-of-the-art imitation learning methods on a simulated clip insertion task, then explore the case of extremely sparse demonstrations on a simulated peg insertion task and finally, demonstrate real-world applicability on its physical counterpart.

4.1. CLIP INSERTION:

In the first task, a simulated robot arm has to insert an elastic clip into a plug, which requires the robot to first flex the clip in order to be able to insert it (see figure 1a ). In real-world insertion tasks, the Figure 2 : a) Depiction of the peg task b) Average success rate and 95% confidence interval of GAIL with 25 demonstrations across 10 runs (evaluated over 100 roll-outs). c) Average success rate of GPRIL across 5 seeds. Unlike GAIL, the performance of GPRIL doesn't drop off when provided with only final states. d) Average success rate and confidence interval of GPRIL. Final performance after 10 6 iterations increases steadily as the number of demonstrated trajectory increases but is unaffected by dropping steps from each demonstration. e) Median length and IQR of trajectories that are successfully inserting the peg. Providing only final states is significantly faster. f) Best seed performance on both variations of peg insertion on the real robot. pose of the robot, the socket, or the grasped object may vary. We capture this variability by mounting the socket on a pan-tilt unit, which is randomized by ±0.8 (pan) and ±0.2 radians (tilt). To perform this behavior, the robot observes proprioceptive features, specifically joint position, velocity and torques as well as the position of the end-effector and the socket orientation as a unit quaternion. The task terminates when the robot leaves the work-space, reaches the goal, or after 50 seconds.

For comparative evaluation, we train a policy network to predict mean and variance, modelling a multivariate normal distribution over target velocities and train it using GPRIL, GAIL as well as behavioral cloning. We record expert demonstrations using tele-operation and normalize observations based on the recorded demonstrations. We then train GPRIL using a single asynchronous simulation and compare against the open source implementation of GAIL1 for which we use 16 parallel simulations. We select the best hyper parameters for GAIL found on a grid around the hyperparameters used by Ho and Ermon (2016) but lower the batch size to 256 as it increases the learning speed and accounts for the significantly slower simulation of the task. We furthermore enable bootstrapping regardless of whether or not the episode terminated. As all discriminator rewards are positive, handling terminal transitions explicitly can induce a bias towards longer episodes. This is beneficial in the domains used by Ho and Ermon but harmful in domains such as ours where the task terminates on success. A detailed list of hyper-parameters can be found in appendix B.

We report final results after convergence and can see in figure 1b that both GAIL and GPRIL outperform behavioral cloning, indicating that generalizing over state-action trajectories requires fewer demonstrations than generalizing over actions alone. Furthermore, we observe a higher success rate using GPRIL and find that policies trained using GPRIL are more likely to retry insertion if the robot slides the clip past the insertion point. To compare sample efficiency of GPRIL to GAIL, we limit the rate at which the asynchronous actor is collecting data. While sample efficiency of GAIL could be increased by decreasing batch size or increasing various learning rates, we found that this can lead to unstable learning performance while reducing the amount of samples required by only a small amount. As can be seen in figure 1c , GPRIL requires several orders of magnitudes fewer environment interactions to learn this task. Finally, we evaluate the case where the expert's actions are unknown. Since the state-space includes information about joint velocities as well as positions, we find that matching the state-distribution is sufficient to solve the task. GPRIL can achieve this by setting β d = 1 and β π = 0 . As can be seen in figure 1d , performance deteriorates only marginally with a similar difference in performance between both methods.

4.2. PEG INSERTION WITH PARTIAL DEMONSTRATIONS:

The second task is a simulated version of the peg-insertion task depicted in figure 2a . In this task, the robot has to insert the peg into the hole, which is again mounted on a pan-tilt that randomly assumes pan and tilt angles varying by 0.4 and 0.1 respectively. Hyperparameters are largely identical and we report minor differences in appendix B. Observation and action space are identical with the exception of the omission of torques from the observation space as they are not necessary to solve this task. We use this task to evaluate the performance of GAIL and GPRIL when learning from only a very limited set of demonstrated states. To this end, we compare three different scenarios in which the demonstrations are sparsified to varying degrees: In the first case, the agent has access to the full state-trajectories of the expert, in the second only every tenth state is available and in the third the agent sees only the final state of each of the 25 trajectories. Being able to learn from only partial demonstrations is a useful benchmark for the effectiveness of imitation learning methods but can also provide a convenient way of providing demonstrations and can free the agent to find more optimal trajectories between states (see for example Akgun et al. (2012) ; Schroecker et al. (2016) ). As can be seen in figures 2d and 2e, GPRIL achieves similar final success rates in all three scenarios while being able to learn a significantly faster insertion policy when learning from final states alone. We find that in the first two scenarios, this holds for GAIL as well as can been in figure 2b while in the third case, GAIL becomes highly unstable and the resulting performance can vary wildly, leading to a low average success rate. We hypothesize that these instabilities are a result of the discriminator overfitting to the very small amount of negative samples in its training data.

4.3. PEG INSERTION ON A PHYSICAL SYSTEM:

In previous sections we demonstrated sample-efficiency that indicates applicability of GPRIL to real-world physical systems. To test this, we evaluate our approach on two variations of the physical peg-insertion task depicted in figure 2a , involving a physical pan-tilt unit which is fixed in one scenario and has pan and tilt angles varying by 0.1 and 0.02 radians in the second scenario. For each scenario, we provide 20 demonstrations using kinesthetic teaching, which constitutes a natural way of recording demonstrations but provides state-trajectories only (Chernova and Thomaz, 2014) . Hyper-parameters are altered from section 4.2 to trade off a small margin of accuracy for higher learning speeds and are reported in appendix B. Note, however, that tuning hyper-parameters precisely is very difficult on a physical system. As can be seen in figure 2f , GPRIL is able to learn a successful insertion policy that generalizes to unseen insertion angles using just a few hours of environment interactions2 . We report best-seed performance as we observe a high amount of variability due to factors outside the agent's control, such as the pan-tilt unit not reporting accurate information after physical contact with the robot. However, we wish to point out that the increased difficulty due to less predictable control is also likely to introduce additional variance that could be reduced further with careful design of exploration noise and other hyper-parameters. We furthermore provide a video of the training procedure and final policy to highlight the efficiency of our method3 .

5. CONCLUSION:

We introduced GPRIL, a novel algorithm for imitation learning which uses generative models to model multi-step predecessor distributions and to perform state-action distribution matching. We show that the algorithm compares favorably with state-of-the-art imitation learning methods, achieving higher or equivalent performance while requiring several orders of magnitude fewer environment samples. Importantly, stability and sample-efficiency of GPRIL are sufficient to enable experiments on a real robot, which we demonstrated on a peg-insertion task with a variable-position socket.

A DERIVATION OF EQUATION 6:

Here, we derive equation 6 which unrolls the recursive definition of ∇ θ log d π θ (s) and rewrites it such that it can be replaced by an expectation over states and actions along trajectories leading to the state s. In section 3.1, we derive a recursive definition of ∇ θ log d π θ (s) which we will restate in more detail:

∇ θ d π θ (s) = ∇ θ ρ π θ (s)p(s t+1 = s|s t = s, a t = a)ds, a d π θ (s)∇ θ log d π θ (s) = ρ π θ (s, a)p(s t+1 = s|s t = s, a t = a)∇ θ log ρ π θ (s, a)ds, a ∇ θ log d π θ (s) = q π θ (s t = s, a t = a|s t+1 = s) (∇ θ log d π θ (s) + ∇ θ log π θ (a|s)) ds, a

We can now unroll this definition:

EQUATION

Note that lim T →∞ q π θ (s t = s, a t = a|s t+T = s) = ρ π θ (s, a) due to Markov chain mixing and, therefore, the second term of the above sum reduces to 0 as

EQUATION

)

By pulling out the sum, we can now marginalize out most variables and shift indices to arrive at the desired conclusion:

EQUATION

C RELATION TO THE POLICY GRADIENT THEOREM

In this section, we outline the relation between the policy gradient theorem (Sutton et al., 1999) and the state-action-distribution gradient derived in section A and show equivalence of the discount factor used in reinforcement learning and the discount factor γ introduced in this work. We first show that the state-action distribution gradient is equal to the policy gradient in the average reward case using a specific reward function. We then derive the γ-discounted approximation of the state-actiondistribution gradient presented in section 3.1 from the policy gradient in the discounted reward framework using γ as a discount factor. While this derivation is more cumbersome than the one presented in the body of the paper, it allows us to gain a better understanding of the meaning of γ since discounting in the reinforcement learning setting is well understood. For notational simplicity, we assume that the state space S and action space A are countable. Note that this is a common assumption made in the field (Tsitsiklis and Van Roy, 1997) and considering that any state-action space is countable and finite when states and actions are represented using a finite number of bits makes it apparent that this assumption does not constitute a simplification in the practical sense.

In this work we propose to follow the gradient ∇ θ log ρ π θ evaluated at demonstrated states and actions, consider thus the gradient evaluated at s, a: 

∇ θ log ρ π θ (

As we can see, the gradient is equivalent to the policy gradient using the reward function R(s, a) := 1(s=s,a=a)

ρ π θ (s,a)

in the average reward framework where θ corresponds to the parameters of the policy at the current iteration. This reward function is not practical as it can be infinitely sparse and furthermore depends on the unknown stationary distribution. However, it allows us to derive the policy gradient using this reward function in the discounted reward setting which constitutes a well understood approximation of the average-reward scenario. Following the notation in Sutton et al. (1999) , we have: 

∇ θ J(π θ ) = E sj ,

We can now replace the expectation over the stationary distribution by an expectation over the path of the agent. This yields a double sum whose order of summation can be changed: While this derivation is less direct than the derivation used in section A, it draws a connection between the problem of matching state-action-distributions and the reinforcement learning problem with a reward that is positive for demonstrated state-action pairs and 0 otherwise. This indicates that the role of the discount factor is similar in both settings: Lower values of γ trade-off accurate matching of distributions for lower variance and as expedience. Agents with low γ will learn policies that recover and reach demonstrated states quicker over policies that are matching the experts state-action distribution more accurately long-term. Furthermore, the long history of successes in reinforcement learning validates the use of γ as an approximation to the true objective which is often more accurately described by the average-reward objective.

∇ θ J(π θ ) = lim

:

https://github.com/openai/baselines/tree/master/baselines/gail

Total time to collect and train on 2000 roll-outs was 18.5 and 16.5 hours on the fixed and changing versions of the task respectively. However, GPRIL converged to good policies significantly sooner.

https://youtu.be/Dm0OCNujEmE

Figures and Tables:

Figure 1: a) Depiction of the clip-insertion task. b) Median final success rate and interquartile range out of 100 roll-outs over 8 seeds. GPRIL achieves the highest success rate followed by GAIL. c) Median final success rate and IQR on clip insertion comparing sample efficiency. GPRIL is able to solve the task using several orders of magnitude fewer environment interactions. d) Comparison on clip insertion trained on states alone. Learning from states alone only slightly affects performance.

s, a) = 1 ρ π θ (s, a) ∇ θ s∈S,a∈A ρ π θ (s, a)1(s = s, a = a) = ∇ θ E s,a∼ρ π θ [R(s, a)] =: ∇ θ J(π θ )

θ log π θ (a j |s j ) T t=0 γ t R(s t+j , a t+j )|π θ π θ (a j |s j )γ t-j R(s t , a t )|π θ j ∇ θ log π θ (a j |s j )R(s t , a t )|π θ order of summation we can replace the outer sum with the expectation over the stationary distributions and use the special nature of the chosen reward function to write the gradient as a conditional expectation:∇ θ J(π θ ) = lim t→∞ j ∇ θ log π θ (a j |s j )R(s t , a t )|π θ j ∇ θ log π θ (a j |s j )|π θ , s t = s, a t = a   (19)Published as a conference paper at ICLR 2019 Finally, we notice that this equation constitutes the discounted version of equation 15, thus we can immediately obtain our estimate of the state-action-distribution gradient:∇ θ J(π θ ) = 1 1γ E s,a∼B π θ (•,•|s) [∇ θ log π θ (a|s)] + ∇ θ log π θ (a|s) ≈ ∇ θ log ρ π θ (s, a)(20)

aj ∼ρ π θ ∇ θ log π θ (a j |s j )E sj+1,aj+1,...

Table content: <table/>

Peer review (iclr2019):
 The submission builds up on recent advances in neural density estimation to develop a new algorithm for imitation learning based on a probabilistic model for predecessor state dynamics. In particular, the method trains masked autoregressive flows as a probabilistic model for state action pairs conditioned on future states. This model is used to estimate the gradient of the stationary distribution of a policies visited states. Finally, the proposed objective uses this estimate and the gradient of the log likelihood of expert actions under the policy to maximise the similarity of the expert’s and agent’s stationary state-action distributions. The proposed method outperforms existing imitation learning approach (GAIL & BC) on 2 simulation-based manipulation tasks. It performs particularly well in terms of sample efficiency. The magnitude of difference between the sample efficiencies of GAIL and the proposed approach seems quite surprising and it would be beneficial if the authors could explicitly state if the measured number of samples include the ones used for training of the probabilistic model as well as the policy (apologies if I have missed a section fulfilling this purpose). While the improvements on the presented experiments are clear, the experimental section represents a small shortcoming of the submitted paper. The 2 experiments (clip and peg insertion) are quite similar in type and to not take into account other common domains e.g. locomotion tasks from the original GAIL paper. Furthermore, an additional comparison to SAIL would be recommended since the approaches are closely related as the authors acknowledge. The provided comparison with different types of available expert data is quite interesting and could possibly be extended to test other state-of-the-art methods (action-free versions of GAIL, AIRL,etc.). Nonetheless, the paper overall presents a strong submission based on novelty & relevance of the proposed method and is recommended for publication. Minor issues: - Related work: improve transitions between the section about trajectory tracking and BC. - Ablation studies with less flexible probabilistic models would strengthen the experiment section further. - Add derivation from Eq. 3 to 4 and 5 to appendix to render the paper more self-contained and easier to access. - A release of the code base would further strengthen the contributions of the submission. General recommendation: - The authors are encouraged to further investigate off-policy corrections for improved convergence.