Memory-Consistent Neural Networks for Imitation Learning
Authors: Kaustubh Sridhar
Souradeep Dutta
Dinesh Jayaraman
James Weimer
Insup Lee
Venue: ICLR 2024 poster (2024)
Keywords: Imitation Learning
Behavior Cloning
Deep Learning
Primary area: reinforcement learning
TL;DR: We develop a method to interpolate between nearest neighbours and neural networks for controlling the sub-optimality gap and improving performance in imitation learning.

Abstract: Imitation learning considerably simplifies policy synthesis compared to alternative approaches by exploiting access to expert demonstrations. For such imitation policies, errors away from the training samples are particularly critical. Even rare slip-ups in the policy action outputs can compound quickly over time, since they lead to unfamiliar future states where the policy is still more likely to err, eventually causing task failures. We revisit simple supervised "behavior cloning" for conveniently training the policy from nothing more than pre-recorded demonstrations, but carefully design the model class to counter the compounding error phenomenon. Our "memory-consistent neural network" (MCNN) outputs are hard-constrained to stay within clearly specified permissible regions anchored to prototypical "memory" training samples. We provide a guaranteed upper bound for the sub-optimality gap induced by MCNN policies. Using MCNNs on 10 imitation learning tasks, with MLP, Transformer, and Diffusion backbones, spanning dexterous robotic manipulation and driving, proprioceptive inputs and visual inputs, and varying sizes and types of demonstration data, we find large and consistent gains in performance, validating that MCNNs are better-suited than vanilla deep neural networks for imitation learning applications. Website: https://sites.google.com/view/mcnn-imitation

Published as a conference paper at ICLR 2024
MEMORY-CONSISTEN NEURAL NETWORKS FOR IMI- TATION LEARNING
Kaustubh Sridhar', Souradeep Dutta', Dinesh Jayaraman's James Weimer Insup Lee University of Pennsylvania, Vanderbilt University (ksr (ksrichar, duttaso, dineshi, weimeri, 1ee) @seas upenn edu
ABSTRACT
Imitation learning considerably simplifies policy synthesis compared alternative approaches by exploiting access to expert demonstrations. For such imitation policies, errors away from the training samples particularly critical. Even rare slip-ups in the policy action outputs can compound quickly over time, since they lead to unfamiliar future states where the policy still more likely to err, eventually causing task failures. We revisit simple supervised "behavior cloning" for conve- niently training the policy from nothing more than pre- pre-recorded demonstrations, but carefully design the model class to counter the compounding error phenomenon. neural network" (MCNN) outputs are hard-constrained stay within clearly specified permissible regions anchored to prototypical "memory" training samples. We provide a guaranteed upper bound for the sub-optimality gap induced by MCNN policies. Using MCNNs on 10 imitation learning tasks, with MLP, Transformer, and Diffusion backbones, spanning dexterous robotic manipu- lation and driving, proprioceptive inputs and visual inputs, and varying sizes and types of demonstration data, we find large and consistent gains in performance, validating that MCNNs better-suited than vanilla deep neural networks for imitation learning applications.
Our "memory-consisere
INTRODUCTION
For sequential decision making problems such as robotic control, imitation learning is an attractive and scalable option for learning decision making policies when expert demonstrations are available as a task specification. Such demonstrations are typically easier to provide than the typical task specification requirements for reinforcement learning and model-based control, namely, dense rewards and good models of the environment. Furthermore, imitation learning is also typically less experience- intensive than reinforcement learning and less expertise-intenat than model-based control.
We consider the simplest and perhaps most widely used imitation learning algorithm, behavior cloning (BC) [29], which reduces policy synthesis to supervised learning over the expert demonstration data. For example, neural network policy for an autonomous car could be trained to mimic human driving actions [2]. While the policy is synthesized with supervised learning, the evaluation setup is very different: rather than merely achieving low average error states from the training data, as common in supervised learning, the trained policy must, when rolled out in the world, successfully accomplish the demonstrated task.
This sequential deployment makes the behavior imitation policy functions away from their training data particularly critical. To see this, observe that during rollout, the policy's own output actions determine its future input states. Task performance is most closely tied to the policy's 's behavior on this self-induced set of states, which particular, a minor error in the policy's 's action output at any time may induce future input state that subtly different from expert states. If the policy behaves erratically under such small deviations, as often does in practice, the situation quickly snowballs into a vicious cycle of compounding errors leading to task failure.
deviate from the training dataset of expert demonstrations, In
Our website: https: //t////ttttt tee. googl e. com/vi ww/p/ennmnneeeeeeeeeeeeeeeeemmmm imi tan

---

Published as a conference paper at ICLR 2024
Past solutions this compounding error problem have focused modifying the be- havior cloning setup, such as by permitting online experience [13, 34]. reward labels [28), queryable experts [37], or modifying the demonstration data collection procedure [19]. Instead, we retain the conveniences of the plain BC setup and focus on design- ing a model class that encourages better behavior beyond the training data, which in turn could boost task performance by mitigating the compounding error phenomenon discussed above. We provide a simple plug-in approach to improve BC with any deep neural network. It is well known that vanilla deep neural networks, only by themselves, can generate large errors when evaluated away from the training points, and even rare errors could derail an entire task rollout. These large errors are particularly evident when the expert demonstrations are few in num-
300%
Median trend
with MCNN over D4RL BC
250%
% Increase in Retur
200%
ber such as in robotics where human demonstrations are essential for imitation learning. To tame these errors, we propose
150%
ememory-consister neural net- works" (MCNN). MCNNs first subsample the dataset into representative prototype "memories" to form the scaffold for the eventual function. They then fit parametric func- tion to the rest of the training data that is hard constraind by the very formulation of the model class, to exactly fit the training data at all the memories, and further, to stay within louble-cone-like zones of controllable shapes and sizes centered at each memory. As a result, an MCNN behaves mostly like a nearest-neighbor function close to memories, and mostly like deep neural network (subject to the double-cone constraints) far from them. All func- tions in this MCNN model class lie within "permissible regions" centered on each memory, meaning that function values away from the training points bounded. Un- der mild assumptions on the expert policy, we show that this property of MCNNs induces an upper bound on the suboptimality of the learned BC policy. Visualizations of MCNNs can be found in Figure 2.
100%
50%
0%
DDRLL-B
10a
10a
10a
Number of Demonstration Trajectories
Figure I: MCNN significantly improves performance on realistic demonstration datasets. We plot the percentage increase return with MCNN over D4RL. BC [9] for various number of demonstrations across many tasks. In this plot, each point is sep- arate MCNN policy. We see significant im- provements in few demonstrations regime where most realistic imitation learning tasks can be found. The choice of model class is crucial in such regimes and MCNN shines. Additional details are in Appendix F.
Using MCNNs on 10 imitation learning tasks, with MLP, Transformer and Diffusion backbones, spanning dexterous robotic manipulation and driving, proprioceptive inputs and visual inputs, and varying sizes and types of demonstration data, we find large and consistent gains in performance, validating that MCNNs are better-suited than vanilla deep neural networks for imitation learning applications. Figure visualizes the percentage increase in return with MCNN policies compared to the vanilla BC results reported in D4RL [9] for various quantities of training demonstrations across tasks. The trend of the median demonstrates that MCNNs are highly effective in the low data regime where generalization to test trajectories is stressed.
2 RELATED WORK
We present a detailed related work discussion in Appendix & summarize closely related work here.
Compounding errors in imitation learning have previously been tackled by permitting online experience [13, 34], reward labels [28], queryable experts [37], modifying the demonstration data collection procedure [19]. Our work is orthogonal to these methods and creates model class that avoids compounding errors by construction. Other works that propose new models for IL such as Implicit BC (IBC) [7]. Behavior Transformer (Be'T) [38], Action Chunking Transformer [48], and Diffusion Policies [44, 1] are orthogonal to our approach. MCNN can be used as a plug-in approach to improve any of these methods. In fact, we show that MCNN with BeT backbone outperforms vanilla BeT and MCNN with a diffusion model outperforms diffusion BC on all tasks in our experiments in Section 5. We also show that MCNN outperforms IBC in Section 5.
Non-parametric and semi-parametrie methods in imitation learning such as nearest neighbors [39], RBFs [35], and SVMs [20] have historically shown competitive performance on various robotic control benchmarks. But, only recently, a semi-parametrr approach consisting of neural networks for representation learning and k-nearest neighbors for control was proposed in Visual Imitation through
2

---

Published as a conference paper at ICLR 2024
Nearest Neighbors (VINN) [26]. This is the closest paper to work and in Section 5, compare with VINN and demonstrate that we outperform their method comprehensively.
Theoretical guarantees the sub-optimality gap in imitation learning with MCNN provided in this paper. Such guarantees are not available with vanilla neural networks. Our theorem builds on earlier work reductions for imitation learning in [36, 31, 2, 32] and leverages intuitions from [25] on bounding the width of the model class.
PROBLEM FORMULATION
Markov Decision Process (MDP) is a tuple = (S, where S C RD the set of states, A is the of actions, P(s'|s,, a) is the probability of transitioning from state when taking action a, R(s, a) is the reward accrued in state upon taking action a, 7 E [0, 1) is the discount factor, and I is the initial state distribution. We assume that the MDP operates over trajectories with finite length H, in an episodic fashion. Additionally, we assume the set of states S be closed and compact. Given policy т S A, the expected cumulative reward accrued over the duration of an episode is given by the following,
8'
J(T) Ex | R(s at)].
(I)
In imitation learning, we assume that there exists expert policy n" unknown to the learner: This policy induces a distribution dz- on the state-action space The learner agent has access to an expert trajectory dataset D = {(s0, ag).
obtained by rollouts on the MDP.
(SN, ax)
drawn from distribution dg., The goal of imitation learning is to estimate policy म. which mimics the expert's policy and reduces the sub-optimality gap: ((x)) )- -(t).
4 APPROACH
Our approach involves developing new model class, memory consistent neural networks (MCNN), and training with supervised learning to clone the expert from the demonstration data. We start by setting up the MCNN model class in Sec 4.1, analyze its theoretical properties for imitation learning in Sec 4.2, and finally describe our behavior cloning algorithm that uses MCNNs in Sec 4.3.
4.1 THE MODEL CLASS: MEMORY-CONSISTENT NEURAL NETWORKS
First, we develop the semi-parametric MCNN model class for imitation learning. MCNNs rely on code-book set of "memories" B = {(s, k which are subsampled from the expert training dataset and summarize it. In practice, such memory code-book can be created using one of various off-the-shelf approaches. We describe our algorithmic choices later in See 4.3. For notational convenience, we describe the approach for a scalar action space, but it is easily generalizable to the vector action spaces we evaluate in our experiments.
Given this memory code-book B, we now define "nearest memory neighbor policy". For a finite set S C S, and an input æ S, we first define its closest element S as, Cs(x) = arg min d(s, x). .
where d some distance metric defined on the space S. We denote by Bls, and BI as the set of all states and actions captured by the memory code-book B. With slight abuse of notation, we denote (ss) as the action assigned by the codebook for a state input s. Using the above, we now define nearest neighbor regression function fNN as the following.
Definition 4.1 (Nearest Memory Neighbor Function). For an input & E S, assume that = CB (x). then N. (x) == B(s').
In other words, the nearest memory neighbor function assigns actions according nearest neighbor look-up in the memory code-book B. We are now ready to define memory-com coonssscon neural networks (MCNN). which permit interpolating between nearest neighbor functions and parametric deep neural network (DNN)-based functions. Let f denote a DNN function parameterized by which maps from MDP states actions.
3

---

Published as a conference paper at ICLR 2024
Pramspe
Nearest welghhor Fancter
Neural Netewoee Function class y otf(
Hemorry-oomneer Neural Net Class
y [*)/==1000
y- 52g0e
0.0
2.5
5.0
7.5
10.0
2.5
5.0
7.5
10.0
0.0
2.5
5,0
7.5
10.0
x
More Examples of MCNN Model Class
Increasing
increasing
The elements of the MCNN model class. In the top row, the left panel shows the nearest memory neighbour component with memories subsampled from the training dataset shown in red circles. The middle panel depicts the constrained neural network function class, where the blue shaded regions represent permissible regions; by design, the function cannot take values outside these shaded regions. Finally, the right panel shows the combined MCNN model class. The size of the permissible regions can be modulated by increasing X (bottom left) by decreasing the number of memories (bottom right). The second shows many such MCNN model families with increasing capacity. For additional plots, Appendix
Figure 2:
Definition 4.2 (Memory-Conenl Consistent Neural Network). A
neural network f fMC is defined using the codebook and DNN function pair (B, )"), and hyperparameters X R+, L R as
MC (I))
J (1r) E (e-Ad(uret) d((.,o))
+ L. =1- o (x)
Memory Neighbour
Constrained Network
where, s' = CBis (x) is the nearest memory to and a R + [-1, that imposes hard limits the outputs of fo, In practice, we use tanh similar functions.
1] is compressive non- non- lineaiiy
We refer the reader Figure 2 to drive the intuition. For inputs that are close to the points in the code- book B, the function predicts values that are similar to the one observed in the training dataset, More concretely, the value predicted by the function is simple mixture: a NN (x) + - a) L a of where the mixing factor E [0, 1] changes in proportion to the distance to the nearest memory. Thus, for points further away more weight is placed on the neural network and the memories have little influence. The degree of permissible deviation from nearest neighbor prediction NNN is controlled by the parameter A. Thus, we obtain a purely nearest neighbor function for = 0, and a vanilla deep neural network function for À cc. Note that the MCNN function values in regions far away from memories are in the set [- L, LJ. For this reason, we normalize output actions to [ L, L] before training MCNN function.
4.2 THEORETICAL ANALYSIS OF MCNNS FOR IMITATION LEARNING
For fixed hyperparameters L, À and memory codebook B, denote by 3, the class of memory- consistent functions outlined in Equation 2. Note that choice of the DÑN function parameters fixes specific function in this class as well.
Assumption 4.3 (Realizability). We assume that the expert policy belongs to the function class 3.
This assumption trivially holds at the memories, where the MCNN exactly reproduces expert actions. For all other points, assume that there exist some parameters 0, which can capture the policy "² with sufficient accuracy, for a choice L and X. For a point x at a distance of d(s', x), the vanilla
e> -(z."))) Without this restriction, might have been able to capture behaviors that went well beyond these ranges. This is reasonable since, expert policies do not make sudden unbounded jumps in their actions. What we propose here is way to enforce this bound using zeroth order nearest neighbor estimate.
DNN can affect the predictions only by an amount of + T.
4

---

Published as a conference paper at ICLR 2024
We analyze the behavior of this function class, and present useful lemmas along the way. For a set of memories present in Bls, we wish to capture the maximum value that the distance term: d(2,, s') can take in Equation 2. To that end, we define the most isolated state the following: Definition 4.4 (Most Isolated State), For a given set of memory points Bis, we define the most isolated state of SB|s := arg max min d(s, d)), and consequently the distance of the most isolated meBls
#ES d(s
min point as B)s mEBle
The distance of the most isolated state captures the degree of emptiness that persists with the current knowledge of the state space due to memory code-book B.
Lemma 4.5. Assume two sets of memory code-books B,, By, such that B, S By, then d/ 2 B,ls
Proof: The proof of the above lemma is straightforward, since the infimum of subset (B,) is larger than the infimum of the original set (B,).
This observation is useful when study the effects of increasing the size of the code-book B. Note, when learning a memory-consisted neural network f0.8 we deploy the standard SGD based training to adjust the parameters 0. The choice of the number of memories in B is kept as hyperparameter. This allows us bound the maximum width of the function class, first described in [25]. We analyze this for single output functions next.
Lemma 4.6. (Width of Function Class) The width of the function class 3. V 01, O2 E θ, and Vs E E, (JM SMC. B) (s)l is upper bounded by: 2L (1 > d
defined as mar 6,88
Proof: Please see Appendix B.
Theorem 4.7. The sub-optimaliiyy gap J(x²) J(A) min min HA -bd1ss)
Proof: Please see Appendix C.
Corollary 4.8. Using Lemma 4.5 we know that B; C Bj. then (1 e dBb 8) 2 (1 A This result in lower performance gap according to Theorem 4.7, when H 2 H&|ALT(1
x dbis). Hence, reflecting the utility of adding more memories in such cases.
Takeaways. We summarize the insights from the above theoretical analysis here. First, our MCNN class of functions is bounded in width (Lemma 4.6) even though it uses a high-capacity function approximator like DNNs. No such bound is available for vanilla neural networks. This translates to bounded sub-opting optimalttlltty gap (Theorem 4.7) also not available in vanilla neural networks. Finally, our Corollary states that can likely gain better imitation learning performance by simply adding more memories (up to a limit).
4.3 ALGORITHM: IMITATION LEARNING WITH MCNN POLICIES
We now describe our algorithm to use MCNNs for imitation learning. The first step in our method is to learn the memory code-book B from the expert trajectory dataset 1.. The goal of Algorithm is to build the nearest memory neighbor function FNN_ This is followed by details on the training aspects of the MCNN parameters from the imitation dataset D. in Algorithm 2.
For building the memory code-book, we leverage an off-the-shelf approach, Neural Gas [8], that selects prototype samples to summarize a dataset. For completeness, we summarize this approach briefly below.
Definition 4.9 (Neural Gas). A neural gas 9 = (N, is composed of the following components,
1. A set .et c S of the nodes of a network. Each node my E W is called a memory in this paper. 2. set & C [mm,,, m2) A22 + j} of edges among pairs of nodes, which encode the topological structure of the data. The edges are unweighted.
Neural gas. The neural gas algorithm [8. 30, 22] is primarily used for unsupervised learning tasks, particularly for data compression or vector quantization. The goal is to group similar data points together based their similarities. The algorithm works by creating set of prototype vectors, also known as codebook vectors or neurons. These vectors represent the clusters in the data space. The algorithm works by adaptively placing prototype vectors in the data space and distributing them like
5

---

Published as a conference paper at ICLR 2024
Algorithm Learning Memories
Input: Offline dataset D = [(*t, ] number of memories m
Output: nearest neighbor based function fNN
A
// Iearns the distribution induced by D 2: Nodes A", D(A") + For each node in N,, find the closest observation in D, and call this N".. Additionally, return the corresponding action taken by the expert in D, denoted by the map D(N")
1: Nodes N, edges & +- NeuralGaaSClusteringg
m)
3: G &- Define neural-gas with nodes N" and edges E.
4: Define memory code-book B using B|s and B = D(N"). Pairing nodes in the neural-gas to its corresponding actions.
Define nearest neighbor function fg along lines described in Definition 4.1 using B.
6: return f N
a gas in order to capture the density. For more details we refer the reader to [8]. We use this in Algorithm (Line 1) get the initial clustering. We now go ahead and outline how "memories" are picked in our case.
Learning memories. Algorithm first uses the neural-gas algorithm pick candidate points (nodes) N in the state space. However, these points could be potentially absent in the dataset D, making it hard to associate the correct action. To remedy this situation, we replace these points with the closest states from the training set as memories. Such memory states come with the corresponding actions taken by the expert. This is then used to define a nearest neighbor function by building the memory codebook B and defining function outlined in Definition 4.1.
Algorithm Behavior Cloning with Memory-Consisee Neural Networks: Training
Input: Dataset D [(ss, a:) N nearest neighbor function f neuural network function (.). batch size, total training steps parameters and
Output: Learned policy fa.8
1: for step= to do
2: Sample batch B from D.
3: Forward propagate (st, a.) ~ B. J
PCC (z) = 1 (e (e-Ad(tx))) d (1-e-Ad(aa)) (x)
where, s' the nearest neighbor in B s. (x) is a lanh-like activation function given by (x(x)) 2 [LeakyReLU, (+))
B) ReLU [[]] and B = max (0, 1 166 J)-
4: Update + V ( C(S (s:), a.) where C the negative log likelihood or mean squared
loss or other loss function.
5: end for
Training MCNNs. Finally, we train MCNN policies through gradient descent on the parameters of the neural network over the expert dataset D.. For the compressive non-linearity, we use the 0g function given in Algorithm 2 which is similar to tanh. We describe this in detail in Appendix F.
Figure 3: The environments here include: Adroit Pen, Hammer, Relocate, and Door [33], CARLA's Town03 and Town04 [3], and Franka Kitchen [11]. The four Adroit environments and Franka Kitchen have proprioceptive observations and CARLA environment has image observations.

---

Published as a conference paper at ICLR 2024
5 EXPERIMENTAL EVALUATION
We now perform a thorough experimental evaluation of MCNN-based behavior cloning in large variety of imitation learning settings.
We test our approach on 10 tasks, in 6 environments: 4 Adroit dextrous manipulation, 1 CARLA environment, and the Franka Kitchen environment as pictured in Figure 3. Demonstration datasets are drawn from D4RL [9] for Adroit and CARLA and from the multimodal relay policy learning dataset [11] for Franka Kitchen. For each Adroit task, we evaluate imitation learning from 2 different experts: (1) small realistic human demonstration datasets ("human') with 25 trajectories per task (5000 transitions), and (2) large demonstration datasets with 5000 trajectories (1 million transitions) from a well-trained RL policy ('expert'). In CARLA, we train on demonstrations from hand-coded expert. The Franka Kitchen demos were collected by humans wearing VR headsets. For observations, we use high- high-dimensional states in Adroit and Franka Kitchen, and 48 x 48 images in CARLA. Action spaces are 24-30 dimensional in the Adroit, 9-D in Franka Kitchen, and 2-D in CARLA. Further, in all Adroit environments, a goal is randomly chosen at reset and goal information is included in the observation vector (more in Appendix F).
Environments and Datasets:
Baselines: We run the following baselines for comparison. (1) Behavior Cloning: We obtain results with a vanilla MLP architecture. The details of the architecture can be found in Appendix F. We report results from our implementation of BC and also report results given in D4RL [9] under the names 'MLP-BC' and 'D4RL BC' respectively. Our BC implementation has only one difference from [9]'s implementation: we normalize the observations. Normalizing observations has been shown improve BC's performance [10]. (2) 1 Nearest Neighbours (1-NN): We set up a simple baseline where the action for any observation in the online evaluation is the action of the closest observation in the training data. In the expert and cloned datasets for each environment, this amounts to having to perform search amongst million datapoints online at every step (which is highly inefficient). (3) Visual Imitation with (k) Nearest Neighbours (VINN) [26]: VINN is a recent method that performs Euclidean kernel weighted average of some k nearest neighbors. In the Adroit case, we directly perform the k nearest neighbors on the raw observation vectors. In the CARLA case, we perform it in the same embedding space that we use create memories (we discuss this embedding space more below). (4) CQL-Sparse (CQL-S): We learn a policy using the CQL offline RL algorithm [18] and sparse reward given for task completion only. (5) Implicit BC (IBC) [7]: We report the results from [7] which performs BC with energy models on the human tasks. (6) Behavior Transformer (BeT-BC) [38]: We train and evaluate behavior transformer using the official implementation on all tasks. (7) Diffusion BC (Diff-BC) [44, 1]: We also train and evaluate a diffusion based BC policy using the implementations in [44, 1]. We provide additional details for all baselines and comprehensive hyperparameter sweeps in Appendix F.
Learning memories and MCNN: We learn memories using the incremental neural gas algorithm for 10 epochs starting from 2.5% of the total dataset to 10% of the total dataset for each task. We update all the transitions in each dataset by appending the closest memory observation and its corresponding target action Algorithm 1). We train the MCNN this dataset following Algorithm 2 for 1 million steps and evaluate on 20 trajectories after training and repeat each experiment for a minimum of seeds. We report results with an MLP, a behavior transformer (BeT), and diffusion policy the underlying neural network under the names 'MNN++MP'' MMNNN+BTT', and "MCNN respectively. We expand on the experimental setup and all hyperparameters in Appendix F.
Embedding CARLA images: In the CARLA tasks, we use an off-the-shelf ResNet34 encoder [2] that has been shown be robust background and environment changes in CARLA to convert the 48 x 48 images to embeddings of size 512. We use this embedding space as the observation space for learning memories and policies.
All our environments come with pre-specified dense task rewards which we use to define performance metrics. We report the cumulative rewards (return) for each task. For the aggregate plots on a set of four tasks, we compute the percentage increase in return of a method over D4RL. BC all four tasks and report the median.
Performance Metrics:
First, for the "human" tasks with the most realistic imitation learning setup, we plot aggregate and taskwise results in Figure 4. In aggregate- we see that MCNN+MLP with fixed hyperparameters performs the best followed by MCNN+Diff and IBC at second place. We also see that the MCNN variants of MLP, BeT and Diffusion consistently outperform the vanilla versions.
Results:
7

---

Published as a conference paper at ICLR 2024
*ggregate human
per humrnn v2
hammeer Purnan v1
relotace human
0oor human vi
150%
3500
403
+240
Modan of inevecsn
203
2900
203
103
183
-30
D4 BC
VIRN
IBC
Dift-Br
MCNN + BeT (Fieed] MCN [Fleed]
+ (Rxed)
7-NN
CXQU
BeTT-BC
MiP--C
Figure 4: Adroit human tasks (25 demos]: Comparison of returns (across 20 evaluation trajectories and 3 random seeds) between baselines and methods (MCNN+Bel MCNN+Diff, and MCNN+MLP). Our MCNN methods use the same fixed set of hyperparameters across tasks.
AQU squtteoopeert
pen-eccert-v11
docresspetsvt
relocate sttpee--pppttt
130:23
hammer-ssper-vt
Wedian of her In
20%6
1900
14039
7300
L2000
000%
58:30
2201
G4RL DC
VRNN COL-S
Dcl BC
MLP-BC
MCMN-DI Dff (Rwed)
MCNN-MLP (Fiec)
C.NN
CNTT
MCNNA BoT (Fowac)
Comparison of returns (across 20 evaluation trajectories and random seeds) between baselines and our methods (MCNN+BeT) MCNN+Diff, and MCNN+MLP), Our MCNN methods use the same fixed set of hyperparameters across all tasks.
Figure 5: Adroit expert tasks [5000 demos]:
carla-lane-o
In pen-human-vl MCNN+MLP outperforms the nearest baseline by 33%. It is also the only method to shoot past the expert ceiling of 100 (depicted by dashed red line). In hammer-human- v1, MMLN is the only method to obtain a positive return outpacing the nearest baseline by an order of magnitude (from -11 to 262). In relocate-human- 1,, MCNN+BeT is the only method to achieve positive return. We attribute, like previous work [24], the stronger performance of MCNN+BeT over other MCNN variants in the relocate task the 'memory' advantage available to transformers that is specifically suited for this task (where the historical states inform whether the ball has been grasped). Lastly, in door-human-v// MCNN+Diff outperforms all but the IBC baseline by an order of magnitude. In this task, where repeated attempts at grasping and opening the door handle are usually required for success, we see methods that enable such repetition (energy models in IBC and diffusion in MCNN+Diff) succeed.
DARL BC
VIN
HCWM--RTT (Roedi
HCWW-MEEP FFand)
Figure 8: CARLA [400 de- mos]: Comparison of return (across evaluation trajectorie and random sceds) between baselines and our methods (MCNN+BeT and MCNN+MLP) For the results shown here, our MCNN methods the same fixed set of hyperparameters across all tasks.
On the expert tasks in Figure 5, even with large amount of data, we see MCNN+MLP come in first outperforming the nearest baseline in the aggregate plot by over 100%. MCNN+Diff comes in second in aggregate with a 40% improvement over the nearest baseline. Here too, MCNN variants outperform the vanilla versions across all tasks. Also, MCNN+MLP and MCNN+Diff are the only methods to exceed the expert ceiling in all four tasks. Across expert tasks both MCNN+MLP and MCNN+Diff perform competitively and obtain up to a 25% improvement over the nearest baseline.
Table 1:
Franka Kitchen [566 demos]: Com- parison of probabilities/success rates of interacting with I to 5 objects (given by pl p5 respectively) in Franka Kitchen hetween haselines and one of our methods (MCNN+Diff).
Franka Kitchen
pl
p2
p3
p4
p5
LSTM GMM [21]
1.0
0.9
0.74
0.34
IRC
0.99
0.87
0.24
BeT-BC
0.99
0.93
0.71
0.44
Duff-C [1]
1.0
1.0
1.0
0.99
0.02
MCNN+DiI
1.0
1.0
1.0
1.0
0.08
8

---

Published as a conference paper at ICLR 2024
3549
54000
2400
1ea0d
2000
14000
1204
Mevraties
222
Figure 6: Number of Memories in MCNN+MLP:
0567
Comparison of returns (across 20 evaluation trajecto- ries) with our method (MCNN+MI. using between find a
Figure 7: Fixed vs Tuned Hyperparameters MCNN+MLP: Comparison of returns (across 20 eval trajectories and 3 seeds) between our method with fixed hyperparameters across all tasks and with hy- perparameters tuned online. Hence, it's possible improve performance with limited online interaction.
2.5% to 80% of the dataset as memories. "sweet spot" for number of memories at 10-20%. We also see the expected decrease to I-NN performance number of memories increases to 100%.
In the high dimensional CARLA lane task as well, we see in Figure 8 that MCNN+MLP outperforms the nearest baseline by 27%. MCNN+BeT comes in a close second.
In Table 1, we compare MCNN+Diff with various baselines in the multimodal Franka Kitchen task. Following previous work [1], we compare the probability of success in interacting with to 5 objects (pI-p5). MCNN+Diff performs similarly to Diff-BC [1] and obtains 100% success on pl to p4. On p5, where baselines have failed to show any success, MCNN+Diff obtains a 4x improvement over the nearest baseline. Additional figures and all means and standard deviations of our results can be found in Appendix G.
Discussion: We identify some high-level trends across results here. For every architecture MIP, BeTT, or Diffusion, plugging in MCNN significantly improves performance in every task. The best-performing method in nearly every task is MCNN-based method. MCNN can even improve the performance of simple MLP architectures to beyond that of more sophisticated recent architectures such as Diffusion models. For example, in pen-human-vl in Figure 4, diffusion outperforms MLP but MCNN+MLP significantly improves upon Diffusion. The above statements remain true across different types and sizes of expert data and across disparate tasks. We discuss some ablations next.
We plot return against number of memories in Figure 6 for MCNN+MLP. It demonstrates the existence of a "sweet spot" for the number of memories around 10-20% of the dataset. This allows for more efficient inference in MCNN than in baselines like VINN and 1-NN. It also shows the expected degradation to 1-NN performance as the number of memories increases towards 100%. Additional discussion on computation cost and improved efficiency of MCNN compared to VINN and I-NN can be found in Appendix F.
Ablations:
In Figure 7, we show that given limited online interaction (20 episodes), selecting the best-performing hyperparameters (A, L, and number memories) further improves MCNN+MLP performance. We show the significance of neural gas-based memories by comparing MCNN+MLP with a version that uses randomly chosen memories in Figure 13 in Appendix G. We observe significant reduction in performance with randomly chosen memories. Neural-gas uses competitive Hebbian learning algorithm which is better at capturing the distribution of training points (creating memories that are "spread out"). This reduces the distance the most isolated state, improving imitation performance. Ablation results on the values of X and other tasks is in Appendix G.
CONCLUSIONS AND LIMITATIONS.
Imitation learning, and in particular behavior cloning, one of the most promising approaches when it comes to transferring complex robotic manipulation skills from experts to embodied agents. In this work, we introduced MCNNs, a semi- parameerettric approach to behavior cloning that significantly increases the performance of behavior cloning methods across diverse realistic tasks and datasets regardless of the underlying architecture (MLP, transformer, or diffusion). While our theoretical and empirical results support the idea that appropriately constraining the function class based on training data memories improves imitation performance, MCNNs are only one heuristic way to accomplish this; it is very likely that there are even better- model classes in this spirit, that we have not explored in this work. Finally, we would like in future work to explore MCNNs as model classes beyond just behavior cloning, such in reinforcement learning and meta-learning.

---

Published as a conference paper at ICLR 2024
Acknowledgements: This work was supported in part by ARO MURI W91 1991NNF20-1-0080) NSF 2143274, and ONR N00014-22-1-2677 Any opinions, findings, conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views the Army Research Office (ARO), Office of Naval Research (ONR), the Department of Defense, or the United States Government.
Reproducibility Statement: To ensure reproducibility, we have released all code for MCNN variants and the baselines our website attps://siis google. cwm/viiiiviio..ooo vivvviee, eww monn-imi tation, We have also described all environments, datasets, baselines, and MCNN variants detail (with their hyperparameters) in Section and Appendix
10

---

Published as a conference paper at ICLR 2024
REFERENCES
[1] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04J33 2023. (Cited on 2, 7, 8, 9, 17, 18, 19)
[2] Felipe Codevilla, Eder Santana, Antonio M López, and Adrien Gaidon. Exploring the limitations of behavior cloning for autonomous driving. In Proceedings of the IEEE/CVE International Conference Computer Vision, pp.
2019. (Cited on 3, 7, 17, 18)
[3] Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. Carla: An open urban driving simulator. In Conference on robot learning, pp. 1-16. PMLR, 2017. (Cited on 6)
[4] Souradeep Dutta, Yahan Yang, Elena Bernardis, Edgar Dobriban, and Insup Lee. Mem- ory classifiers: Two-stage classification for robustness in machine learning. arXiv preprint 2rXiv:2206.05333 2022. (Cited on 18)
[5] Souradeep Dutta, Kaustubh Sridhar, Osbert Bastani, Edgar Dobriban, James Weimer, Insup Lee, and Julia Parish-Morris. Exploring with sticky mittens: Reinforcement learning with expert interventions via option templates. In Conference Robot Learning, pp. 1499-1509, PMLR, 2023. (Cited on 18)
[6] Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12873-12883, 2021. (Cited on 18)
[7] Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on Robot Learning, pp. 158-168. PMLR, 2022. (Cited on 7, 17, 18, 19, 22, 23)
[8] Bernd Fritzke. A growing neural gas network learns topologies. In Proceedings of the 7th International Conference on Neural Information Processing Systems, NIPS' 94, pp. 625-632, Cambridge, MA, USA, 1994. MIT Press. (Cited on 5, 6)
[9] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020. (Cited on 2, 18, 20, 21, 22, 23)
[10] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 10-20132-20144 2021. (Cited on 7)
[11] Abhishek Gupta, Vikash Kumar, Corey Lynch. Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. arXiv preprint arXiv: 1910. 1956, 2019. (Cited on 6, 18)
[12] Philippe Hansen -ansen-Estruch, Amy Zhang, Ashvin Nair, Patrick Yin, and Sergey Levine. Bisimula- tion makes analogies in goal-conditioned reinforcement learning. In International Conference on Machine Learning, pp. 8407-8426. PMLR, 2022. (Cited on 18)
[13] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016. (Cited on 2, 17)
[14] Kuk Jin Jang, Souradeep Dutta, Jean Park, James Weimer, and Insup Lee. Memory classifiers for robust ecg classification against physiological noise. Annual International Conference of the IEEE Engineering in Medicine and Biology Society (EMBC), 2023, (Cited on 18)
[15] Xiayan Ji, Hyonyoung Choi, Oleg Sokolsky, and Insup Lee. Incremental anomaly detection with
guarantee in the internet of medical things. In Proceedings of the Sth ACM/IEE Conference on
Internet of Things Design and Implementation. IoTDI '23, pp. 327-339, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400700378. 3582374, URL https:://///ec. http://oo..oooo////////o///////. 4145 / .e...222...2222. 3582374, (Cited on 18)
doi: 10. 10.1145/3576842
11

---

Published as a conference paper at ICLR 2024
[16] Ramneet Kaur, Susmit Jha, Anirban Roy, Sangdon Park, Edgar Dobriban, Oleg Sokolsky, and Insup Lee. idecode: In-distribution equivariance for conformal out-of-distribution detection. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 7104-7114, 2022. (Cited on 18)
[17] Ramneet Kaur, Kaustubh Sridhar, Sangdon Park, Susmit Jha, Anirban Roy, Oleg Sokolsky, and Insup Lee. Codit: Conformal out-of-distribution detection in time-series data. arXiv preprini arXiv:2207. 17769, 2022. (Cited on 18)
[18] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33: 1179-1191, 2020. (Cited on 7, 18, 19, 22, 23)
[19) Michael I Laskey, Jonathan Lee, Roy Fox, Anca Dragan, and Ken Goldberg. Dart: Noise injection for robust imitation learning. In Conference on robot learning, pp. 143-156. PMLR, 2017. (Cited on 2, 17)
[20] Minwoo Lee and Charles W Anderson. Robust reinforcement learning with relevance vector machines. Robot Learning and Planning (RLP 2016). pp. 5, 2016. (Cited on 2, 17)
[21] Ajay Mandlekar, Danfei Xu, Josiah Wong, Soroush Nasiriany, Chen Wang, Rohun Kulkamn,, Li Fei-Fei, Silvio Savarese, Yuke Zhu, and Roberto Martín-Martín. What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint a0X.v:208..03298. 2021. (Cited on 8)
[22] T.M. Martinetz, S.G. Berkovich, and K.J. Schulten. 'neural-gas' network for vector quantization and its application to time-series prediction. IEEE Transactions on Neural Networks, 4(4): 558-569, 1993. doi: 10.1109/72.23831 1. (Cited on 5)
[23] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. arXiv preprint arXiv:2209. 005888, 2022. (Cited on 18)
[24] Tianwei Ni, Michel Ma. Benjamin Eysenbach, and Pierre-Lue Bacon. When do transformers shine in r!? decoupling memory from credit assignment. arXiv preprint arXiv:2307.0388 2023. (Cited on 8)
[25] Ian Osband and Benjamin Van Roy. Model-based reinforcement learning and the eluder dimen- sion. In Proceedings of the 27th International Conference on Neural Information Processing Systems- Volume 1, NIPS'14 14, pp. 1476, Cambridge, MA, USA, 2014. MIT Press. (Cited on 3, 5, 18)
[26] Jyothish Pari, Nur Muhammad Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The surprising effectiveness of representation learning for visual imitation. arXiv preprini arXiv:2112.01511, 2021. (Cited on 3, 7, 18, 19, 22, 23)
[27] Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, et al. Imitating human behaviour with diffusion models. arXiv preprint arXiv: 2023. (Cited on 17)
[28] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. advantage-weighhed regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00117,, 2019. (Cited on 2, 17)
[29] Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural computation, 3(1):88-97, 1991. (Cited on I)
[30] Y. Prudent and A. Ennaji. An incremental growing neural gas learns topologies. In Proceedings. 2005 IEEE International Joint Conference on Neural Networks, 2005., volume 2, pp. 1211-1216 vol. 2, 2005. doi: 10. 1011111//NNNNN 200.000NNN220...555555 1556.2. (Cited on 5)
[31] Nived Rajaraman, Lin E. Yang, Jiantao Jiao, and Kannan Ramchandran. Toward the fundamental limits of imitation learning. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS 20, Red Hook, NY, USA, 2020. Curran Associates Inc (Cited on 17, 18)
ISBN
12

---

Published as a conference paper at ICLR 2024
[32]
Nived Rajaraman, Yanjun Han, Lin Yang, Jingbo Liu, Jiantao Jiao, and Kannan Ramchandran.
On the value of interaction and function approximation in imitation learning. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp- 1325-1336. Curran Associates, Inc., 2021. URL http:httpt/p/pptttt /ptoccee//c//ceeeeeee www..n....nppae ec wwwpeepppppppppeeeeeeeeeeeeeeeee fi files /pappppprrr r 22022/ fflff/ii/// 09 O=bcc11722.S1eeffeaaaliissssss32
pdf. (Cited on 3, 17, 18)
[33] Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. arXiv preprint arXiv: 10087, 2017. (Cited on 6)
[34] Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. arXiv preprint urXiv::7099 10087, 2017. (Cited on 2, 17)
[35] Aravind Rajeswaran, Kendall Lowrey, Emanuel V Todorov, and Sham M Kakade. Towards generalization and simplicity in continuous control. Advances in Neural Information Processing Systems, 30, 2017. (Cited on 2, 17)
[36] Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pp. 661-668. JMLR Workshop and Conference Proceedings, 2010. (Cited on 3, 17, 18)
[37] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth interna- tional conference on artificial intelligence and statistics, pp. 627-635. JMLR Workshop and Conference Proceedings, 2011. (Cited on 2, 17)
[38] Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya Arty Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning A modes with one stone. Advances in neural information processing systems, 25:22955-229668, 2022. (Cited on 2, 7, 17, 19, 22, 23)
[39] Devavrat Shah and Qiaomin Xie. Q-learning with nearest neighbors. Advances in Neural Information Processing Systems, 31, 2018. (Cited on 2, 17)
[40] Kaustubh Sridhar, Souradeep Dutta, James Weimer, and Insup Lee. Guaranteed conformance of neurosymbolic models to natural constraints. arXiv preprini 1rX1v:2012.01346, 2022. (Cited on 18)
[41] Kaustubh Sridhar, Oleg Sokolsky, Insup Lee, and James Weimer. Improving neural network robustness via persistency of excitation. In 2022 American Control Conference (ACC), pp. 1,5111126.. IEEE, 2022. (Cited on 18)
[42] Yihao Sun. Offlinerl-kit: An elegant pytorch offline reinforcement learning library. https: wttp:uun11///////1nee Kit, 2023. (Cited on 19)
//github. htw.ttt...coom////
[43] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017. (Cited on 18)
[44] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06199 2022. (Cited on 7, 17, 19, 22, 23)
[45] Yahan Yang, Ramneet Kaur, Souradeep Dutta, and Insup Lee. Interpretable detection of distri- bution shifts in learning enabled cyber- physscal systems. In 2022 ACM/IEEE 13th International Conference Cyber-Physical Systems (ICCPS), pp. 225-235. IEEE, 2022. (Cited on 18)
[46] Yahan Yang, Souradeep Dutta, Kuk Jin Jang, Oleg Sokolsky, and Insup Lee. Incremental learning with memory regressors for motion prediction in autonomous racing. In Proceedings of the ACM/IEEE 14th International Conference on Cyber-Physica Systems (with CPS-IoT Week 2023), pp. 264-265, 2023. (Cited on 18)
13

---

Published as a conference paper at ICLR 2024
[47] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint a7X1v:2006.10742, 2020. (Cited on 18)
[48] Tony Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware, arXiv preprini arXiv:2304. 13705, 2023. (Cited on 2, 17)
14

---

Published as a conference paper at ICLR 2024
APPENDIX
A EXAMPLES OF THE MCNN FUNCTION CLASSES
We demonstrate the effects of varying number of memories, X, and L. in Figures 9, 10, and 11 respectively. By increasing or decreasing the number of memories, we directly increase the width of the model class. By increasing X, quicken the interpolation from the nearest neighbor components to the neural network function class. Similarly, by decreasing X, we slow this transition.
Example function classes for A=1, 1=1, and 2 memories
Nearest Neighbor Function y= (xle -1
Neural Network Function Class y = L(1 - e 0) (x)
Memory Consistent Neural Net Class y (x)
5.0
10.0
10.0
5.0
10.0
X
Example function classess for A-1, 1-1, and 3 memories
Nearest Neighbor Function
Neural Network Function Class y = e dls. ((f(x))
Memory-Consistert Neural Glass y = gx00)
y= (xle
10.0
X
Example function classes for Am1. 1=1, and 4 memories
Nearest Neighbar Function
Neural Netwark Function Class y L(1 e-t o(f(x))
Memory-Consistent Neural (x))
(x)e
0
10.0
X
X
Example function classes for A=1, 1=1, and 5 memories
Nocrest Nolghbor Function
Ncural Network Function Class y L(1 e dlv.57 oxf
Memory Consistent
(* -1
(NC
0
0
X
Figure 9: Effects of varying number of memorics (keeping À and fixed) on the MCNN function class.
15

---

Published as a conference paper at ICLR 2024
Example function claoses for As0.5, and memories
Nearest Function (1l
Neural Network Function
(,00
0.0
2.5
5.0
7.5
16.0
00
2.5
5.0
7.5
16.0
2.5
5.0
7.5
10.0
Exampla classas 2-1,, Lm3, memories
Nearest heighter Function
Neura Networe Tunchion Class y-ill (=*")))))""0)))
Memory Consistent Neural Net Class
y- bxie
y Clil0
-2
-2
Prample funetion classes for 2 mm,, and mamortes
Nearest heighbor Function (Mlaeo-aa
Neural Networs Function Class LiL and sh (t0)
Menner Cormmssent Noural Not Cless y 65500
Figure 10: Effects of varying X (keeping L and number of memories fixed) on the MCNN function class.
keighhor -400
Memory-Consistenn y (5000
y-ill a100)
-1
-2
5.0
10.0
3.0
10.0
Example function classes for J-1, L=1, and memories
Nosmesi heighter Punttion Polieie
Neunal Network hunation Class Ull -e |o
Mernory Consistent Neural Net Class y thask
Example function clesars for 2-1,, L-2,
2-1, and mamories
Nevreet Keighter Bunction Aoli 146000
Neunal Networe Funehion Class - mac ec"
Consistent
0.0
2.5
5.U
7.5
10.0
0.0
2.5
5.0
7.5
10.0
0.0
2.5
5.0
7.5
20.0
Figure 11: Effects of varying L (keeping À and number of memories fixed) on the MCNN function class.
16

---

Published as a conference paper at ICLR 2024
B PROOF OF LEMMA 4.6
The width of the function class 3, A 01, 02 E e, and V S, PMC 0
defined
as
(PMC
(s) is upper bounded by 2L x - dais
max 0..4g
Proof.
Using Equation 2, express the difference the following:
18,.8 (1)
MC
I
e e Allos)))) (x)) L ((1 (1----dde.)) (1- (xf(
(3)
= [1 - e - (0(50))x)) - (z)) (x))
(4)
< (1 - - d81s)
10 (x)) (x))
(5)
Now, observing that (o(f - off (x)) < completes the proof.
C PROOF THEOREM 4.7
The sub-optimality gap J(x*) - (fA) S min H||A| (1- -1-e-Ah
dbis
Proof. order to take advantage of well-known results in the imitation learning literature [36, 31, 2, 32]. we restrict ourselves for the purpose of this analysis to the discrete cetion-space A scenario, where the policy S -> ((A). Even still, the intuitions developed through this analysis guide our algorithmic choices in continuous environments. The actions picked in the expert dataset D induce dirac distribution over the actions corresponding to each input state.
Recall that in imitation learning. if the population total variation (TV) risk T(ã,"") < 6, then, J(x)) ((t) S min H, H2 (See [31] Lemma 4.3). We note the following for population TV risk:
1 H
H
T(f, m²) =
Eg s
[V(+(+ss n (+ss)) < H
1
E S.
| T (A)
H
1=1
$ |A
dels
(6) where f². is the empirical distribution induced on state s' obtained by rolling out policy x². For the first inequality in the above derivation, we use Lemma 4.6. Using this, in the performance gap lemma gives us the following:
J( J(+) < min (H, (AAL (1 (1-e-A
D EXTENDED RELATED WORK
Compounding errors in imitation learning have previously been tackled by permitting online experience [13, 34], reward labels [28], queryable experts [37], modifying the demonstration data collection procedure [19]. Our work is orthogonal to these methods and creates model class that avoids compounding errors by construction. Other works that propose new models for IL such as Implicit BC (IBC) [7]. Behavior Transformer (BeT) [38], Action Chunking Transformer [48], and Diffusion Policies [44, 27] are orthogonal to our approach. MCNN can be used plug-in approach to improve any of these methods. In fact, we show that MCNN with BeT backbone outperforms vanilla BeT and MCNN with a diffusion model outperforms diffusion BC on all tasks in experiments in Section 5. We also show that MCNN outperforms IBC in Section 5.
Non-parametric and semi-parametric methods in imitation learning such as nearest neighbors [39], RBFs [35]. and SVMs [20] have historically shown competitive performance various robotic control benchmarks. But, only recently, a semi-parametric approach consisting of neural networks for representation learning and k-nearest neighbors for control proposed in Visual Imitation through
17

---

Published as a conference paper at ICLR 2024
Nearest Neighbors (VINN) [26]. This is the closest paper to work and in Section 5, we compare with VINN and demonstrate that we outperform their method comprehensively.
Theoretical guarantees the sub-optimality gap in imitation learning with MCNN are provided in this paper. Such guarantees are not available with vanilla neural networks. Our theorem builds on earlier work reductions for imitation learning in [36, 31, 2, 32] and leverages intuitions from [25] on bounding the width of the model class.
Learning a codebook of prototypes, like our memories, has been previously explored for image
reconstruction [43, 6], physics-consrrained learning [40], online RL [23, 47, 12, 5], interpretable OOD detection [45, 16, 17, 15], robust classification [4, 14, 41], and motion prediction [46]. The closest related usage of memories that are representative of the topology of the imput space is in [40]. But, here, external information in the form of physics and medical constraints plays key role in enforcing constraints at these pivotal points. In this paper, our prior is simply a kind of "consistency" with no external information utilized.
E COMPUTE
We ran the experiments either two Nvidia GeForce RTX 3090 GPUs (each with 24 GB of memory) or two Nvidia Quadro RTX 6000 GPUs (each with 24 GB of memory). The CPUs used were Intel Xeon Gold processors @ 3 GHz. Videos of our policies for many random goals across environments as well as our code can be found in our website ht https: //sites///eeeeee wwwggo.oooooooooooooo com/ vi view/menn-ime tation.
F DETAILED EXPERIMENTAL SETUP
Additional information about our task definition: Following notation in prior works [9, 18, 7], we define an imitation learning task tuple of (environment, reward function, and demonstration dataset), Since we perform experiments with both "expert" and "human" demonstrations, our experiment list has the following tasks: pen-human-y pen-expert-vl, hammer-human--i hammer- expert-v1, door- hurrhhumnn-l. door-expert-vl, relocate- human-vi 1,, relocate-expeeeeeeeeeeeeeeeeeeee expert-vl, and carla-lane-vO.
relocate-human-v/
Additional information for Adroit environments: All policies learned in the four Adroit envi- ronments are goal-conditioned. The goal positions and orientations are provided as part of the observation vector. Every time that the environment reset (such as for new evaluation rollout), the goal is randomly chosen. These include random goal orientations of the pen, goal locations for the relocate task, door locations in the door task, and nut-and-bolt location in the hammer task. In Adroit human tasks in particular, with only 25 demonstrations, it is not possible to cover every random goal location for each of the four environments (such as every possible pen orientation) in the training dataset. This increases the difficulty of the four adroit tasks. Even with this challenge, MCNN methods outperform baselines across environments and demonstrate their ability to generalize even from small datasets of demonstrations.
We also provide videos of the Adroit human demos in our website. The multimodality in these datasets collected by humans wearing VR gloves visible in these videos. Some examples include (1) in the hammer task, some demonstrations move the hammer above the bolt before hitting it, some move it below, and some directly hit the bolt and (2) in the pen task, some demonstrations have the little finger above the pen and some have it below the pen throughout the episode. MCNN helps address the significant compounding error challenge which MLP-BC fails to address and hence results in MLP having an overall better performance than Diffusion-BCC even though Diffusion better equipped to handle multimodality [1]. Our results in Franka Kitchen (see Appendix G) highlight this where MCNN+Diffusion improves upon the already strong Diffusion-BCC BC baseline.
Additional information for Franka Kitchen environment: this environment, a Franka robot arm interacts with various items in Kitchen such as kettle, microwave, light switch, cabinets, burners, etc. We use the multimodal relay policy learning dataset [11] with 566 demonstrations collected by humans wearing VR headsets interacting with objects (out of total 7 in the environment) in each episode. The goal is to execute as many demonstrated tasks as possible and the metrics capture the probability of successful interactions with various numbers of objects. These metrics, following
18

---

Published as a conference paper at ICLR 2024
previous work [1], have been named pl to p5 for the probability of success in interacting with 1 to 5 objects respectively.
Normalization of inputs: We normalized all observations by subtracting the mean and dividing by standard deviation. We didn't have to normalize the actions as they are already in the range [- 1, 1] but if we are learning transition models instead, the outputs (next state reward) would have to be normalized.
Our tanh-like dynamic activation function: We plot our activation function (.) described in Algorithm 2 in Figure 12. It is exactly like the tanh function for 33 = 0 as seen in Figure 2. Further, as given in Algorithm 2's Line 3, we set / max (0, 1 step )) Hence, 8 = 0 after 100 steps until 1 million steps during training. We also set 8 = 0 during inference. For 3 = O, our activation function reduces -1 for æ < -1, z for -1 S x S 1, and 1 for x > 1. This is exactly like tanh. The reason for using this activation function is the very few initial steps when it is not like tanh where gradients are available beyond [-1, 1] during which time, the neural network component adjusts for the presence of the nearest neighbor component. We use the standard tanh activation function for our reimplementation of BC.
Implementation details for baselines and MCNN+MLP: In BC and MCNN+MLP, we use an MLP with two hidden layers (three total layers) of size [256, 256] for Adroit tasks and [1024, 1024] for CARLA. We use an Adam optimizer with a starting learning rate of 3e 4 and train for I million steps. We simply minimize the mean squared error for training the policies. We use a batch size of 256 throughout. We describe the MCNN-specie memories in another paragraph below. For all other BC hyperparameters, we use the recommended values in the TD3-BC implementation in [42].
hyperparameters, namely X, L,, and number of
For VINN, we set k = 10 in the Euclidean weighted k nearest neighhors algorithm. This value was recommended by the original paper [26).
For BeТ,, we follow the official implementation'. We use 6 layers, 6 heads, and an embedding dimension of 120 in the transformer model. We also 64 clusters for action discretization performed on the actions seen in the first 100 steps. Since the original paper [38] did not run experiments on D4RL tasks, we ran sweep over various choices of number of layers (4, 6), number of heads (4, 6), and embedding dimension (32. 64, 120). Following previous work [7]. we chose the hyperparameters (6 layers, 6 heads, 120 embedding dimension) with the highest average scores on three human tasks. We used these hyperparameters on all other tasks as well. For all other Be'T hyperparameters, we use the recommended values in the official BeT implementation.
For Diffusion-BC we use the official implementation from [44] in Adroit tasks and the official implementation from [1]* in Franka Kitchen tasks. We also use the recommended hyperparameters provided in the code for pen-human-vl in the former repository in all human and expert Adroit tasks. We similarly use the recommended hyperparameters provided in the latter repository along for the kitchen task.
In CQL-Sparse, we use an MLP with three hidden layers (three total layers) of size [256, 256, 256] for Adroit tasks. We give a reward of 1 for the last timestep in both expert and human datasets where each trajectory achieves task completion. Hence, we run CQL-Sparse only on the human and expert datasets where each trajectory has achieved task completion. We a reward of 0 in all preceding timesteps. For all CQL-sparse hyperparameters, we use the recommended values in [18]. We also ran a sweep over the three key hyperparameters of CQL, namely actor learning rate (3e 5, le 4), initial a (5, 10). and Lagrange threshold (5, None). We note that while performing the sweep, if the Lagrange threshold is None. CQL was run with fixed a. Otherwise, a is tuned automatically, as described in Kumar et al. [18), based on the Lagrange threshold and starting from its initial value. We found that the recommended values in [18], i.e. actor learning rate 3e 5, initial a = 5, and Lagrange threshold = 5 performed the best overall.
Implementation details for MCNN + Behavior Transformer (BeT): We train MCNN+BeT following the official BeT implementation described above with one major change to the action offsets output by the BeT model. Let us denote the sequence of input observations as 7 and the action
hhtpps: htttt////sttttt...ooeeeeeaaaaasa
19

---

Published as a conference paper at ICLR 2024
2.0
000)
8-th
1.5
030
0.5
xtx
0.0
-1.0
-2.0
-2.
-1.5
-1.0
-0.5
0.0
0.5
1.0
1.5
2.0
X
Figure 12: Dynamic tanh-like activation function σạ(2)) shown for B and 8 =
offsets output by the BeT model as (Be Then, the output of MCNN+BET is given as follows:
MMC ( = f (r) (6 - - + (1-E-- -()(r)))))) 20 (f
(7)
where 7' is the nearest (memory) sequence to and f N retrieves the corresponding sequence of actions from the codebook. We also multiply and divide by two inside the tanh-like function since each item in the sequence output by the BeT lies in 2_
Implementation details for MCNN + Diffusion: We augment the diffusion process to use memories in every step. Rather than predicting the noise value, we predict the true values. This amounts to minor change to the loss function within the official implementation. We provide this code in our repository.
MCNN-specific hyperparameters: We use a value of L 1.0 for all runs. This is suitable for BC because our actions the range of [-1, 1]. For both MCNN+MLP(Fixed) and MCNN+ we use the same set of hyperparameters across all tasks. In particular, the MCNN-specifh hyperpa- rameter values are À 0.1 and 10% memories. We show an abltaion study of normalized scores with varying number of memories (2.5%-80%) in Figure 14. We also show an ablation study with both varying X (0.1, 1.0, 10.0, 100.0) values and number of memories (2.5%, 5%, and 10%) in the 3D bar charts of Figure 15. We also tabulate the highest value across all MCNN-specific hyperparameters (obtained by evaluating on 20 trajectories) in the MCNN+MLP(Tune))) and MCNN+BeI(Tuned)) columns of Tables 2 and 3.
Implementation details for Figure We run MCNN+MLP experiments on randomly sampled subsets of the human and expert task datasets and report our performance normalized with that of the D4RL BC scores from [9]. This means that for many points in Figure 1 we used a smaller set of datapoints from partieular task for training an MCNN+MLP policy and yet performed better than D4RL BC's score on that task.
Discussion on computational costs: by the nearest neighbors search among the M memories in an observation space Rd, From the neural gas, obtain graph connecting memories. The search for the nearest memory is O(Md) d) for the first observation. For every subsequent observation, we simply perform a nearby breadth- first search starting from the previous memory. This is O(nearby search depth * d) in the average case. Please note that the nearby search depth is kept small (1, 2, 3, 4). The search for the first memory can also be further reduced O log M) with the K-D trees data structure. Alternatively, leveraging the parallel processing power of GPUs, this search can be done very quickly (<1 ms) in practice even without a graph and by simply searching amongst all memories for the closest memory.
The computation cost of MCNN during inference is dominated
Further, since the vanilla neural net is a component of MCNN, we are less efficient at in- ference than the vanilla neural net component alone. But, MCNN is more efficient than baselines like VINN and I-NN since we only have M memories, and this value is set to 10% of the total number of datapoints lower. VINN and 1-NN have to perform nearest neighbors search much larger
20

---

Published as a conference paper at ICLR 2024
dataset than MCNN. Further VINN requires the top neighbors (usually k=10) while require only the nearest memory.
G ADDITIONAL FIGURES AND DETAILED RESULTS TABLES
We tabulate all reward values from our previous bar charts in Table 2. We tabulate all corresponding normalized scores obtained following the methodology in D4RL [9] in Table 3. We plot all the ablation results for replacing neural gas memories with random memories for human and expert tasks in Figure 13. We plot performance variation against number of memories (from 2.5% to 80%) for all human tasks in Figure 14. We also plot comparison of scores with varying Às (0.1, 1.0, 10.0, 100.0) and varying number of memories (2.5%, 5%, and 10%) for each task in the 3D bar charts of Figure 15.
21

---

Published as a conference paper at ICLR 2024
Table 2: Comparison of returns (across 20 evaluation trajectorics and 3 random seeds) between basclines, our reimplementation of BC, and our methods MCNN+MLP, MCNN+BeT, and MCNN+Diff (see smaller table below) on the Adroit and CARLA domains. Here, we show both our MCNN results with the same fixed set of hyperparameters across all tasks (Fixed) and with hyper peerpaaametes. tuned online (Tuned).
MCNN Ber (Tumed Hypers)
2103 = 101
1604 12
4417 = 297
10 4: 08
558
-0.3
-53 #-0.3
-58 =01
-130 +:
1038 =
1089
-228
390
MCNN + Ber
(Fieed Hyyens)
2033 + 1.8
1595 1: 15
3603 - 063
8.2 42
-2.9
0.
902 902 50
-58 001
+:
-150
-229
558
300
MCNN MLP
4051 4051 + 195
2820 # 119
16387 + 902
-4.7 -.7.4.00 03
-6.0 2:47
9.0 9.0 29
3035 3035
-59 00
-465 213
(Tuned
262
-155
466
MCNN + MIPP
(Fiked Hypers))
3947 : 227
16027 2 365
-235 2:33
-4.7 3. 03
-8. 1:0
-54 4 15
3033 0.3
-59 03
5511 -51 210
262 +:
5285
2208
MLF-BC (Relmpl.)
3194 = 127
1806 72
11 = 118
13710 = 2002
5.2 5.22 1.7
.9.0 20
-33 23
2798 55
-59 0.0
-997 4 128
1845
-232
358
Implicn BC [7]
132 132
-0.1 1:21
4- 61
2586
SpL-Se [18] + 011
671 1
2241 0,3
11311 = 502
-20 0.04
.66 0-03
#
2731
MCNN DIY (Tuned FlynerFyyyrrr)) Flyeer))
2345 16005 3568 63 28- 05
1300 28 15421
00 08.87
M0 16.06
12 + 1477
1055 # 1010
0.78+ 382
3157 = AN
1909 a: 35
-232 # 10
-4.7 0M
-000 10
-3 944
2760 7-60-22
-59 =00
315 : 102
16181 16181.622
VINN [26]
197.254
Bacellnes
1490
-230
1283
297
I-NN
3102 275
148
666 + 20
10069 = 770
-1.7 -4.7 0.4
23521
-25 85
2716
-99 29:08
-458 + 179
MCNN DII
(Fixed Hypers)
14:00
35166 16.772 284 64
16088 67.96
2261 1:20
09 59.27
3000 3000 302
1982
3102
1902
-207
1095
348
1 17.89-L
572.252
BeTBB [38]]
+ 117
28
261
2,7 1 16
01
0.1
296
002
1853
1348
1189
2731
225
-54
560
559
327
3194 9611 8611-2044 40
159885 56.00
14044842 131.54
4.314 1.70
12:72
121.774 4494
DIF-BC [44)
41
2800.39+
DARL
[9]
2633
1792
-79
16140
-0.4
-11
969
-59
325
-161
hummer humun-
hammer-exxe mummer-exper--e exper--rl
mamner--lonsee cloned-vi
relocat--iine relecate-hmman- human-- V
relocate- meeat--oxpert expe--eee--tl
relocare-d clneedd-vvv
hummer nummer-humun human-y
cammer-oxper-sr
hammer-<lonese c-noed-v
relocate peecate-hmminn human v
relocate-expert expeep-e---l
relocate- cloned-vv
pan human-n
pen con-expee--v
pan-cloned-vl
door human-- v
door expertx
Moor-clooed--v
carlo-lanexx0
carlu Nown- v
Name
pen- human-ul v
pen-expen-e
nun-cloned-vl
door huur-hhman---- v
door experteppet-v
door-clonad-v
carlo-lame- lame-10
carla town v
22

---

Published as a conference paper at ICLR 2024
MCNN + Ber (Tuned Hypers)
67.34 + 3,30
50.6 1:04
1.11 1:04
35.9 227
0.394 :: 0.16
0.145 + 0.00
0.11 4 000
37.27 77 4.4
4008 1 0.005
Table 3: Comparison of normalized scores (across 20 evaluation trajectories and 3 random seeds) between basclines, our reimplementation of BC, and our methods MCNN+MLP, MCNN+Be1I and DCNN+Diff (sce smaller table below), on the Adroit and CARLA domains. Here, we show both our MCNN results with the same fixed set of hyperparameters across all tasks (Fixed) and with hyper perrraaaeetess tuned online (Tuned).
0,36
38.2
MCNN Ber (Fixed Hypers)
64,98 0.06
50.3 0:52
0.2
1.11 +
29.69 29.69 =5
1 1:000
0.315 1.000
0 0.08
0.08 000
-005 00005
38.1310000 1:006
32
32.0
0353
MCNN MLP
(Tuned Tyyeer))
132.70 = 6,55
91.38 5.00
4.11 1.002
127.49 300 300
0,92 1.447
0.04 1 0.02
1.10
0.01 0.00:0000 1000
2:23 1.00
105.26 1:24
-0.07 #
45,59 1.00
133.70 + 8401
MCNN M MLP
129 20 1 7.63
70:36 2:75
4.11 1:11--022 0.22
124.74 233
0.04 1.0 1 002
-004 -004 001
1,74 051
-010 001
43.14 1
11.52 :: 8.20
(Fixed
10S.
MLP-BC (Reimpl.)
103 4.22
57.36 2.4
107.01 = 15
0.03 0.04
-0.06 00.0
0.13 008
97.2 22 1.13
10:000011 0.01
35.05 3.000
14.95 14.95 3M
103 #
2.22
0.33
Implici BC
1.09 0.39
0.15 1..0
14.22 2228
171
CQL-Span Sparse [18] 01 reward
19.28 1: 0.48
0.26 DIME
88.65 3384
-0.315 1.115:::00 0.01
-0336 0.11
94.9 4.
102.68 2.94
00.81 4: 1:38
0.33 4 0ON
0,34 8-78
0A 1 coN
-0.06 0,01
0.01
(.61 032
95.88 0009 0.09
-0077 -0.77 = 0.00
29.09 1.6
-7.85 # 4.00
VINN
[26]
46.77
MONN DHE
46 5.4
116.5 9.53
12592 4: 0.68
0.17 0.00
108.04 1004
8.64 + 1:11
105.16
311
100.3 1 923
60,60 1.96
79.15 # 5.9
0.04 0:0 001
+0.05 0.0
1.07 0.29
94.39 10
-0.08 0.000
34,03 2.3
13.05 + 700
I-NW
Ours
25.97
1.60
0.2
MCNN + DIY (Fixed Hypers)
114.76 9.35
2.24 1.098
125.21 1 0.2
0.09 0.09 0004
107.08 11038
8.02 4: 195
104.88
78.18
BeT-BC
58.95 = 3.91
0.94
0,56 0100
230 230=20 =30
0.089 1. 0005
1.
0.036 0.03610000 0003
01096 # 0002
20198 87
00.08 0007
[38]
25.0
0,302
320
Baseline
Dir--C
103.96 14.3
0,88 2:03
109.57 4: 033
0,05 00
103.011:0
6.07 1.53
97.27 0:38
DARL
[9]
85.1
56.9
1.5
1256
0.0
-0.1
0.5
34.9
-0.1
-1.8
hummer humun-
hummer-exee mummer-exper-e exper--rl
hommer-cloon cloned-vi
relocate-iin relocate-humann human-- V1
relocate- mncate-oxpert eepe---ee--tl
relocare-clonss -lndd---
mommer-hummn-y/
hommer cloned.vi
relecae-bmmmn hommmmon I
relocate allecat--loneee cooedd-v v
pan human-
pen- con-expen-vv
pon-cloned-vl
door human--) vI
door expertx
Moor-clooed-vl
carla carlu-towoe bown- v
Name
pon- nn--human-vl
pen experx-r--v
pan clonedal
door human-v]
door-cxpert-]
cloo-clemed-vi vI
carln-town--d x0
23

---

Published as a conference paper at ICLR 2024
pon humanv v2
Ooor human
hammer human vl
relcode human v1
expert V1
door export v1
8000
-5
200
3040
2500
-10
3020
2000
Det m
Deb m
-15
1500
1000
-20
3700
-25
-000
3000
Nescal Dat
Random
Reoral Day
Rendom
Nenoral Gat
Nandoo
Nexuall Oess
Remdion
Netral Das
Random
Netral Dat
Random
horm mer expert
relocate expert vl
14500
16000
130M0
M00
14500
4300
4200
13500
4200
13000
Necral O1
Nesural Olt
Rondom
Figure 13: Neural Gas vs Random Memories MCNN+MLP: Comparison of returns (across 20 eval trajectories) between our method with neural gas based memories and randomly chosen memories. This shows that MCNN+MLP performs better with neural gas memories. We attribute this the spread-out nature of neural gas memories that reduces distance the most isolated state, improving imitation performance.
pen pen-hurnan hurnan-v
door-human-v1
3500
3000
Return
2000
Return
1500
1000
-40
500
2.5 gigh 10th 30th " woth s0t both yoth 50t%
2.5high
10%
10"
200
40th
s0ts
60%
70%
80%
Num Memories
Num Memories
hammer-humann-i1
relocate-hue -occatt-human-v1
200
-2
Retur
150
Retum
-3
100
50
shsh gothystt ass yash oott york 80%
25 10th thth 300 with
sort yort 80%
Num Memories
Num Memories
Figure 14: Number of Memories in MCNN Comparison of returns (across 20 evaluation trajectories)
with our method (MCNN+MLP) using 2.5% of the dataset as memories up to 80% of the datasct as memories. We find "sweet spot" for num memories 10-20%. We also see the expected decrease 1 I-NN performance with num memories increases to 100%.
24

---

Published as a conference paper at ICLR 2024
pen-human-v1
pen-expert-v1
pen-cloned-v1
4000
3000
2300
2000
2000
1500 8
2000
1000
1000
1800
500
100.0
100 o
100
2.5
10.0
10.6
10.0
Memories (%6
1.
Memories (%)
Memories (%)
1.0
10.0
10.0
hammer-hhuan--11
hammer-expert-vl
hammer-cloned-11
200
150
1o000
200
2200g
1002
-2002
200
100.0
100
100
2.5
2.5
2.5
Memories (%)
50
T
Memories (%)
5.0
Memories (%)
0.1
0.1
10.0
0.1
relocate-oxpert-1
relocato-clonedvv
-
Return
4000
Retur
Retum
-8
1nco
100.0
100 o
100
10.0
10c
10c
Memories (%6)
1.0
>>
Memories (%)
5.0
10.0 Memories (%)
1.0
10.0
10.0
door-human-v1
door-expert-v1
door-cloned-v1
700n
Return
Retur
80
1000
100.0
100
100.0
10.0
2.5
10.0
Memories
Memories (%)
Memories (%)
01
0.1
carla-lane-vo
carla-town-vO
400
guum
5500
130
550
100
Mcmories (%)
5.
Memories (%)
5.
30.0
10.0
Figure 15: Bar chart of returns for our method against various combinations of À and memories for each task. Each of the 12 bars in each of the 14 subfigures represents the average performance across 20 evaluation trajectories and three seeds. We notice that the best performance can be obtained for values at the middle, iee., X E 10.1, 1.0), for any number of memories. This is where our method can interpolate, by design, between the nearest memories and vanilla BC. These plots also demonstrate that by training MCNNs on offline data for few sets of hyperparameters and simply choosing the best hyperparameter with limited online interaction, we can obtain significant improvements in performance.
25

---

Published as a conference paper at ICLR 2024
Figure 1 values: We provide the mean and standard deviation of the return values for various number of demos for various tasks (across various runs) that were in the scatter plot in Figure Tables and 5. Other points plotted in Figure include the return values (across various runs) for all 9 tasks at the full dataset sizes (25 demos for Adroit human tasks, 5000 for Adroit expert tasks, 400 for the CARLA task) given in Table
Table 4: Returns with various methods on pen-expert-vl and door-expert-vl for subsampled and full datasets (where NA Not Available). These values represent the means and standard deviations obtained from various points (corresponding to various runs) in the scatter plot in Figure
Number of demonstrations
Method
100 domos
500 domes
1000 demos
2000 domos
4000 demos
5000 demos
pen-
MCNN+MLP
2724.61
2931.19
2991.88
3016.65
3025.47
3035
9.8
+
DARL BC
969
door-
MCNN+MLP
NA
3711.79
3807.9
3857.74
3933,75
4051
expert-v1
MCNN+MLP
31.6
5,58
29.16
41.83
4: 195
DARL BC
NA
2633
Table Returns with various methods on pen-human-vl and hammer-human-vl for subsampled and full datasets (where NA Not Available). These values represent the means and standard deviations obtained from various points (corresponding to various runs) in the scatter plot in Figure
Number of demonstrations
Task
Method
demos
10 demos
demos
20 demos
25 demos (full daraset)
pen-
MIP
189.76
128.38
100,93
+
= 728
human-v/
D4RL BC
NA
1122
hemmer--
MCNN+MLP
23.29 5.5
36.19 1 3.27
113.73 1 3.08
134.24 1: 8.0L
233.37 1 3.26
buman-v1
NA
-79
Oversampling memories: We also compare MCNN+MLP with vanilla MLP-BC algorithm that oversamples the memories which comprise 10% of the dataset and report results Table 6. It can be clearly seen that MCNN+MLP significantly outperforms the oversampled MLP-BC for various amounts of oversampling. Moreover, MLP+BC with oversampling performs similar to vanilla MLP+BC.
Table 6: Comparison of MCNN+MLP with vanilla MLP-BC with memories (or codebook entries) oversampled by various amounts. Here, we use Nx sampled' to denote that each memory was sampled N times in each training epoch.
Name
Baselines
Ours
MLP-BC
MCNN + MLP
MCNN MLP
normal
2x sampled
3x sampled
6x sampled
(Fixed Hypers)
(Tuned Hypers)
pen-human-vi
1845 213
2080.13
2093.24
2114.11
3285 1: 209
3405 328
pen pen-expert-vl
3194 + 127
3106.64
3136.44
3285.47 47
3947 +: 227
4051 + 195
hammer-human-se
-11
-10.66
-10,86
-4.33
262
262
hammer expert-v1
13710 + 2002
13447.58
13447.55 55
13970.34
16027 383
16387 392
human-yl
-5.2
-5.94
relocnte-cpert-vl
4361 =
4403.49
4403.49
4445 89
4566 4:47
4566 47
doos-buman-y1
-53
-48.64
-46.79
-35.6
-5.4 1:
9.0 1:
expertvptt
2798 33
2821.83
2851.2
2909.94
3033 +
3035 +
26


Reviews:
[{"review_id": "84bNH9gnWd", "reviewer": "Reviewer_Cf33", "summary": "This paper proposes a method for behavior cloning to address the i.i.d. assumption violation that occurs under supervised learning algorithms. Specifically, a memory-consistent neural network is proposed in which a set of \"memories\" are used to constrain the output of the neural network to stay within permissible regions. The proposed method is compared to baseline imitation learning algorithms across a set of standard manipulation and autonomous driving tasks.", "questions": "1) Given the distance lookups that are necessary for finding neighbors, how does this impact the wallclock time for training and inference?\n2) It seems like in some tasks the model is very sensitive to the number of memories, e.g. carla-town-v0 in Fig. 15. Do you know why this is the case? It seems unintuitive that 5% would perform worse than both 2.5% and 10%. Do  you think that the number of memories needed to attain sufficient performance is related to task/behavior complexity?\n3) It seems like there might be a relationship to critical/important states which exhibit a large difference in the expected reward between the best and worst action, i.e. it is costly to recover from a mistake in a critical state. It may be interesting to do an analysis where you identify the critical states in a task and then ablate them from memory. Do you have any thoughts on whether you can do more intelligent memory selection?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Constraining the output of the neural network such that it stays \"close\" to a specified set of data points is an interesting approach for tackling behavior cloning. Even more so since this can be combined with underlying network architectures.\n* The empirical evaluation is fairly thorough, and shows reasonable performance gains across many of the tasks.\n* The implementation seems straightforward (although this will incur an additional computational cost during training and inference).", "weaknesses": "* As a (semi-)parametric method, the computational cost of training and inference scales with the number of memories. There is a discussion on computational complexity in the appendix, but some analysis on training time would be appreciated here as it's difficult to tell whether this is a significant factor.\n* The performance improvement seems sensitive to the underlying network architecture and the task. E.g. in Fig. 4, different models exhibit different levels of improvement for each task. This could make it difficult to use in cases where we can't sample from the environment.\n* Similarly, it seems that the clustering method requires you to be able to sample from the entire state space, is this the case? If it is not possible to sample from the environment at will, is it possible to do this clustering purely from data? If the method's performance strongly depends on sampling from the environment, then it would only be fair to compare it to interactive imitation learning methods that also sample from the environment like SQIL, GAIL, etc.", "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": "", "extra_scores": {}}, {"review_id": "mqofesTHpO", "reviewer": "Reviewer_kyQU", "summary": "In this work, the authors present their work Memory Consistent Neural Networks (MCNN), that aims to improve the performance of traditional behavior cloning techniques by introducing a novel demonstration handling technique, which aims to reduce compounding errors that can commonly occur in the BC framework. The authors augment the BC procedure by incorporating demonstrations provided by the expert, directly into the policy’s decision step. They accomplish this by utilizing a convex combination between the expert’s action for the most similar state encountered in the demonstration set, and the output of a parametrized model approximating the policy. They offer both intuitive and theoretical motivation for their strategy: 1) that past similar experiences can help ground the system’s response to behaviors that have been displayed by an expert and thus avoid drifting into potentially catastrophic state-space areas and 2) offer theoretical guarantees  that indeed the system’s output is constrained within a desired area according to the provided hyperparameter L. MCNN requires that a subset of the available demos be organized in a particular form of memory buffer, referred to as Memory Code Book throughout the work. . Selected demonstrations are structured into a graph G, where each node is connected to its nearest neighbor according to some provided distance function d, defined over the state space, and the edge stands for the action used to transition to that nearest neighbor. They argue that within this framework, only a subset of the training data is required to craft graph G. This graph G, is then used to craft a function F(x) that locates the nearest known state s to x, and outputs the action a that the expert has demonstrated from that state. Using a convex combination between this action a, and the action a’ the policy network predicts for the input state x, the system finally outputs the action a_final as the predicted action for state x. Finally the policy network’s parameters are updated by standard Stochastic Gradient Descent, using a loss signal over the predicted a_final and the expert action for state x.\nThe method is then compared to several BC based approaches over 5 simulated tasks where reward signals are known and can be used to provide quantitative assessment.", "questions": "1) What distance measure did you use for your problems? How do you see the performance of your method be impacted from a different measure choice? Would a more generic measure adversely affect performance compared to a more tailored one? I.e Aggregated Manhattan distance for images compared to some kernel-based distance.\n\n2) Relevant to the generalization point raised in the weaknesses section, how do you believe would MCNN behave to inputs based on varying levels of state - perturbation?\n\n3) What was the motivation behind a static nearest neighbor function? Have you considered training it as well? (The authors have clarified that the distance used is a generic L2 distance)", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The authors present all the material in a very easy to follow manner. All figures are good quality and the results are well displayed.\n\nThe authors have performed extensive experimentation and comparison to the baselines. In addition, the baselines selected are reasonable as this work aims to improve the behavior of BC methods. Finally, they offer adequate implementation details in their appendix.\n\nTheir method is predicated on a simple yet elegant idea: use expert demonstrations as a human-like memory of experiences to constrain the output of a BC model. In this manner the output is kept within a reasonable distance of the known state space. This, assuming good coverage provided by the expert, can significantly reduce compounding errors that can affect BC techniques (i.e DAGGER). Additionally, the system can potentially better handle not-too-dissimilar unseen states by virtue of projecting them to the known state space via the  nearest neighbor function capability.\n\nSound theoretical justification for the work’s motivation. Essentially the authors argue that by integrating known state-action  information into the decision process directly they can constrain the output range of the system. They also provide reasoning as to why the code book can only be a subset of the total amount demonstrations, as long as it is crafted in a sophisticated manner. That is, the code book can be small enough to facilitate  fast inference, as long as the selected state-actions are good representatives of the known state-action space.", "weaknesses": "The main weakness of this work, as is common with BC approaches, is assumption of representative state-action pairs and optimal behavior provided by the demonstrator. There is no insight as to how the method will behave with reasonably suboptimal demonstrations, as it being a BC based approach has no apparent mechanism that enables it to focus on the better demonstrations of a provided set. \nThere is recent interest in work that can discern between useful demonstrations and harmful / irrelevant ones that can inhibit training. (i.e TREX, or Subdominance minimization (Ziebart 2022)). While understandably not the direct focus of this work, it can be inhibiting for complex tasks to assume a plethora of very carefully curated demonstrations that are both optimal and extremely descriptive of the entire problem space, hence the aforementioned attempts to augment imitation learning with such capabilities.\n\nIn the same direction, there is no experimentation for generalization. For example, with demonstration provided for task A, how would the method perform for a slightly altered task A’ there perhaps the goal has changed slightly, the demonstrations are perturbed by noise or simply slightly different. It is not unreasonable to believe that this method would not perform poorly in such a setting and good results would strengthen the method’s appeal considerably.\n\nFinally, the distance used to select the nearest neighbor, is quite critical for this method and problem specific. For example, the distance needed for an image-based problem could be quite different from a control state-based problem. This can potentially deduct from the appeal of a traditional BC approach that doesn't need such an engineered mechanism. See question 1 and 3 below. (The authors have clarified that the distance used is a generic L2 distance across all problem spaces considered.)\n\nHaving revisited the sections pointed out to by the authors, on the parts that I had concerns, namely the generalization aspect and the problem specific distance metric, I believe this work can have utility in the low data data sample regime, and thus I have updated my rating.", "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": "", "extra_scores": {}}, {"review_id": "vWVb9kgP8N", "reviewer": "Reviewer_5EEF", "summary": "The paper proposes a semi-parametric method for behavior cloning, MCNN,  that combines non-parametric nearest neighbor based policy learning with paramteric neural network based policies. As a result of this amalgam, the authors show that the MCNN class of functions are bounded in width and the suboptimality gap, something that does not exist for vanilla neural networks. Further, the authors provide results across 5 environments where adding MCNN to a new architecture consistently improves the results.", "questions": "It would be great if the authors could address the points mentioned in “Weaknesses”.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper proposes a new class of functions that combines non-parametric nearest neighbor based policy learning with parametric neural network based policies. As a result of this amalgam, the authors show that the MCNN class of functions is bounded in width and the suboptimality gap, something that does not exist for vanilla neural networks.\n- The authors provide results across 5 environments where adding MCNN to a new architecture consistently improves the results.\n- The authors ablate the performance of MCNN across a different number of memories and the optimal number of memories is significantly smaller than the entire dataset. This highlights the efficiency of the non-parametric portion of the algorithm as compared to approaches like VINN and 1-NN.\n- The proposed method also seems to work on image inputs from the CARLA environment.", "weaknesses": "- The authors mention that they provide results on 9 tasks across 5 environments. But I only see 5 tasks, 1 per environment. It would be great if the authors could clarify the 9 tasks that they evaluate on and where they have provided the results.\n- For CARLA, the images have been embedded using a fixed off-the-shelf ResNet34 encoder. This might not be ideal for more complicated visual scenes such as the Franka Kitchen environment used in BeT and Diffusion Policy. It would be great if the authors could evaluate MCNN in this benchmark and provide some comments on possible pretraining approaches for the encoder when such an off-the-shelf encoder is insufficient.\n- *“ MCNN can even improve the performance of simple MLP architectures to beyond that of more sophisticated recent architectures such as Diffusion models.”*: In quite a few cases, MCNN works better with MLP than with BeT or Diffusion Policy. Both BeT and Diffusion Policy were developed to deal with multimodal data distributions. Hence, I wonder if the tasks shown are not multimodal enough. Evaluating MCNN on the Franka Kitchen environment used in BeT and Diffusion Policy would help clarify this.", "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": "", "extra_scores": {}}, {"review_id": "2uRCW9vEK8", "reviewer": "Reviewer_L3an", "summary": "The paper proposes a mechanism to cope with the loss of alignment between training and online evaluation in behavior cloning. It first collects state-action pairs from training datasets and then combines the nearest-neighbor state outputs with regular neural network predicted output, such that the collected output can be chosen for observed states identical to existing ones in the training dataset and the predicted neural network output can be chosen for novel states. The results are based off of 9 tasks across 5 environments and compared to 7 baselines.", "questions": "- In Fig. 1, what are the absolute returns for ~10, ~100, ~1,000 and ~10,000 trajectories?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- I appreciate the effort to formalize most of the definitions included in the paper.\n- The experiments contain a good amount of baselines and model variants.", "weaknesses": "- The proposed method seems equivalent to supervised training of a NN with the codebook entries oversampled. (Edit: this has been clarified by the authors in an additional experiment)\n- It is unclear to me why the clipping function in Fig. 12 should be referred to as \"tanh-like\".\n- The proposed method seems to quickly drop the performance gains as the number of training data increases, which might render it irrelevant for tasks requiring a regular amount of training examples. (Edit: performance in low-data regime seems actually competitive upon revisiting the results)", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)