Learning Efficient Planning-based Rewards for Imitation Learning
Authors: Xingrui Yu
Yueming Lyu
Ivor Tsang
Venue:  (2021)

Abstract: Imitation learning from limited demonstrations is challenging. Most inverse reinforcement learning (IRL) methods are unable to perform as good as the demonstrator, especially in a high-dimensional environment, e.g, the Atari domain. To address this challenge, we propose a novel reward learning method, which streamlines a differential planning module with dynamics modeling. Our method learns useful planning computations with a meaningful reward function that focuses on the resulting region of an agent executing an action. Such a planning-based reward function leads to policies with better generalization ability. Empirical results with multiple network architectures and reward instances show that our method can outperform state-of-the-art IRL methods on multiple Atari games and continuous control tasks. Our method achieves performance that is averagely 1,139.1% of the demonstration. 

Under review as a conference paper at ICLR 2021
LEARNING EFFICIENT PLANNING-BASED REWARDS FOR IMITATION LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Imitation learning from limited demonstrations is challenging. Most inverse rein- forcement learning (IRL) methods are unable to perform as good as the demon- strator, especially in a high- nigh-dimenssonal environment, e.g, the Atari domain. To address this challenge, we propose a novel reward learning method, which stream- lines differential planning module with dynamics modeling. Our method learns useful planning computations with a meaningful reward function that focuses on the resulting region of an agent executing an action. Such planning-based reward function leads to policies with better generalization ability. Empirical results with multiple network architectures and reward instances show that method can outperform state-of-the-art IRL methods on multiple Atari games and continuous control tasks. Our method achieves performance that is averagely 1,139. % of the demonstration.
INTRODUCTION
Imitation learning (IL.) offers an alternative to reinforcement learning (RL) for training an agent, which mimics the demonstrations of an expert and avoids manually designed reward functions. Behavioral cloning (BC) (Pomerleau 1991) is the simplest form of imitation learning, which learns a policy using supervised learning. More advanced methods, inverse reinforcement learning (IRL) (Ng & Russell 2000 Abbeel & Ng. 2004) seeks to recover reward function from the demonstrations and train an RL agent on the recovered reward function. In the maximum entropy variant of IRL, the aim is to find a reward function that makes the demonstrations appear near-optimal on the principle of maximum entropy (Ziebart et al. 2008 2010, Boularias et al. 2011 Finn et al. 2016).
However, most state-of-the-ar IRL. methods fail to meet the performance of demonstrations in high- dimensional environments with limited demonstration data, e-g., a one-life demonstration in Atari domain (Yu et al. 2020). This is due to the main goal of these IRL approaches is to recover reward function that justifies the demonstrations only. The rewards recovered from limited demonstration data would be vulnerable to the overfitting problem. Optimizing these rewards from arbitrary initial policy results in inferior performance. Recently, Yu et al. (2020) proposed generative intrinsic reward learning for imitation learning with limited demonstration data. This method outperforms expert and IRL. methods in several Atari games. Although GIRIL uses the prediction error as curiosity to design the surrogate reward that encourages (pushes) states away from the demonstration and avoids overfitting, the curiosity also results in ambiguous quality of the rewards in the environment.
In this paper, focus on addressing the key issues of previous methods when learning with limited demonstration data, i.e., 1) overfitting problem, and 2) ambiguous quality of the reward function. To address these issues, we propose to learn a straightforward surrogate reward function by learning to plan from the demonstration data, which is more reasonable than the previous intrinsic reward function (i.e., the prediction error between states). Differential planning modules (DPM) potentially useful to achieve this goal, since it learns to map observation to a planning computation for a task, and generates action predictions based on the resulting plan (Tamar et al.) 2016 Nardelli et al.) 2019 Zhang et al. 2020). Value iteration networks (VIN) (Tamar al.) (2016) the representative one, which represents value iteration as a convolutional neural network (CNN). Meaningful reward and value maps have been learned along with the useful planning computation, which leads to policies that generalize well to new tasks. However, due to the inefficiency of summarizing complicated transition dynamics, VIN fails scale up to the Atari domain.

---

Under review as a conference paper at ICLR 2021
To address this challenge, we propose novel method called variational planning-embedded learning (vPERL), which is composed of two submodules: module and the transition dynamics module. We leverage variational objective based the conditional variational autoencoder (VAE) (Sohn et al. 2015) to jointly optimize the two submodules, which greatly improves the generalization ability. This is critical for the success of achieving straightforward and smooth reward function and value function with limited demonstration data.
embedded reward
action fack-tracing
As shown in Figure 1 vPERL learns meaningful reward and value maps that attends the resulting region of the agent executing an action, which indicates meaningful planning computation. However, directly applying VIN in Atari domain in the way of supervised learning (Tamar al., 2016) only learns reward and value maps that attend no specific region, which usually results in no avail.
(a) (b) VPERL Maps.
(c)
Maps.
Figure 1: Visualization of state, reward and value maps of vPERL and VIN on Battle Zone game (the first row) and Breakout game (the second row).
Empirical results show that method outperforms tatte-of-the-ae IRL methods multiple Atari games and continuous control tasks. Remarkably, our methods achieve performance that is up to 58 times of the demonstration. Moreover, the average performance improvement of our method is 1,139. 1 %% of the demonstration over eight Atari games.
2 BACKGROUND AND RELATED LITERATURE
Markoy Decision Process (MDP) 1966) is a standard model for sequential decision making and planning. An MDP M is defined by a tuple (S, A, T, 7), where S is the set of states, A is the set of actions, T x x - IR., is the environment transition distribution, R S IR is the reward function, and 7 E (0, 1) is the discount factor (Puterman 2014). The expected discounted return or value of the policy и is given by (T(s) = E, 7 R[s a s0 = s], where (So; ao, $1, ..) denotes the trajectory, in which the actions selected according 80 ~ To(Sp). a ~ (a))st), and 8141 ~ T($+11|8e az). The goal an MDP to find the optimal policy n² that enables the agent obtain high long-term rewards.
Generative Adversarial Imitation Learning (GAIL) (Ho & Ermon 2016) extends IRL by inte- grating adversarial training technique for distribution matching (Goodfellow al.! 2014). GAIL performs well in low-dine www-dimensional applications, c.g., MuJoCo. However, it does not scale well to high-dimensional scenarios, such Atari games (Brown et al. 2019a). Variational adversarial imitation learning (VAIL) (Peng et al. 2019) improves on GAIL by compressing the information via variational information bottleneck. GAIL and VAIL inherit problems of adversarial training, such as instability in training process, and are vulnerable to overfitting problem when learning with limited demonstration data. We have included both methods as comparisons to vPERL in our experiments.
Generative Intrinsic Reward driven Imitation Learning (GIRIL) YYu et al. 2020 leverage gen- erative model to learn generative intrinsic rewards for better exploration. Though GIRII outperforms previous IRL methods on several Atari games, the reward map of GIRIL is ambiguous and less informative, which results in inconsistent performance improvements in different environments. In contrast, our vPERL learns efficient planning-base based reward that is more straightforward and informative. We have included GIRIL as competitive baseline in our experiments.
Differentiable planning modules perform end-to-end learning of planning computation, which leads to policies that generalize to new tasks. Value iteration (VI) (Bellman, 1957) is well-known method for calculating the optimal value V* and optimal policy x": Vn+1 max, Qn(s,a) a).
2

---

Under review as a conference paper at ICLR 2021
where (s) = = value function Vn in VI converges as n - do to V*, from which the optimal policy may derived as n (s) arg max, 20 (s,a). a).
denotes the Q value in the nth iteration. The
Value iteration networks (VIN) (Tamar ct al. 2016) proposes to embed value iteration (VI) (Bellman 1957) process with recurrent convolutional network, and generalizes well in conventional navigation domains. VIN assumes there is some unknown embedded MDP M where the optimal plan in M contains useful information about the optimal plan in the original MDP M. VIN connects the
two MDPs with a parametric reward function π. = [*)()). Nardelli et al. (0119) proposes value propagation networks (VPN) generalize VIN for better sample complexity by employing value propagation (VProp). Recently, universal value iteration networks (UVIN) extends VIN to spatially variant MDPs (Zhang et al., 2020). Although VIN can be extended to irregular spatial graphs by applying graph convolutional operator (Niu et al. 2018), most of the VIN variants still focus on solving the conventional navigation problems (Zhang et al. 2020).
In this paper, we extend differentiable planning module to learn efficient reward function for imitation learning on limited demonstration data. We dig more on leveraging the learned reward function for imitation learning; while previous related work of VIN focuses more the value function. Therefore, our work is complementary to the research of VIN and its variants. Note that any differentiable planning module can be embedded in our method. As a simple example, utilize the basic VIN as a backhone to build our reward learning module.
VARIATIONAL PLANNING-EMBEDDER REWARD LEARNING
In this section, we introduce our solution, variational Hanninngemmeedee reward learning (vPERL). As illustrated in Figure 2
our reward learning module is composed of two submodules to accomplish planning-embeddee action back-tracing and explicit forward transition dynamics modeling.
vPERL
Vilta
fvi
q+
R
Action Backetrccng
fR
fa
S
St
G2
Transition Modeling
fv
Pre Vile
Po
f'v
Figure 2: Illustration of the proposed vPERL model.
3.1 ACTION BACK- TRACING AND FORWARD DYNAMICS MODELLING IN VPERL
Planning-embedded action back-tracing. Instead of directly applying VIN for policy learning (Tamar et al. 2016), we build our first submodule (ax 814-) for action back-tracing. As illustrated in the top section of Figure 2 embedded MDP M, where fR is a convolutional layer. A VI module fyz takes in the reward map R. and effectively performs K times of VI by recurrently applying convolutional layer Q for K times (Tamar et al. 2016). The Q layer is then max-poot max-pooled to obtain the next-iteration value V. The right-directed circular arrow in a light-blue color denotes the direction of convolutions. Then, we simply obtain the action from the intermediate optimal value V by an action mapping function: a, = (V"). On these terms, we build our planning-embedder action back-tracing submodule, which is formally represented (a, (8t, St-1) fa(fvr(Jn(s) 8t41))). Since the
we first obtain the reward map R = fR(S,, $141) on an
3

---

Under review as a conference paper at ICLR 2021
convolutional kernel is incapable of summarizing the transition dynamics in complex environment, directly training this submodule is still insufficient for learning efficient reward function and planning computation in an environment like Atari domain.
Explicit transition dynamics modeling via inverse VI. To address this, we further build upon another submodule po po(s+11 8t) for explicit transition dynamics modeling. We build the submodule based on the inverse VI module, which is a NN architecture that mimics the process of the inverse version of VI. The implementation of the inverse VI module is straightforward. We first map the action for the intermediate optima value in another embedded MDP M' by a value mapping function: V fv az). Then, we apply the inverse VI module to obtain the reward map R'. The inverse VI module fvr takes in the intermediate value V and recurrently apply a deconvolutional layer Q' for K times on the value to obtain the reward map R'. The left-directed circular arrow in purple color denotes the direction of deconvolutions. To accomplish the transition, we map the obtained R' to the future state by: 8t-1 = fs+((R), The transition modeling is therefore presented as po(st+1 la,, 8t) = (fv (fv. a+))), which is a differentiable submodule, and can be trained simultaneously with the action back-tracing submodule.
Variational solution to vPERL. A variational autoencoder (VAE) (Kingma & Welling 2013) can be defined as being autoencoder whose training is regularised to avoid overfitting and ensure that the latent space has good properties that enable generative process. To avoid the learned planning- based reward overfitting the demonstration, we optimize both submodules in unified variational solution, which follows the formulation of conditional VAE (Sohn et al.] 2013). Conditional VAE conditional generative model for structured output prediction using Gaussian latent variables, which is composed of conditional encoder, decoder and prior. Accordingly, we regard the action back wrek-tracin module q [2] St, St+1) as the encoder, po p0($++1 2, 8t) as the decoder, and po (2) (s(ss)) as the prior. Our vPERL module is maximized with the following objective:
L(s,, (141,8441:40 11 0, ()- (xle.. po (0 [z, 8c)] (L Se+ (1l0000(())) -aKLLgg (g (āelse, 81+1 (
(1)
where = is the latent variable, T (at is the expert policy distribution, à, = (2) is the
transformed latent variable, a is a positive scaling weight. The first two terms on the RHS of Eq. (T)
in the first line denote the evidence lower bound (ELBO) of the conditional VAE (Sohn et al. 2015)) These two terms are critical for our reward learning module to perform planning- -baiid action back- tracing and transition modeling. Additionally, we integrate the third term on the RHS of Eq. (I) in the second line to further boost the action back-tracing. The third term minimizes the KL divergence between the expert policy distribution Te (a; (8x) and the action distribution go (āe 18t, 81+1), 1), where à = Softmax(z)) == is transformed from the latent variable 2. In this way, we train the forward state transition and action back-tracing simultancously.
Algorithm 1 Imitation learning via variational
reward learning (vPERL).
1: Input: Expert demonstration data D = {(s, a
2: Initialize policy and the dual planning networks.
3: for e 1,... E do
4: Sample a batch of demonstration D ~ D.
5: Train vPERL module on D to converge.
6: end for
7: for 1,
-1,...,MMXITER MAXITER do
8: Update policy via any policy gradient method, e.g., PPO on the learned surrogate reward r.-
9: end for
Output: Policy π.
Note that the full objective in Eq. (1) is still variational lower bound of the marginal likelihood 1og(pe(s1+11))
Accordingly, it is reasonable to maximize this as an objective of our reward learning module. By optimizing the objective, we improve the forward state transition and action back-tracing. As a result, our reward learning module efficiently models the transition dynamics of the environment. During training, we use the latent variable calculate the surrogate rewards from the learned reward map. As shown in Figure [ our method leamss meaningful reward map, which highlights the resulting region of agent executing one action.
the intermediate action. After training,
will
4

---

Under review as a conference paper at ICLR 2021
To leverage such meaningful information, we calculate two types of rewards that both correspond to the highlighted informative region, i..,, 7; = RMax = max and 7, = RMeau = meanR, which uses the maximum and mean value of the reward map R. respectively.
Algorithm summarizes the full training procedure of imitation learning via vPERL. The process begins by training vPERL module for E epochs (steps 3-6). In each training epoch, we sample mini-batch demonstration data D with mini-batch size of B and maximize the objective in Eq. (1) Then in steps 7-9, we update the policy и via any policy gradient method, e.g., PPO (Schulman et al. 2017), so to optimize the policy with the learned surrogate reward function rt-
EXPERIMENTS
4.1
ATARI GAMES
We first evaluate our proposed vPERL on one-life demonstration data for eight Atari games within OpenAl Gym (Brockman et al.] 2016). To enable fair comparison, we evaluate our method and all other baselines under the same standard setup, where we train an agent to play Atari games without access to the true reward function (Ibarz et al. 2018 Brown et al. 2019a). The games and demonstration details are provided in Table
A one-life demonstration only contains the states and actions performed by an expert player until they lose their life in game for the first time (Yu et al.] 2020). In contrast, one full- episode demonstration contains states and ac- tions after the expert player loses all available lives in game. Therefore, the one-life demon- stration data is much more limited than one full- episode demonstration. We define three levels of performance: 1) basic one- life demonstration- level ggmmplay up to one life lost ("one-life"), 2) expert-level gameplay up to all-lives lost ("one full-episode"), and 3) beyond expert "better-than-experr performance.
Table 1: Statistics of Atari environments.
Demonstration Length
# Lives
Game
One-life
Full-episode
available
Kung-Fu Master
1.167
3,421
4
Battle Zone
260
1,738
5
Centipede
166
663
3
Seaquest
562
2,252
4
Qtbert
787
1,881
4
Breakout
1,577
2,301
5
Beam Rider
1,875
4,587
3
Space Invaders
697
750
3
Our ultimate goal to train an imitation agent that can achieve a better-than-expers performance with the demonstration data recorded up to the moment of losing their first life in the game.
Demonstrations To generate one-life demonstrations, we trained a PPO (Schulman et al. 2017) agent with the ground-truth reward for 10 million simulation steps. We used PPO implementation with the default haper--arrmene parameters in the repository (Kostrikov) 2018). As Table I shows, the one-life demonstrations are all much shorter than the full full-episode demonstrations, which make for extremely limited training data.
train reward learning module for each game on the one- life demonstration. We set = 10 in vPERL for all of the Atari games. By default, we use a neural network architecture that keeps the size of the reward map and value maps the same as that of the state, which is 84 x 84. We achieve this by using a convolutional kernel of size for each convolutional layer, and applying padding. The corresponding method is called 'VPERL-Large', Additionally, to enable faster learning, we implement our method with another neural network architecture that reduces the size of the reward map and value maps into 18 18. The corresponding method is called 'vPERL-Small" Both vPERI-Large and vPERI-Small can learn meaningful reward map well as useful planning computation. Training was conducted with the Adam optimizer (Kingma Ba 2015) at learning rate of 3e-5 and mini-batch size of 32 for 50,000 epochs. In each training epoch, we sampled a mini-batch of data every four states.
Experimental Setup Our first step
To evaluate the quality of our learned reward, we trained policy to maximize the inferred reward function via PPO. We set a = 100 for training our reward learning module. We trained the PPO on the learned reward function for 50 million simulation steps to obtain final policy. The PPO is trained with a learning rate of 2.5e-4, a discount factor of 0.99, a clipping threshold 0. an entropy coefficient of 0.01, value function coefficient of 0.5, and GAE parameter of 0.95 (Schulman et al..
5

---

Under review as a conference paper at ICLR 2021
2016
We compared imitation performance by our vPERL agent against VIN, two state-of ooatte-of-hh-arr the-hhterr inverse reinforcement learning methods, GAIL (Ho & Ermon 2016) and VAIL (Peng et al. 2019). More details of setup outlined in Appendix F2
Results In Figure 3 we report the performance by normalizing the demonstration performance to Figure 3 shows that vPERL achieves performance that is usu-
PPERL- Lar.-Larg-eaaa Mange-Raan
Normalized Performanes
50-
PPERL-Small-RRu
VPIRL Small-----
CIRIL
ally close or better than that of the expert demonstrator. The most im- pressive one the Centipede game, our vPERL achieves performance that is around 60 times higher than the demonstration. GIRIL achieves the second best performance in Centipede, beating the demonstra- tion by around 30 times. On the Qbert game, vPERL beats all other baselines by a large margin, achiev- ing performance that more than 15 times of the demonstration.
40
VAIL
GALL
Expert
30
Demonstration
20
Kang Pl Mavter
BatleZone
Contiedls
Suupaant
ophent
Breakoil
Deam Rider
Speed Invadars
Figure 3: Performance improvement of vPERL and baselines.
Table 2: Average return of vPERL with two types network architectures (Large and Small) and two types of rewards (RMax and RM GIRIL YYu et al. 2020), VIN and state-of-the-an IRL algorithms GAIL (Ho & Ermon 2016) and VAIL (Peng et al. (019) with one-life demonstration data on eight Atari games. The results shown are the mean performance over five random seeds with setter-than-exper improvements of IL algorithms versus demonstration (Demo.) and expert performance with greater than 100% in bold.
performance in bold. The last two rows show the average performance
Expert
Deno.
PPO-50M
.PERL-Large
sPERL- Small
Imilation Learingg Algorithoss
Game
Average
Aserage
Average
Raxux
Rateaa
Raaus
Ratea
GIRIL
VAIL
GAIL
Kung-Fu Fu Master
23,434.8
6,500.0
38,58000
20,701.0
18,900.
20,500.0
20.518.3
26,742.6
11,751.2
1.3
190.4
Battle Zone
18,137.5
7,000.0
33,400.0
27,500.0
11.000.0
47,200.0
21,899.1
7.372.8
14,222.7
6,200.0
3,372.0
Centipele
4.069.4
360.0
7,701.6
15,516,9
10,448.8
17.371.4
21,113,4
11,446.5
3,936,1
1.306.5
1,907.5 5
Seaguest
1.744.5
1,920.0
706.0
756.0
1.604.2
624.8
476.1
60.0
148.0
Ofbert 13.441.5
8,150.0
26,235.0
71,920.0
165.325.0
58,965,0
135.598.0
13,070.3
8,314.8
25.0
283.4
Breakont
346.4
305.0
$22.9
374.7
420.0
422.1
34888
430.3
370.6
2,0
0.4
Beam Rider
2.447.7
1,332.0
5.146.2
3,264,6
4.590.0
2.960.5
3,8018
1,599.5
1,710.8
286.5
351,4
Space Invaders
734.1
600.0
2,880.
915.2
1,135.0
1,325.5
1,106.5
556.6
646.2
113.9
Average
345.8%
100.%
823.1%
780.5%
913,8%
1.139.1%
544.0%
255.6%
63.3%
85.3%
Improvements (%)
100.0%
47.8%
194.9%
267.9%
210.1%
278.5%
103.9%
72.1%
12.2%
15.2%
A detailed quantitative comparison of IL algorithms is listed in Table 2 We have evaluated four variants of vPERL with two types of network architectures (Large and Small) and surrogate rewards (RMax and R. Both vPERI-Large and VPERI-Small can learn meaningful reward and value maps as well as useful computation in the Atari domain. In Appendix D and E we visualize the learned reward and value maps of vPERL-Large and vPERL-Small, respectively. With such meaningful rewards, the four variants of vPERL outperform the expert demonstrator in six out of eight Atari games. Remarkably, vPERL-Small with RMean achieves an average performance that is 1. 1,139.1 % of the demonstration and 278.5% of the expert over the eight Atari games. Figure 3 shows the bar plot of normalized performance of four vPERL variants against other baselines.
Table shows that VIN is far from achieving demonstration-leww performance, since it is unable to learn useful planning computation as well as the meaningful reward and value maps in the Atari domain. GAIL fails to achieve good imitation learning performance. VAIL manages to exceed the expert performance in one game, i.e., Breakout. GIRIL performs better than previous methods, outperforming the expert demonstrator in three games. The results show that vPERL agent outperforms the expert demonstrator by a large margin in six out of eight games. Figure 4 shows the qualitative comparison of our method SPERR-Smaal-RReess average performance of the expert and demonstration. Additionally in Appendix C.I] our method consistently outperforms expert and other baselines on two more Atari games, Krull and Time Pilot.
GIRIL, VAIL, GAIL, VIN and the

---

Under review as a conference paper at ICLR 2021
VPERL-Sman Ryten
GIRIL
VAIL
GAIL
VIN
Expert
Demonstration
23000
13001
12:01
2000
13X03
590
Kung-Fu Master.
(b) Battle Zone.
(c) Centipede.
(d) Seaquest.
20000
Avple
12
3000
10003
25090
(c) Qhber..
(f) Breakout.
Beam Rider.
(h) Space Invaders.
Figure 4: Average return the number of simulation steps Atari games. The solid lines show the mean performance over five random seeds. The shaded area represents the standard deviation from the mean. The blue dotted line denotes the average return of the expert. The area above the blue dotted line indicates performance beyond the expert.
4.1.1 How DOES vPERL OUTPERFORM PREVIOUS METHODS?
The contributions of each component of vPERL. In this subsection, we study the contribution of each component of our method, i.e. Action tracing submodule, Transition Modeling submodule, directly train the Action Back-tracing and
and the variational objective in Eq. (T) Specifically, Transition Modeling submodules in terms of supervised learning. We used the mean of the reward map and the prediction error of the next state the reward for the former and latter submodules, respectively. To study the contribution of the variational objective, we introduced another baseline, PERL, which trains both submodules as an autoencoder. Table 3 shows quantitative comparison between the performance of vPERL and its components.
Table Average return of vPERL-Small- RMMan and components (i.e., Action Back-tracing and Transition Modeling, and PERL) with one-life demonstration data on eight Atari games. The results shown are the mean performance over five random seeds with better- than than-expert-tett performance in bold.
Expert
-PERI-Small-Raae
Composents VPHRI Smail--SSmll-RRe. Rylons
Game
Average
Average
Avenge
Action Back-tracing
Transition Modeling
PERI.
Supervised PHRI.
Kung- Fu Master
23,434,8
6,500.0
20,518.3
0.0
0.0
0.0
0.0
Buttle Zone
18,137.5
7,000.0
21,899.1
6000.0
0.0
2,000.0
Centipede
4,069.4
360.
21,113.4
151.0
2,216.7
16,479.6
2,485.1
Seaquest
1,744,5
440.0
1,631,2
0.0
0.0
140.0
$0.0
Obeet 13,441.5
8,150.9
135,59880
74,9
150.0
0.0
0.0
346.4
Beam Rider
2,447.7
1,332.0
3,801.8
22.0
784.8
17.6
0.0
Space Invaders
734.1
600.0
1,106.5
285.0
145.5
285.0
0.0
The results show that individual training of each component results in avail. PERL successfully outperforms the demonstration in one game, i.e. Centipede, which indicates the potential advantage of using both submodules. However, PERL fails the other seven games, while vPERL outperforms the demonstration in eight games and outperforms expert in six. The large performance gap between PERL and vPERL indicates the variational objective in Eq. ①1 is important to learn efficient rewards. To further investigate the key reason for why our method works well, we added another baseline supervised PERL, which forces the encoding of PERL to be close to the true action. The supervised PERL fails in all of the games. Comparing with vPERL, we can attribute the critical contribution to the use of the ELBO of conditional VAE, or more specific, the term (L(go (2) (t, 8t+1) ())s in
Eq. (T). It helps vPERL work well and outperform previous methods for two reasons:
1) The generative training of VAE can serve as a good regularization to alleviate the overfitting problem.
7

---

Under review as a conference paper at ICLR 2021
2) The regularization enables vPERL to learn smooth value function and reward function, which consistently provides straightforward and informative rewards for the moving states in the environ- ment.
Empirical evidence:
1) Better generalization ability. The empirical results in Table and Figure show that VIN, GAIL and VAIL are vulnerable to overfitting problem, usually results in no avail and has fewer chances to reach the demonstration--eeel performance. In contrast, our vPERL has better generalization ability and consistently achieves performance that is either close to or better than the expert.
2) Straightforward and informative reward. Figure 5 shows and GIRIL in three Atari games. The reward map of GIRIL can be close to zero (in Battle Zone) and state (in Q*bert) or occasionally informative (in Centipede), which is ambiguous and less informative. In contrast, the reward map of our vPERL is more straightforward, and consistently attends to informative regions the state for all of the games.
shoow the state, the reward maps of vPERL
Sale
VEERI newerd
GIBIL rewertt
Battle Zane
(b) Centipadee
(c) Qrhatt
Figure 5: Visualization of state, vPERL reward map and GIRIL reward map for Atari games.
Our method successfully addresses the two key issues, therefore, it can outperform previous methods.
4.2 CONTINUOUS CONTROL TASKS
We also evaluated our method on continuous control tasks where the state space is now-dimensional and the action space is continuous. The continuous control tasks were from Pybullet Sobullerrenvinoeen environment.
Demonstrations To generate demonstrations, we trained a Proximal Policy Optimization (PPO) agent with the ground-truth reward for 1 million simulation steps. We used the PPO implementation in the repository (Kostrikov) 2018) with the default hyper-parameters for continuous control tasks. In each task, used one demonstration with fixed length of 1,000 for evaluation. The details of experimental setup can be found in Appendix F.1
Results Figure 6 shows that our method vPERL achieves the best imitation perfor- mance in both continuous control tasks, i.e. Inverted Pendulum and Inverted Double Pendulum. Although GIRIL achieves per- formance that is close to the demonstration, efficient planning-base
PPERL
VAIL
GATL
VIN
Demonstration
14:00
reward fune- tion enables vPERL to perform even bet- ter. Other baselines are unable to reach the demonstration-levè performance by learn- ing from only one demonstration. Quanti- tative results are shown Appendix A
(a) Inverted Pendulum.
(b) Inverted Double Pendulum.
Figure 6: Average return vs. number of simulation steps continuous control tasks.
4.3 ABLATION STUDIES
Ablation study on the growing number of demonstrations. Figure shows the average return versus number of full-episode demonstrations on both Atari games and continuous control tasks. The results shows that our method vPERL achieves the highest performance across different numbers of full-episode demonstrations. GIRIL usually comes the second best, and GAIL can achieve good performance with more demonstrations in continuous control tasks. Quantitative results have been shown in Appendix B.I
https//pyballett.rrg
8

---

Under review as a conference paper at ICLR 2021
20003
Aestag Netum
ANNIAA Returm
300
0000
19853
5000
Capart Droweddddee
u822
12388
8000
200
Demonstrations
(a) Centipede.
(b)
Qbert.
(c) Inverted Pendilum
(d) Inverted Double Pendulum.
Figure 7: Average return vs. number of demonstrations on Atari games and continuous control tasks.
Ablation study on the optimality of the demonstrations. Figure 8 shows the average return versus optimality of demonstrations on Atari games and continuous control tasks. In sections 4.1 and 4.2 we trained PPO agents with ground-truth reward for 10 million (10M) steps as the expert for Atari games, and for 1 million (1M) steps as the expert for continuous control tasks. In this ablation, we train PPO agents with 10% and 50% simulation steps of the expert to generate demonstrations with diverse optimality. The results show that vPERL consistently outperforms the expert and demonstrations the demonstrations of different optimality, Quantitative results are shown Appendix B.2
2.000
bboogg Rotur
aF31 Dencrstutinn
2009
10
00
sopert emarttio (millonn))
daps (millona)
Number scpert staps (nallonal
Lumber expert eleps Imillano
(a)
Comipede.
(b)
(c) Inverted Pendulum.
(d) lnverted Double Peadulum.
Figure 8: Average return vs. number of the expert simulation steps (Ng) on Atari games and
continuous control tasks.
Ablation study on the hyper-parametery K. Figure 8 shows the average return versus different choices of K on Atari games and continuous control tasks. The results show that our method not very sensitive to different choices of K. Quantitative results are shown in Appendix B.3
3000
Average Rotrr
152009
AsseRON Raauum
VMELL
123088
800
Uemonammtion
2000
Cholce
Choice of
Choice of
Cholce cfK
Comppade.
(b)
(c) Inverted Pondalurs.
(d) Inverted Double Posdulum.
Figure 9: Average return VS. choices of K on Atari games and continuous control tasks.
CONCLUSION
This paper presents simple but efficient reward learning method, called variational planning- embedded reward learning (vPERL). By simultaneously training a planning-embedded action back- tracing module and transition dynamics module in a unified generative solution, we obtain reward function that is straightforward and informative, and has better generalization ability than previous methods. Informative analysis and empirical evidence support the critical contribution of ELBO regularization term for learning efficient planning-based reward with extremely limited demonstrations, Empirical results show our method outperforms state-of-the-arn imitation learning methods multiple Atari games and continuous control tasks by large margin. Extensive ablation studies show that our method is not very sensitive to the number of demonstrations, optimality of demonstration. and choices of the hyperparameter K. We remain the extension of method to more complex continuous control tasks as future work. Another interesting topic for future investigation would be applying vPERL hard exploration tasks with extremely sparse rewards.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference on Machine Learning, pp- 1, 2004.
Richard Bellman. A markovian decision process. Journal of mathematics and mechanics, pp. 679-684, 1957.
Richard Bellman. Dynamic programming. Science, 153(3731):34-37, 1966.
Abdeslam Boularias, Jens Kober, and Jan Peters. Relative entropy inverse reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pp. 182-189, 2011.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint 1o0ivv:/00.00154000 2016.
Daniel S Brown, Wonjoon Goo, Prabhat Nagarajan, and Scott Nickum. Extrapolating beyond suboptimal demonstrations via inverse reinforcement learning from observations. International Conference Machine Learning, 2019a.
Daniel S Brown, Wonjoon Goo, and Scott Niekum. Lanking-basede reward extrapolation without rankings. In Conference Robot Learning, 2019b.
Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. In International Conference Machine Learning, pp. 49-58, 2016.
Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems, pp. 2672-2680, 2014.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in Neural Information Processing Systems, pp. 4565-4577 2016.
Borja Ibarz, Jan Leike, Tobias Pohlen, Geoffrey Irving, Shane Legg, and Dario Amodei. Reward learning from human preferences and demonstrations in atari. In Advances Neural Information Processing Systems, pp. 801 1-8023, 2018.
Animesh Karnewar. Pytorch implementationss of variational discriminator bottleneck, 2018. URL max htttp:://qithin com /tl akan imx /Vrri at ional 1_riacr___iiii mi nator Baciioo___tt_t___ 1estllercck
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International
Conference on Learning Representations, 2015.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes.
arXiv preprint
arXiv:1312.6114, 2013.
Ilya Kostrikov. Pytorch implementatiions of reinforcement leaming algorithms, 2018. URL https /tt.hhabcccco///sssssrroo/roorrr
acktr gaall
Nantas Nardelli, Gabriel Synnaeve, Zeming Lin, Pushmeet Kohli, Philip HS Torr, and Nicolas Usunier. Value propagation networks. In International Conference on Learning Representations, 2019.
Andrew Ng and Stuart Russell. Algorithms for inverse reinforcement learning. In International Conference Machine Learning, volume 1, pp. 663-670, 2000.
Sufeng Niu, Siheng Chen, Hanyu Guo, Colin Targonski, Melissa C Smith, and Jelena Kovačević. Generalized value iteration networks: Life beyond lattices. In AAAI Conference on Artificial Intelligence, 2018.
Xue Bin Peng, Angjoo Kanazawa, Sam Toyer, Pieter Abbeel, and Sergey Levine. Variational discriminator bottleneck: Improving imitation learning, inverse rl, and gans by constraining information flow. In International Conference Learning Representations, 2019.
10

---

Under review as a conference paper at ICLR 2021
Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural Computation, 3(1):88-91) 1991.
Martin L Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. John Wiley & Sons, 2014.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. In International Conference on Learning Representations, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint
2017.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In Advances in neural information processing systems, pp. 3483-3491, 2015.
Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. In Advances in Neural Information Processing Systems, pp. 2154-2162, 2016.
Xingrui Yu, Yueming Lyu, and Ivor Tsang. Intrinsic reward driven imitation learning via generative model. In International Conference on Machine Learning, pp. 6672-6682, 2020.
Li Zhang, Xin Li, Sen Chen, Hongyu Zang, Jie Huang, and Mingzhong Wang. Universal value iteration networks: When spatially-invariant is not universal. In AAAI Conference on Artificial Intelligence, pp. 2154-2162, 2020.
Brian D Ziebart, Andrew Maas, J Andrew Bagnell, and Anind K Dey. Maximum entropy inverse reinforcement learning. In AAAI Conference on Artificial Intelligence, 2008
Brian D Ziebart, Andrew Bagnell, and Anind K Dey. Modeling interaction via the principle of maximum causal entropy. In International Conference Machine Learning, 2010.

---

Under review as a conference paper at ICLR 2021
A QUANTITATIVE RESULTS OF CONTINUOUS CONTROL TASKS.
Table 4 shows the detailed quantitative comparison of the demonstration and imitation methods. The results shown in the table were the mean performance over three random seeds.
Table Average return of vPERL, GIRIL, VIN and state-of of-tee the arts inverse reinforcement learning algorithms GAIL (Ho & Ermon 2016) and VAIL (Peng al. 2019) with one demonstration data on continuous control tasks. The results shown are the mean performance over random seeds with best imitation performance in bold.
Demonstration
vPERL
Imitation Learning Algorithms
Task
Average
Ratean
GIRIL
VAIL.
GAIL
VIN
Inverted Pendulum
1,000.0
1,00000
993.2
520.1
572.2
21.0
Inverted Double Pendulum
9,355.1
9,357.8
539.8
181.1
B
ABLATION STUDIES
B.1 THE EFFECT OF THE NUMBER FULL-EPISODE DEMONSTRATIONS.
We also evaluated our method with different numbers of full-episode demonstrations on both Atari games and continuous control tasks. Table and 6 the detailed quantitative comparison of imitation learning methods across different numbers of full-episode demonstrations in the Centipede game and Qbert game. The comparisons on two continuous control tasks, Inverted Pendulum and Inverted Double Pendulum, have been shown in Table and Table g
The results shows that our method vPERI achieves the highest performance across different numbers of full-episode demonstrations, and GIRIL usually comes the second best. GAIL is able achieve better performance with the increase the demonstration number both continuous control tasks.
Table 5: Parameter Analysis of the vPERL versus other baselines with different numbers of full- episode demonstrations on Centipede game. The results shown are the mean performance over random seeds with best performance in bold.
# Demonstrations
vPERL
GIRIL
VAIL
GAIL
13,253.2
9,195.1
4,212.8
171.0
5
21,219.2
2,952.5
2,465.6
161.0
10
14,406.2
8.599.6
3.108.0
2.232.4
Table 6: Parameter Analysis of the vPERL versus other baselines with different numbers of full- episode demonstrations on Qbert game. The results shown are the mean performance over 5 random seeds with best performance in bold.
# Demonstrations
vPERL
GIRIL
VAIL
GAIL
200,455.0
11,105.0
13,695.0
0.0
5
200,450.0
13,740.0
13.300.0
0.0
10
192,962.5
12,327.5
10,852.5
0.0
Table 7: Parameter Analysis of the vPERL versus other baselines with different numbers of full- episode demonstrations on Inverted Pendulum task. The results shown are the mean performance over 5 random seeds with best performance bold.
Demonstrations
vPERL
GIRIL
VAIL
GAIL
1,000.0
993.2
520.1
$72.2
1,000.0
1.00.00
68.2
591.2
10
1,000.0
1,000.0
41.8
1,000.0
12

---

Under review as a conference paper at ICLR 2021
Table 8: Parameter Analysis of the vPERI. versus other baselines with different numbers of full- episode demonstrations on InvereedDoublePondullum task. The results shown are the mean perfor- mance over 5 random seeds with best performance in bold.
Demonstrations
vPERL
GIRIL
VAIL
GAIL
9,357.8
9,184.5
539.8
$19.6
9,352.4
9.317.2
372.5
9.342.5
10
9,353.1
9,314.1
274.6
9,352.1
B.2 THE EFFECT OF EXPERT OPTIMALITY.
Table and Table 10 show the average return of vPERL-Small- with demonstrations of different optimality on Atari games and continuous control tasks, respectively. In the experiments, we trained PPO agent with ground-truth reward for 10 million (10M) simulation steps as the expert for Atari games, and million (IM) steps for continuous control tasks. To study the effects optimality of the demonstrations, 5M steps for Atari games, and 0.IM steps and 0.5M steps for continuous control tasks. With the additional PPO agents, generated demonstrations with 10%, 50% optimality of the 10M-step 'Expert' for both Atari games and continuous control tasks.
additionally trained PPO agents with less simulation steps: 1M steps and
The results show that method outperforms expert by large margin in Atari games and reach the demonstration-leww performance in continuous control tasks with demonstrations of different optimality.
Table 9: Average return of -PERL-Small-Reaan with one-life demonstration data of expert policy trained under different simulations steps (Ng=l million, 5 million and million). The results shown the mean performance over five random seeds with better-thand thaert performance in bold.
Expert
Demonstration
VPERL-Small-Ria
Game
Average
Average
Ng=IM
Ny=5M
Ng=10M
Battle Zone
18,137.5
7,000.0
24,800.0
34,897.2
21,899.1
Centipade
4.069.4
360.0
16.343.7
17,223.3
21,113.4
Qteert
13,441.5
8,150.0
147,565.5
200,30000
135,598.0
Beam Rider
2,447.7
1,332.0
3,374.4
3,256.0
3,801.8
Table 10: Average return of vPERL- with one demonstration data of expert policy trained under different simulations steps (N g=.1 million, 0.5 million and I million). The results shown the mean performance over five random seeds with demonstration performance bold.
Demonstration
IPERL-RAs.
Game
Average
Ng=0.1M
Ng=0.5M
Ng=IM
Inverted Pendulum
1,000
1,000.0
1,000.0
1,000.0
Inverted Double Pendulum
9,355.1
9,351.9
9,350.8
9,357.8
B.3 THE EFFECT OF CHOICES OF K.
For the sake of consistency, we set = 10 for experiments on Atari games and continuous control tasks in Section 4 To study the effects of the hyperparameter K, evaluate our method on two Atari games and two continuous control tasks with two additional K (K=5, and K=15).
Table I1 and Table 12 shows the average return of VPERL-Small-RR K on Atari games and continuous control tasks. With the three choices of K, our method consistently outperforms the expert in the Atari games, and reach the best (demonstration-Leve) performance in continuous control tasks. This indicates that our method not very sensitive to the choices of hyperparameter K.
versus different choices of
13

---

Under review as a conference paper at ICLR 2021
Table 11: Average return of -PERL-Small-Raaee with different choices of K, one-life demonstra- tion data. The results shown are the mean performance over five random seeds with netter-than-exper performance in bold.
Expert
Demo.
<PERL- Small- Small-Ratton
Game
Average
Average
K=5
K=10
K=15
Contipede
4,069.4
360.0
19,228.8
21,113.4 135,598.0
18,776.2
Qabert
13,441.5
8,150.0
46,705.5
173.735.0
Table Average return of -PERL-RMenn with different choices of K, on one-life demonstration data. The results shown are the mean performance over five random seeds with demonstration-level performance bold.
Demonstration
vPERL- Rattn
Game
Average
K=5
K=10
K=15
Inverted Pendulum
1,000)
1,000.0 9,358.8
1,000.0
1,000.0
Inverted Double Peadulum
9,355.1
9,357.8
9,351.9
C ADDITIONAL EVALUATION RESULTS
C.I ADDITIONAL ATARI GAMES.
Table 13 shows the average return of vPERL-Small with RMean and other baselines on two additional Atari games, Krull and Time Pilot. The results show that our method outperforms the expert and other baselines by a large margin on both additional Atari games.
Table 13: Average return of vPERL- Small with Rxlean GIRIL Yu et al. 2020). VIN and state- of-the-art IRL algorithms GAIL (Ho & Ermon 2016) and VAIL (Peng et al. 2019) with one-life
demonstration data on additional Atari games. The results shown are the mean performance over five random seeds with than---than-expert performance bold.
Expert
Demo.
SPERL-Small
Imitation Learning Algorithms
Game
Average
Averuge
Ratean
GIRIL
VAIL
GAIL
VIN
Krull
8,262.0
2,826.0
19,412.5 4,400.0
6.515.0
5,819.7
264.0
2,013.3
Time Pilot
4,200.0
100.0
3,680.0
2,080.0
400.0
2,572.0
C.2 ONE-LIFE DEMONSTRATION WITHOUT SCORES AND LIVES ON THE STATES.
To avoid the effects the scores and lives the states of Atari games, we also evaluate our method the "No-score Demo.", which is obtained by masking the game score and number of lives left in the demonstrations (Brown et al., (019a).. Table [4]] compares the average return of vPERL-Small -- with the "Standard Demo." and the "No-score Demo." on e*bert game and Krull game.
Table 14: Average return of vPERL-Small- RMtan with different choices of K on one-life demonstra- tion data. The results shown are the mean performance over five random seeds with better-than-exper performance in bold.
Expert
Demo.
PERL--Sall--FE Small-Fateon
Game
Average
Average
Standard Demo.
Demo.
Krull
8,262.0
2,826.0
19,412.5 135,598.0
27,009.6
Qbbert
13,441.5
8,150.0
200,500.0
The results show that method achieves better performance on the "No-score Demo." than the "Standard Demo.". This indicates the negative effects of the game scores and numbers of left lives on the states of demonstrations. From Figure I and more reward visualization in Section D and we observe that our method learns to attend on the meaningful region in a state and ignore the game score and numbers of left lives automatically. Masking the game score and numbers of left
14

---

Under review as a conference paper at ICLR 2021
lives in the demonstration further alleviates burdens on learning efficient planning computations and planning-bel based rewards for Atari games.
In summary, method can learn to outperform the expert without explicitly access to the true rewards, and does not relied on the game scores and numbers of left lives in the states of demonstra- tions. Furthermore, the results show that the performance of our method can be improved by masking out the game scores and numbers left lives in the demonstrations.
D VISUALIZATION OF REWARD AND VALUE IMAGES OF VPERL-LARGE AND VIN.
In this section, visualize the reward maps and value maps learned by vPERL-Large and VIN on Atari games. Here, both vPERL and VIN are based on large-size VIN architecture. The size of reward map is 84 84. The figures show that the reward and value maps learned by vPERL are much meaningful than that by VIN.
(a) State.
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Visualization of state, reward map and value map on Battle Zone game. (a) The state, (b) the reward map and value map of vPERI,, and (c) the reward map and value map of VIN.
Figure 10:
(a) State.
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Figure 11: Visualization of state, reward map and value map on Centipede game. (a) The state, (b) the reward map and value map of vPERL, and (c) the reward map and value map of VIN.
(a) State.
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Figure Visualization state, reward map and value map on Seaquest game, (a) The state, (b) the reward map and value map of vPERL, and (c) the reward map and value map of VIN.
15

---

Under review as a conference paper at ICLR 2021
(a) State.
(b) vPERI. Reward and Value.
(c) VIN Reward and Value.
Visualization of state, reward map and value map on Qbert game. (a) The state, (b) the reward map and value map of vPERL, and (c) the reward map and value map of VIN.
Figure 13:
1:34 1
(a) State.
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Figure 14: Visualization state, reward map and value map on Breakout game. (a) The state, (b) the reward map and value map of vPERI, and (c) the reward map and value map of VIN.
(a) State
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Figure 15: Visualization of state, reward map and value map Space Invaders game. (a) The state, (b) the reward map and value map vPERL, and (c) the reward map and value map of VIN.
16

---

Under review as a conference paper at ICLR 2021
E VISUALIZATION OF REWARD AND VALUE IMAGES OF VPERL-SMALL AND VIN.
In this section, visualize the reward maps and value maps learned by vPERL and VIN on several Atari games. To enable faster training, here both vPERL and VIN are based on small-size VIN architecture. The size of reward map 18 18. The figures show that the reward and value maps learned by vPERL are much meaningful than that by VIN.
(a) State
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Visualization state, reward map and value map on Kung-Fu Master game. (a) The state, (b) the reward map and value map of vPERL, and (c) the reward map and value map of VIN.
Figure 16:
do
(a) State.
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Figure 17:
Visualization of state, reward map and value map on Battle Zone game. (a) The state, (b) the reward map and value map of vPERL, and (c) the reward map and value map of VIN.
(a) State
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Figure 18: Visualization of state, reward map and value map on Centipede game. (a) The state, (b) the reward map and value map of vPERL, and (c) the reward map and value map of VIN.
(a) State
vPERL Reward and Value.
(c) VIN Reward and Value.
Figure 19: Visualization of state, reward map and value map Seaquest game. (a) The state, (b) the reward map and value map of vPERL, and (c) the reward map and value map of VIN.
17

---

Under review as a conference paper at ICLR 2021
(a) State
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Visualization of state, reward map and value map on Qbert game. (a) The state, (b) the reward map and value map of vPERL, and (c) the reward map and value map of VIN.
Figure 20:
(a) State.
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Visualization of state, reward map and value map Breakout game. (a) The state, (b) the reward map and value map of vPERI, and (c) the reward map and value map of VIN.
Figure 21:
(a) State
vPERL Reward and Value.
VIN Reward and Value.
Figure 22: Visualization of state, reward map and value map Space Invaders game. (a) The state, (b) the reward map and value map vPERL, and (c) reward map and value map of VIN.
(a) State
(b) vPERL Reward and Value.
(c) VIN Reward and Value.
Figure 23: Visualization of state, reward map and value map on Beam Rider game. (a) The state, (b) the reward map and value map of vPERL, and (c) the reward map and value map of VIN.
ADDITIONAL DETAILS OF EXPERIMENTAL SETUPS
F. EXPERIMENTAL SETUP OF CONTINUOUS CONTROL TASKS
Our first step was also to train a reward learning module for each continuous control task on one demonstration. To build reward learning module for continuous tasks, used simple VIN and inverse VIN as the model bases of action back -raccin and transition modeling submodules, respectively. In the simple VIN model, used ID convolutional layer with kernel size of 2 and
18

---

Under review as a conference paper at ICLR 2021
stride of I to implement the function fn. reward map R and Q value Q. To accomplish the action back-tracing, the final value map of VIN was fully connected with a hidden layer with a size of 32. Reversely, used ID deconvolutional layer to implement the inverse VIN model. We kept the size of feature maps in both VIN and inverse VIN unchanged across all the layers. We set K' = 10 for both the VIN and inverse VIN in all tasks. The dimension of latent variable is to the action dimension for each task. Additionally, we used a two-layer feed forward neural network with tanh activation function as policy architecture. The number of hidden unit is set to 100 for all tasks, To extend our method on continuous control tasks, we made minor modification on the training objective. In Atari games, we used the KL divergence to measure the distance between the expert policy distribution and the action distribution in Eq. ( In continuous control tasks, we instead directly treated the latent variable as the back-traced action and used mean squared error to measure the distance between the back-traced action and the true action in the demonstration. We set the scaling weight in Eq. (T to 1.0 for all tasks. Training was conducted with the Adam optimizer (Kingma & Ba 2015) at learning rate of 3e-5 and mini-batch size of 32 for 50, 000 epochs. In each training epoch, sampled a mini-batch of data every 20 states.
To evaluate the quality of our learned reward, used the trained reward learning module to produce rewards, and trained a policy to maximize the inferred reward function via PPO. We trained the PPO on the learned reward function for 5 million simulation steps obtain final policy. The PPO is trained with learning rate of 3e-4, a clipping threshold of 0.1, entropy coefficient of 0.0, a value function coefficient of 0.5, and GAE parameter of 0.95 (Schulman et al. 2016).
For fair comparison, we used the same VIN as the model base for all the baselines. The reward function of GAIL and VAIL was chosen according the original papers (Ho & Ermon 2016 Peng et al. 2019). The information constraint Le in VAIL was set to 0.5 for all tasks. To enable fast training we trained all the imitation methods with 16 parallel processes.
F.2 ADDITIONAL DETAILS OF GAIL AND VAIL
The discriminator for both GAIL and VAIL takes in state (a stack of four frames) and an action
(represented as a 2d one-hot vector with shape of (AA) x 84 x 84), where [AA] is the number of valid discrete actions in each environment) (Brown al. 2019b). The discriminator outputs binary classification value, and - log D(s, a)) is the reward. VAIL was implemented according to the repository of Karnewar (2018). The discriminator network architecture has additional convolutional layer (with kernel size of 4) as the final convolutional layer to encode the latent variable in VAIL. We used the default setting of 0.2 for the information constraint (Karnewar) 20188) PPO with the same hyper-parameters was used to optimize the policy network for all the methods. For both GAIL and VAIL, we trained the discriminator using the Adam optimizer with a learning rate of 0.001. The discriminator was updated at each policy step.
19


Reviews:
[{"review_id": "wdlLyoaGROH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a inverse RL method that utilizes a differential planning module and dynamics modeling to learn reward signals. The learned reward function can be used to train policies that can match or even exceed demonstration, as demonstrated in several Atari games and two simple continuous control tasks.\n\nThe paper demonstrates a novel usage of VIN like planning module for inverse RL and the result is quite compelling. Although the writing can and should be improved to help reader to understand what is going on, including multiple grammatical errors like \" we introduce a variational solution to optimize the both submodules jointly\". \n\nSome technical questions:\n(1)$p_\\theta$ takes $z$ and $s_t$ as input, but in equation (1), it also only takes $s_t$ as input. I am a bit confused here.\n(2) As far as I understand, the learned transition dynamics is only meaningful for the expert action, for example, given $s_t$, if the $a_t$ is completely different from the expert action, it is likely to predict completely wrong $s_{t+1}$, and thus completely wrong reward estimate with the setup in the paper. Since the policy is able to learn much better with the learned reward function, I will imagine it learns to take action that is completely different from the demonstration. Any intuitive explanation why this will work so well? It will be also interesting to see how the learned reward compare to the true reward in those region not covered by the demonstration.\n(3) The continuous control tasks is like a toy example, but I guess the impressive part is that it can learn so well with only 1000 demonstration data. But are there any challenges to scale this to more complicated tasks given enough demonstration? If not, potential limitation should be discussed.", "extra_scores": {}, "extra_text": {}}, {"review_id": "lXOBFeMkZQZ", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper addresses the problem of imitation learning. The authors\npropose a system where you first learn a surrogate reward function\nfrom the demonstrations, and then you do RL to optimize that reward.\nThe surrogate reward function is the result of learned forward and\ninverse dynamics models that are \"planning-embedded\" in the same\nsense as value iteration networks.\n\nI think this paper is quite interesting overall, and the idea seems\nworth pursuing. I know that the general approach to imitation learning\nof first learning a cost function from the demonstrations and then\napplying RL to optimize it has a very rich history, but I have not\nseen prior work that learns the cost function in the manner proposed\nin this paper. Additionally, the experimental results seem good to me,\nespecially on the chosen Atari games. I think the result that vPERL\nachieves performance far greater than that of the demonstrations, in\nsome instances, will be of interest to the community. Therefore, I am\nleaning toward recommending acceptance, but I am not familiar enough\nwith the scope of related literature to be certain.\n\nI do have some questions which I hope the authors can address in the rebuttal.\n\n1. What is special about Kung-Fu Master and Breakout that causes vPERL\nto not perform well? Have the authors tried Atari games other than the\n8 shown? Why were these 8 picked?\n\n2. Regarding Figure 1, I'm a bit unclear about the interpretation. A\nvalue function is a function mapping states to values; what does it\nmean to have a \"value map\" for a single state? Is the idea that these\nare spatial domains where an agent is moving around, and so the map\nyou're showing us implicitly represents the value of a state where the\nagent is at that location?\n\n3. After Equation 1, you say that \"\\pi_E(a_t | s_t) is the expert\npolicy distribution.\" But where do we get this distribution from? We\nonly have a dataset of expert behavior, and states are continuous, so\nwe will never see the same state twice in the data. So how can we\nbuild a distribution over what actions the expert would take from any\ngiven state? Will this distribution just be a Dirac in practice?\n\n4. I am curious how important the use of both forward and inverse\nmodels is in the reward learning module. The authors provide some\njustification why both are needed, but it would be good to support\nthis with an ablation experiment.\n\n5. How sensitive is the approach to the optimality of the\ndemonstrations? That is, in your experiments, you generate the\ndemonstrations from a policy trained on 10 million PPO steps; what if\nthis number of steps were lower?\n\n6. To me, an obvious baseline would be to run PPO with the exact\nconfiguration that you use in your experiments (with 50 million\nsimulation steps), but using the ground truth reward instead of the\nlearned reward. I am curious if the authors have considered this\nbaseline, and how it performs.", "extra_scores": {}, "extra_text": {}}, {"review_id": "cdHV8ibZuy", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a method for inverse reinforcement learning that incorporates a differential planning module. Explicit transition dynamics modeling with inverse value iteration is added to promote meaningful reward learning. Empirical evaluations on several high-dimensional Atari environments and 2 continuous control environments are provided which show improvements over existing inverse reinforcement learning baselines when given only one-life demonstrations. Some visuals are also presented to show that the proposed method is able to learn more meaningful reward maps than previous methods.\n\nPositives:\n1. The proposed method is natural and makes use of a few existing works. The combination of Value Iteration Networks (VIN) and the Generative Intrinsic Reward driven Imitation Learning (GIRIL) is novel. \n\n2. The visual reward maps show that the proposed method learns a more semantically meaningful reward map than VIN. When using PPO on the learned reward function, it learns a better performing policy faster in most cases.\n\n3. The results in continuous control environments demonstrate generalizability.\n\nNegatives:\n1. There is little analysis of the learned reward maps beyond a few plots for visual comparisons. As the proposed method bears close connections to VIN (used as a submodule) and GIRIL (same joint learning pipeline), more detailed comparisons are definitely needed. A few concrete suggestions are: \na) add reward maps for GIRIL and provide qualitative analysis on the difference, e.g., in semantic meanings.\nb) for the proposed method, there are two variants on using the predicted reward maps: mean and max. It is unclear which version of the maps are shown in the paper and the empirical results are not conclusive on which one is better. It would be interesting to analyze the differences between them and, as a bonus, provides a rule-of-thumb on which one to use.\n\n2. There are several places where the writings could be improved. Some missing derivation details make understanding the paper difficult.\n\nQuestions:\n1. Please clarify a few terms in the paper: \na) Embedded MDP in section 3.1 near the bottom of page 3\nb) What is the conditional VAE? And the derivation details of equation (1) in the paragraph about Variational solution to vPERL.\n\n2. It is somewhat counterintuitive to me that vPERL-small generally outperforms vPERL-large since the small version has less state information. Could the authors provide some explanation for this?\n\n3. The baseline method VIN/BC is confusing. Could the authors clarify whether it means Behavioral cloning (as in supervised learning with state-action pairs from demonstrations) or VIN?\n\n======================\nPost-rebuttal comments:\n\nI want to thank the authors for providing clear answers to my questions and comments. I found the answers satisfactory so I raised my score to 6. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "OXKjo6y8msG", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper assumes no access to the reward values and attempts to learn a policy by starting just with one demonstration to define the reward. For obtaining the reward, the authors rely on the ideas from Value Iteration Networks (VIN) method and they add the modules that help to deal with cases with complex transition dynamics. The resulting method is tested on atari domain and on continuous control tasks.\n\nStrong points:\n\n- The ideas from Value Iteration Networks are applied to the environments with complex dynamics.\n- The problem setting is interesting. The restrictive assumption about the amount of demonstrations makes the method practical. \n- The results are promising and seem to significantly outperform the competing baselines from inverse reinforcement learning.\n\nWeak points:\n\n- The paper relies a lot on VIN method ideas without introducing it with sufficient detail. This makes it difficult to follow as many  concepts (and notation) in Section 3 are not properly introduced. Besides, the contributions of this paper are not very clearly stated.\n- The analysis part in 4.1 is not very informative and does not provide *experimental* evidence on *why* the proposed method works well and what issues with previous methods it addresses. For example, it is still not clear to me why the small model performs better than the big model in Table 2. Another example is the lack of ablations to understand the contribution of each of the components of the method. Additionally, the reader would benefit from understanding how the gap between the proposed method and the baselines changes with the growing amount of demonstrations: In particular, does the proposed model address the issues of low data regime or does it help in general?\n- The atari environment actually includes the scores in the input images (clearly in Breakout, Bean Rider, Space Invader, Seaquest, Kung Fu, and maybe in Qbert, BattleZone). Is it fair if no reward is assumed to be available?\n- The paper is quite hard to understand, it immediately jumps into details of the method in Section 3 without providing sufficient overview and intuition into what the method tries to achieve and what the issues with previous works are. \n\n\nI am leaning towards the rejection of this paper. While the experimental results are very encouraging as they are provided for the environments with complex dynamics and with very little expert data, my main concerns are that 1) the method description is currently very hard to understand, the contributions are not stated clearly, and 2) the experimental results do not include any informative analysis to understand the underlying reasons for why the method performs well. \n\nQuestions:\n- In the conclusion, the authors say \"policies that generalise well to new tasks\", what is meant by this and what are the supporting experiments?\n- Why does the small method model work better than the big model?\n- What is the contribution of different components of the method? How sensitive the method is to the choice of hyperparameters, such as K?\n- In Figure 4 it seems that many methods didn't converge to the final score yet. Why was such a number of training steps chosen? What happens if the experiments run for longer? \n- What is the computational complexity of the proposed algorithm?\n- Why is it possible to take average or max reward values of the reward map? How is the reward map idea justified when it is provided not for the map of states but where each map corresponds to one state?\n\nAdditional comments\n- The clarity of writing could be further improved. For example, see the Conclusion section (but it applies in other parts of the paper as well), avoid -> avoids, outperform -> outperforms, limiting and limited are repeated etc. Figure 3 has its labels shifted and Figure 1 seems to have the rows swapped. In Figure 2: is q_\\phi meant? \n- Maybe the analysis part in 4.1 is more suitable for the related work.\n\n=== Post rebuttal ===\n\nI would like to thank the authors for the very detailed response and the improvements in the manuscript. I found the ablation experiments and experiment with \"No scores\" particularly useful. However, I am still a bit confused about better performance of the smaller model. Maybe, more understanding could be gained if there was an additional \"tiny\" model that would show that when going beyond a certain size, the performance degrade. Additionally, if the overfitting is an important issue, some regularisation methods could be explored. \n\nFinally, in the light of many changes in the paper and the original request of all the reviewers to have the writing in the manuscript improved, I think the paper would benefit from another round of reviews. However, I find this method promising and if the paper is not accepted this time, I encourage the authors to re-submit a revised version.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors introduce vPERL, a model that generates an intrinsic reward for imitation learning. vPERL is trained on demonstrations to minimise a variational objective that matches a posterior formed by "action backtracking" and a forward model, with the intrinsic reward coming from the reward map. The authors might be interested in related work on few shot imitation learning: e.g., "One shot imitation learning", Duan et al, 2017, "Watch, try learn: meta-learning from demonstrations and rewards", Zhou et al 2019. As all reviewers pointed out, and I can confirm, the paper is quite tricky to understand in its present form, and would very much benefit the writing being re-visited to more clearly express the ideas within (in particular, section 3, which is the core of the contributions). 

Decision: Reject