Memory-Consistent Neural Networks for Imitation Learning
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
Kaustubh Sridhar, Souradeep Dutta, Dinesh Jayaraman, James Weimer, Insup Lee
# Memory-Consistent Neural Networks for 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 9 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](https://sites.google.com/view/mcnn-imitation)

## 1 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-intensive than model-based control.

We consider the simplest and perhaps most widely used imitation learning algorithm, behavior cloning (BC) [26], which reduces policy synthesis to supervised learning over the expert demonstration data. For example, a 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 on 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 of 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 behavior on this self-induced set of states, which can deviate from the training dataset of expert demonstrations. In particular, a minor error in the policy's action output at any time may induce a future input state that is subtly different from expert states. If the policy behaves erratically under such small deviations, as it often does in practice, the situation quickly snowballs into a vicious cycle of compounding errors leading to task failure.

Past solutions to this compounding error problem have focused on modifying the behavior cloning setup, such as by permitting online experience [12, 31], reward labels
[25], queryable experts [34], or modifying the demonstration data collection procedure [18]. Instead, we retain the conveniences of the plain BC setup and focus on designing 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 number such as in robotics where human demonstrations are essential for imitation learning. To tame these errors, we propose semi-parametric "memory-consistent neural networks" (MCNN). MCNNs first subsample the dataset into representative prototype "memories" to form the scaffold for the eventual function. They then fit a parametric function to the rest of the training data that is hard-constrained by the very formulation of the model class, to exactly fit the training data at all the memories, and further, to stay within double-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 a deep neural network (subject to the double-cone constraints) for from them. All functions in this MCNN model class lie within "permissible regions" centered on each memory, meaning that function values away from the training points are bounded. Under 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.

Using MCNNs on 9 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 1 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 C & summarize closely related work here.

**Compounding errors in imitation learning** have previously been tackled by permitting online experience [12, 31], reward labels [25], queryable experts [34], or modifying the demonstration data collection procedure [18]. Our work is orthogonal to these methods and creates a 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) [35], Action Chunking Transformer [45], and Diffusion Policies [41, 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 a 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-parametric methods in imitation learning** such as nearest neighbors [36], RBFs [32], and SVMs [19] have historically shown competitive performance on various robotic control benchmarks. But, only recently, a semi-parametric approach consisting of neural networks for representation learning and k-nearest neighbors for control was proposed in Visual Imitation through Nearest Neighbors (VINN) [24]. This is the closest paper to our work and in Section 5, we compare with VINN and demonstrate that we outperform their method comprehensively.

Figure 1: **MCNN significantly improves performance on realistic demonstration datasets. We plot the percentage increase in return with MCNN over D4RL BC [9] for various number of demonstrations across many tasks. In this plot, each point is a separate MCNN policy. We see significant improvements in the 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 E.**
**Theoretical guarantees on 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 on reductions for imitation learning in [33, 28, 2, 29] and leverages intuitions from [23] on bounding the width of the model class.

## 3 Problem Formulation

A _Markov Decision Process (MDP)_ is a tuple \(\mathcal{E}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},\gamma,\mathcal{I})\), where \(\mathcal{S}\subseteq\mathbb{R}^{n}\) is the set of states, \(\mathcal{A}\) is the set of actions, \(\mathcal{P}(s^{\prime}|s,a)\) is the probability of transitioning from state \(s\) to \(s^{\prime}\) when taking action \(a\), \(\mathcal{R}(s,a)\) is the reward accreted in state \(s\) upon taking action \(a\), \(\gamma\in[0,1)\) is the discount factor, and \(\mathcal{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 \(\mathcal{S}\) to be closed and compact. Given a policy \(\pi:\mathcal{S}\mapsto\mathcal{A}\), the expected cumulative reward accrued over the duration of an episode is given by the following,

\[J(\pi)=\mathbb{E}_{\pi}\big{[}\sum_{t=1}^{H}\mathcal{R}(s_{t},a_{t})\big{]}. \tag{1}\]

In imitation learning, we assume that there exists an expert policy \(\pi^{*}\) unknown to the learner. This policy induces a distribution \(d_{\pi^{*}}\) on the state-action space \(\mathcal{S}\times\mathcal{A}\) obtained by rollouts on the MDP. The learner agent has access to an expert trajectory dataset \(D=\big{\{}(s_{0},a_{0}),(s_{1},a_{1}),\ldots,(s_{N},a_{N})\big{\}}\) drawn from distribution \(d_{\pi^{*}}\). The goal of imitation learning is to estimate a policy \(\hat{\pi}\), which mimics the expert's policy and reduces the _sub-optimality_ gap: \(J(\pi^{*})-J(\hat{\pi})\).

## 4 Approach

Our approach involves developing a new model class, memory consistent neural networks (MCNN), and training it 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.

### The Model Class: Memory-Consistent Neural Networks

First, we develop the semi-parametric MCNN model class for imitation learning. MCNNs rely on a code-book set of "memories" \(\mathcal{B}:=\{(s_{i},a_{i})\}_{i=1}^{K}\) which are subsampled from the expert training dataset and summarize it. In practice, such a memory code-book can be created using one of various off-the-shelf approaches. We describe our algorithmic choices later in Sec 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 \(\mathcal{B}\), we now define a "nearest memory neighbor policy". For a finite set \(S\subset\mathcal{S}\), and an input \(x\in\mathcal{S}\), we first define its closest element in \(S\) as, \(C_{S}(x)=\operatorname*{arg\,min}_{s\in S}\,d(s,x)\), where \(d\) is some distance metric defined on the space \(\mathcal{S}\). We denote by \(\mathcal{B}|_{S}\), and \(\mathcal{B}|_{A}\) as the set of all states and actions captured by the memory code-book \(\mathcal{B}\). With slight abuse of notation, we denote \(\mathcal{B}(s)\) as the action assigned by the codebook for a state input \(s\). Using the above, we now define a nearest neighbor regression function \(f^{NN}\) as the following,

**Definition 4.1** (Nearest Memory Neighbor Function).: For an input \(x\in\mathcal{S}\), assume that \(s^{\prime}=C_{\mathcal{B}|_{S}}(x)\), then \(f^{NN}(x):=\mathcal{B}(s^{\prime})\).

In other words, the nearest memory neighbor function assigns actions according to a nearest neighbor look-up in the memory code-book \(\mathcal{B}\). We are now ready to define _memory-consistent neural networks_ (MCNN), which permit interpolating between nearest neighbor functions and parametric deep neural network (DNN)-based functions. Let \(f^{\theta}\) denote a DNN function parameterized by \(\theta\), which maps from MDP states to actions.

**Definition 4.2** (Memory-Consistent Neural Network).: A memory-consistent neural network \(f^{MC}\) is defined using the codebook and DNN function pair (\(\mathcal{B},f^{\theta}\)), and hyperparameters \(\lambda\in\mathbb{R}^{+}\), \(L\in\mathbb{R}\) as

\[f^{MC}_{\theta,\mathcal{B}}(x)=\underbrace{f^{NN}(x)\left(e^{-\lambda\ d(x,s^ {\prime})}\right)}_{\text{Nearest Memory Neighbour Function}}+\underbrace{L\left(1-e^{-\lambda\ d(x,s^{\prime})}\right)\,\sigma(f^{ \theta}(x))}_{\text{Constrained Neural Network Function Class}} \tag{2}\]
where, \(s^{\prime}=C_{\mathcal{B}|_{S}}(x)\) is the nearest memory to \(x\), and \(\sigma:\mathbb{R}\mapsto[-1,1]\) is a compressive non-linearity that imposes hard limits on the outputs of \(f^{\theta}\). In practice, we use \(\tanh\) or similar functions.

We refer the reader to Figure 2 to drive the intuition. For inputs that are close to the points in the codebook \(\mathcal{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 a simple mixture: \(\alpha f^{NN}(x)+(1-\alpha)\;\;L\;\sigma(f^{\theta}(x))\;\), where the mixing factor \(\alpha\in[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 \(f^{NN}\) is controlled by the parameter \(\lambda\). Thus, we obtain a purely nearest neighbor function for \(\lambda=0\), and a vanilla deep neural network function for \(\lambda=\infty\). Note that the MCNN function values in regions far away from memories are in the set \([-L,L]\). For this reason, we normalize output actions to \([-L,L]\) before training an MCNN function.

### Theoretical Analysis of MCNNs for Imitation Learning

For fixed hyperparameters \(L\), \(\lambda\) and memory codebook \(\mathcal{B}\), we denote by \(\mathfrak{F}\), the class of memory-consistent functions outlined in Equation 2. Note that a choice of the DNN function parameters \(\theta\) fixes a specific function in this class as well.

**Assumption 4.3** (Realizability).: We assume that the expert policy \(\pi^{*}\) belongs to the function class \(\mathfrak{F}\).

This assumption trivially holds at the memories, where the MCNN exactly reproduces expert actions. For all other points, we assume that there exist some parameters \(\theta\), which can capture the policy \(\pi^{*}\) with sufficient accuracy, for a choice of \(L\) and \(\lambda\). For a point \(x\) at a distance of \(d(s^{\prime},x)\), the vanilla DNN can affect the predictions only by an amount of \(\pm L\left(1-e^{-\lambda\;d(x,s^{\prime})}\right)\). Without this restriction, we 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 a way to enforce this bound using a zeroth order nearest neighbor estimate.

We analyze the behavior of this function class, and present some useful lemmas along the way. For a set of memories present in \(\mathcal{B}|_{S}\), we wish to capture the maximum value that the distance term: \(d(x,s^{\prime})\) can take in Equation 2. To that end, we define the _most isolated state_ as the following:

**Definition 4.4** (Most Isolated State).: For a given set of memory points \(\mathcal{B}|_{S}\), we define the most isolated state \(s^{I}_{\mathcal{B}|_{S}}:=\underset{s\in I}{\operatorname*{arg\,max}}\big{(} \underset{m\in\mathcal{B}|_{S}}{\operatorname*{min}}d(s,m)\big{)}\), and consequently the distance of the most isolated point as \(d^{I}_{\mathcal{B}|_{S}}=\underset{m\in\mathcal{B}|_{S}}{\operatorname*{min}} d(s^{I}_{\mathcal{B}|_{S}},m)\)

Figure 2: **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 the 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 \(\lambda\) (bottom left) or by decreasing the number of memories (bottom right). The second row shows many such MCNN model families with increasing capacity. For additional plots, see Appendix A.**
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 \(\mathcal{B}\).

**Lemma 4.5**.: _Assume two sets of memory code-books \(\mathcal{B}_{i},\ \mathcal{B}_{j}\), such that \(\mathcal{B}_{i}\subseteq\mathcal{B}_{j}\), then \(d^{I}_{\mathcal{B}_{i}|_{S}}\geq d^{I}_{\mathcal{B}_{j}|_{S}}\)_

Proof.: The proof of the above lemma is straightforward, since the infimum of a subset (\(\mathcal{B}_{i}\)) is larger than the infimum of the original set (\(\mathcal{B}_{j}\)). 

This observation is useful when we study the effects of increasing the size of the code-book \(\mathcal{B}\). Note, when learning a memory-consistent neural network \(f^{MC}_{\theta,\mathcal{B}}\), we deploy the standard SGD based training to adjust the parameters \(\theta\). The choice of the number of memories in \(\mathcal{B}\) is kept as a hyperparameter. This allows us to bound the maximum width of the function class, first described in [23]. We analyze this for single output functions next.

**Lemma 4.6**.: _(Width of Function Class) The width of the function class \(\mathfrak{F}\), \(\forall\ \theta_{1},\theta_{2}\in\Theta,\ \text{and}\ \forall s\in \mathcal{S}\), defined as \(\underset{\theta_{i},\theta_{j}}{max}\left|\left(f^{MC}_{\theta_{i},\mathcal{ B}}-f^{MC}_{\theta_{j},\mathcal{B}}\right)\left(s\right)\right|\) is upper bounded by \(:2L\times\left(1-e^{-\lambda\ d^{I}_{\mathcal{B}_{i}|_{S}}}\right)\)_

Proof.: Please see Appendix B.

**Theorem 4.7**.: _The sub-optimality gap \(J(\pi^{*})-J(\hat{\pi})\leq min\{H,H^{2}|\mathcal{A}|L\big{(}1-e^{-\lambda\ d^{I}_{ \mathcal{B}_{i}|_{S}}}\big{)}\}\)_

Proof.: In order to take advantage of well-known results in the imitation learning literature [33, 28, 2, 29], we restrict ourselves for the purpose of this analysis to the discrete action-space \(\mathcal{A}\) scenario, where the policy \(\pi:\mathcal{S}\mapsto\Delta(\mathcal{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 a dirac distribution over the actions corresponding to each input state.

Recall that in imitation learning, if the population total variation (TV) risk \(\mathbb{T}(\hat{\pi},\pi^{*})\leq\epsilon\), then, \(J(\pi^{*})-J(\hat{\pi})\leq min\{H,H^{2}\epsilon\}\) (See [28] Lemma 4.3). We note the following for population TV risk:

\[\mathbb{T}(\hat{\pi},\pi^{*}) =\frac{1}{H}\sum_{t=1}^{H}\mathbb{E}_{s_{t}\sim f^{t}_{\pi^{*}}} \left[TV(\hat{\pi}(\cdot|s_{t}),\pi^{*}(\cdot|s_{t}))\right]\leq\frac{1}{H} \sum_{t=1}^{H}\mathbb{E}_{s_{t}\sim f^{t}_{\pi^{*}}}\bigg{[}|\mathcal{A}|L \big{(}1-e^{-\lambda\ d^{I}_{\mathcal{B}_{i}|_{S}}}\big{)}\bigg{]} \tag{3}\] \[\leq|\mathcal{A}|L\big{(}1-e^{-\lambda\ d^{I}_{\mathcal{B}_{i}|_{ S}}}\big{)}\]

where \(f^{t}_{\pi^{*}}\) is the empirical distribution induced on state \(s^{t}\) obtained by rolling out policy \(\pi^{*}\). 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(\pi^{*})-J(\hat{\pi})\leq min\{H,H^{2}|\mathcal{A}|L\left(1-e^{-\lambda\ d^{I }_{\mathcal{B}_{i}|_{S}}}\right)\}\]

**Corollary 4.8**.: _Using Lemma 4.5 we know that if \(\mathcal{B}_{i}\subseteq\mathcal{B}_{j}\), then \(\big{(}1-e^{-\lambda\ d^{I}_{\mathcal{B}_{i}|_{S}}}\big{)}\geq\big{(}1-e^{- \lambda\ d^{I}_{\mathcal{B}_{j}|_{S}}}\big{)}\). This can result in lower performance gap according to Theorem 4.7, when \(H\geq H^{2}|\mathcal{A}|L\big{(}1-e^{-\lambda\ d^{I}_{\mathcal{B}_{i}|_{S}}} \big{)}\). 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 a _bounded sub-optimality gap_ (_Theorem 4.7_) _also not available in vanilla neural networks._ Finally, our Corollary states that we can likely gain better imitation learning performance by simply adding more memories (up to a limit).

### 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 \(\mathcal{B}\) from the expert trajectory dataset \(D\). The goal of Algorithm 1 is to build the nearest memory neighbor function \(f^{NN}\). This is followed by details on the training aspects of the MCNN parameters from the imitation dataset \(\mathcal{D}_{e}\) 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.
```
1:Offline dataset \(\mathcal{D}=\{(s_{i},a_{i})\}_{i=1}^{N}\), number of memories \(m\)
2:A nearest neighbor based function \(f^{NN}:\mathcal{S}\mapsto\mathcal{A}\)
3:Nodes \(\mathcal{N}\), edges \(\mathcal{E}\leftarrow\) NeuralGasClustering(\(\mathcal{S}\), \(m\)) // learns the distribution induced by \(\mathcal{D}\)
4:Nodes \(\mathcal{N}^{\prime}\), \(\mathcal{D}(\mathcal{N}^{\prime})\leftarrow\) For each node in \(\mathcal{N}\), find the closest observation in \(\mathcal{D}\), and call this \(\mathcal{N}^{\prime}\). Additionally, return the corresponding action taken by the expert in \(\mathcal{D}\), denoted by the map \(\mathcal{D}(\mathcal{N}^{\prime})\)
5:\(\mathcal{G}\leftarrow\) Define neural-gas with nodes \(\mathcal{N}^{\prime}\), and edges \(\mathcal{E}\).
6:Define a memory code-book \(\mathcal{B}\) using \(\mathcal{B}|_{\mathcal{S}}=\mathcal{G}\) and \(\mathcal{B}|_{\mathcal{A}}=\mathcal{D}(\mathcal{N}^{\prime})\). Pairing nodes in the neural-gas to its corresponding actions.
7:Define a nearest neighbor function \(f^{NN}_{\mathcal{B}}\), along the lines described in Definition 4.1 using \(\mathcal{B}\).
8:return\(f^{NN}_{\mathcal{B}}\)
```

**Algorithm 1** Learning Memories

Definition 4.9 (Neural Gas).: A neural gas \(\mathcal{G}\coloneqq(\mathcal{N},\mathcal{E})\), is composed of the following components,

1. A set \(\mathcal{N}\subset\mathcal{S}\) of the nodes of a network. Each node \(m_{i}\in\mathcal{N}\) is called a memory in this paper.
2. A set \(\mathcal{E}\subset\{(m_{i},m_{j})\in\mathcal{N}^{2},i\neq 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, 27, 20] is primarily used for unsupervised learning tasks, particularly for data compression or vector quantization. The goal is to group similar data points together based on their similarities. The algorithm works by creating a 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 a _gas_ in order to capture the density. For more details we refer the reader to [8]. We use this in Algorithm 1 (Line 1) to get the initial clustering. We can now go ahead and outline how "memories" are picked in our case.

Learning memories.Algorithm 1 first uses the neural-gas algorithm to pick candidate points (nodes) \(\mathcal{N}\) in the state space. However, these points could be potentially absent in the dataset \(\mathcal{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 \(\mathcal{B}\) and defining a function as outlined in Definition 4.1.

```
1:Dataset \(\mathcal{D}=\{(s_{i},a_{i})\}_{i=1}^{N}\), nearest neighbor function \(f^{NN}_{\mathcal{B}}\), neural network function \(f^{\theta}(.)\), batch size, total training steps \(T\), parameters \(\lambda\) and \(L\). Output: Learned policy \(f^{MC}_{\mathcal{B},\mathcal{B}}\)
2:for\(\text{step}=1\) to \(T\)do
3: Sample batch \(\mathcal{B}\) from \(\mathcal{D}\).
4: Forward propagate \((s_{i},a_{i})\sim B\), \(f^{MC}_{\mathcal{B},\mathcal{B}}(x)=f^{NN}_{\mathcal{B}}(x)\left(e^{-\lambda \ d(x,s^{\prime})}\right)+L\left(1-e^{-\lambda\ d(x,s^{\prime})}\right)\ \sigma(f^{\theta}(x))\) where, \(s^{\prime}\) is the nearest neighbor of \(x\) in \(\mathcal{B}|_{\mathcal{S}}\), \(\sigma_{\mathcal{B}}(x)\) is a \(\tanh\)-like activation function given by \(\sigma_{\mathcal{B}}(x)=2\left[\text{LeakyReLU}_{\mathcal{B}}\left(\frac{x+1}{ 2}\right)-(1-\beta)\text{ReLU}\left(\frac{x-1}{2}\right)\right]-1\) and \(\beta=\max\left(0,1-\lfloor\frac{4\pi\epsilon\epsilon}{1000}\rfloor\right)\).
5: Update \(\theta\leftarrow\theta-\nabla\operatorname{\mathbb{E}}_{(s_{i},a_{i})\sim b} \mathcal{L}(f^{MC}_{\mathcal{B},\mathcal{B}}(s_{i}),a_{i})\) where \(\mathcal{L}\) is the negative log-likelihood or mean squared loss or other loss function.
6:endfor
```

**Algorithm 2** Behavior Cloning with Memory-Consistent Neural Networks: Training

Training MCNNs.Finally, we train MCNN policies through gradient descent on the parameters \(\theta\) of the neural network over the expert dataset \(D_{e}\). For the compressive non-linearity, we use the \(\sigma_{\beta}\) function given in Algorithm 2 which is similar to \(\tanh\). We describe this in detail in Appendix E.

## 5 Experimental Evaluation

We now perform a thorough experimental evaluation of MCNN-based behavior cloning in a large variety of imitation learning settings.
**Environments and Datasets:** We test our approach on 9 tasks, in 5 environments: 4 Adroit dextrous manipulation and 1 CARLA environment as pictured in Figure 3. Demonstration datasets are drawn from D4RL [9]. 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 400 demonstration trajectories (100K transitions) from a hand-coded expert. For observations, we use high-dimensional states in Adroit pen, hammer, relocate, and door, and \(48\times 48\) images in CARLA. Action spaces are 24-30 dimensional in the Adroit dexterous manipulation environments and 2-D in CARLA.

**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 E. 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 to 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 a search amongst a million datapoints online at every step (which is highly inefficient). **(3) Visual Imitation with (k) Nearest Neighbours (VINN)**[24]: VINN is a recent method that performs a 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 to create memories (we discuss this embedding space more below). **(4) CQL-Sparse (CQL-S)**: We learn a policy using the CQL offline RL algorithm [17] and a 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 (BeI-BC)**[35]: We train and evaluate a behavior transformer using the official implementation on all tasks. **(7) Diffusion BC (Diff-BC)**[41]: We also train and evaluate a diffusion-based BC policy using the implementation in [41]. We provide additional details for all baselines and comprehensive hyperparameter sweeps in Appendix E.

**Learning memories and MCNN:** We learn neural gas memories with 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 (see more in Algorithm 1). We train the MCNN on this dataset following Algorithm 2 for \(1\) million steps and evaluate on 20 trajectories after training and repeat each experiment for a minimum of 3 seeds. We report results with an MLP, a behavior transformer (BeT), and a diffusion policy as the underlying neural network under the names **'MCNN+MLP'**, **'MCNN+BeT'**, and **'MCNN+Diff'** respectively. We expand on the experimental setup and all hyperparameters in Appendix E.

**Embedding CARLA images:** In the CARLA tasks, we use an off-the-shelf ResNet34 encoder [2] that has been shown to be robust to background and environment changes in CARLA to convert the \(48\times 48\) images to embeddings of size \(512\). We use this embedding space as the observation space for learning memories and policies.

**Performance Metrics:** 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

Figure 3: The D4RL [9] environments where we evaluate our trained policies: Adroit Pen, Hammer, Relocate, and Door [30], and CARLA’s Town03 and Town04 [3]. The four Adroit environments have proprioceptive observations and the CARLA environment has image observations.
aggregate plots on a set of four tasks, we compute the percentage increase in return of a method over D4RL BC in all four tasks and report the median.

Results:First, for the "human" tasks with the most realistic imitation learning setup, we plot aggregate and taskwise results in Figure 4. In aggregate-human, 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. In pen-human-v1, MCNN+MLP outperforms the nearest baseline by 33%. It is also the only method to shoot past the expert ceiling of 100 (depicted by a dashed red line). In hammer-human-v1, MCNN+MLP 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-v1, MCNN+BeT is the only method to achieve a positive return. We attribute, like previous work [22], the stronger performance of MCNN+BeT over other MCNN variants in the relocate task to 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-v1, 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.

On the expert tasks in Figure 5, even with a 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.

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. Additional figures and all means and standard deviations of our results can be found in Appendix F.

Figure 4: **Adroit human tasks [25 demos]:** Comparison of returns (across 20 evaluation trajectories and 3 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]:** Comparison of returns (across 20 evaluation trajectories and 3 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.
**Discussion:** We identify some high-level trends across our results here. For every architecture - MLP, BeT, or Diffusion, plugging in MCNN significantly improves performance in every task. The best-performing method in nearly every task is an 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-v1 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.

**Ablations:** 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 1-NN can be found in Appendix E.

In Figure 7, we show that given limited online interaction (here, 20 episodes), simply selecting the best-performing hyperparameters (\(\lambda\), \(L\), and a number of memories) further improves MCNN+MLP performance. We also demonstrate the significance of neural gas-based memories by comparing MCNN+MLP with a version that uses randomly chosen memories in Figure 13 in Appendix F. We observe significantly reduced performance with randomly chosen memories. We attribute this to the competitive Hebbian learning algorithm for the neural gas, which is better at capturing the distribution of training points (and creating memories that are "spread out"). This in turn reduces the distance to the most isolated state, improving imitation performance. Additional results on ablating the values of \(\lambda\) and on all other tasks can be found in Appendix F.

## 6 Conclusions and Limitations.

Imitation learning, and in particular behavior cloning, is 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-parametric 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-designed 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 as in reinforcement learning and meta-learning.

Figure 8: **CARLA [400 demos]**: Comparison of return (across 20 evaluation trajectories and 3 random seeds) between baselines and our methods (MCNN+BeT and MCNN+MLP). For the results shown here, our MCNN methods use the same fixed set of hyperparameters across all tasks.

Figure 6: Number of Memories in MCNN+MLP: Comparison of returns (across 20 evaluation trajectories) with our method (MCNN+MLP) using between 2.5% to 80% of the dataset as memories. We find a "sweet spot" for number of memories at 10-20%. We also see the expected decrease to 1-NN performance as number of memories increases to 100%.

Figure 7: Fixed vs Tuned Hyperparameters in MCNN+MLP: Comparison of returns (across 20 eval trajectories and 3 seeds) between our method with fixed hyperparameters across all tasks and with hyperparameters tuned online. This shows that limited online interaction (20 ep) for finetuning hyperparameters can further improve our method.