Planning for Sample Ef ﬁ cient Imitation Learning
Abstract
Imitation learning is a class of promising policy learning algorithms that is free from many practical issues with reinforcement learning, such as the reward design issue and the exploration hardness.
However, the current imitation algorithm struggles to achieve both high performance and high in-environment sample ef ﬁ ciency simultaneously.
Behavioral Cloning (BC) does not need in-environment interactions, but it suffers from the covariate shift problem which harms its performance.
Adversarial Imitation Learning (AIL) turns imitation learning into a distribution matching problem.
It can achieve better performance on some tasks but it requires a large number of in-environment interactions.
Inspired by the recent success of
Ef ﬁ cientZero in RL, we propose Ef ﬁ cientImitate (EI), a planning-based imitation learning method that can achieve high in-environment sample ef ﬁ ciency and performance simultaneously.
Our algorithmic contribution in this paper is two-fold.
First, we extend AIL into the MCTS-based RL.
Second, we show the seemingly incompatible two classes of imitation algorithms (BC and AIL) can be naturally uni ﬁ ed under our framework, enjoying the bene ﬁ ts of both.
We benchmark our method not only on the state-based DeepMind Control Suite, but also on the image version which many previous works ﬁ nd highly challenging.
Experimental results show that EI achieves state-of-the-art results in performance and sample ef ﬁ ciency.
EI shows over 4x gain in performance in the limited sample setting on state-based and image-based tasks and can solve challenging problems like Humanoid, where previous methods fail with small amount of interactions.
Our code is available at https://github.com/zhaohengyin/EfficientImitate .
Introduction
The real-world sequential decision process in robotics is highly challenging.
Robots have to handle high dimensional input such as images, need to solve long horizon problems, some critical timesteps need highly accurate maneuver and the learning process on the real robot has to be sample ef ﬁ cient.
Imitation learning is a promising approach to solving those problems, given a small dataset of expert demonstrations.
However, current imitation algorithms struggle to achieve these goals simultaneously.
There are two kinds of popular imitation learning algorithms, Behavior Cloning (BC) and Adversarial
Imitation Learning (AIL).
BC formulates imitation learning as a supervised learning problem.
It needs no in-environment samples, but it suffers from the covariate shift issue [40], often leading to test time performance degradation.
Adversarial Imitation Learning (AIL) [16, 6] casts imitation learning as a distribution matching problem.
Though AIL suffers less from the covariate shift problem and can perform better than BC on some domains, it requires impractical number of online interactions [20, 23] and can perform badly on image inputs [37].
These drawbacks heavily limit its application in ﬁ elds like robotics, where physical robot time matters.
In summary, current imitation learning algorithms fail to achieve high online testing performance and high in-environment sample ef ﬁ ciency at the same time.
Further, the two types of imitation algorithms seem to be incompatible since they are based on two completely different training objectives.
Previous work ﬁ nds that it is hard to unify them naively[34].
Inspired by the recent success in sample ef ﬁ cient RL, such as Ef ﬁ cientZero [52], we propose a planning-based imitation algorithm named Ef ﬁ cientImitate (EI) that achieves high test performance and high in-environment sample ef ﬁ ciency at the same time.
Our method extends AIL to a modelbased setting with multi-step losses under the MCTS-based RL framework.
Our algorithm also uni ﬁ es the two types of the previous imitation algorithms (BC and AIL) naturally, thanks to the planning component of our algorithm.
Intuitively, BC gives a coarse solution that is correct most of the time but fails to match the expert’s behavior in the long term.
On the other hand, AIL knows the goal of the learning, i.e., matching state action distribution, but doesn’t give the solution directly.
Our method’s planning component uni ﬁ es those two methods and shows a signi ﬁ cant performance boost, especially in the harder tasks such as Humanoid.
We illustrate the detailed procedure in Figure 1.
We validate our idea not only on state-based tasks but also on image-based tasks, which relatively few previous algorithms can handle them [37].
EI achieves state-of-the-art results in sample ef ﬁ ciency and performance.
It shows over 4x gain in performance in the limited sample setting on state-based and image-based tasks.
On harder tasks such as Humanoid, the gain is even larger.
A by-product of this paper is that we extend the recent Ef ﬁ cientZero algorithm to the continuous action space.
We open-source the code at https://github.com/zhaohengyin/EfficientImitate to facilitate future research.
Our contributions in this paper are summarized as follows.
• We present Ef ﬁ cientImitate, a sample ef ﬁ cient imitation learning algorithm based on MCTS.
• We identify that MCTS can bene ﬁ t from BC by using BC actions during the search, which is crucial for challenging tasks.
Ef ﬁ cientImitate suggests a natural way to unify BC and AIL.
• We conduct experiments in both the state and the image domain to evaluate EI.
Experimental results show that EI can achieve state-of-the-art sample ef ﬁ ciency and performance.
Related Work 2.1 Imitation Learning
Imitation learning (IL) aims to solve the sequential decision-making problems with expert demonstration data.
It has wide applications in games and robotics [4].
Compared with RL, one major bene ﬁ t of IL is that it can avoid the notorious reward design problem.
IL can be divided into two branches: BC [1] and IRL [33].
To solve the covariate shift problem of BC, researchers propose methods like dataset aggregation [40] and noise injection [25].
But these methods either require extra expert queries or exert constraints over the learning process.
A recent variant branch of IRL is the
Adversarial Imitation Learning (AIL) [16, 6].
AIL models IL as a state-action distribution matching problem.
Many works extend AIL by using better distribution metrics such as f -divergence [54] and
Wasserstein distance [5].
Though these methods can learn better reward functions and somewhat speed up the training, they do not directly focus on AIL’s sample ef ﬁ ciency problem.
Some works have drawn attention to sample ef ﬁ ciency, for example [23, 41, 55] propose to use off-policy training in AIL training to reduce sample complexity.
ValueDICE [24] reformulates AIL objective in an of ﬂ ine min-max optimization process, but recent work points out that it is an improved form of BC [29].
VMAIL [37] uses a model-based approach to improve sample ef ﬁ ciency.
It collects latent rollout with a variational model and reduces online interactions.
MoBILE [22] also shows a provably sample ef ﬁ cient model-based imitation approach.
Compared with these methods, our method introduces
MCTS planning to the off-policy imitation learning and uses it to unify BC and AIL to take advantage of both.
The idea of combining BC and AIL helps to improve the sample ef ﬁ ciency can be traced back to
GAIL[16].
GAIL suggests using BC to initialize the policy network, but [34] ﬁ nds that this does not work well because the initialized BC knowledge will be corrupted in AIL training.
Then, [20] proposes to use an annealed (decaying) BC loss to regularize policy training to solve this problem.
But in this work, we ﬁ nd that such a BC regularizer can be harmful to exploration when BC is incorrect.
One concurrent work mitigates this issue by using adpative BC loss and replacing AIL reward with optimal-transport-based imitation reward [13].
We also notice that AlphaGo [46] involves BC in their method, but they focus on RL rather than IL, and BC is only used to initialize the policy.
Different from these methods, EI uses BC actions as candidates in MCTS.
2.2 Sample Ef ﬁ ciency in RL
The sample ef ﬁ ciency problem of imitation learning is closely related to that in RL.
One line of work ﬁ nds that the reward signal is not a good data source for representation learning in RL and is one reason for sample inef ﬁ ciency.
Then they utilize self-supervised representation learning to accelerate representation learning and improve the sample ef ﬁ ciency.
Researchers propose to use contrastive learning [27], consistency-based learning [44, 53, 52], or pretrained representation [45] for this purpose.
Some works also explore the possibility of applying self-supervised representation learning to imitation learning [4].
Another line of work focuses on RL with a learned model, which is promising for sample ef ﬁ cient learning [8, 10, 11, 12, 19, 28, 31, 52, 14].
These approaches usually imagine additional rollouts with the learned model or use it as a more compact environment representation for RL.
Besides, some works also ﬁ nd that data augmentation can effectively improve sample ef ﬁ ciency [26, 50, 49].
EI also bene ﬁ ts from these approaches.
It includes a model and applies representation learning to boost sample ef ﬁ ciency.
In application, people also consider to augment RL with the demonstration [15, 38, 32, 21] to improve the sample ef ﬁ ciency.
This can be viewed as a combination of RL and imitation learning and favor RL on real robots.
We believe that our method can also be extended to this setting.
Background 3.1 Setting
We formalize the sequential decising making problem as Markov Decision Process M = ( S , A , R , T ) .
Here, S is the state space,
A is the action space, R is the reward function, and
T is the transition dynamics.
The agent’s state at timestep t is s t ∈ S .
The agent takes action a t and receives reward r t =
R ( s t , a t ) .
Its state at timestep t + 1 is then s t +1 ∼ T ( s t , a t ) .
The objective of the agent is to maximize the return
P Tt =
γ t r t , where γ is a discount factor.
In the imitation learning problem studied here, the agent has no access to the reward function R and transition dynamics T .
It is provided with a ﬁ xed expert demonstration dataset D = { τ i } .
Here, each τ i =
( s E 0 , a E 0 , s E 1 , a E 1 , ...
s ET , a ET ) is an expert trajectory that can achieve high performance in
M .
The agent can not solicit extra expert demonstrations but can interact with the MDP, observing new states and actions, but not rewards.
In this work, we de ﬁ ne (in-environment) sample ef ﬁ ciency as the number of online interactions during training.
We expect the agent to achieve high performance within a ﬁ xed online sample budget.
3.2 BC and AIL
BC considers imitation learning as a supervised learning problem.
It trains a policy network π to minimize the following loss function:
[Equation]
AIL treats imitation learning as a distribution matching problem.
One typical AIL algorithm is GAIL.
It trains a discriminator D to distinguish the agent generated state-action tuple ( s t , a t ) from those ( s Ei , a Ei ) in the demonstration dataset by minimizing
[Equation]
where ρ is the state-action distribution induced by the agent.
Meanwhile, it trains the agent to maximize the return with respect to the adversarial reward r t = − log(1 − D ( s t , a t )) using any on-policy RL algorithm.
3.3 MuZero and its Extensions
Our planning method is based on MuZero [42] and its extensions.
MuZero learns an environment model for MCTS.
The model consists of an encoder network f , a dynamics network g , and a reward network R .
It operates on abstract states [52].
Concretely, it gets the abstract state h t of the current state s t by h t = f
( s t ) .
It can then predicts the future abstract states recursively by h t
+1 =
g ( h t , a t ) , and the rewards by R ( h t , a t ) .
Besides the model, MuZero also contains a policy network and a value network.
The policy network provides a prior over the actions at each node, and the value network calculates the expected return of the node.
MuZero uses the model, the policy network, and the value network to search for improved policy and value for each state with MCTS.
We refer the readers to the original MuZero paper for details.
Sampled MuZero Sampled MuZero [17] extends MuZero from the discrete action domain to the continuous action domain, which is of our interest in this paper.
At each node s to expand, Sampled
MuZero samples K actions { a i } Ki =1 from current policy π ( a | s ) .
During the search, it selects action a ∗ from the sampled actions that maximize the probabilistic upper con ﬁ dence bound
[Equation]
where ˆ π ( a | s )
= 1 K P i δ ( a, a i ) .
Q ( s, a ) is the current Q -estimation of the pair ( s, a ) .
N ( s, a ) denotes the times that this pair is visited in MCTS.
c ( s ) is a weighting coeffcient.
During policy optimization, MuZero minimizes the Kullback-Leibler divergence between the current policy π and the MCTS statistics π MCTS at the root node D KL ( π MCTS || π ) .
Ef ﬁ cientZero We also apply Ef ﬁ cientZero [52] in this paper.
Ef ﬁ cientZero improves the sample ef ﬁ ciency of MuZero by using a self-supervised representation learning method to regularize the hidden representation.
It uses a SimSiam-style structure [3] to enforce the similarity between the predicted future representation and the real future representation.
Ef ﬁ cientImitate
In this section, we present our Ef ﬁ cientImitate algorithm.
We ﬁ rst present an MCTS-based approach to solving the AIL problem in Section 4.1.
Then we show a simple yet effective method to unify BC and AIL with MCTS in Section 4.2.
We brie ﬂ y discuss the implementation in Section 4.3, and the full details can be found in the Appendix.
4.1 Extending AIL to MCTS-based RL
Traditionally, the adversarial imitation learning (AIL) algorithm trains a discriminator D between the policy samples and the expert samples and uses some form of D , such as − log(1 − D ) , as the reward function.
Then some model-free RL algorithms are used to maximize the cumulative reward.
In
MCTS-based RL algorithms, such as MuZero [42] and Ef ﬁ cientZero [52], the reward function is used in the value target computation and the MCTS search.
The use in value target computation is similar to prior model-free RL algorithms, where the value target is computed with n-step value bootstrapping on the actual observations.
However, during the MCTS search, the rewards are computed on the abstract state obtained by running the forward dynamics function h t
+1 =
g ( h t , a t ) multiple times.
If we were training the discriminator only on actual observations of the expert and the policy rollouts, the discriminator might not generalize well to abstract states outputted by the forward dynamics functions.
Therefore, we train the discriminator with the model-based rollout.
Speci ﬁ cally, we sample sequence ( s t , a t +1 , ..., a t + n ) in replay buffer B and expert sequence ( s Et ′ , a Et ′ , a Et ′ +1 , ..., a Et ′ + n ) in demonstration dataset D and minimizes following multi-step discriminator loss function:
[Equation]
Here, h t + i (and h t ′ + i ) terms are produced by the forward dynamics in Ef ﬁ cientZero (Figure 2).
We use the GAIL transition reward R ( h, a )
=
− log(1 − D ( h, a )) , and then the MCTS planner searches for action that can maximize cumulative GAIL reward to guarantee long-term distribution matching.
Note that V-MAIL also propose a similar discriminator training technique, but under the
Dreamer [11] model.
Besides, since the discriminator’s input is based on the representation rather than raw input, the discriminator should be trained with the encoder jointly.
However, this can lead to a highly nonstationary reward during the bootstrap value calculation.
To mitigate this issue, we also propose to use a target discriminator for bootstrap value calculation.
This can make the training more stable.
Though we use the GAIL reward here, one can also use other kinds of AIL and IRL reward functions proposed in recent research.
Using the GAIL reward can already achieve satisfactory performance in our experiments.
When the real reward presents, one may also combine this into planning [21].
This may favor application scenarios where handcrafting a reward function is not hard.
We do not study this case here and leave it to future work.
4.2 Unifying BC and AIL in MCTS
As discussed in related work, researchers realize that using BC can improve AIL’s sample ef ﬁ ciency by providing a good initial BC policy or constraining the policy to BC.
However, these existing solutions are not good enough in practice.
The main pitfall in these methods is that BC knowledge in the policy network is destined to be forgotten if the policy network is trained with the AIL objective, and then BC will no longer be helpful [34].
We observe that MCTS can naturally unify the BC and AIL methods, enjoying the bene ﬁ t of both and being free from this pitfall.
We propose to plug BC actions into MCTS as candidates at each node and use a planning process to search for an improved solution.
This time, the BC actions are consistently considered throughout the entire training procedure without being forgotten.
Concretely, we train a BC policy π BC and use a mixture policy ˜ π for the sampling at each node in MCTS:
[Equation]
) α is a mixture factor, which is ﬁ xed during training and π is the current policy.
We use α = 0 .in
this paper.
This ensures that a small fraction of action samples are from the BC policy.
During planning, the BC actions are evaluated and will be frequently visited and selected as output if they can lead to long-term distribution matching.
This can then reduce the effort of ﬁ nding good expert-like actions from scratch as desired.
Moreover, another unique advantage of this procedure is that it does not fully trust BC like [20], which forces the output of the policy network to be close to BC.
When
BC is wrong due to covariate shifts or insuf ﬁ cient demonstrations, it can neglect these BC actions and allow the policy to search for better actions.
This ensures that BC does not hurt training.
However, due to the conceptual simplicity, one arising question is whether this approach can be applied to other model-based methods.
Here, we take Dreamer [11] as an example.
Though Dreamer builds a model of the environment, it only uses the model to roll-out the policy for policy optimization.
In other words, the model is not used to evaluate whether a speci ﬁ c BC action is good or not in the long term, so our idea can not be applied directly to Dreamer.
From this example, we see that the core of our idea is to leverage the power of planning, only with which the long-term outcomes of certain (BC) actions can be calculated.
For the training of π BC , we minimize the following multi-step BC objective (Figure 2):
[Equation]
This is to avoid distributional shifts during multi-step prediction in MCTS.
For the training of the policy, we still minimize D KL ( [Equation] | [Equation]
Note that the BC design proposed here is not coupled with AIL.
It can go beyond imitation learning and be applied in other robot learning settings, such as RL with demonstration [38].
4.3 Implementation
We ﬁ rst implement a continuous version Ef ﬁ cientZero for planning, and the details can be found in the Appendix.
The BC policy network is simply a duplicate of the policy network.
The discriminator and BC policy networks share the same encoder network with the policy network and value network.
The overall loss function for optimization is
[Equation]
L EZ is Ef ﬁ cientZero’s loss function (excluding reward loss).
All the networks are trained jointly to minimize this loss function 7.
We use the Reanalyze algorithm [43, 52] for of ﬂ ine training, and we require that all the samples should be reanalyzed.
Experiments
In this section, we evaluate the sample ef ﬁ ciency of the proposed method.
We measure the sample ef ﬁ ciency by evaluating the performance of an algorithm at a small number of online samples.
We also analyze the effect of the BC actions and planning. [Table] [Table] [Table]
[Table]
ValueDICE
[Table]
SQIL
[Table] [Table]
5.1 Setup
We use the DeepMind Control Suite [47] for evaluation.
We use the following tasks: Cartpole
Swingup, Reacher Easy, Ball-in-cup Catch, Finger Spin, Cheetah Run, Walker Walk, Hopper Hop, and Humanoid Walk.
We conduct both state-based and image-based experiments.
Note that many previous imitation learning works use the OpenAI Gym [2] version of these tasks for evaluation.
We ﬁ nd that the DMControl version used here brings extra challenges by using more challenging initial states.
Take the Walker task as an example; the initial state in OpenAI Gym is standing.
However, in
DMControl, the agent’s initial state is lying on the ground, and the agent should also learn to stand up ﬁ rst from very limited data.
For the state-based experiments, we allow 10k-50k online steps in the environment based on the dif ﬁ culty of each task.
Since learning a robust and meaningful visual representation requires more data for image-based experiments, we allow 50k-100k online steps.
Detailed setup will be shown in the result.
We train SAC [9] policies to collect expert demonstrations for imitation learning.
The expert demonstrations are not subsampled.
We use 5 demonstrations in the state-based experiment, except for Reacher and Humanoid, where we use 20 demonstrations.
We use 20 demonstrations in the image-based experiments.
5.2 Baselines
We present several baselines of sample ef ﬁ cient imitation learning.
(1) DAC DAC [23] is an adversarial off-policy imitation learning method.
It matches the distribution of the replay buffer and that of the expert demonstration dataset using the TD3 [7] algorithm.
(2) SQIL SQIL [39] is a non-adversarial off-policy imitation learning method.
It labels all the expert transitions with reward 1 and non-expert transitions with reward 0.
Then it trains a SAC policy over these relabeled data.
SQIL is a regularized form of BC.
(3) ValueDICE ValueDICE [24] considers imitation learning as a distribution matching problem and solves it with a min-max optimization process.
(4) VMAIL
VMAIL [37] is a model-based visual imitation learning method.
It learns a variational model for simulating on-policy rollouts.
We only evaluate VMAIL on the image-based domain, as they did in
� ��� ���� ���� ���� ����� ��� ��� ��� ��� ��� ��� � � � � � � � � � � � ������ � ��� ���� ���� ���� ����� ��� ��� ��� ��� ��� ��� ������� � ��� ���� ���� ���� ����� ��� ��� ��� ��� ��� ��� ������ � ��� ���� ���� ���� ����� ��� ��� ��� ��� ��� ��� ������ ����������������������
� ��� ���� ���� ���� ����� ��� ��� ��� ��� ��� ��� � � � � � � � � � � � ������ � ��� ���� ���� ���� ����� ��� ��� ��� ��� ��� ������� � ��� ���� ���� ���� ����� ��� ��� ��� ��� ��� ��� ������ � ��� ���� ���� ���� ����� ��� ��� ��� ��� ��� ��� ��� ��� ������ ������������������ ��������� the original paper.
Besides the online imitation learning baselines, we also include BC as an of ﬂ ine baseline.
5.3 Results
State-based experiments Table 1 shows the state-based experiments’ evaluation results within the given budget.
We also plot the performance curve of four challenging tasks in Figure 4.
The performance is normalized to 0.0 and 1.0 with respect to the performance of the random agent and the expert.
We ﬁ nd that our proposed method can outperform all the baselines by a large margin.
Except for BC, these baselines methods could hardly learn meaningful behaviors using a limited online budget.
We ﬁ nd that DAC is a strong baseline.
Its performance can grow to the expert in 200k samples in most tasks except Hopper, where it will eventually get stuck.
Our method is much more sample ef ﬁ cient than the best of these baseline methods.
For some tasks like Walker Walk and
Cheetah Run, our method only requires about 20k steps to reach near-expert performance, equivalent to 80 online trajectories (around 0.5 hours in real).
This result is notable for the robotics community.
It shows that online imitation learning is possible with only a handful of trials, and applying it directly on a real locomotion robot is possible.
Image-based experiments So far, image-based tasks are still challenging for adversarial imitation learning algorithms, and the evaluation of most of the prior AIL works is carried out in the state-based tasks.
Table 2 shows the evaluation result within the given budget in the image-based experiments (see
Figure 5 for curves).
Our method can also learn expert behavior given a slightly larger budget.
Still, most of the baselines fail to match experts’ behavior using the given budget.
We notice an inherent dif ﬁ culty in learning a robust visual representation for adversarial training in the limited data set in image-based tasks.
Discriminator can judge whether a behavior is expert-like using various possible features in this case.
Solving this open problem is out of the scope of this paper.
In the presence of such a dif ﬁ culty, EI can still achieve good sample ef ﬁ ciency in most of the tasks.
� ���� ���� ����� ��� ��� ��� ��� ��� ��� ��� ��� ��� � � � � � � � � � � � �������� � ��� ��� ����� ��� ��� ��� ��� ��� ��� ������ � ��� ��� ����� ��� ��� ��� ��� ��� ��� ������� � ��� ��� ����� ��� ��� ��� ��� ��� ��� ������ � ��� ��� ����� ��� ��� ��� ��� ��� ��� ������ �������������������������
5.4 Analysis
Effect of BC We carry out ablation analysis on the BC to see whether it helps in our method.
We set α = 0 to remove all the BC actions and see how our method’s performance and sample ef ﬁ ciency will change.
The result is shown in Figure 6.
We ﬁ nd that the performance of our method degrades after we remove all the BC actions in MCTS.
The effect is task-speci ﬁ c, and we ﬁ nd BC is more helpful in those challenging tasks.
For example, in tasks that are high-dimensional or have a hard exploration process like Humanoid and Hopper, removing BC actions will make the learning process stuck in local minima.
At that local minima, the agent struggles to ﬁ nd the correct action that matches the expert’s behavior.
Although removing BC actions does not trap the agent in local minima in some relatively simpler tasks like Cheetah and Walker, it slows down the training.
It doubles the number of online interactions to reach expert performance.
This result con ﬁ rms that using BC actions can indeed provide a good solution to the distribution matching problem, which can help to speed up learning and improve performance.
We also notice that even when we remove the BC actions, the method is still able to outperform the previous baselines; this suggests that planning with AIL alone is also powerful.
� ���� ���� ���� ���� ���� ����� ��� ��� ��� ��� ��� � � � � � � � � � � � �������� �����������������������������������������������������
Other Ways to use BC We then study another two variants of using BC: (1).
BC-Ann .
This variant does not use BC actions in MCTS but exerts an annealing BC loss to the policy network like [20].
BC-Rep .
This variant does not use BC actions in MCTS but still uses BC loss to regularize the representation.
We test these variants on the Humanoid Walk (Figure 7).
We ﬁ nd that these variants do not lead to an essential improvement.
For BC-Ann, it harms the performance in the early stage (before 100k) since the BC regularization will constrain the agent’s policy near the BC policy, which contains an error and hinders learning.
The agent only starts to acquire meaningful behavior after the regularization decays, but at that time, BC does not help much and can not lead to improvement.
Compared with BC-Ann, BC-Reg is more helpful here.
This is possibly because
BC-Reg makes the encoder focus on more useful features.
However, BC-Reg still gets stuck in a local minimum.
This result suggests that using BC actions directly for exploration can be essential for improving AIL.
Using BC simply as a regularizer may not be the ideal approach though it can be useful sometimes.
Ablation of Planning In this part, we study to what extent planning can help to learn and whether insuf ﬁ cient search in MCTS leads to inferior performance.
We study K , the number of sampled actions at each node, and N , the number of simulations.
The default value of K and N in the previous experiments are 16 and 50 .
We sweep K ∈ { 4 , 8 , 16 , 24 } and N ∈ { 5 , 10 , 25 , 50 } to evaluate their effects.
We collect the result on the state-based Cheetah, Walker, and Hopper task and report the averaged relative performance change at the given budget used in previous experiments (see Table 3).
The general trend is that larger K and N lead to better imitation results.
We ﬁ nd that varying K only affects the performance a little, and K = 4 can also work well.
Compared with K , N has a larger impact.
When the number of simulations becomes small, the performance drops signi ﬁ cantly.
This result also explains why we can achieve a large improvement over DAC even without BC.
[Table]
Discussion
In this paper, we presented Ef ﬁ cientImitate, an MCTS-based imitation learning algorithm.
We extended AIL to a model-based setting and solved it with MCTS in a sample-ef ﬁ cient way.
Moreover, we proposed a method to unify BC and AIL in MCTS, enjoying the bene ﬁ ts of both.
Experimental results in state-based and image-based tasks showed that Ef ﬁ cientImitate can achieve state-of-the-art sample ef ﬁ ciency and performance.
Limitations One limitation of this work is that the computation process of MCTS is more expensive compared with that of the model-free methods, though this is a common issue of model-based methods.
One possible approach to mitigate this issue can be using better MCTS acceleration methods [51].
Besides, in this paper we did not study the long horizon problem with multiple objects, which is a common case in the robotic manipulation.
However, this requires the model to predict the interaction with multiple objects, which is still a challenging open problem in the learning community [36] and orthogonal to our contribution.
We believe that our framework can be combined with the works in this ﬁ eld to handle this challenge.
Future Work There are many problems to study along our direction.
First, since we only use the vanilla
AIL algorithm here, it is interesting to see if using more advanced algorithms such as optimal-transport-based learning [5] will make our algorithm more powerful.
Second, due to the modularity of our method, one can try to extend Ef ﬁ cientImitate to more general settings like RL with demonstration, which will also favor the application scenarios.
Third, in this work we consider an online learning setting, one possible future direction is to study the use of Ef ﬁ cientImitate on the existing of ﬂ ine interaction dataset to further reduce the dependence on in-environment samples.
In conclusion we believe that this work shows a promising direction and opens up new possibilities for model-based methods in robot learning.
Acknowledgments and Disclosure of Funding
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 Arti ﬁ cial Intelligence” (Grant No.
2021AAA0150000).
This work is also supported by a grant from the Guoqiang Institute, Tsinghua University. A Implementation
In this part, we introduce our detailed implementation of Ef ﬁ cientImitate. Our code is available at https://github.com/zhaohengyin/EfficientImitate . A.1 Model Details Our model is based on the Ef ﬁ cientZero [52] model.
It is composed of several neural networks: representation network f , dynamics network g , value network V , policy network π , BC policy network π BC , discriminator network D , projector network, and predictor network.
We use PyTorch [35] to implement these networks.
Their detailed structures are as follows.
A.1.1 State-based experiments
Representation Network
The representation network is an MLP with one hidden layer of size 256.
Its output dimension is 128.
It uses LeakyReLU
[30] as the hidden layer’s activation function.
The output activation function is identity.
Dynamics Network
The dynamics network is an MLP with one hidden layer of size 256.
It concatenates the representation and the action as the input.
Its output dimension is 128 (representation’s dimension).
It uses LeakyReLU
as the hidden layer’s activation function.
The output activation function is identity.
Value Network
The value network is an MLP with one hidden layer of size 128 (256 for Humanoid).
It uses LeakyReLU
as the hidden layer’s activation function.
The output activation function is identity.
We use the categorical value representation introduced in MuZero [42].
The value prediction is discretized into 401 bins to represent the value between [ − 100 , 100] .
Therefore, the output dimension of value network is 401.
Policy & BC Network
The policy and the BC policy networks are both MLPs with one hidden layer of size 128 (256 for Humanoid).
These MLPs use LeakyReLU as the hidden layer’s activation function.
Their output activation functions are identity.
They produce a Squarshed-Normal (TanhNormal) action distribution [9], which is determined by a predicted mean and a predicted logstd.
Therefore, the output dimension is twice of the dimension of the action space.
Discriminator Network
The discriminator network is an MLP with one hidden layer of size 128.
It uses LeakyReLU as hidden layer’s activation function.
The output activation function is sigmoid.
Projector Network The projector network is an MLP with one hidden layer of size 512.
It uses
ReLU [30] as the hidden layer’s activation function.
The output activation function is identity.
The output dimension (projection dimension) is 128.
Predictor Network
The predictor network is an MLP with one hidden layer of size 512.
It uses
ReLU as the hidden layer’s activation function.
The output activation function is identity.
The output dimension is 128.
A.1.2 Image-based experiments
Representation Network
The representation network consists of a four layer convolutional neural network and an MLP.
Its structure is as follows.
• Convolution layer.
Input dim: 12.
Output dim: 32.
Kernel Size: 5.
Stride: 2.
Padding: 2. • ReLU activation.
• Convolution layer.
Input dim: 32.
Output dim: 32.
Kernel Size: 3.
Stride: 2.
Padding: 1. • ReLU activation.
• Convolution layer.
Input dim: 32.
Output dim: 32.
Kernel Size: 3.
Stride: 2.
Padding: 1. • ReLU activation.
• Convolution layer.
Input dim: 32.
Output dim: 32.
Kernel Size: 3.
Stride: 2.
Padding: 1. • Flatten.
• Linear layer.
Output dim: 128. • ReLU activation.
Dynamics Network
The dynamics network is an MLP with one hidden layer of size 256.
It concatenates the representation and the action as the input.
Its output dimension is 128 (representation’s dimension).
It uses ReLU as the hidden layer’s activation function.
The output activation function is also ReLU.
Value Network
The value network is an MLP with two hidden layers of size 100.
It uses ReLU as the hidden layer’s activation function.
The output activation function is identity.
We use the categorical value representation introduced in MuZero.
The value prediction is discretized intobins to represent the value between [ − 40 , 40] .
Therefore, the output dimension of value network is 401.
Policy & BC Network
The policy and the BC policy networks are both MLPs with two hidden layers of size 100.
These MLPs use ReLU as the hidden layer’s activation function.
Their output activation functions are identity.
They produce a Squarshed-Normal action distribution, which is determined by a predicted mean and a predicted logstd.
Therefore, the output dimension of them is twice of the dimension of the action space.
Discriminator Network
The discriminator network is an MLP with one hidden layer of size 100.
It uses ReLU as hidden layer’s activation function.
The output activation function is sigmoid.
Projector Network The projector network is an MLP.
Its structure is as follows.
• Linear layer.
Output dim: 1024.
BatchNorm [18] with momentum 0.1.
ReLU activation.
• Linear layer.
Output dim: 1024.
BatchNorm with momentum 0.1.
ReLU activation.
• Linear layer.
Output dim: 1024.
BatchNorm with momentum 0.1.
Predictor Network
The predictor network is an MLP.
Its structure is as follows.
• Linear layer.
Output dim: 512.
BatchNorm with momentum 0.1.
ReLU activation.
• Linear layer.
Output dim: 1024.
A.2 MCTS Details
Our MCTS implementation is mainly based on the Sampled MuZero [17].
We also apply modi ﬁ cations proposed by the Ef ﬁ cientZero.
The detailed procedure is as follows.
Expansion For the task having a continuous action space, we can not enumerate all the possible actions at a node as the original MCTS algorithm.
To solve this problem, we use the sampling method proposed by the Sampled MuZero.
For the expansion of a node s (in the representation space), we sample K actions { a i } Ki =1 from current policy π ( a | s ) .
In this work we propose to integrate BC actions into MCTS, then we actually sample from
[Equation]
Here α = 0 .
25 is a mixture factor.
Selection For the action selection, we selects action a ∗ from the sampled actions that maximize the probabilistic upper con ﬁ dence bound
[Equation]
where ˆ π ( a | s )
= 1 K
P i δ ( a, a i ) .
Q ( s, a ) is the current Q -estimation of the pair ( s, a ) .
N ( s, a ) denotes the times that this pair is visited in MCTS.
c ( s ) is a weighting coeffcient de ﬁ ned by
[Equation]
where c 1 = 1 . 25 , c 2 = 19625 .
To encourage exploration, we also inject Dirichlet noise to ˆ π ( a | s ) at the root node.
So ˆ π ( a | s ) becomes
[Equation]
Here, ρ = 0 .
25 is a mixture factor.
N D ( ξ ) is the Dirichlet distribution, and ξ is set to 0.3.
At the root node of MCTS, we use the discriminator network D and value network V to calculate
Q -value by its original de ﬁ nition: Q ( s, a ) =
R ( s, a ) + γ V ( g ( s, a )) .
At the other nodes, we use the mean Q -value calculation used by the Ef ﬁ cientZero.
Simulation & Backup The simulation and the backup process is the same as Ef ﬁ cientZero’s implementation, and we refer the readers to Ef ﬁ cientZero for details.
A.3 Training Details and Hyperparameters
Finally, we introduce some important training details and the hyperparameters.
Initialization We initialize the weights and biases of the last layer of policy, BC policy, value, and discriminator network to be zero.
The other parameters are initialized by the default initializers in
PyTorch.
Discriminator
In AIL training, it is very useful to apply the gradient penalty to the discriminator [34].
We also apply gradient penalty to the discriminator network.
Target Update We propose to use a target model for the calculation of policy, value, and AIL reward during reanalyze.
The target model is updated periodically during training subject to an update frequency.
The training hyperparameters used in the state-based experiments are in Table 4.
The training hyperparameters used in the image-based experiments are in Table 5.
B Environment Details B.1 State-based experiments
The setup of each task in the state-based experiments is in Table 6.
B.2 Image-based experiments
The setup of of each task in the image-based experiments is in Table 7.
We use a 48 × 48 resolution in the image-based experiments.
C Other Ablations
We also perform ablations on the target discriminator and the multi-step discriminator loss.
To evaluate the effect of the target discriminator, we use the latest model to calculate AIL reward in the value target.
To evaluate the effect of the multi-step discriminator loss, we replace it with the single-step discriminator loss.
We conduct experiments on the state-based and image-based Walker and Cheetah.
The results are shown in Figure 9.
We ﬁ nd that removing these components will not only lead to instability in training but also harm the performance.
Compared with the target discriminator, the multi-step discriminator loss has a larger impact on the image-based tasks.
D Computation Resources
All of our experiments are conducted on a server with 4 NVIDIA RTX 3090 GPUs, 64 CPU cores, and 256GB RAM.
For the most of state-based and image-based experiments except Humanoid Walk and image-based Hopper Hop, our experiments require 12-18 hours of training.
The main bottleneck is at the Reanalyse [43, 52], where the minibatch cannot be produced and sent to the training loop at a high frequency.
We are improving the computation ef ﬁ ciency by using better parallel computation implementation and applying MCTS speed up techniques.
E Visualization
One approach to interpret the learned model is by the t-SNE [48] plot.
We use the image-based
Walker experiment as an example.
We use the trained model at 100k env steps to generate the state embeddings of one expert trajectory, and in environment trajectory at 0k, 25k, 50k, 75k, and 100k steps.
Then we use t-SNE to visualize the embeddings on the 2D plane.
As is shown in the Figure 8, the agent’s trajectory gradually matches expert’s trajectory (blue) during training.
Moreover, the expert’s trajectory has a circle structure, which represents the periodic pattern of the Walker’s walking behavior.
Therefore, our model can represent the environment in a meaningful way.
[Figure]
��� ��� ��� ��� ��� � � � � � � � � � � � ��������������� � ��� ��� ��� ��� ��� ����� ��� ��� ��� ��� ��� ��� �������������� � ��� ��� ��� ��� ��� ����� ��� ��� ��� ��� ��� ��� ��������������� � ��� ��� ��� ��� ���� ����� ��� ��� ��� ��� ��� ��� �������������� ���������������������������������
[Table]
[Table]
[Table]
[Table]