Constrained Imitation Q-learning with Earth Mover’s Distance reward
fi
.fi}
Abstract We propose constrained Earth Mover’s Distance (CEMD)
Imitation Q-learning that combines exploration of Reinforcement Learning (RL) and the sample efficiency of Imitation Learning (IL).
Sample efficiency makes CEMD suitable for robot learning.
Immediate rewards can be efficiently computed by a greedy Earth Mover’s Distance (EMD) variant between observed state-action pairs and state-actions in the stored expert demonstrations.
In CEMD, we constrain the previously proposed nonstationary greedy EMD reward by proposing a greedy EMD upper bound estimate and a generic Q-learning lower bound.
In PyBullet continuous control benchmarks, CEMD is more sample efficient, achieves higher performance, and yields less variance than its competitors.
Introduction
Recently, approaches combining Reinforcement Learning (RL) and Imitation Learning (IL) have gained momentum as RL and IL have complementary strengths and weaknesses for robot learning [4, 17].
Reinforcement Learning enables robots to improve autonomously, but introduces significant challenges with exploration and stable learning.
Imitation learning provides more stable learning from expert demonstrations, but cannot improve from failures.
It is unclear how the two approaches should be combined to take the best of the both worlds.
There have been multiple attempts to combine RL and IL.
The first methods simply learned preliminary policy or its parts from expert demonstrations, and then further improved the policy based on feedback received from the environment [27, 31].
More recently, methods based on Inverse Reinforcement Learning (IRL) have been proposed [11, 6, 9].
IRL methods iterate between learning the underlying reward function for the expert demonstrations, and running RL to find the optimal policy for the estimated rewards by interacting with the environment.
The IRL methods suffer from sample inefficiency which makes them less suitable for robot learning.
The sample hungry inverse reward optimization can be avoided by defining a computational reward function. A computational reward is computed using the difference between states and actions obtained from the environment and those in the expert demonstrations.
This line of work has been found more sample efficient, and the proposed methods mainly differ in how they compute the reward [1, 8, 4].
The proposed method in this work belongs to this group of methods.
The expert demonstrations are treated as finite distributions, and the reward is computed via distance metric between the expert and agent obversation distributions.
In our work, we adopt the Earth Mover’s Distance (EMD) as the basic framework. EMD has been used in recent works of Imitation Learning [16, 8, 4]
The main contribution of this work is 1) a new RL-based imitation learning method, CEMD, that uses Earth Mover’s Distance (EMD) to compute the RL reward.
The EMD reward reflects how well the obtained actions and states correspond to those in the expert demonstrations.
As the starting point, we adopt the greedy EMD from [4] which can be computed instantly after each new exploration sample.
Our EMD reward is tightened by more strict upper and lower bounds. We 2) derive a greedy EMD reward upper bound and 3) a generic Q-learning lower bound. The lower bound is adapted from [30].
The tight reward constraints make CEMD sample efficient during exploration, reduce its variance, and achieve better overall performance.
These findings are experimentally validated with PyBullet continuous control benchmarks.
Preliminaries 2.1 Reinforcement learning
In the classic RL setting [27], Markov Decision Processes (MDPs) are formalized as a tuple M := ( S, A, r, P, γ, P 0 ) with the state space S and the action space A .
A reward function r : S ×
A → R maps a state-action pair to a reward.
P ( s t +1 | s t , a t ) denotes the probability of transferring from state s t to state s t +1 under action a t .
P 0 ( s ) is the initial state distribution and γ ∈ [0 , 1] is a discount factor.
At each time step t , an agent observes a state s t ∈ S and chooses an action a t ∈ A based on a policy π ( a t | s t ) :
S → P ( A ) .
The return R ( τ ) of an MDP is defined as the discounted sum of rewards r ( s t , a t ) along a trajectory τ :=
( s 0 , a 0 , ..., s t − 1 , a t − 1 , s t ) .
The classic objective of RL is to learn an optimal policy π ∗ that maximizes the expected cumulative discounted reward:
[Equation]
The state-action value function, Q-function Q π ( s, a ) , provides the expected cumulative discounted reward of a state s if an action a is taken and the trajectory τ according to the policy π is followed.
Formally, Q-function is defined as Q π ( s, a ) =
E π
[ R ( τ ) | s, a ] .
The value function obeys the Bellman equation: Q π ( s, a ) =
E s t +1 ∼ P ( s t +1 | s,a ) (cid:2) r
( s, a ) + γ E a t +
∼ π
[ Q π ( s t +1 , a t +1 )] (cid:3) .
The Bellman operator B π : R n → R n describes a single step of dynamic programming [27, 21]:
[Equation]
Q π is the fixed point of the operator B π ( Q ) to which the process Q k +1 :
= B π ( Q ) k converges.
For high-dimensional and continuous state and action spaces, we need to approximate Q ( s, a ) to efficiently solve for Q π .
Let’s denote with θ and ω the parameters of a Q-function Q θ ( s, a ) and the associated parameterized policy π ω , correspondingly.
In addition, a double target function Q ¯ θ is often used to reduce value overestimation [7].
Now, improving Q θ ( s, a ) relies on finding values θ that minimize the Mean Squared Bellman Error (MSBE).
Since the approximation function often takes the form of a neural network, it is convenient to define the error as a loss term:
[Equation]
where Q target =
r ( s, a ) + Q θ ( s t +1 , π ω ( s t +1 )) is the target for minimizing MSBE, or alternatively
Q target =
r ( s, a ) + Q ¯ θ ( s t +1 , π ω ( s t +1 )) in the double target function formulation.
Optimization of the MSBE loss (1) is the fundamental step in many SotA off-policy RL methods [18, 15, 7, 2, 10].
2.2 Earth Mover’s Distance (EMD)
The Earth Mover Distance (EMD, a.k.a Wasserstein metric or distance) [25] is based on the minimal cost that must be paid to transform one distribution into another.
Consider two finite distributions:
X = { ( x 1 , w 1 ) , ..., ( x M , w M ) } with points x i ∈ R d and Y = { ( y 1 , u 1 ) , ..., ( y N , u N ) } with y j ∈ R d , where w i ∈ R and u j ∈ R represent the weights of each point in X and Y .
The flow matrix
F = ( f ij ) ∈
R m × n has elements f ij that denote the amount of weight at x i matched to the amount of weight at y j .
The EMD between two finite distributions is cast as a linear programming problem to find F that minimizes
[Equation]
where d ij = d ( x i , y j ) is the metric distance between x i and y j in the d -dimensional feature (or state-action) space.
In this case, the flow F is a solution to the optimal transport problem such that the overall transportation cost is minimized.
The EMD linear programming problem has the following constraints: (i) f ij ≥ 0 requires the all weights of x i matched to y j be non-negative; (ii) (cid:80) Nj =1 f ij ≤ w i , i = 1 , 2 , ..., M ensures that the transported weight from each x i over all y j does not exceed w i ; (iii) (cid:80)
Mi
=1 f ij ≤ u j , j = 1 , 2 , ..., N ensures that weights over all x i transported to each y j do not exceed u j ; and (iv) (cid:80) i,j f ij = min( (cid:80) w i , (cid:80)
u j ) forces the total flow to be equal to the total weights of the smaller of the two distributions.
Fig.
1(a) visualizes the original EMD that we denote as the optimal coupling EMD .
Greedy EMD
As mentioned in [4], the original optimal coupling EMD has an undesirable property to require that the two distributions X an Y are known in order to solve the optimal weight transportation problem.
Consequently, Q-function updates using the MSBE loss (1) and EMD can be computed only after an exploration episode is completed.
This leads to sample inefficient episodic learning which in [4] is solved by proposing a greedy coupling EMD .
To bring EMD to the Imitation Learning (IL) context, we define X as the agent’s sample trajectory after M samples, and Y of N elements is the static set of trajectories obtained through expert demonstrations.
The sets X and Y are sequences of observed states, { s 0 , s 1 , ..., s N − 1 } for X , or state-action pairs {⟨ s 0 , a 0 ⟩ , ⟨ s 1 , a 1 ⟩ , ..., ⟨ s N − 1 , a N − 1 ⟩} , depending on the problem.
The greedy EMD ( EMD g ) is defined for horizon of M observations in X as
[Equation]
Let’s further assume that N expert trajectory points have the capacity of u i
= 1 N and the mass of the each obtained agent sample is w i = 1 M .
The greedy EMD cost e gi for the sample s i is defined as:
[Equation]
The first greedy EMD constraint ensures that all the agent’s dirt (EMD terminology) appearing at the time step i is transported.
The second constraint guarantees that each of the expert holes cannot take more dirt load than its capacity allows.
Greedy EMD provides a sub-optimal transport solution, and is an upper bound estimate of the true EMD costs e ⋆i [4]:
[Equation]
The RL agent objective is to minimize the upper bound estimate of the EMD distance between the expert and the agent policies.
Figure 1(b) illustrates how the agent mass s πi is assigned to the expert holes s ej under the greedy-coupling EMD.
The reward r i = r
( s, a ) in MSBE loss (1) is computed from e gi using some monotonically decreasing mapping g ( · ) ,
[Equation]
(5) that converts small distances to large rewards.
A common mapping is the exponential function g ( x ) =
e − x .
Method
In [4] the greedy coupling EMD (3) is used to compute the sample distance to the expert demonstrations after each sample, and the computed distance is converted to a greedy reward r g by the exponential mapping (5).
Q-function updates are computed using the Mean Squared Bellman Error (MSBE) loss (1).
We denote this greedy EMD reward [4] loss as L g .
Next, we ropose tighter bounds for the loss by introducing the greedy EMD reward upper bound (Section 3.1) and a general Q-learning lower bound (Section 3.2).
The new loss term is defined in Section [-3.
-]3. ej dirt thus sometimes moves it further closest occupied this example vector s π
temporally s π and therefore closest position s e 2 is occupied and the dirt needs to be moved s π our greedy s moved to the closest s e 2 even if it is already occupied
3.1 Upper bound of the Greedy EMD reward
It turns out that the greedy reward r g has two undesirable properties, i) it is non-stationary, and ii) it underestimates the reward by over-estimating the EMD distance (the proof is in our supplementary material).
To improve the reward computation to boost learning by a tighter estimate, we introduce the lower bound e gLB of the greedy EMD effort e g .
The lower bound is constructed by removing the second constraint in (3).
Lemma
By removing the constraint 2 in Eq.
3, the earth moving effort e gLB for the dirt x t appearing at time step t becomes
[Equation]
and this new effort e gLB is a lower bound for the greedy coupling EMD effort
e g in (3), i.e., ∀
[Equation]
, e gLB
[Equation] e [Equation] By removing the second constraint, the lower bound greedy EMD is allowed to allocate new observations x t to the nearest expert hole even if the hole is already occupied by the previous observations.
This is not possible in the original greedy EMD formulation.
The differences between the optimal EMD, greedy EMD and Lower Bound greedy EMD couplings are visualized in Figure 1.
For Q-learning the main result is that the lower bound greedy EMD formulation produces an upper bound estimate of the greedy EMD reward , i.e.
[Equation]
Using the upper bound greedy EMD reward we define a new loss term L g,LB and combine it with the original greedy loss term L g to establish a tighter EMD loss for Q-function optimization
[Equation]
3.2 Generic Q-learning Lower Bound
In off-policy Reinforcement Learning one of the main challenges is how to facilitate efficient learning from off-policy samples that can be very noisy, especially in the beginning of learning.
One of the recently proposed approaches tries to benefit from a small number of particularly good episodes stored in the learning memory (a replay buffer).
This line of works is often called as Self-Imitation
Learning (SIL) [20, 30].
The main theoretical foundation of the recent SIL methods is in tightening the Q-learning bounds.
Intuitively, we learn only from samples that produce better Q-values than the current lower bound.
This is particularly important in our case as the previously defined greedy EMD Lower Bound (6) forms an upper bound for the greedy EMD.
Therefore, by identifying a similar lower bound, it is possible to tighten the estimator from the both directions.
It turns out that we can form the greedy EMD lower bound by utilizing the generic Q-learning lower bound.
The lower bound Q-learning was studied in [30] where they introduced an n-step bootstrapped Q-value as the general lower bound for the true Q-value under the optimal policy π ∗
[Equation]
In [30] the Q-function Q θ ( s, a ) is updated only when Q θ ( s, a ) is greater than the lower bound
Q LB ( s, a ) , and in our case the rule is formulated inside the loss function,
[Equation]
where || · || + is the max( · , 0) operator.
As shown in [30], this n-step clipped loss brings positive bias for value estimation, and thus makes it suitable for reducing underestimation.
3.3 Constrained EMD Loss
Imitation learning with the proposed constrained EMD (CEMD) Q-estimation uses a combination of the three loss terms that represent the greedy EMD based reward, and its upper and lower bounds,
[Equation]
L g is the greedy EMD reward loss from [4], L g,UB is the greedy EMD reward upper bound loss defined in (9), and L LB is the greedy EMD reward lower bound loss in (10).
The experiments in the next section verify that the compound loss results to better overall performance, better sample efficiency (faster convergence), and smaller variance, than the greedy EMD based reward loss without the bounds.
In addition, CEMD is superior to other competing IL methods.
Experiments 4.1 Tasks and settings
The RL baseline of our method is TD3 [7] which is a popular choice for off-policy RL.
The implementation is based on [22].
Experiments were conducted using the continuous control environments of PyBullet [5].
The task horizon is 1000 steps in all experiments.
For each task, only five expert demonstrations were recorded and each of them is sub-sampled to the horizon length of 50 steps.
The purpose of sub-sampling is to simulate the low data regime. The same setting was used in [4].
4.2 Method comparison
In the first experiment, variants of the proposed CEMD are compared to SotA RL and IL methods.
The tested CEMD variants are: 1) CEMD-1 using the greedy EMD loss as in [4], 2) CEMD-2 using the greedy EMD and the upper bound loss ( λ 1 = 0 . 3 ), 3) CEMD-3 using the greedy EMD and the generic lower bound loss ( λ 2 = 0 . 3 ), and 4) CEMD-4 using all three losses ( λ 1 = λ 2 = 0 . 15 ).
For the generic lower bound loss, L LB , the bootstrapping step size of 6 was used.
We trained all variants with the limit of 1M environment interactions and 10 random seeds.
Imitation Learning research is divided to multiple lines of work, such as Behavior Cloning (BC) [24],
Inverse Reinforcement Learning (IRL) [12, 19, 11, 9, 6], adversarial IL [11, 9, 6], and expert support estimation [26, 13, 3].
We chose the following off-policy methods as strong baselines to compare with: Behavior Cloning (BC) [24], Discriminator-Actor-Critic (DAC) [14], soft-Q Imitation Learning (SQIL) [23], and Adversarial Reward-Moment Imitation Learning [29].
DAC uses an off-policy RL method (TD3) to learn the imitation policy.
SQIL learns to recover expert behavior by balancing the expert transitions with the agent transitions in the replay buffer.
AdRIL was selected for the comparison as it uses trajectory matching for reward computation, and at the moment it is largely considered as the SotA IL method [29].
For BC, SQIL and AdRIL, we used the code provided by the authors of [29, 28] who conducted similar experiments.
All hyper-parameters remain the same as in the experiments in [29].
The same five expert demonstrations were provided to all IL methods.
[Table]
We first evaluated the performance by computing the cumulative test rewards in five environments (see Table 1).
Our CEMD outperformed the original greedy EMD [4] in all tasks.
In HalfCheetah, all variants had similar performance.
Among the variants, the one with all three loss terms was the best for Humanoid, Walker2D [Table] and Hopper.
For the most challenging tasks (Ant and Humanoid), the greedy upper bound variants obtained the best results.
Overall, the proposed CEMD variants achieve performance gains 0.5%-41.0% depending on the task.
Table 2 shows the true EMD distances to the expert demonstrations.
The results of Table 2 verify the results in Table 1.
400 rollouts are collected for each task in this experiment.
CEMD-4 that includes all three loss terms has the smallest EMDs for all tasks.
The importance of the greedy upper and generic lower bound losses is particularly evident for the Humanoid and Ant tasks for which the original greedy EMD loss in [4] remains substantially better than the others.
Clearly, the proposed tightening bounds improve Q-learning using the greedy EMD based loss.
4.3 Empirical verification of the greedy EMD upper bound
It is important to empirically verify that the greedy EMD lower bound estimate defined in (6) and its corresponding reward upper bound r gUB in (8) indeed represent upper bounds for the greedy EMD.
Verification is important as the Q-learning is based on updates using the greedy reward r g .
This was made by running the CEMD method with the original greedy reward in [4] (CEMD-1). During training, number of episodes were unrolled and the average differences between the greedy EMD reward and its upper bound estimates were computed.
At every 1k environment exploration steps, a 10-episode random seed roll-out was performed.
It means that every evaluation stop generates ten trajectories and each has 1000 steps horizon.
Using these steps the greedy reward r g and its greedy upper bound reward r g,UB were computed.
The mean difference between the two, r g,UB − r g , were recorded, and stored as a single evaluation value.
These evaluation values over the increasing number of training episodes should reveal the assymptotic behaviour between the two rewards, and whether the upper bound holds for the greedy estimate.
Figure 2 shows the differences during training.
In all tasks, the differences are negative, verifying that r g,UB indeed is an upper bound estimate.
Moreover, there are clear differences between the two values, that confirms that the upper bound reward affects learning performance.
When the policy gradually improves, and starts to converge near the optimal solution, the difference asymptotically approaches zero, and verifies negligible bias in the estimate. [Table] [Figure] [Figure]
4.4 Benefits of the bounded reward for Q-learning
To verify that the bounded reward converges faster than the unbounded greedy EMD reward the two were compared.
The comparison was made in a realistic case where the number of expert demonstrations was low (five).
We computed the overall task rewards and the behavior cloning errors for the greedy EMD reward r g Q-learning (CEMD-1), and its bounded version.
The mean test rewards, variances, and behavior cloning errors are shown in Figure 3.
The results verify that Imitation Q-learning with bounded rewards converges faster (particularly evident in the
Ant task), and achieves better performance than the unbounded reward.
This systematically holds for all tested environments.
Moreover, the variance is reduced (the standard deviations in Table 1).
Limitations
Similarly to other imitation learning methods, the reliance on expert demonstrations is a limitation of our approach.
In tasks where experts do not provide optimal demonstrations, or where the demonstrations are performed in an environment which is not identical to the actual environment, running RL in the actual environment may be needed.
Nevertheless, imitation learning serves as initialization of the policy.
Since we use RL to minimize the distance to expert demonstrations, the performance depends on the selected RL.
TD3 performed well in the experiments, but other methods should be considered in the future work.
Our approach builds on Earth Mover’s Distance (EMD) bounds.
While the current bounds already provide a performance boost, better bounds should be investigated.
Conclusion
In this work, we proposed an Imitation Q-learning method, CEMD, that learns to solve a control task using expert demonstrations to guide its learning.
The expert demonstrations make Imitation Learning more sample efficient than plain Reinforcement Learning (RL), and therefore more suitable for real robot cases.
Our main contributions are the tighter bounds for Earth Mover’s Distance (EMD)
based reward computation. The bounds boost Q-learning to converge faster, to achieve higher performance, and to reduced variance.
The results were verified by multiple continuous control tasks where the proposed method systematically outperforms its competitors.
All code will be made publicly available Acknowledgments Wenyan Yang received funding from the Doctoral School of Industry Innovations of Tampere
University, and Cargotec Ltd.