CEIP: Combining Explicit and Implicit Priors for Reinforcement Learning with Demonstrations
Authors: Kai Yan
Alex Schwing
Yu-Xiong Wang
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
imitation learning
normalizing flow
action prior

Abstract: Although reinforcement learning has found widespread use in dense reward settings, training autonomous agents with sparse rewards remains challenging. To address this difficulty, prior work has shown promising results when using not only task-specific demonstrations but also task-agnostic albeit somewhat related demonstrations. In most cases, the available demonstrations are distilled into an implicit prior, commonly represented via a single deep net. Explicit priors in the form of a database that can be queried have also been shown to lead to encouraging results. To better benefit from available demonstrations, we develop a method to Combine Explicit and Implicit Priors (CEIP). CEIP exploits multiple implicit priors in the form of normalizing flows in parallel to form a single complex prior. Moreover, CEIP uses an effective explicit retrieval and push-forward mechanism to condition the implicit priors. In three challenging environments, we find the proposed CEIP method to improve upon sophisticated state-of-the-art techniques.

CEIP: Combining Explicit and Implicit Priors for Reinforcement Learning with Demonstrations
Kai Yan
Alexander G. Schwing
Yu-Xiong Wang
University of Illinois Stbana-Champaign
(kaiyan3, aschwing, yxw)@illinois edu https://g wttp.cccmm2003722///////
Abstract
Although reinforcement learning has found widespread use in dense reward set- tings, training autonomous agents with sparse rewards remains challenging. To address this difficulty, prior work has shown promising results when using not only task-specife demonstrations but also task-agnostic albeit somewhat related demonstrations. most cases, the available demonstrations are distilled into an implicit prior, commonly represented via a single deep net. Explicit priors in the form of a database that can be queried have also been shown to lead to encouraging results. To better benefit from available demonstrations, we develop method to Combine Explicit and Implicit Priors (CEIP), CEIP exploits multiple implicit priors in the form of normalizing flows in parallel to form single complex prior. Moreover, CEIP uses an effective explicit retrieval and push-forward mechanism condition the implicit priors. In three challenging environments, find the proposed CEIP method to improve upon sophisticated state-of-the-an techniques.
Introduction
Reinforcement learning (RL) has found widespread use across domains from robotics [57] and game AI [44] to recommender systems [6]. Despite its success, reinforcement learning is also known be sample inefficient. For instance, training a rohot arm with sparse rewards to sort objects from scratch still requires many training steps if it all feasible [46].
To increase the sample efficiency of reinforcement learning, prior work aims to leverage demonstra- tions [4, 34, 40]. These demonstrations can be task task-specific [4, 17], i.e., they directly correspond to and address the task of interest. More recently, the use task-agnostic demonstrations has also been studied [14, 16, 34, 46], showing that demonstrations for loosely related tasks enhance sample efficiency of reinforcement learning agents.
To benefit from either of these two types of demonstrations, most work distills the information within the demonstrations into an implicit prior, by encoding available demonstrations
deep net. For example, SKiLD [34] and FIST [16] use variational auto-encoder (VAE) to encode the "skills," i.e. action sequences, in a latent space, and train a prior conditioned on states based on demonstrations to use the skills. Differently, PARROT [46] adopts a state-conditional normalizing flow to encode transformation from latent space to the actual action space. However, the idea of using the available explicit prior has not received a lot of attention. Explicit priors enable the agent to maintain a database of demonstrations, which can be used to retrieve state-action sequences given agent's s current state. This technique has been utilized in robotics [32, 47] and early attempts of reinforcement learning with demonstrations [4]. It was also implemented as a baseline in [14]. One notable recent exception is FIST [16]. which queries database of demonstrations using the current state to retrieve a likely next state. The use of an explicit prior was shown to greatly enhance the
demonstrations
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

performance. However, FIST uses pure imitation learning without any RL, hence losing the chance for trial and remedy if the imitation is not good enough.
Our key insight is to leverage demonstrations both explicitly and implicitly, thus benefiting from both worlds. To achieve this, we develop CEIP, method which combines explicit and implicit priors. CEIP leverages implicit demonstrations by learning a transformation from a latent space to the real action space via normalizing flows. More importantly, different from prior work, such as PARROT and FIST which combine all the information within a single deep net, CEIP selects the most useful prior by combining multiple flows in parallel to form a single large flow. To benefit from demonstrations explicitly, CEIP augments the input of the normalizing flow with likely future state, which is retrieved via a lookup from a database of transitions. For an effective retrieval, we propose a push-forwar sohh-fowarrd technique which ensures the database to return future states that have not been referred to yet, encouraging the agent to complete the whole trajectory even if fails on a single task.
We evaluate the proposed approach on three challenging environments: fetchreach [36], kitchen [11], and office [45]. In each environment, we study the use of both task-specific and task-agnostic demon- strations. We observe that integrating an explicit prior, especially with our proposed push-forward technique, greatly improves results. Notably, the proposed approach works well on sophisticated long-horizon robotics tasks with few, or sometimes even one task-specific demonstration.
Preliminaries
Reinforcement Learning, Reinforcement learning (RL) aims to train an agent to make the 'best' decision towards completing a particular task in given environment. The environment and the task are often described as Markov Decision Process (MDP), which is defined by tuple AS, A, r, 7)- In timestep of the Markov process, the agent observes the current state 8t E S, and executes an action a E A following some probability distribution, i.e., policy x (a, ((a))s) E (AA), where ((A) denotes the probability simplex over elements in space A. Upon executing action a, the state of the agent changes to $t+1 following the dynamics of the environment, which are governed by the transition function T [st,, at) S x A (SS). Mcanwhile, the agent receives a reward r(st, ar) E R. The agent aims to maximize the cumulative reward Σ, 7"r(s,, az), where y [0, 1] is the discount factor. One complete run in an environment is called an episode, and the corresponding state-action pairs {($1, 01), ($2, 02).... ... } form trajectory T.
Normalizing Flows. A normalizing flow [24] a generative model that transforms elements 20 drawn from a simple distribution Pa. e.g., Gaussian, to elements op drawn from a more complex distribution Pa- For this transformation, a bijective function / is used, i.e., ag = [(=o). The use of a bijective function ensures that the log-likelihood of the more complex distribution at any point is tractable and that samples of such a distribution can be easily generated by taking samples from the simple distribution and pushing them through the flow. Formally, the core idea of normalizing flow can be
05 ()))(a) lommo where |-| is the determinant (guaranteed du
summarized via (a (f (ao)) log
positive by flow designs), is a random variable with the desired more complex distribution, and & is a random variable governed by a simple distribution. To efficiently compute the determinant of the Jacobian matrix of f-1, special constraints are imposed on the form of f. For example, coupling flows like RealNVP [8] and autoregressive flows [31] impose the Jacobian of f to be triangular.
3
CEIP: Combining Explicit and Implicit Priors
3.1 Overview
As illustrated in Fig. 1, our goal is to train autonomous agent to solve challenging tasks despite sparse rewards, such as controlling a robot arm to complete item manipulation tasks (like turning on switch opening cabinet). For this we aim to benefit from available demonstrations. Formally, -TS), }, where T!' is the i-th trajectory of we consider a task-specific dataset Dis = {T|' -TS -
f1,2, 3, ......... n)), where D. = { .i...., Thyy subsumes the demonstration trajectories for the i-th task in the task aagostic dataset. Each trajectory T = {($1, 01), (#2, az),... in the dataset is a state-action pair sequence of a complete episode, where is the state, and is the action. We assume that the number of
the task-specift dataset, and task task-agnostic dataset DIA = Dul

---

Cuderrn Tall-agadd Drone
Traning D03
Conto ning Tows
Downntwerm
Row A Auac!
Follyy PlAMO
Avuai
Roo (ol
1112
Eirose
How A
Deroo ta udkin
Many 20d
sporli
Enwronment
Baak 11
Figure 1: Overview of proposed approach, CEIP. Our approach can be divided into three steps: a) cluster the task-agnostic dataset into different tasks, and then train one flow on each the v lasks of the task-agnostic dataset; b) train a flow on the task-specific dataset, and then train the coefficients to combine the n + flows into one large flow fis. which is the implicit prior; c) conduct reinforcement learning on the target task; for cach timestep. perform datasct lookup in the task-specific dataset to find the state most similar to current state 8, and return the likely next state Seeat in the trajectory, which is the explicit prior.
available task-specific trajectories is very small, i.e., Ei= my >> m, which is common in practice. For readability, we will also refer Drs using D,+1-
Our approach leverages demonstrations implicitly by training a normalizing flow fTs. which trans- forms the probability distribution represented by policy π ss) over simple latent probability space Z, i.c., 2 E Z, into a reasonable expert policy over the space of real-world actions A. As before, s is the current environment state. Thus, the downstream RL agent only needs to learn policy (2)ss) that results in a probability distribution over latent space Z, which is subsequently mapped via the flow frs to a real-world action A. Intuitively, the MDP in the latent space is governed by a less complex probability distribution, making it easier to train because the flow increases the exposure of more likely actions, while reducing the chance that less-likely action is chosen. This is because the flow reduces the probability mass for less likely actions given the current state.
Task-agnosti demonstrations contain useful patterns that may be related to the task at hand. However, not all the task-agnostic data are always equally useful, as different task-agnostic data may require to expose different parts of the action space. Therefore, different from prior work where all data are fed into the same deep net model, we first partition the task-agnostic dataset into different groups according to task similarity so to increase flexibility. For this we use classical k-means algorithm. We then train different flows f. on each of the groups, and finally combine the flows via learned coefficients into single flow fis. Beneficially, this process permits to expose different parts of the action space as needed and according to perceived task similarity.
Lastly, our approach further leverages demonstrations explicitly, by conditioning the flow not only on the current state but also on a likely next state, to better inform the agent of the state it should try to achieve with its current action. In the following, we first discuss the implicit prior of CEIP in Sec. 3.2; afterward we discuss our explicit prior in Sec. 3.3, and the downstream reinforcement learning with both priors in Sec. 3.4.
3.2 Implicit Prior
To better benefit from demonstrations implicitly, use a 1-layer normalizing flow the backbone of our implicit prior. It essentially corresponds to conditioned affine transformation of a Gaussian distribution. We choose a flow-based model instead of a VAE-based one for two reasons: 1) as the dimensionality before and after the transformation via a normalizing flow remains identical and since the flow is invertible, the agent guaranteed to have control over the whole action space. This ensures that all parts of the action space are accessible, which is not guaranteed by VAE-based methods like SKiLD or FIST; normalizing flows, especially coupling flows such as RealNVP [8], can be easily stacked horizontally, so that the combination of parallel flows is also a flow. Among feasible flow models, we found that the simplest 1 llayeer flow suffices to achieve good results, and even more robust in training than a more complex RealNVP. Next, in See. 3.2.1 we first introduce details regarding the normalizing flow fú, before we discuss in Sec. 3.2.2 how to combine the flows into one flow frs applicable to the task for which the task-specifi dataset contains demonstrations.
3

---

u (a, Dj
NN
NN
d())
1I (u,a) Dys
de(1))
l7s
(x)(),(10) O d((w)
NN
a(w).
Σ (eepper(W)) 0 2+
+ 4 a((ww))
Figure 2: An illustration of how combine different flows into one large flow for the task task-specific dataset. Each red block of "NN" stands for neural network. Note that ea ci(u) and (u) are vectors, while pr and a are the i-th dimension of u(() and M(u).
3.2.1 Normalizing Flow Prior. For each task i in the task-agnostic dataset, i.c., for each D.. we train a conditional I-layer normalizing flow f; u) = a which maps a latent space variable 2 E R to an action a E RR,, where q is the number of dimensions of the real-valued action vector. We let refer a conditioning variable. In our case u either the current environment state (if no explicit prior is used) or concatenation of the current and a likely next state [s, Sneal (if an explicit prior is used). Concretely, the formulation of our I-layer flow
(1) where cu(u)= E RR, du(u)= E 19 are trainable deep nets, and refers to the Hadamard product. The exp function is applied elementwise. When training the flow, we sample state-action pairs (without explicit prior) or transitions (with explicit prior) (u, a) from the dataset D,, and maximize the log-likelihood (n.a)DD, log p(a|u); refer to [24] for how to maximize this objective.
fz u) = (x)(c:(u)) @ + du)),
In the discussion above, assume the decomposition of the task-agnostic dataset into tasks to be given. If such a decomposition is not provided (e.g., for the kitchen and office environments in our experiments), perform k-means clustering divide the task-agnostic dataset into different parts. The clustering algorithm operates on the last state of a trajectory, which is used to represent the whole trajectory. The intuition is two-fold. First, for many real-world MDPs, achieving particular terminal state is more important than the actions taken [12]. For example, when we control a robot to pick and place items, we want all target items to reach the right place eventually: however, we do not care too much about the actions taken to achieve this state. Second, among all the states, the final state is often the most informative about the task that the agent has completed. The number of clusters in the k-means algorithm is a hyperparameter, which empirically should be larger than the number of dimensions of the action space. Though we assume the task-agnostic dataset is partitioned into labeled clusters, our experiments show that our approach is robust and good results are achieved even without a precise ground-truth decomposition.
In addition to the clusters in the task-agnostic dataset, we train a flow ( w) = a on the task- specific dataset D = Drs, using the same maximum log- log-likelihood loss, which is optional but always available. This is not necessary when the task is relatively simple and the episodes are short (e.g., the fetchreach environment in the experiment section), but becomes particularly helpful in scenarios where some subtasks of a task sequence only appear in the task-specific dataset (e.g., the kitchen environment).
3.2.2 Few-shot Adaptation. The flow models discussed in Sec. 3.2.1 learn which parts of the action space to more strongly exposed from the latent space. However, not all the flows expose useful parts of the action space for the current state. For example, the target task needs the agent to move its gripper upwards at a particular location, but in the task-agnostic dataset, the robot more often moves the gripper downwards to finish another task. In order to select the most useful prior, we need to tune our set of flows learned on the task- aggostic datasets to the small number of trajectories available in the task-specific dataset. To ensure that this does not lead to overfitting as only a very small number of task-specific trajectories are available, we train a set of coefficients that selects the flow that works the best for the current task. Concretely, given all the trained flows, we train set of coefficients to combine the flows fi to fn trained on the task- data, and also the flow fn+1 trained on the task- task-specific data. The coefficients select from the set of available flows the most useful one. To

---

achieve this, we use the combination flow illustrated in Fig. 2 which is formally specified as follows:
n+1 Σ pilu) px(()exp(c(())) + =
n+1
fis(z;2)=
Σ ((()dd())
Here, m:(u) R. A,(U) R are the i-th entry of the deep nets u(u) E R" M(u) E R'+1, respectively, which yield the coefficients while the deep nets a; and d, are frozen. As before, the exp function is applied elementwise. We use softplus activation and offset at the output of j to force p); (u) 2 10-1 for any for numerical stability. Note that the combined flow frs consisting of multiple 1-layer flows is also 1-layer normalizing flow. Hence, all the compelling properties over VAE-based architectures described at the beginning of Sec. 3.2 remain valid. To train the combined flow, we use the same log likelihood loss ((u.a)>DDs Drs log p(a (u) that for training single flows. Here, optimize the deep nets μu()) and X(u) which parameterize fis-
Obviously, the employed combination of flows can be straightforwardly extended to a more compli- cated flow, e.g.. a RealNVP [8] or Glow [22]. However, we found the discussed simple formulation to work remarkably well and be robust.
3.3 Explicit Prior
Beyond distilling information from demonstrations into deep nets which are then used as implicit priors, we find explicit use of demonstrations to also be remarkably useful. To benefit, we encode future state information into the input of the flow. More specifically, instead of sampling (s. a) -apirs from a dataset D for training the flows, we consider sampling a transition (s, a, Snex) from D. During training, we concatenate s and Spext before feeding it into a flow, i.c., u = [s, Spex] instead of u s.
However, we do not know the future state Saexi when deploying the policy. To obtain an estimate, we use task-specific demonstrations as explicit priors. More formally, we use the trajectories within the task-specific dataset Drs as database. This is manageable as we assume the task-specific dataset to be small. For each environment step of reinforcement learning with current state 8, we perform lookup, where is the query, states Skey in the trajectories are the keys, and their corresponding next state Snext is the value. Concretely, we assume &pext belongs to trajectory T in the task-specific dataset Drs, and define Spett the result of the database retrieval with respect to the given query s, i.c.
Seet argmin Snoz (8ky),, 9, [ Skey
8)² + C (( where
1 if 3s/nex E T, s.t. S'reu no earlier than next in 7 and has been retrieved,
(3)
8( ((0ee))
0 otherwise.
In Eq. (3), C is constant and 5 is the indicator function. We set = [s, Smexi the condition, feed it into the trained flow fTS, and map the latent space element obtained from the RL policy to the real-world action a. The penalty term 5 is a push-forward technique, which aims push the agent to move forward instead of staying put, imposing monotonicity on the retrieved 8next- Consider an agent particular state s and flow Srs, conditioned on [s, Snex]] which maps the chosen action to a real-world action a that does not modify the environment. Without the penalty term, the agent will remain the same state, retrieve the same likely next state, which again maps onto the action that does not change the environment. Intuitively, this term discourages 1) retrieving the same state twice, and 2) returning to earlier states given trajectory. In our experiments, we set C =
3.4 Reinforcement Learning with Priors
Given the implicit and explicit priors, we use RL to train a policy ( ss) to accomplish the target task demonstrated in the task-specific dataset. As shown in Fig. 1, the RL agent receives state and provides a latent space element z. The conditioning variable of the flow retrieved via the dataset lookup described in Sec. 3.3 and the real- world action a is then computed using the flow. Note, our approach is suitable for any RL method, i.c., the policy (2[s) can be trained using any RL algorithm such proximal policy optimization (PPO) [43] soft-actor-critic (SAC) [15].
4 Experiments
In this section, we evaluate our CEIP approach on three challenging environments: fetchreach (Sec. 4.1), kitchen (Sec. 4.2), and office (Sec. 4.3), which are all tasks that manipulate robot arm.
5

---

In each experiment, we study the following questions: 1) Can the algorithm make good use of the demonstrations compared baselines? 2) Are our core design decisions (e.g., state augmentation with explicit prior and the push-forward technique) indeed helpful?
Baselines. We compare the proposed method three baselines: PARROT [46], SKiLD [34], and FIST [16]. In all environments, we use reward as our criteria (higher is better), and the results are averaged over 3 runs for SKiLD (which much slower to train than other baselines) and runs for all other methods unless otherwise mentioned. To differentiate variants of the PARROT baseline and our method, we use suffixes. We use "EX" to refer to variants that use an explicit prior, and "forward" to denote variants with the push-forward technique. For our method, if we train task-specific flow on Dis = Dn+1. we append the abbreviation "TS." " For PARROT, the use of the task- specciiciiicc data is indicated with "TS" and the use of task aggnotti data is indicated with "TA." See Table 3 and Table 4 for precise correspondence.
4.1 FetchReach Environment
Environment Setup. The agent needs to control a robot arm move its gripper a goal location in 3D space, and remain there. During an episode of 40 steps, the agent receives 10-dimensional state about its location and outputs a 4-dimensional action, which indicates the change of coordinates of the agent and the openness of the gripper. It will receive reward of 0 if it arrives and stays in the vicinity of its target. Otherwise, it will receive a reward of 1. This environment is harder version of the FetchReach-vl 1 robotics environment in gym [36], where we increase the average distance of the starting point to the goal, effectively increasing the training difficulty. Moreover, to test the robustness of the algorithm, we sample random action from a normal distribution at the beginning of each episode, which the agent executes for x steps before the episode begins. We use x U5, 20). For simplicity, we denote the goal generated with azimuth nd as "direction d' (e.g., direction 4.5). 4
Dataset Setup. We use trajectories from directions d E 1,..,... 7 as the task-agnostic data. Each task includes 40 trajectories, and each of the trajectories has 40 steps, i.e., 1600 environment steps in total. The task-specific datasets contain directions 4.5, 5.5, 6.5, and 7.5. (The robot cannot reach the other four 5 directions due to physical limits.) For each task-specific dataset, we use 4 trajectories. for a total of 160 environment steps.
Experimental Setup. For fetchreach, we use fully-connected deep net with one hidden layer of width 32 and ReL.U [1] activation as a standard "block" of our algorithm (each block corresponds to red "NN" rectangle in Fig. 2). We have pair of blocks for (s(s) and ds(s) for each flow fi- For flow training, we train 8 flows for 8 directions in the task-agnostic dataset without the explicit prior. We use a batch size of 40 and train for 1000 epochs for both each flow and the combination of flows, with gradient clipping at norm 10-4, learning rate 0.001, and Adam optimizer [21]. We use the model that has the best performance the validation dataset at the end of every epoch. For each dataset, we randomly draw 80% state-action pairs (or transitions in ablation) as the training set and 20% state-action pairs as the validation set. The combination of flow is also a block, which outputs both p(s) and A(s). See the Appendix for the implementation details of SKiLD, FIST, and PARROT. For each method with RL training, we use soft-actor-critic (SAC) [15] with 30K environment steps, a batch size of 256, and 1000 steps of initial random exploration. Unless otherwise noted, all other RL hyperparameters in all experiments use the default values of basel-e [38].
Main Results. Fig. 3 shows the results for different methods without explicit priors or task-specific single flow fa+1. In all four tasks, our method significantly outperforms the other baselines. This indicates that the flow training indeed helps boost the exploration process. Naïve reinforcement learning from scratch fails in most cases, which underscores the necessity of utilizing demonstrations to aid RL exploration. As this is a simple task with only a few wildly varied trajectories. adding flow for the task-specific dataset does not improve our method. Noteworthy, neither SKiLD nor FIST works on fetchreach. Their VAE- base architecture with each action sequence as the agent's 's output ("skill") can not be trained with the little amount of wildly varied data with short horizon. Flow-based models like ours and PARROT, which only consider the action of the current step instead of the action sequence, work better.
'The original PARROT in [46] is essentially PARROT+TA. It is straightforward to use PARROT directly on the task-specific dataset. Hence, we tried PARROT+TS and PARROT+((S+TA) as well.

---

(a) Direction 4.5
(b) Direction 5,5
(c) Direction 6.5
(d) Direction 7.5
Figure 3: Main performance results on the fetchreach environment for different directions, where the lines are the mean reward (higher is better) and shades are the standard deviation. FIST is represented by a dashed line as it docs not require RL.
(a) Direction 4.5
(b) Direction 5.5
(c) Direction 6.5
(d) Direction 7.5
Figure Ablation on the number of flows used our method. observe more flows to lead better performance, likely because expressivity increases which helps in fitting the expert policy.
Are more flows helpful for CEIP? Fig. 4 shows the performance of our method using different number of flows, which are trained on the data of the directions that are the closest to the task-specific direction (e.g., directions and 6 for flows with the target being direction 5.5). The result shows that within reasonable range, increasing the number of flows improves the expressivity and consequently results of our model. See Appendix for more ablation studies.
4.2 Kitchen Environment
Environment Setup. We use the kitchen environment adopted from D4RL [11], which serves as a testbed for many reinforcement learning and imitation learning approaches, like SPiRL [33], SKiLD [34], relay policy learning [14), and FIST [16]. The agent needs to control a 7-DOF robot arm to complete a sequence of four tasks (e.g., move the kettle, or slide the cabinet) in the correct order. The agent will receive +1 reward only upon finishing a task, and 0 otherwise. The action space 9-dimensional and the state space is 60-dimensional. This environment is very challenging, as control of the robot is needed and also long horizon of 280 timesteps is needed. Moreover, there is small noise applied to each agent action, which requires the agent to be robust.
Dataset Setup. We use two dataset settings, which are adopted from SKiLD and FIST (denoted as SiichenSSKLL. and Kitchen-FIST below). In Sitchnn-SKiLD we use 601 teleoperated sequences that perform a variety of task sequences as the task-agnostic dataset, and use only one trajectory for the task lask-specific dataset. In Kitchen-FISS we use part of the aank-sgn dataset (about 200 300 trajectories) that does not contain a particular task in the task-specific dataset, and use only one trajectory for the task-specific dataset. There are two different task-specific datasets in Kitchen- SKiLDD, and four different task-specifi datasets in Kitchen-FIST. The latter is significantly harder, as the agent must learn a new task from very little data. For simplicity, we denote them as "SKiLD-A/B" and "FIST-A/B/C/D" respectively. See Appendix C for details on each task.
Experimental Setup. We use k-means partition the task-agnostic datasets into 24 different clusters, and train 24 flows accordingly. For each flow, we use fully-connected network with 2 hidden layers of width 256 with ReLU activation as a "block" for our algorithm. For the combination of flows, we use fully-connected network with hidden layer of width 64 with ReLU activation. Each layer of deep net (except the output layer) described above has a 1D batchnorm function. The blocks are used analogously to the fetchreach environment. We use batch size of 256 for the task-agnostic dataset and a batch size of 128 for the task-specific dataset. Other training hyperparameters are identical to the fetchreach environment. For each RL training, use proximal policy optimization
7

---

12
SELD
7A
CEP
28033 5000 1VDE 3REDI JONE LU000
(a) Kitchen-SKiLD-A
(c) Kitchen-SKiLD-B
(e) Kitchen-FIST-A
SEAD
IN
12
2000 2009 1000 1880 U500 SO0M 150003000
20000
2200 2000 7500 Jaces Lidon zo00s 1100020000
(b) Kitchen-FIST-B
(d) Kitchen-FIST-C
(f) Kitchen-FIST-D
Figure 5: Comparison Kitchen-SKiLD and Kitchen-FIST environments.
(PPO) [43] for 200K environment steps, with update interval being 2048 (Kitchen-SKiLD) / 4096 Kitchen-FISTT 60 epochs per update, and a batch size of 64.
Main Results. Fig. 5 shows the main results Kitchen-SKiLD and Kitchen-FIST. Our method outperforms all other baselines in all of the 6 settings of the task-agnostic and task-specific datasets. For our method,
the task-specific single flow fn+1, explicit prior, and the push-forward technique. We compare to the original PARROT formulation. See Appendix D for ablation studies of PARROT with explicit prior and our method without fn+1 or explicit prior.
Does CEIP overly rely on the task flow if is used? One concern for our method could be: does the task-specific single flow dominate the model? Theoretically, when all flows are perfect, trivial combination of flows that minimizes the training objective is to set An+1 = 1, Pa+1 = for the task-speciff single flow, and X, O, 14 x 0 for * + 1. To study this concern, of the coefficient μ during an episode in Fig. 6. We observe that the single flow on the tasc-spec dataset is not dominating the combination of the flow, despite being trained the task-specific dataset. The blue curve with legend "TA-8' in Fig. 6 shows the coefficient for the 8th flow trained in the task-agnostic dataset. It exhibits increase of ju at the end of an episode, the last subtask in the target task is more relevant to the prior encoded in the Sth flow. Intuitively, over-reliance in our design (Fig. 8 in the Appendix) is discouraged, because of the softplus function and the positive offset applied on μ.. For over-reliance, all task task-agnostic flows fa with i E 21, n should have coefficient of [4: = O. which is hard approach due the offset of μ and softplus. In fact, degenerated CEIP is essentially PARROT+-S(+EE--orwall which is worse than our method but still a powerful baseline.
plot the change
Is reinforeement learning useful in cases with a perfect initial reward? Fig. shows the episode length of our method on SitcDen-SKiLD-A. Even if the reward is already perfect, the reinforcement learning process is still able to maximize discounted reward, which optimizes the path.
4.3 Office Environment
Environment and Dataset Setups. We follow SKiLD, where robot with 8-dimensional action space and 97 17-dimensional state space needs to put three randomly selected items on a table into three containers in the correct, randomly generated order. The agent will receive a +1 reward when it completes a subtask (e.g.. picking up an item, or dropping the item at the right place), and otherwise. This environment is even harder than the kitchen environment, as the agent must manipulate freely movable objects and the number of possible subtasks in the task-agnostic dataset is much larger than that in the kitchen environment. We use the same task-agnostic dataset SKiLD, which contains
8

---

40%
118
360
1006 10000 76740 Ha00 17960 ralle SPGIN Fallee
(a) Illustration of coefficient
(b) Average episode length
Figure a) Illustration the coefficient change of trained CEIP model during an episode of Kitchen-SKiLDA This CEIP model is trained with the task-specific single flow and without the explicit prior. "TA-8' is the 8-th single flow for the task task-agnostic dataset, and 'TS' is the single flow for the task specific dataset. The grey dotted lines the partition of different subtasks. b) Average episode length of our method the Sitchen-SKiLDD task. The episode ends immediately when all the tasks are completed; thus, shortening length means that RL helps find policy with more efficient completion of tasks.
2400 trajectories with randomized subtasks sampled from script policy. For the task-specific dataset, we use 5 trajectories for particular combination of tasks.
Experimental Setup. Similar to the kitchen environment, we use k-means over the last state of each trajectory and partition the task-agnostic dataset into 24 clusters. The architecture and training paradigm of the flow model are identical to those used in the kitchen environment. For RL training. we use PPO for 2M environment steps, with update interval being 4096 environment steps, 60 epochs per update, and batch size of 64. All other hyperparameters follow the kitchen environment setting. We run each method with 3 different seeds.
Main Results and Ablation. Fig. 7a shows the main result across different methods. Our method with explicit prior, push-forward tech- nique, and task-specific flow outperforms all baselines. FIST works well in this environment, probably because of two reasons: 1) there are sufficient number task-specific trajectories for the VAE-architccture, and 2) the office environment is less noisy than the kitchen environment. However, as FIST does not contain a reinforcement learning stage, it has no chance to improve on decent policy which could have been good start for an RL agent. Fig. 7b shows the ablation of our method. While the task-specific single flow fa+1 does not help in this environment, the explicit prior greatly improves results. Also, as illustrated, the reward curve of the variants with the explicit prior but without the push-forward skill does not grow, which is due to the agent getting stuck described at the end of Sec. 3.3. Fig. 7c shows the ablation result of PARROT, which also emphasizes that the explicit prior and push-forward skill greatly improve results.
(a) Main result
(b) Ours ablation
(c) PARROT ablation
Figure 7: Main result and ablation of our method and PARROT on the office environment.
5
Related Work
Reinforcement Learning with Demonstrations. Using demonstrations improve the sample efficiency of RL is an established direction [13, 26, 41, 55]. Recently, the use of task-agnostic demonstrations has gained popularity, as task-specific data need to be sampled from a particular expert and can be expensive to acquire [33, 34, 46). To utilize the prior, skill-based methods such as SPiRL [33], SKiLD [34], and SIMPL [30] extract action sequences from the dataset with a VAE-based model, while TRAIL [58] recovers transitions from task-agnostic dataset with uniformly randomly sampled action. Our method considers the situation where both task-agnostic and task-specific data exist and significantly improves results over prior work with similar settings, c.g., SKiLD [34].

---

Action Priors. An action prior is a common way to utilize demonstrations for reinforcement learning [34] and imitation learning [16]. Most work uses an implicit prior, where a probability distribution of actions conditioned on state is learned by deep net and then used to rule out unlikely attempts [3], to form a hierarchical structure [34, 46), or to serve as a regularizer for RI. training [34, 40], preventing the agent to stray too far from expert demonstrations. Explicit priors are less explored. They come in the form of nearest neighbors [2] (as in our work) or in the form of locally weighted regression [32]. They are utilized in robotics [2, 32, 42, 47] and early work of RL with demonstrations [4]. Another way to explicitly use demonstrations includes filling the buffer of offline RL algorithms with transitions sampled from expert dataset to help exploration [17, 29, 54]. Different from all such work, we propose a novel way of using both implicit and explicit priors.
Normalizing Flow. Normalizing flows are generative model that can be used for variational inference [23, 52] and density estimation [18, 31] and come in different forms: RealNVP [8], Glow [22], or autoregressive flow [31]. Many methods use normalizing flows in reinforcement learn- ing [20, 27, 28, 48, 49, 56] and imitation learning [5]. However, most prior work uses normalizing flows as strong density estimator to exploit richer class of policies. Most closely related to work is PARROT [46], which trains a single normalizing flow as an implicit prior. Different from our work, PARROT does not differentiate tasks among the task-agnostic dataset and does not use an explicit prior. More importantly, different from prior work, we develop a simple yet effective way to combine flows using learned coefficients. While there are some approaches that combine flows via variational mixtures [7, 35]. they have not been shown to succeed on challenging RL tasks.
Few-shot Generalization. Few-shot generalization [50] is broadly related, as model is first trained across different datasets, and then adapted to new dataset with small sample size. For example, similar to our work, FLUTE [51], SUR [9], and URT [25] use models for multiple datasets, which are then combined via weights for few-shot adaptation. Other methods have shared parameters across different tasks and only used some components within the model for adaptation [10, 37, 39,51, 51, 59]. While most work focuses on classification tasks, we address more complex RL tasks. Also, different from existing work, we found training of independent I-layer flows without shared layers to be more flexible, and free from negative transfer as also reported by [19].
Discussion and Conclusion
We developed CEIP, method for reinforcement learning which combines explicit and implicit priors obtained from task-agnostic and task- demonstrations. For implicit priors normalizing flows. For explicit priors we use a database lookup with a push-forward retrieval. In three challenging environments, show that CEIP improves upon baselines.
Limitations. Limitations of CEIP are as follows: 1) Training time. The use of demonstrations requires training decent number of flows which can be time-consuming. albeit mitigated to some extent by parallel training. 2) Reliance on optimality of expert demonstrations. Similar to prior work like SKiLD [34] and FIST [16]. our method assumes availability of optimal state-action trajectories for the target task. Accuracy of those demonstrations impacts results. Future work will focus on improving robustness and generality. 3) Balance between the degree of freedom and generalization in fitting the flow mixture. Fig. reveals that more degrees of freedom in the flow mixture improve results of CEIP. Our current design uses a linear combination which offers O(n) degrees of freedom (μ and X), where n is the number of flows. In contrast, too many degrees of freedom will result in overfitting. It is interesting future work to study this tradeoff.
Societal impact. Our work helps to train RL agents more efficiently from demonstrations for the same and closely related tasks, particularly when the environment only provides sparse rewards. If successful, this expands the applicability of automation. However, increased automation may also cause job loss which negatively impacts society.
Acknowledgemmntts. This work was supported in part by NSF under Grants 1718221, 2008387, 2045586, 2 2106825, MRI 1725729, NIFA award 2020-67021-32799 the Jump ARCHES endowment through the Health Care Engineering Systems Center, the National Center for Supercomputing Applications (NCSA) at the University of Illinois at Sibana-Champaal through the NCSA Fellows program, and the IBM-IIlinois Discovery Accelerator Institute. We thank NVIDIA for a GPU.
10

---

References
[1] A. Agarap. Deep learning using rectified linear units (ReLU). arXiv/803.08337, 2018.
ArXiv preprint:
[] S. P. Arunachalam, S. Silwal, B. Evans, and I.. Pinto. Dexterous imitation made easy: A learning- based framework for efficient dexterous manipulation. ArXiv preprint: arXiv2208.13251, 2022.
[3] O. Biza, D. Wang, R. Platt, J.-W. van de Meent, and L. L. Wong. Action priors for large action spaces in robotics. AAMAS, 2021.
[4] Brys, A. Harutyunyan, H. B. Suay, S. Chernova, M. E. Taylor, and Nowé. Reinforcement learning from demonstration through shaping. In IJCAI, 2015.
[5] W.-D. Chang, J. C. G. Higuera, S. Fujimoto, D. Meger, and G. Dudek. 11-flow: Imitation learning from observation using normalizing flows. NeurIPS 2021 Workshop on Robot Learning, 2021.
[6] X. Chen, S. Li, H. Li, S. Jiang. Y. Qi. and L. Song. Generative adversarial user model for reinforcement learning based recommendation system. ICML, 2019.
[7] S. Ciobanu. Mixtures normalizing flows. In CAINE, 2021.
[8] L. Dinh, J. Sohl-Dickstein, and Bengio. Density estimation using real NVP. In ICLR, 2017.
[9] N. Dvornik, C. Schmid, and J. Mairal. Selecting relevant features from a multi-domain representation for few-shot classification. In ECCV, 2020.
[10] S. Flennerhag, A. A. Rusu, R. Pascanu, H. Yin, and R. Hadsell. Meta-learning with warped gradient descent. In ICLR, 2020.
[11] J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4RL: Datasets for deep data-driven reinforcement learning. ArXiv preprint: 2.K22000.07211922022
[12] S. Ghasemipour, R. Zemel, and S. Gu. A divergence minimization perspective on imitation learning methods. In CoRL, 2019.
[13] A. Giusti, J. Guzzi, D. C. Ciresan, F.-L. He, J. P. Rodríguez, F. Fontana, M. Faessler, C. Forster, Schmidhuber, G. Caro, D. Scaramuzza, and M. Gambardella. machine learning approach to visual perception of forest trails for mobile robots. IEEE Robotics and Automation Letters, 2016.
[14] A. Gupta, Kumar, C. Lynch, Levine, and K. Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. In CoRL, 2019.
[15] T. Haarmoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critie: Off Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In ICML, 2018
[16] K. Hakhamaneshi, R. Zhao, A. Zhan, P. Abbeel, and M. Laskin. Hierarchical few-shot imitation with skill transition models. In ICLR, 2022.
[17] T. Hester, M. Vecerik, O. Pietquin, M. Lanctot, T. Schaul, B. Piot, D. Horgan, Quan, A. Sendonaris, I. Osband, G. Dulac-Amold, J. Agapiou, J. Z. Leibo, and A. Gruslys. Deep q-learnin from demonstrations. AAAI, 2018.
[18] C.-W. Huang, D. Krueger, A. Lacoste, and Courville. Neural autoregressive flows. In ICML, 2018.
[19] A. Javaloy and I. Valera. Rotograd: Gradient homogenization in multitask learning. In ICLR, 2021.
[20] S. Khader, H. Yin, P. Falco, and D. Kragic. Learning stable normalizing-foot flow control for robotic manipulation. In ICRA, 2021.
[21] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR, 2015.

---

[22] D. P. Kingma and P. Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. In NeurIPS, 2018.
[23] D. Kingma, T. Salimans, R. Jozefowicz, X. Chen, I. Sutskever, and M. Welling, Improved variational inference with inverse autoregressive flow. In NIPS, 2016.
[24] I. Kobyzev, S. Prince, and M. A. Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Transactions Puttern Analysis and Machine Intelligence, 2021.
[25] L. Liu, W. Hamilton, G. Long, J. Jiang, and H. Larochelle. A universal representation trans- former layer for few-shot image classification. In ICLR, 2021.
[26] C. Lynch, M. Khansari, T. Xiao, V. Kumar, J. Tompson, S. Levine, and P. Sermanet. Learning latent plans from play. In CoRL, 2019.
[27] X. Ma. J. K. Gupta, and M. J. Kochenderfer. Normalizing flow policies for multi-agent systems. In GameSec, 2020.
[28] B. Mazoure, Doan, A. Durand, R. D. Hjelm, and J. Pineau. Leveraging exploration in off off-policy algorithms via normalizing flows. In CoRL, 2019.
[29] A. Nair, A. Gupta, M. Dalal, and S. Levine. AWAC: Accelerating online reinforcement learning with offline datasets. ArXiv preprint: arXiv2006.09333 2020.
[30] T. Nam, S.-H. Sun, K. Pertsch, S. J. Hwang. and J. J. Lim. Skill-based meta-reinnfrrcement learning. In ICLR, 2022.
[31] G. Papamakarios, T. Pavlakou, and I. Murray. Masked autoregressive flow for density estimation. In NIPS, 2017.
[32] J. Pari, N. M. Shafiullah, S. P. Arunachalam, and L. Pinto. The surprising effectiveness of representation learning for visual imitation. In RSS, 2022.
[33] K. Pertsch, Y. Lee, and Lim, Accelerating reinforcement learning with learned skill priors. In CoRL., 2020.
[34] K. Pertsch, Y. Lee, Y. Wu, and J. J. Lim. Demonstration-guidde reinforcement learning with learned skills. In CoRI, 2021.
[35] G. G. P. F. Pires and M. A. T. Figueiredo. Variational mixture of normalizing flows In ESANN, 2020.
[36] M. Plappert, M. Andrychowicz, A. Ray, B. McGrew, B. Baker, G. Powell, J. Schneider, J. Tobin, M. Chociej, P. Welinder, V. Kumar, and W. Zaremba. Multi-goal reinforcement learning: Challenging robotics environments and request for research. ArXiv preprint: arXiv1802. 09464, 2018.
[37] J. Puigcerver, C. Riquelme, B. Mustafa, C. Renggli, A. S. Pinto, S. Gelly, D. Keysers, and N. Houlsby. Scalable transfer learning with expert models. In ICLR, 2021.
[38] A. Raffin, A. Hill, A. Gleave, A. Kanervisto, M. Ernestus, and N. Dormann. Stable-baselines:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 2021.
[39] Z. Ren, R. Yeh, and A. G. Schwing. Not all unlabeled data are equal: Learning to weight data in semi-supervised learning. In NeurIPS, 2020.
[40] D. Rengarajan, G. Vaidya, Sarvesh, D. Kalathil, and S. Shakkottai. Reinforcement learning with sparse rewards using guidance from offline demonstration. In ICLR, 2022.
[41] S. Schaal. Learning from demonstration. In NIPS, 1996.
[42] Schaal and C. Atkeson. Robot juggling: implementation of memory-based learning. IFEE Control Systems Magazine, 1994.
12

---

[43] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. ArXiv preprint crXiv:1777.06347 2017.
[44] D. Silver, T. Hubert, J. Schrittwieser, I. Antonoglou, M. Lai, A. Guez, M. Lanctot, L. Sifre, D. Kumaran, T. Graepel, T. P. Lillicrap, K. Simonyan, and D. Hassabis. Mastering chess and shogi by play with general reinforcement learning algorithm. Science, 2018.
[45] A. Singh, A. Yu, J. Yang, J. Zhang, A. Kumar, and S. Levine. Cog: Connecting new skills to past experience with offline reinforcement learning. In CoRL, 2020.
[46] A. Singh, H. Liu, G. Zhou, A. Yu, N. Rhinehart, and Levine. Parrot: Data-driven behavioral priors for reinforcement learning. In ICLR, 2021.
[47] D. Singh Chaplot, R. Salakhutdinov, A. Gupta, and S. Gupta. Neural topological slam for visual navigation. In CVPR, 2020.
[48] Y. Tang and S. Agrawal. Boosting trust region policy optimization by normalizing flows policy, ArXiv preprint: arXiv/809.10326 2018.
[49] A. Touati, H. Satija, J. Romoff, J. Pineau, and P. Vincent. Randomized value functions via multiplicative normalizing flows. In UAI, 2019.
[50] E. Triantatillou, T. Zhu, V. Dumoulin, P. Lamblin, U. Evci, K. Xu, R. Goroshin, C. Gelada, K. Swersky, Manzagol, and H. Larochelle. Meta-dataset: A dataset of datasets for learning to learn from few examples. In ICLR, 2020.
[51] E. Triantafillou, H. Larochelle, R. Zemel, and V. Dumoulin. Learning a universal template for few-shot dataset generalization. In ICMI., 2021.
[52] R. van den Berg, L. Hasenclever, J. Tomezak, and M. Welling. Sylvester normalizing flows for variational inference. In UAI, 2018.
[53] A. van den Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. ArXiv preprint arXiv: 2018.
[54] M. Vecerík, Т. Hester, J. Scholz, F. Wang. O. Pietquin, B. Piot, N. M. O. Heess, T. Rothörl, T. Lampe, and M. A. Riedmiller. Leveraging demonstrations for deep reinforcement learning on robotics problems with sparse rewards. ArXiv preprint: arXiv/ 2017
[55] Y. Wang, C. Xu, B. Du, and H. Lee. Learning to weight imperfect demonstrations. In ICML, 2021.
[56] N. Ward, A. Smofsky, and A. Bose. Improving exploration in soft-actor-crin with normalizing flows policies. In Workshop Invertible Neural Networks and Normalizing Flows in ICML, 2019.
[57] F. Xia, C. Li, R. Martin-Martin, O. Litany, A. Toshev, and S. Savarese. ReLMoGen: Lever- aging motion generation in reinforcement learning for mobile manipulation. ArXiv preprint vrXiv:2088.07792 2020.
[58] M. Yang, S. Levine, and O. Nachum. TRAIL: Nearaoptimal imitation learning with suboptimal data. In NeurlPS, 2021.
[59] L. M. Zintgraf, K. Shiarlis, V. Kurin, K. Hofmann, and Whiteson. Fast context adaptation via meta-learning. In ICML, 2019.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Sec. 6.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See Sec.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes] ensure my paper conforms to the ethics review.
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A] We have no theortical results.
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main exper-
imental results (cither in the supplemental material or as a URL)? [Yes] We include /ttp///9937119////0
them at https: htpps///gthhll
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See our experiment section and Appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes] See Sec. E in the Appendix.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cite the creators of all baselines and testbeds.
(b) Did you mention the license of the assets? [Yes] See the Appendix.
(c) Did you include any new assets either the supplemental material or as URL? [Yes] We include modified version of packages like gym for reproduction of our methods.
(d) Did you discuss whether and how consent obtained from people whose data you's 're using/curating? [Yes] See the Appendix.
(e) Did you discuss whether the data you are using/curating contains personally identifi- able information or offensive content? [Yes] Our data are all simulated and have no personally identifiable information offensive content.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A] We do not crowdsourcing conduct research with human subjects.
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "VJ-TiRi8O6Y", "reviewer": "Reviewer_P2zU", "summary": "The paper proposes a method CEIP that leverages explicit and implicit priors from the demonstration for reinforcement learning. On top of PARROT [45] which uses normalizing flow prior to task agnostic dataset to generate action, CEIP learns coefficients to combine flows to generate one task-specific flow parallel and use (state, predicted next state) pairs as the input of the flow instead of the state. As a result, CEIP achieves better performance than its baselines, PARROT, FIST, and SKiLD. ", "questions": "1. In Line 202, us -> use\n\n2. In Figure 5 (b) and Figure 9, it is hard to differentiate between solid and dotted lines. It is better to use different markers for the same category of EX instead of two lines.\n\n3. In Lines 164-167, “most useful” is a vague term, it would be better to use other expressions such as making the best performance, etc.\n", "limitations": "The authors addressed the limitations and potential negative societal impact of their work. As mentioned in the paper, I am concerned about the training time. I recommend parallelizing the training procedure.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**[Strength]**\n\n1. The paper is well-written and easy to understand. \n\n2. Especially, the detailed description of baselines, environments and other experimental detail make the reader have a better understanding of with the attached codebase.\n\n3. The proposed method clearly outperforms other baselines, PARROT, SKiLD, and FIST in all experiments and the authors also present the contribution of explicit prior by presenting Figure 5(c).\n\n\n**[Weakness]**\n\n1. It seems that there is no difference between task-agnostic flow (f_1 to f_n) and task-specific flow f_{i+1}. The ratio and size of the task-agnostic dataset and task-specific dataset are needed. Moreover, I recommend presenting the coefficient of f_{i+1} in those cases.\n\n2. Another concern is that a detailed description difference between CEIP and PARROT-TS-EX-forward is missing. According to Line 220, PARROT-TS only uses task-specific data while CEIP-TS uses both task-specific and task-agnostic data. Is the dataset size a key factor of the difference or the Few-Shot Adaptation?\n\n\n3. To present the effectiveness of the explicit prior, showing the results of TA-EX, TA-EA-forward in PARROT are better than TS-EX and TS-EX-forward (Figure 5(c)).\n\n\n4. It is good to compare using CEIP with the ground-truth task label (oracle) for not using k-means clustering.\n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "B0L9oRjZ_FG", "reviewer": "Reviewer_N6m9", "summary": "The paper presents a technique for combining \"explicit\" and \"implicit\" priors for reinforcement learning with demonstrations. Implicit priors refer to priors that store the knowledge from demonstration in a neural network of some kind (such as a flow-based generative model), while explicit priors store demonstrations as a database that can be queried when learning a new task. The overall method has several steps: \n\n1. Given a dataset of task-agnostic (TA) demonstrations, cluster them using k-means (where the last state of each trajectory is used as the clustering feature). If the given datasets are already divided into separate tasks, this step can be skipped. \n2. Learn a flow-based prior for each of the cluster from the previous step, and learn a flow-based prior for task-specific trajectories. The number of task-specific trajectories is much less than the total number of trajectories. Often, there's only one task-specific trajectory (one-shot learning). \n3. Learn weights to combine the numerous flow-based models into a single model. These weights are learned by minimizing a loss on  the task-specific dataset, and I have some questions about this step below. \n4. Given a new task, train a policy \\pi(z|s) that controls the prior model, which in turn controls the environment. This is similar to prior work (PARROT). \n5. You are also allowed to access a dataset of expert trajectories on the current task. This dataset is accessed via a retrieval procedure wherein you search for a state that is most similar to the current state, and retrieve the next state. The next state is then passed as an additional input to the flow-based models, allowing the flow-based model to act as a sort of inverse dynamics model -- given state and next state, the flow based model predicts an action. This step is crucial for good performance from what I can tell. \n\nThe authors evaluate their method on simulated robot manipulation tasks. \n\nCurrently, my rating for the paper is somewhat low as I am not able to make sense of the results, and I am hoping the authors are able to answer the questions I asked here so that I can revise my rating accordingly. \n\nPost-rebuttal: I have updated my rating to a borderline accept following the author response, as the authors have addressed a most of my concerns. I would encourage the authors to release their code if they can, as it would make the paper reproducible, and would spur future research in this area. One of my concern was that the results seemed almost too good to be true, and this concern can also be put to rest with a code release. ", "questions": "- Authors use the terms “explicit priors” and “implicit priors” frequently, but these terms are never clearly defined (and they are not standard terms, as far as I know). From my reading, it seems like implicit priors refers to a method that trains a neural network using provided demonstrations, while explicit priors store all demonstrations in memory and can be later accessed using a nearest neighbor style technique. Can you clarify what exactly these two terms mean, and add these definitions towards the beginning of the paper? \n- The authors train several different flows on the task-agnostic and task-specific datasets, and then combine these flows using some coefficients, and these coefficients are learned by minimizing the log-likelihood of the “combined flow” on task-specific trajectories. However, given that one of the flow models was trained on task-specific trajectories alone, this subsequent optimization might result in the coefficient for the task-specific flow model to receive most of the weight. Am I missing something here? Can authors share what coefficients were assigned to the task-specific flow model in their experiments? \n- I understand the motivation behind the push-forward term in the retrieval objective function (Equation 6), but I did not completely follow the definition of the indicator function \\delta. Can the authors add another equation that precisely defines when \\delta = 1? Is it when s_next = s? \n- In Figure 3, why does CEIP have a much higher reward right at the start of training (when steps = 0)? Why do other flow-based methods (such as PARROT) start with a much lower return? In particular, I would expect PARROT-TS to have a good initial performance, given that it’s trained on data for the task at hand (and task is relatively simple). Also, since CEIP trains on both TA and TS data, it would make sense to compare against a version of PARROT that uses both TA and TS. \n- In Figure 4, the learning curve for CEIP for about half the tasks are essentially flat, with maximum performance achieved without any finetuning. Can the authors explain what is happening here? Is imitation learning alone sufficient for solving this task? The results seem very puzzling to me.\n", "limitations": "There is almost no discussion on limitations in the paper (the authors only briefly mention the computational cost of training many different flow models), which is disappointing. The lack of discussion around limitations, combined with results that seem almost too good to be true is a bit concerning. ", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The paper tackles an important problem: how to best make use of task-agnostic and task-specific demonstrations to speed up reinforcement learning of new tasks. \n- The paper obtains some impressive results on an interesting suite of tasks. \n- The paper compares against a number of baselines, competing methods, and ablations. \n\nWeaknesses\n- The overall method is somewhat complex, with a number of moving pieces, and some of the design choices (such as the retrieval system) seem a bit ad-hoc. \n- The experimental results are a bit puzzling, and I will add more details on this in the Questions section below. \n- Most of the analysis / ablations are actually in the Appendix. Overall, the authors spend too much space explaining the method (which could have been shortened, since we don't need all the details on how flow-based models are trained, we can refer to prior work for that), and too little space analyzing the experimental results. \n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "0zlYybwuPBm", "reviewer": "Reviewer_sUgF", "summary": "This submission proposes a skill learning algorithm for offline RL that can leverage task-specific demonstrations as implicit and explicit priors. The skill policy consists of a mixture of normalizing flows (trained on task-agnostic demonstrations) where the mixture weights are fine-tuned on task-specific demonstrations. The flows resemble an inverse model, in that they are conditioned on both the current and next state. The next state forms the explicit prior, as it is queried from the task-specific dataset via nearest-neighbor lookup for each state encountered during training. The explicit prior enables very effective few-shot generalization (i.e., utilizing the task-specific demonstrations without with random latent vectors, as this is what I suppose corresponds to the first data point in the learning curves), and learning a high-level policy to predict latent state flow inputs can further improve performance.", "questions": "- How would simply replaying a task-specific demonstration perform on the benchmarks? Does it already provide optimal reward and trajectories? I assume that fine-tuning the normalizing flow mixture requires ground truth actions of the demonstrations, so the return of the demonstrations themselves could be added as a baseline/topline.\n- How much leverage does the high-level policy actually have over the actions? The flow policy is invertible, but conditioning on current and next state might, depending on the environment, provide all required information to compute an action already. Some analysis on this aspect would be appreciated.", "limitations": "The paper is not clear on the quality of demonstrations that are expected for the method to work effectively.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is very well-written and easy to follow. Clarity in the experimental section could be improved by stating what the abbreviations in ablation studies mean, and the ablations in the appendix are a bit hard to follow. Overall, the idea presented is well-motivated, novel and concerns a topic of high interest. The technique mainly builds upon PARROT and includes tricks for fine-tuning on task-specific demonstrations.\n\nThe experiments clearly show that the proposed method is effective in the settings considered; in fact, there is no setting in which few-shot generalization does not already work much better than all baselines considered. Overall, this suggests a potential limitation that doesn't seem to be acknowledged explicitly in the paper: how much does the method rely on precise task-specific demonstrations being available? What would happen if the environment changes slightly, e.g., the microwave in the kitchen environment would be moved 10cm to the left? I'd be happy to raise my score if this limitation is properly addressed, as I think it would help greatly in judging the method's applicability and impact.", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept