Visual Imitation Learning with Patch Rewards
Authors: Minghuan Liu
Tairan He
Weinan Zhang
Shuicheng YAN
Zhongwen Xu
Venue: ICLR 2023 poster (2023)
Keywords: imitation learning
reinforcement learning

Abstract: Visual imitation learning enables reinforcement learning agents to learn to behave from expert visual demonstrations such as videos or image sequences, without explicit, well-defined rewards. 
Previous reseaches either adopt supervised learning techniques or induce simple and coarse scalar rewards from pixels, neglecting the dense information contained in the image demonstrations.
In this work, we propose to measure the expertise of various local regions of image samples, or called patches, and recover multi-dimensional patch rewards accordingly. 
Patch reward is a more precise rewarding characterization that serves as fine-grained expertise measurement and visual explainability tool.
Specifically, we present Adversarial Imitation Learning with Patch Rewards (PatchAIL), which employs a patch-based discriminator to measure the expertise of different local parts from given images and provide patch rewards.
The patch-based knowledge is also used to regularize the aggregated reward and stabilize the training.
We evaluate our method on the standard pixel-based benchmark DeepMind Control Suite. 
The experiment results have demonstrated that PatchAIL outperforms baseline methods and provides valuable interpretations for visual demonstrations.  

Published as a conference paper at ICLR 2023
VISUAL IMITATION LEARNING WITH PATCH REWARDS
Minghuan Liu Liu! Tairan He', Weinan Zhang! Shuicheng Yan", Zhongwen Xu² Shanghai Jiaotong University, SSea AI Lab
(minghuanliu, whynot, wnzhang) @ @jjtu.ddu.cn, [yansc,xuzw) @ sea.com
ABSTRACT
Visual imitation learning enables reinforcement learning agents to learn to be- have from expert visual demonstrations such as videos or image sequences, with- out explicit, well-defined rewards. Previous research either adopted supervised learning techniques or induce simple and coarse scalar rewards from pixels, ne- glecting the dense information contained in the image demonstrations. In this work, we propose to measure the expertise of various local regions of image sam- ples, or called patches, and recover coalti-dimensional patch rewards accordingly. Patch reward is more precise rewarding characterization that serves as fine- grained expertise measurement and visual explainability tool. Specifically, we present Adversarial Imitation Learning with Patch Rewards (PatchAIL), which employs a patch-based discriminator to measure the expertise of different local parts from given images and provide patch rewards. The patch-based knowledge is also used to regularize the aggregated reward and stabilize the training. We evaluate our method on DeepMind Control Suite and Atari tasks. The experiment results have demonstrated that PatchAlL. outperforms baseline methods and pro- vides valuable interpretations for visual demonstrations. Our codes are available /P
at https:: //github. com/ sai 1--il-- sg/ /aat ChAIL.
INTRODUCTION
Reinforcement learning (RL) has gained encouraging success in various domains, e.g., games (Silver et al., 2017; Yang al., 2022; Vinyals et al., 2019), robotics (Gu et al., 2017) and autonomous driving (Pomerleau, 1991; Zhou et al., 2020), which however heavily relies on well-designed reward functions. Beyond hand-crafting reward functions that are non-trivial, imitation learning (IL) offers a data-driven way to learn behaviors and recover informative rewards from expert demonstrations without access any explicit reward (Arora & Doshi, 2021; Hussein al., 2017). Recently, visual imitation learning (VIL.) (Pathak et al., 2018; Rafailov et al., 2021) has attracted increasing attention, which aims to learn from high-dimensional visual demonstrations like image sequences or videos. Compared with previous IL works tackling low-dimensional inputs, i.e. proprioceptive features like positions, velocitics, and accelerations (Ho Ermon, 2016; Fu et al., 2017; Liu et al., 2021), VIL more general problem in the physical world, such learning cook by watching videos. For humans, it is very easy but it remains a critical challenge for Al agents.
Previous approaches apply supervised learning (i.e. behavior cloning (Pomerleau, 1991; Pathak et al., 2018)) or design reward functions (Reddy et al., 2019; Brantley et al., 2019) to solve VIL problems. Among them, supervised methods often require many expert demonstrations and suffer from compounding error (Ross et al., 2011), while the rest tend to design rewards that are too naive to provide appropriate expertise measurements for efficient learning. Some recent works (Rafailov et al., 2021; Cohen et al., 2021; Haldar et al., 2022) choose to directly build on pixel-based RL methods, i.e. using an encoder, and learn rewards from the latent space. These methods inspect the whole image for deriving the rewards, which tend to give only coarse expertise measurements. Additionally, they focus on behavior learning, and seldom pay any attention to the explainability of the learned behaviors.
In this work, we argue that the scalar-form reward used in previous works only conveys "sparse" information of the full image, and hardly measures local expertise, which limits the learning effi-
The work was done during an internship Sea AT Lab, Singapore. "CCoresponding author.

---

Published as conference paper at ICLR 2023
Observatons
Expert Palicy
Potch Rewards
Observttions
Randam Palicy
Patch Rewords
Cartpole Swingup
Cheetah Run
Finger Spin
Hopper Stand
Quadruped
Walker Stand
Walker Walk
Figure I: Visualization of patch rewards by proposed PatchAIL on seven DeepMind Control Suite tasks, The patch rewards are calculated regarding several stacked frames and are mapped back on each pixel weighted by the spatial attention map. We use the final model trained with around 200M frames. Strikingly, the patch rewards align with the movement of key elements (e.g., moving foot, inclined body) on both expert and random samples, indicating where the high rewards (red areas) low ones (blue areas) come from, compared with the reward backgrounds.
the patch rewards recovered from expert's visual demonstrations, we can easily reason why and where agents are rewarded or penalized in these tasks compared with expert demonstrations without any task prior. Figure best viewed in color. A corresponding video is provided in the demo page com/view view//iww///ee/// patchail/.
https: //sttes....//.../..e/..eees.eett google. com/
ciency. Instead of inducing a scalar reward from the whole image, our intuition is to measure the expertise of different local regions of the given images, or what we call patches, and then recover multi-dimensional rewards, called patch rewards accordingly. Compared with scalar rewards, patch rewards provide a more precise rewarding characterization of visual control tasks, thus supporting efficient IL training and visual explainability of the learned behaviors.
We then propose to achieve more efficient VIL via inferring dense and informative reward signals from images, which also brings better explainability of the learned behaviors. To achieve this goal, we develop an intuitive and principled approach for VIL., termed Adversarial Imitation Learning with Patch rewards (PatchAIL), which produces patch rewards for agents to learn by comparing local regions of agent image samples with visual demonstrations, and also provides interpretability of learned behaviors. Also, it utilizes patch regularization mechanism employ the patch infor- mation and fully stabilize the training. We visualize the patch rewards induced by our PatchAIL in Fig. From the illustration, we can easily interpret where and why agents are rewarded or penalized in these control tasks compared with expert demonstrations without any task prior. These examples show the friendly explainability of our method, and intuitively tell why the algorithm can work well.
We conduct experiments on wide range of DeepMind Control Suite (DMC) (Tassa et al., 2018) tasks along with several Atari tasks, and find our PatchAIL significantly outperforms all the baseline methods, verifying the merit of our PatchAlL for enhancing the learning efficiency of VIL. More- over, we conduct ablations on the pivotal choices of PatchAlL, showing key options in PatchAlL. such the number of patches and aggregation function highly impact the learning performance. To verify its explainability, we further compare the spatial attention maps of PatchAIL to baseline methods, showing PatchAlL is better focusing the correct parts of robot bodies. It is hoped that our results could provide some insights for exploiting the rich information in visual demonstrations benefit VIL and also reveal the potential of using patch-level rewards.
PRELIMINARIES
Markoy decision process. RL problems are typically modeled as a didiccuutted infinite horizon Markov decision process (MDP) M A, P, po, 7, where S denotes the set of states, AA the
2

---

Published as a conference paper at ICLR 2023
[0 the dynamics transition function, Pa S - [0, 1] the initial state distribution, and ? E [0 the discount factor. The agent makes decisions through a policy (als)) : S + A and receives rewards r S x A + R (or S x S - R). For characterizing the statistical properties of a policy interacting with an MDP, the concept of occupancy measure (OM) (Syed al., 2008; Puterman, 2014; Ho & Ermon, 2016) is proposed. Specifically, the state OM is defined as the time-discounted cumulative stationary density over the states under given policy म: (s) = -i=0 P(s s|T), similarly the state-action OM is (s, {(a)s)pr() and the state-transition OM is pr 6') [A (s, Pss a) da. To construct a normalized density, one can take an extra normalization factor Z T-y
action space, P: 8 A S
1 for mathematical convenience.
Adversarial imitation learning. Imitation Learning (IL) (Hussein et al., 2017) aims to Learn from Demonstrations (LfD), and specifically, to learn policy from offline dataset T without ex- plicitly defined rewards. Typically, the expert demonstrations consist of state-action pairs, iee, T" arg min, Es-og 1e (TT((-s) [s), ([]])] > arg min, e (prg (8, a), P (x,8, a)), where denotes some dis- tance metric function. Adversarial imitation learning (AIL) (Ho & Ermon, 2016; Fu et al., 2017: Ghasemipour et al., 2020) is explored to optimize the objective by incorporating the advantage of Generative Adversarial Networks (GAN). In (Ho & Ermon, 2016), the policy learned by RL on the recovered reward is characterized by
T (s,, ). A general IL objective minimizes the state-action OM discrepancy
arg min H(T) + +56 (p= - PTR)
(1)
where s a regularizer, and v= is its convex conjugate. According to Eq. (1), various set- tings of vi can be seen as distance metric leading to various solutions of IL, like any f- al., 2020). In particular, Ho & Ermon (2016) chose a special form of zà and the second term becomes minimizing the JS divergence by training discriminator that differentiates between the agent's action samples and the expert data:
divergence (Ghasemipour
L(D) -s(,,))- log D(s, a)] E(s,a)~<loog (1 D(s,
As the reward can be derived from the discriminator, the policy can be learned with any RL algorithm. When the action information is absent in demonstrations, i.e., T (81, a common practice is to minimize the discrepancy of the state transition OM " arg min_ [ (Pag (8, s') (x,sss')) instead by constructing a discriminator for state transitions (Torabi et al., 2019):
((D) = - D(s, s')] E(s,s")~ (1 D(s,s)))]
In the visual imitation learning problem we study here, visual demonstrations usually do not contain expert actions, shall take the state transition OM form (Eq. (3)) in the solution.
ADVERSARIAL VISUAL IMITATION WITH PATCH REWARDS
The proposed Adversarial Imitation Learning with Patch rewards (PatchAlL.) method is built upon Adversarial Imitation Learning (AIL) to derive informative and meaningful patch rewards while better explaining expert demonstrations. Trained in a GAN-style, the discriminator in AIL works by differentiating the difference between agent samples and expert samples, i.e. through a binary- class classification problem that regards the expert data real (label and the agent data as fake (label 0). This strategy works well when learning from the low-level proprioceptive observations, where the reward given based on the discriminator. However, this is insufficient when data is videos or image sequences. With a scalar label to classify and derive the reward simply based on the whole visual area, the images may easily be regarded as far-from-expert due to small noise or changes in local behavior. Also, there may be multiple regions of images that affect the rewarding of control tasks, while the discriminator may only focus on insufficient regions that support its overall judgment. In addition, such scalar rewards are only meaningful for the agent to learn its policy, and they lack explainability for humans to understand where the expertise comes from, such as which parts of the image contribute to higher reward bring negative effects.
In some recent works (Cohen et al., 2021; Rafailov et al., 2021), the discriminator measures the expertise based on the latent samples encoded by the encoder trained with pixel-based RI. Albeit this
3

---

Published as a conference paper at ICLR 2023
may relieve the sensitivity to local changes, it is not intuitively reasonable. When the discriminator shares the latent space with the critic, the encoder is updated only depending on the critic loss, which is the TD eπrо that is meant to learn correct value prediction depending the reward obtained from the discriminator. On the other hand, the discriminator space induced by the encoder. As both modules must be learned from scratch, such interdependence can cause training instability, and deteriorate the attention to key elements related to the expertise. Besides, it still lacks explainability due to the scalar label.
updated according to the latent
3.1 FROM SCALAR TO PATCHES
We make an intuitive and straightforward modification to the classic way of measuring expertise, i.e. evaluating different local parts the given images by replacing the scalar judgment with patch ones. By evaluating local parts, on one hand, local changes only affect some of the patches rather than the whole, increasing the robustness of the overall judgment. On the other hand, the discriminator is forced to focus on more areas. Also, in this way, the classification result and reward on each patch can measure how each part of images contributes the similarity expert visual demonstrations, which constructs a holistic view for us to understand where the expertise comes from.
Patch discriminator. To implement the idea of the AIL framework, we make the discriminator classify image patches and give patches of reward lieu of scalar rewards by learning a single clas- sification result sual demonstration from experts (s.s') ) ~ Tg. and the sample from the agent policy π, optimize the patch discriminator Dxx with patch labels:
D
PxP
image. Formally, given a vi-
Illustration of the patch discriminator, where each pixel of the right-side mosaic represents the patch logit corresponding local region of the left-side image.
Figure 2:
C(D) = - Ei..s)aa llog D Dpxp(ss s'), E(,,""~a (log (IpxP - Dpxp(s,s'))
(4)
In this formula, the discriminator outputs a matrix of probabilities DpxP S S S [0, 1 PxP, where P is the height/width of each image patch, illustrated in Fig. 2. In other words, the clas- sification loss is computed on each patch logit, rather than on the averaged logit of all patches. Furthermore, we do not have to maintain extra encoder for the discriminator and the discrimina- tor learns everything directly from pixels.
3.2 LEARNING FRAMEWORK
Reward aggregation. Generally, the probabili- ties/logits given by D need to be further transformed by R = h(D) to become a reward function where transforms probabilities to reward signals, e.g. in GAIL (Ho & Ermon, 2016) h
Regularize
Patch
Policy
Critic
log D or R = - log(( (1 D), and R = log D - log(1 - D) in AIRL (Fu al., 2017) and DAC (Kostrikov et al., 2018). In addition, after constructing patch rewards, we need an aggregation function integrate the global information while deriving scalar signals for training RL agents. Most straightforwardly, one can choose mean as the aggregation operator that evaluates the overall expertise compared to expert visual demon- strations for given image sequences, whilst other kinds of operations can also be adopted. Formally, we can derive the reward function from the discriminator output as
Patch Discriminator
Patch Discriminator
Image Encoder
Inference
Export Data
Samp ed Data
Figure 3: The framework of PatchAlL..
R(s,s') Aggr(Rpxxle s')) =
(5)
where Aggr aggregates the patch rewards into scalar signals with operators mean, max, min and median, Different aggregation operators can incorporate different inductive biases into the problem and lead various learning preferences.
4

---

Published as a conference paper at ICLR 2023
Patch regularization. By aggregating the patch rewards into scalars, the rich information tained in the patch rewards is lost. In other words, the scalar rewards again only tell the agent about the overall judgment, and the details such as which part of the image is similar to the expert demon- stration do not help the learning. To solve this issue, we further propose to regularize the agen' 's not only require the
patch rewards by comparing them to the expert's patch rewards. Intuitively, agent to maximize the similarity of the sampled data to the expert but also want the distribution of the agent's 's patch rewards be similar to the expert's. To this end, we construct a normalized probabil- ity mass of a discrete distribution ((PxxP(s, (')) from the output logits of the patch discriminator (for simplicity we abuse the symbol DpxP P as the unnormalized logits by the discriminator in the following). Formally, we want to maximize the similarity Sim Sim(s, 8') to the demonstrations for the observation pair (8, 8'), where the similarity S Sim(s, 8') is defined by the nearest statistical distance to the distributions induced by expert samples:
max (eeid') (sim(s, 8')]
(6)
Sim(s, s') = exp (- min (((Dpxp(ss s')), ADDpx((s. +z))) (sg.*g)
where É represents some distance metrics. However, computing Eq. (6) can be time-consuming
during training. As replacement, we choose to compute simplified similarity Sim(s, s') as the distance to the distribution induced by the averaged logits of expert samples instead:
Sim(s,s') s') exp (- pxpss 5')), A(B(sR,,'))> Dpxp(ss, ('))))
(7)
In the experiment we compare these two choices and find the simplified similarity plays good role in measuring the distance between the expert demonstrations and agent samples. Particularly, if we instantiate é as KL divergence DKL, the similarity is then computed as
P
P
(Dpp((sss Bl A (D (D ((E,, (ss,s'))) st))]]
Sim(s,s') = exp (- ΣΣ AIJ (ix(ppx(ss, log
(8)
1=1 j=1
where A., denotes the probability value of the (i, entry of A(D The similarity can serve as the regularizer to weigh the aggregated reward signals. Considering ways of using the similarity regularizer, we have two variants of the PatchAlL algorithm. The first version is named PatchAIL- PatchAlL-W(eight) where we take the regularizer as the weighting multiplier on the reward:
R(s, 8') = A Sim(s, 8')
((, 8') d'))):
the other variant is named PatchAlL-Bi B(onss, where the regularizer serves as a bonus term:
R(s, s') = X Sim(s, s') (Aggrh((D)))) (Dpxp(ssP
(10)
À is the hyperparameter for balancing the regularization rate. In our experiments, we find that such a regularization helps stabilize the training procedure on some tasks, as shown in Fig.
Implementation of the patch discriminator. In principle, the architecture of the patch discrim- inator is an open choice to any good designs to construct patch rewards, such as vision trans- former (Isola et al., 2017) (explicitly break input images patches) (FCN) (Long et al., 2015) (output matrix implicitly mapped back to patches on input images). In this paper, we choose implement the patch discriminator as a 4-layer FCN, referring to the one in (Zhu et al., 2017). In detail, the discriminator has four convolution layers without ML. connections and keeps the last layer as Sigmoid to give the probabilities of the classification for each patch, where each patch may share some perceptive field with each other. The discriminator classifies if each patch real (i.e,, expert) or fake (i.e., non-expert). As result, we can obtain multi-dimensional rewards by evaluating the expertise of every local region of images.
colly-convolution network
Overall algorithm. In a nutshell, PatchAlL. obtains patch rewards from the patch discriminator, and training agents with reward aggregation and regularization. The same as other AIL methods, PatchAIL works with any off-the-shelf pixel-based RL algorithm. In this paper, we choose DrQ- v2 (Ye et al., 2020) as the underlying component for learning the policy and integrating gradient penalty (Gulrajani et al., 2017) for training the discriminator. The algorithm is trained under an off- policy style inspired by Kostrikov et al. (2018), therefore the agent samples are from the replay buffer and the expert samples are from the demonstration dataset T. An overview of the framework is illustrated in Fig. 3, and the step-by-step algorithm can be found in Appendix A.
5

---

Published as a conference paper at ICLR 2023
4 RELATED WORK
Pixel-hased reinforcement learning. The proposed solution for Visual Imitation Learning (VIL) is built the recent success of pixel-based reinforcement learning. When learning from raw pixels, a good state representation of raw pixels encoded by image encoders is crucial to the learning perfor- mance, and it widely acknowledged that the policy shares the latent representation with the critic, and only the critic loss is used for updating the encoder (Yarats et al., 2019; Laskin et al., 2020a;b; Yarats al., 2020). Besides, many recent works integrate diverse representation learning techniques for pixel-based RI.. For instance, Yarats et al. (2019) trained a regularized autoencoder jointly with the RI. agent; Lee et al. (2020) learned a latent world model on top of VAE features and built la- tent value function; Laskin et al. (2020a) optimized contrastive learning objectives to regularize the latent representations; Laskin et al. (2020b); Yarats et al. (2020; 2021) utilized data augmentation for efficient representation learning. For learning such image encoders, all these methods require well-defined reward function, which is inaccessible for VIL.
Visual imitation learning. In the imitation learning community, most previous solutions are de- signed for state-based settings, where the states/observations are low-level perceptive vector inputs, such as (Ho & Ermon, 2016; Kostrikov al., 2018; 2019; Liu et al., 2021; Reddy et al., 2019; Brantley et al., 2019; Garg et al., 2021; Liu et al., 2022). Among them, some can be directly used for visual imitation. For example, Reddy et al. (2019) allocated a positive reward (+1) to expert data and negative reward (-1) to experience data sampled by the agent; Brantley et al. (2019) used the disagreement among a set of supervised learning policies to construct the reward function. Instead of learning the reward function, Garg et al. (2021) chose to learn Q function from expert data directly, but only tested on discrete Atari games. Besides, Pomerleau (1991); Pathak et al. (2018) designed supervised learning techniques for resolving real-world imitation learning tasks. Recently, Rafailov et al. (2021); Cohen et al. (2021) extended previous adversarial imitation learning solution based on latent samples with existing pixel-based RL algorithms was intuitively unreasonable dis- cussed in Section 3. Based on Cohen et al. (2021), Haldar et al. (2022) further augmented a filtered BC objective into the imitation learning procedure, showing a considerable improvement in effi- ciency and final performance. However, their work requires access to experts' action information, which usually absent image-based demonstrations, where we only have visual observations. In Appendix B.3 we also report experiments of PatchAlL augmented with BC objective assuming the accessibility to the action information, reaching much better efficiency and more stable learning performance. Besides the learnability, none of these works can recover explainable reward signals that interpret the expertise of the demonstration and help people to understand the nature of tasks.
Image patches in literature. Breaking up images into patches is not novel idea in the deep learning community, not even the proposal of the patch discriminator, as it has already been used in Dosovitskiy al. (2020); Zhu et al. (2017); Isola al. (2017). On one hand Li & Wand (2016); Isola et al. (2017); Zhu et al. (2017) argued that using a patch discriminator can better help the generator learn details, where Li & Wand (2016) used it to capture local style statistics for texture synthesis, and Isola et al. (2017); Zhu et al. (2017) used it to model high-frequencies; on the other hand, Dosovitskiy et al. (2020) proposed to utilize the ability of transformer (Vaswani et al., 2017) by splitting the image into patches, mainly for capturing the semantic information in images. Different from them, our usage of image patches along with the patch discriminator is distinct and novel in the context of RL and VII., we are seeking a fine- fine-ggained expertise measurement of the images given a set of visual demonstrations, like where the agent should be rewarded and where not.
EXPERIMENTS
In our experiments, we thoroughly examine our proposed Adversarial Imitation Learning with Patch rewards (PatchAlL) algorithm. We first conduct fair comparison against baseline methods on a wide range of pixel-based benchmarking tasks on Deep Control Suit (Section 5.1,, showing the effectiveness of PatchAIL; next, we compare the spatial attention maps given by PatchAIL with baselines to show the reasonability of the reward given by PatchAIL: finally, we conduct a set of ablation studies to investigate the effect of different choices in PatchAlL. Due to the page limit, leave more results on Atari tasks, additional ablation experiments and more implementation details in Appendix B.

---

Published as a conference paper at ICLR 2023
Table 1: Averaged returns on 7 DMC tasks 500K and 1M environment steps over 5 random seeds using 10 expert trajectories. gray columns represent our methods, and Shared-Enc. A/L is a solution proposed in (Cohen et al., 2021). PatchAlL. and its variants outperform other approaches in data-efficient (500K) and asymptotic performance (IM). Since BC does not need online training, we report their asymptotic performance.
500K
PatchAlL W.o. Reg.
PatchAIL-W
PatchALL-B
Shared--Ene, AlL Independent-Bi Enc. AlL
BC
Expett
Carppole Swingup
184+15
811+233
729-140
091+34
+21+1200
Choctah Rua
413+16
409:29
415+30
533-144
335=
185
890=19
Finger
850+32
271+25
839+23
777+448
547--51
124+1120
976+9
Hopper Stand
235=:128
430199
207::007
2064189
729
386:272
93929
Walker Stand
896+38
902#:34
900+:39
881+15
865+=20
496+700
970=20
121
5131117
20045
455192
166:1110
961=18
Quadruped Run
230#38
239::43
273::17
220::33
263
277:588
547+:36 136
IM step
Cartpole Swingup
854#2
844+9
83
84443
367=99
521=+120
859=00
Cheetah Run
6541124
681121
15
646:119
7911:13
185:1499
190=11
Finger Spin
875+38
855+:3
943:27
826+71
110=54
184:+1200
976=
Hopper Stand
7811:37
693
103119
676:11300
762.:20
386.172
93919
Walker Stand
916:1:1
97312
9721-
8862227
824::25
496:270
970=20
Walker Walk
956.1.
951.13
960112
952.12
945.15
566.1110
961-118
Quadruped Run
4062114
42319
41628
331:33
354234
277:188
1474:336
Baselines. Our experiments cover three kinds of baseline methods.
Shared-Encoder AIL.: the discriminator learns from latent samples as discussed in Section 3, which is also the proposal in Cohen et al. (2021):
Independent-Encode AlL: the discriminators learn from images directly with a scalar label. Behavior cloning (BC): supervised learning for the policy function, which included only for reference it gets access to the expert's 's exact actions.
For fair comparison, all baselines are implemented under the same code base with the same archi- tecture of the policy and Q networks. AlL methods take the same choice of the reward transformation function R log D D-log(1 D) as in Fu et al. (2017); Kostrikov et al. (2018). We emphasize that all seggenee sequences and cannot obtain accurate action information.
agents except BC can only require image
Experimental setup. Our experiments are conducted on 7 tasks in the DeepMind Control (DMC) Suite (Tassa et al., 2018). For expert demonstrations, we use 10 trajectories collected by a DrQ- v2 (Yarats et al., 2021) agent trained using the ground-truth reward and collect. We follow Yarats et al. (2021) to use Random-Shife ft to augment the input samples for training the policy, critic and discriminator for every method tested. The default choice of aggregation function for PatchAlL is mean; and the default architecture of the patch discriminator is a 4-layer FCN with kernels [4x4, 32, 2. 11. [4x4, 64, 1, 1], [4x4, 128, 1, 1] and [4x4, 1, 1. 1] (represented as [size, channel, stride, padding]), whose output has 39 39 patches in total and each has a receptive filed of 22. We put results trained by 10 trajectories in the main text and leave the rest in Appendix B. Since the expert can be as the optimal policy under the ground-truth reward function, i.e., its trajectory distri- bution regarded as proportional to the exponential return (Ho & Ermon, 2016; Liu et al., 2021), we take the ground-truth return the metric for measuring the agents. Throughout experiments, these algorithms are tested using fixed hyperparameters on all tasks and the full set of hyperparameters in Appendix B. All methods use 2, 000 exploration steps which are included in the overall step count. We run 5 seeds in all experiments under each configuration and report the mean and standard deviation during the training procedure.
5.1 IMITATION EVALUATIONS
We compare our proposed PatchAIL and its variants to all baseline methods on DeepMind Control Suit tasks. In Tab. 1, we present the average returns on each task at both 500K and 1M training steps, and in Fig. 4, we show the learning curve of each method. Statistics from Tab. indicate that on almost all testbeds, PatchAIL methods own better imitation efficiency and efficacy. At 500K frames, PatchAIL obtains an averaged performance gain of 44% over Shared-Enconn Encoder AIL and 13% over Crrependent-Encodee AIL; at IM steps, the averaged performance gain over Shared-Encoded AIL becomes 11% and over Independent-Enoee encoder-based AIL baselines work well various tasks. Nevertheless, by learning its represen- tation, independent-encodee AlL is rather unstable in almost all tasks during the learning procedure. Specifically, when it reaches near-expert performance, it starts to fall and becomes worse (such as
AIL is 131%. From Fig. 4, we find that the
7

---

Published as a conference paper at ICLR 2023
Cartpole Swingup
Cheetah Run
Finger Spin
Hopper Stand
800
Episode Return
800
Episode Return
Episode Return
800
M
800
600
600
Episode
400
400
200
200
200
0.0
0.5
1.0
1.5
0.0
1.0
1.5
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
Frames
Frames
Frames
Frames
Walker Stand
Walker Walk
Quadruped Run
1000
1000
Episode Return
Episode Return
Episode Return
PatchALL w.o. (Durs)
800
Pelciell--Weight (Cun)
600
400
Shared Prcoder All
PatchAll Bonus (Ours)
400
400
Indepandent Encoder AlL
200
0.0
1.5
1.0
Frames
leB
Frames
1eB
Frames
led
Figure 4: Learning curves on DMC tasks over 5 random seeds using 10 expert trajectories. The curve and shaded zone represent the mean and the standard deviation, respectively.
Obsrenttions
PotchilL-- WW
Shared Encader
Indegeeneent Enccrder
Walker Stand Figure 5: Selected spatial attention maps showing where discriminators focus to make decision, given on Deepmind Control
Cartpole Swingup
Cheetah Run
Finger Soin
Hopper Stand Quadruped Run
Walker Walk
by discriminators PatchAIL-W shared-encoder AIL and independent-encoodr Suite tasks using the final model trained around 200M frames. Red represents high value while blue is low.
Sartpole-Swingup. Cheetah-Run, Finger-Spin, Walker-Stand, and Walker-Walk) On the contrary, shared-encoder AlL performs better and tends to become steady the state representation is almost unchanged, showing in some tasks, it has less learning efficiency in the early learning stage due to the interdependent up- date of the discriminator and the critic (such Cartpole-Swingup. Finger-Spin and Walker-Walk). In comparison, PatchAIL and its variants are more stable and much more efficient. This obvi- ous Cartpole-Swingun and Finger-Spin. Notably, they have noticeable performance advantage in the Quadruped-Run domain. Moreover, with the proposed patch regulation, PatchAIL becomes more stable, which quite evident the Finger-Spin domain, yet the main performance gain comes from the patch reward structure given the overall results.
late stage of training, where be effective in many cases. However,
5.2 VISUAL EXPLANATIONS
In this part, we focus on the visual explanations, which we deem as superior advantage of PatchAIL compared with baselines. To understand where the discriminators focus, visualize the spatial feature maps of different methods in Fig. 5. is clear that PatchAlL provides complete and holistic concentration the key elements of images. Particularly, notice that the encoder-based
8

---

Published as a conference paper at ICLR 2023
Cheetain Run
Finger Spin
Cheetah Run
Finger Spin
1400
Raturn
Episodd Return
800
Sassors Raturn
800
Epicada Return
003
46 46 Pocches
42 42 PacThes
600
800
600
400
400
403
25 Putchas
Espert
202
Frames
Frames
Frames
Frames
(a) Ablation study for aggregation functions.
(b) Ablation study for the number of patches.
Figure 6: Ablation studies, only tested with PatchAIL w.o. Reg.
baselines both miss key areas on some tasks, such as on Finger-Spin, where both methods seem mainly to make judgments the left finger. Furthermore, shared-encoder AlL also overlooks the leg and independent AIL pays more attention to the background
on Hopper-Stand and Quadruped-RR on Cartpole-Swingup. Cheetah- and Quadruped- These observations meet our assumptions in Section 3. To better comprehend and explain which part of images contribute to the overall expertise. we mapped each patch reward back its corresponding pixels, weighted by the attention map, on both expert data and random data. The results have been illustrated in Section 1, which indicates where the high rewards or low ones come from.
5.3 ABLATION STUDIES
Ablation on aggregation function. In the proposed PatchAlL framework, we learn patch rewards for image data and then aggregate them to obtain a scalar reward for learning the policy by RL algorithms. In this paper, the default choice of the aggregation function is mean, but there leaves wide range of choices. Therefore we conduct ablation experiments on various choices of aggregation functions, like min, max and median. The results on the Finger-Spin and Cheetah- domain are shown in Fig. 6a. It can be seen that mean is tne-of-the-best choice among these aggregation functions, while median and min also preserve well learnability. On the other hand, the max operator tends to fail in both domains, which may be attributed to its over-optimism of bad data.
Ablation on size of patches. We are also interested in the impact of the size of patches along with the receptive field of each patch the performance of PatchAlL. As mentioned above, the default implementation of PatchAlL utilizes a 4-layer FCN, and each layer uses a convolution kernel with a size of 4x4, resulting 39x39 patches in the last layer. Each of them has a 22x 22 reception field. In ablation, we only modify the kernel size as [2, 3, 5, 8] with [46 x 46, 42 x 42, 35 x 35, 25 x 25] patches and reception field sizes of [8 x 8, 15 15, 29 x 29, 50 x 50). keeping the stride and padding size the same as the default setting; besides, we also test special case where each layer of the FCN is 1 x convolution layer, resulting in a pixel level patch reward (84x84 patches and each has a 1 x reception field). The learning curves on the Finger-Spin and Cheetah-Run domain are illustrated in Fig. 6b, where we observe three apparent phenomena: 1) PatchAlL with bigger reception fields behaves with high training efficiency in the early stage; 2) with the biggest reception fields (i.e., 25 x 25 patches with a kernel size of 8 x 8) PatchAlL can be extremely unstable; 3) learning rewards for each pixel fails to obtain well-performed policy; 4) some tasks like Cheetah- Run are less sensitive to the change of patch size than Finger- Spin, which maybe can explain why PatchAlL almost shows advantage in Fig. 4 compared to baselines.
CONCLUSION
We proposed PatchAIL. intuitive and principled learning framework for efficient visual imitation learning (VIL). PatchAlL offers explainable patch-level rewards, which measure the expertise of various parts of given images and this information regularize the aggregated reward and stabilize the training, in place of less informative scalar rewards. We evaluate our method on the widely-used pixel-based benchmark DeepMind Control Suite. The results indicate that PatchAIL supports efficient training that outperforms baseline methods and provides valuable interpretations for VIL. We think that the proposed patch reward is a crucial method for utilizing the wealth of information contained in image observations, and that there is still a great potential for leveraging the pateh-level information to enhance visual reinforcement learning and visual imitation learning.

---

Published as a conference paper at ICLR 2023
7 ETHICS STATEMENT
This submission adheres to and acknowledges the ICLR Code of Ethics and we make sure that we do violate any ethics concerns.
8 REPRODUCIBILITY STATEMENT
All experiments included in this paper conducted over 5 random seeds for stability and reliability. The algorithm outline is included in Appendix A, and the hyperparameters and implementation details are included in Appendix C. Our code has been released to the public.
ACKNOWLEDGEMEN
The SJTU Team supported by Shanghai Municipal Science and Technology Major Project (202ISHZDZX0102 and National Natural Science Foundation of China (62076161). Minghuan Liu is also supported by Wu Wen Jun Honorary Doctoral Scholarship, Institute, SJTU. We sin- cerely thank Quanhong Fu very much for her careful revising of our first manuscripts, and we also thank Qingfeng Lan, Zhijie Lin, Xiao Ma for helpful discussions.
REFERENCES
Saurabh Arora and Prashant Doshi. A survey of inverse reinforcement learning: Challenges, meth- ods and progress. Artificial Intelligence, 297:103500, 2021.
Kianté Brantley, Wen Sun, and Mikael Henaff. eeaagreemenn-rrgglaaiizee imitation learning. In International Conference on Learning Representations, 2019.
Samuel Cohen, Brandon Amos, Marc Peter Deisenroth, Mikael Henaff, Eugene Vinitsky, and Denis Yarats, Imitation learning from pixel observations for continuous control. arXiv preprint, 2021.
Ekin Cubuk, Barret Zoph, Jonathon Shlens, and Quoe Le. Randaugment: Practical automated data augmentation with reduced search space. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition workshops, pp. 702-703, 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.1 1929, 2020.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse rein- forcement learning. arXiv preprint arXiv: 1244, 2017.
Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, and Stefano Ermon. Iq-learn: Inverse soft-q learning for imitation. Advances in Neural Information Processing Systems, 34, 2021.
Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. A divergence minimization perspective on imitation learning methods. In Conference on Robot Learning, pp. 1259-1277. PMLR, 2020.
Shixiang Gu, Ethan Holly. Timothy Lillicrap. and Sergey Levine. Deep reinforcement learning for robotic manipulation with asynchronous off- off-policy updates. In 2017 IEEE international confer- ence robotics and automation (ICRA), pp. 3389-3396. IEEE, 2017.
Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez, Konrad Zolna, Rishabh Agarwal, Josh S Merel, Daniel Mankowitz, Cosmin Paduraru, et al. RI unplugged: A suite of benchmarks for offline reinforcement learning. Advances Neural Information Process- ing Systems, 33:7248-7259, 2020.
10

---

Published as a conference paper at ICLR 2023
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Im- proved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
Siddhant Haldar, Vaibhav Mathur, Denis Yarats, and Lerrel Pinto. Watch and match: Supercharging imitation with regularized optimal transport. arXiv preprini arXiv:2206. 2022.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR), 50,))
Phillip Isola, Jun Yan Zhu, Tinghui Zhou, and Alexci A Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1 1125-1133 2017.
Rohit Jena, Changliu Liu, and Katia Sycara. Augmenting gail with be for sample efficient imitation learning. In Conference on Robot Learning, pp. 80-90. PMLR, 2021.
Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tomp- son. Discriminator-actor.com critic Addressing sample inefficiency and reward bias in adversarial imitation learning. arXiv preprint arXiv: (18880990222) 2018
Iyya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off-policy distribution matching. arXiv preprint arXiv: 2019.
Michael Laskin, Aravind Srinivas, and Pieter Abheel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. In International Conference Machine Learning, pp. 5639- 5650, PMLR, 2020a.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. Advances Neural Information Processing Systems, 1981984499995,, 2020b.
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. Advances Neural Information Pro- cessing Systems, 33:741-752, 2020.
Chuan Li and Michael Wand. Precomputed real-time texture synthesis with markovian generative adversarial networks. In European conference on computer vision, pp. 702-716. Springer. 2016.
Minghuan Liu, Tairan He, Minkai Xu, and Weinan Zhang. Energy-based imitation learning. In International Conference on Autonomous Agents and Multiagent Systems, 2021.
Minghuan Liu, Zhengbang Zhu, Yuzheng Zhuang, Weinan Zhang, Jianye Hao, Yong Yu, and Jun Wang. Plan your target and learn your skills: Transferable state-only imitation learning via de- coupled policy optimization. In Proceedings of the 39th International Conference on Machine Learning, volume 162, pp. 14173-1411 14196. PMLR, 2022.
Jonathan Long, Evan Shelhamer, and Trevor Darrell. Fully convolutional networks for semantic segmentation. In Proceedings the IEEE conference on computer vision and pattern recognition, pp. 3431-3440, 2015.
Deepak Pathak, Parsa Mahmoudich, Guanghao Luo, Pulkit Agrawal, Dian Chen, Yide Shentu, Evan Shelhamer, Jitendra Malik, Alexei A Efros, and Trevor Darrell. Zero-shot visual imitation. In Proceedings of the IEEE conference on computer vision and pattern recognition workshops, pp- 2050-2055 2018.
Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neu- ral Computation, 3(1):88-9)) 97, 1991.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
11

---

Published as a conference paper at ICLR 2023
Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Visual adversarial imitation learning using variational models. Advances in Neural Information Processing Systems, 34:3016- 3028, 2021.
Siddharth Reddy, Anea D Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards. arXiv preprint 1905. 2019.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and struc- tured prediction no-regret online learning. In Proceedings of the fourteenth international con- ference on artificial intelligence and statistics, pp. 627-633- JMLR Workshop and Conference Proceedings, 2011.
Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-basee local- ization. In Proceedings of the IEEE international conference on computer vision, pp. 618-626, 2017.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. Mastering chess and shogi by self-play with a general reinforcement learning algorithm. arXiv preprint arXiv1111.018111 2017.
Umar Syed, Michael Bowling, and Robert E Schapire. Apprenticeship learning using linear pro- gramming. In Proceedings of the 25th international conference on Machine learning, pp. 1032- 1039, 2008
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 00690, 2018
Faraz Torabi, Garrett Warnell, and Peter Stone. Adversarial imitation learning from state-only demonstrations. In Proceedings of the 18th International Conference on Autonomous Agents and MultiAgent Systems, pp. 2229-2231, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Oriol Vinyals, Igor Babuschkin, Wojeiech Czarnekki, Michaël Mathieu, Andrew Dudzik, Juny- oung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782)350-3544 2019.
Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Mare Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In International conference machine learning, pp. 1995-2003, PMLR, 2016.
Guan Yang, Minghuan Liu, Weijun Hong, Weinan Zhang, Fei Fang, Guangjun Zeng, and Yue Lin. Perfectdou: Dominating doudizhu with perfect information distillation. In Advances in Neural Information Processing Systems, 2022.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Im- proving sample efficiency in model-free reinforcement learning from images. arXiv preprint 11100:1110001411, 2019.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference Learning Representa- tions, 2020.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous con- trol: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2 107.099440 2021.
Deheng Ye, Zhao Liu, Mingfei Sun, Bei Shi, Peilin Zhao, Hao Wu, Hongsheng Yu, Shaojie Yang, Xipeng Wu, Qingwei Guo, et al. Mastering complex control in moba games with deep reinforce- ment learning. In Proceedings of the AAAI Conference Artificial Intelligence, 2020.
12

---

Published as a conference paper at ICLR 2023
Ming Zhou, Jun Luo, Julian Villella, Yaodong Yang, David Rusu, Jiayu Miao, Weinan Zhang, Mont- gomery Alban, Iman Fadakar, Zheng Chen, et al. Smarts: Scalable multi-agent reinforcement learning training school for autonomous driving. arXiv preprint arXiv:2010.09776. 2020.
Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pp. 2223-2232, 2017.
13

---

Published as a conference paper at ICLR 2023
Appendices
A ALGORITHM OUTLINE
Algorithm Adversarial Imitation Learning with Patch rewards (PatchAIL)
Input: Empty replay buffer B, patch discriminator DpxP, policy n, expert demonstrations T.
2: for k= 1....... do
3: Collect trajectories {(8, a, s',", s., } using current policy m
4: Store {(s, 8', done)) in B
5: Sample (s, a, 5') B, (sg, s'Ez) D
6: Update the discriminator Dx followed by Eq. (4)
7: Compute the reward r for (s, a, s') with the patch discriminator Dyxr by Eq. (5) or Eq. (9)
Eq. (10)
8: Update the policy To by any off-the-shelf RL algorithm
9: end for
B ADDITIONAL EXPERIMENT RESULTS
B.I THE CHOICE OF OBSERVATION PAIRS
In this section, test the choice of using an observation or observation pairs (8,8") for the discriminator in our experiments. The learning curve of the shared-encoder AIL and PatchAIL is shown in Fig. 7. It seems that this choice makes nearly difference for shared-encoder AlL, as it learns from latent samples. On the contrary, utilizing observation pairs (s, shows better performance for PatchAIL. In experiments, since the visual observation is always composed of stack of images (e.g., a stack of frames in DeepMind Control Suite used in this paper), using an observation pair is just stacking two stacks of images into one bigger stack.
B.2 IMITATION USING FEWER DEMONSTRATIONS
We also evaluate algorithm against baselines using fewer demonstrations. As shown in Fig. 8, we can draw consistent conclusions with the results using 10 trajectories on most tasks. Neverthe- less, with less demonstrations, PatchAIL and its variants appear have no advantage on Cartpole- Swingup with fewer demonstrations, they emerge as winners on Hopper-Stand.
B.3 AUGMENTED WITH BC WITH ONLY DEMONSTRATION
When we can get access the action information, we are able further improve the imitation performance by mimicking the experts. Motivated by Haldar et al. (2022) and Jena et al. (2021), test the performance by augmenting PatchAIL with BC, where the weight of BC loss is annealing following Jena al. (2021). We compare with the recent state-of-the-an work ROT (Haldar et al., 2022), and the two encoder-based baselines in Section 5, and show the results in Fig. We find that
Cartpole Swingup
Cheetah Run
Finger Spin
Hopper Stand
10000
Episode Return
800
Episode Return
800
Episode Return
Episode Return
800
600
600
600
400
200
200
1.5
1.5
0.0
0.5
1.0
1.5
Frames
165
Frames
1e6
Frames
1eB
Frames
1e6
w.o Reg-Dis.s
Reg--DS))
Shared-Encoder AlL-O((ss))
Shared-Encoder AlL-Dis
Expert
Figure 7: Ablation for the choice of observation pairs over 5 random seeds.
14

---

Published as a conference paper at ICLR 2023
Cartpole Swingup
Cheetah
Finger Spin
Hopper Stand
1000
Episode Return
800
Episode Return
800
Episode Return
Episode Return
800
600
800
400
400
400
200
200
200
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
0.0
0.5
3.0
1.5
Frames
1e6
Frames
1e6
Frames
led
Frames
le6
Walker Stand
Walker Walk
Quadruped Run
1000
Episode Return
Episode Return
Episode Return
Reg.
800
MaachAl-Waigge (Durs)
620
600
400
Shared Encoder AIL
PatchAlL Bonus (Ours)
400
400
200
200
Expert
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
0.0
0.5
1.0
2.5
Frames
1e6
Frames
1e6
Frames
166
Figure 8: Learning curves on DMC tasks over random seeds using expert trajectories. The observations
are almost consistent with the results using 10 trajectories except on Cartpole-Swingee and Hopper-Stand.
PatchAIL and its variants show faster convergence speed but ROT seems to have a slightly better final performance, while PatchAlL-W and PatchAIL-B are able to alleviate the unstable problem of PatchAlL. In comparison, encoder-based baselines does not show any advantage when augmented with BC.
Cartpole Swingup
Finger Spin
800
Eols.ode Return
Episode Return
PetchALL
800
600
ROT
Shared Encoder ALL
400
Indlependentt Eneedd All.
400
Expert
200
0
0.25
0.75
0.00
0.50 Frames
0.75
1.00
Frames
166
166
Figure Learning curves by augmenting PatchALL with annealing using expert trajectories over random seeds. PatchALL, PatchAlL-W and PatchAlL-B all show faster convergence speed than all baselines but ROT seems to have slightly better final performance.
B.4 COMPARING THE SIMILARITY FORMULATION
In Section 3.2, we propose a similarity Sim(s,s) s') defined by the nearest statistical distance to the distributions induced by expert samples (will denote as Eq. (6)), and simplified version Sim(s, s') as the distance to the distribution induced by the averaged logits of expert samples for computing efficiency (will denote as Eq. (7)). In this subsection, we compare the effect of these two different similarity version on two representative tasks, Sortpole-Swingun and Finger-Spin using and expert trajectories. The learning results and corresponding averaged similarities during training are shown in Fig. 10. In our experiment, we let all similarity begins from the same initial value (shown in Tab. 2) to eliminate the effect of different value scale. From Fig. 10, we observe two facts: 1) Eq (7) and Eq (6) achieve similar similarity trends and Eq (6) obtains larger similarity values average 2) Eq (6) makes learning less aggressive, leading to slower initial learning (as Finger-Spin), but can be more stable and may converge to better ultimate results, especially when there are fewer expert data.
15

---

Published as a conference paper at ICLR 2023
Carpole Swingup
Finger Spin
Cartpole Swingup
Finger Spin
353
6.37))
Cploode Retum
Eptzoce Rerum
soce
Rest Ai sct?)
302
Pacsel Je cart
wee
Expart
10
15
Frames
Frames
Framas
Frames
Carppole Swingup
Finger Spin
Cortpole Swingup
Finger Spin
235
3.5
Frames
15
0.5
frames
Frames
1.0
1.5
(a) 10 expert trajectories.
(b) expert trajectories.
Figure 10: Comparing different similarity formulation Eq (7) and Eq (6).
Share Encoddr All
Independent: Encoder All
Patch AlL w.o. Reg
Potcholl Weight
FatchAALL Bonus
Epicdd Return
Ep Returrn
600
600
Epands
100
3
aze
200
200
200
200
0.3
0.5
1.5
6.0
1.8
1.5
6.5
3.0
3.2
0.5
1.0
1-5
1.0
140
Figure 11: Different data augmentation strategics for discriminators on Finger- over 5 random seeds.
with 10 expert trajectorics
B.5 DIFFERENT DATA AUGMENTATION STRATEGIES FOR DISCRIMINATORS
In claim, we argue that the discriminator of the encoder-based solutions can be sensitive to changes in local behavior, and may only focus on insufficient regions that support its overall judg- ment. In our experiment, we do observe that both shared-encoder AIL and independent-encooder AIL lost come regions of given image samples. However, in the deep learning and computer vision community, researchers have developed various kinds of data augmentation methods to alleviate such a problem by improving the robustness of the discriminator. Although in our experiments, we have included Random-Shi ft augmentation for the training of the discriminator, which has been shown to be the most useful data augmentation technique in pixel-based reinforcement learn- ing problems (Kostrikov et al., 2019; Laskin et al., 2020b). We still wonder different data aug- mentation strategies can improve VIL performance. To verify this, we try Random Cut Out, Sandom-Crop-Resll ize, and RandAugment (Cubuk et al., 2020) (without color operations since we are augmenting stack of images) strategy using the two encoder- based baselines the Finger-Spin domain, which we think may be the most useful augmentation strategy for the discrim- inator differentiate samples the DeepMind Control Suite testbed. The results shown in Fig. 11. Compared to observations in Fig. 4, it is easy to conclude that augmentation strategies are not the key to making independent-encoder AIL successful, but patch rewards can make discrimi- nator works directly on pixels; in comparison, Random Random-Crop- Resize and Random Shi i f lead share-encoder AIL into similar training efficiency, while RandAugment may promote the ef- ficiency and final performance, causing more instability. And it turns out the best augmentation strategy for PatchAIL and its variants is the default Random-Sh ft, leading stable, efficient, and better performance. Nevertheless, we can still conclude the effectiveness of learning with patch rewards. But this also motivates us that improving the augmentumentttion strategy with shared- encoder AIL may be worth further investigation.
B.6 VISUAL ATTENTION WITH GRADCAM
We also try interpreting the rewards given by discriminators using GradCAM Selvaraju et al. (2017), a popular visual explanation tool. Specifically, we use GradCAM to generate the response score regarding the expert label and the non-expert label for each method. The comparisons are in Fig. 12. We find the results are mainly hard to interpret, indicating that such a tool may not be appropriate to explain the reward provided by AlL methods. But the proposed patch rewards naturally provide
16

---

Published as a conference paper at ICLR 2023
Observations
TIEE 1owand-
ThE reard 2.05
IOcIAWIE 1.23
Label
Expert
PatchAIL-W
Non Expert
Label
AOA
Label Expert
Chared-Encoder
Non Expert
tout EX
LONGINE PAS
Label
P63016 1.40
LONT NJR0
Independent-Enng Encoeed
Label
Expert
Non Expert
Label
Cartpole Swingup
Cheetah Run
Finger Spin
Happer Stand
Quadruped Run
Walker Stand
Walker Walk
Figure 12: Visual interpretation by GradCAM for each method. PatchAlL and shared-encodes AlL show an ad- vantage over independent-enccon encoder in terms of focusing more on the robot and focusing less on the background. This especially obvious five out of seven environments (Cartpole-Swinge Cheetah-Run Quadruped- RUn, Walker-Stand and Walker-Walk).
interpretability. We can only observe from Fig. 12 that PatchAIL and shared-encoder AlL show advantages over independent-eneoder in terms of focusing more on the robot and focusing less on the background. This is especially obvious in five out of seven environments (Cartpole-Swingup, Cheetah-Run, Quadruped-Rung Run, Walker-Stand and Walker- Walk).
17

---

Published as a conference paper at ICLR 2023
Freeway
Boxing
Pong
Krull
Episode Return
-0
Ep spisode Return
Episode Return
10
Episode Return
-0
4000
10
Putchikk
-10
2000
Shared Encader
-50
-20
Frames
16b
Frames
Frames
1e6
Frames
1e6
Figure 13: Learning curves on 4 Atari tasks over random seeds using 20 expert trajectories. The curve and
shaded zone represent the mean and the standard deviation, respectively.
B.7 ATARI EXPERIMENTS
Except for DMC tasks, we test our methods on another commonly used pixel-based reinforcement learning benchmark, Atari. We use the data collected by RLUnplugged Gulcehre et al. (2020) as ex- pert demonstrations. In particular, we choose 20 episodes sampled by the last 3 checkpoints for imi- tation without using any reward and action information (except BC). We utilize Dueling DQN Wang et al. (2016) as our base learning algorithm, whose rewards are provided by discriminators. We test mentioned PatchAIL methods against baseline methods on four tasks, including Freeway, Boxing, Pong and Krull, and show the learning curves in Fig. 13. We train each method in less than 10M frames
converged. The conclusion remains consistent compared with those on DMC tasks. Concretely, PatchAIL methods obviously observed to have better learning efficiency all tasks, and notably have strong improvement on harder tasks.
C IMPLEMENTATION DETAILS
C.I BASELINE ARCHITECTURES
The architecture of two encoder-based Adversarial Imitation Learning (AIL) baselines and the pro- posed PatchAIL solution are illustrated in Fig. 14. As for the detailed encoder architecture, all methods use the default structure for the policy/critic that works in Haldar et al. (2022): 4-layer CNN with kernels [3x3, 32, 2. 0]. [3x3, 32, 1, 0}, [3 x, 32, 1, 0], [3x3, 32, 1, 0] (represented as [size, channel, stride, padding]). resulting in 35 35 feature map and each entry has a receptive field of 15. Note that the encoder-based baselines also utilize such an encoder for the discriminator.
After encoding, the feature maps are flattened as a 39200-length vector, which is then put into a 4- layer MLP actor as and a 4-layer critic. For mccoer--ased baselines, the discriminator also follows the structure in Haldar et al. (2022) as a 3-layer MIP which outputs scalar result. Regarding PatchAIL methods, the discriminator used for DMC domain is 4-layer FCN with default kernels [4x4, 32, 2, 1]. [4x4, 64, 1]. [4x4, 128, 1, 1] and [4x4, I, 1, 1] (represented as [size, channel, stride, paddingl), resulting in 39 39 patches and each has receptive field of 22. As for Atari domain, the discriminator is a 4-layer FCN with default kernels [4x4, 32, 2. 11, [4x4, 64, 2, 1], [4x4, 128, 1, I] and [4x4, 1, 1, (represented [size, channel, stride, paddingl). resulting in 19 x 19 patches and each has a receptive field of 34. In our experience, the discriminators should have reasonable structure (i.e., the resulting patch number and the receptive field size of each patch relative to the task) to achieve good results.
C.2 IMPORTANT HYPERPARAMETERS
We list the key hyperparameters of AIL methods used in our experiment in Tab. 2. The hyperparam- eters are the same for all domains and tasks evaluated in this paper.
18

---

Published as a conference paper at ICLR 2023
Roward
Discriminator
Policy
Cnbc
Discriminator
Policy
Critic
Policy
Crtti
Patch Disortminator
Patch Dicrrminator
Inforcnoe
Image Encoder
Imege
Image Encoder
Image Encoder
Hinenor
Trih
Exper
Sanpled
Dali
(a) Shared-encoder AlL structure. (b) Independent-cnconn encoder AlL structure.
(c) PatchAlL structure.
Architectures for two visual imitation learning baselines (not our proposal) and the proposed method. (a) The discriminator shares the encoder with the policy and the critic. The encoder of the policy/critic is only updated by the critic loss. (b) The encoder of the discriminator is independent of the policy/critic. (c)
Figure 14:
PatchAIL, where the discriminator also independent the critic's encoder.
Method
Parameter
Value
Common
Default replay buffer size Learning rate
150000
1e-4
Discount y
0.99
Frame stack n-step returns
3
Action repeat
2
Mini-batch size
256
Agent update frequency
2
Critic soft-update rate
0.01
Feature dim
50
Hidden dim
1024
Optimizer
Adam
Exploration steps
2000
DDPG exploration schedule
linear(1,0.1 (Meear110.1,5000000
Gradient penalty coefficient
10
Target feature processor update frequency(steps)
20000
Default Reward scale
1.0
Default data augmentation strategy
random shift
PatchAlL-W
initial value
1.3
Replay buffer size
1000000
PatchAIL-B
initial value
0.5
Replay buffer size
1000000
Reward scale
0.5
Table 2: List of hyperparameters for all domains.
19


Reviews:
[{"review_id": "uB1ClZW2gN", "reviewer": "Reviewer_psvi", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "This paper describes an adversarial imitation learning algorithm that learns a discriminator with patch pixel rewards rather than directly from images. Recent work in the vision deep learning community have shown that using patches is effective in embedding pixels. The authors argue that patch pixels rewards are more informative than using a shared encoder for the discriminator and critic or directly learning from images.\n\nComments:\n- Is the encoder layer CNN and kernels the same size for the independent-Encoder and Shared-Encoder compared to the patch (obviously ignore the 1x1 kernel at the end)? \n- What is \\boldmath{R}_{pxp} in equaiton 5?\n- Which of the 7 tasks were used for tuning hyperparameters versus for evaluation? Or where all hyperparameters run on all 7 tasks and the best hyperparameter across tasks chosen?", "summary_of_the_review": "I recommend this paper for marginally above the acceptance threshold. Though the technical novelt is limited, the empirical novelty is not. The authors do a thorough job empirically justifying the decision made in the paper.", "strength_and_weaknesses": "Strength:\n- The idea of training a  discriminator with patches to gather dense rewards straightforward\n- The paper is mostly well-written and easy to understand\n- The authors perform thorough experiments comparing all relevant baselines to showcase the significance of their idea.\n- The authors perform an ablation study in figure 6 (b) on various patches, empirically showing why patch size 39x39 was chosen.\n- Additional ablation studies in the appendix show the breadth of empirical analysis.\n\nWeakness:\n- Though training a discriminator with patches is easy and intuitive, the best way to aggregate patches into a reward is not apparent. The authors proposed three aggregation metrics: mean, max, min, and median. From this set of metrics mean was chosen as the default aggregation metric. But from the ablation study, figure 6 (a), mean seems to perform significantly worse than the other metrics. It is not apparent why mean was chosen.\n- There are 3 variants of PatchAIL: PatchAIL w.o. Reg, PatchAIL-W and PatchAIL-B. It is not apparent which algorithm variant is better and why regularization helps. The authors motivate patch regularization, but the experiment evaluation does not justify the motivation.\n- The approximation of equation (6) with equation (7) is unclear.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "eUABI9s5HVm", "reviewer": "Reviewer_AJEj", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "The paper proposes a general class of reward function defined by image patches in observations, which enables effective learning guidance when training an VIL agent, visual interpretability, and better training stability. The method is evaluated on a range of DeepMind Control tasks and achieves favorable performance compared baseline methods. ", "summary_of_the_review": "The method offers a simple AIL framework for training pixel-based RL agents. The main concern is that the empirical benefits, both in terms of asymptotic performance and sample efficiency, is not high compared to existing works. The paper claims that the method offers better visual explanations, but this particular point is not well-supported by Figure 5 where the baseline method already provides a similar level of interpretability. \n\n--- Post Rebuttal ---\nAfter the rebuttal responses and reading other reviewers' comments, I increased my score from 3 to 5 with more clarifications from the authors. My main concern is that the asymptotic performance gain over the baselines is specific to particular environments, and the visualization analysis is spuriously correlated with motion and does not distinguish positive and negative contributions to the reward. ", "strength_and_weaknesses": "Advantages\n* The method is straightforward and intuitively corresponds to the observation that comparing patches in image observations with the expert demonstration can provide useful learning signals. \n* The method is evaluated on standard pixel-based benchmarks and achieves comparable performance compared to its baselines. \n* The paper offers a discussion on the visual interpretation of the proposed reward function. \n\nWeaknesses\n* The proposed method is highly similar to previous AIL methods. This itself is not a problem if it provides large performance gain. However, gain is small compared to Shared-Encoder AIL according to Figure 4, while introducing several extra parameters such as aggregation function class and patch size. \n* The other claimed advantage of the method is visual interpretability. However, according to Figure 5, both Shared-Encoder as will as the proposed method have their attention maps focusing on the agent.\n\nOthers\n* There are two typos in equation (4). It seems to be (s, s') ~ \\pi and (s, s') ~ \\pi_E. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "X42Y9nWOLr", "reviewer": "Reviewer_nR86", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "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": "This paper investigates the problem of visual adversarial imitation learning. Based on Generative Adversarial Imitation Learning (GAIL), the authors re-formulate the discriminator and make it focus on the local regions of the image. Specifically, the visual demonstrations are divided into patches, the discriminator can predict reward for each patch separately. Empirical results show the performance of PatchAIL can outperform the compared methods in several DMC tasks.", "summary_of_the_review": "Overall, the paper is well-written and easy to follow. The idea is novel and reasonable. Empirical results are good and convincing. However, there are still some concerns about the experiment, as listed in the weaknesses part. As a result, I provide my initial score as *borderline accept*.", "strength_and_weaknesses": "Strengths:\n* This paper investigates a practical setting. Visual adversarial imitation learning is quite difficult and few works can solve this problem.\n* The idea is novel. Though dividing images into patches is nothing new, it is a good try to use it in imitation learning.\n* Good and convincing empirical results.\n\nWeaknesses:\n* In 5.3, the authors claim that using 'mean' is the one-of-the-best choice among aggregation functions while 'max' tends to fail in both domains. But in my opinion, once the authors use 'mean' as the aggregation function, it seems that there is no difference between using the patch discriminator and a vanilla discriminator. Also, in Figure 6(a), the curve clearly shows that using 'mean' as an aggregation function results in the worst performance. This seems contradictory to what the authors describe in the paper.\n* The effect of patch regularization is not clear, as shown in Figure 4.\n* Would you mind providing some intuitive analysis about why a too-large or too-small patch could hurt the performance? And In my opinion, the best patch size for different agents should be various due to their different body sizes in the image. Can the authors elaborate on this point more?\n* Do PatchAIL have to use (s,s') as the input of the discriminator? I think this adds a strong constraint on the used expert demonstrations. I believe there are some works about adversarial imitation learning with observations only use a single state as the discriminator input. What if we use a single (s) or (s,a) as the discriminator input?\n* Experiments can be conducted on other visual imitation learning tasks such as Atari Games.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "SYRraayyzw", "reviewer": "Reviewer_52Vk", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "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": "Visual imitation learning algorithm using patch rewards is proposed. For patch rewards, input images from agent and expert are decomposed into small patches and classified by using multiple patch-wise discriminator. The output values of patch-wise discriminators are postprocessed (by $h$) and aggregated (by $Aggr$, e.g., $\\mathrm{mean}$, $\\mathrm{median}$, $\\mathrm{min}$) to generate a scalar reward that is suitable for the RL inner loop of adversarial imitation learning. Authors also introduce patch regularizers that aim to maintain patch-wise information disappearing through the aforementioned reward aggregation. The proposed algorithm shows its effectiveness in DeepMind Control Suite benchmarks and some amount of visual explanability through learned patched rewards.", "summary_of_the_review": "Overall, I believe this is a good submission and far above the acceptance boundary. I only have a few questions below:\n- The approximation from Eq (6) to Eq (7) is unclear to me, especially how $\\mathrm{min}$ can be converted into $\\mathbb{E}$ in the bracket of Eq (7).\n- Adding empirical analysis w.r.t. the number of expert demonstrations will be interesting. I think that may affect Eq (7) in the sense that the expectation in Eq (7) and its approximation can be affected by the number of expert demonstrations.\n- In Figure 6, (a), it seems that $\\mathrm{mean}$ doesn’t seem to work well, whereas $\\mathrm{max}$ works well. This is different from what authors state in Section 5.3. I guess labels in Figure 6 are simply mistyped. ", "strength_and_weaknesses": "### Strength\n- The idea seems to be easily applicable to any visual imitation learning tasks. \n- One can visualize the gap between agent and expert behaviors through images. \n### Weaknesses\n- Motivation and theoretical derivation on patch regularization is unclear. It seems that patch regularization is effective only for some tasks where algorithms w/o patch regularization is unstable. \n- Analysis on how the number of expert demonstrations affects the performance is missing, e.g., sample efficiency w.r.t. the number of expert demonstrations, whether the patch reward can be effectively trained for small or large numbers of expert demonstrations, etc.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster