# Distilling Morphology-Conditioned Hypernetworks
for Efficient Universal Morphology Control

Zheng Xiong  Risto Vuorio  Jacob Beck  Matthieu Zimmer  Kun Shao  Shimon
Whiteson

###### Abstract

Learning a universal policy across different robot morphologies can
significantly improve learning efficiency and enable zero-shot generalization
to unseen morphologies. However, learning a highly performant universal policy
requires sophisticated architectures like transformers (TF) that have larger
memory and computational cost than simpler multi-layer perceptrons (MLP). To
achieve both good performance like TF and high efficiency like MLP at
inference time, we propose HyperDistill, which consists of: (1) A morphology-
conditioned hypernetwork (HN) that generates robot-wise MLP policies, and (2)
A policy distillation approach that is essential for successful training. We
show that on UNIMAL, a benchmark with hundreds of diverse morphologies,
HyperDistill performs as well as a universal TF teacher policy on both
training and unseen test robots, but reduces model size by 6-14 times, and
computational cost by 67-160 times in different environments. Our analysis
attributes the efficiency advantage of HyperDistill at inference time to
knowledge decoupling, i.e., the ability to decouple inter-task and intra-task
knowledge, a general principle that could also be applied to improve inference
efficiency in other domains.

Machine Learning, ICML

##  1 Introduction

Reinforcement learning (RL) for robotic control has made great progress in
recent years (Levine et al., 2016; Kalashnikov et al., 2018; Andrychowicz et
al., 2020; Brohan et al., 2022). However, generalization across robots remains
a key challenge, as the policy trained for one robot transfers poorly to
another robot with a different morphology, i.e., the topology graph of a robot
and the hardware parameters of each body part. Furthermore, it is too sample
inefficient to train a separate policy for each new robot from scratch.

To tackle this challenge, universal morphology control aims to learn a
universal policy to control different robots. By training on a set of diverse
morphologies, this constitutes a form of multi-task RL (Vithayathil Varghese &
Mahmoud, 2020) where controlling each robot is a separate task. The learned
policy is expected to not only improve learning efficiency on the training
robots, but also enable zero-shot generalization to test robots with unseen
morphologies.

Training a multi-layer perceptron (MLP) policy is usually sufficient to
achieve good performance on a single robot, but often generalizes poorly to
other robots (Wang et al., 2018). While it is feasible to train a multi-robot
MLP policy for universal morphology control, it is significantly outperformed
by graph neural networks (GNN) (Wang et al., 2018; Pathak et al., 2019; Huang
et al., 2020) and transformers (TF) (Kurin et al., 2020; Gupta et al., 2022)
w.r.t. both training and generalization performance. Moreover, TF outperforms
GNN by better modeling interactions between distant nodes in the morphology
graph (Kurin et al., 2020), and has thus been adopted by most recent
approaches (Dong et al., 2022; Furuta et al., 2022; Gupta et al., 2022; Hong
et al., 2022; Trabucco et al., 2022; Chen et al., 2023; Xiong et al., 2023).

However, TF has significantly higher memory, computation, and energy costs
than MLP, all of which are key considerations when deploying the policy on
real-world robots with constrained hardware (Hutter et al., 2016; Zhao et al.,
2021; Brohan et al., 2022; Leal et al., 2023). For example, a state-of-the-art
TF-based method for universal morphology control (Xiong et al., 2023) requires
about 40M FLOPs for a single step of inference on a robot with just 10 limbs,
more than 100 times that of a single-robot MLP policy with similar
performance, and this efficiency gap increases proportionally with the number
of limbs. So a natural question arises: _Can we learn a universal policy with
the performance of TF but the inference efficiency of MLP?_

In this paper, we answer this question affirmatively. To get the best of both
worlds, we introduce a hypernetwork (HN) (Ha et al., 2016), i.e., a network
that takes context features as input to generate the parameters of a base
network. Our key intuition is that, compared to TF, HN can better decouple
inter-task knowledge and intra-task knowledge via its hybrid architecture for
more efficient inference, which we call the _knowledge decoupling hypothesis_.
Under our problem setting, the context-conditioned HN can provide sufficient
model capacity to accommodate inter-robot knowledge, while the generated base
network only needs to encode task-specific knowledge about the robot it
controls. By contrast, TF encodes both kinds of knowledge with a single large
model, which introduces high model redundancy when deployed on a specific
robot.

Specifically, suppose we train a set of single-robot MLP policies
$\left\\{\pi_{k}\right\\}$, each with good performance on a different robot
$k$. Motivated by the intuition that the optimal control policy of a robot
critically depends on its morphology (Gupta et al., 2022; Xiong et al., 2023),
we train an HN that takes the morphology context $c_{k}$ of robot $k$ as input
to predict the corresponding MLP policy parameters $\pi_{k}$. For each robot,
as the morphology context $c_{k}$ is constant, the HN-generated parameters are
also fixed. Consequently, we only need to call it once before deployment to
generate the base MLP, while the HN itself is not needed for policy execution.
This yields a universal policy that works like an MLP at inference time but
still has the potential to achieve good performance on both training and
unseen test robots.

While training such an HN policy via RL is a straightforward option,
empirically we find that it is unstable and significantly underperforms a
universal TF policy. Instead, we adopt a policy distillation (PD) approach by
distilling a universal TF teacher policy into an HN student policy via
behavior cloning (BC). While it is not hard to distill a student policy to
match the teacher’s performance on the training task(s) (Parisotto et al.,
2015; Rusu et al., 2015; Czarnecki et al., 2019), a key challenge in our
problem setting is to maintain the teacher policy’s zero-shot generalization
performance after distillation. In this paper, we identify several critical
algorithmic choices in PD that influence the generalization performance of the
student policy to unseen tasks: (1) The choice of the teacher(s), (2)
Architecture alignment between the teacher and the student, (3) The number of
tasks on which to collect training data for PD, and (4) Regularization in task
space. We believe that these algorithmic choices could serve as general
guidelines for improving task generalization of PD in other domains as well.

We name our approach as _HyperDistill_ to highlight its two key components:
(1) The HN architecture to achieve both good performance and high inference
efficiency via knowledge decoupling, and (2) Training via policy distillation
to successfully learn such an HN policy. We experiment on a challenging
universal morphology control benchmark called UNIMAL (Gupta et al., 2021),
which includes hundreds of diverse morphologies to evaluate both multi-robot
training and zero-shot generalization. HyperDistill achieves performance
similar to the universal TF teacher on both training and unseen test robots,
while significantly reducing the model size by 6-14 times, and computational
cost by 67-160 times in different environments at inference time. The
experimental results further support our knowledge decoupling hypothesis,
which could serve as a general principle to improve inference efficiency in
other domains.

##  2 Background

###  2.1 Problem Formulation

Learning a universal policy to control different morphologies can be seen as
solving a contextual Markov Decision Process (CMDP) (Hallak et al., 2015),
where the task context space $\mathcal{C}$ is defined over all possible
morphology configurations. For robot (task) $k$, we use $\mathcal{S}_{k}$,
$\mathcal{A}_{k}$, $T_{k}$ and $R_{k}$ to represent its state space, action
space, transition function and reward function respectively.

We assume that all the robots are drawn from a modular design space, i.e.,
each robot can be seen as a morphology tree over a set of basic nodes (limbs),
and the node-level state and action space are homogeneous across different
robots’ limbs. Based on this assumption, we have
$\mathcal{S}_{k}=\mathcal{S}_{k}^{1}\times\cdots\times\mathcal{S}_{k}^{N_{k}}$
and
$\mathcal{A}_{k}=\mathcal{A}_{k}^{1}\times\cdots\times\mathcal{A}_{k}^{N_{k}}$,
where $N_{k}$ is the number of limbs in robot $k$. The task context $c_{k}$
includes morphology information about the robot, consisting of node-wise
context features $\\{c_{k}^{i}|i=1,\dots,N_{k}\\}$ (see Appendix A for more
details), and a topology tree of the robot.

We use $s_{k,t},a_{k,t},r_{k,t}$ to represent the state, action and reward at
time step $t$ for robot $k$. The training objective is to learn a universal
policy $\pi_{\theta}(a_{k,t}|s_{k,t},c_{k})$ to maximize the average return
over a set of $K$ training robots, i.e.,
$\max_{\theta}\left[\frac{1}{K}\sum_{k=1}^{K}\sum_{t=0}^{H}r_{k,t}\right]$,
where $H$ is the task horizon for all different robots. In addition to good
training performance, we also expect the learned policy to generalize well on
unseen test morphologies in a zero-shot manner.

###  2.2 Architectures for Universal Morphology Control

To learn an MLP policy for the state and action space defined over a
morphology tree, we first need some way to order the limbs, so that we can
concatenate their node-wise states and actions into single vectors as the
MLP’s input and output. As the morphology of a robot has a tree structure, the
limbs in each robot are usually ordered by some tree traversal methods like
depth-first search (Gupta et al., 2022).

When learning a universal MLP policy, to handle the variable number of limbs
across different robots, it is common to assume a maximal limb number
$N_{\text{max}}$, and zero-pad the state and action vectors to this maximal
length so that the state and action dimensions are aligned across different
robots. As the universal policy also conditions on the morphology context, the
context features $c_{k}^{i}$ are usually concatenated with the state features
$s_{k,t}^{i}$ as node-wise input (Gupta et al., 2022). So the final input to
the MLP policy is
$x=\left[x_{1},x_{2},\dots,x_{N_{k}},\vec{0},\dots,\vec{0}\right]$, where
$x_{i}=[s_{i},c_{i}]$, and we omit the robot index $k$ and time index $t$
thereafter for simplicity of notation. For the input layer of the MLP, we have
$h^{(0)}=\sigma\left(W^{\text{in}}x+b^{\text{in}}\right)=\sigma\left(\sum_{i=1}^{N_{k}}W^{\text{in}}_{i}x_{i}+b^{\text{in}}\right)$,
where $\sigma$ is the activation function, $W_{i}^{\text{in}}$ is the subset
of the weight matrix that encodes the input features of node $i$. For the
$l$-th hidden layer, we have
$h^{(l+1)}=\sigma\left(W^{(l)}h^{(l)}+b^{(l)}\right)$. For the output layer,
we have $a=W^{\text{out}}h^{(L)}+b^{\text{out}}$, which can be decomposed into
node-wise action $a_{i}=W_{i}^{\text{out}}h^{(L)}+b_{i}^{\text{out}}$.

Unlike MLP, GNN (Wu et al., 2020) and TF (Vaswani et al., 2017) can naturally
handle variable numbers of homogeneous elements. GNN can be directly applied
to different morphology graphs, while TF treats each limb as a token. Unlike
in many domains where tokens are processed sequentially, in universal
morphology control, TF processes different nodes in parallel to model their
spatial interactions.

###  2.3 Hypernetworks

A hypernetwork (Ha et al., 2016) is a network that generates the parameters of
a base network conditioned on some context $c$. We can decompose an HN into a
context encoder $f$ and several output heads (parameterized as linear layers)
that take the context embedding $e=f(c)$ as input, and output parameters in
the base network. For example, to generate the parameters of a linear layer
$y=Wx+b$, the HN needs two output heads to generate $W$ and $b$ respectively,
i.e., $W=\text{HN}_{W}(e)$, $b=\text{HN}_{b}(e)$. If $W$ is a $M\times N$
matrix, and the context embedding dimension is $E$, then $\text{HN}_{W}$ is a
linear layer with input dimension $E$ and output dimension $M\times N$, while
$\text{HN}_{b}$ is a linear layer with input dimension $E$ and output
dimension $N$.

##  3 HyperDistill

This section introduces HyperDistill, which achieves both good performance and
high efficiency at inference time. In Section 3.1, we analyze the limitations
of existing methods for universal morphology control, which motivates us to
introduce HN as a solution in Section 3.2. In Section 3.3, we discuss how and
why we train the HN via policy distillation, and analyze some of its critical
algorithmic features.

###  3.1 Limitations of Existing Architectures

We analyze the limitations of two representative architectures for universal
morphology control: TF, which achieves SOTA performance but is expensive to
run at inference time, and MLP, which runs efficiently and achieves good
performance on each single robot but performs poorly in the multi-robot
setting. We first highlight a knowledge decoupling issue that exists in both
TF and MLP, then discuss the lack of order-invariance in MLP.

#### Knowledge Decoupling

Intuitively, knowledge learned by a universal policy can be decomposed into
two parts: inter-task knowledge for generalization across robots, and intra-
task knowledge about how to control a specific robot. Since both TF and MLP
use a single set of parameters to encode both kinds of knowledge, they cannot
decouple them from each other. Consequently, at inference time, there may be a
lot of redundant knowledge in the policy that is irrelevant to solving the
current task, which harms efficiency. It also explains why MLP can perform
well on a single robot but not under the multi-robot setting, as the model
capacity of a compact MLP is sufficient to accommodate task-specific knowledge
of a single robot, but not enough to encode both inter-task and intra-task
knowledge under a multi-robot setting. The same conclusion also holds for TF
when compressing a large TF into a smaller one, as we confirm experimentally
in Section 4. In summary, due to lack of knowledge decoupling, TF and MLP have
to trade off between performance and efficiency at inference time.

#### Order-Invariance Issue with MLP

A further issue with MLP is that it is not invariant to the order of limbs
across robots. TF is order-invariant, as all the limbs across different robots
share the same set of parameters. By contrast, in an MLP, only the limbs with
the same index across different robots share the same subset of parameters in
the input and output layers, so how we order the limbs influences the policy
output. As we do not have a consistent way to order the limbs across different
morphologies, MLP can overfit to spurious patterns in the manually chosen limb
indexing method, harming generalization.

###  3.2 HyperDistill Architecture

To tackle the limitations of TF and MLP, we introduce an HN that takes
morphology context as input to generate an MLP base policy for each robot.
First, it supports knowledge decoupling, as the morphology-conditioned HN
encodes inter-task knowledge, while the base MLP encodes intra-task knowledge.
The expensive HN is called only once to generate a task-specific base policy
for each new task, while the much smaller base network is sufficient to
efficiently solve the task. Second, the base MLP generated by the HN is order-
invariant, as the parameters associated with each limb no longer depend on the
limb index, but only condition on the limb’s context representation.

Figure 1: The architecture of HyperDistill. Different colors highlight the
correspondence between the parameters in the base network and the context
embedding they condition on via HN. We only show one hidden layer in the base
MLP for ease of illustration. More hidden layers can be easily added in a
similar way.

The overall architecture is shown in Figure 1. Unlike prior work in HN, a
unique challenge in applying HN to universal morphology control is that the
input and output dimensions of the MLP base network vary across robots due to
the variable number of limbs. We tackle this challenge by generating limb-wise
parameters for the input and output layers, as the subset of parameters
corresponding to each limb still has fixed dimension based on the modular
space assumption in Section 2.1. Specifically, for the input layer, we have
$h^{(0)}=\sigma\left(\sum_{i}W^{\text{in}}_{i}x_{i}+b^{\text{in}}_{i}\right)$,
where
$W_{i}^{\text{in}}=\text{HN}_{W}^{\text{in}}\left(f\left(c_{i}\right)\right)$,
and $b^{\text{in}}_{i}=\text{HN}_{b}^{\text{in}}\left(f(c_{i})\right)$. As the
whole policy conditions on the morphology context through HN, we no longer
need to concatenate context features $c_{i}$ to the network input, so we have
$x_{i}=s_{i}$. Similarly, for the output layer, we have
$a_{i}=W_{i}^{\text{out}}h^{(L)}+b_{i}^{\text{out}}$, where
$W_{i}^{\text{out}}=\text{HN}_{W}^{\text{out}}\left(f(c_{i})\right)$, and
$b_{i}^{\text{out}}=\text{HN}_{b}^{\text{out}}\left(f(c_{i})\right)$. For the
hidden layers, we first aggregate the context embedding of different limbs
using the mean to get a context embedding for the whole morphology as
$e_{m}=\frac{1}{N}\sum_{i}e_{i}$, then pass $e_{m}$ through HN output heads to
generate the hidden layer parameters, i.e.,
$h^{(l+1)}=\sigma\left(W^{(l)}h^{(l)}+b^{(l)}\right)$, where
$W^{(l)}=\text{HN}_{W}^{(l)}(e_{m})$, and
$b^{(l)}=\text{HN}_{b}^{(l)}(e_{m})$.

#### Context Representation Learning

The quality of the HN-generated policy critically depends on the quality of
the learned context embedding, e.g., if $e_{i}$ of different limbs are too
similar, the policy will generate similar actions across different limbs,
which is unlikely to be a good policy.

The context representation should be both discriminative to encode the diverse
behaviors of different limbs, and generalizable to unseen robots. We adopt two
approaches to achieve this goal. First, we apply some simple transformations
to the context features $c_{i}$ to make them more discriminative across limbs
(see Appendix A). Second, we use a TF as the context encoder to enrich each
limb’s context representation by interacting with other limbs in the robot,
e.g., if two limbs in two different morphologies have the same hardware
configurations, then we cannot tell them apart based on their own context
features alone, but the TF context encoder can learn a distinguishable
representation by further encoding morphology information. Since the TF
context encoder is also a part of the HN, it is not needed at inference time,
unlike a TF policy that uses TF as the controller.

###  3.3 HyperDistill Training via Policy Distillation

In principle, we can train a universal HN policy from scratch via RL. However,
empirically we find that the training process is unstable and the learned
policy significantly underperforms a TF policy, possibly because both HN and
RL are known to be unstable during learning, and combining them together
further exacerbates the optimization challenges.

Instead, we adopt policy distillation (Parisotto et al., 2015; Rusu et al.,
2015) by first training a universal TF policy, then distilling it into an HN
policy via behavior cloning, which replaces RL training with more stable
supervised learning to alleviate the optimization challenge. Moreover, as BC
empirically requires much less time to train than RL, we can reuse the same
pre-trained policy for more efficient evaluation of different algorithmic
choices.

Given a set of training robots and a universal TF policy trained on them as
the teacher $\pi^{T}$, we first collect expert trajectories on each training
robot with $\pi^{T}$ to generate a training dataset $\mathcal{B}^{T}$ for PD.
Then we train an HN student policy $\pi$ by minimizing the KL-divergence
between the action distributions generated by the teacher and student on
transitions randomly sampled from the buffer:

$L_{\pi}=\mathbb{E}_{s,c\sim\mathcal{B}^{T}}\left[KL(\pi^{T}(a|s,c)||\pi(a|s,c)\right].$

While more sophisticated loss functions (Czarnecki et al., 2019) have been
proposed to facilitate PD by collecting online data with the student,
empirically we find that this simple BC loss is sufficient to learn a student
policy that matches the teacher’s performance on the training robots, without
having to collect further samples with the student.

However, we also want the student to zero-shot generalize as well as the
teacher on unseen test robots. While several prior works evaluate zero-shot
generalization of a student policy to unseen tasks in different problem
settings (Chen et al., 2022; Furuta et al., 2022; Wan et al., 2023), none of
them systematically investigate how different algorithmic choices in PD
influence the student’s generalization performance. In this paper, we
highlight four key factors that may influence the generalization gap between
teach and student in PD.

#### The Choice of PD Teacher(s)

Theoretically, we can either train a universal TF teacher on multiple robots,
or train multiple single-robot MLP policies on different morphologies as the
teachers (Furuta et al., 2022). While the average performance of the single-
robot MLPs is similar to or even better than that of a universal TF on the
training robots (Xiong et al., 2023), we hypothesize that the student policy
distilled from a universal teacher policy generalizes better to unseen robots.
Intuitively, as the single-robot teacher policies are trained via RL
separately, they can be dramatically different from each other, which may
exacerbate the discontinuity of the distilled policy in the parameter space
and harm generalization.

#### Student-Teacher Architecture Alignment

We hypothesize that the generalization gap between teacher and student is
smaller when their neural architectures are more aligned (Hao et al., 2023).
Intuitively, the inductive bias introduced by a more aligned architecture
helps the student extrapolate to unseen task context and state in a more
consistent way with the teacher, which helps reduce the generalization gap.

However, as we are trying to distill from a TF to an HN, there is an
unavoidable mismatch between the teacher and student architectures. The two
algorithmic choices discussed next may help compensate for the generalization
gap caused by this architecture misalignment.

#### Number of Tasks for Distillation Training

To better align the teacher and student policies’ behaviors on unseen tasks, a
natural idea is to train the student policy on the teacher’s demonstrations
collected from more tasks, so that the PD training data better covers the task
space. This is different from training the teacher policy on more tasks, which
is an effective but orthogonal way to improve task generalization. Our
approach does not require modifying the teacher’s training process. Instead,
it simply requires collecting data from more tasks with the pre-trained
universal teacher, which introduces little computational overhead. For
clarity, we use “training robots” to denote the robots on which we train the
teacher policy, and “PD robots” to denote the robots on which we collect
expert trajectories for PD. Advanced methods like curriculum learning (Dennis
et al., 2020; Narvekar et al., 2020; Wan et al., 2023) could be utilized to
get a more robust task distribution to further mitigate generalization gap,
which we leave for future work.

#### Regularization in Task Space

Regularization is a common approach to reduce overfitting and improve
generalization (Cobbe et al., 2019). We consider it to be especially important
for HyperDistill, as HN may be more prone to overfitting than other
architectures. In prior work, morphology context is usually just used as an
additional policy input (Gupta et al., 2022), while in HyperDistill, it is
used as HN input to generate the whole control policy. As we only have a few
hundred different robots for training, which is much less than the number of
parameters in the HN, there is a high chance of overfitting. To reduce
overfitting, we apply dropout to the context embedding $e_{i}$ and $e_{m}$.
This can be seen as regularization in task space, which encourages the HN to
learn an ensemble of different MLP policies that can all work well on the same
robot. It can also be seen as domain randomization (Tobin et al., 2017) by
making the generated policy more robust to changes in morphology context. We
leave it for future work to investigate other regularization methods like
weight decay.

##  4 Experiments

Figure 2: The performance of different methods on the test robots in each
environment.

Our experiments aim to answer the following questions:

(1) Can HyperDistill achieve good performance on both training and unseen test
robots? How does it compare to other methods w.r.t. performance and efficiency
at inference time? (Section 4.2)

(2) How do different algorithmic and architecture choices in HyperDistill
influence its training and generalization performance? (Section 4.3)

###  4.1 Experimental Setup

We experiment on the UNIMAL benchmark (Gupta et al., 2021) built upon the
Mujoco simulator (Todorov et al., 2012), which includes 100 training robots
and 100 test robots with diverse morphologies, while new morphologies can be
easily generated via mutation operations supported by the UNIMAL design space.
Following the setup of Gupta et al. (2022), we consider three different
environments with increasing difficulties: (1) Flat terrain (FT): maximize
locomotion distance on a flat floor; (2) Variable terrain (VT): maximize
locomotion distance on a variable terrain randomly reset for each episode. (3)
Obstacle: maximize locomotion distance on a flat terrain while avoiding
randomly positioned obstacles.

#### Teacher Policy and Data Collection

For each environment, we train a universal TF policy on the 100 training
robots as the teacher policy. We adopt ModuMorph (Xiong et al., 2023) as the
TF architecture, as it achieves SOTA performance on the UNIMAL benchmark. To
collect data for policy distillation, we first generate an augmented robot set
of 1,000 PD robots by mutating the 100 training robots. See Appendix B.1 for
how we do the mutation. Then we collect 8,000 transitions from each PD robot,
forming a multi-robot dataset with 8M transitions for policy distillation.

#### Baselines

We compare HyperDistill with the following architectures as the student
policy: (1) TF (oracle): A TF with the same architecture as the teacher. It
serves as a performance upper bound on how well the student can perform
without architecture misalignment. (2) TF (compressed): A TF with a similar
number of parameters as the base MLP in HyperDistill. Nonetheless, the
computational cost of TF (compressed) is still higher due to parameter
sharing. (3) MLP: A multi-robot MLP with the same architecture as the base MLP
in HyperDistill, which is used to validate the advantages of HyperDistill over
MLP as discussed in Section 3.2. We also compare with training an HN policy
via RL (HN-RL) to show the importance of distillation.

#### Distillation Setup

The distillation process runs for 150 epochs, with a mini batch size of 5120.
We use Adam with a learning rate of 0.0003, and clip the gradient norm to 0.5.
We run three random seeds for each method in each environment, and report the
average performance and standard deviation. For HyperDistill, we apply dropout
to context embedding with $p=0.1$. See Appendix B.3 for the size of each
student model in each environment. 111Code will be publicly available upon
acceptance.

###  4.2 Main Results

Figures 2 illustrates how different methods generalize to test robots during
the distillation process. Table 1 compares the model size and FLOPs of
different methods at inference time. Due to space limitations, see Figure 6 in
the appendix for different methods’ performance on the training robots.

Task | Method | Model size | FLOPs
---|---|---|---
|  | Abs. | Rel. | Abs. | Rel.
FT | TF (oracle) | 1.73 M | 14.0 | 39.86 M | 160.8
TF (compressed) | 0.15 M | 1.2 | 1.78 M | 7.2
MLP | 0.23 M | 1.9 | 0.46 M | 1.9
HyperDistill | 0.12 M | 1 | 0.25 M | 1
VT | TF (oracle) | 1.97 M | 6.6 | 40.33 M | 67.2
TF (compressed) | 0.31 M | 1.0 | 3.58 M | 6.0
MLP | 0.41 M | 1.4 | 0.82 M | 1.4
HyperDistill | 0.30 M | 1 | 0.60 M | 1
Obstacle | TF (oracle) | 2.02 M | 6.7 | 40.43 M | 67.4
TF (compressed) | 0.31 M | 1.0 | 3.58 M | 6.0
MLP | 0.41 M | 1.4 | 0.82 M | 1.4
HyperDistill | 0.30 M | 1 | 0.60 M | 1

Table 1: Model size and FLOPs of different methods _at inference time_. Abs.
denotes the absolute value, and Rel. denotes the relative value w.r.t.
HyperDistill. See Appendix B.2 for how we compute the FLOPs, and C.2 for more
analysis on the results in this table.

HyperDistill achieves performance similar to TF (oracle), matching the TF
teacher on both the training and test robots in all the three environments,
while reducing model size by 6-14 times, FLOPs by 67-160 times in different
environments. Furthermore, TF (compressed) cannot match the performance of
HyperDistill in all the three environments, proving that the TF architecture
needs to trade off between performance and efficiency, while HyperDistill can
achieve both at inference time. HyperDistill also significantly outperforms
MLP, which validates the importance of introducing the HN. Moreover, although
TF (compressed) and MLP have similar model size, TF (compressed) still
outperforms MLP, and their performance gap on the test robots is larger than
the gap on the training robots, which may be overfitting caused by the order-
invariance issue of MLP as discussed in Section 3.1.

HN-RL performs poorly on both training and test robots, which reflects the
optimization difficulty of combining HN and RL, and validates the importance
of training via PD.

###  4.3 Ablation Studies

In this section, we investigate how the algorithmic choices in PD influence
generalization performance of the student. Due to space limitations, see
Appendix C.3 for ablations on architecture choices in HyperDistill. To save
computation, we run PD for only 50 epochs as most methods can already converge
within this budget, and experiment only in FT unless otherwise stated.

#### The Choice of PD Teacher(s)

We train (1) A universal TF teacher on all the training robots; (2) A set of
single-robot MLP teachers on each training robot. Then we collect 80,000
transitions from each training robot with the teacher(s) to get
$\mathcal{B}^{T}$. We experiment with both HyperDistill and TF (oracle) to
ensure that the PD choice applies to different student architectures.

Figure 3: The student’s learning curves under different teacher choices. In
the figure legend, “X $\rightarrow$ Y” means that we distill from teacher X
into student Y.

Figure 3 shows how the student policies perform on the training and test
robots with different teacher choices. As expected, when using single-robot
MLP teachers, although the student policies perform well on the training
robots, they generalize much worse than the students distilled from a
universal TF teacher, which validates our hypothesis in Section 3.3.

Figure 4: Final generalization performance of HyperDistill and TF (oracle)
w.r.t. the number of PD robots in different environments.

#### Student-Teacher Architecture Alignment

We reexamine the results in Figure 3 from a different perspective. When using
a universal TF teacher, the TF student has a more aligned architecture and
achieves better generalization performance than HyperDistill. When using
single-robot MLP teachers, while both students generalize much worse,
HyperDistill has a more aligned architecture (as its base MLP has the same
architecture as the teachers), and generalizes better than the TF student.
Moreover, for the same teacher, both students achieve similar performance on
the training robots regardless of their architectures, indicating that the
generalization gap is not caused by the difference in model capacity of
different student models, but is more likely the consequence of architecture
misalignment.

#### Number of Distillation Training Tasks

To validate our hypothesis that collecting distillation data from more robots
can improve the student’s task generalization, we experiment with 100, 500,
and 1,000 PD robots. For a fair comparison, the number of transitions
collected from each robot decreases proportionally so that the total data size
remains unchanged. As shown in Figure 4, HyperDistill’s generalization
performance increases by 6%, 15% and 13% in the three environments as the
number of PD robots increases from 100 to 1,000. There is no significant
improvement in the TF student’s generalization performance due to a ceiling
effect and its better aligned architecture with the TF teacher.

#### Regularization in Task Space

Figure 5: Learning curves of HyperDistill in FT with different ways of
incorporating dropout regularization.

As shown in Figure 5, applying dropout to the context embedding improves
generalization performance by 8.5%, while applying dropout to the base MLP
does not provide significant improvement, which agrees with our intuition that
regularization may be more important in task space than in state space.

##  5 Discussion

We reexamine our knowledge decoupling hypothesis to see if it is supported by
experimental results. First, HyperDistill achieves similar performance as TF
but runs much more efficiently at inference time, which validates the
efficiency benefits provided by knowledge decoupling. Second, as expected, due
to the lack of knowledge decoupling, when we compress a universal TF policy,
the compressed TF does not have enough model capacity to accommodate both
inter-task and intra-task knowledge, as evidenced by its worse performance in
our experiments. The knowledge decoupling hypothesis further suggests that if
we compress a universal TF into a single-robot TF, it should outperform a
compressed universal TF on that specific robot, as we relieve it of the burden
of distilling inter-task knowledge. To validate this idea, we conduct a proof-
of-concept experiment in the Obstacle environment. We randomly sample 10
training robots and distill the universal TF teacher into a single-robot TF
with the same architecture as TF (compressed) for each robot. As expected, the
compressed single-robot TFs achieve an average return of 2345, outperforming
the compressed universal TF with an average return of 2115, but at the cost of
losing the generalization ability to other robots.

In summary, TF needs to trade off between performance and efficiency at
inference time, while HyperDistill can enjoy both thanks to knowledge
decoupling. It resembles related ideas explored in other fields, such as
mixtures-of-experts (Riquelme et al., 2021; Shen et al., 2023) that learn a
large model with high capacity while sparsely activating sub-modules for
different tasks to improve inference efficiency. There is also evidence from
neuroscience that although the human brain contains billions of neurons, it is
still power-efficient because it activates only a small fraction of neurons at
a time to solve a given task (Barth & Poulet, 2012). Consequently, we believe
that our knowledge decoupling hypothesis could serve as a general principle to
improve inference efficiency in other domains as well.

##  6 Related Work

#### Universal Morphology Control

While generalization across robots with variations only in kinematics or
dynamics parameters has been extensively studied in prior work (Tobin et al.,
2017; Peng et al., 2018; Clavera et al., 2019; Ghadirzadeh et al., 2021; Feng
et al., 2022), universal morphology control poses a more challenging task of
generalization across morphologies. While learning an MLP policy with zero-
padding is feasible to work across morphologies, it is significantly
outperformed by GNN (Wang et al., 2018; Pathak et al., 2019; Huang et al.,
2020) and TF (Kurin et al., 2020; Dong et al., 2022; Furuta et al., 2022;
Gupta et al., 2022; Hong et al., 2022; Trabucco et al., 2022; Chen et al.,
2023; Xiong et al., 2023). Furthermore, TF outperforms GNN by better modeling
the interactions between distant limbs in a robot (Kurin et al., 2020), thus
has been adopted by most recent works. However, these more sophisticated
models introduce higher computational cost during deployment, which we aim to
tackle by generating a compact MLP controller at inference time with
morphology-conditioned HN. Similar to our work, ModuMorph (Xiong et al., 2023)
also utilizes HN. However, the main body of ModuMorph is still a TF, while HN
is only used to generate node-wise embedding and output layers, thus cannot
reduce computational cost. By contrast, we use HN to generate an MLP
controller which is much more efficient at inference time.

#### Hypernetworks

A hypernetwork (Ha et al., 2016) takes context features as input to generate
the parameters of a base network. In the context of RL, HNs provide a powerful
way to model the complex dependency between task context and the optimal
control policy for the task (Galanti & Wolf, 2020; Sarafian et al., 2021), and
has been widely adopted in multi-task RL and meta-RL (Yu et al., 2019; Peng et
al., 2021; Beck et al., 2022; Rezaei-Shoshtari et al., 2022; Beck et al.,
2023). While these works mainly utilize the expressive power of HNs to improve
task performance, we focus more on utilizing HN’s knowledge decoupling ability
to improve inference efficiency.

#### Policy Distillation

Distillation (Buciluǎ et al., 2006; Hinton et al., 2015) transfers knowledge
from a teacher model or multiple teacher models to a student model. In the
context of RL, it has been used to compress the policy network (Rusu et al.,
2015), accelerate learning on a new task (Parisotto et al., 2015), and
facilitate policy learning (Lee et al., 2020; Chen et al., 2022; Wan et al.,
2023). However, less attention is paid to measuring and minimizing the
generalization gap between the teacher and the student. Igl et al. (2020)
distill a teacher policy into a student with identical architecture, which is
used as an intermediate step to improve task generalization of a policy
trained via RL. Chen et al. (2022) and Wan et al. (2023) investigate
generalization of a distilled policy to unseen objects (tasks) in robotic
manipulation. However, as the teacher has access to privileged information
while the student does not under their setting, there is always a performance
gap between the teacher and the student, and this gap is larger on the test
tasks than that on the training tasks, which indicates a generalization gap.
Most similar to our work is Furuta et al. (2022), which distill multiple
single-robot MLP teachers into a TF student for universal morphology control.
But they only show the advantage of using TF as the student architecture, but
do not compare with a universal TF teacher to measure generalization gap.
Furthermore, the TF student policy they learn still has efficiency issue
during deployment. To the best of our knowledge, we are the first to
systematically investigate how to mitigate generalization gap in policy
distillation via proper algorithmic choices.

##  7 Conclusion

In this paper, we propose HyperDistill to achieve both good performance and
high efficiency at inference time for universal morphology control. While
training HyperDistill via RL is hard, we take a policy distillation approach,
and systematically investigate how some key algorithmic choices influence task
generalization of the student policy. HyperDistill matches the performance of
the universal TF teacher on both training and test robots, while significantly
reducing memory and computational cost, which supports our hypothesis that
decoupling inter-task and intra-task knowledge can improve inference
efficiency. Given the wide application of TF in foundation models and their
extremely high inference cost (Achiam et al., 2023; Brohan et al., 2022, 2023;
Padalkar et al., 2023), the findings in our paper could serve as a general
principle to reduce the inference cost of TF in other domains for potential
future work.

## Impact Statements

This paper presents work whose goal is to advance the field of Machine
Learning. There are many potential societal consequences of our work, none
which we feel must be specifically highlighted here.

Specifically, as the control policy learned by our method can be deployed on
real-world robots for tasks like locomotion and manipulation, issues like
safety in operation need to be carefully considered, whose impacts have been
extensively discussed in prior work.

## References

  * Achiam et al. (2023) Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_ , 2023.
  * Andrychowicz et al. (2020) Andrychowicz, O. M., Baker, B., Chociej, M., Jozefowicz, R., McGrew, B., Pachocki, J., Petron, A., Plappert, M., Powell, G., Ray, A., et al.  Learning dexterous in-hand manipulation.  _The International Journal of Robotics Research_ , 39(1):3–20, 2020.
  * Barth & Poulet (2012) Barth, A. L. and Poulet, J. F.  Experimental evidence for sparse firing in the neocortex.  _Trends in neurosciences_ , 35(6):345–355, 2012.
  * Beck et al. (2022) Beck, J., Jackson, M. T., Vuorio, R., and Whiteson, S.  Hypernetworks in meta-reinforcement learning.  In _6th Annual Conference on Robot Learning_ , 2022.
  * Beck et al. (2023) Beck, J., Vuorio, R., Xiong, Z., and Whiteson, S.  Recurrent hypernetworks are surprisingly strong in meta-rl.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.
  * Brohan et al. (2022) Brohan, A., Brown, N., Carbajal, J., Chebotar, Y., Dabis, J., Finn, C., Gopalakrishnan, K., Hausman, K., Herzog, A., Hsu, J., et al.  Rt-1: Robotics transformer for real-world control at scale.  _arXiv preprint arXiv:2212.06817_ , 2022.
  * Brohan et al. (2023) Brohan, A., Brown, N., Carbajal, J., Chebotar, Y., Chen, X., Choromanski, K., Ding, T., Driess, D., Dubey, A., Finn, C., et al.  Rt-2: Vision-language-action models transfer web knowledge to robotic control.  _arXiv preprint arXiv:2307.15818_ , 2023.
  * Buciluǎ et al. (2006) Buciluǎ, C., Caruana, R., and Niculescu-Mizil, A.  Model compression.  In _Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining_ , pp. 535–541, 2006.
  * Chen et al. (2020) Chen, D., Lin, Y., Li, W., Li, P., Zhou, J., and Sun, X.  Measuring and relieving the over-smoothing problem for graph neural networks from the topological view.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 34, pp. 3438–3445, 2020.
  * Chen et al. (2023) Chen, R., Han, J., Sun, F., and Huang, W.  Subequivariant graph reinforcement learning in 3d environments.  _arXiv preprint arXiv:2305.18951_ , 2023.
  * Chen et al. (2022) Chen, T., Xu, J., and Agrawal, P.  A system for general in-hand object re-orientation.  In _Conference on Robot Learning_ , pp. 297–307. PMLR, 2022.
  * Clavera et al. (2019) Clavera, I., Nagabandi, A., Liu, S., Fearing, R. S., Abbeel, P., Levine, S., and Finn, C.  Learning to adapt in dynamic, real-world environments through meta-reinforcement learning.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=HyztsoC5Y7.
  * Cobbe et al. (2019) Cobbe, K., Klimov, O., Hesse, C., Kim, T., and Schulman, J.  Quantifying generalization in reinforcement learning.  In _International Conference on Machine Learning_ , pp. 1282–1289. PMLR, 2019.
  * Czarnecki et al. (2019) Czarnecki, W. M., Pascanu, R., Osindero, S., Jayakumar, S., Swirszcz, G., and Jaderberg, M.  Distilling policy distillation.  In _The 22nd international conference on artificial intelligence and statistics_ , pp. 1331–1340. PMLR, 2019.
  * Dennis et al. (2020) Dennis, M., Jaques, N., Vinitsky, E., Bayen, A., Russell, S., Critch, A., and Levine, S.  Emergent complexity and zero-shot transfer via unsupervised environment design.  _Advances in neural information processing systems_ , 33:13049–13061, 2020.
  * Dong et al. (2022) Dong, H., Wang, T., Liu, J., and Zhang, C.  Low-rank modular reinforcement learning via muscle synergy.  _Advances in Neural Information Processing Systems_ , 35:19861–19873, 2022.
  * Feng et al. (2022) Feng, G., Zhang, H., Li, Z., Peng, X. B., Basireddy, B., Yue, L., SONG, Z., Yang, L., Liu, Y., Sreenath, K., et al.  Genloco: Generalized locomotion controllers for quadrupedal robots.  In _6th Annual Conference on Robot Learning_ , 2022.
  * Furuta et al. (2022) Furuta, H., Iwasawa, Y., Matsuo, Y., and Gu, S. S.  A system for morphology-task generalization via unified representation and behavior distillation.  _arXiv preprint arXiv:2211.14296_ , 2022.
  * Galanti & Wolf (2020) Galanti, T. and Wolf, L.  On the modularity of hypernetworks.  _Advances in Neural Information Processing Systems_ , 33:10409–10419, 2020.
  * Ghadirzadeh et al. (2021) Ghadirzadeh, A., Chen, X., Poklukar, P., Finn, C., Björkman, M., and Kragic, D.  Bayesian meta-learning for few-shot policy adaptation across robotic platforms.  In _2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_ , pp. 1274–1280. IEEE, 2021.
  * Gupta et al. (2021) Gupta, A., Savarese, S., Ganguli, S., and Fei-Fei, L.  Embodied intelligence via learning and evolution.  _Nature communications_ , 12(1):1–12, 2021.
  * Gupta et al. (2022) Gupta, A., Fan, L., Ganguli, S., and Fei-Fei, L.  Metamorph: Learning universal controllers with transformers.  _arXiv preprint arXiv:2203.11931_ , 2022.
  * Ha et al. (2016) Ha, D., Dai, A., and Le, Q. V.  Hypernetworks.  _arXiv preprint arXiv:1609.09106_ , 2016.
  * Hallak et al. (2015) Hallak, A., Di Castro, D., and Mannor, S.  Contextual markov decision processes.  _arXiv preprint arXiv:1502.02259_ , 2015.
  * Hao et al. (2023) Hao, Z., Guo, J., Han, K., Tang, Y., Hu, H., Wang, Y., and Xu, C.  One-for-all: Bridge the gap between heterogeneous architectures in knowledge distillation.  _arXiv preprint arXiv:2310.19444_ , 2023.
  * Hinton et al. (2015) Hinton, G., Vinyals, O., and Dean, J.  Distilling the knowledge in a neural network.  2015\.
  * Hobbhahn & Sevilla (2021) Hobbhahn, M. and Sevilla, J.  What’s the backward-forward flop ratio for neural networks?, 2021.  URL https://epochai.org/blog/backward-forward-FLOP-ratio.  Accessed: 2024-02-01.
  * Hong et al. (2022) Hong, S., Yoon, D., and Kim, K.-E.  Structure-aware transformer policy for inhomogeneous multi-task reinforcement learning.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=fy_XRVHqly.
  * Huang et al. (2020) Huang, W., Mordatch, I., and Pathak, D.  One policy to control them all: Shared modular policies for agent-agnostic control.  In _International Conference on Machine Learning_ , pp. 4455–4464. PMLR, 2020.
  * Hutter et al. (2016) Hutter, M., Gehring, C., Jud, D., Lauber, A., Bellicoso, C. D., Tsounis, V., Hwangbo, J., Bodie, K., Fankhauser, P., Bloesch, M., et al.  Anymal-a highly mobile and dynamic quadrupedal robot.  In _2016 IEEE/RSJ international conference on intelligent robots and systems (IROS)_ , pp. 38–44. IEEE, 2016.
  * Igl et al. (2020) Igl, M., Farquhar, G., Luketina, J., Boehmer, W., and Whiteson, S.  Transient non-stationarity and generalisation in deep reinforcement learning.  _arXiv preprint arXiv:2006.05826_ , 2020.
  * Kalashnikov et al. (2018) Kalashnikov, D., Irpan, A., Pastor, P., Ibarz, J., Herzog, A., Jang, E., Quillen, D., Holly, E., Kalakrishnan, M., Vanhoucke, V., et al.  Scalable deep reinforcement learning for vision-based robotic manipulation.  In _Conference on Robot Learning_ , pp. 651–673. PMLR, 2018.
  * Kurin et al. (2020) Kurin, V., Igl, M., Rocktäschel, T., Boehmer, W., and Whiteson, S.  My body is a cage: the role of morphology in graph-based incompatible control.  _arXiv preprint arXiv:2010.01856_ , 2020.
  * Leal et al. (2023) Leal, I., Choromanski, K., Jain, D., Dubey, A., Varley, J., Ryoo, M., Lu, Y., Liu, F., Sindhwani, V., Vuong, Q., et al.  Sara-rt: Scaling up robotics transformers with self-adaptive robust attention.  _arXiv preprint arXiv:2312.01990_ , 2023.
  * Lee et al. (2020) Lee, J., Hwangbo, J., Wellhausen, L., Koltun, V., and Hutter, M.  Learning quadrupedal locomotion over challenging terrain.  _Science robotics_ , 5(47):eabc5986, 2020.
  * Levine et al. (2016) Levine, S., Finn, C., Darrell, T., and Abbeel, P.  End-to-end training of deep visuomotor policies.  _The Journal of Machine Learning Research_ , 17(1):1334–1373, 2016.
  * Narvekar et al. (2020) Narvekar, S., Peng, B., Leonetti, M., Sinapov, J., Taylor, M. E., and Stone, P.  Curriculum learning for reinforcement learning domains: A framework and survey.  _The Journal of Machine Learning Research_ , 21(1):7382–7431, 2020.
  * Padalkar et al. (2023) Padalkar, A., Pooley, A., Jain, A., Bewley, A., Herzog, A., Irpan, A., Khazatsky, A., Rai, A., Singh, A., Brohan, A., et al.  Open x-embodiment: Robotic learning datasets and rt-x models.  _arXiv preprint arXiv:2310.08864_ , 2023.
  * Parisotto et al. (2015) Parisotto, E., Ba, J. L., and Salakhutdinov, R.  Actor-mimic: Deep multitask and transfer reinforcement learning.  _arXiv preprint arXiv:1511.06342_ , 2015.
  * Pathak et al. (2019) Pathak, D., Lu, C., Darrell, T., Isola, P., and Efros, A. A.  Learning to control self-assembling morphologies: a study of generalization via modularity.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Peng et al. (2021) Peng, M., Zhu, B., and Jiao, J.  Linear representation meta-reinforcement learning for instant adaptation.  _arXiv preprint arXiv:2101.04750_ , 2021.
  * Peng et al. (2018) Peng, X. B., Andrychowicz, M., Zaremba, W., and Abbeel, P.  Sim-to-real transfer of robotic control with dynamics randomization.  In _2018 IEEE international conference on robotics and automation (ICRA)_ , pp. 3803–3810. IEEE, 2018.
  * Rezaei-Shoshtari et al. (2022) Rezaei-Shoshtari, S., Morissette, C., Hogan, F. R., Dudek, G., and Meger, D.  Hypernetworks for zero-shot transfer in reinforcement learning.  _arXiv preprint arXiv:2211.15457_ , 2022.
  * Riquelme et al. (2021) Riquelme, C., Puigcerver, J., Mustafa, B., Neumann, M., Jenatton, R., Susano Pinto, A., Keysers, D., and Houlsby, N.  Scaling vision with sparse mixture of experts.  _Advances in Neural Information Processing Systems_ , 34:8583–8595, 2021.
  * Rusu et al. (2015) Rusu, A. A., Colmenarejo, S. G., Gulcehre, C., Desjardins, G., Kirkpatrick, J., Pascanu, R., Mnih, V., Kavukcuoglu, K., and Hadsell, R.  Policy distillation.  _arXiv preprint arXiv:1511.06295_ , 2015.
  * Sarafian et al. (2021) Sarafian, E., Keynan, S., and Kraus, S.  Recomposing the reinforcement learning building blocks with hypernetworks.  In _International Conference on Machine Learning_ , pp. 9301–9312. PMLR, 2021.
  * Shen et al. (2023) Shen, S., Yao, Z., Li, C., Darrell, T., Keutzer, K., and He, Y.  Scaling vision-language models with sparse mixture of experts.  _arXiv preprint arXiv:2303.07226_ , 2023.
  * Tobin et al. (2017) Tobin, J., Fong, R., Ray, A., Schneider, J., Zaremba, W., and Abbeel, P.  Domain randomization for transferring deep neural networks from simulation to the real world.  In _2017 IEEE/RSJ international conference on intelligent robots and systems (IROS)_ , pp. 23–30. IEEE, 2017.
  * Todorov et al. (2012) Todorov, E., Erez, T., and Tassa, Y.  Mujoco: A physics engine for model-based control.  In _2012 IEEE/RSJ international conference on intelligent robots and systems_ , pp. 5026–5033. IEEE, 2012.
  * Trabucco et al. (2022) Trabucco, B., Phielipp, M., and Berseth, G.  Anymorph: Learning transferable polices by inferring agent morphology.  In _International Conference on Machine Learning_ , pp. 21677–21691. PMLR, 2022.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Vithayathil Varghese & Mahmoud (2020) Vithayathil Varghese, N. and Mahmoud, Q. H.  A survey of multi-task deep reinforcement learning.  _Electronics_ , 9(9):1363, 2020.
  * Wan et al. (2023) Wan, W., Geng, H., Liu, Y., Shan, Z., Yang, Y., Yi, L., and Wang, H.  Unidexgrasp++: Improving dexterous grasping policy learning via geometry-aware curriculum and iterative generalist-specialist learning.  _arXiv preprint arXiv:2304.00464_ , 2023.
  * Wang et al. (2018) Wang, T., Liao, R., Ba, J., and Fidler, S.  Nervenet: Learning structured policy with graph neural networks.  In _International conference on learning representations_ , 2018.
  * Wu et al. (2020) Wu, Z., Pan, S., Chen, F., Long, G., Zhang, C., and Philip, S. Y.  A comprehensive survey on graph neural networks.  _IEEE transactions on neural networks and learning systems_ , 32(1):4–24, 2020.
  * Xiong et al. (2023) Xiong, Z., Beck, J., and Whiteson, S.  Universal morphology control via contextual modulation.  _arXiv preprint arXiv:2302.11070_ , 2023.
  * Yu et al. (2019) Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., and Finn, C.  Multi-task reinforcement learning without interference.  In _Proc. Optim. Found. Reinforcement Learn. Workshop NeurIPS_ , 2019.
  * Zhao et al. (2021) Zhao, Y., Gao, Y., Sun, Q., Tian, Y., Mao, L., and Gao, F.  A real-time low-computation cost human-following framework in outdoor environment for legged robots.  _Robotics and Autonomous Systems_ , 146:103899, 2021.

##  Appendix A Morphology Context Features and Transformations

For each robot, its morphology context includes the topology graph of the
robot and limb-wise context features. Limb-wise context features include: (1)
The initially relative position of the limb w.r.t. its parent node; (2) The
initially geometric orientation of the limb w.r.t. its parent node; (3) The
mass and shape parameters of the limb; (4) Parameters about the joints that
connect the limb to its parent node, including joint type, joint range and
axis, and motor gear (Gupta et al., 2022).

As discussed in Section 3.2, the context features need to be distinguishable
enough to learn a good morphology-conditioned HN. Among the original context
features, we find that the relative position feature does not provide
sufficient discrimination, as different limbs, especially symmetric ones
within a robot, can have the same relative position to their parents, but
actually locate in different parts of the robot and play different roles. To
solve this issue, we transform relative position features into absolute
position features to better distinguish different limbs’ positions in the
morphology, which can be easily computed by following the path from the torso
(the root of the morphology tree) to each limb node. Experimental results in
Figure 7 validate the importance of this feature transformation.

##  Appendix B Further Experimental Setup

###  B.1 Generation of PD Robots

We mutate the 100 training robots in the UNIMAL benchmark to get an augmented
set of 1000 PD robots, on which we collect expert data for policy
distillation. Specifically, for each training robot, we first uniformly sample
a number $m$ from $\left[1,2,3\right]$, then sequentially apply $m$ mutation
steps to the robot to get a new morphology. See Gupta et al. (2021) for the
set of feasible mutations allowed in the UNIMAL benchmark. We generate 9
variants for each training robot, yielding an augmented set of 1000 PD robots
in total.

###  B.2 FLOPs Computation

We compute the FLOPs of a linear layer with input dimension $M$ and output
dimension $N$ as $2\times M\times N$ (Hobbhahn & Sevilla, 2021). Then the
FLOPs of an MLP or a TF can be analytically computed by recursively
decomposing its FLOPs into the summation of basic linear layers’ FLOPs. We
omit the FLOPs of other operations like activation functions, layer
normalization and etc., as these operations only have linear complexity, which
is negligible compared to the quadratic complexity of linear layers.

###  B.3 Architecture Details of Different Methods

Table 2 shows the size of different architectures in each environment. For the
TF teacher and TF (oracle), we use the same architecture hyperparameters as in
Xiong et al. (2023). For the base MLP in HyperDistill, we fix the hidden layer
size to 256, and do a grid search over hidden layer number to find the
smallest base MLP that can achieve on-par performance with the teacher policy
in each environment. Then we set MLP’s architecture identical to the base MLP
in HyperDistill. For TF (compressed), as there is a two-layer MLP after each
attention module in the standard TF implementation (Vaswani et al., 2017), we
first reduce the dimension of this MLP’s hidden layer to 128, and reduce the
number of attention heads to 1, then do a grid search over the number of
attention layers to get a TF with a similar number of parameters as
HyperDistill’s base MLP. The token embedding dimension is 128 for all
different TF architectures.

Environment | Base MLP in HyperDistill | TF teacher | TF (compressed)
---|---|---|---
| & MLP | & TF (oracle) |  |  |
| Layer num. | Hidden size | Layer num. | Attention head | Linear hidden size | Layer num. | Attention head | Linear hidden size
FT | 2 | 256 | 5 | 2 | 1024 | 1 | 1 | 128
VT | 3 | 256 | 5 | 2 | 1024 | 2 | 1 | 128
Obstacle | 3 | 256 | 5 | 2 | 1024 | 2 | 1 | 128

Table 2: The size of different models in each environment.

##  Appendix C Further Experimental Results and Analysis

###  C.1 Main Results on the Training Robots

Figure 6 shows how different methods perform on the training robots during the
distillation process. See Section 4.2 for analysis on the main results.

Figure 6: The performance of different methods on the training robots in each
environment.

###  C.2 Further Analysis on Inference Efficiency

As shown in Table 1, HyperDistill’s efficiency advantage is more significant
in the FT environment, as in the other two environments, there is an
additional high-dimensional terrain information input to the policy, which
needs to be processed by a large MLP encoder. This adds a large constant to
the model size and FLOPs of all methods, and thus reduces the efficiency ratio
of HyperDistill to the other methods.

Although MLP and the base network of HyperDistill have the same number of
hidden layers and hidden units, the model size and FLOPs of MLP are still a
bit larger than those of HyperDistill, as it needs to condition on the
morphology context by concatenating limb-wise context features to the policy
input, which introduces some additional cost.

###  C.3 Ablation Study on Context Representation Learning

Figure 7: How context representation learning in HyperDistill influences
generalization performance in FT. Left: context feature transformation. Right:
architecture of the context encoder.

#### Context Feature Transformation

As shown in Figure 7 (left), conducting feature transformations to improve
feature discrimination plays an important role in improving performance.

#### Architecture of Context Encoder

As shown in Figure 7 (right), for different choices of the context encoder,
GNN performs the worst, possibly due to the over-smoothing issue of GNN (Chen
et al., 2020) which harms discrimination of the context embedding. While the
MLP context encoder does not consider node interaction, it still achieves
quite good performance in FT. We thus further compare MLP and TF context
encoder in the other two more challenging environments (Figure 8), and find
that the TF context encoder performs better, which validates the benefits of
modeling node interaction for context representation learning.

Figure 8: How TF and MLP context encoder influence HyperDistill’s
generalization performance. Left: VT; Right: Obstacle.


# Accelerating Gradient-based Meta Learner

Varad Pimpalkhute  IIIT NagpurIndia<EMAIL_ADDRESS>and  Amey
Pandit, Mayank Mishra, Rekha Singhal  TCS ResearchIndia<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
(2018)

###### Abstract.

Meta Learning has been in focus in recent years due to the meta-learner
model’s ability to adapt well and generalize to new tasks, thus, reducing both
the time and data requirements for learning. However, a major drawback of meta
learner is that, to reach to a state from where learning new tasks becomes
feasible with less data, it requires a large number of iterations and a lot of
time. We address this issue by proposing various acceleration techniques to
speed up meta learning algorithms such as MAML (Model Agnostic Meta Learning).
We present 3.73X acceleration on a well known RNN optimizer based meta learner
proposed in literature (Wichrowska et al., 2017). We introduce a novel method
of training tasks in clusters, which not only accelerates the meta learning
process but also improves model accuracy performance.

Meta learning, RNN optimizer, AGI, Performance optimization

††copyright: none††copyright: acmcopyright††journalyear: 2018††doi:
10.1145/1122445.1122456††conference: Woodstock ’18: ACM Symposium on Neural
Gaze Detection; June 03–05, 2018; Woodstock, NY††booktitle: Woodstock ’18: ACM
Symposium on Neural Gaze Detection, June 03–05, 2018, Woodstock, NY††price:
15.00††isbn: 978-1-4503-XXXX-X/18/06

##  1\. Introduction

Artificial general intelligence (AGI) is a hypothetical concept where machines
are capable of learning and thinking like humans. Meta learning also known as
learning to learn, brings us one step closer to achieve AGI by learning on a
task in a manner similar to humans. To give an analogy, in Meta-learning
paradigm, a distribution of related tasks (say racket sports: Badminton,
Tennis etc.) are used to train a model (human) which can use this ”experience”
to learn any new task (Squash) with lesser amount of data (lesser practice of
Squash than a person who doesn’t know any racket sport) and in lesser time.

The meta learning algorithms, like MAML (Finn et al., 2017) (and others (Raghu
et al., 2020; Wichrowska et al., 2017; Snell et al., 2017)), are time and
compute intensive. These algorithms require a lot of iterations over the
available tasks (60000 (Finn et al., 2017)) to make a model converge and be
ready to learn a new task quickly with good accuracy.

In this paper, we present our ongoing work involving several performance
optimization techniques to reduce the training time of metalearning
algorithms. We will discuss these techniques in the context of acceleration of
a gradient based meta learner proposed in (Wichrowska et al., 2017). A
gradient based meta learner learns how gradients change during gradient
descent based optimizations of various learning tasks (learning task involves
training a model). The goal is to enable rapid model parameter updates (faster
than optimizers like SGD, ADAM, RMSProp) after each learning iteration
111There are other approaches of meta learning too, namely, metric based and
loss based meta learners (Hospedales et al., 2020)..

The gradient based meta learner which is our focus in this paper employs a
Hierarchical RNN 222RNN: Recurrent Neural Network architecture to ensembles
the loss landscapes of various small tasks. Figure 1 shows how this meta
learner learnt from a single task’s training process (the same is repeated for
every task). There are two learning loops: a) Task loop \- which updates the
task parameters for every batch of data using hierarchical RNN as optimizer;
and b) Meta loop \- which updates the parameters of RNN optimizer itself using
loss accumulated across all batches of task learning using a RMSProp
optimizer. This process is repeated multiple times (N epochs) for each task
and sequentially for every task available in a single meta iteration. Large
number of such meta iterations are required to let Hierarchical RNN based meta
optimizer converge.

Figure 1. Meta-training on a single task using a Hierarchical RNN-based
optimizer. Figure 2. Earlier Proposed approach in the paper

Figure 2 shows the sequential nature of the meta training approach and the
iterations involved using a flow-chart. We improved on the model performance
by introducing a better and wider meta-optimization pipeline, which
incorporates the optimizer parameters being learned in form of task clusters.
We achieved a speedup of more than 2x with our proposed approach and overall
speedup of 3.73x by also incorporating several coding related optimizations.

The rest of the paper is structured as follows: We briefly discuss the
previous related work in the field of meta learning and accelerations in
Section 2. We revisit definition of tasks in Section 3. We propose the
approaches for accelerating Hierarchical RNN in Section 4. We verify and
discuss the effectiveness of our experiments in Section 5. Finally, we discuss
the ongoing work in Section 6 and conclude the paper in Section 7.

##  2\. Related Work

Gradient based meta-learners proposed in literature (Raghu et al., 2020;
Mishra et al., 2017; Lee et al., 2019; Franceschi et al., 2018) have worked on
improving model accuracy. However, very little work has been done on improving
the training time of the meta learner, where meta-learner especially MAML is a
time hungry algorithm.

Gradient descent on the meta-parameters ($w$) involves computing second-order
derivatives or a simplified and cheaper first order implementation (Finn et
al., 2017; Ravi and Larochelle, 2017; Nichol et al., 2018). (Ravi and
Larochelle, 2017) proposes ANIL which makes use of feature reuse for few-shot
learning. ANIL simplifies on MAML by using almost no inner loop; thus being
more computationally efficient than MAML, getting a speed up as high as 1.8x.
(Nichol et al., 2018) focuses on decreasing the computations and memory by not
unrolling a computational graph or calculating any second derivatives.
(Bollenbacher et al., 2020) propose a distributed framework to accelerate the
learning process of MAML.

In this work, we also aim to accelerate the learning process by proposing
various optimization techniques. Contrary to the above approaches (which
involve changing the mathematical model or framework), our approach is more
intuitive and less complex.

##  3\. Tasks Preparation for meta learning

The Hierarchical RNN based meta learner work proposed in (Wichrowska et al.,
2017) which we intend to accelerate is designed to be an optimizer with rapid
updates. Consequently, it becomes important that the optimizer should work for
any model, may it be CNN, Fully connected, RNN, logistic regression, etc.
Thus, optimizer is trained to give rapid updates by training it on various
instances of these models. The instance of a model is considered as task. The
diversity of tasks helps in generalizing the model well to new tasks. Datasets
to train the task (model) are generated by following the approach mentioned in
paper.

##  4\. Proposed Approaches for Acceleration

The approaches we tried for speeding up Hierarchical RNN based meta learner
can be divided into two categories A) Algorithmic optimizations, and B)
Program optimizations, as shown in Figure 3.

Figure 3. Performance optimizations experimented on RNN optimizer based meta
learner

We focus on discussions of Algorithmic optimizations performed in this
section, while limiting the discussion around program optimizations tried in
experiments section as they are mostly straightforward and simple to
understand.

###  4.1. Parallelization inside Task Loop

The task-loop shown in Figures 1 and 2 involved training over multiple batches
of data sequentially. Instead of training over data batches sequentially we
train over multiple batches in parallel. We calculate each batch’s loss and
parameter gradients and then calculate average of individual parameter’s
gradients. Finally, updating the task parameters using these averaged
parameter gradients. The updated model is used in the next iteration. The
degree of parallelism is configurable and depends on the underlying hardware.

###  4.2. Considering tasks in clusters and training clusters parallely

Although, meta learning approaches employ similar tasks for training the meta
learner to begin with, some tasks in the task set are more similar to each
other than other tasks. For example, a softmax regression task will be more
similar to other softmax regression tasks in terms of how gradients get
updated in each optimisation step than task which involves training a fully
connected neural network. The Hierarchical RNN based meta optimizer approach
doesn’t differentiate between such tasks and while meta training the
”sequence” of tasks is randomly chosen.

We grouped the tasks into clusters333There are various ways to create clusters
of tasks. Thus for simplicity, we consider creation of clusters as a black box
which returns k clusters, each cluster having similar tasks. depending on
their similarity with each other and then we performed meta training across
these clusters parallely. The Hierarchical RNN optimizer’s parameters were
updated by using the averaged parameter gradients across these clusters. This
iteration of meta training parallely over the clusters is repeated for a
number of times (number of meta-iterations is user defined). To properly
utilise the underlying hardware and ensuring that no one branch of
parallelization becomes bottleneck and takes too long we group clusters into
cluster groups and schedule cluster groups as shown in Figure 4.

The approach is that all cluster-groups should take similar time for meta
learning iteration. Figure 4 shows 5 task clusters containing similar tasks
namely Softmax Regression, Quadratic problems, 2D problems, Bowl problems, and
Fully connected problems are scheduled as two cluster groups.

Figure 4. Optimal grouping of task clusters as both the branches are of equal
length.

Figure 5 gives the complete picture of the algorithmic optimization we
performed in form of a flow chart. The original sequential pipeline of meta
training mentioned in flowchart of Figure 2 has been broadened by introducing
parallelizations in task loop (box 8 in flowchart of Figure 5) as well as
across training task-cluster (green and yellow region). The red triangles
indicate the crucial steps where we have placed our optimizations. To
summarize our acceleration approach we train all clusters parallely while
learning on the tasks in each cluster sequentially. Rest of the procedure is
similar to that followed in (Wichrowska et al., 2017).

Figure 5. Our Proposed overall approach

##  5\. Results & Discussion

In this section, we perform experiments to verify the effectiveness of
proposed acceleration techniques. The hardware and software configuration
employed for experiments is mentioned in Table 1. Before performing
experiments we upgraded the baseline code available from prior work
(Wichrowska et al., 2017) to newer libraries.

We performed experiments in a resource constrained setup as can be seen in
Table 1. We believe that a higher configuration setup will exhibit better
acceleration and performance gains due to possibility of higher parallelism.
All the experiments are performed while ensuring that the achieved test
accuracy of 0.65 for test ”task-set” (consisting of softmax regression tasks)
on base model is matched using our proposed approach. We now present details
of experiments performed for each of proposed acceleration technique.

Table 1. Experiment Setup Type of tasks |  Quadractic, Bowl, Fully Connected, Softmax Regression, 2D Problems
---|---
Hardware Setup |  Intel core i5 CPU @2.50GHz, 8GB RAM, 2 physical and 4-logical cores.
Software |  Windows 10 (64-bit) OS, Cython, Tensorflow, Python, Numpy, Pandas, Scikit

###  5.1. Acceleration achieved due to parallelism introduced

Table 2 lists the reduction in overall execution time due to parallelizations
introduced. We varied the number of task clusters used for meta training and
always created 2 cluster groups to keep degree of cluster training
parallelization fixed at 2. This was done due to hardware limitations of our
setup. However, the improvement in terms of time saving is encouraging. The
experiments were speeded by more than 2x (base time vs optimized time) while
the accuracy achieved on test task set remained same.

Table 2. Clustering of tasks for Meta Training. Num Clusters |  Base Time (in sec) |  Optimized Time (in sec)
---|---|---
2 |  1450.15 |  826.53
3 |  6313.8 |  2977.83
4 |  11361.3 |  4918.31
5 |  21932.75 |  12473.1

We also introduced program and code level changes mentioned in Figure 3. Next
we discuss what changes were made and how they improved the performance.

###  5.2. Program and Code Optimizations

There are several straight forward optimization like upgrading the libraries
and platforms used in the available code. For example, upgrading python itself
resulted in 15% speedup. However, to our surprise upgrading Tensorflow Library
didn’t result in any speedup.

We also used vector instructions inside of loops using Numpy arrays and
removed redundant code by fusing some of the operations performed, put
together these Code Optimizations resulted in 1.1X performance improvement.
Cythoning the code (converting Python code to C) shows the most promising
results with an average speed up of 1.33x. The impact of these individual
optimizations is presented as graph in Figure 6. We varied the number of tasks
from 5 to 35 as shown in X-axis. This variation resulted in different number
of task clusters being formed with different number of tasks in them, thus,
covering wider range of scenarios.

Figure 6. Comparison of trend in speed up for optimization techniques.

Finally, we combined all the optimizations to get a speed up of 3.73x. Figure
7 presents the trend in variation in speed up as the number of tasks being
meta trained on are increased. Initially, with lesser number of tasks, the
performance curve shows improvement in performance reaching upto 4x and later
stabilizing at 3.73x.

Figure 7. Trend in overall speed up as number of tasks meta trained on are
increased.

##  6\. on-going and future work

We are experimenting with different ways and granularities in which the task
clusters can be created. We believe grouping similar tasks together in
clusters and training meta learner over them will result in faster convergence
than without clustering. This is because the the gradient updates for tasks in
same cluster will be more similar to each other.

We are also experimenting whether sequence in which tasks-cluster are used for
meta-learning has impact on convergence time of meta learning. A simple
analogy from racket sports is that a person good at Racket-ball will find it
easier to learn Squash as compared to Tennis 444In Squash, Racket-ball the
ball bounces back from the wall and both players are on same court, whereas,
in Tennis the courts are different..

We also experimented replacing GRU cells used in Hierarchical RNN optitimizer
with LSTM in hope of achiving convergence with lesser number of meta
iterations. Although, the convergence happened in lesser number of iterations
than when GRU cells were used, the time per iteration increased considerably,
thus, nullifying the gains. We shelved this idea for future.

##  7\. conclusions

We presented our on-going work involving acceleration of a Hierarchical RNN
based meta learner proposed in (Wichrowska et al., 2017). We changed the
sequential meta learning pipeline discussed in paper to a parallel pipeline by
introducing parallizations inside training of the tasks as well as across
clusters of tasks. We also introduced grouping of tasks clusters for better
scheduling.

Apart from introducing wider metalearning pipeline we also introduced certain
program and code level optimizations. Put together the proposed techniques
resulted in a 3.73X speed-up of the meta-learning algorithm.

## References

  * (1)
  * Bollenbacher et al. (2020) Jan Bollenbacher, Florian Soulier, Beate Rhein, and Laurenz Wiskott. 2020.  Investigating Parallelization of MAML. In _Discovery Science_ , Annalisa Appice, Grigorios Tsoumakas, Yannis Manolopoulos, and Stan Matwin (Eds.). Springer International Publishing, Cham, 294–306.
  * Finn et al. (2017) Chelsea Finn, Pieter Abbeel, and Sergey Levine. 2017\.  Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks.  _CoRR_ abs/1703.03400 (2017).  arXiv:1703.03400 http://arxiv.org/abs/1703.03400
  * Franceschi et al. (2018) L. Franceschi, P. Frasconi, Saverio Salzo, Riccardo Grazzi, and M. Pontil. 2018.  Bilevel Programming for Hyperparameter Optimization and Meta-Learning.  _ArXiv_ abs/1806.04910 (2018).
  * Hospedales et al. (2020) Timothy Hospedales, Antreas Antoniou, Paul Micaelli, and Amos Storkey. 2020.  Meta-learning in neural networks: A survey.  _arXiv preprint arXiv:2004.05439_ (2020).
  * Lee et al. (2019) Kwonjoon Lee, Subhransu Maji, Avinash Ravichandran, and Stefano Soatto. 2019.  Meta-Learning With Differentiable Convex Optimization. In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_.
  * Mishra et al. (2017) Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. 2017.  Meta-Learning with Temporal Convolutions.  _CoRR_ abs/1707.03141 (2017).  arXiv:1707.03141 http://arxiv.org/abs/1707.03141
  * Nichol et al. (2018) Alex Nichol, Joshua Achiam, and John Schulman. 2018\.  On First-Order Meta-Learning Algorithms.  _CoRR_ abs/1803.02999 (2018).  arXiv:1803.02999 http://arxiv.org/abs/1803.02999
  * Raghu et al. (2020) Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. 2020.  Rapid Learning or Feature Reuse? Towards Understanding the Effectiveness of MAML. In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net.  https://openreview.net/forum?id=rkgMkCEtPB
  * Ravi and Larochelle (2017) S. Ravi and H. Larochelle. 2017.  Optimization as a Model for Few-Shot Learning. In _International Conference on Learning Representations_.
  * Snell et al. (2017) Jake Snell, Kevin Swersky, and Richard S. Zemel. 2017\.  Prototypical Networks for Few-shot Learning.  _CoRR_ abs/1703.05175 (2017).  arXiv:1703.05175 http://arxiv.org/abs/1703.05175
  * Wichrowska et al. (2017) Olga Wichrowska, Niru Maheswaranathan, Matthew W. Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando de Freitas, and Jascha Sohl-Dickstein. 2017.  Learned Optimizers that Scale and Generalize. In _Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017_ _(Proceedings of Machine Learning Research, Vol. 70)_ , Doina Precup and Yee Whye Teh (Eds.). PMLR, 3751–3760.  http://proceedings.mlr.press/v70/wichrowska17a.html